#7244·trl

Fused Triton kernels for the loss head

Author: qgallouedecCreated Sep 16, 2026Updated Sep 17, 2026

Feature request

Ship fused Triton kernels for TRL's loss head, with the current torch implementations kept as fallback:

  1. logprob + entropy from logits (selective_log_softmax + entropy_from_logits), used by GRPO/RLOO/DPO and ~15 other trainers
  2. chunked CE returning loss + entropy + token accuracy (SFT loss_type="chunked_nll")
  3. _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_softmax 3.50 ms + entropy_from_logits 9.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 bare lm_head GEMM 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.