Fused Triton kernels for the loss head
Feature request
Ship fused Triton kernels for TRL's loss head, with the current torch implementations kept as fallback:
- logprob + entropy from logits (
selective_log_softmax+entropy_from_logits), used by GRPO/RLOO/DPO and ~15 other trainers - chunked CE returning loss + entropy + token accuracy (SFT
loss_type="chunked_nll") _ChunkedLogProbFunction(KTO, async GRPO)
The loss head is the only hot path TRL owns: transformers already kernelizes the model internals (RMSNorm, SwiGLU, Linear, attention) and nothing on the Hub covers what comes after the decoder.
Measured on one H100, bf16, V=151936, H=4096, 8192 tokens:
- GRPO/RLOO/DPO logprob path:
selective_log_softmax3.50 ms +entropy_from_logits9.33 ms = 12.8 ms per pass over[8, 1024, 151936], plus 2.32 GiB of logits. A 40-line fused Triton kernel (online logsumexp + gather + entropy, forward only) does the same in 0.89 ms, 14x, and lands closer to the fp32 reference than the current bf16 path. GRPO runs this 2 to 4 times per step. - SFT
chunked_nll(the default), chunk 256, fwd+bwd: 167 ms, 13.1 GiB peak. For reference, liger FLCE from the Hub is 192 ms / 6.1 GiB, the naive materialized-logits CE is 59 ms / 19.9 GiB, and the barelm_headGEMM fwd+bwd is 42 ms. A fused CE should sit around 60-80 ms at ~6 GiB, so 2.5x faster than today's default at half its memory. _ChunkedLogProbFunction: 495 ms fwd+bwd vs 58 ms for the naive path, 5.6 vs 10.6 GiB.
Reusing an existing kernel does not work: kernels-community/liger-kernels FLCE returns the loss only, while we need per-token logprobs, entropy, token accuracy, the num_items_in_batch reduction, temperature, logit_scale and final_logit_softcapping.
One thing to settle first: in-tree Triton vs a torch-universal kernel repo on the Hub (build.toml with universal = true, no compile step). The Hub gives rocm/xpu variants and no hard triton dependency, but needs network: get_kernel(repo, version=N) hits the Hub even when the snapshot is already cached, so a default-on path needs a warm-up and a fallback.
Source: huggingface/trl