From 0a3b219a79722980b146bd9edb7a12c3aa892c5d Mon Sep 17 00:00:00 2001 From: specture724 Date: Thu, 10 Sep 2026 16:18:47 +0800 Subject: [PATCH 1/2] feat: add the async GPU connector on the symmetric window GpuAsyncAFDConnector drives AFD's Attention-to-FFN traffic over the symmetric window with no control plane at all. The Attention side routes tokens itself and one-sided-writes the already-grouped rows into every FFN peer's slot; the FFN side learns the layer and the token counts from the arriving payload, so nothing has to be broadcast ahead of a step and the two roles never rendezvous per layer. `plan_dispatch` produces one row per (token, expert slot), sorted by global expert, so a destination's partials are one contiguous run already grouped for its grouped GEMM. Nothing about the routing is ever read back to the host: the per-destination counts go into the slot header on the device, which is what takes the last synchronize off the layer path. The dispatch sequence number lives on the device for the same reason the window's flag protocol uses a constant marker -- a replay runs no Python, so a host-side counter would freeze at whatever was live at capture. AFDAsyncTopology and build_async_topology move out of the NPU CAM connector into connectors/async_topology.py so both async connectors share one role-rank-to-world-rank mapping, and AFD_ASYNC_CONNECTOR becomes AFD_ASYNC_NPU_CONNECTOR now that there is more than one. Co-Authored-By: Claude Opus 5 Signed-off-by: specture724 --- afd_plugin/compat/npu/feature_validation.py | 6 +- afd_plugin/config.py | 23 +- afd_plugin/connectors/async_topology.py | 94 ++ afd_plugin/connectors/factory.py | 5 + afd_plugin/connectors/gpu/__init__.py | 3 +- afd_plugin/connectors/gpu/async_gpu.py | 1276 +++++++++++++++++ afd_plugin/connectors/gpu/async_moe_op.py | 166 +++ afd_plugin/connectors/npu/async_cam.py | 75 +- .../npu/deepseek_v2_async_cam_forward.py | 154 +- afd_plugin/v1/worker/dbo.py | 80 +- .../v1/worker/npu/attention_model_runner.py | 4 +- pyproject.toml | 1 + tests/e2e/async_gpu_connector_e2e.py | 335 +++++ tests/e2e/async_gpu_moe_equivalence.py | 142 ++ tests/unit/config/test_config.py | 23 +- .../unit/connectors/gpu/test_async_moe_op.py | 62 + .../connectors/test_async_gpu_connector.py | 749 ++++++++++ tests/unit/v1/worker/test_dbo.py | 69 +- 18 files changed, 3086 insertions(+), 181 deletions(-) create mode 100644 afd_plugin/connectors/async_topology.py create mode 100644 afd_plugin/connectors/gpu/async_gpu.py create mode 100644 afd_plugin/connectors/gpu/async_moe_op.py create mode 100644 tests/e2e/async_gpu_connector_e2e.py create mode 100644 tests/e2e/async_gpu_moe_equivalence.py create mode 100644 tests/unit/connectors/gpu/test_async_moe_op.py create mode 100644 tests/unit/connectors/test_async_gpu_connector.py diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index 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..c13a7f55 100644 --- a/afd_plugin/config.py +++ b/afd_plugin/config.py @@ -16,14 +16,22 @@ from vllm.config import VllmConfig AFD_ADDITIONAL_CONFIG_KEY: Final[str] = "afd" -AFD_ASYNC_CONNECTOR: Final[str] = "CAMAsyncAFDConnector" +AFD_ASYNC_NPU_CONNECTOR: Final[str] = "CAMAsyncAFDConnector" +AFD_ASYNC_GPU_CONNECTOR: Final[str] = "GpuAsyncAFDConnector" +# Connectors that drive FFN work from the connector receive loop instead of a DP +# metadata control plane, and therefore need the async-DP engine patches. The +# patches are platform-neutral; only the connector below them differs. +AFD_ASYNC_CONNECTORS: Final[frozenset[str]] = frozenset( + {AFD_ASYNC_NPU_CONNECTOR, AFD_ASYNC_GPU_CONNECTOR}, +) AFDRole = Literal["attention", "ffn"] SUPPORTED_AFD_ROLES: Final[tuple[str, ...]] = ("attention", "ffn") SUPPORTED_AFD_CONNECTORS: Final[tuple[str, ...]] = ( "P2pNcclAFDConnector", "CAMP2pAFDConnector", - AFD_ASYNC_CONNECTOR, + AFD_ASYNC_NPU_CONNECTOR, + AFD_ASYNC_GPU_CONNECTOR, ) _ALIASES: Final[dict[str, str]] = { @@ -283,7 +291,7 @@ def is_afd_async_dp(vllm_config: VllmConfig) -> bool: return ( config is not None and config.async_dp - and config.connector == AFD_ASYNC_CONNECTOR + and config.connector in AFD_ASYNC_CONNECTORS ) @@ -307,9 +315,10 @@ def validate_afd_config( "AFD connector must be one of " f"{SUPPORTED_AFD_CONNECTORS!r}, got {config.connector!r}", ) - if config.async_dp and config.connector != AFD_ASYNC_CONNECTOR: + if config.async_dp and config.connector not in AFD_ASYNC_CONNECTORS: raise ValueError( - "AFD async mode requires connector='CAMAsyncAFDConnector'", + "AFD async mode requires one of " + f"{sorted(AFD_ASYNC_CONNECTORS)!r}, got {config.connector!r}", ) if config.connector == "P2pNcclAFDConnector": from afd_plugin.distributed import validate_p2p_topology @@ -331,7 +340,9 @@ def validate_afd_config( __all__ = [ "AFDConfig", - "AFD_ASYNC_CONNECTOR", + "AFD_ASYNC_NPU_CONNECTOR", + "AFD_ASYNC_CONNECTORS", + "AFD_ASYNC_GPU_CONNECTOR", "afd_config_from_mapping", "AFD_ADDITIONAL_CONFIG_KEY", "AFDRole", diff --git a/afd_plugin/connectors/async_topology.py b/afd_plugin/connectors/async_topology.py new file mode 100644 index 00000000..20e267a3 --- /dev/null +++ b/afd_plugin/connectors/async_topology.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Rank layout shared by the asynchronous AFD connectors. + +Both async connectors -- Ascend CAM and CUDA NVSHMEM -- lay their world out +Attention-first and derive expert placement the same way. Keeping that here lets +the CUDA connector reuse it without importing a backend module. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from afd_plugin.config import AFDConfig + +ASYNC_MOE_REQUEST_SPLIT = "request" +ATTN_RANKS_PER_DP_CONFIG_KEY = "attn_ranks_per_dp" + + +@dataclass(frozen=True, slots=True) +class AFDAsyncTopology: + """Role-local and world rank information for one async participant.""" + + role: str + role_rank: int + world_rank: int + attn_size: int + ffn_size: int + expert_per_rank: int + + @property + def world_size(self) -> int: + """Return the total number of Attention and FFN ranks.""" + return self.attn_size + self.ffn_size + + +def build_async_topology( + afd_config: AFDConfig, + role_rank: int, + *, + num_routed_experts: int | None = None, +) -> AFDAsyncTopology: + """Validate role-local rank settings and derive the async world rank. + + The world is Attention-first: Attention role rank ``i`` maps to world rank + ``i`` and FFN role rank ``j`` maps to ``num_attention_ranks + j``. Routed + experts are distributed across FFN ranks using a ceiling division; + production model layouts should keep the routed-expert count divisible by + the FFN rank count. + """ + attn_size = afd_config.num_attention_ranks + ffn_size = afd_config.num_ffn_ranks + if attn_size <= 0 or ffn_size <= 0: + raise ValueError("AFD async topology sizes must be positive") + if role_rank < 0: + raise ValueError(f"AFD async role rank must be non-negative, got {role_rank}") + + if afd_config.role == "attention": + if role_rank >= attn_size: + raise ValueError( + "Attention role rank must be within attention size " + f"(rank={role_rank}, size={attn_size})", + ) + world_rank = role_rank + elif afd_config.role == "ffn": + if role_rank >= ffn_size: + raise ValueError( + "FFN role rank must be within FFN size " + f"(rank={role_rank}, size={ffn_size})", + ) + world_rank = attn_size + role_rank + else: + raise ValueError(f"unknown AFD role {afd_config.role!r}") + + expert_count = num_routed_experts or 1 + expert_per_rank = (expert_count + ffn_size - 1) // ffn_size + return AFDAsyncTopology( + role=afd_config.role, + role_rank=role_rank, + world_rank=world_rank, + attn_size=attn_size, + ffn_size=ffn_size, + expert_per_rank=expert_per_rank, + ) + + +__all__ = [ + "ASYNC_MOE_REQUEST_SPLIT", + "ATTN_RANKS_PER_DP_CONFIG_KEY", + "AFDAsyncTopology", + "build_async_topology", +] diff --git a/afd_plugin/connectors/factory.py b/afd_plugin/connectors/factory.py index 6e94e1ac..8a962d0c 100644 --- a/afd_plugin/connectors/factory.py +++ b/afd_plugin/connectors/factory.py @@ -100,6 +100,11 @@ def parse_connector_extra_info( "afd_plugin.connectors.npu.async_cam", "CAMAsyncAFDConnector", ) +AFDConnectorFactory.register_connector( + "GpuAsyncAFDConnector", + "afd_plugin.connectors.gpu.async_gpu", + "GpuAsyncAFDConnector", +) __all__ = ["AFDConnectorFactory"] diff --git a/afd_plugin/connectors/gpu/__init__.py b/afd_plugin/connectors/gpu/__init__.py index 62314259..77e73746 100644 --- a/afd_plugin/connectors/gpu/__init__.py +++ b/afd_plugin/connectors/gpu/__init__.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project """GPU-specific AFD connector implementations.""" +from afd_plugin.connectors.gpu.async_gpu import GpuAsyncAFDConnector from afd_plugin.connectors.gpu.p2p import P2pNcclAFDConnector -__all__ = ["P2pNcclAFDConnector"] +__all__ = ["GpuAsyncAFDConnector", "P2pNcclAFDConnector"] diff --git a/afd_plugin/connectors/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py new file mode 100644 index 00000000..6af62592 --- /dev/null +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -0,0 +1,1276 @@ +# 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 ``docs/design/rfc_async_gpu_connector.md``. Supported deployment requires +``async=true``, ``compute_gate_on_attention=true``, and a single node. +""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from typing import TYPE_CHECKING, Any, Final + +import torch +from torch import Tensor +from vllm.logger import init_logger + +from afd_plugin.config import AFDConfig +from afd_plugin.config_utils import ( + coerce_extra_bool, + coerce_extra_positive_int, + coerce_extra_str, +) +from afd_plugin.connectors.async_topology import ( + ASYNC_MOE_REQUEST_SPLIT, + build_async_topology, +) +from afd_plugin.connectors.base import AFDConnectorBase, ConnectorExtraInfo +from afd_plugin.connectors.gpu.symm_window import ( + FLAG_SHUTDOWN_BIT, + 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 + + # 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 + self.tp_size = self.extra_info.attn_ranks_per_dp + + self.topology = build_async_topology( + afd_config, + role_rank, + num_routed_experts=self.num_routed_experts, + ) + self.world_rank = self.topology.world_rank + self.attn_size = self.topology.attn_size + self.ffn_size = self.topology.ffn_size + self.expert_per_rank = self.topology.expert_per_rank + self.is_attention = afd_config.role == "attention" + + # 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 + 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, + ) + + +__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/npu/deepseek_v2_async_cam_forward.py b/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py index 383ea56e..82154008 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 @@ -24,6 +24,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 +38,22 @@ ) from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield +afd_async_dispatch, afd_async_recv = register_async_moe_ops() + + +def _dispatches_through_ops(connector: object) -> bool: + """Whether this connector's round trip goes through the opaque ops. + + The ops exist so Dynamo splits at the MoE round trip instead of tracing + into the connector, which is the GPU async connector's requirement. CAM + keeps the direct calls: the ops carry neither router logits nor FlashComm1 + token sharding, and the CAM protocol needs both. + """ + from afd_plugin.connectors.gpu.async_gpu import GpuAsyncAFDConnector + + return isinstance(connector, GpuAsyncAFDConnector) + + if TYPE_CHECKING: from afd_plugin.model_executor.models.deepseek_v2 import ( AFDDeepseekV2DecoderLayer, @@ -119,12 +136,12 @@ 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 + use_ops = _dispatches_through_ops(afd_metadata.connector) # 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 +149,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 +194,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 +265,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 +374,16 @@ def compute_stage_attention( else: stage_forward_context.num_tokens = int(stage.input_tokens) stage_forward_context.pad_size = 0 + # 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. 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. + stage_forward_context.slot_mapping = { + layer_name: mapping[stage.token_slice] + for layer_name, mapping in forward_context.slot_mapping.items() + } expected_tokens = int(stage_hidden_states[stage_idx].shape[0]) log_async_moe_stage_attention( stage_idx, @@ -503,7 +560,10 @@ 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 CUDA's fused_add_rms_norm requires + # contiguous inputs -- it aborts in the kernel rather than falling back. + return hidden_states.contiguous(), residual.contiguous() __all__ = [ diff --git a/afd_plugin/v1/worker/dbo.py b/afd_plugin/v1/worker/dbo.py index d2b9bf71..885988e8 100644 --- a/afd_plugin/v1/worker/dbo.py +++ b/afd_plugin/v1/worker/dbo.py @@ -2,13 +2,66 @@ # 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 ImportError: # not an Ascend build + _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 +103,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/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..0f1effbd --- /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.model_executor.models.gpu.deepseek_v2_attention_gate import ( + compute_attention_gate_moe_ffn, +) + +from afd_plugin.config import AFDConfig +from afd_plugin.connectors.gpu.async_gpu import GpuAsyncAFDConnector +from afd_plugin.connectors.metadata import AFDTransferContext, AFDTransferMetadata + +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..f29742c0 --- /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.model_executor.models.gpu.deepseek_v2_attention_gate import ( + compute_attention_gate_moe_ffn, +) + +from afd_plugin.connectors.gpu.async_gpu import plan_dispatch + +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..6bbd3f79 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,24 @@ def test_async_dp_requires_async_connector(): ) +@pytest.mark.parametrize( + "connector", + ["CAMAsyncAFDConnector", "GpuAsyncAFDConnector"], +) +def test_async_dp_accepts_every_async_connector(connector): + config = parse_afd_config( + { + "afd": { + "connector": connector, + "role": "attention", + "async": True, + }, + }, + ) + assert config.connector == connector + assert config.async_dp + + def test_original_common_afd_field_aliases_are_supported(): raw = { "afd_role": "ffn", diff --git a/tests/unit/connectors/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..053d9dcf --- /dev/null +++ b/tests/unit/connectors/test_async_gpu_connector.py @@ -0,0 +1,749 @@ +# 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, + ), + 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/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 From eb93999845b4ba60b01a56eb600a199f909e3404 Mon Sep 17 00:00:00 2001 From: specture724 Date: Thu, 17 Sep 2026 11:11:57 +0800 Subject: [PATCH 2/2] fix: address review findings on the async GPU connector Review follow-up on #333. [P1] Both E2E scripts imported models.gpu.deepseek_v2_attention_gate, which first existed in #334, so this PR shipped its only runnable proof of the flag protocol as dead code. Move that module here instead: it is 135 lines, has no V4 reference, and is the Attention-side gate the connector's own E2E needs. Also fixes a rename this PR missed. It renamed AFD_ASYNC_CONNECTOR to AFD_ASYNC_NPU_CONNECTOR in config.py but left three call sites importing the old name, so `import afd_plugin.model_executor.models.deepseek_v2` raised ImportError and took eight test modules down with it. The deepseek_v2 call site wanted the set, not the NPU constant: both async connectors run that schedule. Remaining findings: - Branch on a capability, not a type. `isinstance(connector, GpuAsyncAFDConnector)` routed any subclass or future async connector down the CAM direct path; connectors now declare `uses_opaque_moe_ops`, which deletes the helper that wrapped the check. - Give the combine a failure path. Only the constant FLAG_REPLY_READY releases a combine wait, so an FFN rank leaving between a dispatch and its reply stranded Attention forever, possibly inside a graph replay where no Python runs. Shutdown now stamps the flag on every ring, which is safe to repeat because the value is constant, and close() names any un-drained combines rather than clearing them in silence. - Validate the connector's requirements at config time. async=true and compute_gate_on_attention=true are structural, and attn_ranks_per_dp must equal the Attention tensor_parallel_size. - Name and test the NPU ride-alongs. The stage-wise slot_mapping slicing is a real KV-cache fix: extracted as build_stage_slot_mapping and pinned by tests. The .contiguous() rationale now says why the copy is not avoidable, and the Ascend yield import no longer lets a non-ImportError kill dbo.py on every platform. - Docs: connector_contracts.md gains the async-GPU mode row, its configuration and graph-capture paragraphs, and an evidence row. The docstring cited an RFC that is not in the tree; it now cites that page. Unit suite: 2 pre-existing failures (os.pidfd_open missing on this build), down from 13. pre-commit clean over the changed files. Co-Authored-By: Claude Opus 5 --- afd_plugin/config.py | 14 ++ afd_plugin/connectors/base.py | 7 + afd_plugin/connectors/gpu/async_gpu.py | 71 ++++++++- .../model_executor/models/deepseek_v2.py | 13 +- .../model_executor/models/gpu/__init__.py | 3 + .../models/gpu/deepseek_v2_attention_gate.py | 135 ++++++++++++++++++ .../npu/deepseek_v2_async_cam_forward.py | 63 ++++---- .../model_executor/models/npu/deepseek_v4.py | 4 +- afd_plugin/v1/worker/dbo.py | 6 +- docs/design/module/connector_contracts.md | 23 +++ tests/e2e/async_gpu_connector_e2e.py | 6 +- tests/e2e/async_gpu_moe_equivalence.py | 4 +- tests/unit/config/test_config.py | 40 +++++- .../connectors/test_async_gpu_connector.py | 2 + .../test_async_cam_stage_slot_mapping.py | 49 +++++++ .../models/test_deepseek_v2_proxy.py | 4 +- .../models/test_forward_context.py | 2 + 17 files changed, 401 insertions(+), 45 deletions(-) create mode 100644 afd_plugin/model_executor/models/gpu/__init__.py create mode 100644 afd_plugin/model_executor/models/gpu/deepseek_v2_attention_gate.py create mode 100644 tests/unit/model_executor/models/test_async_cam_stage_slot_mapping.py diff --git a/afd_plugin/config.py b/afd_plugin/config.py index c13a7f55..c16c0e80 100644 --- a/afd_plugin/config.py +++ b/afd_plugin/config.py @@ -320,6 +320,20 @@ def validate_afd_config( "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 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/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py index 6af62592..3444bbd3 100644 --- a/afd_plugin/connectors/gpu/async_gpu.py +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -24,14 +24,16 @@ live when it was recorded. The FFN side is driven by a host poll loop and stays eager, so nothing there needs capturing. -See ``docs/design/rfc_async_gpu_connector.md``. Supported deployment requires -``async=true``, ``compute_gate_on_attention=true``, and a single node. +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 Mapping +from collections.abc import Iterable, Mapping from dataclasses import dataclass from datetime import timedelta from typing import TYPE_CHECKING, Any, Final @@ -375,6 +377,10 @@ 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. @@ -406,7 +412,20 @@ def __init__( 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, @@ -573,6 +592,17 @@ def close(self) -> None: 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() @@ -1262,6 +1292,41 @@ def announce_shutdown(self) -> 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__ = [ 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 82154008..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 @@ -41,19 +42,6 @@ afd_async_dispatch, afd_async_recv = register_async_moe_ops() -def _dispatches_through_ops(connector: object) -> bool: - """Whether this connector's round trip goes through the opaque ops. - - The ops exist so Dynamo splits at the MoE round trip instead of tracing - into the connector, which is the GPU async connector's requirement. CAM - keeps the direct calls: the ops carry neither router logits nor FlashComm1 - token sharding, and the CAM protocol needs both. - """ - from afd_plugin.connectors.gpu.async_gpu import GpuAsyncAFDConnector - - return isinstance(connector, GpuAsyncAFDConnector) - - if TYPE_CHECKING: from afd_plugin.model_executor.models.deepseek_v2 import ( AFDDeepseekV2DecoderLayer, @@ -61,6 +49,26 @@ def _dispatches_through_ops(connector: object) -> bool: ) +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, @@ -141,7 +149,12 @@ def run_attention_gate_afd_forward( pending_ffn_recv = False pending_dispatch_ref: torch.Tensor | None = None pending_dispatch_layout: CAMDispatchLayout | None = None - use_ops = _dispatches_through_ops(afd_metadata.connector) + # 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. @@ -374,16 +387,10 @@ def compute_stage_attention( else: stage_forward_context.num_tokens = int(stage.input_tokens) stage_forward_context.pad_size = 0 - # 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. 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. - stage_forward_context.slot_mapping = { - layer_name: mapping[stage.token_slice] - for layer_name, mapping in forward_context.slot_mapping.items() - } + 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, @@ -561,8 +568,12 @@ def _restore_async_moe_stage_state( dim=-1, ) # Splitting the last dimension leaves two interleaved views. The next thing - # to touch them is the final norm, and CUDA's fused_add_rms_norm requires - # contiguous inputs -- it aborts in the kernel rather than falling back. + # 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() 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 885988e8..a3ebd32f 100644 --- a/afd_plugin/v1/worker/dbo.py +++ b/afd_plugin/v1/worker/dbo.py @@ -27,7 +27,11 @@ from afd_plugin.v1.worker.npu.ubatching import ( dbo_yield as _ascend_dbo_yield, ) -except ImportError: # not an Ascend build +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 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/tests/e2e/async_gpu_connector_e2e.py b/tests/e2e/async_gpu_connector_e2e.py index 0f1effbd..0f57c6e3 100644 --- a/tests/e2e/async_gpu_connector_e2e.py +++ b/tests/e2e/async_gpu_connector_e2e.py @@ -36,13 +36,13 @@ import torch import torch.distributed as dist import torch.nn.functional as F -from afd_plugin.model_executor.models.gpu.deepseek_v2_attention_gate import ( - compute_attention_gate_moe_ffn, -) 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 diff --git a/tests/e2e/async_gpu_moe_equivalence.py b/tests/e2e/async_gpu_moe_equivalence.py index f29742c0..4422ab4b 100644 --- a/tests/e2e/async_gpu_moe_equivalence.py +++ b/tests/e2e/async_gpu_moe_equivalence.py @@ -17,12 +17,12 @@ 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, ) -from afd_plugin.connectors.gpu.async_gpu import plan_dispatch - NUM_TOKENS = 32 HIDDEN = 128 INTERMEDIATE = 256 diff --git a/tests/unit/config/test_config.py b/tests/unit/config/test_config.py index 6bbd3f79..6c651a2e 100644 --- a/tests/unit/config/test_config.py +++ b/tests/unit/config/test_config.py @@ -135,16 +135,21 @@ def test_async_dp_requires_async_connector(): @pytest.mark.parametrize( - "connector", - ["CAMAsyncAFDConnector", "GpuAsyncAFDConnector"], + ("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): +def test_async_dp_accepts_every_async_connector(connector, extra): config = parse_afd_config( { "afd": { "connector": connector, "role": "attention", "async": True, + **extra, }, }, ) @@ -152,6 +157,35 @@ def test_async_dp_accepts_every_async_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/test_async_gpu_connector.py b/tests/unit/connectors/test_async_gpu_connector.py index 053d9dcf..9ec6dbcb 100644 --- a/tests/unit/connectors/test_async_gpu_connector.py +++ b/tests/unit/connectors/test_async_gpu_connector.py @@ -669,6 +669,8 @@ def _stage_config(*, use_ubatching, num_ubatches=2, extra=None): 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, 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)