[gau-nernst]

Overview

This PR adds sequence parallelism (SP) to MiniMax-M3 on NVIDIA when running TP + EP with the MegaMoE backend, plus the GEMM-RS kernel for o_proj and the dense MLP down_proj. The design follows the existing Kimi-K3 and DeepSeek-V4.1 SP.

Without SP, MegaMoE already shards tokens across TP ranks internally, but then writes its shard into a zero-filled full-size buffer, and the next RMSNorm all-reduces it. So every layer pays 2 full all-reduces: one after attention, and one on a mostly-zero MoE output. With SP, the residual stream stays sharded by token, and each MoE layer does 1 all-gather before attention and 1 reduce-scatter after o_proj. MegaMoE consumes the shard directly. The 3 dense layers do one more all-gather/reduce-scatter pair around their MLP.

Claims

• SP is enabled automatically for TP > 1 + EP + MegaMoE + PP1, same rule as DeepSeek-V4.1 and Kimi-K3. GEMM-RS is opt-in with VLLM_ENABLE_GEMM_RS=1.
• On 4×GB300 with MiniMaxAI/MiniMax-M3.1-Flash (8k in / 1k out, DSpark spec decode), compared to TP4 + EP without SP:
  • P50 TTFT is 14-20% lower at all concurrencies.
  • Output throughput is +6.5% at concurrency 16 and +13.6% at concurrency 64, a tie at concurrency 4, and -4% at concurrency 1 (see below).
• gsm8k is unchanged.

Validation

All benchmarks were done on a single 4×GB300 node, with f81a414fae as the base and MiniMaxAI/MiniMax-M3.1-Flash (9e7ebac). The serve command is the "Single node (TP4, NVIDIA GB200)" command from the vLLM recipe, plus --no-enable-prefix-caching. 3 setups are compared:

• TP4: the recipe command as-is, on main
• TP4 + EP: -ep --moe-backend deep_gemm_mega_moe, on main
• TP4 + EP + SP: same flags with VLLM_ENABLE_GEMM_RS=1, on this PR

Serve command (TP4 + EP + SP)

VLLM_FLOAT32_MATMUL_PRECISION=high VLLM_ENABLE_GEMM_RS=1 \
vllm serve MiniMaxAI/MiniMax-M3.1-Flash \
  -tp 4 -ep --moe-backend deep_gemm_mega_moe \
  --reasoning-parser minimax_m3 \
  --block-size 128 \
  --kv-cache-dtype nvfp4 \
  --attention-config '{"indexer_kv_dtype": "nvfp4", "use_prefill_query_quantization": true}' \
  --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \
  --speculative-config '{"method":"dspark","num_speculative_tokens":7,"attention_backend":"FLASHINFER_TS","kv_cache_dtype":"fp8","draft_sample_method":"probabilistic","rejection_sample_method":"block"}' \
  --max-num-seqs 256 \
  --disable-hybrid-kv-cache-manager \
  --no-enable-prefix-caching

E2E: SpeedBench throughput_8k

VLLM_USE_RUST_BENCH=1 vllm bench serve --backend openai-chat \
  --dataset-name speed-bench --speed-bench-config throughput_8k \
  --speed-bench-max-input-len 8192 --speed-bench-output-len 1024 --ignore-eos \
  --num-warmups 5 --sweep-max-concurrency 1,4,16,64 --sweep-num-prompts-factor 10

The full sweep was run twice for each setup, restarting the server in between. Each cell is the range over the 2 runs; best (by mean) in bold.

Output throughput

Concurrency	TP4	TP4 + EP	TP4 + EP + SP (this PR)
1	410–411 tok/s	370 tok/s	355 tok/s
4	1,022–1,058 tok/s	928–950 tok/s	945–972 tok/s
16	2,212–2,245 tok/s	2,254–2,260 tok/s	2,397–2,410 tok/s
64	3,937–3,938 tok/s	4,299–4,338 tok/s	4,896–4,914 tok/s

P50 TTFT

Concurrency	TP4	TP4 + EP	TP4 + EP + SP (this PR)
1	164–166 ms	140 ms	111–114 ms
4	190–193 ms	163–167 ms	134–137 ms
16	200–202 ms	179–184 ms	151–152 ms
64	289–329 ms	266–280 ms	233–236 ms

P50 TPOT

Concurrency	TP4	TP4 + EP	TP4 + EP + SP (this PR)
1	2.53 ms	2.68 ms	2.76 ms
4	3.65–3.79 ms	4.30–4.53 ms	3.98–4.28 ms
16	6.55–6.64 ms	6.52–6.81 ms	6.12–6.68 ms
64	15.38–15.71 ms	14.05–14.08 ms	12.16–12.39 ms

As expected, pure TP4 is the fastest at low concurrency, and EP pulls ahead at high concurrency. SP lowers TTFT at every concurrency, with no overlap between the ranges, and widens EP's lead in throughput at concurrency 16 and 64. At concurrency 4 and for TPOT at concurrency 16, the SP and EP ranges overlap, so those are a tie.

At concurrency 1, TPOT is 3% higher with SP. This comes from spec decode acceptance (3.10 vs 3.20), not from slower steps: TPOT × acceptance length is the same (8.56 ms vs 8.58 ms). Concurrency 1 only has 10 prompts, so a small numerical difference in the outputs is enough to change the acceptance length. It was identical in both runs.

Acceptance length (mean of 2 runs)

Concurrency	TP4	TP4 + EP	TP4 + EP + SP (this PR)
1	3.02	3.20	3.10
4	3.08	3.11	3.09
16	3.14	3.15	3.17
64	3.15	3.14	3.16

Accuracy

gsm8k, 5-shot, chat completions, temperature 0. One eval per server start (run 1 / run 2).

Metric	TP4	TP4 + EP	TP4 + EP + SP (this PR)
flexible-extract	95.53 / 95.75	95.45 / 95.98	96.13 / 95.45
strict-match	95.00 / 95.45	95.22 / 95.68	95.91 / 95.22

GEMM-RS microbenchmark

torchrun --nproc-per-node=4 -- benchmarks/kernels/benchmark_gemm_rs_ar.py \
  --model minimax_m3 --backend flashinfer_cutedsl \
  --m 128 512 768 1024 1152 1280 1536 2048 4096 8192 16384 32768

TP4, N=6144. Baseline is the existing unfused path in vLLM (MXFP8 GEMM + sp_reduce_scatter).

M	o_proj (K=2048) baseline	o_proj GEMM-RS	down_proj (K=3072) baseline	down_proj GEMM-RS
128	36.6 us	39.8 us (0.92x)	55.9 us	60.5 us (0.92x)
512	61.8 us	63.3 us (0.98x)	64.6 us	64.9 us (1.00x)
768	68.2 us	67.4 us (1.01x)	68.7 us	66.7 us (1.03x)
1024	79.0 us	76.9 us (1.03x)	79.0 us	72.9 us (1.08x)
1536	101.5 us	84.0 us (1.21x)	105.3 us	82.2 us (1.28x)
2048	114.1 us	90.9 us (1.26x)	119.4 us	89.2 us (1.34x)
4096	179.3 us	133.7 us (1.34x)	191.6 us	129.9 us (1.47x)
8192	271.3 us	211.6 us (1.28x)	299.7 us	210.0 us (1.43x)
16384	475.0 us	363.6 us (1.31x)	558.9 us	374.4 us (1.49x)
32768	952.6 us	680.1 us (1.40x)	1146.9 us	704.8 us (1.63x)

M=512-1536 was repeated two more times. GEMM-RS was faster at M ≥ 1024 every time, while M=768 was mixed, so GEMM-RS is only used for M ≥ 1024.

Above 1024 tokens, the fused dense GEMM1 hands down_proj an MXFP8 QuantizedActivation, so the BF16-input numbers above include an activation quantization that production skips. Re-measured with the same pre-quantized input for both paths (range over 2 runs):

M	down_proj (K=3072) baseline	down_proj GEMM-RS
1152	66.4–74.4 us	61.4–66.2 us (1.08–1.12x)
1280	72.8–74.9 us	63.6–65.7 us (1.14x)
1536	88.9–91.4 us	65.1–66.7 us (1.37x)
2048	109.7–110.3 us	76.9–78.2 us (1.41–1.43x)
4096	173.8–175.8 us	111.3–114.0 us (1.54–1.56x)
8192	279.4–279.7 us	190.6–192.1 us (1.45–1.47x)
16384	534.0–534.7 us	343.3–345.2 us (1.55–1.56x)
32768	1083.8–1087.4 us	650.4–652.4 us (1.66–1.67x)

GEMM-RS stays faster at every M above the threshold.

Tests

• tests/kernels/test_gemm_rs_ar.py: extended with a pre-quantized MXFP8 input for both MXFP8 backends. Passes on 2×GB300. CI runs it in the 2×B200 distributed step.
• E2E SP needs 4 SM100 GPUs in one NVLink domain plus the M3.1 checkpoint, which CI does not have. It is covered by the benchmark and gsm8k runs above.

Details

Changes in vllm/models/minimax_m3/nvidia/model.py:

• Model forward: shard the hidden states after the embedding, and shard the MegaMoE padding mask once. At the end, all-gather the output together with the EAGLE3/DSpark aux hidden states in 1 collective.
• Attention: all-gather the input, reduce-scatter the output (o_proj + sp_reduce_scatter, or GEMM-RS).
• MegaMoE: takes the shard as-is (input_is_sequence_sharded), so no shard/gather/zero-fill. The existing use_sequence_parallel flag in MiniMaxM3MoE is renamed to shard_replicated_input, since it means the opposite of model-level SP.
• Dense MLP (3 of 60 layers): all-gather, gate_up (keeps the fused dense GEMM1) and down_proj, then reduce-scatter.
• Fusions: the AR + Gemma RMSNorm fusions are skipped under SP, since there is nothing left to all-reduce.

GEMM-RS with a quantized input Above its token threshold, the fused dense GEMM1 outputs an MXFP8 QuantizedActivation instead of BF16. So GemmRsAr.apply() now accepts it and skips its own quantization. The scales are already in the F8_128x4 swizzled layout that the kernel expects.

Limitations and follow-ups

• MTP blocks and PP > 1 don't use SP.
• The dense MLP stays TP-sharded instead of adding Kimi-K3's replicated option, since it is only 3 layers.
• The all-gathers move BF16. Gathering the MXFP8 qkv_proj input instead would halve the bytes; this is left for a follow-up.

AI assistance (Claude) was used to write this change and to run the validation above.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[gau-nernst]

/ci run

[github-actions]

✅ Triggered Buildkite CI #98614 for commit fce9a0df590f.

ℹ️ This PR is 13 commits behind upstream main. Merged upstream main into this branch as fce9a0df590f before starting CI. Pull the branch before pushing more changes.

[gau-nernst]

/ci retry

[github-actions]

✅ Triggered Buildkite CI #98679 for commit e4b6967f8a46, running 3 failed step(s) from Buildkite CI #98614.