[Perf][DSV4.1] Pad shared experts for native MegaMoE fusion

[gcanlin]

Purpose

#56217

DeepSeek-V4.1-Flash's 2304-wide shared expert cannot fuse with MegaMoE's routed experts padded to 2560. Pad shared weights with zeros and unit scales to enable fusion in all 40 layers, preserving checkpoint-shaped parameters and the unsupported-layout fallback.

Extends #53040. Checks of #45861 and related open PRs found no fix for this padding gap; #53567 and #54049 cover different paths.

AI assistance was used for implementation, testing, and documentation.

Test Plan

Deploy

Use a fresh server per arm: VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION=1 for serial shared experts, 0 for fused. The omitted load-time hook verifies fusion state on all 8 ranks.

export MODEL=/path/to/DeepSeek-V4.1-Flash
export VLLM_USE_V2_MODEL_RUNNER=1
export GLOO_SOCKET_IFNAME=bond0 NCCL_SOCKET_IFNAME=bond0 NCCL_DEBUG=WARN
export VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION=0
uv run --active --no-project python -m vllm.entrypoints.cli.main serve "$MODEL" \
  --served-model-name dsv41-bench --host 127.0.0.1 --port 8032 \
  --tensor-parallel-size 8 --enable-expert-parallel --distributed-executor-backend mp \
  --moe-backend deep_gemm_mega_moe \
  --tokenizer-mode deepseek_v41 --reasoning-parser deepseek_v41 \
  --attention-backend FLASHINFER_MLA_SPARSE_DSV41 --kv-cache-dtype fp8 --block-size 128 \
  --max-model-len 16384 --max-num-batched-tokens 8192 --max-num-seqs 256 \
  --gpu-memory-utilization 0.9 --seed 2026 --generation-config vllm \
  --no-enable-prefix-caching --no-enable-flashinfer-autotune --enable-prompt-tokens-details \
  --compilation-config '{"cudagraph_capture_sizes":[1,2,4,8,16,32,64,128,256],"max_cudagraph_capture_size":256}'

Benchmark

Workloads follow #53040, plus 8192/1024 at C16. V4.1 uses block size 128 and FP8 indexer cache.

# Use a separate result directory for each arm.
export RESULT_DIR=results/dsv41-fused
mkdir -p "$RESULT_DIR"
bench() {
  uv run --active --no-project python -m vllm.entrypoints.cli.main bench serve \
    --backend vllm --base-url http://127.0.0.1:8032 --endpoint /v1/completions \
    --model dsv41-bench --tokenizer "$MODEL" --tokenizer-mode deepseek_v41 \
    --trust-remote-code --dataset-name random \
    --random-input-len "$1" --random-output-len "$2" \
    --max-concurrency "$3" --num-prompts "$4" \
    --random-range-ratio 0 --random-prefix-len 0 --request-rate inf \
    --ignore-eos --temperature 0 --seed 42 --num-warmups 2 \
    --percentile-metrics ttft,tpot,itl,e2el --metric-percentiles 50,95,99 \
    --save-result --save-detailed --disable-tqdm \
    --result-dir "$RESULT_DIR" --result-filename "$5.json"
}
for repeat in 1 2 3; do bench 128 256 1 8 "decode-128-256-c1-r$repeat"; done
for repeat in 1 2 3; do bench 1024 128 1 8 "decode-1024-128-c1-r$repeat"; done
for repeat in 1 2 3; do bench 1024 128 64 128 "balanced-1024-128-c64-r$repeat"; done
for repeat in 1 2 3; do bench 8192 32 16 32 "prefill-8192-32-c16-r$repeat"; done
bench 8192 1024 16 100 kimi-8192-1024-c16-r1

GSM8K evaluation reproduction

Use lm-eval 0.4.12, the same fusion flag, and a fresh server per arm.

export MODEL=/path/to/DeepSeek-V4.1-Flash
export VLLM_USE_V2_MODEL_RUNNER=1
export GLOO_SOCKET_IFNAME=bond0 NCCL_SOCKET_IFNAME=bond0
export NCCL_DEBUG=WARN
export VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION=0
uv run --active --no-project python -m vllm.entrypoints.cli.main serve "$MODEL" \
  --served-model-name dsv41-gsm8k --host 127.0.0.1 --port 8031 \
  --tensor-parallel-size 8 --enable-expert-parallel --distributed-executor-backend mp \
  --moe-backend deep_gemm_mega_moe \
  --tokenizer-mode deepseek_v41 --reasoning-parser deepseek_v41 \
  --attention-backend FLASHINFER_MLA_SPARSE_DSV41 --kv-cache-dtype fp8 \
  --max-model-len 16384 --max-num-batched-tokens 8192 --max-num-seqs 96 \
  --gpu-memory-utilization 0.9 --no-enable-prefix-caching \
  --no-enable-flashinfer-autotune --generation-config vllm --seed 123 \
  --default-chat-template-kwargs '{"thinking": false}'

Run the full test split in the lm-eval environment, with separate output directories:

OPENAI_API_KEY=EMPTY uv run --active --no-project python -m lm_eval run \
  --model local-chat-completions \
  --model_args model=dsv41-gsm8k,base_url=http://127.0.0.1:8031/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 results/gsm8k-fused

Test Result

31 tests passed on B200. Ruff, typos, mypy 3.10/3.12, and applicable pre-commit checks passed; nonmatching Actionlint and requirements-compilation hooks were skipped.

All 8 ranks confirmed 0 → 40 fused layers. Model-load memory increased from 41.47 to 43.34 GiB/rank (+1.87 GiB) because the original shared parameters remain allocated.

8×B200, TP8/EP/SP, FP8 KV, 16K context. vllm bench serve: 2 warmups/run, median of 3 runs except 8192/1024 (1 run), no profiler. All 1,256 requests succeeded with the requested token counts. Both arms used source 30aa0ade39ab6625674865f578fae19a49cf3e43 + this patch, native build e7edf17cea217e52701f913cd8491fcacf2d9490, and PyTorch 2.13.0+cu130.

Input / output	Concurrency	Requests × runs	Output tok/s: serial → fused	Change	Mean TPOT ms: serial → fused	Change
128 / 256	1	8 × 3	122.6 → 134.3	+9.6%	8.115 → 7.402	-8.8%
1024 / 128	1	8 × 3	109.2 → 121.0	+10.9%	8.178 → 7.455	-8.8%
1024 / 128	64	128 × 3	2853.9 → 3359.2	+17.7%	16.851 → 14.567	-13.6%
8192 / 32	16	32 × 3	192.7 → 226.1	+17.3%	48.804 → 41.806	-14.3%
8192 / 1024	16	100 × 1	1292.1 → 1428.8	+10.6%	10.449 → 9.449	-9.6%

Input / output, concurrency	Mean TTFT ms: serial → fused	Change	Median ITL ms: serial → fused	Change
128 / 256, C1	18.57 → 17.80	-4.1%	8.107 → 7.392	-8.8%
1024 / 128, C1	133.85 → 110.82	-17.2%	8.162 → 7.439	-8.9%
1024 / 128, C64	706.90 → 558.69	-21.0%	11.038 → 10.285	-6.8%
8192 / 32, C16	1109.92 → 937.88	-15.5%	9.170 → 8.387	-8.5%
8192 / 1024, C16	912.06 → 815.41	-10.6%	9.082 → 8.357	-8.0%

Negative latency changes mean lower latency. C1 has only 8 requests/run, so tail-latency estimates are limited.

C64's first run hit quantization JIT in both arms. Excluding it from both changes the throughput gain from +17.7% to +16.6%.

GSM8K metric	Serial shared	Fused shared	Change (percentage points)
flexible-extract	96.21% (1269/1319)	96.13% (1268/1319)	-0.076
strict-match	78.17% (1031/1319)	79.23% (1045/1319)	+1.061

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

[gcanlin]

cc @WoosukKwon @zyongye. This PR is for DSv4.1 specific shared expert dim 2304. Padding to 2560 can enable the native shared expert fusion.