Revert "[Core] Replace routing replay with device cache and async D2H pipeline" (vllm-project#39917) (vllm-project#42434)

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

Author: aoshen02

docs/training/routed_experts_replay.md

@@ -1,162 +1,244 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import types
+from types import SimpleNamespace
+from unittest.mock import patch
 
 import pytest
 import torch
 
+from vllm.distributed.eplb.eplb_state import EplbLayerState
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
+    RoutedExpertsCapturer,
+)
+from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
+
 pytestmark = pytest.mark.cpu_test
 
+_REC_MODULE = "vllm.model_executor.layers.fused_moe.routed_experts_capturer"
 
-def test_bind_routing_capture_to_model_sets_layer_view(monkeypatch):
-    import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
-    import vllm.model_executor.layers.fused_moe.routed_experts_capturer as rec_mod
 
-    class _DummyMoEConfig:
-        is_sequence_parallel = False
-        dp_size = 1
+def _capturer_with_buffer(
+    *,
+    max_tokens: int = 8,
+    num_layers: int = 4,
+    num_experts_per_tok: int = 2,
+    dp_rank: int = 0,
+) -> RoutedExpertsCapturer:
+    c = RoutedExpertsCapturer()
+    c.dp_rank = dp_rank
+    c._device_buffer = torch.full(
+        (max_tokens, num_layers, num_experts_per_tok),
+        -1,
+        dtype=torch.int32,
+    )
+    return c
 
-    class _DummyQuantMethod:
-        supports_internal_mk = True
 
-    class DummyFusedMoE:
-        _routing_replay_out: torch.Tensor
+class DummyRouter(BaseRouter):
+    @property
+    def routing_method_type(self) -> RoutingMethodType:
+        return RoutingMethodType.FUSED_TOPK
 
-        def __init__(self, moe_layer_id):
-            self.moe_layer_id = moe_layer_id
-            self.moe_config = _DummyMoEConfig()
-            self.quant_method = _DummyQuantMethod()
+    def _compute_routing(
+        self, hidden_states, router_logits, indices_type, *, input_ids=None
+    ):
+        topk_ids = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64)
+        topk_weights = torch.ones_like(topk_ids, dtype=torch.float32)
+        return topk_weights, topk_ids
 
-    monkeypatch.setattr(fused_moe_layer, "FusedMoE", DummyFusedMoE)
+    def _apply_eplb_mapping(self, topk_ids: torch.Tensor) -> torch.Tensor:
+        # Make mapping observable without requiring CUDA EPLB path.
+        return topk_ids + 10
 
-    num_layers, num_tokens, top_k = 4, 8, 2
-    buffer = torch.zeros((num_layers, num_tokens, top_k), dtype=torch.int16)
 
-    class DummyDeviceCache:
-        def __init__(self, buf):
-            self.buffer = buf
+def _make_router(eplb_state: EplbLayerState | None = None) -> DummyRouter:
+    return DummyRouter(
+        top_k=2,
+        global_num_experts=16,
+        eplb_state=eplb_state,
+        indices_type_getter=None,
+    )
 
-    class DummyCapturer:
-        def get_device_cache(self):
-            return DummyDeviceCache(buffer)
 
-    monkeypatch.setattr(rec_mod, "get_global_experts_capturer", lambda: DummyCapturer())
+def test_base_router_capture_pre_eplb_mapping():
+    router = _make_router()
+    captured = []
 
-    m0 = DummyFusedMoE(moe_layer_id=0)
-    m2 = DummyFusedMoE(moe_layer_id=2)
+    def capture_fn(ids):
+        captured.append(ids.clone())
 
-    class DummyModel:
-        def modules(self):
-            return iter([m0, m2])
+    router.set_capture_fn(capture_fn)
+    topk_weights, topk_ids = router.select_experts(
+        hidden_states=torch.empty(1),
+        router_logits=torch.empty(1),
+    )
 
-    rec_mod.bind_routing_capture_to_model(DummyModel())
+    assert topk_weights.shape == topk_ids.shape
+    assert len(captured) == 1
+    assert torch.equal(captured[0], torch.tensor([[1, 2], [3, 4]]))
+    assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]]))
 
-    assert torch.equal(m0._routing_replay_out, buffer[0])
-    assert torch.equal(m2._routing_replay_out, buffer[2])
 
+def test_base_router_capture_with_eplb_enabled():
+    eplb_state = EplbLayerState()
+    eplb_state.expert_load_view = torch.zeros(32, dtype=torch.int64)
+    eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1)
+    eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64)
+    eplb_state.should_record_tensor = torch.ones((), dtype=torch.bool)
+    router = _make_router(eplb_state=eplb_state)
 
-def test_bind_routing_capture_to_model_noop_when_disabled(monkeypatch):
-    import vllm.model_executor.layers.fused_moe.routed_experts_capturer as rec_mod
+    captured = []
 
-    class DummyCapturer:
-        def get_device_cache(self):
-            return None
+    def capture_fn(ids):
+        captured.append(ids.clone())
 
-    monkeypatch.setattr(rec_mod, "get_global_experts_capturer", lambda: DummyCapturer())
+    router.set_capture_fn(capture_fn)
+    _, topk_ids = router.select_experts(
+        hidden_states=torch.empty(1),
+        router_logits=torch.empty(1),
+    )
 
-    class DummyModel:
-        def modules(self):
-            return iter([])
+    assert len(captured) == 1
+    # Capture should see logical ids pre-EPLB mapping.
+    assert torch.equal(captured[0], torch.tensor([[1, 2], [3, 4]]))
+    # Our DummyRouter mapping adds +10.
+    assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]]))
 
-    rec_mod.bind_routing_capture_to_model(DummyModel())
 
+def test_gpu_model_runner_binds_router_capture(monkeypatch):
+    from vllm.v1.worker import gpu_model_runner as gmr
 
-# =========================================================================
-# Tests for device-cache routing replay architecture
-# =========================================================================
+    class DummyFusedMoE:
+        def __init__(self):
+            self.layer_id = 7
+            self.router = _make_router()
 
+    class DummyCapturer:
+        def __init__(self):
+            self.calls = []
 
-class TestRoutedExpertsDeviceCache:
-    """Tests for _RoutedExpertsDeviceCache (GPU buffer for routing data)."""
+        def capture(self, layer_id, topk_ids):
+            self.calls.append((layer_id, topk_ids))
 
-    def test_allocation_shape_and_dtype(self):
-        """Device cache allocates (L, N, K) int16 buffer."""
-        from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-            _RoutedExpertsDeviceCache,
-        )
+    dummy_module = DummyFusedMoE()
 
-        cache = _RoutedExpertsDeviceCache(
-            num_hidden_layers=40,
-            max_num_batched_tokens=8192,
-            num_experts_per_tok=8,
-            device="cpu",
-        )
-        assert cache.buffer.shape == (40, 8192, 8)
-        assert cache.buffer.dtype == torch.int16
+    # Patch the runtime import inside _bind_routed_experts_capturer.
+    import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
 
-    def test_per_layer_view_is_contiguous(self):
-        """buffer[layer_id] gives contiguous (N, K) view for FlashInfer."""
-        from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-            _RoutedExpertsDeviceCache,
-        )
+    monkeypatch.setattr(fused_moe_layer, "FusedMoE", DummyFusedMoE)
 
-        cache = _RoutedExpertsDeviceCache(
-            num_hidden_layers=40,
-            max_num_batched_tokens=8192,
-            num_experts_per_tok=8,
-            device="cpu",
+    dummy_self = types.SimpleNamespace(
+        compilation_config=types.SimpleNamespace(
+            static_forward_context={"dummy": dummy_module}
         )
-        layer_view = cache.buffer[0]
-        assert layer_view.is_contiguous()
-        assert layer_view.shape == (8192, 8)
+    )
 
+    capturer = DummyCapturer()
+    gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)
 
-class TestRoutedExpertsHostCache:
-    """Tests for _RoutedExpertsHostCache (per-request numpy buffer)."""
+    assert dummy_module.router.capture_fn is not None
+    dummy_module.router.capture_fn(torch.tensor([[5, 6]]))
 
-    def test_sentinel_initialization(self):
-        """Host cache initializes with zeros by default."""
-        import numpy as np
+    assert len(capturer.calls) == 1
+    layer_id, topk_ids = capturer.calls[0]
+    assert layer_id == 7
+    assert torch.equal(topk_ids, torch.tensor([[5, 6]]))
 
-        from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-            _RoutedExpertsHostCache,
-        )
 
-        cache = _RoutedExpertsHostCache(
-            num_hidden_layers=40,
-            num_experts_per_tok=8,
-            max_model_len=1024,
-        )
-        buf = cache.get_or_grow_buffer("req1", max_pos=100)
-        assert buf.dtype == np.int16
-        assert (buf == 0).all(), "Host cache must initialize with zeros"
-
-    def test_grow_preserves_existing_data(self):
-        """Growing the buffer preserves previously written data."""
-        from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-            _RoutedExpertsHostCache,
-        )
+def test_gpu_model_runner_binding_stage(monkeypatch):
+    from vllm.v1.worker import gpu_model_runner as gmr
 
-        cache = _RoutedExpertsHostCache(
-            num_hidden_layers=40,
-            num_experts_per_tok=8,
-            max_model_len=1024,
-        )
-        buf = cache.get_or_grow_buffer("req1", max_pos=50)
-        buf[0, 0, 0] = 42
-        buf2 = cache.get_or_grow_buffer("req1", max_pos=200)
-        assert buf2[0, 0, 0] == 42, "Data lost during buffer grow"
-
-    def test_free_request_removes_buffer(self):
-        """Freeing a request removes its buffer."""
-        from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-            _RoutedExpertsHostCache,
-        )
+    class DummyFusedMoE:
+        def __init__(self):
+            self.layer_id = 11
+            self.router = _make_router()
+
+    class DummyCapturer:
+        def __init__(self):
+            self.calls = []
+
+        def capture(self, layer_id, topk_ids):
+            self.calls.append((layer_id, topk_ids))
+
+    dummy_module = DummyFusedMoE()
+
+    import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
+
+    monkeypatch.setattr(fused_moe_layer, "FusedMoE", DummyFusedMoE)
 
-        cache = _RoutedExpertsHostCache(
-            num_hidden_layers=40,
-            num_experts_per_tok=8,
-            max_model_len=1024,
+    dummy_self = types.SimpleNamespace(
+        compilation_config=types.SimpleNamespace(
+            static_forward_context={"dummy": dummy_module}
         )
-        cache.get_or_grow_buffer("req1", max_pos=50)
-        cache.free_request("req1")
-        assert cache.get_buffer("req1") is None
+    )
+
+    # Before binding, no capture hook.
+    assert dummy_module.router.capture_fn is None
+
+    capturer = DummyCapturer()
+    gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)
+
+    # After binding, hook should exist and be callable.
+    assert callable(dummy_module.router.capture_fn)
+    dummy_module.router.capture_fn(torch.tensor([[9, 10]]))
+    assert len(capturer.calls) == 1
+
+
+def test_routed_experts_capturer_single_dp_no_metadata():
+    """dp_metadata is None: capture writes the full topk_ids rows."""
+    capturer = _capturer_with_buffer(dp_rank=0)
+    topk = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32)
+    ctx = SimpleNamespace(dp_metadata=None)
+    with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
+        capturer.capture(layer_id=0, topk_ids=topk)
+    assert torch.equal(capturer._device_buffer[:3, 0, :], topk)
+    assert capturer._device_buffer[3, 0, 0].item() == -1
+
+
+def test_routed_experts_capturer_dp_naive_concatenated_all_ranks():
+    """n == sum(num_tokens_dp): slice this rank's segment from concatenated topk."""
+    capturer = _capturer_with_buffer(dp_rank=1)
+    num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
+    ctx = SimpleNamespace(
+        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    )
+    # Concatenated order: rank0 rows then rank1 rows.
+    topk = torch.tensor(
+        [[0, 1], [2, 3], [10, 11], [12, 13], [14, 15]], dtype=torch.int32
+    )
+    with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
+        capturer.capture(layer_id=0, topk_ids=topk)
+    want = topk[2:5]
+    assert torch.equal(capturer._device_buffer[:3, 0, :], want)
+
+
+def test_routed_experts_capturer_dp_modular_local_tokens():
+    """n == token_num_per_dp: topk is already local to this DP rank."""
+    capturer = _capturer_with_buffer(dp_rank=1)
+    num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
+    ctx = SimpleNamespace(
+        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    )
+    topk = torch.tensor([[10, 11], [12, 13], [14, 15]], dtype=torch.int32)
+    with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
+        capturer.capture(layer_id=0, topk_ids=topk)
+    assert torch.equal(capturer._device_buffer[:3, 0, :], topk)
+
+
+def test_routed_experts_capturer_dp_unexpected_batch_raises():
+    """Mismatch between topk batch dim and DP layout: fail fast."""
+    capturer = _capturer_with_buffer(dp_rank=0)
+    num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
+    ctx = SimpleNamespace(
+        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    )
+    # total=5, local=2: n=1 matches neither naive (5) nor modular (2).
+    topk = torch.tensor([[1, 2]], dtype=torch.int32)
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        pytest.raises(AssertionError, match="unexpected topk_ids batch dim"),
+    ):
+        capturer.capture(layer_id=0, topk_ids=topk)
+    assert capturer._device_buffer[0, 0, 0].item() == -1