[Perf] Warm up DeepGEMM MegaMoE kernels before serving (#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>

tests/model_executor/test_deep_gemm_warmup.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import sys
 from types import SimpleNamespace
 
 import pytest
@@ -88,3 +89,53 @@ def test_kernel_registers_itself_as_warmup_provider(is_bmm) -> None:
 
     provider = getattr(layer, "deep_gemm_warmup_provider", None)
     assert provider is (None if is_bmm else kernel)
+
+
+def test_mega_moe_predicate_matches_only_native_mega_moe_layers(monkeypatch) -> None:
+    """Without the MegaMoE model modules loaded, nothing matches and nothing is
+    imported; with them loaded, only DeepSeek-V4 and Kimi K3 MegaMoE layers match."""
+    name = "vllm.models.deepseek_v4.nvidia.model"
+    kimi_name = "vllm.models.kimi_k3.nvidia.model"
+    monkeypatch.delitem(sys.modules, name, raising=False)
+    monkeypatch.delitem(sys.modules, kimi_name, raising=False)
+    other = SimpleNamespace(use_native_mega_moe=True, use_mega_moe=True)
+    assert not deep_gemm_warmup._mega_moe_may_use_deep_gemm(other)
+    assert name not in sys.modules and kimi_name not in sys.modules
+
+    moe_cls = type("DeepseekV4MoE", (), {})
+    monkeypatch.setitem(sys.modules, name, SimpleNamespace(DeepseekV4MoE=moe_cls))
+    native, fused = moe_cls(), moe_cls()
+    native.use_native_mega_moe, fused.use_native_mega_moe = True, False
+    assert deep_gemm_warmup._mega_moe_may_use_deep_gemm(native)
+    assert not deep_gemm_warmup._mega_moe_may_use_deep_gemm(fused)
+    assert not deep_gemm_warmup._mega_moe_may_use_deep_gemm(other)
+
+    kimi_cls = type("KimiMoE", (), {})
+    monkeypatch.setitem(sys.modules, kimi_name, SimpleNamespace(KimiMoE=kimi_cls))
+    kimi_mega, kimi_fused = kimi_cls(), kimi_cls()
+    kimi_mega.use_mega_moe, kimi_fused.use_mega_moe = True, False
+    assert deep_gemm_warmup._mega_moe_may_use_deep_gemm(kimi_mega)
+    assert not deep_gemm_warmup._mega_moe_may_use_deep_gemm(kimi_fused)
+
+
+def test_mega_moe_m_values_reach_mega_gate(monkeypatch) -> None:
+    """Each bf16_mega_gate config is warmed with an M above the GateLinear
+    threshold (1 for 128 experts), so the forward actually reaches the kernel."""
+    dg = SimpleNamespace(
+        get_bf16_mega_gate_config=lambda m, *_: {"num_sms": -(-m // 16)},
+        get_block_m_for_mega_moe=lambda *_: 16,
+    )
+    monkeypatch.setattr(deep_gemm_warmup, "_import_deep_gemm", lambda: dg)
+    monkeypatch.setattr(deep_gemm_warmup, "_mega_moe_uses_mega_gate", lambda _: True)
+    monkeypatch.setattr(
+        deep_gemm_warmup, "get_ep_group", lambda: SimpleNamespace(world_size=8)
+    )
+    buffer = SimpleNamespace(num_max_tokens_per_rank=64)
+    module = SimpleNamespace(
+        gate=SimpleNamespace(input_size=4096, output_size=128),
+        experts=SimpleNamespace(
+            top_k=6, num_experts=128, get_symm_buffer=lambda: buffer
+        ),
+        use_sequence_parallel=False,
+    )
+    assert deep_gemm_warmup._get_fp8_fp4_mega_moe_m_values(module, 40) == [16, 32, 40]

vllm/model_executor/warmup/deep_gemm_warmup.py

@@ -5,11 +5,19 @@ DeepGEMM JIT's the kernels. The warmup aims to JIT all the kernels that would
 be used during model execution beforehand.
 """
 
+import sys
+from typing import TYPE_CHECKING, TypeGuard
+
 import torch
 from tqdm import tqdm
 
 import vllm.envs as envs
-from vllm.distributed.parallel_state import get_dp_group, is_global_first_rank
+from vllm.distributed.parallel_state import (
+    get_dp_group,
+    get_ep_group,
+    is_global_first_rank,
+)
+from vllm.forward_context import ForwardContext, override_forward_context
 from vllm.model_executor.kernels.linear.scaled_mm.deep_gemm import (
     DeepGemmFp8BlockScaledMMKernel,
 )
@@ -23,6 +31,7 @@ from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
 )
 from vllm.tracing import instrument
 from vllm.utils.deep_gemm import (
+    _import_deep_gemm,
     fp8_gemm_nt,
     get_mk_alignment_for_contiguous_layout,
     m_grouped_fp8_gemm_nt_contiguous,
@@ -32,6 +41,10 @@ from vllm.utils.math_utils import cdiv
 from vllm.utils.platform_utils import num_compute_units
 from vllm.v1.worker.workspace import current_workspace_manager
 
+if TYPE_CHECKING:
+    from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MoE
+    from vllm.models.kimi_k3.nvidia.model import KimiMoE
+
 
 def _generate_optimal_warmup_m_values(
     max_tokens: int, n: int, device: torch.device
@@ -144,6 +157,22 @@ def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool:
     return isinstance(fused_experts, (DeepGemmExperts, TritonOrDeepGemmExperts))
 
 
+# Model construction already imports the model modules. Avoid importing them here.
+def _mega_moe_uses_mega_gate(module: torch.nn.Module) -> "TypeGuard[DeepseekV4MoE]":
+    # DeepSeek-V4 routes with bf16_mega_gate, Kimi K3 with its GateLinear.
+    dsv4 = sys.modules.get("vllm.models.deepseek_v4.nvidia.model")
+    return dsv4 is not None and isinstance(module, dsv4.DeepseekV4MoE)
+
+
+def _mega_moe_may_use_deep_gemm(
+    module: torch.nn.Module,
+) -> "TypeGuard[DeepseekV4MoE | KimiMoE]":
+    if _mega_moe_uses_mega_gate(module):
+        return module.use_native_mega_moe
+    kimi = sys.modules.get("vllm.models.kimi_k3.nvidia.model")
+    return kimi is not None and isinstance(module, kimi.KimiMoE) and module.use_mega_moe
+
+
 FP8_GEMM_NT_WARMUP_CACHE: set[torch.Size] = set()
 
 
@@ -305,6 +334,110 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
             GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE.add(w.size())
 
 
+FP8_FP4_MEGA_MOE_WARMUP_CACHE: set[tuple] = set()
+
+
+def _get_fp8_fp4_mega_moe_warmup_key(module: "DeepseekV4MoE | KimiMoE") -> tuple:
+    gate = module.gate
+    return (
+        gate.output_size,
+        getattr(gate, "tid2eid", None) is not None,
+        gate.e_score_correction_bias is not None,
+        getattr(module, "image_sentinel_lo", 0),
+        module.experts.has_fused_shared_experts,
+    )
+
+
+def _get_fp8_fp4_mega_moe_m_values(
+    module: "DeepseekV4MoE | KimiMoE", max_tokens: int
+) -> list[int]:
+    """Get the largest M per distinct fp8_fp4_mega_moe config and, where used,
+    bf16_mega_gate config."""
+    dg = _import_deep_gemm()
+    gate, experts = module.gate, module.experts
+    uses_mega_gate = _mega_moe_uses_mega_gate(module)
+    num_max_tokens_per_rank = experts.get_symm_buffer().num_max_tokens_per_rank
+    if module.use_sequence_parallel:
+        max_tokens = cdiv(max_tokens, module.tp_size)
+
+    m_values: dict[tuple, int] = {}
+    for num_tokens in range(1, max_tokens + 1):
+        gate_config = (
+            dg.get_bf16_mega_gate_config(
+                num_tokens, gate.input_size, gate.output_size, experts.top_k
+            )
+            if uses_mega_gate
+            else {}
+        )
+        block_m = dg.get_block_m_for_mega_moe(
+            get_ep_group().world_size,
+            experts.num_experts,
+            num_max_tokens_per_rank,
+            num_tokens,
+            experts.top_k,
+            "fp8xfp4",
+        )
+        # Keep the largest M: DeepSeek-V4 only calls bf16_mega_gate above a
+        # small token threshold, and a config's largest M is always above it.
+        m_values[(tuple(gate_config.items()), block_m)] = num_tokens
+    return sorted(m_values.values())
+
+
+def _deepgemm_fp8_fp4_mega_moe_warmup(
+    module: "DeepseekV4MoE | KimiMoE",
+    max_tokens: int,
+    pbar: tqdm | None = None,
+):
+    from vllm.models.deepseek_v4.nvidia.model import (
+        prepare_mega_gate_routing_metadata,
+    )
+
+    key = _get_fp8_fp4_mega_moe_warmup_key(module)
+    if key in FP8_FP4_MEGA_MOE_WARMUP_CACHE:
+        return
+
+    m_values = _get_fp8_fp4_mega_moe_m_values(module, max_tokens)
+    max_m = m_values[-1]  # sorted
+
+    device = module.experts.get_symm_buffer().x.device
+    input_ids = torch.zeros(max_m, device=device, dtype=torch.int64)
+    hidden_states = torch.zeros(
+        (max_m, module.gate.input_size), device=device, dtype=torch.bfloat16
+    )
+    is_padding = torch.zeros(max_m + 1, device=device, dtype=torch.bool)
+    # With padding skip on, sequence parallel gives ranks > 0 a padding mask
+    # slice at any alignment.
+    sp_masked = module.use_sequence_parallel and envs.VLLM_MOE_SKIP_PADDING
+    offsets = (0, 1) if sp_masked else (0,)
+
+    # Run the whole layer so the gate sees the same optional inputs as serving.
+    for num_tokens in m_values:
+        inputs: tuple = (hidden_states[:num_tokens],)
+        if _mega_moe_uses_mega_gate(module):
+            ids = input_ids[:num_tokens]
+            inputs += (
+                ids,
+                prepare_mega_gate_routing_metadata(
+                    ids,
+                    has_hash_routing=module.gate.tid2eid is not None,
+                    image_sentinel_base_id=module.image_sentinel_lo or None,
+                ),
+            )
+        for offset in offsets:
+            forward_context = ForwardContext(
+                no_compile_layers={},
+                attn_metadata={},
+                slot_mapping={},
+                is_padding=is_padding[offset : offset + num_tokens],
+            )
+            with override_forward_context(forward_context):
+                module(*inputs)
+        if pbar is not None:
+            pbar.update(1)
+
+    FP8_FP4_MEGA_MOE_WARMUP_CACHE.add(key)
+
+
 def deepgemm_fp8_gemm_nt_warmup(
     model: torch.nn.Module, max_tokens: int, pbar: tqdm | None = None
 ):
@@ -330,11 +463,21 @@ def deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
         )
 
 
+@torch.inference_mode()
+def deepgemm_fp8_fp4_mega_moe_warmup(
+    model: torch.nn.Module, max_tokens: int, pbar: tqdm | None = None
+):
+    for m in model.modules():
+        if _mega_moe_may_use_deep_gemm(m):
+            _deepgemm_fp8_fp4_mega_moe_warmup(m, max_tokens, pbar=pbar)
+
+
 def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int:
     seen_fp8_sizes: set[torch.Size] = set(FP8_GEMM_NT_WARMUP_CACHE)
     seen_grouped_sizes: set[torch.Size] = set(
         GROUPED_FP8_GEMM_NT_CONTIGUOUS_WARMUP_CACHE
     )
+    seen_mega_moe_keys: set[tuple] = set(FP8_FP4_MEGA_MOE_WARMUP_CACHE)
 
     total = 0
     for m in model.modules():
@@ -355,6 +498,11 @@ def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int:
             if w2.size() not in seen_grouped_sizes:
                 total += n_values
                 seen_grouped_sizes.add(w2.size())
+        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))
+                seen_mega_moe_keys.add(key)
     return total
 
 
@@ -369,6 +517,8 @@ def deep_gemm_warmup(model: torch.nn.Module, max_tokens: int):
         with tqdm(total=total, desc="DeepGEMM warmup") as pbar:
             deepgemm_fp8_gemm_nt_warmup(model, max_tokens, pbar)
             deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(model, max_tokens, pbar)
+            deepgemm_fp8_fp4_mega_moe_warmup(model, max_tokens, pbar)
     else:
         deepgemm_fp8_gemm_nt_warmup(model, max_tokens, None)
         deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(model, max_tokens, None)
+        deepgemm_fp8_fp4_mega_moe_warmup(model, max_tokens, None)

vllm/models/kimi_k3/nvidia/model.py

@@ -566,6 +566,7 @@ class KimiMoE(nn.Module):
         )
         self.latent_moe_use_norm = config.latent_moe_use_norm
         self.tp_size = get_tensor_model_parallel_world_size()
+        self.use_sequence_parallel = use_sequence_parallel
         self.routed_scaling_factor = config.routed_scaling_factor
         self.moe_renormalize = moe_renormalize
         self.use_grouped_topk = config.use_grouped_topk