perf(megatron): autotune fused LM-head Triton kernels

[dyurk-lila]

Summary

The fused Triton LLM head kernel was originally vendored from VERL with basically no changes. This PR does make two performance-improving changes, without changing any of the underlying computation logic:

• Use triton autotune to optimize block sizes, K tiles, pipeline stages, and warp counts rather than keeping one fixed set for every model/context len/hardware.
• Add an option to exclude entropy calculation entirely if the caller doesn't want to know entropy anyway

Other changes in the PR include:

• Make autotuned epilogues non-mutating. Triton executes every candidate during tuning, so overwriting candidate inputs can silently corrupt the selected result.
• Fix an edge case where the kernels could NaN if all incoming logits were strongly negative
• Remove vendored reductions and backward strategies that no SkyRL path can reach.
• Extend the existing standalone fused-LM-head benchmark and GPU tests.

Note: the LoC seems big, but most deletions are removing unused VERL code and most additions are either in the benchmark script or just adding the triton autotune config lines. The scope of relevant logic changes is relatively small.

H100 benchmark

The standalone sweep used one 8x NVIDIA H100 80GB node with 180 CPU requested and limited. It covered four public LM-head configurations, total token counts from 16K through 256K, and every (TP, CP) combination in {1,2,4,8} with TP*CP <= 8. The 200 logical cases collapse to 104 unique rank-local (M,H,V) shapes. This is a synthetic kernel-shape sweep; it does not consume dataset examples.

Compared with the prior fixed five-stage mainloop:

Metric across 104 local shapes	Autotuned
Mean latency reduction	5.57%
Median latency reduction	3.63%
p90 latency reduction	14.47%
Maximum latency reduction	26.82%
Shapes faster	88 / 104
Worst regression	1.40%

The largest observed win was the Nemotron Ultra 550B shape at local (M=8192,H=8192,V=131072), from 34.91 ms to 25.54 ms. These are isolated fused-mainloop results on one GPU generation, not end-to-end training claims; the candidate set intentionally retains schedule diversity for other models and accelerators.

Compile and tuning behavior

A specialization probe measured:

• 20,000 tokens, new 32K bucket: 13.27 s cold compile+tune;
• 20,001 tokens, same bucket: 6.13 ms, with no exact-length recompile or retune;
• 40,000 tokens, new 64K bucket: 4.66 s retune.

Crossing a power-of-two bucket intentionally retunes because the winning geometry can change. Broad cold tuning is nontrivial (14.9 s median first call in this sweep), so long-running jobs should warm expected buckets before timing.

Correctness

Autotuning originally exposed a real hazard: an epilogue overwrote reductions that later candidates consumed. The final kernels use distinct raw, reduced, and result buffers.

Two 8-GPU H100 gates passed against independently materialized fp32 references:

• every one of the 14 forward candidates matched, with maximum absolute error 9.54e-7;
• odd/non-divisible shapes, TP=8 reductions, entropy compatibility, and split-N backward matched with maximum errors 1.91e-6 for reduced log-probs/entropy, 2.34e-5 for hidden gradients, and 2.80e-5 for weight gradients.

Additional hardening initializes log-sum-exp from the true row maximum (-inf) for strongly negative rows and releases the final partial-split d-logits staging buffer before its projections. A pinned-config audit found bit-identical results before and after the lifetime change across 24 TP/dtype/shape cases.

Vendored VERL history

VERL used @triton.autotune, but each config list contained one entry, so it performed no search. The kernel originated in VERL PR #462; its comments, discussion, reviews, and commit messages do not explain the singleton sets. The original benchmark covered one bf16 workload, which may explain the narrow schedule but is not documented upstream rationale.

---

Note

Medium Risk
Changes sit on the Megatron fused LM-head hot path (forward/backward numerics, TP reductions, and first-call Triton compile/autotune latency); fixes reduce NaN risk but new tuning behavior can affect performance and timing.

Overview
Reworks the vendored fused LM-head Triton path for real schedule search and a log-prob-only training path, plus correctness fixes on the softmax epilogue.

Autotuning: Forward mainloop, epilogues (including TP), and split-N backward now use multi-candidate @triton.autotune with cache_results=True. Cache keys use power-of-two token buckets (num_tokens_bucket) while launches still use the exact token count (do_not_specialize) to limit recompile churn.

Entropy optional: compute_entropy gates entropy buffers and kernel work; FusedLinearLogprobTriton sets compute_entropy=False on forward/backward. Scalar CE reductions and the unused VERL LinearCrossEntropy wrapper / alternate backward kernels are removed; backward is only the split-N d-logits path.

Correctness / memory: Epilogues write final log-probs (and entropy) to separate result buffers so autotune replays do not corrupt reductions. Log-sum-exp uses -inf init and padding for max reductions. The last backward split drops the full-width d-logits staging buffer before matmul projections.

Tooling: bench_fused_linear_logprob.py adds --autotune-sweep, entropy-vs-logprob probes, and stricter multi-rank correctness. GPU tests cover bucket math, autotune metadata, strongly negative logits, and adapter compute_entropy=False.

^{Reviewed by Cursor Bugbot for commit cfbaff0. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request refactors and optimizes the vendored Triton fused linear-cross-entropy / log-prob kernel. Key improvements include introducing a TP-aware adapter, implementing cached and bucketed autotuning to prevent recompilation churn, compiling out entropy-only computations for log-prob-only paths, and correcting the epilogue's log-sum-exp shift to use the true row maximum (with -inf padding) to handle strongly negative logits. It also adds comprehensive benchmarks and tests. The review feedback correctly points out that the rank parameter should be added to the do_not_specialize list in both the forward mainloop and backward split-N kernels to avoid redundant compilations across different GPU ranks in a tensor-parallel group.

[gemini-code-assist]

The rank parameter is an integer that varies across different GPU ranks in a tensor-parallel group. Since it is not included in do_not_specialize, Triton will specialize the kernel for each unique rank value (e.g., 0 through 7 for TP=8). This leads to redundant compilations across ranks, significantly increasing the cold compile/tuning time at the start of distributed training. Since rank is only used in a simple arithmetic expression to shift vocabulary indices and does not affect loop bounds or tile sizes, specializing it provides no performance benefit. Adding "rank" to do_not_specialize will prevent these redundant compilations.

Suggested change

	@triton.jit(do_not_specialize=["num_tokens", "num_tokens_bucket"])
	@triton.jit(do_not_specialize=["num_tokens", "num_tokens_bucket", "rank"])

[gemini-code-assist]

Similarly to the forward mainloop kernel, the rank parameter in the backward split-N kernel is not in do_not_specialize. This causes Triton to compile a separate kernel for each tensor-parallel rank, multiplying the compilation overhead in distributed settings. Adding "rank" to do_not_specialize here will avoid redundant compilations across ranks.

Suggested change

	@triton.jit(do_not_specialize=["split_idx", "num_tokens", "num_tokens_bucket"])
	@triton.jit(do_not_specialize=["split_idx", "num_tokens", "num_tokens_bucket", "rank"])

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit cfbaff0. Configure here.}

[cursor]

Sweep omits required entropy constexpr

Medium Severity

_raw_launch and _autotuned_launch never pass COMPUTE_ENTROPY, a required constexpr on efficient_entropy_kernel_general_mainloop. The autotune sweep and specialization probe call these helpers first, so they fail to launch instead of measuring schedules. The production adapter always supplies this flag.

 

^{Reviewed by Cursor Bugbot for commit cfbaff0. Configure here.}