Conversation
added 3 commits
August 7, 2026 04:35
Wires FlashInfer's Blackwell MNNVL CuTe DSL backend (flashinfer-ai/flashinfer#4358) into --flashinfer-allreduce-fusion-backend as "cute-dsl", and uses the fused MoE-finalize pattern it provides. Two fusion patterns share one compiled workspace: * AllReduce + residual + RMSNorm, the drop-in for the mnnvl backend. * MoE finalize + shared-expert add + AllReduce + residual + RMSNorm. The TRT-LLM MoE runner already supports deferring its finalize; instead of materializing it with a separate kernel, the MoE layer hands the operands to the next layer's input RMSNorm, where they fold into the collective. The workspace compiles for a single static shape and rendezvouses symmetric memory, so it is built once during warmup and never inside a forward pass; any shape it cannot serve falls back to the unfused path, and a deferred finalize that reaches a non-fusing path is materialized in place. FlashInfer ships routing profiles for GB300 H=8192 / top-k=10 only, so the profile is synthesized for the running model. Its HT persistent kernel also constrains the shard split to divide the hidden size, which 512 threads x 2 vectors does not at H=6144; _ht_shard_split re-derives the split from the kernel's divisibility rules and reproduces FlashInfer's shipped values at H=8192. GLM-5.2-NVFP4, TP8 on 8xB300, bs=1 in=10000 out=512: output throughput 160.8 -> 170.4 tok/s (+6.0%), ITL 6.22 -> 5.87 ms. In the decode profile the comm+finalize block drops 25.6% (36.6 -> 27.2 ms over 25 steps) and the separate sglang::moeFinalizeKernel all but disappears (7556 -> 103 us). GSM8K 5-shot over 400 questions: 0.958. Prefill is unchanged: a 16k-token chunk is above FUSE_ALLREDUCE_MAX_BATCH_SIZE, so every backend falls back to NCCL there.
`_forward_with_allreduce_fusion` is only reached with the producer's all-reduce still owed: every caller (LayerCommunicator.prepare_attn and prepare_mlp, nemotron's input_norm_maybe_fuse_allreduce) takes an explicit all-reduce path when it does not expect the fusion to run. When flashinfer_allreduce_residual_rmsnorm returned None the CUDA path fell through to a plain norm and the reduction was silently skipped. Previously unreachable in practice, because the gate that skipped the producer's all-reduce and the fused call itself tested the same conditions. The cute-dsl backend breaks that symmetry: its workspace is compiled for one static shape and legitimately declines anything outside it.
FlashInfer ships routing profiles for GB300 H=8192 only. The LL/BT crossovers and the LL publish granularity do not transfer to H=6144: at TP8 the shipped LL<=15 bound sends m=16..64 to BT, which carries a ~13.6 us floor there against LL's 5.9 us. Also adds SGLANG_USE_CUTEDSL_ATTN_ALLREDUCE_FUSION to restrict the backend to the MoE output group.
added 6 commits
August 8, 2026 22:02
This reverts commit 997823e.
SGLANG_USE_CUTEDSL_ATTN_ALLREDUCE_FUSION=0 keeps the attention output all-reduce on mnnvl while the MoE output, where the deferred finalize folds in, stays on cute-dsl.
This reverts commit af84e49.
SGLANG_TEST_CUTEDSL_STOCK=1 disables the measured bound tables and the measured low-latency all-reduce tuning, restoring the heuristic defaults. SGLANG_TEST_CUTEDSL_LL_AR_EPT overrides publish_elements_per_thread.
The backend served the MoE output group and the attention output group alike; the toggle that narrowed it to the MoE group only existed to A/B the two, so cute_dsl_serves_group collapses into cute_dsl_backend_selected and its always-taken branches fold into their call sites.
ywgrit
pushed a commit
to ywgrit/sglang
that referenced
this pull request
Aug 30, 2026
Transfer ownership of the final DeepSeek/GLM NextN MoE TP reduction to the shared-head RMSNorm when the existing FlashInfer fusion backend can serve the shape. This removes the standalone post-experts collective/norm boundary while leaving ordinary target-model layers on their existing next-layer consumer path. Keep the optimization fail-closed: only pure TP with TP>1, a supported backend, a non-scattered MLP layout, and no incompatible CP, DP-attention, MoE-CP all-gather, or hybrid EP+TP mode can publish the ownership marker. Marked partial output without a residual fails loudly. If the FlashInfer runtime declines after ownership has moved, materialize the owed AllReduce before the ordinary norm. This fallback is the same independently identified fix in draft PR sgl-project#34134 by b8zhong; it is included here because the new final-norm consumer otherwise has the same correctness obligation. Validation on 2x H20 with real DeepSeek-V3-0324 NextN weights passed output allclose (max abs 0.015625, rtol/atol 1e-2), removed all 30 explicit reductions per measured rank, and replaced them with 30 fused final-norm calls. CUDA Graph replay and Nsight traces passed. The isolated BF16 hidden=7168 operator is 1.77x-2.49x faster for 1-128 tokens; the real batch-1 NextN core gain is below stable measurement resolution, consistent with the boundary's roughly 0.33% baseline share. Focused CPU validation: 8 tests plus 9 subtests passed in the matching SGLang environment; the runtime-decline regression test also passes on the rebased main checkout. Ruff selected checks, isort, Black, py_compile, and git diff checks pass.
Collaborator
Author
|
No need |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
FlashInfer provides an MNNVL CuTe DSL all-reduce fusion backend. It fuses AllReduce + Residual + RMSNorm, and additionally folds the MoE finalize and the shared-expert add into the collective when the MoE runner defers them.
Its shipped routing profiles cover GB300 at H=8192 and K=10 only. The low latency/balanced protocol crossovers in those profiles do not transfer to GLM-5.2 (H=6144).
Modifications
cute-dslto--flashinfer-allreduce-fusion-backend.publish_threads * elements_per_thread; the shipped value of 8 yields eight blocks at H=8192 but only six at H=6144.Accuracy Tests
GSM8K, 200 questions, TP8:
auto)Speed Tests and Profiling
nvidia/GLM-5.2-NVFP4, TP8, B300, 1024 in / 1024 out, decode throughput relative to the MNNVL backend.Total decode throughput (
Output token throughput) and the ratio against the MNNVL backend.Command:
CI States
Pending.
CI States
Latest PR Test (Base): ❌ Run #31552123681
Latest PR Test (Extra): ❌ Run #31552123516