[Model] Support R3 capture with DeepGEMM MegaMoE (#51831)
Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02 committed
Aug 12, 2026 at 16:07 UTC
4eef91c03cad3f080e46db1eeeadf121c9e6c136
4 files changed
+82
tests/models/test_deepseek_v4_mega_moe.py
+58
@@ -6,6 +6,9 @@ from types import SimpleNamespace
6
import pytest
7
import torch
8
9
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
10
+ bind_routed_experts_capturer,
11
+)
12
from vllm.models.deepseek_v4.nvidia.model import (
13
DeepseekV4MegaMoEExperts,
14
make_deepseek_v4_expert_params_mapping,
@@ -44,6 +47,61 @@ def test_deepseek_v4_mega_moe_ue8m0_uint8_to_float():
47
assert decoded[3].item() == 2.0
48
49
50
+@pytest.mark.parametrize("use_kimi", [False, True])
51
+def test_deep_gemm_mega_moe_capture_precedes_eplb(monkeypatch, use_kimi):
52
+ experts_cls = DeepseekV4MegaMoEExperts
53
+ if use_kimi:
54
+ from vllm.models.kimi_k3.nvidia.model import KimiK3MegaMoEExperts
55
+
56
+ experts_cls = KimiK3MegaMoEExperts
57
+
58
+ experts = experts_cls.__new__(experts_cls)
59
+ torch.nn.Module.__init__(experts)
60
+ if use_kimi:
61
+ experts.synchronize_first_launch = lambda: None
62
+ experts.prefix = "model.layers.3.ffn.experts"
63
+ experts.max_num_tokens = 4
64
+ experts.capture_fn = None
65
+ experts.get_symm_buffer = lambda: object()
66
+ experts.eplb_state = SimpleNamespace(
67
+ logical_to_physical_map=torch.empty(1),
68
+ expert_load_view=torch.empty(1),
69
+ logical_replica_count=torch.empty(1),
70
+ should_record_tensor=torch.empty(1),
71
+ num_unpadded_tokens_tensors=None,
72
+ )
73
+
74
+ topk_ids = torch.tensor([[1, 2], [3, 4]])
75
+ captured: list[tuple[int, torch.Tensor]] = []
76
+ bind_routed_experts_capturer(
77
+ SimpleNamespace(modules=lambda: [experts]),
78
+ SimpleNamespace(capture=lambda layer_id, ids: captured.append((layer_id, ids))),
79
+ )
80
+
81
+ class MappingReached(Exception):
82
+ pass
83
+
84
+ def map_ids(**kwargs):
85
+ assert captured == [(3, topk_ids)]
86
+ raise MappingReached
87
+
88
+ monkeypatch.setattr(
89
+ f"{experts_cls.__module__}.eplb_map_to_physical_and_record",
90
+ map_ids,
91
+ )
92
+ monkeypatch.setattr(
93
+ "vllm.utils.deep_gemm._import_deep_gemm", lambda: SimpleNamespace()
94
+ )
95
+
96
+ with pytest.raises(MappingReached):
97
+ experts(
98
+ torch.empty(2, 8),
99
+ torch.empty(2, 2),
100
+ topk_ids,
101
+ activation_clamp=None,
102
+ )
103
+
104
+
105
def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
106
vllm_config = SimpleNamespace(
107
scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
+13
@@ -6,6 +6,9 @@
6
from __future__ import annotations
7
8
import logging
9
+from collections.abc import Callable
10
+from functools import partial
11
+from typing import Protocol, runtime_checkable
12
13
import numpy as np
14
import torch
@@ -20,6 +23,12 @@ from vllm.v1.outputs import RoutedExpertsTensors
23
logger = logging.getLogger(__name__)
24
25
26
+@runtime_checkable
27
+class RoutedExpertsCaptureSource(Protocol):
28
+ layer_id: int
29
+ capture_fn: Callable[[torch.Tensor], None] | None
30
+
31
+
32
def _get_routed_experts_shape(vllm_config: VllmConfig) -> tuple[int, int, int]:
33
model_config = vllm_config.model_config
34
num_layers = model_config.get_total_num_hidden_layers()
@@ -234,6 +243,10 @@ def bind_routed_experts_capturer(
243
244
num_bound = 0
245
for module in model.modules():
246
+ if isinstance(module, RoutedExpertsCaptureSource):
247
+ module.capture_fn = partial(capturer.capture, module.layer_id)
248
+ num_bound += 1
249
+ continue
250
if not isinstance(module, MoERunner):
251
continue
252
layer_id = module.layer_id
vllm/models/deepseek_v4/nvidia/model.py
+8
@@ -183,6 +183,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
183
):
184
super().__init__()
185
self.prefix = prefix
186
+ self.capture_fn: Callable[[torch.Tensor], None] | None = None
187
self.num_experts = num_experts
188
self.num_local_experts = num_local_experts
189
self.experts_start_idx = experts_start_idx
@@ -436,6 +437,10 @@ class DeepseekV4MegaMoEExperts(nn.Module):
437
def update_expert_map(self) -> None:
438
pass
439
440
+ @property
441
+ def layer_id(self) -> int:
442
+ return extract_layer_index(self.prefix)
443
+
444
def forward(
445
self,
446
hidden_states: torch.Tensor,
@@ -464,6 +469,9 @@ class DeepseekV4MegaMoEExperts(nn.Module):
469
if is_padding is not None:
470
is_padding = is_padding[:num_tokens]
471
472
+ if self.capture_fn is not None:
473
+ self.capture_fn(topk_ids)
474
+
475
# EPLB: map logical expert IDs to physical replicas and record load.
476
eplb_state = self.eplb_state
477
if eplb_state.logical_to_physical_map is not None:
vllm/models/kimi_k3/nvidia/model.py
+3
@@ -383,6 +383,9 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts):
383
if is_padding is not None:
384
is_padding = is_padding[:num_tokens]
385
386
+ if self.capture_fn is not None:
387
+ self.capture_fn(topk_ids)
388
+
389
eplb_state = self.eplb_state
390
if eplb_state.logical_to_physical_map is not None:
391
assert eplb_state.expert_load_view is not None