Fix MiniMax-M3 allreduce fusion correctness

[JustinTong0323]

Summary

Enable FlashInfer TRTLLM allreduce fusion by default for MiniMax-M3 while preserving generation correctness.

Root cause:

• The FlashInfer AR kernel was not the correctness culprit. Layer-by-layer fused-vs-native checks showed the fused AR+Gemma RMSNorm op stayed finite and close to native on identical inputs.
• SGLang's MiniMax-M3 forward path was deferring sparse MoE output allreduce into the next layer norm through the generic LayerCommunicator path.
• vLLM's MiniMax-M3 forward path only defers dense MLP output into the next norm; sparse MoE layers immediately all-reduce. This patch matches that contract.

Changes:

• Forward use_attn_tp_group correctly from GemmaRMSNorm.forward_with_allreduce_fusion.
• For MiniMax-M3 sparse MoE layers, force immediate MLP/MoE allreduce instead of deferring into the next input norm. Dense MLP layers still keep the fused deferred path.
• Add MiniMaxM3SparseForCausalLM and MiniMaxM3SparseForConditionalGeneration to the FlashInfer allreduce-fusion auto-enable allowlist.

Correctness

H200 TP8, MiniMax-M3, FlashInfer AR fusion ON.

GSM8K stop-rate validation, no CUDA graph:

• Artifact: /scratch/ar-dense-only-val/smoke100.jsonl.summary.json
• n_total=100, n_ok=100, n_err=0
• stop_rate_pct=100.0
• runaway_rate_pct=0.0
• ct_p50=223, ct_p99=599, ct_max=650
• gsm8k_acc_pct=96.0

GSM8K stop-rate validation, CUDA graph enabled:

• Artifact: /scratch/ar-dense-only-graph-val/smoke100.jsonl.summary.json
• n_total=100, n_ok=100, n_err=0
• stop_rate_pct=100.0
• runaway_rate_pct=0.0
• ct_p50=226, ct_p99=656, ct_max=756
• gsm8k_acc_pct=95.0

The validation criterion here is finish_reason=stop / stop-rate, not GSM8K accuracy alone.

Performance

H200 TP8, CUDA graph enabled, SGLang bench_serving, random workload:

• --random-input-len 1024
• --random-output-len 512
• --num-prompts 128
• --max-concurrency 32

Fusion ON artifact:

• /scratch/ar-fusion-perf/fusion_on/bench.jsonl
• Log confirms default enable: Auto-enabling FlashInfer AllReduce Fusion on SM90/SM10X for MiniMaxM3SparseForConditionalGeneration

Native baseline artifact:

• /scratch/ar-fusion-perf/native_baseline/bench.jsonl
• Launched with --enforce-disable-flashinfer-allreduce-fusion

Arm	Output throughput	Mean TPOT	Median TPOT	P99 TPOT	Median TTFT
Fusion ON	1011.78 tok/s	29.37 ms	29.69 ms	39.55 ms	146.19 ms
Native baseline	1023.78 tok/s	29.07 ms	29.45 ms	39.68 ms	147.25 ms

Delta, fusion ON vs native:

• Output throughput: -12.00 tok/s (-1.17%)
• Mean TPOT: +0.30 ms (+1.02%)
• Median TPOT: +0.23 ms (+0.79%)

This is neutral/slightly negative on H200, where MiniMax-M3 is still on the Triton MoE fallback. The larger AR-fusion gain is expected on B200/MSA and is intentionally left as follow-up.

Checks

• python3 -m py_compile python/sglang/srt/layers/layernorm.py python/sglang/srt/models/minimax_m3.py python/sglang/srt/server_args.py
• pre-commit run --files python/sglang/srt/layers/layernorm.py python/sglang/srt/models/minimax_m3.py python/sglang/srt/server_args.py

Follow-up

• B200/MSA correctness and performance validation.

[JustinTong0323]

B200 follow-up validation for PR #40 (fix/minimax-m3-ar-fusion, head 6820ebf74a) on m3-pcg-b200.

Setup:

• Source: PR Fix MiniMax-M3 allreduce fusion correctness #40 diff applied on /sgl-workspace/sglang over base 56e2903
• Model: /scratch/m3-mxfp8
• Hardware: 8x B200
• Launch flags used for both arms:
  • --tp-size 8
  • --attention-backend fa4
  • --page-size 128
  • --moe-runner-backend deep_gemm
  • --chunked-prefill-size 8192
  • --reasoning-parser minimax-m3
  • CUDA graph enabled
• Fusion-on log confirmed FlashInfer TRTLLM AR workspace init for TP8.
• Native baseline log confirmed auto-enable first, then forced disable via --enforce-disable-flashinfer-allreduce-fusion.

Correctness, full GSM8K finish-reason tally:

• Artifact: /scratch/pr40-b200/fusion_on/gsm8k_stoprate.jsonl.summary.json
• n_total=1319, n_ok=1319, n_err=0
• finish_counts={"stop": 1316, "length": 3}
• stop_rate_pct=99.7726
• runaway_rate_pct=0.2274
• ct_p50=227, ct_p99=1048, ct_max=8192
• gsm8k_acc_pct=90.22

This is close to the target, but not 0-runaway: 3 full-GSM8K requests hit finish_reason=length at 8192 tokens. The three length cases were indices 255, 301, and 840; each had 30k+ reasoning characters and little/no final content, so these are real reasoning runaways rather than answer-format issues.

Perf workload:

• sglang.bench_serving
• --backend sglang-oai
• --random-input-len 2048
• --random-output-len 256
• --random-range-ratio 1
• --max-concurrency 64
• --num-prompts 320
• --flush-cache
• Seeds: 11, 12

Perf artifacts:

• Fusion ON:
  • /scratch/pr40-b200/fusion_on/bench_seed11.jsonl
  • /scratch/pr40-b200/fusion_on/bench_seed12.jsonl
• Native baseline:
  • /scratch/pr40-b200/native_baseline/bench_seed11.jsonl
  • /scratch/pr40-b200/native_baseline/bench_seed12.jsonl

Arm	Seed	Output throughput	Mean TPOT	Median TPOT	P99 TPOT	Median TTFT
Fusion ON	11	1729.95 tok/s	28.08 ms	27.99 ms	37.72 ms	2070.38 ms
Fusion ON	12	1817.53 tok/s	28.05 ms	28.55 ms	36.28 ms	1709.57 ms
Native baseline	11	2087.76 tok/s	24.15 ms	24.27 ms	29.29 ms	1677.86 ms
Native baseline	12	1991.45 tok/s	24.69 ms	24.58 ms	30.70 ms	1954.64 ms

Averages:

• Fusion ON: 1773.74 tok/s, mean TPOT 28.07 ms, median TPOT 28.27 ms, P99 TPOT 37.00 ms
• Native baseline: 2039.60 tok/s, mean TPOT 24.42 ms, median TPOT 24.43 ms, P99 TPOT 30.00 ms

Delta, fusion ON vs native baseline:

• Output throughput: -265.86 tok/s (-13.03%)
• Mean TPOT: +3.65 ms (+14.93%)

So on this B200/MSA setup, PR #40 preserves near-complete stop behavior but does not show a perf win for FlashInfer AR fusion; native allreduce is faster in the measured workload.

[JustinTong0323]

Follow-up sparse full-fusion investigation conclusion: keep PR #40 dense-only.

The attempt to recover full AR-fusion speed by also deferring the 57 sparse-MoE layers is blocked with the current FlashInfer TRTLLM fused AR+GemmaNorm kernel. The root cause is not another SGLang integration knob: the fused kernel's residual path diverges from the native all-reduce + GemmaRMSNorm path even when fed the identical input tensor.

Evidence from the B200 full-forward fused-vs-native trace:

• Trace artifact: /scratch/m3-ar-trace-20260615-184638/trace/rank*.jsonl
• Repro response: /scratch/m3-ar-trace-20260615-184638/one_resp.json
• The served repro request hit finish_reason=length at completion_tokens=128, already repeating text after answering 2+3=5.
• First sparse-MoE deferred fused call:
  • layer_id=3
  • is_sparse=True
  • group=moe_tp
  • shape=[186, 6144]
  • dtype=torch.bfloat16
  • x_contiguous=True
  • residual_max_diff=2.0
  • norm_max_diff=0.00390625
• The diffs accumulate across sparse layers, reaching residual max diffs up to 32.0 and norm max diffs up to 0.25.
• End-to-end full-fusion correctness remains broken: the B200 normal-TP full-fusion run was 3% stop / 97% runaway.

Dense layers tolerate the same class of fused-vs-native bf16 residual differences, which is why PR #40 works. Sparse MoE does not: the residual feeds the sparse MoE/router path, and that path is bf16-residual-sensitive enough that the fused-kernel delta propagates into runaway generation.

Also ruled out:

• Normal TP-group routing for the sparse deferred AR. That still produced the same runaway class (3% stop / 97% runaway on B200).
• A deeper raw-partial / integration-side deferred-consumption idea. The precision trace compares the fused kernel against native on the exact same input tensor, so the divergence exists before any downstream SGLang-specific interpretation can fix it.

This makes full sparse-MoE fusion an upstream FlashInfer kernel correctness problem, not an SGLang forward-contract fix. The safe ceiling for this PR is the dense-only behavior: defer/fuse the dense MLP layers, keep sparse MoE all-reduced immediately.

This also matches vLLM's MiniMax-M3 implementation in vllm-project/vllm#45381: the decoder sets input-allreduce fusion only when the previous layer is not MoE (not _is_moe_layer(config, layer_id - 1)). vLLM independently avoids fusing sparse-MoE output into the next input norm.

Validated PR #40 ceiling:

• H200: 100% stop-rate with fusion ON.
• B200: 99.77% stop-rate with fusion ON.
• Perf: neutral rather than the hoped-for full-fusion gain, because only the 3 dense layers are safely fusable today.

Conclusion: PR #40 should stay dense-only. Full sparse-MoE AR fusion should be revisited only after FlashInfer TRTLLM fused AR+GemmaNorm can match native residual math for these MiniMax-M3 sparse-MoE tensors.