Skip to content

[Perf][DSV4.1] Prefer MegaMoE for supported local expert parallelism - #57220

Closed
WoosukKwon wants to merge 1 commit into
vllm-project:mainfrom
WoosukKwon:dsv41-default-megamoe
Closed

WoosukKwon wants to merge 1 commit into
vllm-project:mainfrom
WoosukKwon:dsv41-default-megamoe

Conversation

@WoosukKwon

Copy link
Copy Markdown
Collaborator

Purpose

For DeepSeek V4.1, resolve --moe-backend auto to deep_gemm_mega_moe for supported local expert-parallel deployments. Previously auto went through generic FusedMoE selection and never reached the model's MegaMoE path. In particular, TP4/DP1 with expert parallelism did not enable the model's sequence parallelism; the new selection enables it without requiring an explicit backend flag.

Resolve the choice during model configuration, before construction selects sequence-parallel and shared-expert paths. Preserve explicit backend choices. Restrict automatic selection to SM100-family CUDA devices with DeepGEMM enabled/available, MXFP4 experts, supported routing and dimensions, and a fully NVLink-connected local EP group. Leave other configurations on the existing selection path, including PP/CP, remote DP, Ray, speculative decoding and microbatching.

This changes backend selection only; it does not modify kernels or attention/indexer defaults. This should land after #57204 to avoid making the existing padding/duplicate-storage overhead the default. That padding change is not included in this diff.

Duplicate-work check

Searched open issues and PRs for DeepSeek/MegaMoE/default and auto-MoE selection, including PRs referencing #57204. No open PR implements this model-specific default selection. #40833 integrates generic MegaMoE support, #54049 integrates a FlashInfer MegaMoE implementation, and #57204 changes native MegaMoE padding; those are separate changes.

Test Plan

.venv/bin/python -m pytest tests/test_config.py -k deepseek_v41_auto_moe -q
.venv/bin/pre-commit run --files vllm/model_executor/models/config.py tests/test_config.py
.venv/bin/pre-commit run mypy-3.12 --files vllm/model_executor/models/config.py tests/test_config.py --hook-stage manual
Real-checkpoint evaluation commands
export VLLM_USE_V2_MODEL_RUNNER=1 VLLM_USE_RUST_FRONTEND=0
.venv/bin/python -m vllm.entrypoints.cli.main serve deepseek-ai/DeepSeek-V4.1-Flash \
  --revision dba1be0a40aa45a94ad051997016db3960a90277 \
  --served-model-name dsv41-gsm8k --host 127.0.0.1 --port 8003 \
  --tensor-parallel-size 4 --enable-expert-parallel --moe-backend auto \
  --tokenizer-mode deepseek_v41 --reasoning-parser deepseek_v41 \
  --kv-cache-dtype fp8 --max-model-len 1048576 --max-num-seqs 64 \
  --max-num-batched-tokens 8192 --no-enable-prefix-caching --language-model-only \
  --generation-config vllm --seed 123 \
  --default-chat-template-kwargs '{"thinking":false}' \
  --compilation-config '{"mode":0,"cudagraph_mode":"FULL_DECODE_ONLY"}' \
  --engram-config '{"cpu_offload":true}'

After the server is ready:

.venv/bin/python -m lm_eval run --model local-chat-completions \
  --model_args model=dsv41-gsm8k,base_url=http://127.0.0.1:8003/v1/chat/completions,num_concurrent=64,timeout=600,max_retries=2,max_gen_toks=1024,max_length=16384,seed=123 \
  --tasks gsm8k --num_fewshot 5 --batch_size 1 \
  --apply_chat_template --fewshot_as_multiturn \
  --gen_kwargs max_gen_toks=1024 temperature=0.0 do_sample=False "chat_template_kwargs={'thinking': False}" \
  --seed 123,123,123,123 --log_samples --output_path gsm8k-results

Repeat on parent commit 1085b64425 with --moe-backend auto for the original selection path. For this comparison, the parent model-configuration file was restored while all other runtime source files were identical. The PR configuration file was restored after evaluation.

Test Result

25 configuration cases passed on GB200. Applicable pre-commit hooks and mypy 3.12 passed. Coverage includes eligible selection, explicit overrides, hardware/dependency/topology guards and unsupported model/parallel configurations.

Real checkpoint: deepseek-ai/DeepSeek-V4.1-Flash, revision dba1be0a40aa45a94ad051997016db3960a90277, four GB200 GPUs, TP4/DP1/EP4, max_model_len=1048576, FULL decode graphs, lm-eval 0.4.12. Full GSM8K, five-shot multi-turn chat, thinking disabled, greedy generation, 1024 output tokens, seed 123.

Selection Correct / 1319 (flexible extraction) Accuracy Strict-match accuracy
Parent auto → FlashInfer TRT-LLM MXFP4/MXFP8 1272 96.4367% 79.8332%
Proposed auto → DeepGEMM MegaMoE 1269 96.2092% 78.3927%

All 1319 document, prompt and target hashes match. Eight baseline errors became correct and eleven baseline-correct answers became errors: a net difference of three questions in this single paired run. This does not establish statistical equivalence across workloads.

Both runs retain upstream attention/indexer defaults: the server selected FLASHMLA_MEGA_ATTN_DSV41 with the NVFP4 attention cache format and the FP8 indexer. No separate attention, indexer or padding patches were included. The 1M model limit does not mean GSM8K tested 1M-token prompts. Startup and graph capture succeeded for both arms.

Validated source commit cac37db059, parent 1085b64425, with precompiled CUDA/runtime assets from that exact parent commit. The initial checkout reused older runtime assets and failed before serving; the results above come from the matched runtime after those setup failures were resolved.

Memory tradeoff: these TP4/EP4 runs reported 81.46 GiB/GPU allocated during model loading for MegaMoE versus 71.03 GiB/GPU for the former default. This is the model-loading allocation reported by vLLM, not total serving peak memory; automatic MegaMoE selection can leave less room for KV cache. Source accounting attributes about 7.03 GiB/GPU to routed FP4 weight padding (2304 → 2560), 1.32 GiB/GPU to retained original shared weights alongside padded fused copies, and 0.99 GiB/GPU to shared-weight replication versus TP4 sharding, before scale/layout differences. #57204 removes the first two causes; the full 10.43 GiB gap should not be interpreted as intrinsic MegaMoE overhead. These component estimates are calculated from checkpoint dimensions, not a runtime tensor census.

AI assistance

Codex assisted with implementation, validation, and this description. Woosuk Kwon requested the default change after reviewing the TEP backend-selection and sequence-parallel behavior.

Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@mergify mergify Bot added the DSv4.1 Related to DeepSeek-V4.1 models label Sep 16, 2026
@WoosukKwon WoosukKwon closed this Sep 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

DSv4.1 Related to DeepSeek-V4.1 models

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant