Skip to content

[Perf] Warm up DeepGEMM MegaMoE kernels before serving - #60468

Merged
DarkLight1337 merged 2 commits into
vllm-project:mainfrom
simon-veitner-redhat:mega-moe-jit-warmup
Oct 13, 2026
Merged

DarkLight1337 merged 2 commits into
vllm-project:mainfrom
simon-veitner-redhat:mega-moe-jit-warmup

Conversation

@simon-veitner-redhat

@simon-veitner-redhat simon-veitner-redhat commented Oct 7, 2026 •

Copy link
Copy Markdown
Contributor

Overview

The DeepGEMM warmup covers fp8_gemm_nt and the grouped MoE GEMM. It does not cover fp8_fp4_mega_moe or the bf16_mega_gate that 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 to deep_gemm_warmup.py, following the structure of the two existing DeepGEMM warmups there.

Claims

  • No serving-time compiles of the MegaMoE kernels: fp8_fp4_mega_moe, bf16_mega_gate and the Triton input-staging kernel. Before: 60 DeepGEMM compiles in one DeepSeek-V4 benchmark run. For Kimi K3, the 2 SiTU fp8_fp4_mega_moe compiles during serving that cost 2.6-2.9 s of TTFT each go to 0, at about +10 s of startup.
  • MTP speculative decoding needs nothing extra. The DeepSeek-V4 MTP drafter's MegaMoE layer uses the same kernels as the target's, and there were 0 DeepGEMM compiles during serving with MTP on.
  • DeepSeek-V4: the first request at a new prefill size no longer stalls for 3.5 to 6 s. It now takes about the same time as a repeat. The exception is a one-time 60-85 ms communication-buffer setup on the server's first large prefill, which is not a compile (see Details).
  • Cost: for DeepSeek-V4, the first startup on a cold DeepGEMM cache takes about 4 to 5 min longer (+247 to +299 s on fully empty caches). A warm restart takes the same time as before. For Kimi K3, it is about +10 s.

Validation

Setup (DeepSeek-V4). DeepSeek-V4-Flash on 2x B300, TP2 + EP2, --moe-backend deep_gemm_mega_moe, max_num_batched_tokens 16384, --max-model-len 16384, prefix caching off, Model Runner V2. vLLM precompiled wheel 3ca00a82 (cu130) with its bundled DeepGEMM.

Each run used a fresh DG_JIT_CACHE_DIR seeded 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_tokens 4, streamed time to first token (TTFT).

  1. Six fixed lengths (100 to 8000 tokens). Each gets a first request, then 3 repeats with new random tokens.
  2. 40 seeded random lengths from 51 to 15504 tokens. The warmup did not pick these.
  3. 32 concurrent requests.

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.

main this PR
sm100_bf16_mega_gate 58 0
sm100_fp8_fp4_mega_moe 2 0
other DeepGEMM families 0 0
Triton _prepare_megamoe_inputs_kernel ≥1 (see note) 0

On 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):

prompt tokens main, first main, repeat this PR, first this PR, repeat
1000 36 36 36 36
2000 3503 106 112 103
4000 5913 119 178 117
8000 3489 136 141 136
  • 40 random lengths: 24 of the 40 requests on main took 3.5 to 6.0 s. The slowest on this PR took 0.27 s. Total TTFT went from 90.1 s to 6.1 s.
  • 32 concurrent requests: worst TTFT went from 8.33 s to 1.47 s.

Startup, init engine (profile, KV cache, warmup, graph capture):

main this PR
cold mega-kernel cache (all other kernels cached) 141 s 428 s
fully empty caches, no MTP (separate tester run) 444 s 691 s (+247 s)
fully empty caches, with MTP (separate tester run) 406 s 705 s (+299 s)
warm restart, same cache dir 63 s 63 s

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):

MegaMoE compiles during serving main this PR
seeded cache (mega kernels cold), 1 run not run 0
fully empty caches, 1 run (separate tester) 2 fp8_fp4_mega_moe, 10 bf16_mega_gate, 6 Triton staging 0

On 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_tokens 8192 (2048 per rank), fp8 KV, no prefix caching, no spec decode.

  • Each run started with empty DeepGEMM, Triton, Inductor and vLLM caches. One run per build.
  • Workload: one 3000-token request, a second 3000-token request with new tokens, then 137, 1500, 5000 and 8000 tokens.
main this PR
DeepGEMM compiles during serving 2 (fp8_fp4_mega_moe SiTU, block_m 128 and 192) 0
Triton _prepare_megamoe_inputs_kernel compiled during serving yes (ranks 1-3, first request) no
TTFT at 1500 / 5000 tokens 2734 / 3139 ms 153 / 244 ms
init engine 46 s* 56-57 s

*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 kernel layer_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_einsum in o_proj (sm100_fp8_gemm_1d1d): 2 kernels.
  • sm100_tf32_hc_prenorm_gemm: 8 kernels without MTP, 6 with MTP.
  • The sparse indexer's 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.py adds two CPU tests:

  • The layer predicate is False for other models, and neither model module is imported. It is True only for DeepSeek-V4 and Kimi K3 MegaMoE layers.
  • The warmup sizes for each bf16_mega_gate config 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
# server (baseline: same command on main)
export VLLM_USE_DEEP_GEMM=1 DG_JIT_CACHE_DIR=<fresh dir>
vllm serve deepseek-ai/DeepSeek-V4-Flash \
  --trust-remote-code --kv-cache-dtype fp8 --block-size 256 \
  --enable-expert-parallel --tensor-parallel-size 2 \
  --attention_config.indexer_kv_dtype=mxfp4 \
  --moe-backend deep_gemm_mega_moe --tokenizer-mode deepseek_v4 \
  --max-model-len 16384 --no-enable-prefix-caching --jit-monitor-verbose --port 8234

# client: fixed lengths x (1 first + 3 repeats), 40 random lengths, 32-request burst;
# records TTFT and the new DeepGEMM cache entries per request
SWEEP=40 python bench.py 8234 $DG_JIT_CACHE_DIR results.json
bench.py
"""First-hit benchmark. Phase 1 (fixed): for each prompt length, one request
at a time: a first hit, then REPEATS more with fresh random tokens. Phase 2
(sweep, optional): random lengths the warmup did not pick, then a concurrent
burst (mixed prefill + decode). Records TTFT and every DeepGEMM JIT cache
entry created during serving, per kernel family."""
import collections, concurrent.futures as cf, json, os, random, re, sys, time
import requests

port, dgdir, out = sys.argv[1], sys.argv[2], sys.argv[3]
LENGTHS = [int(x) for x in os.environ.get("LENGTHS", "100,400,1000,2000,4000,8000").split(",")]
REPEATS = int(os.environ.get("REPEATS", "3"))
SWEEP = int(os.environ.get("SWEEP", "0"))
URL = f"http://localhost:{port}/v1/completions"
MODEL = "deepseek-ai/DeepSeek-V4-Flash"
cache = os.path.join(dgdir, "cache")

def entries():
    return set(os.listdir(cache))

def family(e):
    m = re.match(r"(?:kernel\.)?(.+?)\.[0-9a-f]{32}$", e)
    return m.group(1) if m else e

def fam_count(es):
    return dict(collections.Counter(family(e) for e in es))

def ttft(prompt_ids, max_tokens=4):
    t0 = time.perf_counter()
    r = requests.post(URL, json=dict(model=MODEL, prompt=prompt_ids, max_tokens=max_tokens,
                                     temperature=0, stream=True, ignore_eos=True),
                      stream=True, timeout=3600)
    first = None
    for line in r.iter_lines():
        if line and line.startswith(b"data: ") and line != b"data: [DONE]" and first is None:
            first = time.perf_counter() - t0
    return first, time.perf_counter() - t0

rng = random.Random(0)
def prompt(L):
    return [rng.randrange(1000, 100000) for _ in range(L)]

start = entries()
print(f"entries at ready: {len(start)} {fam_count(start)}", flush=True)
rows = []
ttft(prompt(8))  # HTTP/tokenizer warm; 4 per-rank tokens -> no mega gate
for L in LENGTHS:
    for i in range(1 + REPEATS):
        before = entries()
        t, e2e = ttft(prompt(L))
        new = entries() - before
        rows.append(dict(phase="fixed", prompt_len=L, rep="first" if i == 0 else f"repeat{i}",
                         ttft_s=t, e2e_s=e2e, new=fam_count(new)))
        print(json.dumps(rows[-1]), flush=True)
after_fixed = entries()
sweep_lengths = []
if SWEEP:
    srng = random.Random(1)
    sweep_lengths = sorted(srng.randrange(17, 16380) for _ in range(SWEEP))
    for L in sweep_lengths:
        before = entries()
        t, e2e = ttft(prompt(L))
        rows.append(dict(phase="sweep", prompt_len=L, rep="first", ttft_s=t, e2e_s=e2e,
                         new=fam_count(entries() - before)))
        print(json.dumps(rows[-1]), flush=True)
    # concurrent burst: 32 requests, mixed lengths, 64 output tokens -> prefill + decode mix
    burst = [prompt(srng.randrange(17, 6000)) for _ in range(32)]
    before = entries(); t0 = time.perf_counter()
    with cf.ThreadPoolExecutor(32) as ex:
        res = list(ex.map(lambda p: ttft(p, 64), burst))
    rows.append(dict(phase="burst", prompt_len=-1, rep="32 concurrent", ttft_s=max(r[0] for r in res),
                     e2e_s=time.perf_counter() - t0, new=fam_count(entries() - before)))
    print(json.dumps(rows[-1]), flush=True)
end = entries()
summary = dict(rows=rows, entries_at_ready=fam_count(start), new_during_fixed=fam_count(after_fixed - start),
               new_during_serving=fam_count(end - start), new_entries=sorted(end - start),
               sweep_lengths=sweep_lengths)
json.dump(summary, open(out, "w"), indent=1)
print("new during serving:", summary["new_during_serving"])

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_m in {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:

  • A layer predicate that matches DeepSeek-V4's DeepseekV4MoE and Kimi K3's KimiMoE when they use the native MegaMoE backend.
  • A cache constant.
  • _get_fp8_fp4_mega_moe_m_values: for every per-rank size, it asks DeepGEMM for the block_m and, 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 calls bf16_mega_gate above 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 and VLLM_DEEP_GEMM_WARMUP=full would add nothing.
    • Kimi K3 routes with its GateLinear, which the bf16x3 router warmup already covers, so for Kimi only block_m matters.
    • Kimi's experts call the SiTU variant of fp8_fp4_mega_moe, which picks block_m with the same heuristic.
  • A private worker and a public loop, an elif branch in _count_warmup_iterations, and calls in deep_gemm_warmup.

The worker runs the whole MoE layer at each size, not the raw kernels:

  • The gate gets the same optional inputs it gets in serving.
  • Every EP rank runs the same sizes in the same order, which the collective MegaMoE kernel needs.
  • It runs under a minimal forward context with a padding mask, because the layer reads is_padding from the forward context. This makes the Triton input-staging kernel compile the variants serving uses.
  • Under sequence parallelism with VLLM_MOE_SKIP_PADDING on (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) with getattr, 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.py gains one line, which stores use_sequence_parallel on KimiMoE as DeepseekV4MoE already does. The warmup needs it to compute the per-rank token count.

Startup trade-off: where the cold warmup time goes
  • About 84% of the cold warmup is serial nvcc time for 136 gate kernels (1.7 s each).
  • About 14% is about 24 Triton staging variants.
  • About 3% is 3 MegaMoE kernels.
  • The layer forwards and the size search take under 1 s.

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.

  • Not run: Model Runner V1, DP > 1, other EP sizes, EPLB, and routed-experts capture. With EPLB or routed-experts capture on, the warmup forwards would pass dummy tokens to the load recording and the capture hook.

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.

BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing (anything written below this line will be removed by GitHub Actions)

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

)


def deepgemm_fp8_fp4_mega_moe_warmup(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

warp with torch.inference_mode()

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

@LopezCastroRoberto LopezCastroRoberto Oct 9, 2026 •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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,)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: I think the offset loop doubles work when padding is off.

Suggested change
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,)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +175 to +176
dsv4 = sys.modules["vllm.models.deepseek_v4.nvidia.model"]
return isinstance(module, dsv4.DeepseekV4MoE)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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))

@LopezCastroRoberto LopezCastroRoberto Oct 9, 2026 •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>
@LopezCastroRoberto LopezCastroRoberto added warmup Related to the warmup system ready ONLY add when PR is ready to merge/full CI is needed labels Oct 13, 2026
@github-actions

Copy link
Copy Markdown

✅ @simon-veitner-redhat, CI is now available for this PR.

  • /ci run starts upstream CI; /amd-ci run starts AMD CI only.
  • Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch at your own risk.
  • /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.

@LopezCastroRoberto LopezCastroRoberto left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks!

@LopezCastroRoberto

Copy link
Copy Markdown
Contributor

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #94444 for commit 71c61dc99458.

@DarkLight1337
DarkLight1337 merged commit 253e02b into vllm-project:main Oct 13, 2026
147 checks passed
AndreasKaratzas pushed a commit to AndreasKaratzas/vllm that referenced this pull request Oct 16, 2026
…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>
michael-gpu pushed a commit to practical-inference/vllm-ampere-extended that referenced this pull request Oct 20, 2026
…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>
@simon-veitner-redhat
simon-veitner-redhat deleted the mega-moe-jit-warmup branch October 21, 2026 09:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

k3 kimi ready ONLY add when PR is ready to merge/full CI is needed warmup Related to the warmup system

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants