perf(mtp): fuse trailing TP all-reduce with final RMSNorm

[ywgrit]

Proposed PR: Fuse the trailing NextN TP AllReduce with final RMSNorm

Motivation

The final DeepSeek/GLM NextN decoder layer currently materializes its post-MoE TP reduction and then executes the shared-head residual RMSNorm as a separate boundary. SGLang already fuses the same pattern between ordinary decoder layers, but deliberately excludes the last layer because there is no next decoder-layer consumer.

This change gives the NextN shared-head norm an explicit final consumer path. The producer may defer the reduction only when a strict pure-TP gate proves that the consumer can take ownership; every path must still reduce exactly once.

Changes

• Reuse LayerCommunicator's existing fusion gate through a separate final-norm consumer API.
• Permit only the last NextN layer in pure TP with TP > 1, a supported backend, full/non-scattered layout, and no incompatible DP-attention, CP, MoE-CP all-gather, A2A, or hybrid EP+TP mode.
• Publish the existing _sglang_needs_allreduce_fusion ownership marker from the final NextN producer and consume it in DeepseekModelNextN.shared_head.norm with the MoE TP group.
• Fail loudly if a marked partial result reaches the final norm without a residual.
• If the FlashInfer runtime declines after ownership has moved, materialize the owed AllReduce before the ordinary norm.

The last fallback is the same correctness issue independently identified and implemented by @b8zhong in Draft PR #34134. It is included here because the new final-norm consumer has the same ownership obligation. If maintainers prefer #34134 to land first, I can remove these five lines and rebase this PR on it.

Correctness

On 2×H20 with the official 11.7GB lmsys/DeepSeek-V3-0324-NextN checkpoint:

• BF16 output allclose passed with max absolute difference 0.015625 (rtol=atol=1e-2).
• Baseline executed one explicit post-experts reduction per iteration; candidate executed one fused final-norm collective and zero explicit reductions. Across the measured run, 30 explicit reductions/rank became 30 fused calls/rank.
• CUDA Graph capture and replay passed with changing inputs.
• Negative gates cover TP1, pure EP, hybrid EP+TP, scattered input/MLP, unavailable backend, MoE-CP all-gather, and non-last layers.
• A mocked runtime-decline regression failed before the fallback (zero reductions) and passes after it (one explicit reduction).

Focused CPU results in the matching SGLang environment: 8 passed plus 9 subtests. The rebased runtime-decline test additionally passes on current main. Ruff (F401/F821/UP037), isort, Black, py_compile, and git diff --check pass.

Performance

Hardware: 2×NVIDIA H20, TP2, BF16, hidden size 7168.

Measurement	Split baseline	Fused candidate	Result
AR → RMSNorm operator, 1–128 tokens	about 10–29 us	about 5–12 us	1.77×–2.49×
Real NextN batch-1 boundary	about 10 us	about 5 us	about 2.0×, saves about 5 us
Real NextN core	about 3 ms	run-to-run noise dominates	no stable measurable E2E gain

Nsight confirms that the standalone collective/norm boundary is replaced by the fused final-norm call. The boundary is only about 0.33% of the batch-1 NextN core, so Amdahl predicts about 0.17% whole-core gain; I therefore do not claim the noisy aggregate 1.57% observation as a stable result or claim full-model GLM-5.2 serving improvement.

Scope

Production diff: four files, +49/-6; no new kernel, CLI flag, workspace, or generic communication protocol. The full GLM-5.2 Hopper FP8 recipe requires approximately TP8/89GB per GPU, so the available 2×/4×H20 setup cannot provide a truthful full-model TTFT/ITL result. This PR reports the exact real NextN-layer path and operator/boundary evidence only.

---

CI States

Latest PR Test (Base): ❌ Run #33301037660
Latest PR Test (Extra): ❌ Run #33301037436
Latest PR Test (AMD ROCm 7.2): ❌ Run #33301037503

[ywgrit]

Hi @b8zhong, this PR is now ready for review. It reuses the existing FlashInfer collective primitive to fuse the MTP trailing TP AllReduce with the final residual + RMSNorm, with exactly-once collective ownership and fail-closed fallback.

Because the runtime-decline fallback overlaps with five lines in your Draft #34134, the PR body explicitly credits that implementation and calls out the overlap. The current one-commit patch changes four production files (+49/-6) and has exact-head 2×H20 NextN, collective-count, CUDA Graph, Nsight, and Humanize validation. Operator speedup is 1.77×–2.49×, while the measured batch-1 Amdahl ceiling is only about 0.17%, which is disclosed rather than presented as stable E2E gain.

Could you please advise whether you prefer this focused MTP consumer change to keep the credited fallback, or to rebase on #34134 after that PR lands?

[b8zhong]

Hi, I think no need to merge this change if it provides no measureable peformance difference. What do you think?

[ywgrit]

Hi, I think no need to merge this change if it provides no measureable peformance difference. What do you think?

I agree that the isolated boundary speedup does not translate into a stable measurable end-to-end improvement for the current batch-1 NextN workload. I will close this change.