[Comm] Share the CuTe DSL AR fusion core and wire it for DeepSeek-V3 / GLM-5.x

[kpham-sgl]

Motivation

The FlashInfer MNNVL CuTe DSL AR fusion landed with the Qwen3.5 stack: it folds the MoE finalize, the shared-expert add and the post-experts all-reduce into the next layer's input RMSNorm. Only Qwen3.5 can reach it. This gives the DeepSeek-V3 family the same path — which covers GLM-5.x unchanged, since GlmMoeDsaForCausalLM subclasses DeepseekV2ForCausalLM and reuses its decoder layer and MoE block.

Draft: none of this has run on hardware. See Testing status.

Modifications

1. Extract the shared core (layers/moe/cutedsl_ar_fusion.py).

qwen35_flashinfer_fusion.py was 347 lines of which almost everything is architecture-agnostic — the handoff record, the workspace handle, and the LayerCommunicator hooks and eligibility predicates. Exactly one thing in it is Qwen-specific: gemma_weight, because Qwen3.5 normalizes with GemmaRMSNorm and DeepSeek-V3/GLM use a plain RMSNorm. It appeared at four sites (a hasattr guard and a read in prepare_attn, a read in prepare_mlp, a hasattr in should_use_all_reduce_rms_norm).

The generic parts move out, and the per-architecture surface becomes one class attribute:

layers/moe/cutedsl_ar_fusion.py            368 lines — generic core, no model names
  ├── qwen35_flashinfer_fusion.py           42 lines — NORM_WEIGHT_ATTR = "gemma_weight"
  └── deepseek_flashinfer_fusion.py         43 lines — NORM_WEIGHT_ATTR = "weight"

All four call sites now go through a _norm_gamma(layernorm) hook. qwen35_flashinfer_fusion.py keeps Qwen35MoeFinalizeHandoff / Qwen35FlashInferFusionService as aliases, so qwen2_moe.py, qwen3_5.py and the registered test_qwen35_flashinfer_fusion.py import exactly what they did before — no behaviour change on the Qwen path.

Two things that looked Qwen-specific and are not, now documented as such:

• weight_bias=0.0 was justified as "GemmaRMSNorm.gemma_weight is already checkpoint weight + 1". The value is right for both: the kernel computes x * (gamma + weight_bias) and neither flavour wants a bias.
• "the Qwen model consumes the final layer's handoff with its final GemmaRMSNorm" — generic; any model with a final norm does this.

2. Wire DeepSeek-V3 / GLM-5.x.

• DeepseekV2MoE.forward threads defer_finalize into the dual-stream path, which returns the unfinalized handoff instead of calling finalize_flashinfer_trtllm_deferred_output. Unlike Qwen3.5's shared expert, this family's carries no sigmoid gate, so it goes over unmodified — no _gate_shared_output_out_of_place equivalent is needed.
• A TP1-replicated shared expert is refused: _shared_expert_tp1 adds the shared output after the all-reduce precisely because it is replicated, so folding it into the fused add would count it once per rank.
• The decoder layer asks the communicator for eligibility and returns the handoff untouched — postprocess_layer and the _sglang_needs_allreduce_fusion flag both want a real tensor.
• The model installs one workspace handle across its layers (compiled per hidden/top_k/eps, which every MoE layer shares), and DeepseekV2ForCausalLM forwards the pre-capture hook, since BaseRunner looks it up on the top-level module.

Off unless SGLANG_FLASHINFER_MNNVL_CUTEDSL_AR_FUSION=1, matching the Qwen path. Per-layer eligibility is inherited unchanged: no DP attention, no a2a backend, no CP, unscattered input, tp_size > 1.

Testing status

Nothing has run on hardware — the kernels need a Blackwell MNNVL fabric and this was developed without a GPU. What has been checked:

• All touched modules byte-compile; pre-commit (isort/ruff/black/codespell) clean.
• A new CPU-registered test (test/registered/unit/layers/moe/test_cutedsl_ar_fusion.py) covering the extraction: each family resolves its own RMSNorm weight, _norm_gamma returns None on the wrong flavour so the predicates decline rather than raise, both communicators subclass the shared core, and the Qwen aliases still resolve.
• The existing test_qwen35_flashinfer_fusion.py imports are unchanged by design.

Before this leaves draft: GSM8K + decode throughput on GLM-5.x on B300/GB300 with the env var on, against the same run with it off; and a Qwen3.5 run to confirm the extraction is a no-op for it.

Known gaps

• Tuning is inherited. FlashInfer's shipped routing profiles cover GB300 TP8/TP16 at H=8192, K=10. GLM-5.2 is H=6144 — Integrate the CuTe DSL allreduce fusion from FlashInfer and tune for GLM-5.2 #34134 measured the LL/BT crossovers there and found the shipped bounds do not transfer. That tuning is not in this PR; GLM will run on the inherited GB300 bounds until it is ported over.
• Prefill-CP layers are excluded — the DSACPLayerCommunicator branch is untouched, so dsa_enable_prefill_cp / mla_enable_prefill_cp layers keep the ordinary path.
• Only the dual-stream MoE path defers today (it is the only one that produces a deferred payload on main); forward_normal is unchanged.

Related

• Integrate the CuTe DSL allreduce fusion from FlashInfer and tune for GLM-5.2 #34134 (@b8zhong) integrates the same backend for GLM-5.2 with its own module and a cute-dsl value on --flashinfer-allreduce-fusion-backend, plus a measured B300 tuning table. It predates the Qwen stack landing, so it now duplicates the in-tree module; its tuning is the part worth carrying over here.
• support qwen 3.8 #34585 (@Qiaolin-Yu) is where the Qwen half came from.

---

CI States

Latest PR Test (Base): ❌ Run #33438515949
Latest PR Test (Extra): ❌ Run #33438515579
Latest PR Test (AMD ROCm 7.2): ❌ Run #33438515720