[DSV4][Kernel] Fuse shared experts into MegaMoE

[gcanlin]

Summary

Part of #45861.

This PR fuses DeepSeek V4's replicated FP8 shared expert into DeepGEMM's persistent SM100 MegaMoE kernel in the NVIDIA-specific model path.

Before this change, every MoE layer launched the routed FP4 MegaMoE kernel and then ran the shared FP8 gate/up and down projections serially, followed by a separate add. The new path lets the native SM100 kernel schedule shared L1, routed dispatch/MMA, shared L2, and the final FP32 accumulation together. It retains the checkpoint's FP8 shared weights and produces one BF16 output store.

Test Plan

vllm serve /mnt/models/deepseek-ai/DeepSeek-V4-Flash-0731 \
  --served-model-name dsv4 \
  --host 127.0.0.1 --port 8000 \
  --tensor-parallel-size 8 --enable-expert-parallel \
  --moe-backend deep_gemm_mega_moe \
  --tokenizer-mode deepseek_v4 \
  --kv-cache-dtype fp8 --block-size 256 \
  --attention-config.indexer_kv_dtype=mxfp4 \
  --max-model-len 16384 --max-num-batched-tokens 8192 \
  --gpu-memory-utilization 0.9 --seed 2026 \
  --compilation-config \
  '{"cudagraph_capture_sizes":[1,2,4,8,16,32,64,128,256],"max_cudagraph_capture_size":256}'

Test Results

Batch-size-1 latency and throughput

Workload and metric	Baseline	Fused	Change
128 input / 256 output: output throughput	130.02 tok/s	149.50 tok/s	+14.98%
128 input / 256 output: mean TPOT	7.653 ms	6.643 ms	-13.19%
128 input / 256 output: median ITL	7.577 ms	6.567 ms	-13.32%
128 input / 256 output: mean TTFT	17.32 ms	18.17 ms	+4.93%
1024 input / 128 output: output throughput	125.19 tok/s	142.32 tok/s	+13.68%
1024 input / 128 output: mean TPOT	7.158 ms	6.295 ms	-12.07%
1024 input / 128 output: median ITL	7.090 ms	6.223 ms	-12.23%
1024 input / 128 output: mean TTFT	113.13 ms	99.81 ms	-11.77%

Balanced workload: 128 requests, 1024 input, 128 output, concurrency 64

Metric	Baseline	Fused	Change
Output throughput	3,415.45 tok/s	3,737.98 tok/s	+9.44%
Total token throughput	30,739.04 tok/s	33,641.81 tok/s	+9.44%
Mean TTFT	588.52 ms	530.15 ms	-9.92%
Mean TPOT	14.095 ms	12.935 ms	-8.23%
Median ITL	9.644 ms	8.939 ms	-7.31%

Prefill-heavy workload: 32 requests, 8192 input, 32 output, concurrency 16

Metric	Baseline	Fused	Change
Output throughput	234.09 tok/s	249.42 tok/s	+6.55%
Total token throughput	60,160.73 tok/s	64,099.90 tok/s	+6.55%
Mean TTFT	901.81 ms	836.36 ms	-7.26%
Mean TPOT	40.584 ms	38.458 ms	-5.24%
Median ITL	8.256 ms	7.493 ms	-9.24%

Accuracy

  |Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
  |-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
  |gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9484|±  |0.0061|
  |     |       |strict-match    |     5|exact_match|↑  |0.9484|±  |0.0061|

Log

One of them

INFO 08-20 11:04:59 [utils.py:90] Sampling input_len from [128, 128] and output_len from [256, 256]
Starting initial single prompt test run...
Skipping endpoint ready check.
Warming up with 2 requests...
Warmup run completed.
Starting main benchmark run...
Traffic request rate: inf
Burstiness factor: 1.0 (Poisson process)
Maximum request concurrency: 1
tip: install termplotlib and gnuplot to plot the metrics
============ Serving Benchmark Result ============
Successful requests:                     8         
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  13.70     
Total input tokens:                      1024      
Total generated tokens:                  2048      
Request throughput (req/s):              0.58      
Output token throughput (tok/s):         149.52    
Peak output token throughput (tok/s):    151.00    
Peak concurrent requests:                2.00      
Total token throughput (tok/s):          224.27    
---------------Time to First Token----------------
Mean TTFT (ms):                          17.68     
Median TTFT (ms):                        17.68     
P99 TTFT (ms):                           18.01     
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          6.64      
Median TPOT (ms):                        6.64      
P99 TPOT (ms):                           6.65      
---------------Inter-token Latency----------------
Mean ITL (ms):                           6.64      
Median ITL (ms):                         6.57      
P99 ITL (ms):                            6.95      
==================================================

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

@WoosukKwon @zyongye @yewentao256 Could you please take a look when you have time? This is an optimization for DSV4 about MegaMoE on sm100. Thanks!

[WoosukKwon]

/ci run

[github-actions]

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

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• /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.

[github-actions]

✅ Triggered Buildkite CI #84861 for commit 8b598c61b3ad.

[WoosukKwon]

Thanks for the PR! It's regretful that we didn't ship this earlier 😭