Chirag Jain
50896ec574
Make nvcc threads configurable via environment variable (#885)
|
пре 9 месеци | |
|---|---|---|
| .. | ||
| README.md | пре 1 година | |
| interface.cpp | пре 2 година | |
| setup.py | пре 9 месеци | |
| xentropy_kernel.cu | пре 2 година | |
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.