[simon-veitner-redhat]

Requested by Simon

Overview

This PR enables the deep_gemm_mega_moe MoE backend on SM90 (H200) for MXFP4 routed experts. It uses DeepGEMM's new SM90 FP8 x FP4 Mega MoE kernel, which fuses dispatch, L1 GEMM, SwiGLU with FP8 quant, L2 GEMM and combine into one persistent kernel over NVLink symmetric memory.

Depends on vllm-project/DeepGEMM#44. The DeepGEMM pin in cmake/external_projects/deepgemm.cmake must be bumped once that PR lands. Until then the SM90 test cases skip.

Claims

• moe_backend=deep_gemm_mega_moe works on SM90 with MXFP4 experts, with the same functionality as SM100 within one NVLink domain.
• On DeepSeek-V4.1-Flash, single node DP8+EP8, output throughput is 17-26% higher than main's default MoE path and 12-19% higher than main with deepep_v2, at concurrency 1 to 1024.
• On two nodes, as two DP8+EP8 replicas, output throughput is 16-27% higher than main's default and 10-16% higher than main with deepep_v2.
• In a 2-node P/D AgentX setup, ITL p50 drops by 10-11% and TTFT p50 by 12-17%.
• Greedy generations are coherent and agree in content with main's default backend.
• SM90 activations are staged with one FP32 scale per 128 channels, the format DeepGEMM#44's SM90 kernel now takes. Against the earlier revision of both PRs this is 1-2% faster under load, and GSM8K is unchanged.

Validation

All runs on 8x H200, DeepSeek-V4.1-Flash, MXFP4 experts.

Tests

• tests/kernels/moe/test_deep_gemm_mega_moe.py, tests/models/test_deepseek_v4_mega_moe.py on the final revision: 120 passed, 12 skipped (the skips are SM100-only NVFP4 cases). They include the FP32 staging cases, bit-exact against DeepGEMM's per_token_cast_to_fp8.
• tests/model_executor/test_deep_gemm_warmup.py and the FlashInfer EP tests: 29 passed.
• CI: these files run only in the B200 "DeepSeek V4 Kernels" step, so CI covers SM100 only. The SM90 cases skip unless DeepGEMM has the SM90 kernel, so they need the pin bump and an H200 step; the results above are from local H200 runs.

Final revision vs the measured one

The final revision (FP32 activation scales and the review fixes) was compared end to end with the revision measured below, on one node with three fresh servers each: +1.2% output tok/s at c=1024, +1.9% total tok/s and −2.2% mean TTFT on a prefill-heavy load (4096 in / 16 out, c=256), and +0.5% TPOT at c=1 (8.55 vs 8.51 ms). GSM8K (lm_eval, 5-shot, 1319 questions) is 92.93 vs 92.72, equal within the ±0.5 pt resolution. The throughput tables below were measured on the earlier revision (see Reproduction).

Throughput, single node DP8+EP8 (random 1024 input / 256 output, prompts = max(32, 4 x concurrency))

Cells are output tok/s / mean TPOT (ms). main default is HUMMING with block-FP8 activations (#61875). Mega MoE and main default use allgather_reducescatter. main deepep_v2 is main's fastest all2all backend here. Each cell is the mean of two runs, one on each of two nodes at the same time; the two runs agree within 7%, and most within 2%. The exception is deepep_v2 at c=16, which is one run because the other run had a one-off 1.4 s TTFT stall.

Concurrency	Mega MoE	main default	main deepep_v2	Δ tok/s vs default	Δ tok/s vs deepep_v2	Δ TPOT vs default	Δ TPOT vs deepep_v2
1	115 / 8.4	92 / 10.4	97 / 10.0	+26%	+19%	-19%	-16%
16	1359 / 10.8	1102 / 13.1	1156 / 12.5	+23%	+18%	-18%	-14%
64	4029 / 13.9	3402 / 16.3	3513 / 15.8	+18%	+15%	-15%	-12%
256	9424 / 22.7	7640 / 28.1	7988 / 26.7	+23%	+18%	-19%	-15%
512	13380 / 32.0	11478 / 37.6	11953 / 35.9	+17%	+12%	-15%	-11%
1024	18566 / 45.9	14733 / 59.9	16013 / 53.4	+26%	+16%	-23%	-14%

Throughput, two nodes (2x8 H200, same workload, concurrency is the total)

Mega MoE needs EP within one NVLink domain, so every arm runs as two DP8+EP8 replicas with half the requests each. Cells are the summed output tok/s / mean TPOT (ms), one run per row. Main as one DP16+EP16 server with deepep_v2 internode was not measured.

Concurrency	Mega MoE	main default	main deepep_v2	Δ tok/s vs default	Δ tok/s vs deepep_v2	Δ TPOT vs default	Δ TPOT vs deepep_v2
256	13191 / 16.9	10938 / 20.2	11439 / 19.3	+21%	+15%	-16%	-12%
512	20238 / 21.5	16565 / 26.5	17503 / 24.8	+22%	+16%	-19%	-13%
1024	28716 / 30.0	24847 / 34.9	26004 / 32.9	+16%	+10%	-14%	-9%
2048	39740 / 42.9	31253 / 56.0	34297 / 49.7	+27%	+16%	-23%	-14%

AgentX P/D, 2 nodes

With the prefill and decode roles on separate nodes, Mega MoE beats main HUMMING + deepep_v2 by +1.5% / +4.3% / +4.0% tok/s/GPU at c=128 / 256 / 512. It also cuts ITL p50 by 10-11% and TTFT p50 by 12-17%. The agentic trace has idle gaps, and they cap throughput. Harness: agentx-harness at 89b2186 (a public fork of SemiAnalysisAI/agentx-harness), scenario agentx, dataset semianalysis_cc_traces_weka_062126.

c	Mega MoE tok/s/GPU	main tok/s/GPU	ITL p50 ms (mega / main)	TTFT p50 s (mega / main)
128	248.9	245.2	6.06 / 6.76	0.68 / 0.77
256	485.3	465.4	7.56 / 8.51	0.90 / 1.08
512	894.0	859.6	9.54 / 10.69	1.31 / 1.51

Reproduction

Mega MoE is vLLM main 9633997 with this PR applied and DeepGEMM at the SM90 PR head (f4685f2), the revision before the FP32 activation scales. main is vLLM main 9633997.

# Mega MoE
vllm serve deepseek-ai/DeepSeek-V4.1-Flash --tokenizer-mode deepseek_v41 --trust-remote-code --max-model-len 8192 \
  --enable-expert-parallel --data-parallel-size 8 --moe-backend deep_gemm_mega_moe --all2all-backend allgather_reducescatter
# main default: the same command without --moe-backend
# main deepep_v2:
EP_DISABLE_GIN=1 VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE=0 VLLM_HUMMING_MOE_GEMM_TYPE=indexed \
vllm serve deepseek-ai/DeepSeek-V4.1-Flash --tokenizer-mode deepseek_v41 --trust-remote-code --max-model-len 8192 \
  --enable-expert-parallel --data-parallel-size 8 --all2all-backend deepep_v2

# Client, after a 256-prompt warm-up (two nodes: one client per replica at concurrency C/2)
vllm bench serve --model deepseek-ai/DeepSeek-V4.1-Flash --tokenizer-mode deepseek_v41 --trust-remote-code \
  --dataset-name random --random-input-len 1024 --random-output-len 256 \
  --num-prompts $((C < 8 ? 32 : 4 * C)) --max-concurrency $C --seed 0

Correctness

Greedy samples on five prompts (temperature 0, max 4096 tokens) from Mega MoE and main default, single node DP8+EP8. All outputs are coherent and finish normally. Each pair reaches the same answer and differs only in wording and small details; for example, main's code returns tuples. On the multi-turn prompt, Mega MoE recommends Day 3 and main recommends Day 2, and each calls the other day a close second. On the final revision, 10 further prompts (each sent greedy twice and sampled twice, on three fresh servers) all finish with correct answers, the same as on the earlier revision.

Generation samples (final answers; reasoning omitted, long answers cut after the first paragraph)

Math. Solve step by step: A train leaves at 3:15 pm traveling 84 km/h. Another leaves the same station at 4:00 pm traveling 105 km/h on the same track in the same direction. At what time does the second train catch the first? Give the final answer on its own line.

Mega MoE:

The first train has a 45-minute head start.

• 45 minutes = 0.75 hours
• Distance first train travels before 4:00 pm:
  (84 \times 0.75 = 63) km
• Relative speed of second train:
  (105 - 84 = 21) km/h
• Time for second train to catch up:
  (63 \div 21 = 3) hours

So, 3 hours after 4:00 pm is 7:00 pm.

7:00 pm

main default:

The first train has a 45-minute head start.

• 45 minutes = 0.75 hours
• Distance first train travels before second starts:
  (84 \times 0.75 = 63) km
• Relative speed of second train compared to first:
  (105 - 84 = 21) km/h
• Time for second train to catch up:
  (63 \div 21 = 3) hours

The second train leaves at 4:00 pm, so it catches the first at:

7:00 PM

Code. Write a Python function merge_intervals(intervals) that merges overlapping closed intervals and returns them sorted. Include a short docstring.

Mega MoE:

def merge_intervals(intervals):
    """Merge overlapping closed intervals and return them sorted by start."""
    if not intervals:
        return []

    intervals = sorted(intervals, key=lambda x: x[0])
    merged = [list(intervals[0])]

    for start, end in intervals[1:]:
        if start <= merged[-1][1]:  # closed intervals overlap at endpoints
            merged[-1][1] = max(merged[-1][1], end)
        else:
            merged.append([start, end])

    return merged

main default:

def merge_intervals(intervals):
    """Merge overlapping closed intervals and return them sorted by start."""
    if not intervals:
        return []

    intervals = sorted(intervals, key=lambda x: x[0])
    merged = [list(intervals[0])]

    for start, end in intervals[1:]:
        last_start, last_end = merged[-1]

        # Closed intervals overlap if start <= last_end
        if start <= last_end:
            merged[-1][1] = max(last_end, end)
        else:
            merged.append([start, end])

    return [tuple(interval) for interval in merged]

Factual. Why is the sky blue but sunsets red? Explain in four sentences.

Mega MoE:

Sunlight contains all colors, and air molecules scatter shorter blue wavelengths much more strongly than longer red ones. During the day, this scattered blue light reaches our eyes from all parts of the sky, making it look blue. At sunset, sunlight passes through more atmosphere, so most blue light is scattered away and the remaining light appears red or orange. The same Rayleigh scattering process therefore creates a blue daytime sky and red sunsets.

main default:

Sunlight contains all colors, but air molecules scatter short wavelengths like blue much more strongly than long wavelengths like red. During the day, this scattered blue light reaches our eyes from all parts of the sky, making the sky appear blue. At sunset, sunlight passes through a much longer path of atmosphere, so most blue light is scattered away before it reaches us. The remaining light that travels more directly to our eyes is dominated by longer wavelengths—reds and oranges—so the Sun and sky near the horizon look red.

Multi-turn chat. Three turns. User: "I am planning a 3-day trip to Kyoto in November. Suggest one main sight per day." Assistant: "Day 1: Fushimi Inari Taisha. Day 2: Arashiyama bamboo grove and Tenryu-ji. Day 3: Kiyomizu-dera and the Higashiyama streets." User: "Which of those days is best for autumn leaves, and what time of day should I go to avoid crowds?"

Mega MoE:

Best day for autumn leaves: Day 3 — Kiyomizu-dera + Higashiyama is the strongest overall for classic Kyoto autumn foliage. Day 2 — Arashiyama/Tenryu-ji is a very close second and arguably better if you want river, bamboo, and garden scenery. Day 1 — Fushimi Inari is not primarily a autumn-leaves destination.

[...]

main default:

Best day for autumn leaves: Day 2 — Arashiyama + Tenryu-ji.
Tenryu-ji’s garden and the Arashiyama area are classic Kyoto momiji spots. Day 3 — Kiyomizu-dera + Higashiyama is a very close second and can be spectacular in late November. Day 1 — Fushimi Inari is mainly about the torii gates, not autumn leaves.

[...]

Long-form (about 300 words). Write about 300 words, for a first-year computer science student, on how a hash table works: hashing, collisions (chaining and open addressing), load factor and resizing.

Mega MoE:

A hash table stores key-value pairs and aims for average O(1) lookup, insert, and delete. It uses a hash function to turn a key—such as a student ID or string—into an integer, then maps that integer to an array index: index = hash(key) % table_size. A good hash function spreads keys evenly and is fast to compute.

[...]

main default:

A hash table stores key-value pairs so you can find, insert, and delete items quickly—usually O(1) on average. It uses a hash function to turn a key (like a student ID or string) into an integer, then maps that integer to an index in an array of buckets: index = hash(key) % table_size. A good hash function is fast and spreads keys uniformly. Ideally, different keys get different indices, so lookup is just one array access.

[...]

Kernel time (8 ranks)

DeepSeek-V4.1-Flash shape, slowest rank, median of 3 runs. T=8192 reaches about 675 TFLOPS per rank. The kernel's decode tuning cuts kernel time by 2.0-2.7% at T=1-128 and by 8.5% at T=512 against the untuned kernel; prefill sizes are unchanged. With FP32 activation scales, the final revision is a further 0.5-0.8% faster at T=1-512 and 4.1% faster at T=8192 than the revision in this table (mean over ranks, 3 interleaved reps).

Tokens per rank	BLOCK_M	Mega MoE (µs)
1	16	85.2
8	16	212.7
32	16	282.4
128	32	321.2
512	64	570.6
2048	128	1693
8192	128	6015

Details

• deep_gemm_mega_moe.py: on SM90, activations are dispatched as FP8 with one FP32 scale per 128 channels (MEGA_MOE_MXFP4_SM90). prepare_megamoe_inputs computes it as amax * (1 / 448) with a correctly rounded inverse, so scales and values match DeepGEMM's per_token_cast_to_fp8 bit for bit, and stores its bits in the int32 x_sf buffer and in the shared-expert scale layout. SM100 keeps UE8M0. The weight loader relayouts routed FP4 shards and repacks their scales with DeepGEMM's SM90 transforms. Fused FP8 shared experts are requantized once all their shards are loaded.
• BLOCK_M hint: every rank of a launch passes num_tokens_for_block_m, the mean tokens per EP rank of the step. All ranks then pick the same BLOCK_M, so idle DP ranks do not run small blocks next to busy ones. Without the hint, ragged steps are 1.5-1.9x slower on the slowest rank.
• DeepSeek V4 routes with GateLinear and fused_topk_bias off SM100, because bf16_mega_gate does not exist there. Warmup skips those configs on SM90.
• SM100 behaviour is unchanged.

Follow-ups

• Bump the DeepGEMM pin after the DeepGEMM PR merges.
• A visual walkthrough of the kernel is available on request.

---

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.

This PR was prepared with Claude Code and reviewed by a human (@simon-veitner-redhat).

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

[kabishou11]

We mapped this onto an offline DeepInfer Hopper engine that already runs MXFP4 experts through DeepEP V2 + humming (gs32), not through DeepGEMM Mega MoE.

Our vendor pin is deepseek-ai/DeepGEMM @ 1fc9c66; a CPU-safe probe there still reports SM100 fp8_fp4_mega_moe only, no sm90_fp8_fp4_mega_moe kernel files, and we reject A/B until that SM90 path is actually in the pin. We are treating the DP8+EP8 throughput comparison in this PR as an author claim for the vLLM stack, not as a number to copy, and we will not bump our DeepGEMM pin without a humming baseline JSONL on H200.

Question: once vllm-project/DeepGEMM#44 lands, is the SM90 kernel expected to upstream into deepseek-ai/DeepGEMM, or should downstreams plan on the vLLM fork pin indefinitely?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @simon-veitner-redhat.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork