diff --git a/afd_plugin/__init__.py b/afd_plugin/__init__.py index 0a20a472..2c889d8d 100644 --- a/afd_plugin/__init__.py +++ b/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 diff --git a/afd_plugin/model_executor/routing_simulator.py b/afd_plugin/model_executor/routing_simulator.py new file mode 100644 index 00000000..ac9cd9bd --- /dev/null +++ b/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", +] diff --git a/docs/design/module/compatibility_and_patches.md b/docs/design/module/compatibility_and_patches.md index dd461d55..b5174cca 100644 --- a/docs/design/module/compatibility_and_patches.md +++ b/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`. | diff --git a/tests/unit/model_executor/test_routing_simulator.py b/tests/unit/model_executor/test_routing_simulator.py new file mode 100644 index 00000000..6e676861 --- /dev/null +++ b/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)