[dyurk-lila]

Note

This PR depends on the packing helpers in packing_utils (get_packing_align_size_sequence, packed_segment_layout) added by the packing stack #2369 → #2372. It will not pass until that stack is merged and this stack (#2373 → #2521) is rebased onto main.

What does this PR do?

Skip inactive d-logits tiles and compute hidden-state gradients through ordered active spans in the block-sparse fused LM head.

How it works

A shared packed/unpacked metadata builder derives the forward mask and backward spans from the same TP/CP-aware token layout. The backward kernel processes only active d-logits tiles, and the hidden-gradient projection uses those spans without reordering tokens.

Unpacked active spans require an unpadded, CP=1 row layout; the call rejects spans when padding or CP sharding changes their coordinates. Packed spans retain the CP-aware path.

The final vocabulary split allocates only its live width. Weight gradients continue to use dense cuBLAS; this PR does not replace that projection with sparse accumulation.

Usage

The optimization is used by the triton_block_sparse backend introduced in #2375.

Performance

Recorded isolated LM-head measurements used one 8× NVIDIA H100 80 GB node with 180 CPU requested and limited. The 58-point matrix covered four model shapes at TP4/CP2 and 16K–256K tokens, plus every (TP, CP) layout with TP × CP <= 8 for a representative 16K-token shape.

Optimization	Recorded latency reduction	Baseline and method
Skip inactive d-logits tiles	57/58 wins; mean 5.81%, median 6.25%, range −1.44% to +10.46%	Contribution estimated by matched ratio-of-ratios
Ordered-span dHidden	57/58 wins; mean 4.93%, median 5.22%, range −1.36% to +8.64%	Same-run A/B against forward + d-logits skipping
Live-width final vocabulary split	16/16 wins; 8.8–15.2% backward-only reduction	Same-run A/B against padded staging and copy

The composed d-logits + span contribution was an estimated 9.93% mean, 10.48% median incremental latency reduction over the forward-only optimization in #2375 (range −1.62% to +16.25%). The full path still beat dense Triton in the one incremental regression. The live-width result uses a separate comparison and should not be added to that estimate.

These are kernel-level results on one GPU generation, not end-to-end training speedups or a universal win for every mask. Weight-gradient projection remains dense cuBLAS. The live-width allocation also avoids padding the final temporary buffer, but the recorded figures above measure latency, not a quantified peak-memory reduction. Recorded benchmarks predate the current rebase and were not rerun for it.

Test coverage

CPU active-span tests cover layout reconstruction and inactive regions. GPU tests compare values and gradients with dense references across partial splits, masks, and tensor-parallel layouts. The no-grad entropy parity test permits sub-micro rounding differences while retaining its independent fp64 reference.

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 8 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 #2375. 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
Touches fused LM-head backward and gradient layout under TP/CP packing; wrong spans would silently zero gradients, though guarded by assertions and extensive GPU correctness tests on an optional backend.

Overview
Adds block-sparse backward for the Triton fused LM head by deriving a forward active mask and coalesced active spans (gaps ≤128) from the RL loss mask, including packed TP/CP/FP8 layouts via new active_spans helpers.

The Triton backward now skips inactive d-logits tiles, runs span-local dHidden GEMMs instead of a full dense projection when spans are provided, and can size the final vocab-split staging buffer to the live split width. FusedLinearLogprobTriton and the Megatron fused logprob entry points accept optional active_spans; unpacked spans are restricted to unpadded CP=1 layouts.

Training wiring builds mask/spans on CPU before the microbatch device copy, passes them through megatron_model_wrapper into fused packed/unpacked paths, and ensures loss_mask is forwarded when using triton_block_sparse. Benchmarks and CPU/GPU tests cover span layout, correctness vs dense references, entropy gradients, and the narrow-tail optimization.

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

[gemini-code-assist]

Code Review

This pull request introduces active spans optimization for the fused linear logprob Triton kernel, allowing the backward pass to restrict the d_hidden projection to active regions. It adds a new active_spans.py module to build and coalesce active spans, updates the Triton kernel to support these spans and narrow dlogits tail staging, and integrates the optimization across model wrappers, workers, and trainers. A review comment suggests optimizing the dense backward path in fused_linear_logprob_triton.py by replacing += torch.matmul with torch.addmm to avoid temporary tensor allocations on subsequent splits.

[gemini-code-assist]

Using torch.addmm instead of += torch.matmul avoids allocating a temporary tensor for the intermediate product on every split after the first one. This significantly reduces peak memory usage and improves performance, especially for large sequence lengths and vocabularies.

Suggested change

	d_hidden += torch.matmul(split_d_logits, split_weight)
	torch.addmm(d_hidden, split_d_logits, split_weight, out=d_hidden)

[greptile-apps]

Confidence Score: 4/5

[High risk] Adds block-sparse optimization to the fused language model head backward pass.

The PR should not merge until the partial-split wide-tail comparison launches a valid kernel grid.

Findings

1.  Wide-tail grid exceeds live split ▶

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Loss mask and packed layout] --> B[Host active metadata]
  B --> C[Forward active mask]
  B --> D[CP-local ordered spans]
  C --> E[Fused d-logits kernel]
  E --> F[Dense weight gradient]
  E --> G[Span-restricted hidden gradient]
  D --> G

Loading

_{Reviews (1) · Last reviewed commit: "test(megatron): allow bf16 entropy-mode ..."}

[greptile-apps]

Wide-tail grid exceeds live split

When narrow_dlogits_tail=False and the final vocabulary split is partial, this grid launches programs for all 9,504 staging columns, while the kernel lays out programs using only the smaller live split_width. The extra programs can reach a group with no valid token tiles and make the kernel’s program-ID division invalid. As a result, the new --narrow-dlogits-tail-probe can fail instead of comparing the two staging widths. Size the launch grid from the live split width; the staging buffer can remain wide.

Suggested change

	def d_logits_grid(meta):
	return (triton.cdiv(num_tokens, meta["BLOCK_SIZE_M"]) * triton.cdiv(staging_width, meta["BLOCK_SIZE_N"]),)
	def d_logits_grid(meta):
	return (triton.cdiv(num_tokens, meta["BLOCK_SIZE_M"]) * triton.cdiv(split_width, meta["BLOCK_SIZE_N"]),)