[xsxszab]

Overview

MiniMax-M3 on NVIDIA fails at startup when MegaMoE, TP > 1 and EP are combined and qkv_proj takes MXFP8 input (e.g. MiniMaxAI/MiniMax-M3.1-Flash). With sequence parallelism (SP, #64659), the input norm output is all-gathered before attention; since #64852 that output is a QuantizedActivation, which the all-gather cannot take. This PR keeps the norm output in BF16 under SP and lets qkv_proj quantize after the gather.

Claims

• Fixes AttributeError: 'QuantizedActivation' object has no attribute 'contiguous' in the profile run for TP > 1 + EP + --moe-backend deep_gemm_mega_moe (SP is enabled automatically in that setup).
• No change when SP is off.
• With the fix, GSM8K is 0.958 (0.964 with VLLM_ENABLE_GEMM_RS=1), in line with the 95.5–96.1% that [Perf][MiniMax-M3] Sequence parallelism and GEMM-RS for TP + MegaMoE #64659 reports for TP4 with and without SP.

Validation

MiniMaxAI/MiniMax-M3.1-Flash on 4× B300, main at 39dbf51839. GSM8K is the full 1319 questions, 5-shot, chat mode with thinking off.

Setup	Code	Result	GSM8K
TP4 + EP + MegaMoE (SP on)	main	crashes in the profile run	–
TP4 + EP + MegaMoE (SP on)	this PR	starts	0.958
TP4 + EP + MegaMoE (SP on), VLLM_ENABLE_GEMM_RS=1	this PR	starts	0.964

Commands

The serve command is the one from #64659's validation, without the NVFP4 KV and indexer caches: their Q8KV4 prefill kernel needs CUDA 13.4+, and this cluster has 13.1. The crash path does not depend on the KV cache dtype.

VLLM_FLOAT32_MATMUL_PRECISION=high \
vllm serve MiniMaxAI/MiniMax-M3.1-Flash \
  -tp 4 -ep --moe-backend deep_gemm_mega_moe \
  --reasoning-parser minimax_m3 \
  --block-size 128 \
  --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

python tests/evals/gsm8k/gsm8k_eval.py --port <port> --max-concurrency 64 \
  --use-chat-completions --max-tokens 1024 \
  --chat-template-kwargs '{"enable_thinking": false}'

No unit test is added: the failure needs SP end to end (4 SM100 GPUs in one NVLink domain and the M3.1 checkpoint), which CI does not have, as noted in #64659. pre-commit run --hook-stage manual passes on the changed file.

Details

Root cause. forward_input_norm asks the norm for MXFP8 output whenever qkv_proj can consume it, assuming the activation goes straight from the norm to qkv_proj. Under SP, forward_after_input_norm all-gathers it first, and sp_all_gather only handles tensors. QuantizedActivation is a dataclass holding the E4M3 data and UE8M0 scales, so the custom all-gather fails on .contiguous().

Scope. Only the plain norm branch is reachable under SP: MegaMoE never defers its finalize (no MoEOutput input), and SP layers never defer the FFN all-reduce, so fuse_input_allreduce stays off. With SP off, the condition evaluates exactly as before. The dense MLP already handles SP this way: it all-gathers BF16 and quantizes after the gather.

Tradeoff. Under SP the fused norm+quant is not used. The data flow is the same as when #64659 measured SP (before #64852), with the BF16 norm now on #64852's Triton kernel. Gathering the MXFP8 data and scales instead would halve the all-gather bytes, but the 128x4 swizzled scales cannot be concatenated by rows: it needs unswizzled scales from the norm kernel and a swizzle after the gather. That is left as a follow-up.

Duplicate check. Searched open and closed PRs and issues for the error message, QuantizedActivation + all-gather, and MiniMax SP + MXFP8; no existing fix.

AI assistance (Claude Code) was used to find the root cause, write the change and run the validation; I reviewed every changed line.

---

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.

[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.

[github-actions]

✅ @xsxszab, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[xsxszab]

/ci run

[github-actions]

✅ Triggered Buildkite CI #99686 for commit d27a35c0ff70.

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