[Model] Support R3 capture with DeepGEMM MegaMoE (#51831)

Signed-off-by: aoshen02 <aoshen@inferact.ai>

tests/models/test_deepseek_v4_mega_moe.py

@@ -6,6 +6,9 @@ from types import SimpleNamespace
 import pytest
 import torch
 
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
+    bind_routed_experts_capturer,
+)
 from vllm.models.deepseek_v4.nvidia.model import (
     DeepseekV4MegaMoEExperts,
     make_deepseek_v4_expert_params_mapping,
@@ -44,6 +47,61 @@ def test_deepseek_v4_mega_moe_ue8m0_uint8_to_float():
     assert decoded[3].item() == 2.0
 
 
+@pytest.mark.parametrize("use_kimi", [False, True])
+def test_deep_gemm_mega_moe_capture_precedes_eplb(monkeypatch, use_kimi):
+    experts_cls = DeepseekV4MegaMoEExperts
+    if use_kimi:
+        from vllm.models.kimi_k3.nvidia.model import KimiK3MegaMoEExperts
+
+        experts_cls = KimiK3MegaMoEExperts
+
+    experts = experts_cls.__new__(experts_cls)
+    torch.nn.Module.__init__(experts)
+    if use_kimi:
+        experts.synchronize_first_launch = lambda: None
+    experts.prefix = "model.layers.3.ffn.experts"
+    experts.max_num_tokens = 4
+    experts.capture_fn = None
+    experts.get_symm_buffer = lambda: object()
+    experts.eplb_state = SimpleNamespace(
+        logical_to_physical_map=torch.empty(1),
+        expert_load_view=torch.empty(1),
+        logical_replica_count=torch.empty(1),
+        should_record_tensor=torch.empty(1),
+        num_unpadded_tokens_tensors=None,
+    )
+
+    topk_ids = torch.tensor([[1, 2], [3, 4]])
+    captured: list[tuple[int, torch.Tensor]] = []
+    bind_routed_experts_capturer(
+        SimpleNamespace(modules=lambda: [experts]),
+        SimpleNamespace(capture=lambda layer_id, ids: captured.append((layer_id, ids))),
+    )
+
+    class MappingReached(Exception):
+        pass
+
+    def map_ids(**kwargs):
+        assert captured == [(3, topk_ids)]
+        raise MappingReached
+
+    monkeypatch.setattr(
+        f"{experts_cls.__module__}.eplb_map_to_physical_and_record",
+        map_ids,
+    )
+    monkeypatch.setattr(
+        "vllm.utils.deep_gemm._import_deep_gemm", lambda: SimpleNamespace()
+    )
+
+    with pytest.raises(MappingReached):
+        experts(
+            torch.empty(2, 8),
+            torch.empty(2, 2),
+            topk_ids,
+            activation_clamp=None,
+        )
+
+
 def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
     vllm_config = SimpleNamespace(
         scheduler_config=SimpleNamespace(max_num_batched_tokens=4),

vllm/model_executor/layers/fused_moe/routed_experts_capturer.py

@@ -6,6 +6,9 @@
 from __future__ import annotations
 
 import logging
+from collections.abc import Callable
+from functools import partial
+from typing import Protocol, runtime_checkable
 
 import numpy as np
 import torch
@@ -20,6 +23,12 @@ from vllm.v1.outputs import RoutedExpertsTensors
 logger = logging.getLogger(__name__)
 
 
+@runtime_checkable
+class RoutedExpertsCaptureSource(Protocol):
+    layer_id: int
+    capture_fn: Callable[[torch.Tensor], None] | None
+
+
 def _get_routed_experts_shape(vllm_config: VllmConfig) -> tuple[int, int, int]:
     model_config = vllm_config.model_config
     num_layers = model_config.get_total_num_hidden_layers()
@@ -234,6 +243,10 @@ def bind_routed_experts_capturer(
 
     num_bound = 0
     for module in model.modules():
+        if isinstance(module, RoutedExpertsCaptureSource):
+            module.capture_fn = partial(capturer.capture, module.layer_id)
+            num_bound += 1
+            continue
         if not isinstance(module, MoERunner):
             continue
         layer_id = module.layer_id

vllm/models/deepseek_v4/nvidia/model.py

@@ -183,6 +183,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
     ):
         super().__init__()
         self.prefix = prefix
+        self.capture_fn: Callable[[torch.Tensor], None] | None = None
         self.num_experts = num_experts
         self.num_local_experts = num_local_experts
         self.experts_start_idx = experts_start_idx
@@ -436,6 +437,10 @@ class DeepseekV4MegaMoEExperts(nn.Module):
     def update_expert_map(self) -> None:
         pass
 
+    @property
+    def layer_id(self) -> int:
+        return extract_layer_index(self.prefix)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -464,6 +469,9 @@ class DeepseekV4MegaMoEExperts(nn.Module):
             if is_padding is not None:
                 is_padding = is_padding[:num_tokens]
 
+        if self.capture_fn is not None:
+            self.capture_fn(topk_ids)
+
         # EPLB: map logical expert IDs to physical replicas and record load.
         eplb_state = self.eplb_state
         if eplb_state.logical_to_physical_map is not None:

vllm/models/kimi_k3/nvidia/model.py

@@ -383,6 +383,9 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts):
             if is_padding is not None:
                 is_padding = is_padding[:num_tokens]
 
+        if self.capture_fn is not None:
+            self.capture_fn(topk_ids)
+
         eplb_state = self.eplb_state
         if eplb_state.logical_to_physical_map is not None:
             assert eplb_state.expert_load_view is not None