[Feature] Add deterministic forced routing for CUDA MoE benchmarks

afd_plugin/__init__.py

@@ -114,6 +114,12 @@ def register_afd() -> None:
             exc_info=True,
         )
 
+    from afd_plugin.model_executor.routing_simulator import (
+        register_afd_balanced_routing_strategy,
+    )
+
+    register_afd_balanced_routing_strategy()
+
     try:
         from afd_plugin.v1.worker.dbo import register_dbo_yield_custom_op
 

afd_plugin/model_executor/routing_simulator.py

@@ -0,0 +1,178 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project
+"""Deterministic routing strategy for controlled CUDA MoE benchmarks.
+
+Select ``afd_balanced`` with vLLM's
+``VLLM_MOE_ROUTING_SIMULATION_STRATEGY`` environment variable. The optional
+``AFD_BENCHMARK_FORCE_LB_TOPN_PER_RANK`` variable limits the local expert pool
+on every EP rank; ``0`` selects all local experts.
+
+The strategy returns normalized uniform weights and deterministic expert IDs.
+It changes model outputs and is intended only for benchmark and profiling runs.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+import torch
+from vllm.distributed.parallel_state import get_ep_group
+from vllm.model_executor.layers.fused_moe.router.routing_simulator_router import (
+    RoutingSimulator,
+    RoutingStrategy,
+)
+
+AFD_BALANCED_ROUTING_STRATEGY = "afd_balanced"
+_TOPN_PER_RANK_ENV = "AFD_BENCHMARK_FORCE_LB_TOPN_PER_RANK"
+_DETERMINISTIC_SEED = 1024
+
+
+@dataclass(frozen=True)
+class _RoutingConfig:
+    num_experts: int
+    ep_size: int
+    ep_rank: int
+    top_k: int
+    topn_per_rank: int
+
+
+@dataclass(frozen=True)
+class _RoutingBuffers:
+    weights: torch.Tensor
+    expert_ids: torch.Tensor
+
+
+def _validate_config(config: _RoutingConfig) -> None:
+    if config.num_experts <= 0:
+        raise ValueError("num_experts must be positive")
+    if config.ep_size <= 0:
+        raise ValueError("ep_size must be positive")
+    if not 0 <= config.ep_rank < config.ep_size:
+        raise ValueError("ep_rank must be within the expert-parallel group")
+    if config.top_k <= 0:
+        raise ValueError("top_k must be positive")
+    if config.num_experts % config.ep_size != 0:
+        raise ValueError("num_experts must be divisible by ep_size")
+
+    local_experts = config.num_experts // config.ep_size
+    selected_per_rank = config.topn_per_rank or local_experts
+    if not 0 < selected_per_rank <= local_experts:
+        raise ValueError(f"{_TOPN_PER_RANK_ENV} must be between 0 and {local_experts}")
+    if config.top_k > selected_per_rank * config.ep_size:
+        raise ValueError("selected expert pool must contain at least top_k experts")
+
+
+def _build_expert_cycle(
+    config: _RoutingConfig,
+    device: torch.device,
+) -> torch.Tensor:
+    _validate_config(config)
+    local_experts = config.num_experts // config.ep_size
+
+    if config.topn_per_rank:
+        expert_cycle = torch.cat(
+            [
+                torch.arange(
+                    rank * local_experts,
+                    rank * local_experts + config.topn_per_rank,
+                    device=device,
+                    dtype=torch.int32,
+                )
+                for rank in range(config.ep_size)
+            ]
+        )
+    else:
+        generator = torch.Generator(device="cpu")
+        generator.manual_seed(_DETERMINISTIC_SEED)
+        expert_cycle = torch.randperm(
+            config.num_experts,
+            generator=generator,
+            device="cpu",
+            dtype=torch.int32,
+        ).to(device=device, non_blocking=True)
+
+    source_rank_offset = config.ep_rank * local_experts
+    return (expert_cycle + source_rank_offset) % config.num_experts
+
+
+def _build_routing_buffers(
+    config: _RoutingConfig,
+    max_tokens: int,
+    device: torch.device,
+    indices_dtype: torch.dtype,
+) -> _RoutingBuffers:
+    expert_cycle = _build_expert_cycle(config, device)
+    total_ids = max_tokens * config.top_k
+    repeat_count = (total_ids + expert_cycle.numel() - 1) // expert_cycle.numel()
+    expert_ids = expert_cycle.repeat(repeat_count)[:total_ids].reshape(
+        max_tokens,
+        config.top_k,
+    )
+    expert_ids = expert_ids.to(dtype=indices_dtype)
+    weights = torch.full(
+        (max_tokens, config.top_k),
+        1.0 / config.top_k,
+        device=device,
+        dtype=torch.float32,
+    )
+    return _RoutingBuffers(weights=weights, expert_ids=expert_ids)
+
+
+class AFDBalancedRoutingStrategy(RoutingStrategy):
+    """Generate deterministic, EP-balanced routing for performance tests."""
+
+    def __init__(self, topn_per_rank: int | None = None) -> None:
+        if topn_per_rank is None:
+            topn_per_rank = int(os.environ.get(_TOPN_PER_RANK_ENV, "0"))
+        self.topn_per_rank = topn_per_rank
+        self._buffers: dict[
+            tuple[_RoutingConfig, torch.device, torch.dtype],
+            _RoutingBuffers,
+        ] = {}
+
+    def route_tokens(
+        self,
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        top_k: int,
+        indices_type: torch.dtype | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        ep_group = get_ep_group()
+        config = _RoutingConfig(
+            num_experts=router_logits.shape[-1],
+            ep_size=ep_group.world_size,
+            ep_rank=ep_group.rank_in_group,
+            top_k=top_k,
+            topn_per_rank=self.topn_per_rank,
+        )
+        indices_dtype = indices_type or torch.int32
+        key = (config, hidden_states.device, indices_dtype)
+        buffers = self._buffers.get(key)
+        max_tokens = hidden_states.shape[0]
+        if buffers is None or buffers.expert_ids.shape[0] < max_tokens:
+            buffers = _build_routing_buffers(
+                config,
+                max_tokens,
+                hidden_states.device,
+                indices_dtype,
+            )
+            self._buffers[key] = buffers
+
+        num_tokens = hidden_states.shape[0]
+        return buffers.weights[:num_tokens], buffers.expert_ids[:num_tokens]
+
+
+def register_afd_balanced_routing_strategy() -> None:
+    """Register the opt-in AFD strategy with vLLM's routing simulator."""
+    RoutingSimulator.register_strategy(
+        AFD_BALANCED_ROUTING_STRATEGY,
+        AFDBalancedRoutingStrategy(),
+    )
+
+
+__all__ = [
+    "AFD_BALANCED_ROUTING_STRATEGY",
+    "AFDBalancedRoutingStrategy",
+    "register_afd_balanced_routing_strategy",
+]

docs/design/module/compatibility_and_patches.md

@@ -75,6 +75,7 @@ compatibility evidence rather than a released package or container tag.
 | --- | --- | --- |
 | Version policy | [`compat/vllm.py`](../../../afd_plugin/compat/vllm.py) | [`test_package.py`](../../../tests/unit/package/test_package.py) |
 | Core vLLM patches | [`compat/patches/`](../../../afd_plugin/compat/patches) | [`tests/unit/compat/patches/`](../../../tests/unit/compat/patches) |
+| CUDA benchmark routing | [`model_executor/routing_simulator.py`](../../../afd_plugin/model_executor/routing_simulator.py) | [`test_routing_simulator.py`](../../../tests/unit/model_executor/test_routing_simulator.py) |
 | NPU adapters | [`compat/npu/`](../../../afd_plugin/compat/npu) | [`test_runtime.py`](../../../tests/unit/compat/test_runtime.py), [`test_ascend_ops.py`](../../../tests/unit/compat/test_ascend_ops.py) |
 | NPU patch paths | [`compat/patches/npu/`](../../../afd_plugin/compat/patches/npu) | [`test_force_load_balance.py`](../../../tests/unit/compat/patches/test_force_load_balance.py), [`test_npu_runtime.py`](../../../tests/unit/v1/worker/test_npu_runtime.py) |
 | CUDA DBO Attention metadata workaround | [`attention_model_runner.py`](../../../afd_plugin/v1/worker/attention_model_runner.py) | restoration and error-path coverage in [`test_attention_model_runner.py`](../../../tests/unit/v1/worker/test_attention_model_runner.py) |
@@ -86,11 +87,12 @@ The vLLM plugin entry point applies compatibility in this order:
 1. perform the non-strict vLLM version check;
 2. import `async_dp_engine`, `async_dp_forward_context`, `config_validation`,
    and `engine_core` in one best-effort block;
-3. register the plugin-owned DBO yield operator;
-4. call the idempotent Ascend runtime facade, which installs the NPU platform
+3. register the plugin-owned CUDA benchmark routing strategy;
+4. register the plugin-owned DBO yield operator;
+5. call the idempotent Ascend runtime facade, which installs the NPU platform
    config wrapper when vLLM-Ascend is importable;
-5. import the force-load-balance patch only when vLLM-Ascend is discoverable;
-6. register model mappings, which is required for registration to complete.
+6. import the force-load-balance patch only when vLLM-Ascend is discoverable;
+7. register model mappings, which is required for registration to complete.
 
 Python module import provides process-level one-time execution in the ordinary
 path. Some patches also preserve originals or set explicit sentinels, but this
@@ -119,6 +121,7 @@ These modules adapt upstream behavior without replacing a global symbol:
 
 | Adapter | Current purpose |
 | --- | --- |
+| [`model_executor/routing_simulator.py`](../../../afd_plugin/model_executor/routing_simulator.py) | Registers the opt-in `afd_balanced` strategy through vLLM's public `RoutingSimulator` interface. Select it with `VLLM_MOE_ROUTING_SIMULATION_STRATEGY=afd_balanced`; `AFD_BENCHMARK_FORCE_LB_TOPN_PER_RANK` optionally limits each rank's expert pool. It provides deterministic source-rank phases and normalized weights for controlled CUDA performance tests without replacing router internals. EPLB remains outside the validated scope. |
 | [`compat/npu/runtime_config.py`](../../../afd_plugin/compat/npu/runtime_config.py) | Mirrors vLLM-Ascend's non-SP all-to-all backend rewrite for custom AFD workers and reports the active NPU ubatch count. |
 | [`compat/npu/feature_validation.py`](../../../afd_plugin/compat/npu/feature_validation.py) | Parses connector-owned typed extra information through the factory and fails before execution for unsupported NPU connector, quantization, graph, DBO, gate, or async MoE combinations. |
 | [`compat/npu/forward_context.py`](../../../afd_plugin/compat/npu/forward_context.py) | Enters the pinned Ascend forward context for connector-driven FFN compute and installs AFD metadata in `additional_kwargs`. |

tests/unit/model_executor/test_routing_simulator.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project
+
+from types import SimpleNamespace
+
+import pytest
+
+torch = pytest.importorskip("torch")
+pytest.importorskip("vllm")
+
+from vllm.model_executor.layers.fused_moe.router.router_factory import (  # noqa: E402
+    create_fused_moe_router,
+)
+from vllm.model_executor.layers.fused_moe.router.routing_simulator_router import (  # noqa: E402
+    RoutingSimulatorRouter,
+)
+
+from afd_plugin.model_executor import routing_simulator as routing  # noqa: E402
+
+pytestmark = pytest.mark.vllm_runtime
+
+
+def _config(**overrides: int) -> routing._RoutingConfig:
+    values = {
+        "num_experts": 16,
+        "ep_size": 4,
+        "ep_rank": 0,
+        "top_k": 8,
+        "topn_per_rank": 2,
+        **overrides,
+    }
+    return routing._RoutingConfig(**values)
+
+
+def test_routing_cycles_are_deterministic_and_balanced() -> None:
+    full_config = _config(topn_per_rank=0)
+    first = routing._build_expert_cycle(full_config, torch.device("cpu"))
+    second = routing._build_expert_cycle(full_config, torch.device("cpu"))
+    assert torch.equal(first, second)
+    assert sorted(first.tolist()) == list(range(full_config.num_experts))
+
+    target_ranks = torch.cat(
+        [
+            routing._build_routing_buffers(
+                _config(num_experts=32, ep_rank=rank, topn_per_rank=4),
+                max_tokens=1,
+                device=torch.device("cpu"),
+                indices_dtype=torch.int32,
+            ).expert_ids.div(8, rounding_mode="floor")
+            for rank in range(4)
+        ]
+    )
+    assert torch.bincount(target_ranks.flatten(), minlength=4).tolist() == [8] * 4
+
+
+def test_invalid_routing_configs_are_rejected() -> None:
+    invalid_configs = (
+        _config(num_experts=0),
+        _config(ep_size=0),
+        _config(ep_rank=4),
+        _config(top_k=0),
+        _config(num_experts=15),
+        _config(topn_per_rank=5),
+        _config(ep_size=1, topn_per_rank=4),
+    )
+    for config in invalid_configs:
+        with pytest.raises(ValueError):
+            routing._build_expert_cycle(config, torch.device("cpu"))
+
+
+def test_strategy_returns_normalized_weights_and_reuses_buffers(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(
+        routing,
+        "get_ep_group",
+        lambda: SimpleNamespace(world_size=2, rank_in_group=0),
+    )
+    strategy = routing.AFDBalancedRoutingStrategy(topn_per_rank=2)
+
+    weights, expert_ids = strategy.route_tokens(
+        torch.empty((2, 8)),
+        torch.empty((2, 8)),
+        top_k=4,
+    )
+    reused_weights, reused_ids = strategy.route_tokens(
+        torch.empty((1, 8)),
+        torch.empty((1, 8)),
+        top_k=4,
+    )
+
+    assert weights.tolist() == [[0.25] * 4] * 2
+    assert expert_ids.tolist() == [[0, 1, 4, 5]] * 2
+    assert expert_ids.dtype == torch.int32
+    assert weights.data_ptr() == reused_weights.data_ptr()
+    assert expert_ids.data_ptr() == reused_ids.data_ptr()
+
+
+def test_strategy_registers_with_vllm_and_preserves_capture(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv(
+        "VLLM_MOE_ROUTING_SIMULATION_STRATEGY",
+        routing.AFD_BALANCED_ROUTING_STRATEGY,
+    )
+    monkeypatch.setenv("AFD_BENCHMARK_FORCE_LB_TOPN_PER_RANK", "2")
+    monkeypatch.setattr(
+        routing,
+        "get_ep_group",
+        lambda: SimpleNamespace(world_size=2, rank_in_group=0),
+    )
+    monkeypatch.setattr(routing.RoutingSimulator, "_routing_strategies", {})
+    routing.register_afd_balanced_routing_strategy()
+
+    router = create_fused_moe_router(top_k=4, global_num_experts=8)
+    captured_ids: list[torch.Tensor] = []
+    router.set_capture_fn(captured_ids.append)
+    _, expert_ids = router.select_experts(
+        torch.empty((1, 8)),
+        torch.empty((1, 8)),
+    )
+
+    assert isinstance(router, RoutingSimulatorRouter)
+    assert expert_ids.tolist() == [[0, 1, 4, 5]]
+    assert len(captured_ids) == 1
+    assert torch.equal(captured_ids[0], expert_ids)