You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[Feat][Kernel] Fused CE LogProb without materializing logits + fused backward #155
Avoid landing large [B, S, V] logits for large-vocabulary models, so the saved memory can be used for larger batches or longer CoT responses. The forward path streams vocab blocks with online softmax; the backward path recomputes tiles to trade compute for memory instead of storing full logits/probabilities.
Implementation
CUDA: use SM90 TMA/WGMMA-oriented streaming where available, with SM80 cp.async/mma.sync fallback; keep online-softmax state in registers and shared memory.
Triton: provide a portable semantic baseline and tolerance test target for CUDA/ROCm native implementations.
Activity