From 107a0022f73214dc61d0394718177eb5bfbc58db Mon Sep 17 00:00:00 2001 From: specture724 Date: Thu, 13 Aug 2026 12:40:37 +0800 Subject: [PATCH 1/2] init: first impliment for async GPU connector Signed-off-by: specture724 --- afd_plugin/compat/npu/feature_validation.py | 4 +- afd_plugin/config.py | 23 +- afd_plugin/connectors/async_topology.py | 94 ++ afd_plugin/connectors/factory.py | 5 + afd_plugin/connectors/gpu/__init__.py | 3 +- afd_plugin/connectors/gpu/async_gpu.py | 889 ++++++++++++++++++ afd_plugin/connectors/gpu/nvshmem_rt.py | 270 ++++++ afd_plugin/connectors/gpu/symm_window.py | 528 +++++++++++ afd_plugin/connectors/npu/async_cam.py | 75 +- .../model_executor/models/deepseek_v2.py | 29 +- .../model_executor/models/gpu/__init__.py | 3 + .../models/gpu/deepseek_v2_attention_gate.py | 126 +++ .../npu/deepseek_v2_async_cam_forward.py | 4 +- .../v1/worker/attention_model_runner.py | 80 +- afd_plugin/v1/worker/ffn_model_runner.py | 77 +- afd_plugin/v1/worker/ffn_worker.py | 24 +- .../v1/worker/npu/attention_model_runner.py | 4 +- pyproject.toml | 1 + .../deepseek_v2_lite/1a1f_eager_async.sh | 128 +++ .../deepseek_v2_lite/2a2f_eager_async.sh | 129 +++ tests/e2e/async_gpu_connector_e2e.py | 241 +++++ tests/unit/config/test_config.py | 20 +- .../connectors/test_async_gpu_connector.py | 314 +++++++ .../models/test_deepseek_v2_proxy.py | 4 +- .../models/test_forward_context.py | 7 +- .../v1/worker/test_attention_model_runner.py | 42 +- tests/unit/v1/worker/test_ffn_model_runner.py | 34 +- 27 files changed, 3025 insertions(+), 133 deletions(-) create mode 100644 afd_plugin/connectors/async_topology.py create mode 100644 afd_plugin/connectors/gpu/async_gpu.py create mode 100644 afd_plugin/connectors/gpu/nvshmem_rt.py create mode 100644 afd_plugin/connectors/gpu/symm_window.py create mode 100644 afd_plugin/model_executor/models/gpu/__init__.py create mode 100644 afd_plugin/model_executor/models/gpu/deepseek_v2_attention_gate.py create mode 100644 recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/1a1f_eager_async.sh create mode 100644 recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/2a2f_eager_async.sh create mode 100644 tests/e2e/async_gpu_connector_e2e.py create mode 100644 tests/unit/connectors/test_async_gpu_connector.py diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index 823144b7..a07f3c4e 100644 --- a/afd_plugin/compat/npu/feature_validation.py +++ b/afd_plugin/compat/npu/feature_validation.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING from afd_plugin.config import ( - AFD_ASYNC_CONNECTOR, + AFD_ASYNC_NPU_CONNECTOR, AFDConfig, is_afd_async_dp, parse_afd_config, @@ -34,7 +34,7 @@ def fail_if_unsupported_npu_afd_features( vllm_config, ) - if afd_config.connector == AFD_ASYNC_CONNECTOR: + if afd_config.connector == AFD_ASYNC_NPU_CONNECTOR: _fail_if_unsupported_npu_afd_async_features( vllm_config, afd_config, diff --git a/afd_plugin/config.py b/afd_plugin/config.py index 256590ab..c13a7f55 100644 --- a/afd_plugin/config.py +++ b/afd_plugin/config.py @@ -16,14 +16,22 @@ from vllm.config import VllmConfig AFD_ADDITIONAL_CONFIG_KEY: Final[str] = "afd" -AFD_ASYNC_CONNECTOR: Final[str] = "CAMAsyncAFDConnector" +AFD_ASYNC_NPU_CONNECTOR: Final[str] = "CAMAsyncAFDConnector" +AFD_ASYNC_GPU_CONNECTOR: Final[str] = "GpuAsyncAFDConnector" +# Connectors that drive FFN work from the connector receive loop instead of a DP +# metadata control plane, and therefore need the async-DP engine patches. The +# patches are platform-neutral; only the connector below them differs. +AFD_ASYNC_CONNECTORS: Final[frozenset[str]] = frozenset( + {AFD_ASYNC_NPU_CONNECTOR, AFD_ASYNC_GPU_CONNECTOR}, +) AFDRole = Literal["attention", "ffn"] SUPPORTED_AFD_ROLES: Final[tuple[str, ...]] = ("attention", "ffn") SUPPORTED_AFD_CONNECTORS: Final[tuple[str, ...]] = ( "P2pNcclAFDConnector", "CAMP2pAFDConnector", - AFD_ASYNC_CONNECTOR, + AFD_ASYNC_NPU_CONNECTOR, + AFD_ASYNC_GPU_CONNECTOR, ) _ALIASES: Final[dict[str, str]] = { @@ -283,7 +291,7 @@ def is_afd_async_dp(vllm_config: VllmConfig) -> bool: return ( config is not None and config.async_dp - and config.connector == AFD_ASYNC_CONNECTOR + and config.connector in AFD_ASYNC_CONNECTORS ) @@ -307,9 +315,10 @@ def validate_afd_config( "AFD connector must be one of " f"{SUPPORTED_AFD_CONNECTORS!r}, got {config.connector!r}", ) - if config.async_dp and config.connector != AFD_ASYNC_CONNECTOR: + if config.async_dp and config.connector not in AFD_ASYNC_CONNECTORS: raise ValueError( - "AFD async mode requires connector='CAMAsyncAFDConnector'", + "AFD async mode requires one of " + f"{sorted(AFD_ASYNC_CONNECTORS)!r}, got {config.connector!r}", ) if config.connector == "P2pNcclAFDConnector": from afd_plugin.distributed import validate_p2p_topology @@ -331,7 +340,9 @@ def validate_afd_config( __all__ = [ "AFDConfig", - "AFD_ASYNC_CONNECTOR", + "AFD_ASYNC_NPU_CONNECTOR", + "AFD_ASYNC_CONNECTORS", + "AFD_ASYNC_GPU_CONNECTOR", "afd_config_from_mapping", "AFD_ADDITIONAL_CONFIG_KEY", "AFDRole", diff --git a/afd_plugin/connectors/async_topology.py b/afd_plugin/connectors/async_topology.py new file mode 100644 index 00000000..20e267a3 --- /dev/null +++ b/afd_plugin/connectors/async_topology.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Rank layout shared by the asynchronous AFD connectors. + +Both async connectors -- Ascend CAM and CUDA NVSHMEM -- lay their world out +Attention-first and derive expert placement the same way. Keeping that here lets +the CUDA connector reuse it without importing a backend module. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from afd_plugin.config import AFDConfig + +ASYNC_MOE_REQUEST_SPLIT = "request" +ATTN_RANKS_PER_DP_CONFIG_KEY = "attn_ranks_per_dp" + + +@dataclass(frozen=True, slots=True) +class AFDAsyncTopology: + """Role-local and world rank information for one async participant.""" + + role: str + role_rank: int + world_rank: int + attn_size: int + ffn_size: int + expert_per_rank: int + + @property + def world_size(self) -> int: + """Return the total number of Attention and FFN ranks.""" + return self.attn_size + self.ffn_size + + +def build_async_topology( + afd_config: AFDConfig, + role_rank: int, + *, + num_routed_experts: int | None = None, +) -> AFDAsyncTopology: + """Validate role-local rank settings and derive the async world rank. + + The world is Attention-first: Attention role rank ``i`` maps to world rank + ``i`` and FFN role rank ``j`` maps to ``num_attention_ranks + j``. Routed + experts are distributed across FFN ranks using a ceiling division; + production model layouts should keep the routed-expert count divisible by + the FFN rank count. + """ + attn_size = afd_config.num_attention_ranks + ffn_size = afd_config.num_ffn_ranks + if attn_size <= 0 or ffn_size <= 0: + raise ValueError("AFD async topology sizes must be positive") + if role_rank < 0: + raise ValueError(f"AFD async role rank must be non-negative, got {role_rank}") + + if afd_config.role == "attention": + if role_rank >= attn_size: + raise ValueError( + "Attention role rank must be within attention size " + f"(rank={role_rank}, size={attn_size})", + ) + world_rank = role_rank + elif afd_config.role == "ffn": + if role_rank >= ffn_size: + raise ValueError( + "FFN role rank must be within FFN size " + f"(rank={role_rank}, size={ffn_size})", + ) + world_rank = attn_size + role_rank + else: + raise ValueError(f"unknown AFD role {afd_config.role!r}") + + expert_count = num_routed_experts or 1 + expert_per_rank = (expert_count + ffn_size - 1) // ffn_size + return AFDAsyncTopology( + role=afd_config.role, + role_rank=role_rank, + world_rank=world_rank, + attn_size=attn_size, + ffn_size=ffn_size, + expert_per_rank=expert_per_rank, + ) + + +__all__ = [ + "ASYNC_MOE_REQUEST_SPLIT", + "ATTN_RANKS_PER_DP_CONFIG_KEY", + "AFDAsyncTopology", + "build_async_topology", +] diff --git a/afd_plugin/connectors/factory.py b/afd_plugin/connectors/factory.py index 6e94e1ac..8a962d0c 100644 --- a/afd_plugin/connectors/factory.py +++ b/afd_plugin/connectors/factory.py @@ -100,6 +100,11 @@ def parse_connector_extra_info( "afd_plugin.connectors.npu.async_cam", "CAMAsyncAFDConnector", ) +AFDConnectorFactory.register_connector( + "GpuAsyncAFDConnector", + "afd_plugin.connectors.gpu.async_gpu", + "GpuAsyncAFDConnector", +) __all__ = ["AFDConnectorFactory"] diff --git a/afd_plugin/connectors/gpu/__init__.py b/afd_plugin/connectors/gpu/__init__.py index 62314259..77e73746 100644 --- a/afd_plugin/connectors/gpu/__init__.py +++ b/afd_plugin/connectors/gpu/__init__.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project """GPU-specific AFD connector implementations.""" +from afd_plugin.connectors.gpu.async_gpu import GpuAsyncAFDConnector from afd_plugin.connectors.gpu.p2p import P2pNcclAFDConnector -__all__ = ["P2pNcclAFDConnector"] +__all__ = ["GpuAsyncAFDConnector", "P2pNcclAFDConnector"] diff --git a/afd_plugin/connectors/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py new file mode 100644 index 00000000..add5c69b --- /dev/null +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -0,0 +1,889 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""NVSHMEM-backed asynchronous connector for CUDA AFD. + +``GpuAsyncAFDConnector`` is the CUDA counterpart of ``CAMAsyncAFDConnector``: +Attention ranks run MoE routing, write routed tokens one-sided into the FFN +ranks' symmetric windows, and later reduce the weighted expert output; FFN ranks +poll their window, run their local experts, and write the result back. There is +no DP metadata control plane (``control_plane`` stays ``None``) and FFN work is +driven directly by the connector receive loop, so Attention DP replicas never +wait for each other. + +The world is Attention-first, ``[A0, A1, ..., F0, F1, ...]``, matching +``CAMAsyncAFDConnector``. Every Attention rank routes to every FFN rank, so an +FFN window holds one region per Attention rank and vice versa. + +See ``docs/design/rfc_async_gpu_connector.md``. Supported deployment requires +``async=true``, ``compute_gate_on_attention=true``, eager execution, prefill +only, and a single node. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from typing import TYPE_CHECKING, Any, Final + +import torch +from torch import Tensor +from vllm.logger import init_logger + +from afd_plugin.config import AFDConfig +from afd_plugin.config_utils import ( + coerce_extra_bool, + coerce_extra_positive_int, + coerce_extra_str, +) +from afd_plugin.connectors.async_topology import ( + ASYNC_MOE_REQUEST_SPLIT, + build_async_topology, +) +from afd_plugin.connectors.base import AFDConnectorBase, ConnectorExtraInfo +from afd_plugin.connectors.gpu.symm_window import ( + FLAG_SHUTDOWN_BIT, + SlotLayout, + SymmWindow, + encode_header, +) +from afd_plugin.connectors.metadata import ( + AFDA2FTransferPayload, + AFDF2ATransferPayload, + AFDTransferContext, + AFDTransferMetadata, + AFDTransferState, +) +from afd_plugin.distributed import init_afd_process_group + +if TYPE_CHECKING: + from torch.distributed.distributed_c10d import ProcessGroup + from vllm.config import VllmConfig + +AFD_ASYNC_GPU_GROUP_NAME = "afd_async_gpu" + +_GPU_ASYNC_EXTRA_CONFIG_FIELDS: Final[frozenset[str]] = frozenset( + { + "attn_ranks_per_dp", + "ring_depth", + "routed_cap_multiplier", + "recv_poll_timeout_ms", + "async_moe_ubatching", + "async_moe_num_ubatches", + "async_moe_split", + }, +) + +# Name the logger inside vLLM's tree: vLLM installs its handler on the "vllm" +# logger only, so a bare ``afd_plugin.*`` logger propagates to a handler-less +# root and every line is dropped -- which is how the window summary, the only +# report of a multi-GiB allocation, stayed invisible. +logger = init_logger(f"vllm.{__name__}") + + +def _coerce_extra_float(value: Any, *, field_name: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float, str)): + raise TypeError(f"{field_name} must be a number, got {type(value).__name__}") + result = float(value) + if result <= 0.0: + raise ValueError(f"{field_name} must be positive, got {result}") + return result + + +@dataclass(frozen=True) +class GpuAsyncExtraInfo(ConnectorExtraInfo): + """Typed async GPU connector configuration. + + Attributes: + attn_ranks_per_dp: Number of Attention ranks in each data-parallel group. + ring_depth: Slots per peer region. Derived from the send-then-recv + invariant, not a performance knob: an Attention rank has at most one + in-flight request per ``(peer, stage)``, so ``num_stages`` suffices. + routed_cap_multiplier: Headroom over the balanced routed-token estimate. + Real gates are not balanced -- DeepSeek-V2-Lite at 2 FFN ranks was + observed 1.33x over the even split -- so the default leaves room. + The true worst case is ``ffn_size`` (every partial to one rank). + recv_poll_timeout_ms: Idle poll timeout on the FFN loop; bounds shutdown + response time. + async_moe_ubatching: Whether request-boundary async MoE ubatching is used. + async_moe_num_ubatches: Number of stages used by async MoE ubatching. + async_moe_split: Boundary at which async MoE work is split. + """ + + attn_ranks_per_dp: int = 1 + ring_depth: int = 0 + routed_cap_multiplier: float = 2.0 + recv_poll_timeout_ms: int = 50 + async_moe_ubatching: bool = False + async_moe_num_ubatches: int = 2 + async_moe_split: str = ASYNC_MOE_REQUEST_SPLIT + + @classmethod + def from_mapping(cls, raw: Mapping[str, Any] | None) -> GpuAsyncExtraInfo: + if raw is None: + raw = {} + if not isinstance(raw, Mapping): + raise TypeError( + f"{cls.__name__} connector_extra_config must be a mapping, " + f"got {type(raw).__name__}", + ) + unknown = sorted( + str(key) for key in raw if key not in _GPU_ASYNC_EXTRA_CONFIG_FIELDS + ) + if unknown: + raise ValueError( + "unknown AFD async GPU connector_extra_config field(s): " + + ", ".join(unknown), + ) + + ubatching = coerce_extra_bool( + raw.get("async_moe_ubatching", False), + field_name="async_moe_ubatching", + ) + num_ubatches = coerce_extra_positive_int( + raw.get("async_moe_num_ubatches", 2), + field_name="async_moe_num_ubatches", + ) + # Ring depth follows the number of live stages unless pinned explicitly. + ring_depth = coerce_extra_positive_int( + raw.get("ring_depth", num_ubatches if ubatching else 1), + field_name="ring_depth", + ) + return cls( + attn_ranks_per_dp=coerce_extra_positive_int( + raw.get("attn_ranks_per_dp", 1), + field_name="attn_ranks_per_dp", + ), + ring_depth=ring_depth, + routed_cap_multiplier=_coerce_extra_float( + raw.get("routed_cap_multiplier", 2.0), + field_name="routed_cap_multiplier", + ), + recv_poll_timeout_ms=coerce_extra_positive_int( + raw.get("recv_poll_timeout_ms", 50), + field_name="recv_poll_timeout_ms", + ), + async_moe_ubatching=ubatching, + async_moe_num_ubatches=num_ubatches, + async_moe_split=coerce_extra_str( + raw.get("async_moe_split", ASYNC_MOE_REQUEST_SPLIT), + field_name="async_moe_split", + ), + ) + + def to_mapping(self) -> dict[str, Any]: + return { + "attn_ranks_per_dp": self.attn_ranks_per_dp, + "ring_depth": self.ring_depth, + "routed_cap_multiplier": self.routed_cap_multiplier, + "recv_poll_timeout_ms": self.recv_poll_timeout_ms, + "async_moe_ubatching": self.async_moe_ubatching, + "async_moe_num_ubatches": self.async_moe_num_ubatches, + "async_moe_split": self.async_moe_split, + } + + +class ConnectorShutdown(RuntimeError): # noqa: N818 + """Raised on the FFN loop when a peer announced shutdown.""" + + +@dataclass(slots=True) +class GpuAsyncTransferState(AFDTransferState): + """FFN-side state carried from dispatch recv through combine send. + + ``region``/``ring`` locate the window slot so ``send_ffn_work_item_output`` + can write back to the originating Attention rank and release the slot; + ``route_table`` is echoed so the Attention side can scatter the result. + """ + + region: int = 0 + ring: int = 0 + src_role_rank: int = 0 + layer_idx: int = 0 + stage_idx: int = 0 + seq: int = 0 + num_tokens: int = 0 + routed_tokens: int = 0 + shared_tokens: int = 0 + group_list: Tensor | None = None + route_table: Tensor | None = None + shared_idx: Tensor | None = None + expand_x_shared: Tensor | None = None + + +@dataclass(slots=True) +class GpuAsyncFFNWorkItem: + """Normalized FFN-side work item produced by a window arrival.""" + + hidden_states: Tensor + context: AFDTransferContext + recv_output: AFDA2FTransferPayload + layer_idx: int + stage_idx: int + num_tokens: int + total_num_tokens: int + shared_num_tokens: int + + +@dataclass(slots=True) +class _PendingDispatch: + """Attention-side record of one in-flight layer, popped by combine recv.""" + + context: AFDTransferContext + topk_weights: Tensor + num_tokens: int + ring: int + seq: int + expected_ffn: list[int] + + +def plan_dispatch( + topk_ids: Tensor, + *, + ffn_size: int, + expert_per_rank: int, +) -> tuple[Tensor, Tensor, Tensor]: + """Cluster ``(token, topk_slot)`` partials by destination FFN rank. + + Sorting by the global expert id groups partials by destination rank and, in + the same pass, by local expert inside each destination -- the two orderings + the receiver needs. Returns ``(route_table, counts, offsets)`` where + ``route_table[i] = (token_idx, topk_slot)`` in that order, ``counts`` holds + per-global-expert partial counts padded to ``ffn_size * expert_per_rank``, + and ``offsets`` is the exclusive prefix sum of ``counts``. + """ + num_slots = topk_ids.shape[1] + flat = topk_ids.reshape(-1).to(torch.int64) + order = torch.argsort(flat, stable=True) + counts = torch.bincount(flat, minlength=ffn_size * expert_per_rank) + offsets = torch.cumsum(counts, dim=0) - counts + route_table = torch.stack( + (order // num_slots, order % num_slots), + dim=1, + ).to(torch.int32) + return route_table, counts, offsets + + +def combine_scatter( + accumulator: Tensor, + *, + routed_out: Tensor, + route_table: Tensor, + topk_weights: Tensor, +) -> None: + """Weighted-scatter one FFN rank's routed output into ``accumulator``.""" + token_idx = route_table[:, 0].to(torch.int64) + slot_idx = route_table[:, 1].to(torch.int64) + weights = topk_weights[token_idx, slot_idx].to(accumulator.dtype) + accumulator.index_add_( + 0, + token_idx, + routed_out.to(accumulator.dtype) * weights.unsqueeze(1), + ) + + +class GpuAsyncAFDConnector(AFDConnectorBase): + """NVSHMEM symmetric-window asynchronous connector for CUDA AFD.""" + + control_plane = None + + @classmethod + def parse_extra_config( + cls, + raw: Mapping[str, Any] | None, + ) -> GpuAsyncExtraInfo: + return GpuAsyncExtraInfo.from_mapping(raw) + + def __init__( + self, + rank: int, + local_rank: int, + vllm_config: VllmConfig, + afd_config: AFDConfig, + role_rank: int, + ) -> None: + super().__init__(rank, local_rank, vllm_config, afd_config, role_rank) + self._initialized = False + hf_config = vllm_config.model_config.hf_config + self.hidden_size = hf_config.hidden_size + self.topk = hf_config.num_experts_per_tok + self.num_routed_experts = hf_config.n_routed_experts + self.payload_dtype = vllm_config.model_config.dtype + self.max_seq_len = vllm_config.scheduler_config.max_num_batched_tokens + self.tp_size = self.extra_info.attn_ranks_per_dp + + self.topology = build_async_topology( + afd_config, + role_rank, + num_routed_experts=self.num_routed_experts, + ) + self.world_rank = self.topology.world_rank + self.attn_size = self.topology.attn_size + self.ffn_size = self.topology.ffn_size + self.expert_per_rank = self.topology.expert_per_rank + self.is_attention = afd_config.role == "attention" + + self.ring_depth = self.extra_info.ring_depth + # Every Attention rank routes to every FFN rank, so a window carries one + # region per opposite-role peer. Both roles allocate the larger of the + # two so the symmetric allocation matches. + self.num_regions = max(self.attn_size, self.ffn_size) + routed_cap = int( + -(-self.max_seq_len * self.topk // self.ffn_size) + * self.extra_info.routed_cap_multiplier, + ) + self.routed_cap = max(1, routed_cap) + self.token_cap = max(1, self.max_seq_len) + self.layout = SlotLayout.build( + expert_per_rank=self.expert_per_rank, + routed_cap=self.routed_cap, + token_cap=self.token_cap, + hidden_size=self.hidden_size, + payload_itemsize=torch.empty(0, dtype=self.payload_dtype).element_size(), + ) + + self.pg: ProcessGroup | None = None + self.window: SymmWindow | None = None + self._seq = 0 + self._pending: dict[int, list[_PendingDispatch]] = {} + self._free_rings: dict[int, list[int]] = {} + + @property + def is_initialized(self) -> bool: + return self._initialized + + def init_afd_connector(self) -> None: + """Collectively create the AFD world group and the symmetric window. + + All Attention and FFN ranks must call this with identical rendezvous and + topology settings; the window allocation is symmetric, so a mismatched + size fails here rather than corrupting a later transfer. + """ + if self._initialized: + return + + self.pg = init_afd_process_group( + backend="nccl", + init_method=f"tcp://{self.afd_config.host}:{self.afd_config.port}", + world_size=self.topology.world_size, + rank=self.world_rank, + group_name=AFD_ASYNC_GPU_GROUP_NAME, + timeout=timedelta(minutes=30), + ) + device = torch.device("cuda", self.local_rank) + self.window = SymmWindow( + num_regions=self.num_regions, + ring_depth=self.ring_depth, + layout=self.layout, + payload_dtype=self.payload_dtype, + device=device, + group=self.pg, + rank=self.world_rank, + world_size=self.topology.world_size, + ) + for stage in range(max(1, self.extra_info.async_moe_num_ubatches)): + self._free_rings[stage] = list(range(self.ring_depth)) + logger.info( + "AFD async GPU window ready: role=%s role_rank=%d world_rank=%d/%d " + "regions=%d rings=%d routed_cap=%d slot=%.1fMiB total=%.1fMiB", + self.afd_config.role, + self.role_rank, + self.world_rank, + self.topology.world_size, + self.num_regions, + self.ring_depth, + self.routed_cap, + self.layout.slot_bytes / 2**20, + self.window.total_bytes / 2**20, + ) + self._initialized = True + + def close(self) -> None: + if self.window is not None: + self.window.close() + self.window = None + if self.pg is not None: + import torch.distributed as dist + + dist.destroy_process_group(self.pg) + self.pg = None + self._pending.clear() + self._free_rings.clear() + self._initialized = False + + def select_experts(self, **kwargs: Any) -> tuple[Tensor, Tensor]: + """Run vLLM's grouped top-k on the Attention side. + + ``compute_gate_topk`` delegates expert selection to the connector so the + CAM and CUDA paths can share one gate; this is the CUDA half. + """ + from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( + grouped_topk, + ) + + if kwargs.get("mix_placement"): + raise RuntimeError( + "AFD async GPU connector does not support mix_placement", + ) + return grouped_topk( + hidden_states=kwargs["hidden_states"], + gating_output=kwargs["router_logits"], + topk=kwargs["top_k"], + renormalize=kwargs["renormalize"], + num_expert_group=kwargs.get("num_expert_group", 0), + topk_group=kwargs.get("topk_group", 0), + scoring_func=kwargs.get("scoring_func", "softmax"), + routed_scaling_factor=kwargs.get("routed_scaling_factor", 1.0), + e_score_correction_bias=kwargs.get("e_score_correction_bias"), + ) + + def _require_initialized(self) -> SymmWindow: + if not self._initialized or self.window is None: + raise RuntimeError("AFD async GPU connector is not initialized") + return self.window + + # ================================================================== + # Attention-side data path + # ================================================================== + + def send_attn_output( + self, + hidden_states: Tensor, + context: AFDTransferContext, + **kwargs: Any, + ) -> None: + """Route this layer's tokens and write them into every FFN window. + + ``topk_ids``/``topk_weights`` come from the Attention-side gate. Weights + stay local -- only the routed activations and their route table go on the + wire, and the weighting happens in ``recv_ffn_output``. + """ + window = self._require_initialized() + topk_ids: Tensor | None = kwargs.get("topk_ids") + topk_weights: Tensor | None = kwargs.get("topk_weights") + if topk_ids is None or topk_weights is None: + raise RuntimeError( + "AFD async GPU send_attn_output requires topk_ids and " + "topk_weights from the Attention-side gate", + ) + metadata = context.metadata + num_tokens = metadata.total_tokens + if hidden_states.shape[0] != num_tokens: + raise ValueError( + f"hidden_states has {hidden_states.shape[0]} rows but metadata " + f"expects {num_tokens}", + ) + if tuple(topk_ids.shape) != (num_tokens, self.topk): + raise ValueError( + f"topk_ids shape must be ({num_tokens}, {self.topk}), " + f"got {tuple(topk_ids.shape)}", + ) + + stage_idx = metadata.stage_idx + rings = self._free_rings.setdefault(stage_idx, list(range(self.ring_depth))) + if not rings: + raise RuntimeError( + f"AFD async GPU ring exhausted on stage {stage_idx}; the " + "send-then-recv invariant was violated or the topology config " + "does not match the actual peer count", + ) + ring = rings.pop(0) + self._seq += 1 + + route_table, counts, _ = plan_dispatch( + topk_ids, + ffn_size=self.ffn_size, + expert_per_rank=self.expert_per_rank, + ) + # One D2H per send: offsets are a prefix sum, cheaper to redo on host + # than to fetch a second tensor. + counts_host = counts.cpu().tolist() + offsets_host = [0] * len(counts_host) + for i in range(1, len(counts_host)): + offsets_host[i] = offsets_host[i - 1] + counts_host[i - 1] + token_ids = route_table[:, 0].to(torch.int64) + + # Every FFN rank gets a slot even when routing sends it nothing, and it + # replies to every slot, so a reply is expected from all of them. + # Expecting only the ranks that received data leaves the empty rank's + # reply unmatched and its ring slot never released -- which is what a + # single-token decode hits, since both the routed segment and the + # round-robin shared slice can come out empty for one rank. + expected_ffn = list(range(self.ffn_size)) + for ffn_rank in range(self.ffn_size): + base = ffn_rank * self.expert_per_rank + expert_counts = counts_host[base : base + self.expert_per_rank] + start = offsets_host[base] + routed_tokens = sum(expert_counts) + segment = slice(start, start + routed_tokens) + + # Shared-expert tokens are split round-robin across FFN ranks. + shared_idx = torch.arange( + ffn_rank, + num_tokens, + self.ffn_size, + device=hidden_states.device, + dtype=torch.int32, + ) + header = encode_header( + self.layout, + seq=self._seq, + src_role_rank=self.role_rank, + layer_idx=metadata.layer_idx, + stage_idx=stage_idx, + num_tokens=num_tokens, + routed_tokens=routed_tokens, + shared_tokens=int(shared_idx.numel()), + topk=self.topk, + flags=0, + expert_counts=expert_counts, + ) + window.write_slot( + peer=self.attn_size + ffn_rank, + region=self.role_rank, + ring=ring, + header=header, + route_table=route_table[segment], + routed_x=hidden_states.index_select(0, token_ids[segment]), + shared_idx=shared_idx, + shared_x=hidden_states.index_select(0, shared_idx.to(torch.int64)), + ) + + logger.debug( + "AFD dispatch sent: A%d layer=%d stage=%d tokens=%d ring=%d " + "awaiting_ffn=%s", + self.role_rank, + metadata.layer_idx, + stage_idx, + num_tokens, + ring, + expected_ffn, + ) + self._pending.setdefault(stage_idx, []).append( + _PendingDispatch( + context=context, + topk_weights=topk_weights, + num_tokens=num_tokens, + ring=ring, + seq=self._seq, + expected_ffn=expected_ffn, + ), + ) + + def recv_ffn_output( + self, + ref_tensor: Tensor, + ubatch_idx: int = 0, + **kwargs: Any, + ) -> Tensor: + """Wait for this layer's expert output and reduce it back to ``[B, H]``.""" + + window = self._require_initialized() + queue = self._pending.get(ubatch_idx) + if not queue: + raise RuntimeError( + f"AFD async GPU recv_ffn_output has no pending dispatch on " + f"stage {ubatch_idx}", + ) + pending = queue.pop(0) + + accumulator = torch.zeros( + (pending.num_tokens, self.hidden_size), + dtype=torch.float32, + device=ref_tensor.device, + ) + outstanding = set(pending.expected_ffn) + while outstanding: + arrived = window.poll() + if arrived is None: + continue + header = arrived.header + if header.is_shutdown: + raise ConnectorShutdown( + f"FFN rank {header.src_role_rank} announced shutdown", + ) + if header.src_role_rank not in outstanding: + raise RuntimeError( + "AFD async GPU combine received an unexpected FFN rank " + f"{header.src_role_rank}; expected one of {sorted(outstanding)}", + ) + if header.echo_seq != pending.seq: + raise RuntimeError( + "AFD async GPU combine answered dispatch seq " + f"{header.echo_seq} while waiting on {pending.seq} " + f"(F{header.src_role_rank}, layer {header.layer_idx}); the " + "pending FIFO and the wire have diverged", + ) + outstanding.discard(header.src_role_rank) + logger.debug( + "AFD combine recv: A%d <- F%d layer=%d routed=%d still_waiting=%s", + self.role_rank, + header.src_role_rank, + header.layer_idx, + header.routed_tokens, + sorted(outstanding), + ) + + if header.routed_tokens: + combine_scatter( + accumulator, + routed_out=window.local_routed( + arrived.region, + arrived.ring, + header.routed_tokens, + ), + route_table=window.local_route_table( + arrived.region, + arrived.ring, + header.routed_tokens, + ), + topk_weights=pending.topk_weights, + ) + if header.shared_tokens: + accumulator.index_add_( + 0, + window.local_shared_idx( + arrived.region, + arrived.ring, + header.shared_tokens, + ).to(torch.int64), + window.local_shared( + arrived.region, + arrived.ring, + header.shared_tokens, + ).to(accumulator.dtype), + ) + + self._free_rings.setdefault(ubatch_idx, []).append(pending.ring) + return accumulator.to(ref_tensor.dtype) + + # ================================================================== + # FFN-side data path + # ================================================================== + + def recv_attn_output( + self, + ubatch_idx: int = 0, + **kwargs: Any, + ) -> AFDA2FTransferPayload: + """Block until one Attention rank's routed tokens arrive. + + The layer index, token counts, and per-expert group list all come from + the arrived slot header; the FFN side knows none of them beforehand. + """ + window = self._require_initialized() + timeout_ms = int(kwargs.get("timeout_ms", 0)) + deadline = None + if timeout_ms: + import time + + deadline = time.monotonic() + timeout_ms / 1000.0 + + while True: + arrived = window.poll() + if arrived is not None: + break + if deadline is not None: + import time + + if time.monotonic() >= deadline: + raise TimeoutError("AFD async GPU dispatch recv timed out") + + header = arrived.header + if header.is_shutdown: + raise ConnectorShutdown( + f"Attention rank {header.src_role_rank} announced shutdown", + ) + + states = GpuAsyncTransferState( + region=arrived.region, + ring=arrived.ring, + seq=header.seq, + src_role_rank=header.src_role_rank, + layer_idx=header.layer_idx, + stage_idx=header.stage_idx, + num_tokens=header.num_tokens, + routed_tokens=header.routed_tokens, + shared_tokens=header.shared_tokens, + group_list=torch.tensor( + header.expert_counts, + dtype=torch.int64, + device=torch.device("cuda", self.local_rank), + ), + route_table=window.local_route_table( + arrived.region, + arrived.ring, + header.routed_tokens, + ), + shared_idx=window.local_shared_idx( + arrived.region, + arrived.ring, + header.shared_tokens, + ), + expand_x_shared=window.local_shared( + arrived.region, + arrived.ring, + header.shared_tokens, + ), + ) + logger.debug( + "AFD dispatch recv: F%d <- A%d layer=%d stage=%d routed=%d shared=%d " + "region=%d ring=%d", + self.role_rank, + header.src_role_rank, + header.layer_idx, + header.stage_idx, + header.routed_tokens, + header.shared_tokens, + arrived.region, + arrived.ring, + ) + metadata = AFDTransferMetadata.create_ffn_metadata( + layer_idx=header.layer_idx, + stage_idx=header.stage_idx, + seq_lens=[max(1, header.routed_tokens)], + ) + return AFDA2FTransferPayload( + hidden_states=window.local_routed( + arrived.region, + arrived.ring, + header.routed_tokens, + ), + context=AFDTransferContext(metadata=metadata, states=states), + ) + + def send_ffn_output( + self, + ffn_output: Tensor, + context: AFDTransferContext, + **kwargs: Any, + ) -> None: + """Write expert output back to the originating Attention rank.""" + window = self._require_initialized() + states = context.states + if not isinstance(states, GpuAsyncTransferState): + raise RuntimeError( + "AFD async GPU send_ffn_output requires GpuAsyncTransferState", + ) + shared_output: Tensor | None = kwargs.get("shared_output") + self._seq += 1 + header = encode_header( + self.layout, + seq=self._seq, + src_role_rank=self.role_rank, + layer_idx=states.layer_idx, + stage_idx=states.stage_idx, + num_tokens=states.num_tokens, + routed_tokens=states.routed_tokens, + shared_tokens=states.shared_tokens if shared_output is not None else 0, + topk=self.topk, + flags=0, + expert_counts=list(header_counts(states)), + echo_seq=states.seq, + ) + window.write_slot( + peer=states.src_role_rank, + region=self.role_rank, + ring=states.ring, + header=header, + route_table=states.route_table, + routed_x=ffn_output, + shared_idx=states.shared_idx if shared_output is not None else None, + shared_x=shared_output, + ) + + # ================================================================== + # Connector-driven FFN loop + # ================================================================== + + def recv_ffn_work_item( + self, + *, + stage_idx: int, + max_num_tokens: int, + ) -> GpuAsyncFFNWorkItem: + """Receive and normalize one connector-driven FFN dispatch item.""" + recv_output = self.recv_attn_output( + ubatch_idx=stage_idx, + timeout_ms=self.extra_info.recv_poll_timeout_ms, + ) + states = recv_output.context.states + assert isinstance(states, GpuAsyncTransferState) + return GpuAsyncFFNWorkItem( + hidden_states=recv_output.hidden_states, + context=recv_output.context, + recv_output=recv_output, + layer_idx=states.layer_idx, + stage_idx=states.stage_idx, + num_tokens=states.routed_tokens, + total_num_tokens=states.num_tokens, + shared_num_tokens=states.shared_tokens, + ) + + def send_ffn_work_item_output( + self, + work_item: GpuAsyncFFNWorkItem, + ffn_output: Tensor | AFDF2ATransferPayload, + ) -> Tensor: + """Return one work item's expert output to its Attention rank.""" + if isinstance(ffn_output, AFDF2ATransferPayload): + routed = ffn_output.routed_output + shared = ffn_output.shared_output + else: + routed = ffn_output + shared = None + self.send_ffn_output(routed, work_item.context, shared_output=shared) + return routed + + def announce_shutdown(self) -> None: + """Tell every opposite-role peer to leave its receive loop.""" + window = self._require_initialized() + peers = ( + range(self.attn_size, self.attn_size + self.ffn_size) + if self.is_attention + else range(self.attn_size) + ) + self._seq += 1 + header = encode_header( + self.layout, + seq=self._seq, + src_role_rank=self.role_rank, + layer_idx=0, + stage_idx=0, + num_tokens=0, + routed_tokens=0, + shared_tokens=0, + topk=self.topk, + flags=FLAG_SHUTDOWN_BIT, + expert_counts=[0] * self.expert_per_rank, + ) + for peer in peers: + window.write_slot( + peer=peer, + region=self.role_rank, + ring=0, + header=header, + route_table=None, + routed_x=None, + shared_idx=None, + shared_x=None, + ) + + +def header_counts(states: GpuAsyncTransferState) -> list[int]: + """Echo the per-expert group list back on the combine header.""" + if states.group_list is None: + return [] + return states.group_list.cpu().tolist() + + +__all__ = [ + "AFD_ASYNC_GPU_GROUP_NAME", + "ConnectorShutdown", + "GpuAsyncAFDConnector", + "GpuAsyncExtraInfo", + "GpuAsyncFFNWorkItem", + "GpuAsyncTransferState", + "combine_scatter", + "plan_dispatch", +] diff --git a/afd_plugin/connectors/gpu/nvshmem_rt.py b/afd_plugin/connectors/gpu/nvshmem_rt.py new file mode 100644 index 00000000..d2b3b058 --- /dev/null +++ b/afd_plugin/connectors/gpu/nvshmem_rt.py @@ -0,0 +1,270 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Minimal NVSHMEM host-library binding for the async GPU connector. + +``torch.distributed._symmetric_memory`` cannot serve AFD: its NVSHMEM backend +bootstraps on the *default* process group and carves teams out of +``NVSHMEM_TEAM_WORLD`` with ``nvshmem_team_split_strided``. Each AFD role runs +as its own ``vllm serve`` whose default group covers only that role's ranks, so +the cross-role AFD group is never a strided subset of it and team creation +fails. + +Bootstrapping NVSHMEM ourselves from a unique id exchanged over the AFD group's +store makes the AFD world *be* ``NVSHMEM_TEAM_WORLD``, so no team split is +needed. The host library's ABI is pinned by static asserts in its own headers +(``uniqueid`` 128 B, ``init_attr`` 144 B, version = ``(1 << 16) + sizeof``), +which is why ctypes is enough and ``csrc/gpu/`` stays empty. +""" + +from __future__ import annotations + +import ctypes +import os +from typing import TYPE_CHECKING, Final + +import torch + +if TYPE_CHECKING: + from torch.distributed.distributed_c10d import ProcessGroup, Store + +UNIQUEID_PADDING: Final[int] = 124 +# 128 (init_args) - 4 (version) - 24 (uid_args) - 4 (trailing alignment) +INIT_ARGS_PADDING: Final[int] = 96 +NVSHMEMX_INIT_WITH_UNIQUEID: Final[int] = 1 << 3 +_UID_STORE_KEY: Final[str] = "afd_nvshmem_uid" +_LIB_RELATIVE: Final[str] = "nvidia/nvshmem/lib/libnvshmem_host.so.3" + + +class _UniqueId(ctypes.Structure): + _fields_ = ( + ("version", ctypes.c_int), + ("internal", ctypes.c_char * UNIQUEID_PADDING), + ) + + +class _UniqueIdArgs(ctypes.Structure): + _fields_ = ( + ("version", ctypes.c_int), + ("id", ctypes.POINTER(_UniqueId)), + ("myrank", ctypes.c_int), + ("nranks", ctypes.c_int), + ) + + +class _InitArgs(ctypes.Structure): + _fields_ = ( + ("version", ctypes.c_int), + ("uid_args", _UniqueIdArgs), + ("content", ctypes.c_char * INIT_ARGS_PADDING), + ) + + +class _InitAttr(ctypes.Structure): + _fields_ = ( + ("version", ctypes.c_int), + ("mpi_comm", ctypes.c_void_p), + ("args", _InitArgs), + ) + + +def _find_host_library() -> str: + """Locate ``libnvshmem_host.so.3`` next to the installed nvshmem wheel.""" + import site + import sysconfig + + roots = [sysconfig.get_paths()["purelib"], *site.getsitepackages()] + for root in roots: + candidate = os.path.join(root, _LIB_RELATIVE) + if os.path.exists(candidate): + return candidate + raise RuntimeError( + "AFD async GPU connector requires the NVSHMEM host library; " + f"{_LIB_RELATIVE} was not found under {roots}. Install " + "nvidia-nvshmem-cu13 matching the installed torch build.", + ) + + +def _load_library() -> ctypes.CDLL: + lib = ctypes.CDLL(_find_host_library(), mode=ctypes.RTLD_GLOBAL) + lib.nvshmemx_get_uniqueid.argtypes = [ctypes.POINTER(_UniqueId)] + lib.nvshmemx_get_uniqueid.restype = ctypes.c_int + lib.nvshmemx_set_attr_uniqueid_args.argtypes = [ + ctypes.c_int, + ctypes.c_int, + ctypes.POINTER(_UniqueId), + ctypes.POINTER(_InitAttr), + ] + lib.nvshmemx_set_attr_uniqueid_args.restype = ctypes.c_int + lib.nvshmemx_hostlib_init_attr.argtypes = [ + ctypes.c_uint, + ctypes.POINTER(_InitAttr), + ] + lib.nvshmemx_hostlib_init_attr.restype = ctypes.c_int + lib.nvshmem_malloc.argtypes = [ctypes.c_size_t] + lib.nvshmem_malloc.restype = ctypes.c_void_p + lib.nvshmem_ptr.argtypes = [ctypes.c_void_p, ctypes.c_int] + lib.nvshmem_ptr.restype = ctypes.c_void_p + lib.nvshmem_my_pe.restype = ctypes.c_int + lib.nvshmem_n_pes.restype = ctypes.c_int + return lib + + +# NVSHMEM initialization is process-global: a process joins exactly one NVSHMEM +# world, so every connector in it shares this state. +_lib: ctypes.CDLL | None = None +_initialized_world: tuple[int, int] | None = None + + +def init(pg: ProcessGroup, rank: int, world_size: int) -> None: + """Join the NVSHMEM world described by ``pg``, once per process. + + Rank 0 mints the unique id and publishes it on the group's store; every rank + then initializes with the same id, so NVSHMEM's PE numbering equals the AFD + world rank. + """ + global _lib, _initialized_world + + if _initialized_world is not None: + if _initialized_world != (rank, world_size): + raise RuntimeError( + "NVSHMEM is already initialized in this process as " + f"rank {_initialized_world[0]} of {_initialized_world[1]}; " + f"cannot re-initialize as rank {rank} of {world_size}", + ) + return + + from torch.distributed.distributed_c10d import _get_process_group_store + + lib = _load_library() + store: Store = _get_process_group_store(pg) + + unique_id = _UniqueId() + unique_id.version = (1 << 16) + ctypes.sizeof(_UniqueId) + if rank == 0: + if lib.nvshmemx_get_uniqueid(ctypes.byref(unique_id)) != 0: + raise RuntimeError("nvshmemx_get_uniqueid failed") + store.set(_UID_STORE_KEY, bytes(memoryview(unique_id).cast("B"))) + else: + raw = store.get(_UID_STORE_KEY) + ctypes.memmove(ctypes.byref(unique_id), raw, ctypes.sizeof(_UniqueId)) + + attr = _InitAttr() + attr.version = (1 << 16) + ctypes.sizeof(_InitAttr) + attr.args.version = (1 << 16) + ctypes.sizeof(_InitArgs) + attr.args.uid_args.version = (1 << 16) + ctypes.sizeof(_UniqueIdArgs) + if ( + lib.nvshmemx_set_attr_uniqueid_args( + rank, + world_size, + ctypes.byref(unique_id), + ctypes.byref(attr), + ) + != 0 + ): + raise RuntimeError("nvshmemx_set_attr_uniqueid_args failed") + if ( + lib.nvshmemx_hostlib_init_attr( + NVSHMEMX_INIT_WITH_UNIQUEID, + ctypes.byref(attr), + ) + != 0 + ): + raise RuntimeError("nvshmemx_hostlib_init_attr failed") + + actual_pe, actual_world = lib.nvshmem_my_pe(), lib.nvshmem_n_pes() + if (actual_pe, actual_world) != (rank, world_size): + raise RuntimeError( + f"NVSHMEM PE numbering does not match the AFD world: got PE " + f"{actual_pe} of {actual_world}, expected {rank} of {world_size}", + ) + _lib = lib + _initialized_world = (rank, world_size) + + +def is_initialized() -> bool: + return _initialized_world is not None + + +def _require_lib() -> ctypes.CDLL: + if _lib is None: + raise RuntimeError("NVSHMEM is not initialized; call init() first") + return _lib + + +def malloc(nbytes: int) -> int: + """Allocate a symmetric buffer. Collective: every PE must call it alike.""" + pointer = _require_lib().nvshmem_malloc(nbytes) + if not pointer: + raise RuntimeError( + f"nvshmem_malloc({nbytes}) returned NULL; raise " + "NVSHMEM_SYMMETRIC_SIZE or lower the window capacity", + ) + return int(pointer) + + +def peer_ptr(local_ptr: int, pe: int) -> int: + """Map a peer's copy of a symmetric allocation into this process.""" + pointer = _require_lib().nvshmem_ptr(ctypes.c_void_p(local_ptr), pe) + if not pointer: + raise RuntimeError( + f"nvshmem_ptr returned NULL for PE {pe}: no direct peer access. " + "The async GPU connector requires PEs reachable over NVLink/P2P; " + "cross-node placement is not supported.", + ) + return int(pointer) + + +class _DeviceBuffer: + """Hand a raw device pointer to torch via ``__cuda_array_interface__``.""" + + def __init__(self, pointer: int, nbytes: int) -> None: + self.__cuda_array_interface__ = { + "data": (pointer, False), + "shape": (nbytes,), + "typestr": "|u1", + "version": 3, + "strides": None, + } + + +def tensor_from_ptr( + base_ptr: int, + *, + byte_offset: int, + sizes: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """View symmetric memory as a tensor without copying. + + The buffer is exposed as bytes and then reinterpreted, because + ``__cuda_array_interface__`` has no type string for dtypes like bfloat16. + """ + itemsize = torch.empty(0, dtype=dtype).element_size() + numel = 1 + for size in sizes: + numel *= size + if numel == 0: + # A zero-length __cuda_array_interface__ buffer is rejected by the CUDA + # runtime (cudaErrorInvalidValue); an empty slot is legitimate whenever + # routing sends a peer nothing, so hand back a plain empty tensor. + return torch.empty(sizes, dtype=dtype, device=device) + nbytes = numel * itemsize + if byte_offset % itemsize: + raise ValueError( + f"byte offset {byte_offset} is not aligned to {itemsize}-byte {dtype}", + ) + raw = torch.as_tensor( + _DeviceBuffer(base_ptr + byte_offset, nbytes), + device=device, + ) + return raw.view(dtype).reshape(sizes) + + +__all__ = [ + "init", + "is_initialized", + "malloc", + "peer_ptr", + "tensor_from_ptr", +] diff --git a/afd_plugin/connectors/gpu/symm_window.py b/afd_plugin/connectors/gpu/symm_window.py new file mode 100644 index 00000000..af920960 --- /dev/null +++ b/afd_plugin/connectors/gpu/symm_window.py @@ -0,0 +1,528 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Symmetric-memory window used by the async GPU AFD connector. + +The window mirrors the CAM shared-window substrate: every rank allocates an +identically sized symmetric buffer, senders write one-sided into the receiver's +buffer, and arrival is announced by a magic-stamped flag word that the receiver +polls. Only the allocation is collective; once it is done a sender never needs +the receiver to participate. + +Layout of one window:: + + [flag[num_regions * ring_depth] | slot(0, 0) | slot(0, 1) | ...] + +``slot(region, ring)`` holds a fixed header followed by the routed/shared +payloads. Dispatch (A -> F) and combine (F -> A) use the same slot layout, so a +single spec sizes both directions. + +Flag words are written *after* the payload on the same stream. Same-stream +device-to-device copies complete in issue order, so a visible flag implies a +complete payload. That holds for NVLink-mapped peer memory; a cross-node +transport would need an explicit fence here. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch + +from afd_plugin.connectors.gpu import nvshmem_rt + +if TYPE_CHECKING: + from torch.distributed.distributed_c10d import ProcessGroup + +# Header words shared by dispatch and combine, followed by expert_counts. +HEADER_MAGIC = 0x41464447 # "AFDG" +HEADER_VERSION = 1 +_H_MAGIC = 0 +_H_VERSION = 1 +_H_SEQ = 2 +_H_SRC_ROLE_RANK = 3 +_H_LAYER_IDX = 4 +_H_STAGE_IDX = 5 +_H_NUM_TOKENS = 6 +_H_ROUTED_TOKENS = 7 +_H_SHARED_TOKENS = 8 +_H_TOPK = 9 +_H_FLAGS = 10 +_H_ECHO_SEQ = 11 # combine only: the dispatch seq being answered +HEADER_FIXED_WORDS = 12 + +FLAG_EMPTY = 0 +FLAG_SHUTDOWN_BIT = 1 << 1 + +# Every field starts on this boundary so a byte offset stays divisible by the +# element size of whichever dtype views it. +_FIELD_ALIGN = 256 + + +def _align(offset: int) -> int: + return (offset + _FIELD_ALIGN - 1) // _FIELD_ALIGN * _FIELD_ALIGN + + +@dataclass(frozen=True, slots=True) +class SlotLayout: + """Byte offsets and element counts for the fields inside one slot.""" + + header_words: int + routed_cap: int + token_cap: int + hidden_size: int + payload_itemsize: int + + header_off: int + route_table_off: int + routed_x_off: int + shared_idx_off: int + shared_x_off: int + slot_bytes: int + + @classmethod + def build( + cls, + *, + expert_per_rank: int, + routed_cap: int, + token_cap: int, + hidden_size: int, + payload_itemsize: int, + ) -> SlotLayout: + header_words = HEADER_FIXED_WORDS + expert_per_rank + header_off = 0 + route_table_off = _align(header_off + header_words * 4) + routed_x_off = _align(route_table_off + 2 * routed_cap * 4) + shared_idx_off = _align( + routed_x_off + routed_cap * hidden_size * payload_itemsize + ) + shared_x_off = _align(shared_idx_off + token_cap * 4) + slot_bytes = _align(shared_x_off + token_cap * hidden_size * payload_itemsize) + return cls( + header_words=header_words, + routed_cap=routed_cap, + token_cap=token_cap, + hidden_size=hidden_size, + payload_itemsize=payload_itemsize, + header_off=header_off, + route_table_off=route_table_off, + routed_x_off=routed_x_off, + shared_idx_off=shared_idx_off, + shared_x_off=shared_x_off, + slot_bytes=slot_bytes, + ) + + +def encode_header( + layout: SlotLayout, + *, + seq: int, + src_role_rank: int, + layer_idx: int, + stage_idx: int, + num_tokens: int, + routed_tokens: int, + shared_tokens: int, + topk: int, + flags: int, + expert_counts: list[int], + echo_seq: int = 0, +) -> torch.Tensor: + """Build the fixed header for one slot as a CPU int32 tensor.""" + expert_per_rank = layout.header_words - HEADER_FIXED_WORDS + if len(expert_counts) != expert_per_rank: + raise ValueError( + f"expert_counts must have {expert_per_rank} entries, " + f"got {len(expert_counts)}", + ) + header = torch.zeros(layout.header_words, dtype=torch.int32) + header[_H_MAGIC] = HEADER_MAGIC + header[_H_VERSION] = HEADER_VERSION + header[_H_SEQ] = seq + header[_H_SRC_ROLE_RANK] = src_role_rank + header[_H_LAYER_IDX] = layer_idx + header[_H_STAGE_IDX] = stage_idx + header[_H_NUM_TOKENS] = num_tokens + header[_H_ROUTED_TOKENS] = routed_tokens + header[_H_SHARED_TOKENS] = shared_tokens + header[_H_TOPK] = topk + header[_H_FLAGS] = flags + header[_H_ECHO_SEQ] = echo_seq + if expert_per_rank: + header[HEADER_FIXED_WORDS:] = torch.tensor(expert_counts, dtype=torch.int32) + return header + + +@dataclass(frozen=True, slots=True) +class SlotHeader: + """Decoded slot header.""" + + seq: int + src_role_rank: int + layer_idx: int + stage_idx: int + num_tokens: int + routed_tokens: int + shared_tokens: int + topk: int + flags: int + echo_seq: int + expert_counts: list[int] + + @property + def is_shutdown(self) -> bool: + return bool(self.flags & FLAG_SHUTDOWN_BIT) + + +def decode_header(header: torch.Tensor) -> SlotHeader: + """Decode a CPU int32 header tensor, validating magic and version.""" + values = header.tolist() + if values[_H_MAGIC] != HEADER_MAGIC: + raise RuntimeError( + f"AFD async GPU header magic mismatch: got {values[_H_MAGIC]:#x}, " + f"expected {HEADER_MAGIC:#x}", + ) + if values[_H_VERSION] != HEADER_VERSION: + raise RuntimeError( + f"AFD async GPU header version {values[_H_VERSION]} is not supported " + f"(expected {HEADER_VERSION})", + ) + return SlotHeader( + seq=values[_H_SEQ], + src_role_rank=values[_H_SRC_ROLE_RANK], + layer_idx=values[_H_LAYER_IDX], + stage_idx=values[_H_STAGE_IDX], + num_tokens=values[_H_NUM_TOKENS], + routed_tokens=values[_H_ROUTED_TOKENS], + shared_tokens=values[_H_SHARED_TOKENS], + topk=values[_H_TOPK], + flags=values[_H_FLAGS], + echo_seq=values[_H_ECHO_SEQ], + expert_counts=values[HEADER_FIXED_WORDS:], + ) + + +@dataclass(slots=True) +class ArrivedSlot: + """One arrival found by ``SymmWindow.poll``.""" + + region: int + ring: int + header: SlotHeader + + +class SymmWindow: + """Symmetric receive window plus the one-sided writes that fill peers'.""" + + def __init__( + self, + *, + num_regions: int, + ring_depth: int, + layout: SlotLayout, + payload_dtype: torch.dtype, + device: torch.device, + group: ProcessGroup, + rank: int, + world_size: int, + ) -> None: + if num_regions <= 0 or ring_depth <= 0: + raise ValueError("num_regions and ring_depth must be positive") + self.num_regions = num_regions + self.ring_depth = ring_depth + self.layout = layout + self.payload_dtype = payload_dtype + self.device = device + self.rank = rank + + self.num_flags = num_regions * ring_depth + self._flag_bytes = _align(self.num_flags * 4) + self.total_bytes = self._flag_bytes + self.num_flags * layout.slot_bytes + + nvshmem_rt.init(group, rank, world_size) + self._base = nvshmem_rt.malloc(self.total_bytes) + # Peer mappings are stable for the life of the allocation, so resolve + # them once instead of per transfer. + self._peer_base = { + pe: (self._base if pe == rank else nvshmem_rt.peer_ptr(self._base, pe)) + for pe in range(world_size) + } + self.local_bytes_view().zero_() + torch.cuda.synchronize() + + # The layout is static, so every window view is built once and then + # sliced. Rebuilding them per transfer meant a __cuda_array_interface__ + # import on every field of every message, which dominated the data path. + self._view_cache: dict[tuple[int, int, int, int], torch.Tensor] = {} + self._flag_cache: dict[tuple[int, int], torch.Tensor] = {} + self._flags_local = nvshmem_rt.tensor_from_ptr( + self._base, + byte_offset=0, + sizes=(self.num_flags,), + dtype=torch.int32, + device=device, + ) + + # Pinned staging keeps header transfers asynchronous; a pageable source + # forces a blocking copy, and there is one header per peer per layer. + # One row per (peer, ring): a single shared row would be overwritten on + # the host by the next peer in the send loop while its own asynchronous + # copy was still in flight, delivering another peer's token counts. + self._header_send = torch.zeros( + (world_size, ring_depth, layout.header_words), + dtype=torch.int32, + ).pin_memory() + self._header_recv = torch.zeros( + layout.header_words, + dtype=torch.int32, + ).pin_memory() + # Host mirror of the local flag array; one D2H per poll refreshes it. + self._flag_host = torch.zeros(self.num_flags, dtype=torch.int32).pin_memory() + self._seen = [FLAG_EMPTY] * self.num_flags + + def local_bytes_view(self) -> torch.Tensor: + return nvshmem_rt.tensor_from_ptr( + self._base, + byte_offset=0, + sizes=(self.total_bytes,), + dtype=torch.uint8, + device=self.device, + ) + + def _slot_byte_off(self, region: int, ring: int) -> int: + return ( + self._flag_bytes + + (region * self.ring_depth + ring) * self.layout.slot_bytes + ) + + def _capacity_view( + self, + peer: int, + region: int, + ring: int, + field_off: int, + ) -> torch.Tensor: + """Return the cached full-capacity view of one slot field.""" + key = (peer, region, ring, field_off) + view = self._view_cache.get(key) + if view is not None: + return view + + layout = self.layout + hidden = layout.hidden_size + if field_off == layout.header_off: + sizes, dtype = (layout.header_words,), torch.int32 + elif field_off == layout.route_table_off: + sizes, dtype = (layout.routed_cap, 2), torch.int32 + elif field_off == layout.routed_x_off: + sizes, dtype = (layout.routed_cap, hidden), self.payload_dtype + elif field_off == layout.shared_idx_off: + sizes, dtype = (layout.token_cap,), torch.int32 + elif field_off == layout.shared_x_off: + sizes, dtype = (layout.token_cap, hidden), self.payload_dtype + else: + raise ValueError(f"unknown slot field offset {field_off}") + + view = nvshmem_rt.tensor_from_ptr( + self._peer_base[peer], + byte_offset=self._slot_byte_off(region, ring) + field_off, + sizes=sizes, + dtype=dtype, + device=self.device, + ) + self._view_cache[key] = view + return view + + def _view( + self, + peer: int, + region: int, + ring: int, + field_off: int, + sizes: tuple[int, ...], + dtype: torch.dtype, + ) -> torch.Tensor: + # Slicing a cached capacity view costs no CUDA calls, unlike importing + # a fresh pointer for every field of every message. + return self._capacity_view(peer, region, ring, field_off)[: sizes[0]] + + def _flag_view(self, peer: int, flag_idx: int) -> torch.Tensor: + key = (peer, flag_idx) + view = self._flag_cache.get(key) + if view is None: + view = nvshmem_rt.tensor_from_ptr( + self._peer_base[peer], + byte_offset=flag_idx * 4, + sizes=(1,), + dtype=torch.int32, + device=self.device, + ) + self._flag_cache[key] = view + return view + + # ------------------------------------------------------------------ + # Send side: every write targets ``peer``'s window, one-sided. + # ------------------------------------------------------------------ + + def write_slot( + self, + *, + peer: int, + region: int, + ring: int, + header: torch.Tensor, + route_table: torch.Tensor | None, + routed_x: torch.Tensor | None, + shared_idx: torch.Tensor | None, + shared_x: torch.Tensor | None, + ) -> None: + """Write one slot into ``peer``'s window, then stamp its flag. + + The flag copy is issued last on the same stream, so a peer that observes + the flag also observes the payload. + """ + layout = self.layout + if routed_x is not None and routed_x.shape[0] > layout.routed_cap: + raise RuntimeError( + f"routed tokens {routed_x.shape[0]} exceed routed_cap " + f"{layout.routed_cap}; raise routed_cap_multiplier", + ) + if shared_x is not None and shared_x.shape[0] > layout.token_cap: + raise RuntimeError( + f"shared tokens {shared_x.shape[0]} exceed token_cap " + f"{layout.token_cap}", + ) + + # Stage through pinned memory so the copy is asynchronous: a pageable + # source would force a blocking transfer, and there is one header per + # peer per layer. + staging = self._header_send[peer][ring] + staging.copy_(header) + self._capacity_view(peer, region, ring, layout.header_off).copy_( + staging, + non_blocking=True, + ) + + if route_table is not None and route_table.numel(): + n = route_table.shape[0] + self._view( + peer, + region, + ring, + layout.route_table_off, + (n, 2), + torch.int32, + ).copy_(route_table, non_blocking=True) + if routed_x is not None and routed_x.numel(): + n = routed_x.shape[0] + self._view( + peer, + region, + ring, + layout.routed_x_off, + (n, layout.hidden_size), + self.payload_dtype, + ).copy_(routed_x, non_blocking=True) + if shared_idx is not None and shared_idx.numel(): + n = shared_idx.shape[0] + self._view( + peer, + region, + ring, + layout.shared_idx_off, + (n,), + torch.int32, + ).copy_(shared_idx, non_blocking=True) + if shared_x is not None and shared_x.numel(): + n = shared_x.shape[0] + self._view( + peer, + region, + ring, + layout.shared_x_off, + (n, layout.hidden_size), + self.payload_dtype, + ).copy_(shared_x, non_blocking=True) + + seq = int(header[_H_SEQ].item()) + flag_idx = region * self.ring_depth + ring + self._flag_view(peer, flag_idx).fill_(seq) + + # ------------------------------------------------------------------ + # Receive side. + # ------------------------------------------------------------------ + + def poll(self) -> ArrivedSlot | None: + """Return the first slot whose flag advanced past what we consumed. + + ponytail: host-side poll, one D2H per call. Correct but it burns a + synchronize per attempt; replace with a device-side ``wait_any`` kernel + spinning on the flag array when the poll shows up in a profile. + """ + self._flag_host.copy_(self._flags_local, non_blocking=False) + host = self._flag_host.tolist() + for idx in range(self.num_flags): + if host[idx] != self._seen[idx]: + self._seen[idx] = host[idx] + region, ring = divmod(idx, self.ring_depth) + return ArrivedSlot( + region=region, + ring=ring, + header=self.read_header(region, ring), + ) + return None + + def read_header(self, region: int, ring: int) -> SlotHeader: + # Reuse the pinned mirror instead of allocating a fresh host tensor on + # every arrival. + self._header_recv.copy_( + self._capacity_view(self.rank, region, ring, self.layout.header_off), + ) + return decode_header(self._header_recv) + + def local_route_table(self, region: int, ring: int, count: int) -> torch.Tensor: + return self._view( + self.rank, + region, + ring, + self.layout.route_table_off, + (count, 2), + torch.int32, + ) + + def local_routed(self, region: int, ring: int, count: int) -> torch.Tensor: + return self._view( + self.rank, + region, + ring, + self.layout.routed_x_off, + (count, self.layout.hidden_size), + self.payload_dtype, + ) + + def local_shared_idx(self, region: int, ring: int, count: int) -> torch.Tensor: + return self._view( + self.rank, + region, + ring, + self.layout.shared_idx_off, + (count,), + torch.int32, + ) + + def local_shared(self, region: int, ring: int, count: int) -> torch.Tensor: + return self._view( + self.rank, + region, + ring, + self.layout.shared_x_off, + (count, self.layout.hidden_size), + self.payload_dtype, + ) + + def close(self) -> None: + # ponytail: the symmetric allocation is left to process teardown. + # nvshmem_free is collective, so freeing here would need both roles to + # shut down in lockstep; add it if windows are ever recreated in-process. + self._peer_base = {} diff --git a/afd_plugin/connectors/npu/async_cam.py b/afd_plugin/connectors/npu/async_cam.py index 00dd128a..57370875 100644 --- a/afd_plugin/connectors/npu/async_cam.py +++ b/afd_plugin/connectors/npu/async_cam.py @@ -45,6 +45,12 @@ coerce_extra_positive_int, coerce_extra_str, ) +from afd_plugin.connectors.async_topology import ( + ASYNC_MOE_REQUEST_SPLIT, + ATTN_RANKS_PER_DP_CONFIG_KEY, + AFDAsyncTopology, + build_async_topology, +) from afd_plugin.connectors.base import ( AFDConnectorBase, ConnectorExtraInfo, @@ -64,9 +70,7 @@ AFD_ASYNC_CAM_GROUP_NAME = "afd_async_cam" CAM_COMM_ID = 0 -ATTN_RANKS_PER_DP_CONFIG_KEY = "attn_ranks_per_dp" ASYNC_MOE_NUM_STAGES = 2 -ASYNC_MOE_REQUEST_SPLIT = "request" ASYNC_MOE_TOKEN_SPLIT = "token" _AFD_ASYNC_EXTRA_CONFIG_FIELDS: Final[frozenset[str]] = frozenset( @@ -190,23 +194,6 @@ class AFDAsyncFFNWorkItem: shared_num_tokens: int -@dataclass(frozen=True, slots=True) -class AFDAsyncTopology: - """Role-local and HCCL-world rank information for one CAM participant.""" - - role: str - role_rank: int - world_rank: int - attn_size: int - ffn_size: int - expert_per_rank: int - - @property - def world_size(self) -> int: - """Return the total number of Attention and FFN ranks.""" - return self.attn_size + self.ffn_size - - class CAMAsyncAFDConnector(AFDConnectorBase): """CAM-backed asynchronous connector for Ascend NPU AFD. @@ -850,56 +837,6 @@ def _log_cam_op_values(op_name: str, label: str, **kwargs: object) -> None: logger.warning("AFD CAM %s %s:\n%s", op_name, label, "\n".join(lines)) -def build_async_topology( - afd_config: AFDConfig, - role_rank: int, - *, - num_routed_experts: int | None = None, -) -> AFDAsyncTopology: - """Validate role-local rank settings and derive the CAM HCCL world rank. - - The world is Attention-first: Attention role rank ``i`` maps to world rank - ``i`` and FFN role rank ``j`` maps to - ``num_attention_ranks + j``. Routed experts are distributed across FFN - ranks using a ceiling division; production model layouts should keep the - routed-expert count divisible by the FFN rank count. - """ - attn_size = afd_config.num_attention_ranks - ffn_size = afd_config.num_ffn_ranks - if attn_size <= 0 or ffn_size <= 0: - raise ValueError("AFD async topology sizes must be positive") - if role_rank < 0: - raise ValueError(f"AFD async role rank must be non-negative, got {role_rank}") - - if afd_config.role == "attention": - if role_rank >= attn_size: - raise ValueError( - "Attention role rank must be within attention size " - f"(rank={role_rank}, size={attn_size})", - ) - world_rank = role_rank - elif afd_config.role == "ffn": - if role_rank >= ffn_size: - raise ValueError( - "FFN role rank must be within FFN size " - f"(rank={role_rank}, size={ffn_size})", - ) - world_rank = attn_size + role_rank - else: - raise ValueError(f"unknown AFD role {afd_config.role!r}") - - expert_count = num_routed_experts or 1 - expert_per_rank = (expert_count + ffn_size - 1) // ffn_size - return AFDAsyncTopology( - role=afd_config.role, - role_rank=role_rank, - world_rank=world_rank, - attn_size=attn_size, - ffn_size=ffn_size, - expert_per_rank=expert_per_rank, - ) - - def _validate_topk_payload( topk_ids: Tensor, topk_weights: Tensor | None, diff --git a/afd_plugin/model_executor/models/deepseek_v2.py b/afd_plugin/model_executor/models/deepseek_v2.py index 90ee2462..5f717f6f 100644 --- a/afd_plugin/model_executor/models/deepseek_v2.py +++ b/afd_plugin/model_executor/models/deepseek_v2.py @@ -21,7 +21,7 @@ from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.models import deepseek_v2 as native -from afd_plugin.config import AFD_ASYNC_CONNECTOR, parse_afd_config +from afd_plugin.config import AFD_ASYNC_CONNECTORS, parse_afd_config from afd_plugin.connectors import ( AFDExpertRoutingSpec, AFDF2ATransferPayload, @@ -512,7 +512,8 @@ def compute_attn_output( topk_weights = None topk_ids = None router_logits = None - # NPU-only: Attention-side gate/topk is implemented in the NPU helper. + # The gate helper delegates expert selection to the connector, so both + # platforms share it despite the module's location. if self.compute_gate_on_attention and self.is_moe_layer: from afd_plugin.model_executor.models.npu import ( deepseek_v2_attention_gate, @@ -568,8 +569,23 @@ def compute_ffn_output( ) return output if self.compute_gate_on_attention: - raise RuntimeError( - "GPU Attention-side gate must call compute_experts_output", + if group_list is None: + # Without a group list the caller is the control-plane path, + # which routes on this side and must use compute_experts_output. + raise RuntimeError( + "GPU Attention-side gate must call compute_experts_output", + ) + # Token-level dispatch: rows arrive pre-routed and grouped by local + # expert, so only the grouped GEMM is left to run here. + from afd_plugin.model_executor.models.gpu import ( + deepseek_v2_attention_gate as gpu_attention_gate, + ) + + return gpu_attention_gate.compute_attention_gate_moe_ffn( + self, + hidden_states=hidden_states, + group_list=group_list, + expand_x_shared=expand_x_shared, ) hidden_states = self.mlp(hidden_states) if ( @@ -703,7 +719,10 @@ def forward( intermediate_tensors: native.IntermediateTensors | None, inputs_embeds: torch.Tensor | None = None, ) -> torch.Tensor | native.IntermediateTensors: - if self.afd_config.connector == AFD_ASYNC_CONNECTOR: + if self.afd_config.connector in AFD_ASYNC_CONNECTORS: + # The schedule below is platform-neutral -- it only drives the model + # and the connector interface -- so both async connectors share it + # despite the module still living under the npu package. from afd_plugin.model_executor.models.npu import ( deepseek_v2_async_cam_forward, ) diff --git a/afd_plugin/model_executor/models/gpu/__init__.py b/afd_plugin/model_executor/models/gpu/__init__.py new file mode 100644 index 00000000..6a15edcf --- /dev/null +++ b/afd_plugin/model_executor/models/gpu/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""GPU-specific AFD model wrappers.""" diff --git a/afd_plugin/model_executor/models/gpu/deepseek_v2_attention_gate.py b/afd_plugin/model_executor/models/gpu/deepseek_v2_attention_gate.py new file mode 100644 index 00000000..75764c81 --- /dev/null +++ b/afd_plugin/model_executor/models/gpu/deepseek_v2_attention_gate.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Attention-side gate helpers for DeepSeek-V2 on CUDA. + +The async GPU connector dispatches tokens that are already routed: one row per +``(token, topk_slot)`` partial, grouped by local expert, with a ``group_list`` +of per-expert counts. The FFN side therefore must not run routing again -- it +only needs the grouped GEMM over its local experts. + +That shape is a ``topk == 1`` problem: give every arriving row its own expert id +and a unit weight, and vLLM's ``fused_experts`` computes exactly the local +expert output. Topk weighting stays on the Attention side, applied during +combine, matching the NPU path. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts + +from afd_plugin.connectors.metadata import AFDF2ATransferPayload + +if TYPE_CHECKING: + from torch import nn + + +def compute_attention_gate_moe_ffn( + layer: nn.Module, + *, + hidden_states: torch.Tensor, + group_list: torch.Tensor, + expand_x_shared: torch.Tensor | None = None, +) -> AFDF2ATransferPayload: + """Run this rank's local experts over pre-routed tokens. + + Args: + layer: The AFD DeepSeek decoder layer owning the MoE module. + hidden_states: ``[num_partials, hidden]`` rows sorted by local expert. + group_list: ``[expert_per_rank]`` per-expert row counts. Must sum to + ``hidden_states.shape[0]``. + expand_x_shared: Optional ``[num_shared, hidden]`` shared-expert rows. + + Returns: + Routed output in the same row order as ``hidden_states``, plus the + shared-expert output when the model has shared experts. + """ + # ``mlp.experts`` is a MoERunner; the weight parameters live on its + # RoutedExperts, and the shared experts hang off the runner rather than + # off the MoE module. + runner = layer.mlp.experts + routed_experts = runner.routed_experts + counts = group_list.to(torch.int64) + num_rows = int(hidden_states.shape[0]) + if int(counts.sum()) != num_rows: + raise ValueError( + f"group_list sums to {int(counts.sum())} but hidden_states has " + f"{num_rows} rows", + ) + + num_local_experts = counts.numel() + if num_rows == 0: + # Routing can leave a peer with nothing -- common in decode, where a + # single token's topk may land entirely on the other FFN rank. + routed_output = hidden_states.new_empty((0, hidden_states.shape[-1])) + shared = runner._shared_experts + return AFDF2ATransferPayload( + routed_output=routed_output, + shared_output=( + shared._layer(expand_x_shared) + if shared is not None + and expand_x_shared is not None + and expand_x_shared.shape[0] > 0 + else None + ), + ) + expert_ids = torch.repeat_interleave( + torch.arange( + num_local_experts, + device=hidden_states.device, + dtype=torch.int32, + ), + counts, + ).unsqueeze(1) + # Unit weights: the real topk weighting happens in the connector's combine. + unit_weights = torch.ones( + (num_rows, 1), + dtype=torch.float32, + device=hidden_states.device, + ) + + routed_output = fused_experts( + hidden_states, + routed_experts.w13_weight, + routed_experts.w2_weight, + unit_weights, + expert_ids, + global_num_experts=num_local_experts, + expert_map=None, + ) + + shared_output = None + shared_experts = runner._shared_experts + if shared_experts is not None and expand_x_shared is not None: + # Call the wrapped MLP rather than SharedExperts.forward: the wrapper is + # a stateful scheduler for the runner's own multi-stream pipeline and + # returns None unless its expected ordering matches. AFD feeds shared + # tokens as their own batch, so that machinery does not apply. + shared_output = shared_experts._layer(expand_x_shared) + + # Mirrors the NPU gate path: scale the routed branch unless fp16, where the + # native model instead scales the shared branch down. + routed_scaling_factor = runner.routed_scaling_factor + if hidden_states.dtype != torch.float16: + routed_output = routed_output * routed_scaling_factor + elif shared_output is not None: + shared_output = shared_output * (1.0 / routed_scaling_factor) + + return AFDF2ATransferPayload( + routed_output=routed_output, + shared_output=shared_output, + ) + + +__all__ = ["compute_attention_gate_moe_ffn"] diff --git a/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py b/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py index 97f8d3d8..5222a446 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py @@ -170,7 +170,9 @@ def run_attention_gate_afd_forward( # start the connector-driven FFN loop for a matching profile request. # Launching CAM collectives here would therefore block on unmatched # synthetic routing metadata; only the local Attention profile is run. - if forward_context.in_profile_run: + # ``in_profile_run`` is set by the Ascend forward context only; the CUDA + # one has no such field, so its absence means "not a profile run". + if getattr(forward_context, "in_profile_run", False): continue dispatch_payload = prepare_cam_dispatch_payload( diff --git a/afd_plugin/v1/worker/attention_model_runner.py b/afd_plugin/v1/worker/attention_model_runner.py index 6a5f055e..bb1aae43 100644 --- a/afd_plugin/v1/worker/attention_model_runner.py +++ b/afd_plugin/v1/worker/attention_model_runner.py @@ -51,6 +51,36 @@ from vllm.v1.core.sched.output import SchedulerOutput +@contextmanager +def _dp_batch_coordination_disabled(disabled: bool): + """Skip vLLM's cross-DP batch agreement for connector-driven runs. + + ``GPUModelRunner._determine_batch_execution_and_padding`` all-reduces the + batch shape across the DP group whenever ``data_parallel_size > 1``. Async + AFD deliberately lets each Attention replica advance on its own, so an idle + replica never joins that collective and a busy one blocks in it forever -- + which is where a 2A2F run hangs before it reaches the first MoE layer. + + Returning the single-rank answer (``num_tokens_across_dp=None``) makes the + upstream function skip its DP-padding branch entirely, exactly as it does + for ``data_parallel_size == 1``. + """ + if not disabled: + yield + return + + original = gpu_model_runner.coordinate_batch_across_dp + + def _single_rank_coordination(*_args: Any, cudagraph_mode: int, **_kwargs: Any): + return False, None, cudagraph_mode + + gpu_model_runner.coordinate_batch_across_dp = _single_rank_coordination + try: + yield + finally: + gpu_model_runner.coordinate_batch_across_dp = original + + class AFDAttentionModelRunner(GPUModelRunner): """Attention model runner that injects AFD metadata into forward context.""" @@ -76,10 +106,6 @@ def __init__( self.afd_config, ) self.connector.init_afd_connector() - # TODO: Async GPU connector will be supported in the future - assert self.connector.control_plane is not None, ( - "GPU model runner only supports control-plane-driven connectors" - ) self._is_warmup = False self._afd_is_graph_capturing = False self._afd_pending_metadata: AFDForwardContextMetadata | None = None @@ -125,9 +151,8 @@ def _send_dp_metadata( dp_metadata: DPMetadata | AFDDPMetadata | None, ubatch_slices: Any, ) -> None: - assert self.connector.control_plane is not None, ( - "_send_dp_metadata needs control plane driven connectors" - ) + if self.connector.control_plane is None: + return if ubatch_slices and len(ubatch_slices) > 1: dp_metadata_list = { @@ -334,25 +359,28 @@ def _determine_batch_execution_and_padding( torch.Tensor | None, CUDAGraphStat | None, ]: - ( - cudagraph_mode, - batch_descriptor, - should_ubatch, - num_tokens_across_dp, - cudagraph_stats, - ) = super()._determine_batch_execution_and_padding( - num_tokens, - num_reqs, - num_scheduled_tokens_np, - max_num_scheduled_tokens, - use_cascade_attn, - allow_microbatching, - force_eager, - force_uniform_decode, - force_has_lora, - force_num_active_loras, - num_encoder_reqs, - ) + with _dp_batch_coordination_disabled( + self.connector.control_plane is None, + ): + ( + cudagraph_mode, + batch_descriptor, + should_ubatch, + num_tokens_across_dp, + cudagraph_stats, + ) = super()._determine_batch_execution_and_padding( + num_tokens, + num_reqs, + num_scheduled_tokens_np, + max_num_scheduled_tokens, + use_cascade_attn, + allow_microbatching, + force_eager, + force_uniform_decode, + force_has_lora, + force_num_active_loras, + num_encoder_reqs, + ) args = ( num_tokens, diff --git a/afd_plugin/v1/worker/ffn_model_runner.py b/afd_plugin/v1/worker/ffn_model_runner.py index b0b248b8..cd83f150 100644 --- a/afd_plugin/v1/worker/ffn_model_runner.py +++ b/afd_plugin/v1/worker/ffn_model_runner.py @@ -30,6 +30,11 @@ AFDControlPayload, AFDDPMetadata, ) +from afd_plugin.connectors.gpu.async_gpu import ( + ConnectorShutdown, + GpuAsyncTransferState, +) +from afd_plugin.connectors.metadata import AFDF2ATransferPayload from afd_plugin.v1.worker.attention_model_runner import ( fail_if_unsupported_ubatching, ) @@ -53,10 +58,12 @@ class GPUFFNModelRunner(LoRAModelRunnerMixin): """FFN model runner for AFD GPU execution. - FFN steps are driven by the connector control plane rather than the vLLM - scheduler. GPU only supports control-plane-driven connectors, so the runner - asserts ``connector.control_plane is not None`` at construction; connectors - without a control plane (``control_plane is None``) are not supported. + FFN steps are driven by the connector rather than the vLLM scheduler, in one + of two ways. Control-plane connectors receive broadcast DP metadata and then + walk every layer in lockstep with the Attention side. Connectors without a + control plane (``control_plane is None``) instead pull one work item at a + time from their receive loop, learning the layer and token counts from the + arriving payload; see ``execute_connector_driven_step``. """ afd_expected_role = "ffn" @@ -80,10 +87,9 @@ def __init__(self, vllm_config: VllmConfig, device: object) -> None: vllm_config, self.afd_config, ) - # TODO: Async GPU connector will be supported in the future - assert self.connector.control_plane is not None, ( - "GPU model runner only supports control-plane-driven connectors" - ) + # A connector without a control plane drives FFN steps from its own + # receive loop instead of from broadcast DP metadata. + self.is_connector_driven = self.connector.control_plane is None self.model: Any | None = None self.model_memory_usage = 0 @@ -237,6 +243,61 @@ def _ffn_forward( self.connector.send_ffn_output(rank_ffn_output, context) return rank_ffn_output + def execute_connector_driven_step(self) -> None: + """Drain whatever the connector has already received, then return. + + Returning on an idle poll rather than blocking forever is what lets the + worker loop observe its shutdown event. The batch size below is only a + drain granularity: successive work items may belong to different layers + of different Attention replicas. + """ + step_afd_gpu_profiler(self.prof) + self._ffn_forward_connector_driven() + + def _ffn_forward_connector_driven( + self, + ) -> torch.Tensor | AFDF2ATransferPayload | None: + stage_idx = 0 + rank_ffn_output = None + connector = self.connector + max_items = max(1, int(self.num_layers)) + + with _ffn_forward_context(self.vllm_config) as forward_context: + for _ in range(max_items): + try: + work_item = connector.recv_ffn_work_item( + stage_idx=stage_idx, + max_num_tokens=self.vllm_config.scheduler_config.max_num_batched_tokens, + ) + except TimeoutError: + # Nothing pending; hand control back so the worker loop can + # check for shutdown. + return rank_ffn_output + except ConnectorShutdown: + raise + + states = work_item.context.states + if not isinstance(states, GpuAsyncTransferState): + raise RuntimeError( + "async GPU FFN work item requires GpuAsyncTransferState", + ) + metadata = work_item.context.metadata + forward_context.dp_metadata = None + forward_context.additional_kwargs["afd_metadata"] = metadata + _set_moe_layer_index(forward_context, work_item.layer_idx) + + rank_ffn_output = self.model.compute_ffn_output( + hidden_states=work_item.hidden_states, + layer_idx=work_item.layer_idx, + group_list=states.group_list, + expand_x_shared=states.expand_x_shared, + ) + rank_ffn_output = connector.send_ffn_work_item_output( + work_item, + rank_ffn_output, + ) + return rank_ffn_output + def _execute_eager_mode( self, hidden_states: torch.Tensor, diff --git a/afd_plugin/v1/worker/ffn_worker.py b/afd_plugin/v1/worker/ffn_worker.py index 5d52179b..6c281b12 100644 --- a/afd_plugin/v1/worker/ffn_worker.py +++ b/afd_plugin/v1/worker/ffn_worker.py @@ -13,6 +13,7 @@ from vllm.v1.worker.gpu_worker import Worker from vllm.v1.worker.worker_base import CompilationTimes +from afd_plugin.connectors.gpu.async_gpu import ConnectorShutdown from afd_plugin.model_executor.models.model_utils import get_afd_model_config from afd_plugin.v1.worker.attention_model_runner import fail_if_unsupported_ubatching from afd_plugin.v1.worker.ffn_model_runner import GPUFFNModelRunner @@ -127,6 +128,13 @@ def ffn_worker_loop() -> None: try: self._run_ffn_server_loop() except Exception as exc: + shutdown_event = self._ffn_shutdown_event + if shutdown_event is not None and shutdown_event.is_set(): + logger.debug( + "AFD FFN receive loop stopped during shutdown", + exc_info=True, + ) + return self._ffn_loop_error = exc logger.exception("AFD FFN worker loop failed") @@ -147,11 +155,17 @@ def _run_ffn_server_loop(self) -> None: while not event.is_set(): if self.model_runner.connector.control_plane is None: - raise NotImplementedError( - "GPU FFN only supports control-plane-driven connectors; " - "connectors without a control plane (control_plane is None) " - "are not supported.", - ) + # Connector-driven: the step returns on an idle poll, so the + # loop gets to re-check the shutdown event. No device-wide + # synchronize here -- it would serialize every receive against + # the previous compute and erase the overlap this path exists + # for; ordering is carried by the connector's own streams. + try: + self.model_runner.execute_connector_driven_step() + except ConnectorShutdown: + logger.info("AFD FFN loop exiting: peer announced shutdown") + return + continue payload = self.model_runner.connector.control_plane.recv_dp_metadata_list() dp_metadata_list = payload.dp_metadata_list diff --git a/afd_plugin/v1/worker/npu/attention_model_runner.py b/afd_plugin/v1/worker/npu/attention_model_runner.py index 645f059b..4cdfb171 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner.py @@ -74,7 +74,7 @@ stop_afd_npu_profiler, ) from afd_plugin.config import ( - AFD_ASYNC_CONNECTOR, + AFD_ASYNC_NPU_CONNECTOR, AFDConfig, parse_afd_config, ) @@ -143,7 +143,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.afd_config, ) self.afd_async_extra_info = AFDAsyncExtraInfo() - if afd_config.connector == AFD_ASYNC_CONNECTOR: + if afd_config.connector == AFD_ASYNC_NPU_CONNECTOR: connector_extra_info = self.connector.extra_info if not isinstance(connector_extra_info, AFDAsyncExtraInfo): raise TypeError( diff --git a/pyproject.toml b/pyproject.toml index 4a82c38a..d89ba163 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -95,6 +95,7 @@ select = [ "ISC", "SIM", ] +extend-ignore = ["N812"] [tool.ruff.lint.per-file-ignores] "afd_plugin/compat/patches/**/*.py" = [ diff --git a/recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/1a1f_eager_async.sh b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/1a1f_eager_async.sh new file mode 100644 index 00000000..b3b74068 --- /dev/null +++ b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/1a1f_eager_async.sh @@ -0,0 +1,128 @@ +#!/usr/bin/env bash +# 1A1F async GPU connector, eager, prefill-only. +# +# Launch under a GPU reservation, which sets CUDA_VISIBLE_DEVICES: +# gpu run --gpu-ids 3,7 -- bash recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/1a1f_eager_async.sh +# +# The two roles are separate vllm serve processes: the AFD process group hosts +# its own TCPStore, which cannot be created under a single torchrun/torchelastic +# launcher. +set -u + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} +LOG_DIR=${LOG_DIR:-.} +mkdir -p "$LOG_DIR" +export VLLM_USE_V2_MODEL_RUNNER=0 +# Single node over NVLink: skip the IB transport probe. +export NVSHMEM_REMOTE_TRANSPORT=${NVSHMEM_REMOTE_TRANSPORT:-none} +# Two servers on one box spawn a lot of threads; the HF tokenizer's rayon pool +# is the first thing to fail when thread creation gets refused. +export TOKENIZERS_PARALLELISM=${TOKENIZERS_PARALLELISM:-false} +export RAYON_NUM_THREADS=${RAYON_NUM_THREADS:-2} +export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4} + +IFS=',' read -r -a DEVICES <<< "${CUDA_VISIBLE_DEVICES:-0,1}" +if [ "${#DEVICES[@]}" -lt 2 ]; then + echo "need 2 visible GPUs, got ${#DEVICES[@]}: ${CUDA_VISIBLE_DEVICES:-unset}" >&2 + exit 1 +fi +ATTN_DEVICES="${DEVICES[0]}" +FFN_DEVICES="${DEVICES[1]}" +echo "attention on ${ATTN_DEVICES}, ffn on ${FFN_DEVICES}" + +# Lower this when sharing a box: vLLM refuses to start if the desired +# fraction exceeds what is actually free. +GPU_MEM_UTIL=${GPU_MEM_UTIL:-0.9} +# Prefill batch size drives whether each MoE call clears the compute-bound +# inflection point, so it is the knob to raise when benchmarking. +MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-512} +MAX_NUM_SEQS=${MAX_NUM_SEQS:-8} +AFD_PORT=${AFD_PORT:-6275} +API_PORT=${API_PORT:-18311} + +CUDA_VISIBLE_DEVICES="$ATTN_DEVICES" uv run vllm serve "$MODEL_PATH" \ + --data-parallel-size 1 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --additional-config '{ + "afd": { + "role": "attention", + "connector": "GpuAsyncAFDConnector", + "async": true, + "compute_gate_on_attention": true, + "host": "127.0.0.1", + "port": '"$AFD_PORT"', + "num_attention_ranks": 1, + "num_ffn_ranks": 1 + } + }' \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \ + --api-server-count 1 \ + --gpu-memory-utilization "$GPU_MEM_UTIL" \ + --enforce-eager \ + --host 127.0.0.1 \ + --port "$API_PORT" \ + --trust-remote-code > "$LOG_DIR/attn.log" 2>&1 & +ATTN_PID=$! + +CUDA_VISIBLE_DEVICES="$FFN_DEVICES" uv run vllm serve "$MODEL_PATH" \ + --data-parallel-size 1 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --additional-config '{ + "afd": { + "role": "ffn", + "connector": "GpuAsyncAFDConnector", + "async": true, + "compute_gate_on_attention": true, + "host": "127.0.0.1", + "port": '"$AFD_PORT"', + "num_attention_ranks": 1, + "num_ffn_ranks": 1 + } + }' \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \ + --api-server-count 1 \ + --gpu-memory-utilization "$GPU_MEM_UTIL" \ + --enforce-eager \ + --host 127.0.0.1 \ + --port "$API_PORT" \ + --trust-remote-code > "$LOG_DIR/ffn.log" 2>&1 & +FFN_PID=$! + +cleanup() { + kill "$ATTN_PID" "$FFN_PID" 2>/dev/null + wait "$ATTN_PID" "$FFN_PID" 2>/dev/null +} +trap cleanup EXIT + +for _ in $(seq 1 "${READY_TIMEOUT:-600}"); do + if curl -sf "http://127.0.0.1:$API_PORT/health" > /dev/null 2>&1; then + echo "server ready on http://127.0.0.1:$API_PORT" + echo + echo "curl -s http://127.0.0.1:$API_PORT/v1/completions \\" + echo " -H 'Content-Type: application/json' \\" + echo " -d '{\"model\":\"$MODEL_PATH\",\"prompt\":\"The capital of France is\",\"max_tokens\":16,\"temperature\":0}'" + echo + if [ -n "${SMOKE:-}" ]; then + curl -s "http://127.0.0.1:$API_PORT/v1/completions" \ + -H 'Content-Type: application/json' \ + -d '{"model":"'"$MODEL_PATH"'","prompt":"The capital of France is", + "max_tokens":16,"temperature":0}' + echo + exit 0 + fi + # Stay up so the servers can take requests; Ctrl-C tears both down. + wait "$ATTN_PID" "$FFN_PID" + exit 0 + fi + if ! kill -0 "$ATTN_PID" 2>/dev/null || ! kill -0 "$FFN_PID" 2>/dev/null; then + echo "a server exited early; see $LOG_DIR/attn.log and $LOG_DIR/ffn.log" >&2 + exit 1 + fi + sleep 1 +done +echo "timed out waiting for the server" >&2 +exit 1 diff --git a/recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/2a2f_eager_async.sh b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/2a2f_eager_async.sh new file mode 100644 index 00000000..01cd83dd --- /dev/null +++ b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/2a2f_eager_async.sh @@ -0,0 +1,129 @@ +#!/usr/bin/env bash +# 2A2F async GPU connector, eager, prefill-only. +# +# Launch under a GPU reservation, which sets CUDA_VISIBLE_DEVICES: +# gpu run --gpus 4 -- bash recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/2a2f_eager_async.sh +# +# The two roles are separate vllm serve processes: the AFD process group hosts +# its own TCPStore, which cannot be created under a single torchrun/torchelastic +# launcher. +set -u + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} +LOG_DIR=${LOG_DIR:-.} +mkdir -p "$LOG_DIR" +export VLLM_USE_V2_MODEL_RUNNER=0 +# Single node over NVLink: skip the IB transport probe. +export NVSHMEM_REMOTE_TRANSPORT=${NVSHMEM_REMOTE_TRANSPORT:-none} +# Two servers on one box spawn a lot of threads; the HF tokenizer's rayon pool +# is the first thing to fail when thread creation gets refused. +export TOKENIZERS_PARALLELISM=${TOKENIZERS_PARALLELISM:-false} +export RAYON_NUM_THREADS=${RAYON_NUM_THREADS:-2} +export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4} + +# Split the reserved devices in half: first two Attention, last two FFN. +IFS=',' read -r -a DEVICES <<< "${CUDA_VISIBLE_DEVICES:-0,1,2,3}" +if [ "${#DEVICES[@]}" -lt 4 ]; then + echo "need 4 visible GPUs, got ${#DEVICES[@]}: ${CUDA_VISIBLE_DEVICES:-unset}" >&2 + exit 1 +fi +ATTN_DEVICES="${DEVICES[0]},${DEVICES[1]}" +FFN_DEVICES="${DEVICES[2]},${DEVICES[3]}" +echo "attention on ${ATTN_DEVICES}, ffn on ${FFN_DEVICES}" + +# Lower this when sharing a box: vLLM refuses to start if the desired +# fraction exceeds what is actually free. +GPU_MEM_UTIL=${GPU_MEM_UTIL:-0.9} +# Prefill batch size drives whether each MoE call clears the compute-bound +# inflection point, so it is the knob to raise when benchmarking. +MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-512} +MAX_NUM_SEQS=${MAX_NUM_SEQS:-8} +AFD_PORT=${AFD_PORT:-6271} +API_PORT=${API_PORT:-18307} + +CUDA_VISIBLE_DEVICES="$ATTN_DEVICES" uv run vllm serve "$MODEL_PATH" \ + --data-parallel-size 2 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --additional-config '{ + "afd": { + "role": "attention", + "connector": "GpuAsyncAFDConnector", + "async": true, + "compute_gate_on_attention": true, + "host": "127.0.0.1", + "port": '"$AFD_PORT"', + "num_attention_ranks": 2, + "num_ffn_ranks": 2 + } + }' \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \ + --api-server-count 1 \ + --gpu-memory-utilization "$GPU_MEM_UTIL" \ + --enforce-eager \ + --host 127.0.0.1 \ + --port "$API_PORT" \ + --trust-remote-code > "$LOG_DIR/attn.log" 2>&1 & +ATTN_PID=$! + +CUDA_VISIBLE_DEVICES="$FFN_DEVICES" uv run vllm serve "$MODEL_PATH" \ + --data-parallel-size 2 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --additional-config '{ + "afd": { + "role": "ffn", + "connector": "GpuAsyncAFDConnector", + "async": true, + "compute_gate_on_attention": true, + "host": "127.0.0.1", + "port": '"$AFD_PORT"', + "num_attention_ranks": 2, + "num_ffn_ranks": 2 + } + }' \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \ + --api-server-count 1 \ + --gpu-memory-utilization "$GPU_MEM_UTIL" \ + --enforce-eager \ + --host 127.0.0.1 \ + --port "$API_PORT" \ + --trust-remote-code > "$LOG_DIR/ffn.log" 2>&1 & +FFN_PID=$! + +cleanup() { + kill "$ATTN_PID" "$FFN_PID" 2>/dev/null + wait "$ATTN_PID" "$FFN_PID" 2>/dev/null +} +trap cleanup EXIT + +for _ in $(seq 1 "${READY_TIMEOUT:-600}"); do + if curl -sf "http://127.0.0.1:$API_PORT/health" > /dev/null 2>&1; then + echo "server ready on http://127.0.0.1:$API_PORT" + echo + echo "curl -s http://127.0.0.1:$API_PORT/v1/completions \\" + echo " -H 'Content-Type: application/json' \\" + echo " -d '{\"model\":\"$MODEL_PATH\",\"prompt\":\"The capital of France is\",\"max_tokens\":16,\"temperature\":0}'" + echo + if [ -n "${SMOKE:-}" ]; then + curl -s "http://127.0.0.1:$API_PORT/v1/completions" \ + -H 'Content-Type: application/json' \ + -d '{"model":"'"$MODEL_PATH"'","prompt":"The capital of France is", + "max_tokens":16,"temperature":0}' + echo + exit 0 + fi + # Stay up so the servers can take requests; Ctrl-C tears both down. + wait "$ATTN_PID" "$FFN_PID" + exit 0 + fi + if ! kill -0 "$ATTN_PID" 2>/dev/null || ! kill -0 "$FFN_PID" 2>/dev/null; then + echo "a server exited early; see $LOG_DIR/attn.log and $LOG_DIR/ffn.log" >&2 + exit 1 + fi + sleep 1 +done +echo "timed out waiting for the server" >&2 +exit 1 diff --git a/tests/e2e/async_gpu_connector_e2e.py b/tests/e2e/async_gpu_connector_e2e.py new file mode 100644 index 00000000..3b9cb2be --- /dev/null +++ b/tests/e2e/async_gpu_connector_e2e.py @@ -0,0 +1,241 @@ +"""End-to-end pass over the async GPU connector's public API, two processes. + +Rank 0 runs the Attention side (``send_attn_output`` / ``recv_ffn_output``), +rank 1 runs the FFN side (``recv_ffn_work_item`` / ``send_ffn_work_item_output``) +with the real grouped-GEMM helper. Everything between the gate and the combined +result is exercised: routing, one-sided dispatch, local expert compute, the +write-back, and the weighted reduction. + +Run with two GPUs:: + + python tests/e2e/async_gpu_connector_e2e.py + +Deliberately *not* launched with torchrun. ``init_afd_process_group`` builds its +own TCPStore on the AFD port, and under torchelastic every rank is forced to +``is_master=False`` (``torch/distributed/rendezvous.py:188``), so no rank hosts +the store and the group never forms. Production launches the two roles as +separate ``vllm serve`` processes, which this mirrors. +""" + +import multiprocessing as mp +import sys +from types import SimpleNamespace + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from afd_plugin.config import AFDConfig +from afd_plugin.connectors.gpu.async_gpu import GpuAsyncAFDConnector +from afd_plugin.connectors.metadata import AFDTransferContext, AFDTransferMetadata +from afd_plugin.model_executor.models.gpu.deepseek_v2_attention_gate import ( + compute_attention_gate_moe_ffn, +) + +NUM_TOKENS = 48 +HIDDEN = 128 +INTERMEDIATE = 256 +TOPK = 4 +NUM_EXPERTS = 8 +NUM_LAYERS = 3 +PORT = 29655 +WORLD_PORT = 29656 +SCALING = 1.7 + + +def build_connector(role: str, local_rank: int) -> GpuAsyncAFDConnector: + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + hf_config=SimpleNamespace( + hidden_size=HIDDEN, + num_experts_per_tok=TOPK, + n_routed_experts=NUM_EXPERTS, + ), + dtype=torch.bfloat16, + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=NUM_TOKENS), + additional_config={ + "afd": { + "role": role, + "connector": "GpuAsyncAFDConnector", + "async": True, + "compute_gate_on_attention": True, + "num_attention_ranks": 1, + "num_ffn_ranks": 1, + "port": PORT, + "connector_extra_config": {"ring_depth": 1}, + }, + }, + ) + afd_config = AFDConfig( + role=role, + connector="GpuAsyncAFDConnector", + async_dp=True, + compute_gate_on_attention=True, + num_attention_ranks=1, + num_ffn_ranks=1, + host="127.0.0.1", + port=PORT, + ) + connector = GpuAsyncAFDConnector( + rank=local_rank, + local_rank=local_rank, + vllm_config=vllm_config, + afd_config=afd_config, + role_rank=0, + ) + connector.init_afd_connector() + return connector + + +def make_weights(device): + generator = torch.Generator(device="cpu").manual_seed(11) + w13 = ( + torch.randn(NUM_EXPERTS, 2 * INTERMEDIATE, HIDDEN, generator=generator) + / HIDDEN**0.5 + ).to(device, torch.bfloat16) + w2 = ( + torch.randn(NUM_EXPERTS, HIDDEN, INTERMEDIATE, generator=generator) + / INTERMEDIATE**0.5 + ).to(device, torch.bfloat16) + return w13, w2 + + +def make_layer_inputs(layer_idx, device): + gen = torch.Generator(device="cpu").manual_seed(100 + layer_idx) + x = torch.randn(NUM_TOKENS, HIDDEN, generator=gen).to(device, torch.bfloat16) + topk_ids = torch.stack( + [torch.randperm(NUM_EXPERTS, generator=gen)[:TOPK] for _ in range(NUM_TOKENS)], + ).to(device, torch.int32) + topk_weights = torch.rand(NUM_TOKENS, TOPK, generator=gen).to(device) + return x, topk_ids, topk_weights + + +def reference_moe(x, w13, w2, topk_ids, topk_weights): + out = torch.zeros(x.shape[0], HIDDEN, dtype=torch.float32, device=x.device) + for token in range(x.shape[0]): + for slot in range(topk_ids.shape[1]): + expert = int(topk_ids[token, slot]) + hidden = x[token].to(torch.float32) @ w13[expert].to(torch.float32).T + gate, up = hidden.chunk(2, dim=-1) + y = (F.silu(gate) * up) @ w2[expert].to(torch.float32).T + out[token] += float(topk_weights[token, slot]) * y * SCALING + return out + + +def init_world(rank: int) -> None: + """Mimic a single `vllm serve`: a private default group of size 1. + + The connector bootstraps NVSHMEM on the AFD group itself, so the default + group deliberately does *not* span both roles -- that is the topology the + real deployment has. + """ + dist.init_process_group( + "nccl", + init_method=f"tcp://127.0.0.1:{WORLD_PORT + rank}", + world_size=1, + rank=0, + ) + + +def run_attention(rank: int) -> None: + init_world(0) + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + connector = build_connector("attention", rank) + w13, w2 = make_weights(device) + + for layer_idx in range(NUM_LAYERS): + x, topk_ids, topk_weights = make_layer_inputs(layer_idx, device) + context = AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=layer_idx, + stage_idx=0, + seq_len=NUM_TOKENS, + ), + ) + connector.send_attn_output( + x, + context, + topk_ids=topk_ids, + topk_weights=topk_weights, + ) + got = connector.recv_ffn_output(ref_tensor=x, ubatch_idx=0) + expected = reference_moe(x, w13, w2, topk_ids, topk_weights) + torch.testing.assert_close( + got.to(torch.float32), + expected, + rtol=8e-2, + atol=8e-2, + ) + print( + f"[A] layer {layer_idx}: combined output matches reference MoE", flush=True + ) + + print("PASS: async GPU connector end-to-end", flush=True) + connector.close() + + +def run_ffn(rank: int) -> None: + init_world(1) + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + connector = build_connector("ffn", rank) + w13, w2 = make_weights(device) + layer = SimpleNamespace( + mlp=SimpleNamespace( + experts=SimpleNamespace( + routed_experts=SimpleNamespace(w13_weight=w13, w2_weight=w2), + _shared_experts=None, + routed_scaling_factor=SCALING, + ), + ), + ) + + for _ in range(NUM_LAYERS): + while True: + try: + work_item = connector.recv_ffn_work_item( + stage_idx=0, + max_num_tokens=NUM_TOKENS, + ) + break + except TimeoutError: + continue + states = work_item.context.states + payload = compute_attention_gate_moe_ffn( + layer, + hidden_states=work_item.hidden_states, + group_list=states.group_list, + expand_x_shared=None, + ) + connector.send_ffn_work_item_output(work_item, payload) + print( + f"[F] layer {work_item.layer_idx}: served " + f"{work_item.num_tokens} routed tokens", + flush=True, + ) + + connector.close() + + +def main() -> None: + if torch.cuda.device_count() < 2: + raise SystemExit("this test needs two visible GPUs") + mp.set_start_method("spawn", force=True) + procs = [ + mp.Process(target=run_ffn, args=(1,)), + mp.Process(target=run_attention, args=(0,)), + ] + for proc in procs: + proc.start() + failed = False + for proc in procs: + proc.join(timeout=300) + if proc.exitcode != 0: + failed = True + sys.exit(1 if failed else 0) + + +if __name__ == "__main__": + main() diff --git a/tests/unit/config/test_config.py b/tests/unit/config/test_config.py index 7728bd8a..fa77ba54 100644 --- a/tests/unit/config/test_config.py +++ b/tests/unit/config/test_config.py @@ -119,7 +119,7 @@ def test_parse_async_dp_config_from_async_alias(): def test_async_dp_requires_async_connector(): - with pytest.raises(ValueError, match="requires connector='CAMAsyncAFDConnector'"): + with pytest.raises(ValueError, match="AFD async mode requires one of"): parse_afd_config( { "afd": { @@ -131,6 +131,24 @@ def test_async_dp_requires_async_connector(): ) +@pytest.mark.parametrize( + "connector", + ["CAMAsyncAFDConnector", "GpuAsyncAFDConnector"], +) +def test_async_dp_accepts_every_async_connector(connector): + config = parse_afd_config( + { + "afd": { + "connector": connector, + "role": "attention", + "async": True, + }, + }, + ) + assert config.connector == connector + assert config.async_dp + + def test_original_common_afd_field_aliases_are_supported(): raw = { "afd_role": "ffn", diff --git a/tests/unit/connectors/test_async_gpu_connector.py b/tests/unit/connectors/test_async_gpu_connector.py new file mode 100644 index 00000000..bf320309 --- /dev/null +++ b/tests/unit/connectors/test_async_gpu_connector.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Unit tests for the async GPU connector's wire format and routing math.""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("torch") + +import torch # noqa: E402 + +from afd_plugin.connectors.factory import AFDConnectorFactory # noqa: E402 +from afd_plugin.connectors.gpu.async_gpu import ( # noqa: E402 + GpuAsyncAFDConnector, + GpuAsyncExtraInfo, + combine_scatter, + plan_dispatch, +) +from afd_plugin.connectors.gpu.symm_window import ( # noqa: E402 + FLAG_SHUTDOWN_BIT, + HEADER_FIXED_WORDS, + SlotLayout, + decode_header, + encode_header, +) + + +@pytest.fixture +def layout() -> SlotLayout: + return SlotLayout.build( + expert_per_rank=4, + routed_cap=100, + token_cap=32, + hidden_size=8, + payload_itemsize=2, + ) + + +# ---------------------------------------------------------------------- +# Config +# ---------------------------------------------------------------------- + + +def test_connector_is_registered_and_has_no_control_plane(): + connector_cls = AFDConnectorFactory.get_connector_class("GpuAsyncAFDConnector") + assert connector_cls is GpuAsyncAFDConnector + assert connector_cls.control_plane is None + + +def test_ring_depth_defaults_to_the_number_of_live_stages(): + assert GpuAsyncExtraInfo.from_mapping(None).ring_depth == 1 + assert GpuAsyncExtraInfo.from_mapping({"async_moe_ubatching": True}).ring_depth == 2 + assert GpuAsyncExtraInfo.from_mapping({"ring_depth": 4}).ring_depth == 4 + + +def test_unknown_extra_config_field_is_rejected(): + with pytest.raises(ValueError, match="unknown AFD async GPU"): + GpuAsyncExtraInfo.from_mapping({"nope": 1}) + + +def test_routed_cap_multiplier_must_be_positive(): + with pytest.raises(ValueError, match="routed_cap_multiplier"): + GpuAsyncExtraInfo.from_mapping({"routed_cap_multiplier": 0}) + + +# ---------------------------------------------------------------------- +# Slot layout +# ---------------------------------------------------------------------- + + +def test_slot_fields_are_disjoint_and_fit_inside_the_slot(layout: SlotLayout): + assert layout.header_words == HEADER_FIXED_WORDS + 4 + assert layout.route_table_off >= layout.header_off + layout.header_words * 4 + assert layout.routed_x_off >= layout.route_table_off + 2 * 100 * 4 + assert layout.shared_idx_off >= layout.routed_x_off + 100 * 8 * 2 + assert layout.shared_x_off >= layout.shared_idx_off + 32 * 4 + assert layout.slot_bytes >= layout.shared_x_off + 32 * 8 * 2 + + +def test_every_field_offset_is_viewable_as_int32_and_payload(layout: SlotLayout): + # get_buffer takes an element offset, so a byte offset that is not a + # multiple of the element size would silently land on the wrong address. + for offset in ( + layout.header_off, + layout.route_table_off, + layout.routed_x_off, + layout.shared_idx_off, + layout.shared_x_off, + ): + assert offset % 4 == 0 + assert offset % layout.payload_itemsize == 0 + + +# ---------------------------------------------------------------------- +# Header codec +# ---------------------------------------------------------------------- + + +def test_header_round_trip(layout: SlotLayout): + header = encode_header( + layout, + seq=7, + src_role_rank=3, + layer_idx=11, + stage_idx=1, + num_tokens=32, + routed_tokens=90, + shared_tokens=16, + topk=6, + flags=0, + expert_counts=[10, 20, 30, 30], + ) + decoded = decode_header(header) + assert decoded.seq == 7 + assert decoded.src_role_rank == 3 + assert decoded.layer_idx == 11 + assert decoded.stage_idx == 1 + assert decoded.num_tokens == 32 + assert decoded.routed_tokens == 90 + assert decoded.shared_tokens == 16 + assert decoded.topk == 6 + assert decoded.expert_counts == [10, 20, 30, 30] + assert sum(decoded.expert_counts) == decoded.routed_tokens + assert not decoded.is_shutdown + + +def test_shutdown_flag_survives_the_round_trip(layout: SlotLayout): + header = encode_header( + layout, + seq=8, + src_role_rank=0, + layer_idx=0, + stage_idx=0, + num_tokens=0, + routed_tokens=0, + shared_tokens=0, + topk=6, + flags=FLAG_SHUTDOWN_BIT, + expert_counts=[0, 0, 0, 0], + ) + assert decode_header(header).is_shutdown + + +def test_corrupt_magic_is_rejected(layout: SlotLayout): + header = encode_header( + layout, + seq=1, + src_role_rank=0, + layer_idx=0, + stage_idx=0, + num_tokens=1, + routed_tokens=0, + shared_tokens=0, + topk=6, + flags=0, + expert_counts=[0, 0, 0, 0], + ) + header[0] = 0 + with pytest.raises(RuntimeError, match="magic mismatch"): + decode_header(header) + + +def test_expert_counts_length_must_match_the_layout(layout: SlotLayout): + with pytest.raises(ValueError, match="expert_counts"): + encode_header( + layout, + seq=1, + src_role_rank=0, + layer_idx=0, + stage_idx=0, + num_tokens=1, + routed_tokens=0, + shared_tokens=0, + topk=6, + flags=0, + expert_counts=[1, 2], + ) + + +# ---------------------------------------------------------------------- +# Routing +# ---------------------------------------------------------------------- + +_NUM_TOKENS = 7 +_TOPK = 3 +_FFN_SIZE = 2 +_EXPERT_PER_RANK = 4 +_HIDDEN = 5 + + +@pytest.fixture +def routing_inputs(): + generator = torch.Generator().manual_seed(0) + num_experts = _FFN_SIZE * _EXPERT_PER_RANK + topk_ids = torch.stack( + [ + torch.randperm(num_experts, generator=generator)[:_TOPK] + for _ in range(_NUM_TOKENS) + ], + ).to(torch.int32) + hidden_states = torch.randn(_NUM_TOKENS, _HIDDEN, generator=generator) + topk_weights = torch.rand(_NUM_TOKENS, _TOPK, generator=generator) + return topk_ids, hidden_states, topk_weights + + +def test_every_partial_is_routed_exactly_once(routing_inputs): + topk_ids, _, _ = routing_inputs + route_table, counts, _ = plan_dispatch( + topk_ids, + ffn_size=_FFN_SIZE, + expert_per_rank=_EXPERT_PER_RANK, + ) + assert route_table.shape == (_NUM_TOKENS * _TOPK, 2) + assert int(counts.sum()) == _NUM_TOKENS * _TOPK + seen = {(token_idx, slot) for token_idx, slot in route_table.tolist()} + assert len(seen) == _NUM_TOKENS * _TOPK + + +def test_each_destination_segment_is_grouped_by_local_expert(routing_inputs): + topk_ids, _, _ = routing_inputs + route_table, counts, offsets = plan_dispatch( + topk_ids, + ffn_size=_FFN_SIZE, + expert_per_rank=_EXPERT_PER_RANK, + ) + counts_host, offsets_host = counts.tolist(), offsets.tolist() + for ffn_rank in range(_FFN_SIZE): + base = ffn_rank * _EXPERT_PER_RANK + cursor = offsets_host[base] + for local_expert in range(_EXPERT_PER_RANK): + for _ in range(counts_host[base + local_expert]): + token_idx, slot = route_table[cursor].tolist() + assert int(topk_ids[token_idx, slot]) == base + local_expert + cursor += 1 + assert cursor == offsets_host[base] + sum( + counts_host[base : base + _EXPERT_PER_RANK], + ) + + +def test_identity_experts_recombine_to_the_weighted_sum(routing_inputs): + topk_ids, hidden_states, topk_weights = routing_inputs + route_table, counts, offsets = plan_dispatch( + topk_ids, + ffn_size=_FFN_SIZE, + expert_per_rank=_EXPERT_PER_RANK, + ) + counts_host, offsets_host = counts.tolist(), offsets.tolist() + accumulator = torch.zeros(_NUM_TOKENS, _HIDDEN, dtype=torch.float32) + for ffn_rank in range(_FFN_SIZE): + base = ffn_rank * _EXPERT_PER_RANK + start = offsets_host[base] + total = sum(counts_host[base : base + _EXPERT_PER_RANK]) + segment = route_table[start : start + total] + combine_scatter( + accumulator, + routed_out=hidden_states.index_select(0, segment[:, 0].to(torch.int64)), + route_table=segment, + topk_weights=topk_weights, + ) + expected = hidden_states * topk_weights.sum(dim=1, keepdim=True) + torch.testing.assert_close(accumulator, expected.to(torch.float32)) + + +def test_routing_handles_experts_not_divisible_by_ffn_size(): + # expert_per_rank is a ceiling division, so the padded tail must stay empty + # instead of silently absorbing real partials. + ffn_size, expert_per_rank, num_experts = 3, 2, 5 + topk_ids = torch.tensor([[0, 4], [1, 3], [2, 4]], dtype=torch.int32) + _, counts, _ = plan_dispatch( + topk_ids, + ffn_size=ffn_size, + expert_per_rank=expert_per_rank, + ) + assert counts.numel() == ffn_size * expert_per_rank + assert int(counts.sum()) == topk_ids.numel() + assert int(counts[num_experts:].sum()) == 0 + + +def test_routing_can_leave_one_destination_empty(): + """A single decode token's topk can land entirely on one FFN rank. + + The peer that gets nothing must still see a well-formed, empty segment -- + zero-length windows are what crashed a 2A2F decode step. + """ + ffn_size, expert_per_rank = 2, 4 + # Every partial targets experts owned by FFN rank 0. + topk_ids = torch.tensor([[0, 1, 2]], dtype=torch.int32) + route_table, counts, offsets = plan_dispatch( + topk_ids, + ffn_size=ffn_size, + expert_per_rank=expert_per_rank, + ) + counts_host, offsets_host = counts.tolist(), offsets.tolist() + + base_zero = 0 + assert sum(counts_host[base_zero : base_zero + expert_per_rank]) == 3 + base_one = expert_per_rank + empty_total = sum(counts_host[base_one : base_one + expert_per_rank]) + assert empty_total == 0 + empty_segment = route_table[ + offsets_host[base_one] : offsets_host[base_one] + empty_total + ] + assert empty_segment.shape == (0, 2) + + # Combining an empty segment must be a no-op, not an error. + accumulator = torch.zeros(1, 4, dtype=torch.float32) + combine_scatter( + accumulator, + routed_out=torch.zeros(0, 4), + route_table=empty_segment, + topk_weights=torch.ones(1, 3), + ) + assert torch.count_nonzero(accumulator) == 0 diff --git a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py index b650f480..6c6a2cd1 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py @@ -8,7 +8,7 @@ pytest.importorskip("vllm") nn = torch.nn -from afd_plugin.config import AFD_ASYNC_CONNECTOR, AFDConfig # noqa: E402 +from afd_plugin.config import AFD_ASYNC_NPU_CONNECTOR, AFDConfig # noqa: E402 from afd_plugin.model_executor.models import deepseek_v2 as adapter # noqa: E402 @@ -247,7 +247,7 @@ def async_forward(*args): nn.Module.__init__(model) model.afd_config = AFDConfig( role="attention", - connector=AFD_ASYNC_CONNECTOR, + connector=AFD_ASYNC_NPU_CONNECTOR, ) positions = torch.arange(1) diff --git a/tests/unit/model_executor/models/test_forward_context.py b/tests/unit/model_executor/models/test_forward_context.py index 113066a7..d9ec7181 100644 --- a/tests/unit/model_executor/models/test_forward_context.py +++ b/tests/unit/model_executor/models/test_forward_context.py @@ -346,14 +346,19 @@ def test_deepseek_compute_gate_on_attention_selects_backend_boundary(): assert "self.mlp = AFDDeepseekV2RemoteExpertsMoE(" in source assert "self.mlp = GateOnlyRemoteMoE(" in source assert 'prefix=f"{prefix}.mlp"' in source + # The gate/topk helper delegates expert selection to the connector, so both + # platforms share it; only the FFN-side MoE compute stays platform-split. assert ( - "# NPU-only: Attention-side gate/topk is implemented in the NPU helper." + "# The gate helper delegates expert selection to the connector, so both" in source ) assert ( "# NPU-only: gated MoE FFN compute consumes Attention-side topk payloads." in source ) + # CUDA reaches its own grouped-GEMM entry point only once tokens arrive + # pre-routed; without a group list the control-plane path still applies. + assert "gpu_attention_gate.compute_attention_gate_moe_ffn(" in source def test_async_moe_pipeline_preserves_stage_order(monkeypatch): diff --git a/tests/unit/v1/worker/test_attention_model_runner.py b/tests/unit/v1/worker/test_attention_model_runner.py index ee915cd3..5deec4f4 100644 --- a/tests/unit/v1/worker/test_attention_model_runner.py +++ b/tests/unit/v1/worker/test_attention_model_runner.py @@ -352,13 +352,21 @@ def test_phase5_allows_two_way_ubatching_but_rejects_other_counts(): ) -def _ubatch_runner(uniform_decode, **parallel_overrides): +_DUMMY_CONTROL_PLANE = object() + + +def _ubatch_runner( + uniform_decode, *, control_plane=_DUMMY_CONTROL_PLANE, **parallel_overrides +): runner = object.__new__(AFDAttentionModelRunner) runner.vllm_config = SimpleNamespace( parallel_config=_parallel_config(**parallel_overrides), ) runner.uniform_decode_query_len = 1 runner._is_uniform_decode = lambda **_kwargs: uniform_decode + # The override consults the connector to decide whether cross-DP batch + # coordination applies; a non-None control plane keeps upstream behaviour. + runner.connector = SimpleNamespace(control_plane=control_plane) return runner @@ -988,3 +996,35 @@ def test_afd_rank_raises_for_out_of_range_dp2_tp2(monkeypatch): with pytest.raises(ValueError, match="out of range"): resolve_role_rank(vllm_config, config) + + +def test_connector_driven_runs_skip_cross_dp_batch_coordination(): + """An idle Attention replica never joins the DP all-reduce. + + Async AFD lets each replica advance alone, so a busy replica must not block + in ``coordinate_batch_across_dp`` waiting for one that never steps. + """ + import vllm.v1.worker.gpu_model_runner as gpu_model_runner + + from afd_plugin.v1.worker.attention_model_runner import ( + _dp_batch_coordination_disabled, + ) + + original = gpu_model_runner.coordinate_batch_across_dp + + with _dp_batch_coordination_disabled(True): + assert gpu_model_runner.coordinate_batch_across_dp is not original + result = gpu_model_runner.coordinate_batch_across_dp( + num_tokens_unpadded=8, + parallel_config=None, + allow_microbatching=False, + num_tokens_padded=8, + uniform_decode=True, + cudagraph_mode=0, + ) + # num_tokens_across_dp None makes upstream skip its DP-padding branch. + assert result == (False, None, 0) + assert gpu_model_runner.coordinate_batch_across_dp is original + + with _dp_batch_coordination_disabled(False): + assert gpu_model_runner.coordinate_batch_across_dp is original diff --git a/tests/unit/v1/worker/test_ffn_model_runner.py b/tests/unit/v1/worker/test_ffn_model_runner.py index bf018e01..0743547f 100644 --- a/tests/unit/v1/worker/test_ffn_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_model_runner.py @@ -20,6 +20,7 @@ AFDTransferContext, AFDTransferMetadata, ) +from afd_plugin.connectors.gpu.async_gpu import ConnectorShutdown # noqa: E402 from afd_plugin.v1.worker.cuda_graph import make_ffn_graph_key # noqa: E402 from afd_plugin.v1.worker.ffn_model_runner import ( # noqa: E402 GPUFFNModelRunner, @@ -654,18 +655,45 @@ def test_ffn_worker_reports_zero_compilation_times(): assert compilation_times.encoder == 0.0 -def test_ffn_worker_loop_rejects_connector_without_control_plane(): +def test_ffn_worker_loop_drives_connector_without_control_plane(): worker = object.__new__(AFDFFNWorker) event = threading.Event() + steps = [] + + def execute_connector_driven_step(): + steps.append(1) + # The connector-driven step returns on an idle poll; the loop must come + # back to the shutdown event rather than block forever. + if len(steps) == 3: + event.set() worker._ffn_shutdown_event = event worker.device = SimpleNamespace(type="cpu") worker.model_runner = SimpleNamespace( connector=_ConnectorDrivenFakeConnector(), + execute_connector_driven_step=execute_connector_driven_step, + ) + + worker._run_ffn_server_loop() + + assert len(steps) == 3 + + +def test_ffn_worker_loop_exits_cleanly_when_peer_announces_shutdown(): + worker = object.__new__(AFDFFNWorker) + + def execute_connector_driven_step(): + raise ConnectorShutdown("peer left") + + worker._ffn_shutdown_event = threading.Event() + worker.device = SimpleNamespace(type="cpu") + worker.model_runner = SimpleNamespace( + connector=_ConnectorDrivenFakeConnector(), + execute_connector_driven_step=execute_connector_driven_step, ) - with pytest.raises(NotImplementedError, match="control-plane-driven"): - worker._run_ffn_server_loop() + # A peer shutdown is an ordinary exit, not a loop failure. + worker._run_ffn_server_loop() def test_ffn_worker_loop_logs_unexpected_thread_errors(caplog): From 906ee5831e528a29182044289676e52fcfa13dd6 Mon Sep 17 00:00:00 2001 From: specture724 Date: Thu, 13 Aug 2026 17:18:29 +0800 Subject: [PATCH 2/2] fix: flash_comm_v1_enabled for NPU only Signed-off-by: specture724 --- .../models/npu/deepseek_v2_async_cam_forward.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py b/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py index 5222a446..a195ee54 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py @@ -180,7 +180,14 @@ def run_attention_gate_afd_forward( topk_weights, topk_ids, router_logits, - use_sequence_parallel=forward_context.flash_comm_v1_enabled, + # Ascend-only field: FlashComm1 leaves each Attention TP rank with a + # disjoint token shard. The CUDA forward context has no such field, + # so its absence means the token dimension is still replicated. + use_sequence_parallel=getattr( + forward_context, + "flash_comm_v1_enabled", + False, + ), ) metadata = AFDTransferMetadata.create_attention_metadata( layer_idx=layer.layer_idx, @@ -229,7 +236,11 @@ def run_async_moe_ubatch_afd_forward( """Run the two-stage async MoE ubatch pipeline used by async CAM.""" forward_context = get_forward_context() - runtime_sequence_parallel = bool(forward_context.flash_comm_v1_enabled) + runtime_sequence_parallel = getattr( + forward_context, + "flash_comm_v1_enabled", + False, + ) if runtime_sequence_parallel != async_moe_ubatch_metadata.use_sequence_parallel: raise RuntimeError( "Async CAM stage layout does not match the current FlashComm1 "