[dyurk-lila]

What does this PR do?

Autotune the fused LM-head Triton kernels across bounded shape buckets and skip entropy work for log-probability-only calls.

TLDR: select launch geometry for each workload class while keeping autotuning inputs immutable and log-sum-exp stable for strongly negative logits.

How it works

• Tune tile dimensions, pipeline stages, and warp counts instead of using one fixed launch configuration.
• Bucket token counts so neighboring lengths reuse compiled kernels and tuning results.
• Keep raw, reduced, and result buffers separate: Triton executes multiple candidates during tuning, so an epilogue must not overwrite inputs needed by later candidates.
• Start log-sum-exp from the true row maximum, including rows where every logit is negative.
• Remove unreachable vendored reductions and backward strategies, and release the final partial d-logits buffer before its projections.

H100 benchmark

The standalone sweep used one 8x NVIDIA H100 80GB node with 180 CPU requested and limited, Torch 2.11.0+cu130, and Triton 3.6.0. 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
Regressions greater than 1%	4 / 104
Worst regression	1.40%

Mean reductions by LM-head shape were 1.90% for (H=4096,V=154880), 2.02% for (4096,131072), 7.49% for (7168,163840), and 10.89% for (8192,131072).

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.

Dense log-prob-only specialization

A separate recorded sweep on one 8× NVIDIA H100 80 GB node, with 180 CPU requested and limited, compared the general entropy-capable launch with the log-prob-only specialization used by the SkyRL adapter. It timed the combined fused forward/backward call, with no loss masking. The 22 points covered four LM-head shapes at TP4/CP2 and 16K–256K total tokens, plus (H=4096,V=131072) at TP1/CP1 and TP8/CP1 with 16K tokens.

All 22/22 points were faster: 7.06% mean, 7.17% median runtime reduction, with a 3.28–10.86% range.

LM-head shape (hidden, vocabulary) at TP4/CP2	Mean reduction	Range
Nemotron Super 120B (4096, 131072)	9.00%	8.17–10.06%
Nemotron Ultra 550B (8192, 131072)	5.24%	3.28–9.56%
GLM-5.3 Flash (4096, 154880)	8.53%	6.97–10.86%
Kimi K3 (7168, 163840)	4.78%	3.95–5.61%

The additional TP checks reduced runtime by 7.30% at TP1 and 10.11% at TP8. Removing entropy state saved 128 KiB–2 MiB of measured peak allocation as rank-local token count increased. These results measure the specialization separately from the autotuning comparison above; the percentages should not be added together. They are isolated kernel forward/backward results on H100, not whole-training throughput or other-GPU claims. Recorded benchmarks predate the current rebase and were not rerun for it.

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 and 514 s for the largest unsharded Kimi case in this sweep), so long-running jobs should warm expected buckets before timing.

Validation

Before the short-tail metrics follow-up, focused CPU suites passed in both mirrors: 415 tests across collation, DP sampling, worker batching, layout, active spans, and Megatron correctness (13 Megatron-only skips); 40 bin-packer tests; and 3 unpacked active-span guard tests with megatron-core installed. Ruff and Black passed. Seven targeted Triton value/gradient tests passed on one H100; the full GPU suite and benchmarks were not rerun.

After the metrics follow-up, 39 focused SFT CPU tests passed in each mirror, including synchronous and asynchronous short-tail logging; Ruff, Black, and secret detection passed on the changed files.

Review order

Part 5 of 9 in a linked series. The stack is based on main; each later PR targets the preceding layer. Review and merge from bottom to top.

Depends on #2372. Review only this layer.

1. perf(sft): use modified first-fit decreasing packing #2369 — use modified first-fit decreasing packing
2. perf(train): make sequence packing alignment-aware #2370 — make sequence packing alignment-aware
3. refactor(megatron): share packed segment layout derivation #2371 — share packed segment layout derivation
4. perf(sft): align packed bin counts to data parallelism #2372 — align packed bin counts to data parallelism
5. perf(megatron): autotune fused LM-head Triton kernels #2373 — autotune fused LM-head Triton kernels
6. feat(megatron): fuse LM-head entropy for RL #2374 — fuse LM-head entropy for RL
7. feat(megatron): add block-sparse fused LM-head forward #2375 — add block-sparse fused LM-head forward
8. perf(megatron): add block-sparse fused LM-head backward #2376 — add block-sparse fused LM-head backward
9. perf(megatron): fuse sample-support candidate projection with Triton #2521 — fuse sample-support candidate projection with Triton

---

Note

Medium Risk
Changes core fused LM-head Triton math, tensor-parallel epilogues, and backward memory behavior; incorrect numerics or autotune cache behavior would affect training loss and gradients, though tests and benchmarks target these paths.

Overview
This PR retunes SkyRL’s fused LM-head Triton path for Megatron: forward/backward/epilogue kernels now autotune over expanded tile schedules with cached, power-of-two token buckets (exact num_tokens still drives launches), and the SkyRL adapter calls the stack with compute_entropy=False so entropy-only work is compiled out on the log-prob-only training path.

Numerics and autotune safety are tightened: epilogues initialize and pad max reductions with -inf so all-negative logits stay finite, and final log-probs/entropy are written to separate result buffers so repeated tuning cannot corrupt inputs. Backward is simplified to the split-vocab d_logits path only (unused VERL reductions/wrappers removed), with an early release of the staging d_logits buffer on the last split to cut peak memory.

Benchmarks gain --autotune-sweep and an entropy-vs-logprob-only probe; GPU tests add schedule/bucket invariants and a strongly negative logits regression case.

^{Reviewed by Cursor Bugbot for commit ae7d851. 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 Triton fused linear log-prob kernel, introducing cached and bucketed autotuning to prevent recompilation churn, compiling out entropy-only operations when unused, and ensuring numerical stability for strongly negative logits by using the true row maximum for epilogue shifts. It also updates the benchmark suite and adds comprehensive tests. The review feedback correctly identifies an optimization opportunity in the backward pass: when vocab_size is smaller than 9504, hardcoding vocab_per_split to 9504 causes unnecessary memory allocation and redundant copy overhead. Implementing the suggested min(vocab_size, 9504) limit will improve efficiency.

[gemini-code-assist]

When vocab_size is smaller than 9504 (which is very common per rank when using Tensor Parallelism), hardcoding vocab_per_split = 9504 leads to allocating a significantly larger d_logits staging buffer than necessary. Furthermore, because split_width will be less than vocab_per_split, the code will fall into the else block and perform an expensive .contiguous() copy to slice the tensor, before freeing the original.\n\nBy setting vocab_per_split = min(vocab_size, 9504), we can avoid allocating excess columns, eliminate the redundant copy overhead, and significantly reduce peak GPU memory usage.

    vocab_per_split = min(vocab_size, 9504)\n    num_splits = (vocab_size + vocab_per_split - 1) // vocab_per_split\n    d_logits = torch.empty((num_tokens, vocab_per_split), device=hidden.device, dtype=hidden.dtype).contiguous()

[greptile-apps]

Confidence Score: 4/5

[High risk] Rewrites Triton kernel autotuning and scheduling logic.

The PR appears safe to merge, with non-blocking benchmark-fidelity and regression-coverage improvements recommended.

Findings

1.  Sweep measures different kernel ▶
2.  TP path lacks regression coverage ▶

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[SkyRL fused LM-head adapter] -->|compute_entropy=False| B[Production mainloop]
  C[Autotune sweep] -->|COMPUTE_ENTROPY=True| D[Measured mainloop]
  B --> E[Log-probabilities]
  D --> F[Reported winning schedule]

Loading

_{Reviews (1) · Last reviewed commit: "fix(megatron): pass entropy mode in auto..."}

[greptile-apps]

Sweep measures different kernel The sweep launches candidates with COMPUTE_ENTROPY=True, but the SkyRL adapter uses compute_entropy=False. These kernels do different work, so the reported winning schedule and latency gains may not represent the production path. Measuring the log-probability-only specialization would make the results useful for that path.

[greptile-apps]

TP path lacks regression coverage This negative-logit test uses the local, entropy-enabled forward path. Tensor parallelism uses a separate epilogue, and the production adapter disables entropy, so the test would not catch a regression in that combination. A TP log-probability-only case would cover the changed production path.

Knowledge Base Used: Backend workers and distributed strategies