[Kimi K2.5] Fuse deferred MoE finalize into next-layer AR + residual + RMSNorm

[kpham-sgl]

Summary

Fuse the deferred MoE finalize into the next layer's input-RMSNorm via flashinfer kMoEFinalizeARResidualRMSNorm (pattern 7): weighted expert combine + shared-expert add + all-reduce + residual + RMSNorm in one kernel, while keeping residual_out available for Eagle3 aux capture. Gated by SGLANG_ENABLE_MOE_DEFERRED_FINALIZE. Builds on #28343.

Key modifications

• deepseek_v2.py: when deferred finalize + AR fusion are active, carry the un-finalized permuted gemm2_out across the layer boundary with a finalize-inputs bundle attached.
• communicator.py / layernorm.py: input-RMSNorm consumes the bundle through new RMSNorm.forward_with_moe_finalize_allreduce_fusion (falls back to separate finalize + AR+RMSNorm fusion if unavailable).
• flashinfer_trtllm.py: FlashInferTrtllmMoeFinalizeFusionBundle to package deferred outputs.
• flashinfer_comm_fusion.py: new fused custom op; reshape expanded_idx_to_permuted_idx/expert_scale_factor to [token, top_k] (kernel derives top_k from size(-1) — flat index caused OOB); force trtllm backend (pattern 7 is trtllm-only); lamport workspace sizing + MAX_COMM_SIZE/capacity fallbacks.
• model_runner.py: pre-size the MoE-TP workspace for the permuted buffer before CUDA-graph capture (never grow inside captured graphs).

Note: pattern 7 is trtllm-only, so the fusion is intra-node (single NVLink domain).

HumanEval (164, K2.5-NVFP4 + EAGLE3, attn_tp8_moe_tp8)

Metric	Fused	Baseline
Accept length	3.02	3.03
Successful	164/164	164/164

Accept length is unchanged within noise — residual_out is correctly preserved for Eagle3 (a broken residual would collapse acceptance toward ~1.0).

Numerics & speed vs the separate-finalize fallback

Compared the fused kernel against the fallback path (deferred finalize + kARResidualRMSNorm, which also emits residual_out):

• Correctness: no structural/indexing bug. 8-rank kernel equivalence vs fp32 truth — residual/norm mean rel-error 3.4% / 3.8% (fused) vs 1.9% / 1.9% (fallback): the fused kernel carries ~2× the bf16 rounding error (it accumulates the expert combine inside the bf16 AR pipeline; fp32_acc doesn't reach that step). Negligible on real data — HumanEval accept length 3.02 fused == 3.02 fallback; only OOD random-token GSP shows a small dip (2.18 vs 2.23).
• Speed (graph-captured decode, 8×B200, per MoE layer): fused is slower at every batch size — bs=1 15.5 vs 14.4 µs (+1.0), bs=16 38.4 vs 26.4 µs (+12.0). Confirmed the kernel's allreduce_in is the padded permuted gemm2_out (verified runtime rows 256 → 6272, i.e. 64–128× the token count); this is trtllm's do_finalize=False output that the index map gathers from — the intended input, not a mis-sized/bugged input. The cross-rank AR is over the finalized T rows though (timing is invariant to the permuted-buffer size; AR-only is ~12–22 µs, latency-bound), so the slowdown is not AR bandwidth — it's the fused kernel's in-kernel finalize being ~2–4× the cost of the standalone finalize kernel (4–6 µs).

Takeaway: under CUDA-graph decode the cheaper fallback is faster, more accurate, and already preserves residual_out, so it is preferred there; the fused kernel only wins when launch-bound (eager).

Made with Cursor

---

CI States

Latest PR Test (Base): ❌ Run #27725890409
Latest PR Test (Extra): ❌ Run #27725890323

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[kpham-sgl]

flashinfer-ai/flashinfer#4358 supersede this API