Repository navigation
[Perf] Warm up DeepGEMM MegaMoE kernels before serving - #60468
Conversation
84d0f2a to
40ad220
Compare
| ) | ||
|
|
||
|
|
||
| def deepgemm_fp8_fp4_mega_moe_warmup( |
There was a problem hiding this comment.
warp with torch.inference_mode()
There was a problem hiding this comment.
Done in ba3bd55. deepgemm_fp8_fp4_mega_moe_warmup is now decorated with @torch.inference_mode(), like the other warmups that run a full module forward. Nothing up the call chain enables it, and on Kimi K3 the routed-expert RMSNorm weight requires grad, so the warmup was building autograd graphs that serving never does.
This reply was generated with Claude and reviewed by a human.
| experts.top_k, | ||
| "fp8xfp4", | ||
| ) | ||
| m_values.setdefault((tuple(gate_config.items()), block_m), num_tokens) |
There was a problem hiding this comment.
The PR reports +247 s on a cold start (+299 s with MTP), mostly from about 58 gate configs. There is arelax mode which trims M values for the other DeepGEMM warmups.
I think a cheap relax behavior can be interesting here too, maybe:
key = (block_m,) if envs.VLLM_DEEP_GEMM_WARMUP == "relax" else (tuple(gate_config.items()), block_m)
m_values.setdefault(key, num_tokens)There was a problem hiding this comment.
I looked into this, but I don't think relax fits here. In this file relax reduces the number of M values launched while still reaching every kernel config; it doesn't reduce the number of kernels compiled. The mega MoE M set is already one M per distinct (gate_config, block_m) kernel, so there is nothing left for relax to trim and full == relax for this path.
The cold cost is almost entirely JIT: ~136 bf16_mega_gate compiles take ~229 s, while the layer forwards take under 1 s (a warm restart runs the whole bar in under 1 s). A block_m-only key (at most ~6 values) would skip ~120 gate compiles, and since relax is the default, serving would pay them at ~1.7 s each on first requests, which is exactly the TTFT stall this PR removes (burst max TTFT ~3.7–4.0 s on main vs ~0.27 s with the PR). The cost is paid once per empty DeepGEMM JIT cache, and VLLM_DEEP_GEMM_WARMUP=skip already opts out.
This reply was generated with Claude and reviewed by a human.
| ) | ||
| is_padding = torch.zeros(max(m_values) + 1, device=device, dtype=torch.bool) | ||
| # Sequence parallel gives ranks > 0 a padding mask slice at any alignment. | ||
| offsets = (0, 1) if module.use_sequence_parallel else (0,) |
There was a problem hiding this comment.
nit: I think the offset loop doubles work when padding is off.
| offsets = (0, 1) if module.use_sequence_parallel else (0,) | |
| offsets = (0, 1) if module.use_sequence_parallel and envs.VLLM_MOE_SKIP_PADDING else (0,) |
There was a problem hiding this comment.
Done in ba3bd55. Offset 1 now only runs when sequence parallel and VLLM_MOE_SKIP_PADDING are both on; with padding skip off the expert path gets is_padding=None, so both offsets launched identical kernels. I split the condition into a sp_masked variable to stay under the 88-column limit.
This reply was generated with Claude and reviewed by a human.
| dsv4 = sys.modules["vllm.models.deepseek_v4.nvidia.model"] | ||
| return isinstance(module, dsv4.DeepseekV4MoE) |
There was a problem hiding this comment.
dsv4 = sys.modules.get("vllm.models.deepseek_v4.nvidia.model")
return dsv4 is not None and isinstance(module, dsv4.DeepseekV4MoE)nit: same as _mega_moe_may_use_deep_gemm?
There was a problem hiding this comment.
Done in ba3bd55. _mega_moe_uses_mega_gate now uses the same guarded sys.modules.get(...) lookup, and _mega_moe_may_use_deep_gemm reuses it instead of repeating the DSV4 check. The hard index was only safe because Kimi's model module happens to import DSV4 at top level.
This reply was generated with Claude and reviewed by a human.
| elif _mega_moe_may_use_deep_gemm(m): | ||
| key = _get_fp8_fp4_mega_moe_warmup_key(m) | ||
| if key not in seen_mega_moe_keys: | ||
| total += len(_get_fp8_fp4_mega_moe_m_values(m, max_tokens)) |
There was a problem hiding this comment.
nit: I think it's cheap, but it seems like the enumeration runs twice per key? _get_fp8_fp4_mega_moe_m_values runs once during counting (here) and again during the warmup itself. A single result could be stored and reused, if it is worth it.
There was a problem hiding this comment.
True, it runs twice (once for the progress-bar total, once for the warmup). I measured ~41 ms per key at 8192 tokens on B300 (~5 µs/token), next to seconds of JIT per kernel. _count_warmup_iterations on main already recomputes the fp8_gemm_nt and grouped GEMM lists the same way, so I kept this path consistent rather than adding a cache for it alone. Happy to memoize all three in a follow-up if you'd like.
This reply was generated with Claude and reviewed by a human.
| max_tokens = cdiv(max_tokens, module.tp_size) | ||
|
|
||
| m_values: dict[tuple, int] = {} | ||
| for num_tokens in range(1, max_tokens + 1): |
There was a problem hiding this comment.
My agent flagged a potential issue. Could you double-check whether it’s actually valid? It seems plausible to me, which is exactly why I’m not a fan of these token-based warmups... :))
DeepSeek-V4 only calls bf16_mega_gate above a hardcoded threshold, at model.py:1043-1044:
gate_threshold = 1 if self.gate.weight.shape[0] == 128 else 16
if hidden_states.shape[0] <= gate_threshold:
router_logits, _ = self.gate(hidden_states) # GateLinear + fused_topk_bias
...
else:
... bf16_mega_gate(...)The warmup side. It picks the smallest token count for each config key, at deep_gemm_warmup.py:367-383:
for num_tokens in range(1, max_tokens + 1):
gate_config = (dg.get_bf16_mega_gate_config(num_tokens, ...) if uses_mega_gate else {})
block_m = dg.get_block_m_for_mega_moe(...)
m_values.setdefault((tuple(gate_config.items()), block_m), num_tokens)How the gap happens. Suppose a (gate_config, block_m) key spans token counts 1..N with N > 16. Its representative is then 1. The forward at deep_gemm_warmup.py:433 takes the GateLinear branch, so bf16_mega_gate is never compiled for that config. On DSv4-Flash with EP2 (256 experts, top-6), block_m=16 covers about 1–181 tokens. Any gate config that is constant across 16→17 would be missed. CUDA-graph capture probably compiles most of those sizes, so the leftover gap is eager sizes above the capture limit.
Fix. Make “mega-gate is actually reached” part of the key, and stop hardcoding the threshold in two places:
# model.py — expose it
def mega_gate_threshold(self) -> int:
return 1 if self.gate.weight.shape[0] == 128 else 16
# deep_gemm_warmup.py
threshold = module.mega_gate_threshold() if uses_mega_gate else 0
for num_tokens in range(1, max_tokens + 1):
uses_gate_kernel = uses_mega_gate and num_tokens > threshold
gate_config = dg.get_bf16_mega_gate_config(...) if uses_gate_kernel else {}
...
m_values.setdefault((uses_gate_kernel, tuple(gate_config.items()), block_m), num_tokens)There was a problem hiding this comment.
Thanks, verified: the mechanism is real. On current DeepGEMM the gate config changes exactly at 17 tokens, so the 256/384-expert models (threshold 16) don't miss anything, but a 128-expert gate (threshold 1, e.g. the DSpark drafter) misses the config used by sizes 2..16, because its representative is M=1 which takes the GateLinear branch.
Rather than exposing the threshold, ba3bd55 keeps the largest M per (gate_config, block_m) key. The forward condition is num_tokens > threshold, so any config the forward actually passes to bf16_mega_gate has its largest M above the threshold. That avoids duplicating the constant and adds no keys. I swept 288 configs (hidden 4096/7168, 128/256/384 experts, top_k 6/8, EP 1–8, SP) against the real heuristics on B300: largest-M misses 0 bf16_mega_gate/fp8_fp4_mega_moe kernels, while smallest-M missed a gate kernel in all 96 128-expert configs. Added test_mega_moe_m_values_reach_mega_gate for it. Agreed that CUDA-graph capture covers most sizes otherwise; this closes the eager gap. Re-validated DSV4 with/without MTP: no MegaMoE JIT after warmup, GSM8K 0.985.
This reply was generated with Claude and reviewed by a human.
The DeepGEMM warmup covered fp8_gemm_nt and the grouped MoE GEMM, but not fp8_fp4_mega_moe or the bf16_mega_gate that routes into it. Both pick their JIT config from the per-rank token count, so prefill sizes outside the CUDA graph capture sizes compiled on first use while serving: about 1.7 s per kernel, 3.5 to 6 s added to a request. For each native MegaMoE layer (DeepSeek-V4, Kimi K3), ask DeepGEMM for the block_m and, where the layer routes with bf16_mega_gate, the gate config of every per-rank token count, keep one count per distinct config, and run the layer once at each. Running the whole layer instead of the raw kernels gives the gate the same optional inputs as in serving, and every EP rank runs the collective MegaMoE kernel in the same order. The layer runs under a forward context with a padding mask, at both mask alignments under sequence parallelism, so the Triton input staging kernel is compiled for the variants serving uses. Deduplicating on DeepGEMM's own configs is exact, so VLLM_DEEP_GEMM_WARMUP=full would add no kernels. The two config getters come from _import_deep_gemm, as in the model code. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Simon Veitner <sveitner@redhat.com>
40ad220 to
ba3bd55
Compare
|
✅ @simon-veitner-redhat, CI is now available for this PR.
|
LopezCastroRoberto
left a comment
There was a problem hiding this comment.
LGTM, thanks!
|
/ci run |
|
✅ Triggered Buildkite CI #94444 for commit |
…60468) Signed-off-by: Simon Veitner <sveitner@redhat.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Signed-off-by: Andreas Karatzas <AndreasKaratzas@users.noreply.github.com>
…60468) Signed-off-by: Simon Veitner <sveitner@redhat.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com>
Overview
The DeepGEMM warmup covers
fp8_gemm_ntand the grouped MoE GEMM. It does not coverfp8_fp4_mega_moeor thebf16_mega_gatethat routes into it. With--moe-backend deep_gemm_mega_moe, DeepSeek-V4 and Kimi K3 therefore JIT-compile these kernels on the first request at each new prefill size. This PR adds them todeep_gemm_warmup.py, following the structure of the two existing DeepGEMM warmups there.Claims
fp8_fp4_mega_moe,bf16_mega_gateand the Triton input-staging kernel. Before: 60 DeepGEMM compiles in one DeepSeek-V4 benchmark run. For Kimi K3, the 2 SiTUfp8_fp4_mega_moecompiles during serving that cost 2.6-2.9 s of TTFT each go to 0, at about +10 s of startup.Validation
Setup (DeepSeek-V4). DeepSeek-V4-Flash on 2x B300, TP2 + EP2,
--moe-backend deep_gemm_mega_moe,max_num_batched_tokens16384,--max-model-len 16384, prefix caching off, Model Runner V2. vLLM precompiled wheel3ca00a82(cu130) with its bundled DeepGEMM.Each run used a fresh
DG_JIT_CACHE_DIRseeded with every non-mega DeepGEMM kernel, so only the mega kernels start cold.Workload. One request at a time: token-ID prompts to
/v1/completions,max_tokens4, streamed time to first token (TTFT).JIT compiles during serving. Counted as new entries in the DeepGEMM cache dir between "Application startup complete" and the end of the workload, across all kernel families. Triton, CuTeDSL and TileLang compiles were counted with
--jit-monitor-verbose.sm100_bf16_mega_gatesm100_fp8_fp4_mega_moe_prepare_megamoe_inputs_kernelOn main, the Triton row comes from the default (non-verbose) JIT monitor, which logs only the first compile per kernel, so the true count is unknown. The previous revision of this PR, which warmed without a padding mask, compiled 3 variants for just three prompt lengths in a verbose run.
A separate single-GPU check called both kernels for every per-rank token count from 1 to 8192 after the warmup's representative sizes. It produced 0 new cache entries. An independent rerun of the patched server, with 20 prompt lengths, 16 of them chosen to give unaligned per-rank chunks, saw 0 JIT events.
TTFT, first request vs median of 3 repeats (ms):
Startup,
init engine(profile, KV cache, warmup, graph capture):KV cache sizing is unchanged: 163.13 GiB and 465,145 tokens in both.
MTP (DeepSeek-V4,
--speculative_config.method=mtp --speculative_config.num_speculative_tokens=2, same setup and workload):fp8_fp4_mega_moe, 10bf16_mega_gate, 6 Triton stagingOn this PR's commit with MTP: the server was healthy, GSM8K (200 questions, 5-shot) scored 0.970, and mean MTP acceptance length was 2.62 on GSM8K. On the random-token benchmark, acceptance was 1.61, against 1.55 on the previous revision. DeepGEMM compiles during serving: 0.
The drafter needs no extra warmup. Its MegaMoE layer reuses the target's kernels, and the tester's tracer saw no DeepGEMM compile in the drafter during serving on either build.
Kimi K3 (
mgoin/Kimi-K3-pruned75, a pruned checkpoint): 4x B300, TP4 + EP4 with sequence parallelism,max_num_batched_tokens8192 (2048 per rank), fp8 KV, no prefix caching, no spec decode.fp8_fp4_mega_moeSiTU,block_m128 and 192)_prepare_megamoe_inputs_kernelcompiled during servinginit engine*The main run had a warm FlashInfer JIT cache from an earlier run. The PR runs had one too, so the +10 s compares like with like.
On both builds, two non-MegaMoE kernels still compile during serving: the Inductor kernel
triton_poi_fused_mul_sigmoid_0, on the first prefill above the capture sizes, and the KDA gated RMSNorm Triton kernellayer_norm_gated_fwd_kernel. They are outside this PR.Fully cold cache. With every DeepGEMM and Triton cache empty (not just the mega kernels), three other DeepGEMM kernel families still compile on first use with this PR. They are outside MegaMoE and left for follow-ups:
fp8_einsumino_proj(sm100_fp8_gemm_1d1d): 2 kernels.sm100_tf32_hc_prenorm_gemm: 8 kernels without MTP, 6 with MTP.sm100_mqa_logits: 1 kernel.With MTP, the spec-decode rejection-sampler Triton kernels also compile on first use on both builds. That is unrelated to this PR.
Tests.
tests/model_executor/test_deep_gemm_warmup.pyadds two CPU tests:bf16_mega_gateconfig are above the gate's token threshold, so the forward reaches the kernel.All 8 tests pass.
CI gap. This path needs SM100 GPUs, DeepSeek-V4 or Kimi K3 with the native MegaMoE backend, and expert parallelism. CI has no such setup, so the validation above was run by hand.
Repro commands
bench.py
Details
Root cause. Both kernels choose a JIT config from the per-rank token count:
bf16_mega_gate: block tokens, MMA CTAs, split-K, expert groups and stages (heuristics/mega_gate.hpp). That is 87 configs for per-rank sizes 17 to 8192.fp8_fp4_mega_moe:block_min {16, 32, 64, 128, 192, 240}.Startup only compiles the sizes that the profile run and CUDA graph capture hit, here 512 or fewer tokens per rank. Every other prefill size compiles on first use, about 1.7 s per kernel. DeepSeek-V4 needs two gate kernels per size, because hash-routed and bias-routed layers pass different optional inputs.
Design. The warmup lives in
deep_gemm_warmup.py, with the same parts as the existing DeepGEMM warmups:DeepseekV4MoEand Kimi K3'sKimiMoEwhen they use the native MegaMoE backend._get_fp8_fp4_mega_moe_m_values: for every per-rank size, it asks DeepGEMM for theblock_mand, for DeepSeek-V4, the gate config. It keeps the largest size of each distinct config (95 sizes for DeepSeek-V4 here), because DeepSeek-V4 only callsbf16_mega_gateabove a small token threshold and a config's largest size is always above it. The configs come from DeepGEMM's own getters, so the dedup is exact andVLLM_DEEP_GEMM_WARMUP=fullwould add nothing.GateLinear, which the bf16x3 router warmup already covers, so for Kimi onlyblock_mmatters.fp8_fp4_mega_moe, which picksblock_mwith the same heuristic.elifbranch in_count_warmup_iterations, and calls indeep_gemm_warmup.The worker runs the whole MoE layer at each size, not the raw kernels:
is_paddingfrom the forward context. This makes the Triton input-staging kernel compile the variants serving uses.VLLM_MOE_SKIP_PADDINGon (the default), ranks above 0 get mask slices at any alignment, so each size runs once with an aligned slice and once with an unaligned one.The predicate looks the model modules up in
sys.modules, as other model-specific warmups in this package do (kimi_k3_triton_warmup.py,qwen4_exp_qsa_warmup.py), so other models import neither model. The cache key reads two DeepSeek-V4-only attributes (tid2eid,image_sentinel_lo) withgetattr, because they do not exist on Kimi K3. The two config getters are reached through_import_deep_gemm, as in the model code.models/kimi_k3/nvidia/model.pygains one line, which storesuse_sequence_parallelonKimiMoEasDeepseekV4MoEalready does. The warmup needs it to compute the per-rank token count.Startup trade-off: where the cold warmup time goes
nvcctime for 136 gate kernels (1.7 s each).Compiling in parallel threads crashed inside DeepGEMM's JIT, which has no locking and no compile-only API. Parallel compilation would need DeepGEMM changes. vLLM keeps the DeepGEMM cache in
~/.cache/vllm/deep_gemm, so the cold cost is paid once per machine and DeepGEMM version.Remaining one-time first-request cost
The first eager prefill whose reduce-scatter falls in the MNNVL multimem size range sets up a symmetric-memory buffer (
_init_mnnvl_multimem_reduce_scatter_buffer). This adds 60-85 ms to that one request (+61, +66 and +83 ms in three runs). CUDA graph capture skips this setup. It is paid once per server and is not a compile. A torch profile and a request-order test confirm both points.Limitations. DeepSeek-V4 was validated only on 2x B300 with TP2 + EP2 and Model Runner V2, with and without MTP. Kimi K3 was validated only with a pruned checkpoint on 4x B300, TP4 + EP4, Model Runner V2.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (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.
BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing (anything written below this line will be removed by GitHub Actions)