Repository navigation
[Perf][DSV4.1] Prefer MegaMoE for supported local expert parallelism - #57220
Closed
WoosukKwon wants to merge 1 commit into
Closed
WoosukKwon wants to merge 1 commit into
WoosukKwon wants to merge 1 commit into
Conversation
Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
For DeepSeek V4.1, resolve
--moe-backend autotodeep_gemm_mega_moefor supported local expert-parallel deployments. Previouslyautowent 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
Real-checkpoint evaluation commands
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-resultsRepeat on parent commit
1085b64425with--moe-backend autofor 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, revisiondba1be0a40aa45a94ad051997016db3960a90277, 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.auto→ FlashInfer TRT-LLM MXFP4/MXFP8auto→ DeepGEMM MegaMoEAll 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_DSV41with 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, parent1085b64425, 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.