[MoE] Revert platform hook for routed-expert capture rows (#62624) (#63013)

Signed-off-by: Ronald1995 <ronaldautomobile@163.com>

docs/design/plugin_system.md

@@ -105,7 +105,6 @@ Every plugin has three parts:
     - `check_and_update_config`: This function is called very early in the vLLM's initialization process. It's used for plugins to update the vllm configuration. For example, the block size, graph mode config, etc., can be updated in this function. The most important thing is that the **worker_cls** should be set in this function to let vLLM know which worker class to use for the worker process.
     - `get_attn_backend_cls`: This function should return the attention backend class's fully qualified name.
     - `get_device_communicator_cls`: This function should return the device communicator class's fully qualified name.
-    - (optional) `get_routed_experts_local_rows`: If the platform's MoE preparation pads or shards the router input before expert selection, return this DP rank's routing rows for routed-expert capture. The default `None` keeps the generic DP/TP restoration.
 
 4. Implement the worker class `MyDummyWorker` in `my_dummy_worker.py`. The worker class should inherit from [WorkerBase][vllm.v1.worker.worker_base.WorkerBase]. Please follow the interface to implement the functions one by one. Basically, all interfaces in the base class should be implemented, since they are called here and there in vLLM. To make sure a model can be executed, the basic functions should be implemented are:
 

tests/model_executor/test_routed_experts_capture.py

@@ -31,7 +31,6 @@ from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
     bind_routed_experts_capturer,
 )
 from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
-from vllm.platforms.interface import Platform
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheGroupSpec,
@@ -329,55 +328,6 @@ def test_monolithic_capture_survives_kernel_rebuilds():
         forward(_ToyMonolithicExperts(supports_capture=False))
 
 
-@pytest.mark.parametrize(
-    "rows, top_k",
-    [([[7, 8], [9, 10]], 2), ([], 2), ([[0]], 1)],
-    ids=["multiple-elements", "idle-rank", "zero-expert-id"],
-)
-def test_platform_rows_override_preserves_tensor_results(monkeypatch, rows, top_k):
-    """Platform rows bypass shape inference, even when empty or all zero."""
-    capturer = _capturer_with_buffer(dp_rank=1, tp_size=2, num_experts_per_tok=top_k)
-    expected = torch.tensor(rows, dtype=torch.int32).reshape(-1, top_k)
-    ctx = _forward_context(batch_descriptor=BatchDescriptor(num_tokens=5))
-    topk = torch.full((7, top_k), -2, dtype=torch.int32)
-
-    class CustomPlatform(Platform):
-        @classmethod
-        def get_routed_experts_local_rows(cls, ids, context, dp_rank, tp_size):
-            assert ids is topk and context is ctx
-            assert (dp_rank, tp_size) == (1, 2)
-            return expected
-
-    monkeypatch.setattr(f"{_REC_MODULE}.current_platform", CustomPlatform())
-    with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
-        capturer.capture(0, topk)
-
-    torch.testing.assert_close(capturer.device_buffer[: len(expected), 0], expected)
-    assert (capturer.device_buffer[len(expected) :] == -1).all()
-
-
-def test_platform_rows_none_restores_default_layout(monkeypatch):
-    """The default hook still gathers and trims TP shards."""
-    capturer = _capturer_with_buffer(tp_size=2)
-    local_shard = torch.tensor([[10, 11], [12, 13]], dtype=torch.int32)
-    gathered = torch.tensor([[10, 11], [12, 13], [14, 15], [-1, -1]], dtype=torch.int32)
-    ctx = _forward_context(
-        dp_metadata=DPMetadata(num_tokens_across_dp_cpu=torch.tensor([3]))
-    )
-    monkeypatch.setattr(f"{_REC_MODULE}.current_platform", Platform())
-    tp_group = Mock()
-    tp_group.all_gather.return_value = gathered
-    with (
-        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
-        patch(f"{_REC_MODULE}.get_tp_group", return_value=tp_group),
-    ):
-        capturer.capture(0, local_shard)
-
-    tp_group.all_gather.assert_called_once()
-    torch.testing.assert_close(capturer.device_buffer[:3, 0], gathered[:3])
-    assert (capturer.device_buffer[3:] == -1).all()
-
-
 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)

vllm/model_executor/layers/fused_moe/routed_experts_capturer.py

@@ -91,12 +91,7 @@ class RoutedExpertsCapturer:
                 f"layer count {self.device_buffer.shape[1]}"
             )
 
-        ctx = get_forward_context()
-        local_topk_ids = current_platform.get_routed_experts_local_rows(
-            topk_ids, ctx, self.dp_rank, self.tp_size
-        )
-        if local_topk_ids is None:
-            local_topk_ids = self._restore_local_topk_ids(topk_ids, ctx)
+        local_topk_ids = self._restore_local_topk_ids(topk_ids, get_forward_context())
         self.device_buffer[: len(local_topk_ids), layer_id] = local_topk_ids
 
     def _restore_local_topk_ids(

vllm/platforms/interface.py

@@ -21,7 +21,6 @@ if TYPE_CHECKING:
 
     from vllm.config import VllmConfig
     from vllm.config.kernel import IrOpPriorityConfig
-    from vllm.forward_context import ForwardContext
     from vllm.inputs import EngineInput
     from vllm.pooling_params import PoolingParams
     from vllm.sampling_params import SamplingParams
@@ -1428,26 +1427,6 @@ class Platform:
         """Set some additional forward context for the current platform if needs."""
         return {}
 
-    @classmethod
-    def get_routed_experts_local_rows(
-        cls,
-        topk_ids: torch.Tensor,
-        forward_context: "ForwardContext",
-        dp_rank: int,
-        tp_size: int,
-    ) -> torch.Tensor | None:
-        """Return this DP rank's routing rows after a platform-specific prepare.
-
-        Overrides must pick collectives from rank-invariant state (idle ranks
-        included) and issue device work on the current stream without host
-        reads of device values, so capture stays graph-capture/replay safe.
-        Capture currently runs from the Model Runner V2 aux-output connector.
-
-        Return the rows in token order, or ``None`` for the default
-        restoration. Empty tensors are valid.
-        """
-        return None
-
     @classmethod
     def num_compute_units(cls, device_id: int = 0) -> int:
         """Get the number of compute units for the current platform.