Sākums Izpētīt Palīdzība
Reģistrēties Pierakstīties
david
/
flash-attention
spogulis no https://github.com/Dao-AILab/flash-attention
1
0
Atdalīts 0
Faili Problēmas 0 Vikivietne
Koks: 9b6cba16c1
Atzari Tagi
flash-attention
/
csrc
/
xentropy
Chirag Jain 50896ec574 Make nvcc threads configurable via environment variable (#885) 9 mēneši atpakaļ
..
README.md 5400fdc4ac [CE] Implement CrossEntropyLoss in Triton 1 gadu atpakaļ
interface.cpp dff68c2b22 Add smoothing for CrossEntropyParallel, rename to CrossEntropyLoss 2 gadi atpakaļ
setup.py 50896ec574 Make nvcc threads configurable via environment variable (#885) 9 mēneši atpakaļ
xentropy_kernel.cu dff68c2b22 Add smoothing for CrossEntropyParallel, rename to CrossEntropyLoss 2 gadi atpakaļ

README.md

This CUDA extension implements optimized cross-entropy loss, adapted from Apex's Xentropy. We make it work for bfloat16 and support in-place backward to save memory.

It has only been tested on A100s.

cd csrc/xentropy && pip install .

As of 2023-09-15, this extension is no longer used in the FlashAttention repo. We've instead switched to a Triton-based implementation. See the CrossEntropyLoss module for more details.