diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 25b431429b..25ea9952c9 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/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), diff --git a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py index 618e82bea9..800351f46e 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py +++ b/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 diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index e102b6ddd3..ed7ee11ce5 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/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: diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 0d98bccf06..59682fe32a 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/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