diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index a8278e93..2fec965a 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, @@ -38,7 +38,7 @@ def fail_if_unsupported_npu_afd_features( if is_dsv4: _fail_if_unsupported_dsv4_connector(afd_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, @@ -130,7 +130,7 @@ def _fail_if_unsupported_dsv4_async_features( def _fail_if_unsupported_dsv4_connector(afd_config: AFDConfig) -> None: - if afd_config.connector != AFD_ASYNC_CONNECTOR: + if afd_config.connector != AFD_ASYNC_NPU_CONNECTOR: raise RuntimeError("DSV4 NPU AFD supports only CAMAsyncAFDConnector") diff --git a/afd_plugin/config.py b/afd_plugin/config.py index 256590ab..c16c0e80 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,10 +315,25 @@ 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 == AFD_ASYNC_GPU_CONNECTOR: + # Both of these are structural, not preferences. FFN steps come off the + # connector receive loop, which only the async-DP engine patches drive, + # and the wire carries topk chosen on the Attention side, so there is no + # FFN-side router to fall back on. Reject the combinations here rather + # than as a startup hang or a missing-gate crash mid-run. + if not config.async_dp: + raise ValueError( + f"{AFD_ASYNC_GPU_CONNECTOR} requires async=true", + ) + if not config.compute_gate_on_attention: + raise ValueError( + f"{AFD_ASYNC_GPU_CONNECTOR} requires compute_gate_on_attention=true", + ) if config.connector == "P2pNcclAFDConnector": from afd_plugin.distributed import validate_p2p_topology @@ -331,7 +354,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/base.py b/afd_plugin/connectors/base.py index 95201b83..fcabb1ab 100644 --- a/afd_plugin/connectors/base.py +++ b/afd_plugin/connectors/base.py @@ -60,6 +60,13 @@ class documents the common runtime contract shared by those implementations. control_plane: AFDControlPlane | None = None attn_size: int = 0 ffn_size: int = 0 + # Whether the MoE round trip goes through the opaque dispatch/recv ops + # rather than direct connector calls. The ops exist so Dynamo splits the + # graph at the round trip instead of tracing into the connector; a + # connector that needs that split sets this True. Connectors whose protocol + # carries data the ops do not -- router logits, FlashComm1 token sharding -- + # keep the direct calls and leave it False. + uses_opaque_moe_ops: bool = False @classmethod @abstractmethod 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..3444bbd3 --- /dev/null +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -0,0 +1,1341 @@ +# 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. + +The Attention-side data path is CUDA-graph capturable, under the same +``FULL_DECODE_ONLY`` policy the rest of AFD accepts. Two things make it so, and +both are visible in the flag protocol: the dispatch sequence number lives in a +device tensor the graph advances, so a replay stamps a fresh number even though +it runs no Python; and a reply signals with the constant ``FLAG_REPLY_READY``, +which the Attention rank resets inside the graph once it has consumed the slot, +because a captured stream wait can only ever compare against the value that was +live when it was recorded. The FFN side is driven by a host poll loop and stays +eager, so nothing there needs capturing. + +See the async-GPU rows in ``docs/design/module/connector_contracts.md``. +Supported deployment requires ``async=true``, +``compute_gate_on_attention=true``, and a single node; the first two are +enforced in ``validate_afd_config``. +""" + +from __future__ import annotations + +import os +from collections.abc import Iterable, 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, + H_ROUTED_TOKENS, + H_SEGMENT_START, + HEADER_FIXED_WORDS, + SlotLayout, + SymmWindow, + encode_header, + fill_header_prefix, +) +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" + +# What an FFN reply stamps on the flag the Attention rank is waiting for, and +# what that rank puts back once it has consumed the slot. +# +# It is a constant because the wait has to survive CUDA graph capture. The +# awaited value becomes an immediate inside the captured graph, so a reply that +# signalled with a rising sequence number would be waited for with whatever +# number happened to be live at capture time -- and on the second replay the +# flag still holds it, so the wait falls straight through and the combine reads +# the previous layer's data. A constant marker plus an in-graph reset makes +# every replay identical, which is exactly what a graph needs. +FLAG_REPLY_READY: Final[int] = 1 + +_GPU_ASYNC_EXTRA_CONFIG_FIELDS: Final[frozenset[str]] = frozenset( + { + "attn_ranks_per_dp", + "ring_depth", + "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__}") + +# Bring-up instrument: setting AFD_ASYNC_DEBUG=1 logs magnitude checks for the +# first MoE layers on both roles, pinning where a zero output first appears on +# the async data path (received dispatch, FFN reply, Attention combine). +AFD_ASYNC_DEBUG_ENV: Final[str] = "AFD_ASYNC_DEBUG" +AFD_ASYNC_DEBUG_LAYERS: Final[int] = 3 + + +def _debug_norm(name: str, tensor: Tensor | None) -> str: + if tensor is None or tensor.numel() == 0: + return f"{name}=empty" + return f"{name}={float(tensor.float().abs().mean()):.6e}" + + +@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. + The default here counts only this connector's own MoE stages; the + connector raises it for vLLM's ubatching, which it cannot see from + the extra config alone. + ring_depth_pinned: Whether the extra config named a ring depth. An + explicit one is honoured as given; an implied one is free to grow. + 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 + ring_depth_pinned: bool = False + 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, + ring_depth_pinned="ring_depth" in raw, + 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, + "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. + + ``group_list`` is a device view of the arrived header's trailing words, + which is what the grouped GEMM reads; the host-side copy of those counts + went away with the combine header that used to carry them back. + + ``expand_idx`` and ``weights`` describe this rank's own run of partials: + which token each one reads on the way in, and what to weight it by when the + expert output is reduced back to one row per token on the way out. + """ + + region: int = 0 + ring: int = 0 + src_role_rank: int = 0 + layer_idx: int = 0 + stage_idx: int = 0 + num_tokens: int = 0 + routed_tokens: int = 0 + shared_tokens: int = 0 + group_list: Tensor | None = None + # True when recv_attn_output gathered straight into the caller's buffer, so + # the rows are already staged and must not be copied again. + staged_routed: bool = False + expand_idx: Tensor | None = None + weights: 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. + + Every reply carries a full batch of rows, already weighted and summed over + that rank's experts, so combine is a plain add at matching row positions and + needs no index from the dispatch. + + ``shared_slices[r]`` is the contiguous token range whose shared-expert + output that rank returns. It is recorded here because combine must know the + shape of a reply *before* it arrives: that is what lets the wait happen on a + stream instead of on the host, which cannot then be told what turned up. + """ + + context: AFDTransferContext + shared_slices: list[slice] + num_tokens: int + ring: int + expected_ffn: list[int] + + +@dataclass(frozen=True, slots=True) +class DispatchPlan: + """One layer's routing plan, already in the order every consumer wants. + + Every field is indexed by partial -- one entry per ``(token, topk slot)`` -- + and sorted by global expert, so a destination's partials are one contiguous + run, grouped by local expert inside it, which feeds the grouped GEMM + directly. + + Nothing here is ever read back to the host. The three per-destination + vectors go into the slot headers on the device, and the arrays are shipped + whole, so the sender never needs to know how the routing came out. + + Attributes: + counts: Partials per global expert, padded to + ``ffn_size * expert_per_rank``. These are the header's group list. + routed_per_rank: Partials each FFN rank receives. + segment_start: Where each FFN rank's run of partials begins. + expand_idx: Per partial, the token it carries. + weights: Per partial, its topk weight. + """ + + counts: Tensor + routed_per_rank: Tensor + segment_start: Tensor + expand_idx: Tensor + weights: Tensor + + +def plan_dispatch( + topk_ids: Tensor, + topk_weights: Tensor, + *, + ffn_size: int, + expert_per_rank: int, + expert_bounds_probe: Tensor | None = None, +) -> DispatchPlan: + """Cluster ``(token, topk_slot)`` partials by destination expert. + + Sorting by the global expert id groups partials by destination rank and, in + the same pass, by local expert inside each destination, which is every + grouping any consumer needs. One sort and a bounds lookup is the whole plan. + + Every step is a whole-tensor op and the host learns nothing here, by + design: a readback of the routing would block the send path behind the + device once per MoE layer. + """ + num_slots = topk_ids.shape[1] + # Sort the expert ids at their own width. Widening to int64 first cost a + # full copy of the partials and then doubled the bytes the sort moves; the + # ids index at most ffn_size * expert_per_rank experts, which never needs + # 64 bits. Measured on the real shape (2048 tokens, topk 6): 0.129ms per + # call for the widened sort against 0.087ms for this one. + flat = topk_ids.reshape(-1) + sorted_experts, order = torch.sort(flat, stable=True) + + # Where each expert's partials start and stop, read off the sorted ids. + # torch.bincount would do this in one call, but it sizes its output from the + # data's maximum and so copies that maximum to the host -- a pageable + # readback measured at 783us per call, once per MoE layer, which was the + # largest single host cost on the Attention rank. searchsorted needs no such + # thing, because the expert count is known, and it hands back the offsets + # that would otherwise be a second pass. + # The probe is the same every call, so the caller keeps one rather than + # allocating and filling an arange per MoE layer. + if expert_bounds_probe is None: + expert_bounds_probe = torch.arange( + ffn_size * expert_per_rank + 1, + device=flat.device, + dtype=flat.dtype, + ) + bounds = torch.searchsorted(sorted_experts, expert_bounds_probe) + offsets = bounds[:-1] + counts = bounds[1:] - offsets + + # A destination reads the whole batch out of its slot, so a partial names + # its token directly and needs no rebasing onto shipped rows. + expand_idx = torch.div(order, num_slots, rounding_mode="floor").to(torch.int32) + weights = topk_weights.reshape(-1)[order].to(torch.float32) + + # A rank owns a contiguous block of experts, so its partials begin where its + # first expert's do and run to the end of its last. + return DispatchPlan( + counts=counts, + routed_per_rank=counts.view(ffn_size, expert_per_rank).sum(1), + segment_start=offsets.view(ffn_size, expert_per_rank)[:, 0], + expand_idx=expand_idx, + weights=weights, + ) + + +class GpuAsyncAFDConnector(AFDConnectorBase): + """NVSHMEM symmetric-window asynchronous connector for CUDA AFD.""" + + control_plane = None + # Dynamo must split the graph at the MoE round trip instead of tracing into + # the connector, so this connector's dispatch and receive go through the + # opaque ops. + uses_opaque_moe_ops = True + + # The base builds this in __init__ from parse_extra_config; the narrowed + # annotation is what lets mypy see the async connector's own fields. + extra_info: GpuAsyncExtraInfo + + @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 + # Combine has to know whether a reply carries shared-expert rows before + # it arrives, and only the model config says so. + self.has_shared_experts = bool(hf_config.n_shared_experts) + self.payload_dtype = vllm_config.model_config.dtype + self.max_seq_len = vllm_config.scheduler_config.max_num_batched_tokens + # attn_ranks_per_dp is how the connector shards an Attention replica, + # and vLLM's tensor_parallel_size is how the model is sharded. They name + # the same split, so a mismatch sends dispatches sized for one world + # into ranks laid out for the other. The NPU twin enforces this in + # feature_validation; do it here, where both numbers first meet. + self.tp_size = self.extra_info.attn_ranks_per_dp + if afd_config.role == "attention": + tensor_parallel_size = int(vllm_config.parallel_config.tensor_parallel_size) + if self.tp_size != tensor_parallel_size: + raise ValueError( + "AFD async GPU attn_ranks_per_dp must equal Attention " + f"tensor_parallel_size, got attn_ranks_per_dp={self.tp_size} " + f"and tensor_parallel_size={tensor_parallel_size}", + ) + + 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" + + # Stages that can be in flight at once, from either splitter. vLLM's + # own ubatching (DBO) drives two forwards concurrently and stamps each + # with its ubatch index, which lands here as the stage -- so it needs + # rings just as much as this connector's async_moe_ubatching does. + # Counting only the latter gave both DBO ubatches ring 0, where the + # second dispatch overwrote the first's slot and flag and the reply the + # first was waiting for never came: the two ubatch threads then + # deadlocked, one inside recv_ffn_output and one waiting to be yielded + # to. + parallel_config = vllm_config.parallel_config + vllm_stages = ( + max(1, int(parallel_config.num_ubatches)) + if bool(parallel_config.use_ubatching) + else 1 + ) + connector_stages = ( + max(1, self.extra_info.async_moe_num_ubatches) + if self.extra_info.async_moe_ubatching + else 1 + ) + self.num_stages = max(vllm_stages, connector_stages) + # An implied ring depth grows to fit; a pinned one is taken as given so + # a deliberate choice still fails loudly below rather than being + # silently overruled. + self.ring_depth = ( + self.extra_info.ring_depth + if self.extra_info.ring_depth_pinned + else max(self.extra_info.ring_depth, self.num_stages) + ) + if self.ring_depth < self.num_stages: + raise ValueError( + f"ring_depth {self.ring_depth} cannot serve " + f"{self.num_stages} async MoE stages; each stage needs a slot " + "of its own", + ) + # 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) + # Payload rows are distinct tokens, so a batch is their bound no matter + # how skewed the gate is. Only the 4-byte-per-partial index arrays need + # the every-partial-to-one-rank worst case. + self.token_cap = max(1, self.max_seq_len) + self.partial_cap = max(1, self.max_seq_len * self.topk) + # Shared-expert rows are split across the FFN ranks, so a slot holds a + # fraction of the batch, not all of it -- and none at all when the model + # has no shared experts. Both roles derive this from the same config, so + # the symmetric allocation still matches. + self.shared_cap = ( + -(-self.token_cap // self.ffn_size) if self.has_shared_experts else 0 + ) + self.layout = SlotLayout.build( + expert_per_rank=self.expert_per_rank, + partial_cap=self.partial_cap, + token_cap=self.token_cap, + shared_cap=self.shared_cap, + hidden_size=self.hidden_size, + payload_itemsize=torch.empty(0, dtype=self.payload_dtype).element_size(), + ) + + # Device headers, one buffer per (layer_idx, num_tokens) dispatch shape. + self._header_device: dict[tuple[int, int], Tensor] = {} + self._seq_device: Tensor | None = None + # The searchsorted probe for the dispatch plan: identical every call, + # so it is built once instead of per MoE layer. + self._expert_bounds_probe: Tensor | None = None + + self.pg: ProcessGroup | None = None + self.window: SymmWindow | None = None + self._pending: dict[int, list[_PendingDispatch]] = {} + # Dispatch payloads parked by afd_async_dispatch for the deferred + # afd_async_recv one layer later, keyed by connector stage. The two + # DBO ubatch threads share the connector, so the stage key is what + # keeps one half's payload out of the other's receive. + self.pending_cam_dispatches: dict[int, object] = {} + self._free_rings: dict[int, list[int]] = {} + + @property + def is_initialized(self) -> bool: + return self._initialized + + def _rings_for_stage(self, stage_idx: int) -> list[int]: + """Ring slots this stage owns. + + A ring names a window slot, so stages must not share one: two stages of + the same layer are in flight at the same time, and handing both the same + slot means the second dispatch overwrites the first's payload and flag, + after which the reply the first is waiting for never arrives. + """ + first = stage_idx % self.num_stages + return list(range(first, self.ring_depth, self.num_stages)) + + 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) + # The dispatch sequence number lives on the device so that a CUDA graph + # can advance it. An FFN rank recognizes an arrival by its flag holding + # something other than the value it last consumed, so the number has to + # rise on every dispatch -- including on every replay of a captured + # graph, which runs no Python at all and would otherwise stamp the one + # value that was live when the graph was recorded. + self._seq_device = torch.zeros((), dtype=torch.int32, device=device) + 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(self.num_stages): + self._free_rings[stage] = self._rings_for_stage(stage) + logger.info( + "AFD async GPU window ready: role=%s role_rank=%d world_rank=%d/%d " + "regions=%d rings=%d partial_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.partial_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 + # Teardown ordering: an Attention rank drains its pending combines + # before the FFN world closes. Dropping them here is the last resort -- + # the rows are gone, and only the FFN side's shutdown release keeps the + # waiter from hanging -- so say so rather than clearing in silence. + outstanding = sum(len(queue) for queue in self._pending.values()) + if outstanding: + logger.warning( + "AFD async GPU closing with %d un-combined dispatch(es); " + "their FFN replies are discarded", + outstanding, + ) + self._pending.clear() + self._free_rings.clear() + self._header_device.clear() + self._seq_device = None + 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 + + def _shared_slice(self, ffn_rank: int, num_tokens: int) -> slice: + """Token range whose shared-expert output ``ffn_rank`` owns. + + Contiguous chunks, so a rank's slice is a view: no index to build, none + to gather through, and none to put on the wire. Empty when the model has + no shared experts, which is what keeps the payload off the wire and the + field out of the slot. + + The header's ``shared_tokens`` and the rows actually written have to + agree, and they are produced by different callers, so both come from + here rather than from two copies of the arithmetic. + """ + if not self.has_shared_experts: + return slice(0, 0) + return slice( + ffn_rank * num_tokens // self.ffn_size, + (ffn_rank + 1) * num_tokens // self.ffn_size, + ) + + def _headers_for_shape( + self, + *, + layer_idx: int, + num_tokens: int, + ) -> Tensor: + """Device headers for one dispatch shape, prefix already filled. + + The prefix -- magic, version, layer, batch size, flags -- is fixed by + ``(layer_idx, num_tokens)``, so it is written once here and then left + alone. A dispatch only rewrites the routing tail. + + That the prefix never changes is what lets it leave the layer path. It + used to be copied down from the host on every dispatch, not to carry + anything new but to overwrite the previous layer's prefix in a buffer + every layer shared. One buffer per shape removes the overwrite, and + with it a strided host-to-device copy per MoE layer -- and, under a + CUDA graph, a copy node that re-shipped a constant on every replay. + + Costs ``ffn_size * header_words`` int32s per distinct shape: MoE layers + times the token counts they run at, a few KiB in practice. + """ + key = (layer_idx, num_tokens) + headers = self._header_device.get(key) + if headers is not None: + return headers + + assert self.window is not None + device = self.window.device + # torch.compiler.is_compiling() first, and not merely for speed: under + # Dynamo the capture query below is a torch.* op returning a bool, + # which cannot be traced into an FX graph and aborts compilation. It + # folds to a constant while tracing, so the branch drops out there and + # still guards the eager capture path, which is the one that can hit it. + if ( + not torch.compiler.is_compiling() + and device.type == "cuda" + and torch.cuda.is_current_stream_capturing() + ): + # Allocating here would come from the graph's private pool and the + # prefix copy would be recorded against a host buffer that is gone + # by the first replay. Warmup runs the same shapes capture does, so + # reaching this means the shape was never warmed. + raise RuntimeError( + "AFD async GPU dispatch shape " + f"(layer={layer_idx}, num_tokens={num_tokens}) was first seen " + "during CUDA graph capture; warm it before capturing", + ) + # Outside inference mode on purpose. This cache outlives the call that + # fills it, and the first call for a shape can land inside vLLM's + # inference-mode forward -- which brands the tensor an inference tensor, + # so the next dispatch's write to the routing tail raises "Inplace + # update to inference tensor outside InferenceMode". That is exactly + # what a compiled prefill hits, because AOT compilation moves the first + # touch of each shape inside the compiled region. + with torch.inference_mode(False): + headers = torch.empty( + (self.ffn_size, self.layout.header_words), + dtype=torch.int32, + device=device, + ) + # Written on the device, not copied down. The copy this replaces was + # from pageable memory and therefore synchronous, and a blocking CUDA + # call here deadlocks under vLLM's ubatching: this runs inside a ubatch + # thread, which cannot block before reaching its next yield without + # stranding the peer thread waiting to be yielded to. + fill_header_prefix( + headers, + layer_idx=layer_idx, + num_tokens=num_tokens, + flags=0, + ) + self._header_device[key] = headers + return headers + + def _headers_for( + self, + *, + layer_idx: int, + num_tokens: int, + plan: DispatchPlan, + ) -> Tensor: + """Assemble one dispatch header per FFN rank, on the device. + + Only the routing tail is written here, straight from the plan, without + ever leaving the device. Reading the routing back to encode it on the + host instead was the last synchronize on the layer path, and a profile + put it at 392us a call once per MoE layer -- not the copy, but the host + waiting for everything queued ahead of it, which capped how far ahead of + the device the host could ever get. + """ + headers = self._headers_for_shape( + layer_idx=layer_idx, + num_tokens=num_tokens, + ) + headers[:, H_ROUTED_TOKENS] = plan.routed_per_rank + headers[:, H_SEGMENT_START] = plan.segment_start + headers[:, HEADER_FIXED_WORDS:] = plan.counts.view( + self.ffn_size, + self.expert_per_rank, + ) + return headers + + # ================================================================== + # 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", + ) + if ( + os.environ.get(AFD_ASYNC_DEBUG_ENV) + and context.metadata.layer_idx < AFD_ASYNC_DEBUG_LAYERS + ): + logger.info( + "AFD debug A%d send layer=%d tokens=%d %s %s", + self.role_rank, + context.metadata.layer_idx, + int(hidden_states.shape[0]), + _debug_norm("hidden", hidden_states), + _debug_norm("w", topk_weights), + ) + 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}", + ) + expected_shape = (num_tokens, self.topk) + if tuple(topk_ids.shape) != expected_shape: + raise ValueError( + f"topk_ids shape must be {expected_shape}, got {tuple(topk_ids.shape)}", + ) + # The weights are flattened alongside the ids to give each partial its + # own weight, so a mismatched shape would silently misalign them. + if tuple(topk_weights.shape) != expected_shape: + raise ValueError( + f"topk_weights shape must be {expected_shape}, " + f"got {tuple(topk_weights.shape)}", + ) + + stage_idx = metadata.stage_idx + rings = self._free_rings.setdefault( + stage_idx, + self._rings_for_stage(stage_idx), + ) + 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) + # Bump the counter on the device, not on the host. Under CUDA graph + # capture the host runs this function once and every later replay runs + # only the recorded kernels, so a Python counter would freeze at the + # captured value and stamp every FFN flag with it -- and an FFN rank + # notices a dispatch precisely by its flag changing. + assert self._seq_device is not None + self._seq_device += 1 + + if self._expert_bounds_probe is None: + self._expert_bounds_probe = torch.arange( + self.ffn_size * self.expert_per_rank + 1, + device=topk_ids.device, + dtype=topk_ids.dtype, + ) + plan = plan_dispatch( + topk_ids, + topk_weights, + ffn_size=self.ffn_size, + expert_per_rank=self.expert_per_rank, + expert_bounds_probe=self._expert_bounds_probe, + ) + # 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)) + shared_slices: list[slice] = [] + headers = self._headers_for( + layer_idx=metadata.layer_idx, + num_tokens=num_tokens, + plan=plan, + ) + for ffn_rank in range(self.ffn_size): + # Round-robin needed an arange, a gather and a whole slot field per + # peer per layer to achieve the balance this contiguous split gets + # from a view. + shared = self._shared_slice(ffn_rank, num_tokens) + shared_slices.append(shared) + # Everything but the shared slice goes out whole. The index arrays + # cost a fraction of a percent of the slot, and the payload rows a + # destination does not need are the few tokens none of whose topk + # slots landed on it -- 1.6% of them at 2A2F. Sizing either to the + # routing is what used to make the host wait for the device here. + window.write_slot( + peer=self.attn_size + ffn_rank, + region=self.role_rank, + ring=ring, + header=headers[ffn_rank], + expand_idx=plan.expand_idx, + weights=plan.weights, + routed_x=hidden_states, + shared_x=hidden_states[shared], + flag_value=self._seq_device, + ) + + logger.debug( + "AFD dispatch sent: A%d layer=%d stage=%d tokens=%d ring=%d " + "partials=%d awaiting_ffn=%s", + self.role_rank, + metadata.layer_idx, + stage_idx, + num_tokens, + ring, + num_tokens * self.topk, + expected_ffn, + ) + self._pending.setdefault(stage_idx, []).append( + _PendingDispatch( + context=context, + shared_slices=shared_slices, + num_tokens=num_tokens, + ring=ring, + expected_ffn=expected_ffn, + ), + ) + + def recv_ffn_output( + self, + ref_tensor: Tensor, + ubatch_idx: int = 0, + **kwargs: Any, + ) -> Tensor: + """Queue this layer's combine and return, without waiting for the data. + + The waiting happens on the stream: every reply is answered into a known + slot, carries a known number of rows, and stamps the dispatch sequence + it answers, so the whole combine can be enqueued before any of it has + arrived. The host goes straight on to the next layer, which is the point + -- polling for the reply here left the GPU with nothing queued behind + the wait, and a profile found 86% of kernel launches executing within + 5us of being issued because of it. + + Nothing reports what turned up, so there is no per-arrival header check + any more. The stream wait replaces it: it blocks on one specific slot + being marked ready, where the poll took whatever had landed and had to + check afterwards that it was the right thing. + """ + + 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) + + # Accumulate in the payload dtype. Each row takes at most one + # contribution per FFN rank plus its shared one -- the topk sum already + # happened on the FFN side -- so there is little left for a wider + # accumulator to protect, and float32 cost a widening pass over every + # arriving block plus a narrowing one on the way out. + accumulator = torch.zeros( + (pending.num_tokens, self.hidden_size), + dtype=self.payload_dtype, + device=ref_tensor.device, + ) + for ffn_rank in pending.expected_ffn: + # An FFN rank replies into the region it owns, on the ring the + # dispatch used, and marks it with FLAG_REPLY_READY. + window.stream_wait(ffn_rank, pending.ring, FLAG_REPLY_READY) + # A reply is a whole batch, already weighted and summed over that + # rank's experts, with a zero row wherever the rank held none of a + # token's experts. Row i answers token i, so this is a plain add: + # the scatter it replaces was the second largest kernel on the rank. + accumulator += window.local_routed( + ffn_rank, + pending.ring, + pending.num_tokens, + ) + shared = pending.shared_slices[ffn_rank] + shared_tokens = shared.stop - shared.start + if self.has_shared_experts and shared_tokens: + accumulator[shared] += window.local_shared( + ffn_rank, + pending.ring, + shared_tokens, + ) + # Arm the flag for the next use of this slot, after the reads above + # and before any dispatch that could reuse the ring -- all on this + # stream, in that order. The peer cannot re-raise it early either: + # it only replies to a dispatch, and that dispatch is enqueued + # behind this reset. + window.clear_flag(ffn_rank, pending.ring) + if ( + os.environ.get(AFD_ASYNC_DEBUG_ENV) + and pending.context.metadata.layer_idx < AFD_ASYNC_DEBUG_LAYERS + ): + logger.info( + "AFD debug A%d combine layer=%d ring=%d tokens=%d %s", + self.role_rank, + pending.context.metadata.layer_idx, + pending.ring, + pending.num_tokens, + _debug_norm("acc", accumulator), + ) + logger.debug( + "AFD combine queued: A%d layer=%d stage=%d ring=%d from=%s", + self.role_rank, + pending.context.metadata.layer_idx, + ubatch_idx, + pending.ring, + pending.expected_ffn, + ) + + self._free_rings.setdefault(ubatch_idx, []).append(pending.ring) + return accumulator + + # ================================================================== + # FFN-side data path + # ================================================================== + + def recv_attn_output( + self, + ubatch_idx: int = 0, + routed_out: Tensor | None = None, + **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. + + The slot holds the sender's whole batch and every sender's partials, so + this rank takes the run of partials the header points it at and gathers + the tokens they name -- a local gather that replaces both the duplicate + rows the sender used to put on the wire and the readback it needed to + size a per-destination slice. + + ``routed_out`` is where that gather lands. A caller that already has a + destination -- a CUDA graph's input buffer, which must be written in + place -- passes it and gets the rows staged for free, because the + gather had to write somewhere regardless. Without it the gather + allocates, as it always did. Too small a buffer is not an error: the + gather allocates and the caller sees rows it did not stage, which is + the same fallback an oversized work item takes anyway. + """ + window = self._require_initialized() + timeout_ms = int(kwargs.get("timeout_ms", 0)) + arrived = window.wait(timeout_s=timeout_ms / 1000.0 if timeout_ms else None) + if arrived is None: + raise TimeoutError("AFD async GPU dispatch recv timed out") + + header = arrived.header + # Four fields the wire no longer carries, because this side can work + # them out. The slot's own region names the Attention rank that wrote + # it; ``_rings_for_stage`` hands stage s the rings congruent to s, so + # the ring names the stage; and the shared-token split is the same + # function of the batch size on both sides, evaluated here for this + # rank's own slice. + src_role_rank = arrived.region + stage_idx = arrived.ring % self.num_stages + shared = self._shared_slice(self.role_rank, header.num_tokens) + shared_tokens = shared.stop - shared.start + if header.is_shutdown: + raise ConnectorShutdown( + f"Attention rank {src_role_rank} announced shutdown", + ) + # A slot is peer-written memory, so its routing tail gets one sanity + # check before the grouped GEMM trusts it. Checking the decoded host + # values here is free; the equivalent check on the device tensor cost a + # synchronize per work item. + if sum(header.expert_counts) != header.routed_tokens: + raise RuntimeError( + f"AFD async GPU dispatch header from A{src_role_rank} " + f"has expert counts summing to {sum(header.expert_counts)} but " + f"declares {header.routed_tokens} routed tokens", + ) + + expand_idx = window.local_expand_idx( + arrived.region, + arrived.ring, + header.routed_tokens, + header.segment_start, + ).to(torch.int64) + states = GpuAsyncTransferState( + region=arrived.region, + ring=arrived.ring, + src_role_rank=src_role_rank, + layer_idx=header.layer_idx, + stage_idx=stage_idx, + num_tokens=header.num_tokens, + routed_tokens=header.routed_tokens, + shared_tokens=shared_tokens, + group_list=window.local_expert_counts(arrived.region, arrived.ring), + expand_idx=expand_idx, + weights=window.local_weights( + arrived.region, + arrived.ring, + header.routed_tokens, + header.segment_start, + ), + expand_x_shared=window.local_shared( + arrived.region, + arrived.ring, + 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, + src_role_rank, + header.layer_idx, + stage_idx, + header.routed_tokens, + shared_tokens, + arrived.region, + arrived.ring, + ) + metadata = AFDTransferMetadata.create_ffn_metadata( + layer_idx=header.layer_idx, + stage_idx=stage_idx, + seq_lens=[max(1, header.routed_tokens)], + ) + arrived_rows = window.local_routed( + arrived.region, + arrived.ring, + header.num_tokens, + ) + if routed_out is not None and header.routed_tokens <= routed_out.shape[0]: + gathered = routed_out[: header.routed_tokens] + torch.index_select(arrived_rows, 0, expand_idx, out=gathered) + states.staged_routed = True + else: + gathered = arrived_rows.index_select(0, expand_idx) + return AFDA2FTransferPayload( + hidden_states=gathered, + context=AFDTransferContext(metadata=metadata, states=states), + ) + + def send_ffn_output( + self, + ffn_output: Tensor, + context: AFDTransferContext, + **kwargs: Any, + ) -> None: + """Reduce expert output back to one row per token and write it back. + + Every partial of a token that landed on this rank is weighted and summed + here, into the token's own row of a full batch. Tokens this rank held no + expert for keep their zero row, which is what lets the Attention side + add replies together without an index. Doing the reduction on this side + keeps the duplicates off the wire. + + Both the weighting and the sum stay in the payload dtype. Widening to + float32 first cost two extra passes over ``[partials, hidden]`` and made + the scatter move twice the bytes, which a profile showed costing almost + as much GPU time as the expert GEMM itself -- to protect a sum of at + most ``topk`` terms whose result goes on the wire narrowed anyway. + """ + window = self._require_initialized() + states = context.states + if not isinstance(states, GpuAsyncTransferState): + raise RuntimeError( + "AFD async GPU send_ffn_output requires GpuAsyncTransferState", + ) + if states.expand_idx is None or states.weights is None: + raise RuntimeError( + "AFD async GPU send_ffn_output requires the dispatch expansion", + ) + reduced = torch.zeros( + (states.num_tokens, self.hidden_size), + dtype=self.payload_dtype, + device=ffn_output.device, + ) + if states.routed_tokens: + weighted = ffn_output * states.weights.unsqueeze(1).to(ffn_output.dtype) + reduced.index_add_(0, states.expand_idx, weighted) + + shared_output: Tensor | None = kwargs.get("shared_output") + if ( + os.environ.get(AFD_ASYNC_DEBUG_ENV) + and states.layer_idx < AFD_ASYNC_DEBUG_LAYERS + ): + logger.info( + "AFD debug F%d reply layer=%d ring=%d routed_rows=%d %s %s", + self.role_rank, + states.layer_idx, + states.ring, + states.routed_tokens, + _debug_norm("routed", reduced), + _debug_norm("shared", shared_output), + ) + # ``reduced`` is float32 and the slot is the payload dtype; the copy + # inside write_slot casts on its way into the peer window, so the + # narrowing costs no extra pass over the rows. + window.write_slot( + peer=states.src_role_rank, + region=self.role_rank, + ring=states.ring, + # No header. The rank waiting for this reply knows its shape before + # it exists -- that is the premise of waiting on a stream rather + # than polling -- so it never reads one, and every field a header + # could carry it either chose itself or can derive. Writing one was + # a copy per peer per layer that nothing consumed. + header=None, + expand_idx=None, + weights=None, + routed_x=reduced, + shared_x=shared_output, + # A constant marker: the Attention rank has to name the value it + # waits for before the reply exists, and a captured graph freezes + # that name, so it can only ever be a constant. + flag_value=FLAG_REPLY_READY, + ) + + # ================================================================== + # Connector-driven FFN loop + # ================================================================== + + def recv_ffn_work_item( + self, + *, + stage_idx: int, + max_num_tokens: int, + routed_out: Tensor | None = None, + ) -> GpuAsyncFFNWorkItem: + """Receive and normalize one connector-driven FFN dispatch item.""" + recv_output = self.recv_attn_output( + ubatch_idx=stage_idx, + routed_out=routed_out, + timeout_ms=self.extra_info.recv_poll_timeout_ms, + ) + states = recv_output.context.states + assert isinstance(states, GpuAsyncTransferState) + if ( + os.environ.get(AFD_ASYNC_DEBUG_ENV) + and states.layer_idx < AFD_ASYNC_DEBUG_LAYERS + ): + logger.info( + "AFD debug F%d recv layer=%d ring=%d routed=%d shared=%d %s", + self.role_rank, + states.layer_idx, + states.ring, + states.routed_tokens, + states.shared_tokens, + _debug_norm("hidden", recv_output.hidden_states), + ) + 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) + ) + # The same counter the dispatches stamp: a peer spots this exactly the + # way it spots a dispatch, by the flag differing from what it consumed + # last, so the two must not hand out the same number twice. + assert self._seq_device is not None + self._seq_device += 1 + header = encode_header( + self.layout, + layer_idx=0, + num_tokens=0, + routed_tokens=0, + flags=FLAG_SHUTDOWN_BIT, + expert_counts=[0] * self.expert_per_rank, + ) + for peer in peers: + # Header only: the shutdown bit is the whole message, so every + # payload field is empty and write_slot skips it. + window.write_slot( + peer=peer, + region=self.role_rank, + ring=0, + header=header, + expand_idx=None, + weights=None, + routed_x=None, + shared_x=None, + flag_value=self._seq_device, + ) + if not self.is_attention: + self._release_outstanding_replies(window, peers) + + def _release_outstanding_replies( + self, + window: SymmWindow, + peers: Iterable[int], + ) -> None: + """Free any Attention rank waiting on a reply this rank will not send. + + A combine wait is released only by ``FLAG_REPLY_READY`` landing on this + rank's region, so an FFN rank that leaves between a dispatch and its + reply strands its waiter forever -- and the waiter may be stranded + inside a graph replay, where no Python runs to notice the shutdown that + was just announced. Stamping the flag on every ring frees it. + + Re-stamping a ring that was already answered is harmless: the reply flag + is a constant, so the second write writes the value already there. A + ring that was never answered releases its waiter over a stale payload, + which is the right trade at teardown -- the numbers are discarded on the + way out, and the alternative is a rank that never exits. + """ + for peer in peers: + for ring in range(self.ring_depth): + window.write_slot( + peer=peer, + region=self.role_rank, + ring=ring, + header=None, + expand_idx=None, + weights=None, + routed_x=None, + shared_x=None, + flag_value=FLAG_REPLY_READY, + ) + + +__all__ = [ + "AFD_ASYNC_GPU_GROUP_NAME", + "ConnectorShutdown", + "DispatchPlan", + "GpuAsyncAFDConnector", + "GpuAsyncExtraInfo", + "GpuAsyncFFNWorkItem", + "GpuAsyncTransferState", + "plan_dispatch", +] diff --git a/afd_plugin/connectors/gpu/async_moe_op.py b/afd_plugin/connectors/gpu/async_moe_op.py new file mode 100644 index 00000000..c8de3f41 --- /dev/null +++ b/afd_plugin/connectors/gpu/async_moe_op.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Dynamo-opaque ops for the async GPU connector's MoE dispatch and receive. + +Enabling CUDA graphs on the Attention side puts vLLM's AOT compilation in +front of the model, and tracing used to run straight through the MoE proxy +into the connector -- raw NVSHMEM pointer views, host-side caches, ctypes +driver calls -- and abort. The data path itself is capture-safe (the flag +protocol keeps replays correct; see the ``async_gpu`` module docs); the +obstacle was purely Python visibility. + +Two ops carve the round trip at the points the model already defers across: +``afd_async_dispatch`` sends one layer's tokens and stashes the payload the +receive will need; ``afd_async_recv`` consumes the stash, waits for the reply +on the stream, and restores the model layout. Dynamo splits at both, and the +Python between them -- the deferred-receive bookkeeping -- is plain control +flow it can trace. During cooperative capture the two ubatch threads run the +impls once each, and the kernel order their alternation produced is what the +graph replays. +""" + +from __future__ import annotations + +import torch +from vllm.utils.torch_utils import direct_register_custom_op + +_DISPATCH_OP_NAME = "afd_async_dispatch" +_RECV_OP_NAME = "afd_async_recv" +_REGISTERED = False + + +def _resolve_stage_idx(afd_metadata) -> int: + # vLLM tracks the ubatch by thread, not on the forward context; under + # cooperative capture each thread resolves its own stage here, and the + # kernel order that produces is what the captured graph replays. + from afd_plugin.v1.worker.dbo import current_dbo_ubatch_id + + dbo_ubatch_id = current_dbo_ubatch_id() + return afd_metadata.stage_idx if dbo_ubatch_id is None else int(dbo_ubatch_id) + + +def _dispatch_impl( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + layer_idx: int, +) -> torch.Tensor: + # Deferred imports: the connector sits below the models that trace this + # op, and the forward-context helpers pull the model package in. + from afd_plugin.connectors import AFDTransferContext, AFDTransferMetadata + from afd_plugin.model_executor.models import ( + get_afd_metadata_from_forward_context, + ) + from afd_plugin.model_executor.models.npu.async_cam_layout import ( + prepare_cam_dispatch_payload, + ) + + afd_metadata = get_afd_metadata_from_forward_context() + if afd_metadata is None: + raise RuntimeError("afd_async_dispatch requires AFD forward metadata") + stage_idx = _resolve_stage_idx(afd_metadata) + afd_metadata.stage_idx = stage_idx + connector = afd_metadata.connector + # FlashComm1 token sharding is Ascend-only: the CUDA forward context + # carries no such field, so the dispatch always sees a replicated token + # dimension and the payload passes through unchanged. + payload = prepare_cam_dispatch_payload( + hidden_states, + topk_weights, + topk_ids, + None, + use_sequence_parallel=False, + ) + metadata = AFDTransferMetadata.create_attention_metadata( + layer_idx=layer_idx, + stage_idx=stage_idx, + seq_len=int(payload.hidden_states.shape[0]), + ) + connector.send_attn_output( + payload.hidden_states, + AFDTransferContext(metadata=metadata), + topk_weights=payload.topk_weights, + topk_ids=payload.topk_ids, + ) + connector.pending_cam_dispatches[stage_idx] = payload # type: ignore[attr-defined] + return payload.hidden_states + + +def _dispatch_fake( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + layer_idx: int, +) -> torch.Tensor: + return hidden_states + + +def _recv_impl(hidden_states: torch.Tensor) -> torch.Tensor: + from afd_plugin.model_executor.models import ( + get_afd_metadata_from_forward_context, + ) + from afd_plugin.model_executor.models.npu.async_cam_layout import ( + restore_cam_dispatch_output, + ) + + afd_metadata = get_afd_metadata_from_forward_context() + if afd_metadata is None: + raise RuntimeError("afd_async_recv requires AFD forward metadata") + stage_idx = _resolve_stage_idx(afd_metadata) + connector = afd_metadata.connector + payload = connector.pending_cam_dispatches.pop( # type: ignore[attr-defined] + stage_idx, None + ) + if payload is None: + raise RuntimeError( + f"AFD async receive on stage {stage_idx} has no pending dispatch", + ) + local_ffn_output = connector.recv_ffn_output( + ref_tensor=hidden_states, + ubatch_idx=stage_idx, + ) + return restore_cam_dispatch_output(local_ffn_output, payload.layout) + + +def _recv_fake(hidden_states: torch.Tensor) -> torch.Tensor: + return torch.empty_like(hidden_states) + + +def _register_once(op_name: str, op_func, fake_impl) -> None: + try: + direct_register_custom_op( + op_name=op_name, + op_func=op_func, + mutates_args=[], + fake_impl=fake_impl, + # The body is plain Python that hands the tensors to the connector, + # which does its own device work, so there is nothing + # backend-specific to specialise. Registering under the platform + # dispatch key instead would leave the op undefined for CPU, where + # the NPU forward's unit tests exercise this path. + dispatch_key="CompositeExplicitAutograd", + ) + except RuntimeError as exc: + # A prior import of this module can leave the op in torch's + # process-global registry while this module's flag is back to False; + # reuse the registered op instead of redefining it. + if "already" not in str(exc).lower(): + raise + + +def register_async_moe_ops() -> tuple: + """Register both ops once and return their callable handles.""" + global _REGISTERED + if not _REGISTERED: + _register_once(_DISPATCH_OP_NAME, _dispatch_impl, _dispatch_fake) + _register_once(_RECV_OP_NAME, _recv_impl, _recv_fake) + _REGISTERED = True + return ( + getattr(torch.ops.vllm, _DISPATCH_OP_NAME), + getattr(torch.ops.vllm, _RECV_OP_NAME), + ) + + +register_async_moe_ops() + +__all__ = ["register_async_moe_ops"] diff --git a/afd_plugin/connectors/npu/async_cam.py b/afd_plugin/connectors/npu/async_cam.py index 6ec5e705..0c569815 100644 --- a/afd_plugin/connectors/npu/async_cam.py +++ b/afd_plugin/connectors/npu/async_cam.py @@ -46,6 +46,12 @@ coerce_extra_str, coerce_optional_extra_positive_int, ) +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, @@ -68,9 +74,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( @@ -204,23 +208,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. @@ -879,56 +866,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 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..21e4f86b 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, @@ -165,7 +165,11 @@ def __init__( super().__init__(layer_idx=layer_idx) self.is_internal_router = is_internal_router - def forward( + # The base proxy stands in for a plain MLP and this one for vLLM's + # FusedMoE, so the two forwards match different upstream call shapes and + # cannot be substituted for each other. Nothing calls them polymorphically; + # the inheritance is for _send_and_receive, not for forward. + def forward( # type: ignore[override] self, hidden_states: torch.Tensor, router_logits: torch.Tensor, @@ -703,7 +707,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: + # Both async connectors run this schedule: it is the connector-driven + # shape, not an Ascend one. The module still lives under models/npu/ + # because CAM got here first. + if self.afd_config.connector in AFD_ASYNC_CONNECTORS: 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..1734e184 --- /dev/null +++ b/afd_plugin/model_executor/models/gpu/deepseek_v2_attention_gate.py @@ -0,0 +1,135 @@ +# 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 its own weight, and vLLM's ``fused_experts`` computes exactly the local +expert output, already scaled, in the epilogue it runs anyway. +""" + +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]) + 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 + ), + ) + # ``output_size`` keeps this off the device: without it repeat_interleave + # reads the counts back to the host to size its output, which is a + # synchronize on every work item. It also enforces what the removed + # ``counts.sum() == num_rows`` check used to, raising if they disagree. + expert_ids = torch.repeat_interleave( + torch.arange( + num_local_experts, + device=hidden_states.device, + dtype=torch.int32, + ), + counts, + output_size=num_rows, + ).unsqueeze(1) + # Mirrors the NPU gate path: scale the routed branch unless fp16, where the + # native model instead scales the shared branch down. ``fused_experts`` + # multiplies every row by its topk weight in an epilogue it runs anyway, so + # feeding the factor in there costs nothing, where scaling its output cost a + # full pass over ``[num_partials, hidden]``. The topk weighting itself is + # not applied here -- the connector owns it, and applies it when it reduces + # the partials back to one row per token. + routed_scaling_factor = runner.routed_scaling_factor + scale_shared_instead = hidden_states.dtype == torch.float16 + row_weights = torch.full( + (num_rows, 1), + 1.0 if scale_shared_instead else routed_scaling_factor, + dtype=torch.float32, + device=hidden_states.device, + ) + + routed_output = fused_experts( + hidden_states, + routed_experts.w13_weight, + routed_experts.w2_weight, + row_weights, + expert_ids, + global_num_experts=num_local_experts, + expert_map=None, + ) + + shared_output = None + shared_experts = runner._shared_experts + # A dispatch's round-robin shared slice is empty for a rank whenever the + # ubatch holds fewer tokens than FFN ranks, so zero rows are a normal item + # shape, not a missing payload. + if ( + shared_experts is not None + and expand_x_shared is not None + and expand_x_shared.shape[0] > 0 + ): + # 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) + + if scale_shared_instead and 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 383ea56e..b17c6d01 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 @@ -4,6 +4,7 @@ from __future__ import annotations +from collections.abc import Mapping from copy import copy from itertools import islice from typing import TYPE_CHECKING @@ -24,6 +25,7 @@ AFDTransferContext, AFDTransferMetadata, ) +from afd_plugin.connectors.gpu.async_moe_op import register_async_moe_ops from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context from afd_plugin.model_executor.models.npu.async_cam_layout import ( AsyncMoeUbatchMetadata, @@ -37,6 +39,9 @@ ) from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield +afd_async_dispatch, afd_async_recv = register_async_moe_ops() + + if TYPE_CHECKING: from afd_plugin.model_executor.models.deepseek_v2 import ( AFDDeepseekV2DecoderLayer, @@ -44,6 +49,26 @@ ) +def build_stage_slot_mapping( + slot_mapping: Mapping[str, torch.Tensor], + token_slice: slice, +) -> dict[str, torch.Tensor]: + """This ubatch stage's slice of every layer's KV slot mapping. + + KV-cache writes are indexed per token, so a stage's slot mapping is its + slice of the batch's. Leaving the full-batch mapping in place makes the + attention layer write this stage's rows into the whole batch's slots and + corrupt the cache. + + Only the non-sequence-parallel layout calls this. Under sequence + parallelism the stage slice is in global coordinates and does not index the + rank-local mapping, so that layout keeps the parent's mapping unsliced. + """ + return { + layer_name: mapping[token_slice] for layer_name, mapping in slot_mapping.items() + } + + def run_model_forward( model: AFDDeepseekV2Model, input_ids: torch.Tensor | None, @@ -119,12 +144,17 @@ def run_attention_gate_afd_forward( ) -> tuple[torch.Tensor, torch.Tensor | None]: """Run the Attention-side gate AFD path used by async CAM.""" - afd_connector = afd_metadata.connector - forward_context = get_forward_context() - stage_idx = afd_metadata.stage_idx + # Stage resolution (vLLM tracks the DBO ubatch by thread) lives inside the + # dispatch/receive ops, which the compiled graph treats as opaque. pending_ffn_recv = False - pending_dispatch_layout: CAMDispatchLayout | None = None pending_dispatch_ref: torch.Tensor | None = None + pending_dispatch_layout: CAMDispatchLayout | None = None + # Read the capability off the connector, not its type: a subclass of the + # GPU connector, or a future async one, would otherwise fall through to the + # CAM direct path and silently lose the graph split it asked for. CAM keeps + # the direct calls -- the ops carry neither router logits nor FlashComm1 + # token sharding, and the CAM protocol needs both. + use_ops = afd_metadata.connector.uses_opaque_moe_ops # Async CAM profile forwards are a distributed startup contract: every # Attention rank pairs CAM I/O with the FFN daemon to initialize resources. @@ -132,18 +162,27 @@ def run_attention_gate_afd_forward( islice(model.layers, model.start_layer, model.end_layer), ): if layer_offset > 0 and pending_ffn_recv: - if pending_dispatch_layout is None or pending_dispatch_ref is None: + if pending_dispatch_ref is None: raise RuntimeError("Async CAM receive is missing its dispatch layout") - local_ffn_output = afd_connector.recv_ffn_output( - ref_tensor=pending_dispatch_ref, - ubatch_idx=stage_idx, - ) - hidden_states = restore_cam_dispatch_output( - local_ffn_output, - pending_dispatch_layout, - ) + if use_ops: + # Opaque op: waits for the reply on the stream and restores the + # model layout. Tracing splits here; the payload it needs was + # stashed by the dispatch op one layer earlier. + hidden_states = afd_async_recv(pending_dispatch_ref) + else: + if pending_dispatch_layout is None: + raise RuntimeError( + "Async CAM receive is missing its dispatch layout", + ) + hidden_states = restore_cam_dispatch_output( + afd_metadata.connector.recv_ffn_output( + ref_tensor=pending_dispatch_ref, + ubatch_idx=afd_metadata.stage_idx, + ), + pending_dispatch_layout, + ) + pending_dispatch_layout = None pending_ffn_recv = False - pending_dispatch_layout = None pending_dispatch_ref = None if not layer.is_moe_layer: @@ -168,45 +207,62 @@ def run_attention_gate_afd_forward( llama_4_scaling, ) - dispatch_payload = prepare_cam_dispatch_payload( - hidden_states, - topk_weights, - topk_ids, - router_logits, - use_sequence_parallel=forward_context.flash_comm_v1_enabled, - ) - metadata = AFDTransferMetadata.create_attention_metadata( - layer_idx=layer.layer_idx, - stage_idx=stage_idx, - seq_len=int(dispatch_payload.hidden_states.shape[0]), - ) - context = AFDTransferContext(metadata=metadata) - afd_connector.send_attn_output( - dispatch_payload.hidden_states, - context, - topk_weights=dispatch_payload.topk_weights, - topk_ids=dispatch_payload.topk_ids, - router_logits=dispatch_payload.router_logits, - ) + if use_ops: + # Opaque op: builds the rank-local payload and sends it. Tracing + # splits here -- everything inside (NVSHMEM views, host caches, + # ctypes waits) is invisible to Dynamo and recorded once per + # capture. + hidden_states = afd_async_dispatch( + hidden_states, + topk_weights, + topk_ids, + layer.layer_idx, + ) + pending_dispatch_ref = hidden_states + else: + dispatch_payload = prepare_cam_dispatch_payload( + hidden_states, + topk_weights, + topk_ids, + router_logits, + use_sequence_parallel=get_forward_context().flash_comm_v1_enabled, + ) + afd_metadata.connector.send_attn_output( + dispatch_payload.hidden_states, + AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=layer.layer_idx, + stage_idx=afd_metadata.stage_idx, + seq_len=int(dispatch_payload.hidden_states.shape[0]), + ), + ), + topk_weights=dispatch_payload.topk_weights, + topk_ids=dispatch_payload.topk_ids, + router_logits=dispatch_payload.router_logits, + ) + pending_dispatch_layout = dispatch_payload.layout + pending_dispatch_ref = dispatch_payload.hidden_states pending_ffn_recv = True - pending_dispatch_layout = dispatch_payload.layout - pending_dispatch_ref = dispatch_payload.hidden_states hidden_states = maybe_apply_dbo_yield( hidden_states, role="attention", ) if pending_ffn_recv: - if pending_dispatch_layout is None or pending_dispatch_ref is None: + if pending_dispatch_ref is None: raise RuntimeError("Async CAM receive is missing its dispatch layout") - local_ffn_output = afd_connector.recv_ffn_output( - ref_tensor=pending_dispatch_ref, - ubatch_idx=stage_idx, - ) - hidden_states = restore_cam_dispatch_output( - local_ffn_output, - pending_dispatch_layout, - ) + if use_ops: + hidden_states = afd_async_recv(pending_dispatch_ref) + else: + if pending_dispatch_layout is None: + raise RuntimeError("Async CAM receive is missing its dispatch layout") + hidden_states = restore_cam_dispatch_output( + afd_metadata.connector.recv_ffn_output( + ref_tensor=pending_dispatch_ref, + ubatch_idx=afd_metadata.stage_idx, + ), + pending_dispatch_layout, + ) return hidden_states, residual @@ -222,7 +278,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 " @@ -327,6 +387,10 @@ def compute_stage_attention( else: stage_forward_context.num_tokens = int(stage.input_tokens) stage_forward_context.pad_size = 0 + stage_forward_context.slot_mapping = build_stage_slot_mapping( + forward_context.slot_mapping, + stage.token_slice, + ) expected_tokens = int(stage_hidden_states[stage_idx].shape[0]) log_async_moe_stage_attention( stage_idx, @@ -503,7 +567,14 @@ def _restore_async_moe_stage_state( (hidden_width, residual_width), dim=-1, ) - return hidden_states, residual + # Splitting the last dimension leaves two interleaved views. The next thing + # to touch them is the final norm, and the fused add-RMSNorm kernel requires + # contiguous inputs on both backends -- it aborts in the kernel rather than + # falling back. The copy is not avoidable by ordering: a split along the + # last dimension is never contiguous, so the alternative is not a cheaper + # copy but a correctness bug. It is one pass over the final hidden state per + # forward, not per layer. + return hidden_states.contiguous(), residual.contiguous() __all__ = [ diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4.py b/afd_plugin/model_executor/models/npu/deepseek_v4.py index 32a74a6b..b555f006 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v4.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v4.py @@ -22,7 +22,7 @@ from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.sequence import IntermediateTensors -from afd_plugin.config import AFD_ASYNC_CONNECTOR, parse_afd_config +from afd_plugin.config import AFD_ASYNC_NPU_CONNECTOR, parse_afd_config from afd_plugin.connectors import ( AFDExpertRoutingSpec, AFDF2ATransferPayload, @@ -315,7 +315,7 @@ class AFDDeepseekV4Model(native.DeepseekV4Model): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): afd_config = parse_afd_config(vllm_config, validate=False) if ( - afd_config.connector == AFD_ASYNC_CONNECTOR + afd_config.connector == AFD_ASYNC_NPU_CONNECTOR and not afd_config.compute_gate_on_attention ): raise ValueError( diff --git a/afd_plugin/v1/worker/dbo.py b/afd_plugin/v1/worker/dbo.py index d2b9bf71..a3ebd32f 100644 --- a/afd_plugin/v1/worker/dbo.py +++ b/afd_plugin/v1/worker/dbo.py @@ -2,13 +2,70 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project """Small DBO helpers used by AFD runtime/model wrappers.""" +from collections.abc import Callable + import torch from vllm.utils.torch_utils import direct_register_custom_op -from vllm.v1.worker.ubatching import dbo_enabled, dbo_yield +from vllm.v1.worker.ubatching import ( + dbo_current_ubatch_id, + dbo_enabled, + dbo_yield, +) + +# Resolve the Ascend yield once. This used to be imported inside the op body, +# which runs once per MoE layer: on a CUDA build the module is absent, Python +# does not cache a failed import, and so every call re-walked the import +# machinery. A profile of an Attention rank put that at 833us per call and 258ms +# of a 1403ms window -- the single largest host cost on the layer path, for an +# import that can never succeed there. +_ascend_dbo_enabled: Callable[[], bool] | None +_ascend_dbo_yield: Callable[[], None] | None +try: + from afd_plugin.v1.worker.npu.ubatching import ( + dbo_enabled as _ascend_dbo_enabled, + ) + from afd_plugin.v1.worker.npu.ubatching import ( + dbo_yield as _ascend_dbo_yield, + ) +except Exception: # noqa: BLE001 -- not an Ascend build, or a broken one + # ImportError is the ordinary case on a CUDA build. Anything else means the + # Ascend module is present but failed while importing, and letting that + # escape would take `import afd_plugin.v1.worker.dbo` down with it -- on + # every platform, for a yield that only the Ascend path ever calls. + _ascend_dbo_enabled = None + _ascend_dbo_yield = None _AFD_DBO_YIELD_OP_REGISTERED = False +def current_dbo_ubatch_id() -> int | None: + """Which ubatch this thread is running, or ``None`` when DBO is off. + + The forward context does not carry it. vLLM tracks the ubatch by thread + rather than on the context, so reading ``forward_context.ubatch_idx`` + silently finds nothing and both DBO halves look like stage 0 -- and so + claim the same connector stage, the same window slot, and the same flag, + where the second dispatch overwrites the first and the reply the first is + waiting for never comes. + """ + if not dbo_enabled(): + return None + # During Dynamo tracing this whole read is off limits: vLLM marks + # dbo_current_ubatch_id as skipped, so tracing it raises Unsupported + # before any real code runs. There is no ubatch to report while tracing + # anyway -- the caller keeps the stage it already had. + if torch.compiler.is_compiling(): + return None + try: + return int(dbo_current_ubatch_id()) + except KeyError: + # This thread registered no ubatch, yet some thread did -- the model + # is running on a side thread while the two ubatch threads hold the + # map. There is no ubatch to report; the caller keeps the stage it + # already had. + return None + + def maybe_apply_dbo_yield( tensor: torch.Tensor, *, @@ -50,27 +107,20 @@ def afd_manual_dbo_yield_fake(x: torch.Tensor) -> None: def _yield_if_dbo_enabled() -> None: - try: - from afd_plugin.v1.worker.npu.ubatching import ( - dbo_enabled as ascend_dbo_enabled, - ) - from afd_plugin.v1.worker.npu.ubatching import ( - dbo_yield as ascend_dbo_yield, - ) - except ImportError: - ascend_dbo_enabled = None - ascend_dbo_yield = None - if ( - ascend_dbo_enabled is not None - and ascend_dbo_yield is not None - and ascend_dbo_enabled() + _ascend_dbo_enabled is not None + and _ascend_dbo_yield is not None + and _ascend_dbo_enabled() ): - ascend_dbo_yield() + _ascend_dbo_yield() return if dbo_enabled(): dbo_yield() -__all__ = ["maybe_apply_dbo_yield", "register_dbo_yield_custom_op"] +__all__ = [ + "current_dbo_ubatch_id", + "maybe_apply_dbo_yield", + "register_dbo_yield_custom_op", +] diff --git a/afd_plugin/v1/worker/npu/attention_model_runner.py b/afd_plugin/v1/worker/npu/attention_model_runner.py index 15f53e9b..922914b9 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner.py @@ -75,7 +75,7 @@ stop_afd_npu_profiler, ) from afd_plugin.config import ( - AFD_ASYNC_CONNECTOR, + AFD_ASYNC_NPU_CONNECTOR, AFDConfig, parse_afd_config, ) @@ -146,7 +146,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/docs/design/module/connector_contracts.md b/docs/design/module/connector_contracts.md index 489f8a63..35655daa 100644 --- a/docs/design/module/connector_contracts.md +++ b/docs/design/module/connector_contracts.md @@ -61,6 +61,7 @@ depend on role worker implementations. | CUDA P2P | [`gpu/p2p.py`](../../../afd_plugin/connectors/gpu/p2p.py), [`topology.py`](../../../afd_plugin/distributed/topology.py) | [`test_p2p_connector.py`](../../../tests/unit/connectors/test_p2p_connector.py), [DeepSeek-V2-Lite E2E](../../../tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py) | | Ascend CAMP2P | [`npu/camp2p.py`](../../../afd_plugin/connectors/npu/camp2p.py) | [`test_camp2p_connector.py`](../../../tests/unit/connectors/test_camp2p_connector.py), [DeepSeek-V2-Lite E2E](../../../tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py) | | Ascend CAM async | [`npu/async_cam.py`](../../../afd_plugin/connectors/npu/async_cam.py) | [`test_async_cam_connector.py`](../../../tests/unit/connectors/test_async_cam_connector.py), [`test_async_cam_npu.py`](../../../tests/e2e/models/deepseek_v2_lite/test_async_cam_npu.py) | +| Async GPU connector | [`gpu/async_gpu.py`](../../../afd_plugin/connectors/gpu/async_gpu.py), [`async_topology.py`](../../../afd_plugin/connectors/async_topology.py) | [`test_async_gpu_connector.py`](../../../tests/unit/connectors/test_async_gpu_connector.py), [`async_gpu_connector_e2e.py`](../../../tests/e2e/async_gpu_connector_e2e.py), [`async_gpu_moe_equivalence.py`](../../../tests/e2e/async_gpu_moe_equivalence.py) | | NVSHMEM symmetric window | [`gpu/symm_window.py`](../../../afd_plugin/connectors/gpu/symm_window.py), [`gpu/nvshmem_rt.py`](../../../afd_plugin/connectors/gpu/nvshmem_rt.py), [`gpu/cuda_rt.py`](../../../afd_plugin/connectors/gpu/cuda_rt.py) | [`test_symm_window.py`](../../../tests/unit/connectors/gpu/test_symm_window.py) (slot layout and header codec, CPU-only) | | Process-group construction | [`afd_process_group.py`](../../../afd_plugin/distributed/afd_process_group.py) | Connector initialization tests plus platform E2E paths | @@ -110,6 +111,7 @@ synchronous NPU runtime requires both common and connector-local values to be | `P2pNcclAFDConnector` | CUDA | FFN ranks, then Attention ranks | `P2pNcclAFDControlPlane`; stage DP metadata over a separate NCCL group | `connector.control_plane is not None` | | `CAMP2pAFDConnector` | Ascend | FFN ranks, then Attention ranks | `CAMP2pAFDControlPlane`; stage DP metadata over Gloo plus HCCL data groups | `connector.control_plane is not None` | | `CAMAsyncAFDConnector` | Ascend | Attention ranks, then FFN ranks | `None`; routing/token metadata travels with CAM dispatch payloads | `connector.control_plane is None` | +| `GpuAsyncAFDConnector` | CUDA | Attention ranks, then FFN ranks | `None`; routing and token metadata travel in the symmetric window's slot header | `connector.control_plane is None` | The CUDA P2P mapping requires `num_attention_ranks >= num_ffn_ranks` and an integral A/F ratio. Each FFN rank @@ -118,6 +120,27 @@ requires at least as many Attention ranks as FFN ranks; its control and HCCL groups remain connector-owned. CAM async maps role ranks directly into an Attention-first world and distributes routed experts across FFN ranks. +The async GPU connector maps role ranks into an Attention-first world like CAM +async, and distributes routed experts across FFN ranks. Its connector-owned +configuration is `attn_ranks_per_dp`, `ring_depth`, `routed_cap_multiplier`, +`recv_poll_timeout_ms`, `async_moe_ubatching`, `async_moe_num_ubatches` and +`async_moe_split`; unknown fields are rejected. `attn_ranks_per_dp` must equal +the Attention deployment's `tensor_parallel_size`, and `async=true` plus +`compute_gate_on_attention=true` are structural requirements rather than +defaults -- `validate_afd_config` rejects the other combinations, because FFN +steps come off the connector receive loop and the wire carries topk chosen on +the Attention side. + +It is the only connector whose Attention-side data path is CUDA-graph +capturable. That is what fixes the flag protocol's shape: a captured stream wait +compares against a value baked in at capture time, so a reply signals with the +constant `FLAG_REPLY_READY` and the reader resets the flag in band once it has +consumed the slot, making every replay identical. The FFN side is driven by a +host poll loop and stays eager. Because only that constant releases a combine +wait, an FFN rank that leaves between a dispatch and its reply stamps +`FLAG_REPLY_READY` on every ring as it announces shutdown; an Attention rank is +expected to drain its pending combines before the FFN world closes. + These are current implementation facts, not approved long-term extension contracts. Individual connectors remain sections of this document until they have independent stable ownership. diff --git a/pyproject.toml b/pyproject.toml index 39fb96b6..24a74fa1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -97,6 +97,7 @@ select = [ "ISC", "SIM", ] +extend-ignore = ["N812"] [tool.ruff.lint.per-file-ignores] "afd_plugin/compat/patches/**/*.py" = [ diff --git a/tests/e2e/async_gpu_connector_e2e.py b/tests/e2e/async_gpu_connector_e2e.py new file mode 100644 index 00000000..0f57c6e3 --- /dev/null +++ b/tests/e2e/async_gpu_connector_e2e.py @@ -0,0 +1,335 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + +"""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. + +It runs twice over: first eagerly, then with the Attention half captured into a +CUDA graph and replayed. The replays are what pin the flag protocol down. A +replay runs no Python, so a dispatch sequence number kept on the host and a +stream wait told to expect one would both freeze at whatever was live during +capture -- the wait would then find a flag already holding the value it wants, +fall straight through, and combine the previous replay's data. That failure is +silent, so each replay is fed different tokens and checked against its own +reference, which is what turns it into an assertion. + +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 +# One eager pass on a side stream before capture, the usual prerequisite: the +# caching allocator and the window's cached views have to be warm, and a graph +# cannot be captured off a cold stream. +GRAPH_WARMUPS = 1 +GRAPH_REPLAYS = 4 +# Capture itself records kernels without running them, so it asks nothing of +# the FFN rank; only the eager layers, the warmup and the replays do. +FFN_WORK_ITEMS = NUM_LAYERS + GRAPH_WARMUPS + GRAPH_REPLAYS + + +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, + # This test's FFN loop runs routed experts only. + n_shared_experts=0, + ), + 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_graph_phase(connector, device, w13, w2) -> None: + """Capture one dispatch, then replay it against tokens it never saw. + + A graph replays out of the buffers it was captured with, so the tokens and + the routing are copied into fixed tensors rather than handed in as new + ones. Everything else is the same call pair the eager loop above makes. + """ + x = torch.zeros(NUM_TOKENS, HIDDEN, dtype=torch.bfloat16, device=device) + topk_ids = torch.zeros(NUM_TOKENS, TOPK, dtype=torch.int32, device=device) + topk_weights = torch.zeros(NUM_TOKENS, TOPK, dtype=torch.float32, device=device) + + def dispatch(): + context = AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=0, + stage_idx=0, + seq_len=NUM_TOKENS, + ), + ) + connector.send_attn_output( + x, + context, + topk_ids=topk_ids, + topk_weights=topk_weights, + ) + return connector.recv_ffn_output(ref_tensor=x, ubatch_idx=0) + + def load(iteration: int): + """Fill the captured buffers with one iteration's own tokens.""" + tokens, ids, weights = make_layer_inputs(NUM_LAYERS + iteration, device) + x.copy_(tokens) + topk_ids.copy_(ids) + topk_weights.copy_(weights) + return reference_moe(tokens, w13, w2, ids, weights) + + warmup_stream = torch.cuda.Stream(device=device) + warmup_stream.wait_stream(torch.cuda.current_stream(device)) + for iteration in range(GRAPH_WARMUPS): + expected = load(iteration) + with torch.cuda.stream(warmup_stream): + got = dispatch() + warmup_stream.synchronize() + torch.testing.assert_close( + got.to(torch.float32), + expected, + rtol=8e-2, + atol=8e-2, + ) + print(f"[A] graph warmup {iteration}: eager combine matches", flush=True) + torch.cuda.current_stream(device).wait_stream(warmup_stream) + torch.cuda.synchronize(device) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + combined = dispatch() + print("[A] captured the attention dispatch", flush=True) + + for replay in range(GRAPH_REPLAYS): + expected = load(GRAPH_WARMUPS + replay) + graph.replay() + torch.cuda.synchronize(device) + torch.testing.assert_close( + combined.to(torch.float32), + expected, + rtol=8e-2, + atol=8e-2, + ) + print(f"[A] replay {replay}: combine matches this replay's tokens", flush=True) + + +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 + ) + + run_graph_phase(connector, device, w13, w2) + + print("PASS: async GPU connector end-to-end, eager and CUDA graph", 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(FFN_WORK_ITEMS): + 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/e2e/async_gpu_moe_equivalence.py b/tests/e2e/async_gpu_moe_equivalence.py new file mode 100644 index 00000000..4422ab4b --- /dev/null +++ b/tests/e2e/async_gpu_moe_equivalence.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + +"""Numerical equivalence of the async GPU dispatch/compute/combine chain. + +Run with one GPU:: + + python tests/e2e/async_gpu_moe_equivalence.py + +The connector routes tokens itself and hands the FFN side pre-grouped rows, so +the expert compute runs as a topk==1 problem with unit weights and the real +topk weighting is applied during combine. This checks that the whole chain +reproduces a naive per-token MoE reference. +""" + +from types import SimpleNamespace + +import torch +import torch.nn.functional as F + +from afd_plugin.connectors.gpu.async_gpu import plan_dispatch +from afd_plugin.model_executor.models.gpu.deepseek_v2_attention_gate import ( + compute_attention_gate_moe_ffn, +) + +NUM_TOKENS = 32 +HIDDEN = 128 +INTERMEDIATE = 256 +TOPK = 4 +NUM_EXPERTS = 8 +FFN_SIZE = 2 +EXPERT_PER_RANK = NUM_EXPERTS // FFN_SIZE + + +def reference_moe(x, w13, w2, topk_ids, topk_weights): + """Naive per-token MoE: sum over each token's topk experts.""" + 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 + return out + + +def main() -> None: + torch.manual_seed(0) + device = torch.device("cuda", 0) + dtype = torch.bfloat16 + + x = torch.randn(NUM_TOKENS, HIDDEN, device=device, dtype=dtype) + w13 = ( + torch.randn( + NUM_EXPERTS, + 2 * INTERMEDIATE, + HIDDEN, + device=device, + dtype=dtype, + ) + / HIDDEN**0.5 + ) + w2 = ( + torch.randn( + NUM_EXPERTS, + HIDDEN, + INTERMEDIATE, + device=device, + dtype=dtype, + ) + / INTERMEDIATE**0.5 + ) + topk_ids = torch.stack( + [torch.randperm(NUM_EXPERTS)[:TOPK] for _ in range(NUM_TOKENS)], + ).to(device, torch.int32) + topk_weights = torch.rand(NUM_TOKENS, TOPK, device=device, dtype=torch.float32) + + expected = reference_moe(x, w13, w2, topk_ids, topk_weights) + + # --- what the connector + FFN runner actually do ----------------------- + plan = plan_dispatch( + topk_ids, + topk_weights, + ffn_size=FFN_SIZE, + expert_per_rank=EXPERT_PER_RANK, + ) + # A destination locates its own partials from the two header words the + # sender fills on the device; nothing about the routing is read back here, + # which is what the send path does now too. + starts = plan.segment_start.cpu().tolist() + routed = plan.routed_per_rank.cpu().tolist() + + accumulator = torch.zeros(NUM_TOKENS, HIDDEN, dtype=torch.float32, device=device) + for ffn_rank in range(FFN_SIZE): + base = ffn_rank * EXPERT_PER_RANK + group_list = plan.counts[base : base + EXPERT_PER_RANK] + partials = slice(starts[ffn_rank], starts[ffn_rank] + routed[ffn_rank]) + expand = plan.expand_idx[partials].to(torch.int64) + + # This FFN rank owns experts [base, base + EXPERT_PER_RANK). + layer = SimpleNamespace( + mlp=SimpleNamespace( + experts=SimpleNamespace( + routed_experts=SimpleNamespace( + w13_weight=w13[base : base + EXPERT_PER_RANK].contiguous(), + w2_weight=w2[base : base + EXPERT_PER_RANK].contiguous(), + ), + _shared_experts=None, + routed_scaling_factor=1.0, + ), + ), + ) + # The whole batch crosses the wire; the FFN side gathers one row per + # partial from it before the grouped GEMM, which applies the partial + # weights in its own epilogue. + payload = compute_attention_gate_moe_ffn( + layer, + hidden_states=x.index_select(0, expand), + group_list=group_list, + expand_x_shared=None, + ) + # ...then weights and reduces back to one row per token, in the payload + # dtype, leaving a zero row for tokens this rank held no expert for. + reduced = torch.zeros( + NUM_TOKENS, HIDDEN, dtype=payload.routed_output.dtype, device=device + ) + weighted = payload.routed_output * plan.weights[partials].unsqueeze(1).to( + payload.routed_output.dtype + ) + reduced.index_add_(0, expand, weighted) + accumulator += reduced.to(torch.float32) + + diff = (accumulator - expected).abs() + rel = diff.max() / expected.abs().max() + print(f"max abs diff={diff.max():.4e} max rel={rel:.4e}") + torch.testing.assert_close(accumulator, expected, rtol=6e-2, atol=6e-2) + print("PASS: dispatch -> local experts -> combine matches naive MoE") + + +if __name__ == "__main__": + main() diff --git a/tests/unit/config/test_config.py b/tests/unit/config/test_config.py index 7728bd8a..6c651a2e 100644 --- a/tests/unit/config/test_config.py +++ b/tests/unit/config/test_config.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + from __future__ import annotations from types import SimpleNamespace @@ -119,7 +122,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 +134,58 @@ def test_async_dp_requires_async_connector(): ) +@pytest.mark.parametrize( + ("connector", "extra"), + [ + ("CAMAsyncAFDConnector", {}), + # The GPU connector has no FFN-side router, so the gate is not optional. + ("GpuAsyncAFDConnector", {"compute_gate_on_attention": True}), + ], +) +def test_async_dp_accepts_every_async_connector(connector, extra): + config = parse_afd_config( + { + "afd": { + "connector": connector, + "role": "attention", + "async": True, + **extra, + }, + }, + ) + assert config.connector == connector + assert config.async_dp + + +@pytest.mark.parametrize( + ("afd", "expected"), + [ + ( + {"async": False, "compute_gate_on_attention": True}, + "requires async=true", + ), + ( + {"async": True, "compute_gate_on_attention": False}, + "requires compute_gate_on_attention=true", + ), + ], +) +def test_gpu_async_rejects_the_combinations_it_cannot_run(afd, expected): + # Both are structural: FFN steps come off the connector receive loop, which + # only the async-DP patches drive, and topk is chosen on the Attention side. + # Without this the failure is a startup hang or a missing-gate crash. + with pytest.raises(ValueError, match=expected): + parse_afd_config( + { + "afd": { + "connector": "GpuAsyncAFDConnector", + "role": "attention", + **afd, + }, + }, + ) + + def test_original_common_afd_field_aliases_are_supported(): raw = { "afd_role": "ffn", diff --git a/tests/unit/connectors/gpu/test_async_moe_op.py b/tests/unit/connectors/gpu/test_async_moe_op.py new file mode 100644 index 00000000..d14b4185 --- /dev/null +++ b/tests/unit/connectors/gpu/test_async_moe_op.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Tests for the async GPU connector's Dynamo-opaque dispatch/receive ops.""" + +from __future__ import annotations + +import torch +from torch._subclasses.fake_tensor import FakeTensorMode + +from afd_plugin.connectors.gpu.async_moe_op import ( + register_async_moe_ops, +) + + +def test_async_moe_ops_are_registered_once() -> None: + first = register_async_moe_ops() + second = register_async_moe_ops() + assert first == second + + +def test_async_moe_ops_fake_impls_shape() -> None: + dispatch, receive = register_async_moe_ops() + with FakeTensorMode(): + hidden = torch.randn(6, 16) + ids = torch.zeros(6, 2, dtype=torch.int64) + weights = torch.ones(6, 2) + sent = dispatch(hidden, weights, ids, 0) + out = receive(sent) + assert sent.shape == hidden.shape + assert out.shape == hidden.shape + + +def test_async_moe_ops_are_dynamo_opaque() -> None: + """Tracing must split at the ops, not reach into the connector. + + Before the ops existed, tracing the async MoE forward ran into the + connector's Python -- NVSHMEM pointer views, host caches, ctypes -- and + Dynamo raised Unsupported, first on a logger call. Exporting with fake + tensors asserts the tracer captures the dispatch/receive as opaque calls + with the deferred-receive control flow in between. The ops only register + CUDA kernels, so nothing here executes them. + """ + dispatch, receive = register_async_moe_ops() + + def proxy(hidden: torch.Tensor, ids: torch.Tensor, weights: torch.Tensor): + sent = dispatch(hidden, weights, ids, 3) + returned = receive(sent) + return dispatch(returned, weights, ids, 4) + + with FakeTensorMode(): + hidden = torch.randn(4, 8) + ids = torch.zeros(4, 2, dtype=torch.int64) + weights = torch.ones(4, 2) + graph_module, _ = torch._dynamo.export(proxy)(hidden, ids, weights) + + op_targets = [ + str(node.target) + for node in graph_module.graph.nodes + if node.op == "call_function" + ] + assert sum("afd_async_dispatch" in t for t in op_targets) == 2 + assert sum("afd_async_recv" in t for t in op_targets) == 1 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..9ec6dbcb --- /dev/null +++ b/tests/unit/connectors/test_async_gpu_connector.py @@ -0,0 +1,751 @@ +# 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 inspect +from functools import partial +from itertools import pairwise +from types import SimpleNamespace + +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 + FLAG_REPLY_READY, + GpuAsyncAFDConnector, + GpuAsyncExtraInfo, + GpuAsyncTransferState, + _PendingDispatch, + plan_dispatch, +) +from afd_plugin.connectors.gpu.symm_window import ( # noqa: E402 + HEADER_FIXED_WORDS, + HEADER_HOST_WORDS, + SlotLayout, + SymmWindow, + decode_header, +) +from afd_plugin.connectors.metadata import ( # noqa: E402 + AFDTransferContext, + AFDTransferMetadata, +) + + +@pytest.fixture +def layout() -> SlotLayout: + # shared_cap is deliberately not token_cap: the shared field is sized by the + # per-rank split, so a layout that quietly reused token_cap for it would + # otherwise still satisfy every assertion below. + return SlotLayout.build( + expert_per_rank=4, + partial_cap=100, + token_cap=32, + shared_cap=8, + 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}) + + +@pytest.mark.parametrize("ffn_size", [1, 2, 3, 4, 6]) +@pytest.mark.parametrize("num_tokens", [0, 1, 5, 7, 64, 513]) +def test_shared_split_tiles_the_batch_within_its_capacity( + ffn_size: int, + num_tokens: int, +): + # _shared_slice only reads these two attributes, so the bound can be checked + # without a device or a vLLM config behind it. + connector = SimpleNamespace(has_shared_experts=True, ffn_size=ffn_size) + slices = [ + GpuAsyncAFDConnector._shared_slice(connector, rank, num_tokens) + for rank in range(ffn_size) + ] + + # Every shared token is computed exactly once: the slices tile [0, n). + assert slices[0].start == 0 + assert slices[-1].stop == num_tokens + for earlier, later in pairwise(slices): + assert earlier.stop == later.start + + # And none of them can overflow the field the slot reserves for them, which + # is what lets shared_cap be a fraction of the batch rather than all of it. + shared_cap = -(-num_tokens // ffn_size) + assert max(s.stop - s.start for s in slices) <= shared_cap + + +def test_shared_split_is_empty_without_shared_experts(): + connector = SimpleNamespace(has_shared_experts=False, ffn_size=4) + for rank in range(4): + assert GpuAsyncAFDConnector._shared_slice(connector, rank, 64) == slice(0, 0) + + +def test_shutdown_announcement_matches_the_window_write_signature(layout: SlotLayout): + # Nothing calls announce_shutdown yet, so no runtime path would notice it + # passing a keyword write_slot does not take. Binding the arguments it + # actually sends against the real signature is what catches that. + calls: list[dict] = [] + window = SimpleNamespace(write_slot=lambda **kwargs: calls.append(kwargs)) + connector = SimpleNamespace( + _require_initialized=lambda: window, + attn_size=2, + ffn_size=3, + is_attention=True, + # A device tensor in the real thing; the CPU one behaves the same here + # and keeps the flag stamp on the same code path as a dispatch. + _seq_device=torch.zeros((), dtype=torch.int32), + layout=layout, + role_rank=1, + topk=6, + expert_per_rank=layout.header_words - HEADER_FIXED_WORDS, + ) + + GpuAsyncAFDConnector.announce_shutdown(connector) + + # One message per opposite-role peer, each carrying the shutdown bit. + assert len(calls) == 3 + signature = inspect.signature(SymmWindow.write_slot) + for kwargs in calls: + signature.bind(window, **kwargs) + assert decode_header(kwargs["header"]).is_shutdown + + +# ---------------------------------------------------------------------- +# 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 _destination_slices(plan, ffn_size, expert_per_rank): + """Walk the plan the way an FFN rank does: one run of partials each. + + A destination is told where its run starts and how long it is, and reads the + index arrays the sender shipped whole -- there is no per-destination slicing + on the send side any more, which is what removed the readback. + """ + routed = plan.routed_per_rank.tolist() + starts = plan.segment_start.tolist() + for ffn_rank in range(ffn_size): + yield ffn_rank, slice(starts[ffn_rank], starts[ffn_rank] + routed[ffn_rank]) + + +def test_every_partial_is_routed_exactly_once(routing_inputs): + topk_ids, _, topk_weights = routing_inputs + plan = plan_dispatch( + topk_ids, + topk_weights, + ffn_size=_FFN_SIZE, + expert_per_rank=_EXPERT_PER_RANK, + ) + assert plan.expand_idx.shape == (_NUM_TOKENS * _TOPK,) + assert plan.weights.shape == (_NUM_TOKENS * _TOPK,) + assert int(plan.counts.sum()) == _NUM_TOKENS * _TOPK + assert int(plan.routed_per_rank.sum()) == _NUM_TOKENS * _TOPK + # The runs must tile the array end to end, or a partial is read twice or not + # at all: they are the only thing a destination gets to locate itself by. + cursor = 0 + for _, partials in _destination_slices(plan, _FFN_SIZE, _EXPERT_PER_RANK): + assert partials.start == cursor + cursor = partials.stop + assert cursor == _NUM_TOKENS * _TOPK + + +def test_each_destination_segment_is_grouped_by_local_expert(routing_inputs): + topk_ids, _, topk_weights = routing_inputs + plan = plan_dispatch( + topk_ids, + topk_weights, + ffn_size=_FFN_SIZE, + expert_per_rank=_EXPERT_PER_RANK, + ) + counts = plan.counts.tolist() + for ffn_rank, partials in _destination_slices( + plan, + _FFN_SIZE, + _EXPERT_PER_RANK, + ): + base = ffn_rank * _EXPERT_PER_RANK + # The token behind each partial, in the order the receiver sees them. + tokens = plan.expand_idx[partials] + cursor = 0 + for local_expert in range(_EXPERT_PER_RANK): + for _ in range(counts[base + local_expert]): + token_idx = int(tokens[cursor]) + assert base + local_expert in topk_ids[token_idx].tolist() + cursor += 1 + assert cursor == partials.stop - partials.start + + +def test_every_partial_names_a_token_of_this_batch(routing_inputs): + topk_ids, _, topk_weights = routing_inputs + plan = plan_dispatch( + topk_ids, + topk_weights, + ffn_size=_FFN_SIZE, + expert_per_rank=_EXPERT_PER_RANK, + ) + # Destinations read the whole batch out of the slot, so an index is only in + # range if it names a row of it. + assert int(plan.expand_idx.min()) >= 0 + assert int(plan.expand_idx.max()) < _NUM_TOKENS + for ffn_rank, partials in _destination_slices(plan, _FFN_SIZE, _EXPERT_PER_RANK): + base = ffn_rank * _EXPERT_PER_RANK + expected = { + token + for token in range(_NUM_TOKENS) + for expert in topk_ids[token].tolist() + if base <= expert < base + _EXPERT_PER_RANK + } + assert set(plan.expand_idx[partials].tolist()) == expected + + +def test_identity_experts_recombine_to_the_weighted_sum(routing_inputs): + """The full chain: ship the batch, expand, weight, reduce, add.""" + topk_ids, hidden_states, topk_weights = routing_inputs + plan = plan_dispatch( + topk_ids, + topk_weights, + ffn_size=_FFN_SIZE, + expert_per_rank=_EXPERT_PER_RANK, + ) + accumulator = torch.zeros(_NUM_TOKENS, _HIDDEN, dtype=torch.float32) + for _, partials in _destination_slices(plan, _FFN_SIZE, _EXPERT_PER_RANK): + expand = plan.expand_idx[partials].to(torch.int64) + # What the FFN side does with the batch it was sent. + expanded = hidden_states.index_select(0, expand) + reduced = torch.zeros(_NUM_TOKENS, _HIDDEN, dtype=torch.float32) + reduced.index_add_(0, expand, expanded * plan.weights[partials].unsqueeze(1)) + # A reply is a whole batch, so combine adds it without an index. + accumulator += reduced + 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) + plan = plan_dispatch( + topk_ids, + torch.ones(3, 2), + ffn_size=ffn_size, + expert_per_rank=expert_per_rank, + ) + assert plan.counts.numel() == ffn_size * expert_per_rank + assert int(plan.counts.sum()) == topk_ids.numel() + assert int(plan.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) + plan = plan_dispatch( + topk_ids, + torch.ones(1, 3), + ffn_size=ffn_size, + expert_per_rank=expert_per_rank, + ) + slices = list(_destination_slices(plan, ffn_size, expert_per_rank)) + + _, busy_partials = slices[0] + assert busy_partials.stop - busy_partials.start == 3 + # One token, three of its partials: it is sent once, read three times. + assert plan.expand_idx[busy_partials].tolist() == [0, 0, 0] + + _, empty_partials = slices[1] + assert empty_partials.stop - empty_partials.start == 0 + + # Reducing an empty destination must be a no-op, not an error. + accumulator = torch.zeros(1, 4, dtype=torch.float32) + empty = plan.expand_idx[empty_partials].to(torch.int64) + accumulator.index_add_(0, empty, torch.zeros(0, 4)) + assert torch.count_nonzero(accumulator) == 0 + + +# ---------------------------------------------------------------------- +# Flag protocol +# +# The Attention side runs inside a CUDA graph, which records kernels once and +# replays them without running any Python. Two properties keep that correct and +# neither is visible from the eager path alone, so they are pinned here: +# a dispatch's sequence number must come off a device tensor the graph advances, +# and a reply must be waited for by a constant that the same graph resets. +# ---------------------------------------------------------------------- + + +class _RecordingWindow: + """Window stand-in that records the flag traffic in the order it is issued.""" + + def __init__(self, hidden_size: int, payload_dtype: torch.dtype) -> None: + self.hidden_size = hidden_size + self.payload_dtype = payload_dtype + self.calls: list[tuple] = [] + + def write_slot(self, **kwargs) -> None: + self.calls.append( + ("write_slot", kwargs["peer"], kwargs["flag_value"], kwargs["header"]), + ) + + def stream_wait(self, region: int, ring: int, value) -> None: + self.calls.append(("stream_wait", region, ring, value)) + + def clear_flag(self, region: int, ring: int) -> None: + self.calls.append(("clear_flag", region, ring)) + + def local_routed(self, region: int, ring: int, count: int) -> torch.Tensor: + self.calls.append(("local_routed", region, ring)) + return torch.zeros(count, self.hidden_size, dtype=self.payload_dtype) + + def local_shared(self, region: int, ring: int, count: int) -> torch.Tensor: + self.calls.append(("local_shared", region, ring)) + return torch.zeros(count, self.hidden_size, dtype=self.payload_dtype) + + +def _combining_connector(window, *, ffn_size: int, has_shared_experts: bool): + """Minimal stand-in carrying only what recv_ffn_output reads.""" + return SimpleNamespace( + _require_initialized=lambda: window, + _pending={}, + _free_rings={}, + has_shared_experts=has_shared_experts, + hidden_size=window.hidden_size, + payload_dtype=window.payload_dtype, + ffn_size=ffn_size, + role_rank=0, + ) + + +def _pending(*, ffn_size: int, num_tokens: int, ring: int) -> _PendingDispatch: + metadata = AFDTransferMetadata.create_attention_metadata( + layer_idx=0, + stage_idx=0, + seq_len=num_tokens, + ) + return _PendingDispatch( + context=AFDTransferContext(metadata=metadata), + shared_slices=[ + slice(r * num_tokens // ffn_size, (r + 1) * num_tokens // ffn_size) + for r in range(ffn_size) + ], + num_tokens=num_tokens, + ring=ring, + expected_ffn=list(range(ffn_size)), + ) + + +def test_combine_waits_on_the_constant_reply_marker(): + # A per-dispatch number here would be baked into the captured graph, and + # every later replay would find the flag already at or above it. + window = _RecordingWindow(hidden_size=8, payload_dtype=torch.float32) + connector = _combining_connector(window, ffn_size=2, has_shared_experts=False) + connector._pending[0] = [_pending(ffn_size=2, num_tokens=4, ring=1)] + + GpuAsyncAFDConnector.recv_ffn_output( + connector, + ref_tensor=torch.zeros(4, 8), + ubatch_idx=0, + ) + + waits = [call for call in window.calls if call[0] == "stream_wait"] + assert waits == [ + ("stream_wait", 0, 1, FLAG_REPLY_READY), + ("stream_wait", 1, 1, FLAG_REPLY_READY), + ] + + +@pytest.mark.parametrize("has_shared_experts", [False, True]) +def test_combine_clears_each_flag_only_after_it_has_read_the_slot(has_shared_experts): + # The reset re-arms the slot for the next dispatch. Issued before the reads + # it would race the peer's next reply; left out it would let the following + # replay fall straight through a flag that is still raised. + window = _RecordingWindow(hidden_size=8, payload_dtype=torch.float32) + connector = _combining_connector( + window, + ffn_size=2, + has_shared_experts=has_shared_experts, + ) + connector._pending[0] = [_pending(ffn_size=2, num_tokens=4, ring=0)] + + GpuAsyncAFDConnector.recv_ffn_output( + connector, + ref_tensor=torch.zeros(4, 8), + ubatch_idx=0, + ) + + for ffn_rank in range(2): + own = [call for call in window.calls if call[1] == ffn_rank] + assert own[0][0] == "stream_wait" + assert own[-1] == ("clear_flag", ffn_rank, 0) + assert "local_routed" in {call[0] for call in own} + if has_shared_experts: + assert "local_shared" in {call[0] for call in own} + + +def test_combine_releases_the_ring_it_consumed(): + window = _RecordingWindow(hidden_size=8, payload_dtype=torch.float32) + connector = _combining_connector(window, ffn_size=1, has_shared_experts=False) + connector._pending[0] = [_pending(ffn_size=1, num_tokens=4, ring=3)] + + GpuAsyncAFDConnector.recv_ffn_output( + connector, + ref_tensor=torch.zeros(4, 8), + ubatch_idx=0, + ) + + assert connector._free_rings[0] == [3] + + +def _replying_connector(window): + layout = SlotLayout.build( + expert_per_rank=2, + partial_cap=16, + token_cap=4, + shared_cap=0, + hidden_size=8, + payload_itemsize=4, + ) + return SimpleNamespace( + _require_initialized=lambda: window, + layout=layout, + role_rank=1, + topk=2, + hidden_size=8, + payload_dtype=torch.float32, + ) + + +def test_reply_stamps_the_constant_marker_not_a_sequence_number(): + window = _RecordingWindow(hidden_size=8, payload_dtype=torch.float32) + connector = _replying_connector(window) + states = GpuAsyncTransferState( + region=0, + ring=2, + src_role_rank=0, + layer_idx=0, + stage_idx=0, + num_tokens=4, + routed_tokens=2, + shared_tokens=0, + expand_idx=torch.tensor([0, 1]), + weights=torch.ones(2), + ) + + GpuAsyncAFDConnector.send_ffn_output( + connector, + torch.zeros(2, 8), + AFDTransferContext( + metadata=AFDTransferMetadata.create_ffn_metadata( + layer_idx=0, + stage_idx=0, + seq_lens=[2], + ), + states=states, + ), + ) + + assert [call[:3] for call in window.calls] == [ + ("write_slot", 0, FLAG_REPLY_READY), + ] + + +def test_reply_carries_no_header(): + # The Attention rank knows a reply's shape before it exists -- that is the + # premise of waiting on a stream instead of polling -- so it never reads + # one. Writing a header here was a copy per peer per MoE layer that nothing + # consumed. + window = _RecordingWindow(hidden_size=8, payload_dtype=torch.float32) + connector = _replying_connector(window) + states = GpuAsyncTransferState( + region=0, + ring=2, + src_role_rank=0, + layer_idx=0, + stage_idx=0, + num_tokens=4, + routed_tokens=2, + shared_tokens=0, + expand_idx=torch.tensor([0, 1]), + weights=torch.ones(2), + ) + + GpuAsyncAFDConnector.send_ffn_output( + connector, + torch.zeros(2, 8), + AFDTransferContext( + metadata=AFDTransferMetadata.create_ffn_metadata( + layer_idx=0, + stage_idx=0, + seq_lens=[2], + ), + states=states, + ), + ) + + assert [call[3] for call in window.calls] == [None] + # write_slot has to accept that, not just tolerate it by luck. + assert ( + inspect.signature(SymmWindow.write_slot).parameters["header"].annotation + == "torch.Tensor | None" + ) + + +# ---------------------------------------------------------------------- +# Dispatch header assembly +# +# The prefix is constant for a dispatch shape, so it belongs off the layer +# path: one buffer per shape, written once. Rebuilding it per dispatch cost a +# strided host-to-device copy every MoE layer, and inside a captured graph it +# cost a copy node that re-shipped that constant on every replay. +# ---------------------------------------------------------------------- + + +def _dispatch_plan(ffn_size: int, expert_per_rank: int, fill: int): + experts = ffn_size * expert_per_rank + return SimpleNamespace( + counts=torch.full((experts,), fill, dtype=torch.int32), + routed_per_rank=torch.full( + (ffn_size,), fill * expert_per_rank, dtype=torch.int32 + ), + segment_start=torch.arange(ffn_size, dtype=torch.int32), + expand_idx=torch.zeros(1, dtype=torch.int32), + weights=torch.zeros(1, dtype=torch.float32), + ) + + +def _dispatching_connector(layout: SlotLayout, *, ffn_size: int): + connector = SimpleNamespace( + window=SimpleNamespace(device=torch.device("cpu")), + layout=layout, + ffn_size=ffn_size, + expert_per_rank=layout.header_words - HEADER_FIXED_WORDS, + _header_device={}, + ) + # _headers_for reaches back through self for the per-shape buffer. + connector._headers_for_shape = partial( + GpuAsyncAFDConnector._headers_for_shape, + connector, + ) + return connector + + +def test_dispatch_headers_are_built_once_per_shape(layout: SlotLayout): + connector = _dispatching_connector(layout, ffn_size=2) + first = GpuAsyncAFDConnector._headers_for_shape( + connector, layer_idx=3, num_tokens=16 + ) + again = GpuAsyncAFDConnector._headers_for_shape( + connector, layer_idx=3, num_tokens=16 + ) + other_layer = GpuAsyncAFDConnector._headers_for_shape( + connector, layer_idx=4, num_tokens=16 + ) + other_size = GpuAsyncAFDConnector._headers_for_shape( + connector, layer_idx=3, num_tokens=32 + ) + + assert again is first, "same shape must reuse the buffer, not rebuild it" + assert other_layer is not first + assert other_size is not first + assert len(connector._header_device) == 3 + + decoded = decode_header(first[0].cpu()) + assert decoded.layer_idx == 3 + assert decoded.num_tokens == 16 + assert not decoded.is_shutdown + + +def test_dispatch_writes_only_the_routing_tail(layout: SlotLayout): + # The point of the per-shape buffer: a dispatch must leave the prefix + # alone. A sentinel there survives if -- and only if -- nothing rewrites it. + connector = _dispatching_connector(layout, ffn_size=2) + expert_per_rank = connector.expert_per_rank + headers = GpuAsyncAFDConnector._headers_for_shape( + connector, layer_idx=1, num_tokens=8 + ) + sentinel = torch.arange(HEADER_HOST_WORDS, dtype=torch.int32) + headers[:, :HEADER_HOST_WORDS] = sentinel + + for fill in (1, 2): + out = GpuAsyncAFDConnector._headers_for( + connector, + layer_idx=1, + num_tokens=8, + plan=_dispatch_plan(2, expert_per_rank, fill), + ) + assert out is headers + assert torch.equal(out[0, :HEADER_HOST_WORDS], sentinel) + assert torch.equal( + out[:, HEADER_FIXED_WORDS:], + torch.full((2, expert_per_rank), fill, dtype=torch.int32), + ) + + +def test_a_shape_first_seen_during_capture_is_refused(layout: SlotLayout, monkeypatch): + # Allocating there would come from the graph's private pool and the prefix + # copy would be recorded against a host buffer freed before the first + # replay -- silent corruption. Warmup runs the shapes capture runs. + connector = _dispatching_connector(layout, ffn_size=2) + connector.window = SimpleNamespace(device=torch.device("cuda", 0)) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + + with pytest.raises(RuntimeError, match="during CUDA graph capture"): + GpuAsyncAFDConnector._headers_for_shape(connector, layer_idx=0, num_tokens=8) + + +# ---------------------------------------------------------------------- +# Ring allocation across stages +# +# A ring names a window slot. Two stages sharing one means the second dispatch +# overwrites the first's payload and flag, and the reply the first waits for +# never arrives -- which deadlocked both DBO ubatch threads, one inside +# recv_ffn_output and one waiting to be yielded to. +# ---------------------------------------------------------------------- + + +def _stage_config(*, use_ubatching, num_ubatches=2, extra=None): + afd_raw = { + "connector": "GpuAsyncAFDConnector", + "role": "attention", + "num_attention_ranks": 1, + "num_ffn_ranks": 1, + "compute_gate_on_attention": True, + "connector_extra_config": extra or {}, + } + return SimpleNamespace( + additional_config={"afd": afd_raw}, + parallel_config=SimpleNamespace( + use_ubatching=use_ubatching, + num_ubatches=num_ubatches, + # attn_ranks_per_dp must equal this; both default to 1. + tensor_parallel_size=1, + ), + model_config=SimpleNamespace( + dtype=torch.bfloat16, + hf_config=SimpleNamespace( + hidden_size=16, + num_experts_per_tok=2, + n_routed_experts=4, + n_shared_experts=0, + ), + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=8), + ) + + +def _connector_for(config): + from afd_plugin.config import afd_config_from_mapping + + return GpuAsyncAFDConnector( + rank=0, + local_rank=0, + vllm_config=config, + afd_config=afd_config_from_mapping( + config.additional_config["afd"], + validate=False, + ), + role_rank=0, + ) + + +def test_vllm_ubatching_gets_a_ring_per_stage(): + # DBO drives two forwards at once and stamps each with its ubatch index, + # which arrives here as the stage. Counting only this connector's own + # splitter left both on ring 0. + connector = _connector_for(_stage_config(use_ubatching=True)) + + assert connector.num_stages == 2 + assert connector.ring_depth >= 2 + assert set(connector._rings_for_stage(0)).isdisjoint( + connector._rings_for_stage(1), + ) + + +def test_a_single_stage_still_needs_only_one_ring(): + connector = _connector_for(_stage_config(use_ubatching=False)) + + assert connector.num_stages == 1 + assert connector.ring_depth == 1 + + +def test_a_pinned_ring_depth_too_small_for_the_stages_is_refused(): + # Silently growing it would overrule a deliberate choice; deadlocking on it + # is worse. Fail at construction with the arithmetic in the message. + with pytest.raises(ValueError, match="cannot serve"): + _connector_for( + _stage_config(use_ubatching=True, extra={"ring_depth": 1}), + ) + + +def test_header_cache_survives_being_first_filled_in_inference_mode(): + """A cached header must stay writable after the call that created it. + + The cache outlives its creating call, and the first dispatch for a shape + can land inside vLLM's inference-mode forward. A tensor allocated there is + an inference tensor, and the next dispatch's write to the routing tail + raises "Inplace update to inference tensor outside InferenceMode" -- which + is what a compiled prefill hit, because AOT compilation moves the first + touch of each shape inside the compiled region. + """ + with torch.inference_mode(): + headers = torch.empty((2, 8), dtype=torch.int32) + assert headers.is_inference(), "guard premise: this is what we must avoid" + + with torch.inference_mode(), torch.inference_mode(False): + safe = torch.empty((2, 8), dtype=torch.int32) + assert not safe.is_inference() + # Writable afterwards, which is all the dispatch path needs. + safe[:, 0] = 1 + assert int(safe[0, 0]) == 1 diff --git a/tests/unit/model_executor/models/test_async_cam_stage_slot_mapping.py b/tests/unit/model_executor/models/test_async_cam_stage_slot_mapping.py new file mode 100644 index 00000000..a44ac2fb --- /dev/null +++ b/tests/unit/model_executor/models/test_async_cam_stage_slot_mapping.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Per-stage KV slot mapping for the async-CAM ubatch path. + +The bug this pins is silent: with the full-batch mapping left in place, every +stage writes its rows into the whole batch's KV slots, so the cache is corrupt +but nothing raises. +""" + +import torch + +from afd_plugin.model_executor.models.npu.deepseek_v2_async_cam_forward import ( + build_stage_slot_mapping, +) + + +def test_each_stage_gets_its_own_slice(): + slot_mapping = {"layer.0": torch.arange(8), "layer.1": torch.arange(8) + 100} + + first = build_stage_slot_mapping(slot_mapping, slice(0, 3)) + second = build_stage_slot_mapping(slot_mapping, slice(3, 8)) + + assert torch.equal(first["layer.0"], torch.tensor([0, 1, 2])) + assert torch.equal(second["layer.0"], torch.tensor([3, 4, 5, 6, 7])) + # Every layer is sliced, not just the first. + assert torch.equal(first["layer.1"], torch.tensor([100, 101, 102])) + assert torch.equal(second["layer.1"], torch.tensor([103, 104, 105, 106, 107])) + + +def test_stages_partition_the_batch_without_overlap(): + slot_mapping = {"layer.0": torch.arange(6)} + slices = [slice(0, 2), slice(2, 4), slice(4, 6)] + + rows = torch.cat( + [build_stage_slot_mapping(slot_mapping, s)["layer.0"] for s in slices], + ) + + # Concatenating the stages must reproduce the batch exactly: no row written + # twice, none dropped. + assert torch.equal(rows, slot_mapping["layer.0"]) + + +def test_the_parent_mapping_is_left_alone(): + slot_mapping = {"layer.0": torch.arange(4)} + + build_stage_slot_mapping(slot_mapping, slice(0, 2)) + + assert torch.equal(slot_mapping["layer.0"], torch.arange(4)) + assert list(slot_mapping) == ["layer.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 0a1d1f9a..7f34ca2b 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py @@ -11,7 +11,7 @@ pytest.importorskip("vllm") from torch import nn # noqa: E402 -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 @@ -266,7 +266,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 86ff567c..136aee08 100644 --- a/tests/unit/model_executor/models/test_forward_context.py +++ b/tests/unit/model_executor/models/test_forward_context.py @@ -247,6 +247,8 @@ def recv_ffn_output(ref_tensor, ubatch_idx): connector = SimpleNamespace( send_attn_output=send_attn_output, recv_ffn_output=recv_ffn_output, + # CAM keeps the direct calls; the opaque ops are the GPU connector's. + uses_opaque_moe_ops=False, ) afd_metadata = SimpleNamespace(connector=connector, stage_idx=0) diff --git a/tests/unit/v1/worker/test_dbo.py b/tests/unit/v1/worker/test_dbo.py index 6c6b972f..75657283 100644 --- a/tests/unit/v1/worker/test_dbo.py +++ b/tests/unit/v1/worker/test_dbo.py @@ -1,8 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + from __future__ import annotations import builtins -import sys -from types import SimpleNamespace import pytest @@ -14,7 +15,7 @@ def test_maybe_apply_dbo_yield_uses_custom_op(monkeypatch): - calls = [] + calls: list[object] = [] tensor = object() monkeypatch.setattr( @@ -84,22 +85,14 @@ def fail_on_ascend_import(name, globals=None, locals=None, fromlist=(), level=0) def test_dbo_yield_prefers_plugin_ascend_context(monkeypatch): calls = [] - monkeypatch.setitem( - sys.modules, - "afd_plugin.v1.worker.npu.ubatching", - SimpleNamespace( - dbo_enabled=lambda: True, - dbo_yield=lambda: calls.append("ascend"), - ), - ) - monkeypatch.setitem( - sys.modules, - "vllm.v1.worker.ubatching", - SimpleNamespace( - dbo_enabled=lambda: True, - dbo_yield=lambda: calls.append("vllm"), - ), - ) + # The Ascend yield is resolved once at import, so patch the resolved names + # rather than sys.modules: re-importing per call cost 833us of host time on + # a CUDA build, where the import can only ever fail. + monkeypatch.setattr(dbo, "_ascend_dbo_enabled", lambda: True) + monkeypatch.setattr(dbo, "_ascend_dbo_yield", lambda: calls.append("ascend")) + monkeypatch.setattr(dbo, "dbo_enabled", lambda: True) + monkeypatch.setattr(dbo, "dbo_yield", lambda: calls.append("vllm")) + dbo._yield_if_dbo_enabled() assert calls == ["ascend"] @@ -108,25 +101,31 @@ def test_dbo_yield_prefers_plugin_ascend_context(monkeypatch): def test_dbo_yield_falls_back_to_vllm_context(monkeypatch): calls = [] - monkeypatch.setitem( - sys.modules, - "afd_plugin.v1.worker.npu.ubatching", - SimpleNamespace( - dbo_enabled=lambda: False, - dbo_yield=lambda: calls.append("ascend"), - ), - ) - monkeypatch.setitem( - sys.modules, - "vllm.v1.worker.ubatching", - SimpleNamespace( - dbo_enabled=lambda: True, - dbo_yield=lambda: calls.append("vllm"), - ), - ) + monkeypatch.setattr(dbo, "_ascend_dbo_enabled", lambda: False) + monkeypatch.setattr(dbo, "_ascend_dbo_yield", lambda: calls.append("ascend")) monkeypatch.setattr(dbo, "dbo_enabled", lambda: True) monkeypatch.setattr(dbo, "dbo_yield", lambda: calls.append("vllm")) dbo._yield_if_dbo_enabled() assert calls == ["vllm"] + + +def test_ubatch_id_is_none_when_dbo_is_off(monkeypatch): + # Off DBO the caller keeps whatever stage it already had; returning 0 here + # would silently pin every dispatch to stage 0. + from afd_plugin.v1.worker import dbo as dbo_module + + monkeypatch.setattr(dbo_module, "dbo_enabled", lambda: False) + assert dbo_module.current_dbo_ubatch_id() is None + + +def test_ubatch_id_comes_from_the_thread_not_the_forward_context(monkeypatch): + # vLLM never sets ubatch_idx on the forward context, so reading it there + # made both DBO halves look like stage 0 -- one window slot for two + # concurrent dispatches, the second overwriting the first's flag. + from afd_plugin.v1.worker import dbo as dbo_module + + monkeypatch.setattr(dbo_module, "dbo_enabled", lambda: True) + monkeypatch.setattr(dbo_module, "dbo_current_ubatch_id", lambda: 1) + assert dbo_module.current_dbo_ubatch_id() == 1