[kashif]

Summary

Fixes the softmax collapse from A Fix for Softmax and Cross-Entropy in the kernels. Once the target logit leads the logsumexp of the rest by more than ~16.6 (fp32), softmax(x_y) rounds to 1, so softmax - 1 and lse - x_y become exactly 0 while the other classes keep their gradients. The gradient no longer sums to zero, and the loss reads 0. This happens for very confident tokens, which are common in LLM training.

• CE / FLCE, CE+TVD, vocab-parallel CE: also accumulate the sum without the target in the first pass, then use dx_y = -sum_others / d and loss = softplus(log(sum_others) + m - x_y). Vocab-parallel all-reduces both sums in one call.
• GRPO: the forward stores an accurate logp instead of lse (the backward rebuilds lse from it) and uses -expm1(logp) for 1 - p_y. fused_selective_log_softmax returns the accurate logp too and drops its unused lse= arg.
• fused KL / JSD, softmax: the gradient at the argmax is taken from the other entries (paper Eq. 7). Mixed-beta JSD isn't affected.
• softmax: the multi-block kernels didn't compile (tl.float32(...)) and were unreachable. Rows wider than 65536 now go through them instead of raising.

log1p / expm1 / target_nll helpers in ops/utils.py only use tl.log / tl.exp, so no libdevice is needed.

Details

New tests build rows with the target 5 to 40 above the rest and compare against an fp64 reference written as softplus(lse_others - x_y), since plain fp64 F.cross_entropy collapses too past ~36.7. All of them fail on main.

Overhead on a 3090 Ti, bf16, H=4096, fwd+bwd: fused KL is 2-3% slower, fused JSD 1-2%. I didn't benchmark the CE / GRPO / softmax changes.

Not touched: the cutile / cute / flydsl / ascend variants, and the chunked distillation losses (torch autograd, same as torch).

Testing Done

• Hardware Type: RTX 3090 Ti (also ran on an AMD Radeon 890M / ROCm)
• run make test to ensure correctness (ran the suites of the touched ops on both GPUs: CE, FLCE, GRPO, CE+TVD, selective log softmax, vocab-parallel CE, fused KL / JSD, JSD, KL, softmax, neighborhood attention)
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[sniper35]

Some independent measurements of the near-optimum JSD gradient error this PR removes on the fused Triton path, plus two places where the same residual survives — in case they help the review. The numbers below come from an op-for-op CPU emulation of the kernel arithmetic (fp32; the same formulas as ops/jsd.py and the host-side log-softmax backward in fused_linear_jsd.py), not from running the Triton kernels.

• Mechanism on main. ops/jsd.py L59–70 (and the identical forms in every jsd_loss_and_grad backend) emit dX = loss + Q (β = 1) / dX = −P (β = 0) from an in-kernel exp(log_softmax), and fused_linear_jsd.py L229–232 then applies the log-softmax backward with a second softmax computed by torch.softmax outside the kernel. The "+Q" and "−Q·ΣdX" terms cancel exactly in exact arithmetic but not for two differently rounded Q's, leaving ≈ 2–3 ulp·Q per entry even when student_logits == teacher_logits. With that emulation on Qwen2.5-0.5B logits (fp32 inputs, i.e. after fix: request float32 logits from the fused JSD projection GEMM #1433): spurious gradient norm 1.5e-5 per token at the optimum (max 1e-4, true value 0); at a 1e-2 logit gap 23 % of tokens have > 10 % relative gradient error for β ∈ {0, 1}; synthetic margin-8 rows: 100 % error with 56 % sign flips. β ∈ (0, 1) is benign (≤ 1e-4), consistent with "Mixed-beta JSD isn't affected".
• Not covered by this PR (as its description notes): the standalone LigerJSDFunction path, where PyTorch's own log_softmax backward supplies the second softmax against the kernel's Q·(X − Y) + Q; and the CuTe DSL / cuTile / Ascend jsd_loss_and_grad backends (_cutedsl/jsd.py L245–264, _cutile/jsd.py L100–125, _ascend/ops/jsd.py L64–92), which carry the same dx = loss + q / dx = −p forms — CuTe DSL is the default on sm_90+ when nvidia-cutlass-dsl is installed. A student == teacher ⇒ grad == 0 test on LigerJSDLoss (log-probs in, through autograd) would catch the standalone case.
• Adjacent, unrelated to softmax collapse: ops/tvd.py L90–98 loads p/q without the .to(tl.float32) cast that kl_div.py, jsd.py, fused_linear_kl_div.py and fused_ce_tvd.py apply, so bf16 inputs are differenced in bf16 (loss only; the gradient is sign-based). Happy to send a one-line PR if wanted.

[vaibhavjindal]

cc @arde171 @yueyiming2009 @Mecoli1219 @PrathyushaPolepalli