From 896f30ba159658e1613adafb61808f005518a34c Mon Sep 17 00:00:00 2001 From: specture724 Date: Thu, 13 Aug 2026 12:40:37 +0800 Subject: [PATCH 01/18] init: first impliment for async GPU connector Signed-off-by: specture724 --- afd_plugin/compat/npu/feature_validation.py | 4 +- afd_plugin/config.py | 23 +- afd_plugin/connectors/async_topology.py | 94 ++ afd_plugin/connectors/factory.py | 5 + afd_plugin/connectors/gpu/__init__.py | 3 +- afd_plugin/connectors/gpu/async_gpu.py | 889 ++++++++++++++++++ afd_plugin/connectors/gpu/nvshmem_rt.py | 270 ++++++ afd_plugin/connectors/gpu/symm_window.py | 528 +++++++++++ afd_plugin/connectors/npu/async_cam.py | 75 +- .../model_executor/models/deepseek_v2.py | 29 +- .../model_executor/models/gpu/__init__.py | 3 + .../models/gpu/deepseek_v2_attention_gate.py | 126 +++ .../v1/worker/attention_model_runner.py | 84 +- afd_plugin/v1/worker/ffn_model_runner.py | 77 +- afd_plugin/v1/worker/ffn_worker.py | 24 +- .../v1/worker/npu/attention_model_runner.py | 4 +- pyproject.toml | 1 + .../deepseek_v2_lite/1a1f_eager_async.sh | 128 +++ .../deepseek_v2_lite/2a2f_eager_async.sh | 129 +++ tests/e2e/async_gpu_connector_e2e.py | 241 +++++ tests/e2e/async_gpu_moe_equivalence.py | 121 +++ tests/e2e/async_gpu_window_roundtrip.py | 179 ++++ tests/unit/config/test_config.py | 20 +- .../connectors/test_async_gpu_connector.py | 314 +++++++ .../models/test_deepseek_v2_proxy.py | 4 +- .../models/test_forward_context.py | 7 +- .../v1/worker/test_attention_model_runner.py | 41 +- tests/unit/v1/worker/test_ffn_model_runner.py | 34 +- 28 files changed, 3324 insertions(+), 133 deletions(-) create mode 100644 afd_plugin/connectors/async_topology.py create mode 100644 afd_plugin/connectors/gpu/async_gpu.py create mode 100644 afd_plugin/connectors/gpu/nvshmem_rt.py create mode 100644 afd_plugin/connectors/gpu/symm_window.py create mode 100644 afd_plugin/model_executor/models/gpu/__init__.py create mode 100644 afd_plugin/model_executor/models/gpu/deepseek_v2_attention_gate.py create mode 100644 recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/1a1f_eager_async.sh create mode 100644 recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/2a2f_eager_async.sh create mode 100644 tests/e2e/async_gpu_connector_e2e.py create mode 100644 tests/e2e/async_gpu_moe_equivalence.py create mode 100644 tests/e2e/async_gpu_window_roundtrip.py create mode 100644 tests/unit/connectors/test_async_gpu_connector.py diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index 823144b7..a07f3c4e 100644 --- a/afd_plugin/compat/npu/feature_validation.py +++ b/afd_plugin/compat/npu/feature_validation.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING from afd_plugin.config import ( - AFD_ASYNC_CONNECTOR, + AFD_ASYNC_NPU_CONNECTOR, AFDConfig, is_afd_async_dp, parse_afd_config, @@ -34,7 +34,7 @@ def fail_if_unsupported_npu_afd_features( vllm_config, ) - if afd_config.connector == AFD_ASYNC_CONNECTOR: + if afd_config.connector == AFD_ASYNC_NPU_CONNECTOR: _fail_if_unsupported_npu_afd_async_features( vllm_config, afd_config, diff --git a/afd_plugin/config.py b/afd_plugin/config.py index 256590ab..c13a7f55 100644 --- a/afd_plugin/config.py +++ b/afd_plugin/config.py @@ -16,14 +16,22 @@ from vllm.config import VllmConfig AFD_ADDITIONAL_CONFIG_KEY: Final[str] = "afd" -AFD_ASYNC_CONNECTOR: Final[str] = "CAMAsyncAFDConnector" +AFD_ASYNC_NPU_CONNECTOR: Final[str] = "CAMAsyncAFDConnector" +AFD_ASYNC_GPU_CONNECTOR: Final[str] = "GpuAsyncAFDConnector" +# Connectors that drive FFN work from the connector receive loop instead of a DP +# metadata control plane, and therefore need the async-DP engine patches. The +# patches are platform-neutral; only the connector below them differs. +AFD_ASYNC_CONNECTORS: Final[frozenset[str]] = frozenset( + {AFD_ASYNC_NPU_CONNECTOR, AFD_ASYNC_GPU_CONNECTOR}, +) AFDRole = Literal["attention", "ffn"] SUPPORTED_AFD_ROLES: Final[tuple[str, ...]] = ("attention", "ffn") SUPPORTED_AFD_CONNECTORS: Final[tuple[str, ...]] = ( "P2pNcclAFDConnector", "CAMP2pAFDConnector", - AFD_ASYNC_CONNECTOR, + AFD_ASYNC_NPU_CONNECTOR, + AFD_ASYNC_GPU_CONNECTOR, ) _ALIASES: Final[dict[str, str]] = { @@ -283,7 +291,7 @@ def is_afd_async_dp(vllm_config: VllmConfig) -> bool: return ( config is not None and config.async_dp - and config.connector == AFD_ASYNC_CONNECTOR + and config.connector in AFD_ASYNC_CONNECTORS ) @@ -307,9 +315,10 @@ def validate_afd_config( "AFD connector must be one of " f"{SUPPORTED_AFD_CONNECTORS!r}, got {config.connector!r}", ) - if config.async_dp and config.connector != AFD_ASYNC_CONNECTOR: + if config.async_dp and config.connector not in AFD_ASYNC_CONNECTORS: raise ValueError( - "AFD async mode requires connector='CAMAsyncAFDConnector'", + "AFD async mode requires one of " + f"{sorted(AFD_ASYNC_CONNECTORS)!r}, got {config.connector!r}", ) if config.connector == "P2pNcclAFDConnector": from afd_plugin.distributed import validate_p2p_topology @@ -331,7 +340,9 @@ def validate_afd_config( __all__ = [ "AFDConfig", - "AFD_ASYNC_CONNECTOR", + "AFD_ASYNC_NPU_CONNECTOR", + "AFD_ASYNC_CONNECTORS", + "AFD_ASYNC_GPU_CONNECTOR", "afd_config_from_mapping", "AFD_ADDITIONAL_CONFIG_KEY", "AFDRole", diff --git a/afd_plugin/connectors/async_topology.py b/afd_plugin/connectors/async_topology.py new file mode 100644 index 00000000..20e267a3 --- /dev/null +++ b/afd_plugin/connectors/async_topology.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Rank layout shared by the asynchronous AFD connectors. + +Both async connectors -- Ascend CAM and CUDA NVSHMEM -- lay their world out +Attention-first and derive expert placement the same way. Keeping that here lets +the CUDA connector reuse it without importing a backend module. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from afd_plugin.config import AFDConfig + +ASYNC_MOE_REQUEST_SPLIT = "request" +ATTN_RANKS_PER_DP_CONFIG_KEY = "attn_ranks_per_dp" + + +@dataclass(frozen=True, slots=True) +class AFDAsyncTopology: + """Role-local and world rank information for one async participant.""" + + role: str + role_rank: int + world_rank: int + attn_size: int + ffn_size: int + expert_per_rank: int + + @property + def world_size(self) -> int: + """Return the total number of Attention and FFN ranks.""" + return self.attn_size + self.ffn_size + + +def build_async_topology( + afd_config: AFDConfig, + role_rank: int, + *, + num_routed_experts: int | None = None, +) -> AFDAsyncTopology: + """Validate role-local rank settings and derive the async world rank. + + The world is Attention-first: Attention role rank ``i`` maps to world rank + ``i`` and FFN role rank ``j`` maps to ``num_attention_ranks + j``. Routed + experts are distributed across FFN ranks using a ceiling division; + production model layouts should keep the routed-expert count divisible by + the FFN rank count. + """ + attn_size = afd_config.num_attention_ranks + ffn_size = afd_config.num_ffn_ranks + if attn_size <= 0 or ffn_size <= 0: + raise ValueError("AFD async topology sizes must be positive") + if role_rank < 0: + raise ValueError(f"AFD async role rank must be non-negative, got {role_rank}") + + if afd_config.role == "attention": + if role_rank >= attn_size: + raise ValueError( + "Attention role rank must be within attention size " + f"(rank={role_rank}, size={attn_size})", + ) + world_rank = role_rank + elif afd_config.role == "ffn": + if role_rank >= ffn_size: + raise ValueError( + "FFN role rank must be within FFN size " + f"(rank={role_rank}, size={ffn_size})", + ) + world_rank = attn_size + role_rank + else: + raise ValueError(f"unknown AFD role {afd_config.role!r}") + + expert_count = num_routed_experts or 1 + expert_per_rank = (expert_count + ffn_size - 1) // ffn_size + return AFDAsyncTopology( + role=afd_config.role, + role_rank=role_rank, + world_rank=world_rank, + attn_size=attn_size, + ffn_size=ffn_size, + expert_per_rank=expert_per_rank, + ) + + +__all__ = [ + "ASYNC_MOE_REQUEST_SPLIT", + "ATTN_RANKS_PER_DP_CONFIG_KEY", + "AFDAsyncTopology", + "build_async_topology", +] diff --git a/afd_plugin/connectors/factory.py b/afd_plugin/connectors/factory.py index 6e94e1ac..8a962d0c 100644 --- a/afd_plugin/connectors/factory.py +++ b/afd_plugin/connectors/factory.py @@ -100,6 +100,11 @@ def parse_connector_extra_info( "afd_plugin.connectors.npu.async_cam", "CAMAsyncAFDConnector", ) +AFDConnectorFactory.register_connector( + "GpuAsyncAFDConnector", + "afd_plugin.connectors.gpu.async_gpu", + "GpuAsyncAFDConnector", +) __all__ = ["AFDConnectorFactory"] diff --git a/afd_plugin/connectors/gpu/__init__.py b/afd_plugin/connectors/gpu/__init__.py index 62314259..77e73746 100644 --- a/afd_plugin/connectors/gpu/__init__.py +++ b/afd_plugin/connectors/gpu/__init__.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project """GPU-specific AFD connector implementations.""" +from afd_plugin.connectors.gpu.async_gpu import GpuAsyncAFDConnector from afd_plugin.connectors.gpu.p2p import P2pNcclAFDConnector -__all__ = ["P2pNcclAFDConnector"] +__all__ = ["GpuAsyncAFDConnector", "P2pNcclAFDConnector"] diff --git a/afd_plugin/connectors/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py new file mode 100644 index 00000000..add5c69b --- /dev/null +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -0,0 +1,889 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""NVSHMEM-backed asynchronous connector for CUDA AFD. + +``GpuAsyncAFDConnector`` is the CUDA counterpart of ``CAMAsyncAFDConnector``: +Attention ranks run MoE routing, write routed tokens one-sided into the FFN +ranks' symmetric windows, and later reduce the weighted expert output; FFN ranks +poll their window, run their local experts, and write the result back. There is +no DP metadata control plane (``control_plane`` stays ``None``) and FFN work is +driven directly by the connector receive loop, so Attention DP replicas never +wait for each other. + +The world is Attention-first, ``[A0, A1, ..., F0, F1, ...]``, matching +``CAMAsyncAFDConnector``. Every Attention rank routes to every FFN rank, so an +FFN window holds one region per Attention rank and vice versa. + +See ``docs/design/rfc_async_gpu_connector.md``. Supported deployment requires +``async=true``, ``compute_gate_on_attention=true``, eager execution, prefill +only, and a single node. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import timedelta +from typing import TYPE_CHECKING, Any, Final + +import torch +from torch import Tensor +from vllm.logger import init_logger + +from afd_plugin.config import AFDConfig +from afd_plugin.config_utils import ( + coerce_extra_bool, + coerce_extra_positive_int, + coerce_extra_str, +) +from afd_plugin.connectors.async_topology import ( + ASYNC_MOE_REQUEST_SPLIT, + build_async_topology, +) +from afd_plugin.connectors.base import AFDConnectorBase, ConnectorExtraInfo +from afd_plugin.connectors.gpu.symm_window import ( + FLAG_SHUTDOWN_BIT, + SlotLayout, + SymmWindow, + encode_header, +) +from afd_plugin.connectors.metadata import ( + AFDA2FTransferPayload, + AFDF2ATransferPayload, + AFDTransferContext, + AFDTransferMetadata, + AFDTransferState, +) +from afd_plugin.distributed import init_afd_process_group + +if TYPE_CHECKING: + from torch.distributed.distributed_c10d import ProcessGroup + from vllm.config import VllmConfig + +AFD_ASYNC_GPU_GROUP_NAME = "afd_async_gpu" + +_GPU_ASYNC_EXTRA_CONFIG_FIELDS: Final[frozenset[str]] = frozenset( + { + "attn_ranks_per_dp", + "ring_depth", + "routed_cap_multiplier", + "recv_poll_timeout_ms", + "async_moe_ubatching", + "async_moe_num_ubatches", + "async_moe_split", + }, +) + +# Name the logger inside vLLM's tree: vLLM installs its handler on the "vllm" +# logger only, so a bare ``afd_plugin.*`` logger propagates to a handler-less +# root and every line is dropped -- which is how the window summary, the only +# report of a multi-GiB allocation, stayed invisible. +logger = init_logger(f"vllm.{__name__}") + + +def _coerce_extra_float(value: Any, *, field_name: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float, str)): + raise TypeError(f"{field_name} must be a number, got {type(value).__name__}") + result = float(value) + if result <= 0.0: + raise ValueError(f"{field_name} must be positive, got {result}") + return result + + +@dataclass(frozen=True) +class GpuAsyncExtraInfo(ConnectorExtraInfo): + """Typed async GPU connector configuration. + + Attributes: + attn_ranks_per_dp: Number of Attention ranks in each data-parallel group. + ring_depth: Slots per peer region. Derived from the send-then-recv + invariant, not a performance knob: an Attention rank has at most one + in-flight request per ``(peer, stage)``, so ``num_stages`` suffices. + routed_cap_multiplier: Headroom over the balanced routed-token estimate. + Real gates are not balanced -- DeepSeek-V2-Lite at 2 FFN ranks was + observed 1.33x over the even split -- so the default leaves room. + The true worst case is ``ffn_size`` (every partial to one rank). + recv_poll_timeout_ms: Idle poll timeout on the FFN loop; bounds shutdown + response time. + async_moe_ubatching: Whether request-boundary async MoE ubatching is used. + async_moe_num_ubatches: Number of stages used by async MoE ubatching. + async_moe_split: Boundary at which async MoE work is split. + """ + + attn_ranks_per_dp: int = 1 + ring_depth: int = 0 + routed_cap_multiplier: float = 2.0 + recv_poll_timeout_ms: int = 50 + async_moe_ubatching: bool = False + async_moe_num_ubatches: int = 2 + async_moe_split: str = ASYNC_MOE_REQUEST_SPLIT + + @classmethod + def from_mapping(cls, raw: Mapping[str, Any] | None) -> GpuAsyncExtraInfo: + if raw is None: + raw = {} + if not isinstance(raw, Mapping): + raise TypeError( + f"{cls.__name__} connector_extra_config must be a mapping, " + f"got {type(raw).__name__}", + ) + unknown = sorted( + str(key) for key in raw if key not in _GPU_ASYNC_EXTRA_CONFIG_FIELDS + ) + if unknown: + raise ValueError( + "unknown AFD async GPU connector_extra_config field(s): " + + ", ".join(unknown), + ) + + ubatching = coerce_extra_bool( + raw.get("async_moe_ubatching", False), + field_name="async_moe_ubatching", + ) + num_ubatches = coerce_extra_positive_int( + raw.get("async_moe_num_ubatches", 2), + field_name="async_moe_num_ubatches", + ) + # Ring depth follows the number of live stages unless pinned explicitly. + ring_depth = coerce_extra_positive_int( + raw.get("ring_depth", num_ubatches if ubatching else 1), + field_name="ring_depth", + ) + return cls( + attn_ranks_per_dp=coerce_extra_positive_int( + raw.get("attn_ranks_per_dp", 1), + field_name="attn_ranks_per_dp", + ), + ring_depth=ring_depth, + routed_cap_multiplier=_coerce_extra_float( + raw.get("routed_cap_multiplier", 2.0), + field_name="routed_cap_multiplier", + ), + recv_poll_timeout_ms=coerce_extra_positive_int( + raw.get("recv_poll_timeout_ms", 50), + field_name="recv_poll_timeout_ms", + ), + async_moe_ubatching=ubatching, + async_moe_num_ubatches=num_ubatches, + async_moe_split=coerce_extra_str( + raw.get("async_moe_split", ASYNC_MOE_REQUEST_SPLIT), + field_name="async_moe_split", + ), + ) + + def to_mapping(self) -> dict[str, Any]: + return { + "attn_ranks_per_dp": self.attn_ranks_per_dp, + "ring_depth": self.ring_depth, + "routed_cap_multiplier": self.routed_cap_multiplier, + "recv_poll_timeout_ms": self.recv_poll_timeout_ms, + "async_moe_ubatching": self.async_moe_ubatching, + "async_moe_num_ubatches": self.async_moe_num_ubatches, + "async_moe_split": self.async_moe_split, + } + + +class ConnectorShutdown(RuntimeError): # noqa: N818 + """Raised on the FFN loop when a peer announced shutdown.""" + + +@dataclass(slots=True) +class GpuAsyncTransferState(AFDTransferState): + """FFN-side state carried from dispatch recv through combine send. + + ``region``/``ring`` locate the window slot so ``send_ffn_work_item_output`` + can write back to the originating Attention rank and release the slot; + ``route_table`` is echoed so the Attention side can scatter the result. + """ + + region: int = 0 + ring: int = 0 + src_role_rank: int = 0 + layer_idx: int = 0 + stage_idx: int = 0 + seq: int = 0 + num_tokens: int = 0 + routed_tokens: int = 0 + shared_tokens: int = 0 + group_list: Tensor | None = None + route_table: Tensor | None = None + shared_idx: Tensor | None = None + expand_x_shared: Tensor | None = None + + +@dataclass(slots=True) +class GpuAsyncFFNWorkItem: + """Normalized FFN-side work item produced by a window arrival.""" + + hidden_states: Tensor + context: AFDTransferContext + recv_output: AFDA2FTransferPayload + layer_idx: int + stage_idx: int + num_tokens: int + total_num_tokens: int + shared_num_tokens: int + + +@dataclass(slots=True) +class _PendingDispatch: + """Attention-side record of one in-flight layer, popped by combine recv.""" + + context: AFDTransferContext + topk_weights: Tensor + num_tokens: int + ring: int + seq: int + expected_ffn: list[int] + + +def plan_dispatch( + topk_ids: Tensor, + *, + ffn_size: int, + expert_per_rank: int, +) -> tuple[Tensor, Tensor, Tensor]: + """Cluster ``(token, topk_slot)`` partials by destination FFN rank. + + Sorting by the global expert id groups partials by destination rank and, in + the same pass, by local expert inside each destination -- the two orderings + the receiver needs. Returns ``(route_table, counts, offsets)`` where + ``route_table[i] = (token_idx, topk_slot)`` in that order, ``counts`` holds + per-global-expert partial counts padded to ``ffn_size * expert_per_rank``, + and ``offsets`` is the exclusive prefix sum of ``counts``. + """ + num_slots = topk_ids.shape[1] + flat = topk_ids.reshape(-1).to(torch.int64) + order = torch.argsort(flat, stable=True) + counts = torch.bincount(flat, minlength=ffn_size * expert_per_rank) + offsets = torch.cumsum(counts, dim=0) - counts + route_table = torch.stack( + (order // num_slots, order % num_slots), + dim=1, + ).to(torch.int32) + return route_table, counts, offsets + + +def combine_scatter( + accumulator: Tensor, + *, + routed_out: Tensor, + route_table: Tensor, + topk_weights: Tensor, +) -> None: + """Weighted-scatter one FFN rank's routed output into ``accumulator``.""" + token_idx = route_table[:, 0].to(torch.int64) + slot_idx = route_table[:, 1].to(torch.int64) + weights = topk_weights[token_idx, slot_idx].to(accumulator.dtype) + accumulator.index_add_( + 0, + token_idx, + routed_out.to(accumulator.dtype) * weights.unsqueeze(1), + ) + + +class GpuAsyncAFDConnector(AFDConnectorBase): + """NVSHMEM symmetric-window asynchronous connector for CUDA AFD.""" + + control_plane = None + + @classmethod + def parse_extra_config( + cls, + raw: Mapping[str, Any] | None, + ) -> GpuAsyncExtraInfo: + return GpuAsyncExtraInfo.from_mapping(raw) + + def __init__( + self, + rank: int, + local_rank: int, + vllm_config: VllmConfig, + afd_config: AFDConfig, + role_rank: int, + ) -> None: + super().__init__(rank, local_rank, vllm_config, afd_config, role_rank) + self._initialized = False + hf_config = vllm_config.model_config.hf_config + self.hidden_size = hf_config.hidden_size + self.topk = hf_config.num_experts_per_tok + self.num_routed_experts = hf_config.n_routed_experts + self.payload_dtype = vllm_config.model_config.dtype + self.max_seq_len = vllm_config.scheduler_config.max_num_batched_tokens + self.tp_size = self.extra_info.attn_ranks_per_dp + + self.topology = build_async_topology( + afd_config, + role_rank, + num_routed_experts=self.num_routed_experts, + ) + self.world_rank = self.topology.world_rank + self.attn_size = self.topology.attn_size + self.ffn_size = self.topology.ffn_size + self.expert_per_rank = self.topology.expert_per_rank + self.is_attention = afd_config.role == "attention" + + self.ring_depth = self.extra_info.ring_depth + # Every Attention rank routes to every FFN rank, so a window carries one + # region per opposite-role peer. Both roles allocate the larger of the + # two so the symmetric allocation matches. + self.num_regions = max(self.attn_size, self.ffn_size) + routed_cap = int( + -(-self.max_seq_len * self.topk // self.ffn_size) + * self.extra_info.routed_cap_multiplier, + ) + self.routed_cap = max(1, routed_cap) + self.token_cap = max(1, self.max_seq_len) + self.layout = SlotLayout.build( + expert_per_rank=self.expert_per_rank, + routed_cap=self.routed_cap, + token_cap=self.token_cap, + hidden_size=self.hidden_size, + payload_itemsize=torch.empty(0, dtype=self.payload_dtype).element_size(), + ) + + self.pg: ProcessGroup | None = None + self.window: SymmWindow | None = None + self._seq = 0 + self._pending: dict[int, list[_PendingDispatch]] = {} + self._free_rings: dict[int, list[int]] = {} + + @property + def is_initialized(self) -> bool: + return self._initialized + + def init_afd_connector(self) -> None: + """Collectively create the AFD world group and the symmetric window. + + All Attention and FFN ranks must call this with identical rendezvous and + topology settings; the window allocation is symmetric, so a mismatched + size fails here rather than corrupting a later transfer. + """ + if self._initialized: + return + + self.pg = init_afd_process_group( + backend="nccl", + init_method=f"tcp://{self.afd_config.host}:{self.afd_config.port}", + world_size=self.topology.world_size, + rank=self.world_rank, + group_name=AFD_ASYNC_GPU_GROUP_NAME, + timeout=timedelta(minutes=30), + ) + device = torch.device("cuda", self.local_rank) + self.window = SymmWindow( + num_regions=self.num_regions, + ring_depth=self.ring_depth, + layout=self.layout, + payload_dtype=self.payload_dtype, + device=device, + group=self.pg, + rank=self.world_rank, + world_size=self.topology.world_size, + ) + for stage in range(max(1, self.extra_info.async_moe_num_ubatches)): + self._free_rings[stage] = list(range(self.ring_depth)) + logger.info( + "AFD async GPU window ready: role=%s role_rank=%d world_rank=%d/%d " + "regions=%d rings=%d routed_cap=%d slot=%.1fMiB total=%.1fMiB", + self.afd_config.role, + self.role_rank, + self.world_rank, + self.topology.world_size, + self.num_regions, + self.ring_depth, + self.routed_cap, + self.layout.slot_bytes / 2**20, + self.window.total_bytes / 2**20, + ) + self._initialized = True + + def close(self) -> None: + if self.window is not None: + self.window.close() + self.window = None + if self.pg is not None: + import torch.distributed as dist + + dist.destroy_process_group(self.pg) + self.pg = None + self._pending.clear() + self._free_rings.clear() + self._initialized = False + + def select_experts(self, **kwargs: Any) -> tuple[Tensor, Tensor]: + """Run vLLM's grouped top-k on the Attention side. + + ``compute_gate_topk`` delegates expert selection to the connector so the + CAM and CUDA paths can share one gate; this is the CUDA half. + """ + from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( + grouped_topk, + ) + + if kwargs.get("mix_placement"): + raise RuntimeError( + "AFD async GPU connector does not support mix_placement", + ) + return grouped_topk( + hidden_states=kwargs["hidden_states"], + gating_output=kwargs["router_logits"], + topk=kwargs["top_k"], + renormalize=kwargs["renormalize"], + num_expert_group=kwargs.get("num_expert_group", 0), + topk_group=kwargs.get("topk_group", 0), + scoring_func=kwargs.get("scoring_func", "softmax"), + routed_scaling_factor=kwargs.get("routed_scaling_factor", 1.0), + e_score_correction_bias=kwargs.get("e_score_correction_bias"), + ) + + def _require_initialized(self) -> SymmWindow: + if not self._initialized or self.window is None: + raise RuntimeError("AFD async GPU connector is not initialized") + return self.window + + # ================================================================== + # Attention-side data path + # ================================================================== + + def send_attn_output( + self, + hidden_states: Tensor, + context: AFDTransferContext, + **kwargs: Any, + ) -> None: + """Route this layer's tokens and write them into every FFN window. + + ``topk_ids``/``topk_weights`` come from the Attention-side gate. Weights + stay local -- only the routed activations and their route table go on the + wire, and the weighting happens in ``recv_ffn_output``. + """ + window = self._require_initialized() + topk_ids: Tensor | None = kwargs.get("topk_ids") + topk_weights: Tensor | None = kwargs.get("topk_weights") + if topk_ids is None or topk_weights is None: + raise RuntimeError( + "AFD async GPU send_attn_output requires topk_ids and " + "topk_weights from the Attention-side gate", + ) + metadata = context.metadata + num_tokens = metadata.total_tokens + if hidden_states.shape[0] != num_tokens: + raise ValueError( + f"hidden_states has {hidden_states.shape[0]} rows but metadata " + f"expects {num_tokens}", + ) + if tuple(topk_ids.shape) != (num_tokens, self.topk): + raise ValueError( + f"topk_ids shape must be ({num_tokens}, {self.topk}), " + f"got {tuple(topk_ids.shape)}", + ) + + stage_idx = metadata.stage_idx + rings = self._free_rings.setdefault(stage_idx, list(range(self.ring_depth))) + if not rings: + raise RuntimeError( + f"AFD async GPU ring exhausted on stage {stage_idx}; the " + "send-then-recv invariant was violated or the topology config " + "does not match the actual peer count", + ) + ring = rings.pop(0) + self._seq += 1 + + route_table, counts, _ = plan_dispatch( + topk_ids, + ffn_size=self.ffn_size, + expert_per_rank=self.expert_per_rank, + ) + # One D2H per send: offsets are a prefix sum, cheaper to redo on host + # than to fetch a second tensor. + counts_host = counts.cpu().tolist() + offsets_host = [0] * len(counts_host) + for i in range(1, len(counts_host)): + offsets_host[i] = offsets_host[i - 1] + counts_host[i - 1] + token_ids = route_table[:, 0].to(torch.int64) + + # Every FFN rank gets a slot even when routing sends it nothing, and it + # replies to every slot, so a reply is expected from all of them. + # Expecting only the ranks that received data leaves the empty rank's + # reply unmatched and its ring slot never released -- which is what a + # single-token decode hits, since both the routed segment and the + # round-robin shared slice can come out empty for one rank. + expected_ffn = list(range(self.ffn_size)) + for ffn_rank in range(self.ffn_size): + base = ffn_rank * self.expert_per_rank + expert_counts = counts_host[base : base + self.expert_per_rank] + start = offsets_host[base] + routed_tokens = sum(expert_counts) + segment = slice(start, start + routed_tokens) + + # Shared-expert tokens are split round-robin across FFN ranks. + shared_idx = torch.arange( + ffn_rank, + num_tokens, + self.ffn_size, + device=hidden_states.device, + dtype=torch.int32, + ) + header = encode_header( + self.layout, + seq=self._seq, + src_role_rank=self.role_rank, + layer_idx=metadata.layer_idx, + stage_idx=stage_idx, + num_tokens=num_tokens, + routed_tokens=routed_tokens, + shared_tokens=int(shared_idx.numel()), + topk=self.topk, + flags=0, + expert_counts=expert_counts, + ) + window.write_slot( + peer=self.attn_size + ffn_rank, + region=self.role_rank, + ring=ring, + header=header, + route_table=route_table[segment], + routed_x=hidden_states.index_select(0, token_ids[segment]), + shared_idx=shared_idx, + shared_x=hidden_states.index_select(0, shared_idx.to(torch.int64)), + ) + + logger.debug( + "AFD dispatch sent: A%d layer=%d stage=%d tokens=%d ring=%d " + "awaiting_ffn=%s", + self.role_rank, + metadata.layer_idx, + stage_idx, + num_tokens, + ring, + expected_ffn, + ) + self._pending.setdefault(stage_idx, []).append( + _PendingDispatch( + context=context, + topk_weights=topk_weights, + num_tokens=num_tokens, + ring=ring, + seq=self._seq, + expected_ffn=expected_ffn, + ), + ) + + def recv_ffn_output( + self, + ref_tensor: Tensor, + ubatch_idx: int = 0, + **kwargs: Any, + ) -> Tensor: + """Wait for this layer's expert output and reduce it back to ``[B, H]``.""" + + window = self._require_initialized() + queue = self._pending.get(ubatch_idx) + if not queue: + raise RuntimeError( + f"AFD async GPU recv_ffn_output has no pending dispatch on " + f"stage {ubatch_idx}", + ) + pending = queue.pop(0) + + accumulator = torch.zeros( + (pending.num_tokens, self.hidden_size), + dtype=torch.float32, + device=ref_tensor.device, + ) + outstanding = set(pending.expected_ffn) + while outstanding: + arrived = window.poll() + if arrived is None: + continue + header = arrived.header + if header.is_shutdown: + raise ConnectorShutdown( + f"FFN rank {header.src_role_rank} announced shutdown", + ) + if header.src_role_rank not in outstanding: + raise RuntimeError( + "AFD async GPU combine received an unexpected FFN rank " + f"{header.src_role_rank}; expected one of {sorted(outstanding)}", + ) + if header.echo_seq != pending.seq: + raise RuntimeError( + "AFD async GPU combine answered dispatch seq " + f"{header.echo_seq} while waiting on {pending.seq} " + f"(F{header.src_role_rank}, layer {header.layer_idx}); the " + "pending FIFO and the wire have diverged", + ) + outstanding.discard(header.src_role_rank) + logger.debug( + "AFD combine recv: A%d <- F%d layer=%d routed=%d still_waiting=%s", + self.role_rank, + header.src_role_rank, + header.layer_idx, + header.routed_tokens, + sorted(outstanding), + ) + + if header.routed_tokens: + combine_scatter( + accumulator, + routed_out=window.local_routed( + arrived.region, + arrived.ring, + header.routed_tokens, + ), + route_table=window.local_route_table( + arrived.region, + arrived.ring, + header.routed_tokens, + ), + topk_weights=pending.topk_weights, + ) + if header.shared_tokens: + accumulator.index_add_( + 0, + window.local_shared_idx( + arrived.region, + arrived.ring, + header.shared_tokens, + ).to(torch.int64), + window.local_shared( + arrived.region, + arrived.ring, + header.shared_tokens, + ).to(accumulator.dtype), + ) + + self._free_rings.setdefault(ubatch_idx, []).append(pending.ring) + return accumulator.to(ref_tensor.dtype) + + # ================================================================== + # FFN-side data path + # ================================================================== + + def recv_attn_output( + self, + ubatch_idx: int = 0, + **kwargs: Any, + ) -> AFDA2FTransferPayload: + """Block until one Attention rank's routed tokens arrive. + + The layer index, token counts, and per-expert group list all come from + the arrived slot header; the FFN side knows none of them beforehand. + """ + window = self._require_initialized() + timeout_ms = int(kwargs.get("timeout_ms", 0)) + deadline = None + if timeout_ms: + import time + + deadline = time.monotonic() + timeout_ms / 1000.0 + + while True: + arrived = window.poll() + if arrived is not None: + break + if deadline is not None: + import time + + if time.monotonic() >= deadline: + raise TimeoutError("AFD async GPU dispatch recv timed out") + + header = arrived.header + if header.is_shutdown: + raise ConnectorShutdown( + f"Attention rank {header.src_role_rank} announced shutdown", + ) + + states = GpuAsyncTransferState( + region=arrived.region, + ring=arrived.ring, + seq=header.seq, + src_role_rank=header.src_role_rank, + layer_idx=header.layer_idx, + stage_idx=header.stage_idx, + num_tokens=header.num_tokens, + routed_tokens=header.routed_tokens, + shared_tokens=header.shared_tokens, + group_list=torch.tensor( + header.expert_counts, + dtype=torch.int64, + device=torch.device("cuda", self.local_rank), + ), + route_table=window.local_route_table( + arrived.region, + arrived.ring, + header.routed_tokens, + ), + shared_idx=window.local_shared_idx( + arrived.region, + arrived.ring, + header.shared_tokens, + ), + expand_x_shared=window.local_shared( + arrived.region, + arrived.ring, + header.shared_tokens, + ), + ) + logger.debug( + "AFD dispatch recv: F%d <- A%d layer=%d stage=%d routed=%d shared=%d " + "region=%d ring=%d", + self.role_rank, + header.src_role_rank, + header.layer_idx, + header.stage_idx, + header.routed_tokens, + header.shared_tokens, + arrived.region, + arrived.ring, + ) + metadata = AFDTransferMetadata.create_ffn_metadata( + layer_idx=header.layer_idx, + stage_idx=header.stage_idx, + seq_lens=[max(1, header.routed_tokens)], + ) + return AFDA2FTransferPayload( + hidden_states=window.local_routed( + arrived.region, + arrived.ring, + header.routed_tokens, + ), + context=AFDTransferContext(metadata=metadata, states=states), + ) + + def send_ffn_output( + self, + ffn_output: Tensor, + context: AFDTransferContext, + **kwargs: Any, + ) -> None: + """Write expert output back to the originating Attention rank.""" + window = self._require_initialized() + states = context.states + if not isinstance(states, GpuAsyncTransferState): + raise RuntimeError( + "AFD async GPU send_ffn_output requires GpuAsyncTransferState", + ) + shared_output: Tensor | None = kwargs.get("shared_output") + self._seq += 1 + header = encode_header( + self.layout, + seq=self._seq, + src_role_rank=self.role_rank, + layer_idx=states.layer_idx, + stage_idx=states.stage_idx, + num_tokens=states.num_tokens, + routed_tokens=states.routed_tokens, + shared_tokens=states.shared_tokens if shared_output is not None else 0, + topk=self.topk, + flags=0, + expert_counts=list(header_counts(states)), + echo_seq=states.seq, + ) + window.write_slot( + peer=states.src_role_rank, + region=self.role_rank, + ring=states.ring, + header=header, + route_table=states.route_table, + routed_x=ffn_output, + shared_idx=states.shared_idx if shared_output is not None else None, + shared_x=shared_output, + ) + + # ================================================================== + # Connector-driven FFN loop + # ================================================================== + + def recv_ffn_work_item( + self, + *, + stage_idx: int, + max_num_tokens: int, + ) -> GpuAsyncFFNWorkItem: + """Receive and normalize one connector-driven FFN dispatch item.""" + recv_output = self.recv_attn_output( + ubatch_idx=stage_idx, + timeout_ms=self.extra_info.recv_poll_timeout_ms, + ) + states = recv_output.context.states + assert isinstance(states, GpuAsyncTransferState) + return GpuAsyncFFNWorkItem( + hidden_states=recv_output.hidden_states, + context=recv_output.context, + recv_output=recv_output, + layer_idx=states.layer_idx, + stage_idx=states.stage_idx, + num_tokens=states.routed_tokens, + total_num_tokens=states.num_tokens, + shared_num_tokens=states.shared_tokens, + ) + + def send_ffn_work_item_output( + self, + work_item: GpuAsyncFFNWorkItem, + ffn_output: Tensor | AFDF2ATransferPayload, + ) -> Tensor: + """Return one work item's expert output to its Attention rank.""" + if isinstance(ffn_output, AFDF2ATransferPayload): + routed = ffn_output.routed_output + shared = ffn_output.shared_output + else: + routed = ffn_output + shared = None + self.send_ffn_output(routed, work_item.context, shared_output=shared) + return routed + + def announce_shutdown(self) -> None: + """Tell every opposite-role peer to leave its receive loop.""" + window = self._require_initialized() + peers = ( + range(self.attn_size, self.attn_size + self.ffn_size) + if self.is_attention + else range(self.attn_size) + ) + self._seq += 1 + header = encode_header( + self.layout, + seq=self._seq, + src_role_rank=self.role_rank, + layer_idx=0, + stage_idx=0, + num_tokens=0, + routed_tokens=0, + shared_tokens=0, + topk=self.topk, + flags=FLAG_SHUTDOWN_BIT, + expert_counts=[0] * self.expert_per_rank, + ) + for peer in peers: + window.write_slot( + peer=peer, + region=self.role_rank, + ring=0, + header=header, + route_table=None, + routed_x=None, + shared_idx=None, + shared_x=None, + ) + + +def header_counts(states: GpuAsyncTransferState) -> list[int]: + """Echo the per-expert group list back on the combine header.""" + if states.group_list is None: + return [] + return states.group_list.cpu().tolist() + + +__all__ = [ + "AFD_ASYNC_GPU_GROUP_NAME", + "ConnectorShutdown", + "GpuAsyncAFDConnector", + "GpuAsyncExtraInfo", + "GpuAsyncFFNWorkItem", + "GpuAsyncTransferState", + "combine_scatter", + "plan_dispatch", +] diff --git a/afd_plugin/connectors/gpu/nvshmem_rt.py b/afd_plugin/connectors/gpu/nvshmem_rt.py new file mode 100644 index 00000000..d2b3b058 --- /dev/null +++ b/afd_plugin/connectors/gpu/nvshmem_rt.py @@ -0,0 +1,270 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Minimal NVSHMEM host-library binding for the async GPU connector. + +``torch.distributed._symmetric_memory`` cannot serve AFD: its NVSHMEM backend +bootstraps on the *default* process group and carves teams out of +``NVSHMEM_TEAM_WORLD`` with ``nvshmem_team_split_strided``. Each AFD role runs +as its own ``vllm serve`` whose default group covers only that role's ranks, so +the cross-role AFD group is never a strided subset of it and team creation +fails. + +Bootstrapping NVSHMEM ourselves from a unique id exchanged over the AFD group's +store makes the AFD world *be* ``NVSHMEM_TEAM_WORLD``, so no team split is +needed. The host library's ABI is pinned by static asserts in its own headers +(``uniqueid`` 128 B, ``init_attr`` 144 B, version = ``(1 << 16) + sizeof``), +which is why ctypes is enough and ``csrc/gpu/`` stays empty. +""" + +from __future__ import annotations + +import ctypes +import os +from typing import TYPE_CHECKING, Final + +import torch + +if TYPE_CHECKING: + from torch.distributed.distributed_c10d import ProcessGroup, Store + +UNIQUEID_PADDING: Final[int] = 124 +# 128 (init_args) - 4 (version) - 24 (uid_args) - 4 (trailing alignment) +INIT_ARGS_PADDING: Final[int] = 96 +NVSHMEMX_INIT_WITH_UNIQUEID: Final[int] = 1 << 3 +_UID_STORE_KEY: Final[str] = "afd_nvshmem_uid" +_LIB_RELATIVE: Final[str] = "nvidia/nvshmem/lib/libnvshmem_host.so.3" + + +class _UniqueId(ctypes.Structure): + _fields_ = ( + ("version", ctypes.c_int), + ("internal", ctypes.c_char * UNIQUEID_PADDING), + ) + + +class _UniqueIdArgs(ctypes.Structure): + _fields_ = ( + ("version", ctypes.c_int), + ("id", ctypes.POINTER(_UniqueId)), + ("myrank", ctypes.c_int), + ("nranks", ctypes.c_int), + ) + + +class _InitArgs(ctypes.Structure): + _fields_ = ( + ("version", ctypes.c_int), + ("uid_args", _UniqueIdArgs), + ("content", ctypes.c_char * INIT_ARGS_PADDING), + ) + + +class _InitAttr(ctypes.Structure): + _fields_ = ( + ("version", ctypes.c_int), + ("mpi_comm", ctypes.c_void_p), + ("args", _InitArgs), + ) + + +def _find_host_library() -> str: + """Locate ``libnvshmem_host.so.3`` next to the installed nvshmem wheel.""" + import site + import sysconfig + + roots = [sysconfig.get_paths()["purelib"], *site.getsitepackages()] + for root in roots: + candidate = os.path.join(root, _LIB_RELATIVE) + if os.path.exists(candidate): + return candidate + raise RuntimeError( + "AFD async GPU connector requires the NVSHMEM host library; " + f"{_LIB_RELATIVE} was not found under {roots}. Install " + "nvidia-nvshmem-cu13 matching the installed torch build.", + ) + + +def _load_library() -> ctypes.CDLL: + lib = ctypes.CDLL(_find_host_library(), mode=ctypes.RTLD_GLOBAL) + lib.nvshmemx_get_uniqueid.argtypes = [ctypes.POINTER(_UniqueId)] + lib.nvshmemx_get_uniqueid.restype = ctypes.c_int + lib.nvshmemx_set_attr_uniqueid_args.argtypes = [ + ctypes.c_int, + ctypes.c_int, + ctypes.POINTER(_UniqueId), + ctypes.POINTER(_InitAttr), + ] + lib.nvshmemx_set_attr_uniqueid_args.restype = ctypes.c_int + lib.nvshmemx_hostlib_init_attr.argtypes = [ + ctypes.c_uint, + ctypes.POINTER(_InitAttr), + ] + lib.nvshmemx_hostlib_init_attr.restype = ctypes.c_int + lib.nvshmem_malloc.argtypes = [ctypes.c_size_t] + lib.nvshmem_malloc.restype = ctypes.c_void_p + lib.nvshmem_ptr.argtypes = [ctypes.c_void_p, ctypes.c_int] + lib.nvshmem_ptr.restype = ctypes.c_void_p + lib.nvshmem_my_pe.restype = ctypes.c_int + lib.nvshmem_n_pes.restype = ctypes.c_int + return lib + + +# NVSHMEM initialization is process-global: a process joins exactly one NVSHMEM +# world, so every connector in it shares this state. +_lib: ctypes.CDLL | None = None +_initialized_world: tuple[int, int] | None = None + + +def init(pg: ProcessGroup, rank: int, world_size: int) -> None: + """Join the NVSHMEM world described by ``pg``, once per process. + + Rank 0 mints the unique id and publishes it on the group's store; every rank + then initializes with the same id, so NVSHMEM's PE numbering equals the AFD + world rank. + """ + global _lib, _initialized_world + + if _initialized_world is not None: + if _initialized_world != (rank, world_size): + raise RuntimeError( + "NVSHMEM is already initialized in this process as " + f"rank {_initialized_world[0]} of {_initialized_world[1]}; " + f"cannot re-initialize as rank {rank} of {world_size}", + ) + return + + from torch.distributed.distributed_c10d import _get_process_group_store + + lib = _load_library() + store: Store = _get_process_group_store(pg) + + unique_id = _UniqueId() + unique_id.version = (1 << 16) + ctypes.sizeof(_UniqueId) + if rank == 0: + if lib.nvshmemx_get_uniqueid(ctypes.byref(unique_id)) != 0: + raise RuntimeError("nvshmemx_get_uniqueid failed") + store.set(_UID_STORE_KEY, bytes(memoryview(unique_id).cast("B"))) + else: + raw = store.get(_UID_STORE_KEY) + ctypes.memmove(ctypes.byref(unique_id), raw, ctypes.sizeof(_UniqueId)) + + attr = _InitAttr() + attr.version = (1 << 16) + ctypes.sizeof(_InitAttr) + attr.args.version = (1 << 16) + ctypes.sizeof(_InitArgs) + attr.args.uid_args.version = (1 << 16) + ctypes.sizeof(_UniqueIdArgs) + if ( + lib.nvshmemx_set_attr_uniqueid_args( + rank, + world_size, + ctypes.byref(unique_id), + ctypes.byref(attr), + ) + != 0 + ): + raise RuntimeError("nvshmemx_set_attr_uniqueid_args failed") + if ( + lib.nvshmemx_hostlib_init_attr( + NVSHMEMX_INIT_WITH_UNIQUEID, + ctypes.byref(attr), + ) + != 0 + ): + raise RuntimeError("nvshmemx_hostlib_init_attr failed") + + actual_pe, actual_world = lib.nvshmem_my_pe(), lib.nvshmem_n_pes() + if (actual_pe, actual_world) != (rank, world_size): + raise RuntimeError( + f"NVSHMEM PE numbering does not match the AFD world: got PE " + f"{actual_pe} of {actual_world}, expected {rank} of {world_size}", + ) + _lib = lib + _initialized_world = (rank, world_size) + + +def is_initialized() -> bool: + return _initialized_world is not None + + +def _require_lib() -> ctypes.CDLL: + if _lib is None: + raise RuntimeError("NVSHMEM is not initialized; call init() first") + return _lib + + +def malloc(nbytes: int) -> int: + """Allocate a symmetric buffer. Collective: every PE must call it alike.""" + pointer = _require_lib().nvshmem_malloc(nbytes) + if not pointer: + raise RuntimeError( + f"nvshmem_malloc({nbytes}) returned NULL; raise " + "NVSHMEM_SYMMETRIC_SIZE or lower the window capacity", + ) + return int(pointer) + + +def peer_ptr(local_ptr: int, pe: int) -> int: + """Map a peer's copy of a symmetric allocation into this process.""" + pointer = _require_lib().nvshmem_ptr(ctypes.c_void_p(local_ptr), pe) + if not pointer: + raise RuntimeError( + f"nvshmem_ptr returned NULL for PE {pe}: no direct peer access. " + "The async GPU connector requires PEs reachable over NVLink/P2P; " + "cross-node placement is not supported.", + ) + return int(pointer) + + +class _DeviceBuffer: + """Hand a raw device pointer to torch via ``__cuda_array_interface__``.""" + + def __init__(self, pointer: int, nbytes: int) -> None: + self.__cuda_array_interface__ = { + "data": (pointer, False), + "shape": (nbytes,), + "typestr": "|u1", + "version": 3, + "strides": None, + } + + +def tensor_from_ptr( + base_ptr: int, + *, + byte_offset: int, + sizes: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """View symmetric memory as a tensor without copying. + + The buffer is exposed as bytes and then reinterpreted, because + ``__cuda_array_interface__`` has no type string for dtypes like bfloat16. + """ + itemsize = torch.empty(0, dtype=dtype).element_size() + numel = 1 + for size in sizes: + numel *= size + if numel == 0: + # A zero-length __cuda_array_interface__ buffer is rejected by the CUDA + # runtime (cudaErrorInvalidValue); an empty slot is legitimate whenever + # routing sends a peer nothing, so hand back a plain empty tensor. + return torch.empty(sizes, dtype=dtype, device=device) + nbytes = numel * itemsize + if byte_offset % itemsize: + raise ValueError( + f"byte offset {byte_offset} is not aligned to {itemsize}-byte {dtype}", + ) + raw = torch.as_tensor( + _DeviceBuffer(base_ptr + byte_offset, nbytes), + device=device, + ) + return raw.view(dtype).reshape(sizes) + + +__all__ = [ + "init", + "is_initialized", + "malloc", + "peer_ptr", + "tensor_from_ptr", +] diff --git a/afd_plugin/connectors/gpu/symm_window.py b/afd_plugin/connectors/gpu/symm_window.py new file mode 100644 index 00000000..af920960 --- /dev/null +++ b/afd_plugin/connectors/gpu/symm_window.py @@ -0,0 +1,528 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Symmetric-memory window used by the async GPU AFD connector. + +The window mirrors the CAM shared-window substrate: every rank allocates an +identically sized symmetric buffer, senders write one-sided into the receiver's +buffer, and arrival is announced by a magic-stamped flag word that the receiver +polls. Only the allocation is collective; once it is done a sender never needs +the receiver to participate. + +Layout of one window:: + + [flag[num_regions * ring_depth] | slot(0, 0) | slot(0, 1) | ...] + +``slot(region, ring)`` holds a fixed header followed by the routed/shared +payloads. Dispatch (A -> F) and combine (F -> A) use the same slot layout, so a +single spec sizes both directions. + +Flag words are written *after* the payload on the same stream. Same-stream +device-to-device copies complete in issue order, so a visible flag implies a +complete payload. That holds for NVLink-mapped peer memory; a cross-node +transport would need an explicit fence here. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch + +from afd_plugin.connectors.gpu import nvshmem_rt + +if TYPE_CHECKING: + from torch.distributed.distributed_c10d import ProcessGroup + +# Header words shared by dispatch and combine, followed by expert_counts. +HEADER_MAGIC = 0x41464447 # "AFDG" +HEADER_VERSION = 1 +_H_MAGIC = 0 +_H_VERSION = 1 +_H_SEQ = 2 +_H_SRC_ROLE_RANK = 3 +_H_LAYER_IDX = 4 +_H_STAGE_IDX = 5 +_H_NUM_TOKENS = 6 +_H_ROUTED_TOKENS = 7 +_H_SHARED_TOKENS = 8 +_H_TOPK = 9 +_H_FLAGS = 10 +_H_ECHO_SEQ = 11 # combine only: the dispatch seq being answered +HEADER_FIXED_WORDS = 12 + +FLAG_EMPTY = 0 +FLAG_SHUTDOWN_BIT = 1 << 1 + +# Every field starts on this boundary so a byte offset stays divisible by the +# element size of whichever dtype views it. +_FIELD_ALIGN = 256 + + +def _align(offset: int) -> int: + return (offset + _FIELD_ALIGN - 1) // _FIELD_ALIGN * _FIELD_ALIGN + + +@dataclass(frozen=True, slots=True) +class SlotLayout: + """Byte offsets and element counts for the fields inside one slot.""" + + header_words: int + routed_cap: int + token_cap: int + hidden_size: int + payload_itemsize: int + + header_off: int + route_table_off: int + routed_x_off: int + shared_idx_off: int + shared_x_off: int + slot_bytes: int + + @classmethod + def build( + cls, + *, + expert_per_rank: int, + routed_cap: int, + token_cap: int, + hidden_size: int, + payload_itemsize: int, + ) -> SlotLayout: + header_words = HEADER_FIXED_WORDS + expert_per_rank + header_off = 0 + route_table_off = _align(header_off + header_words * 4) + routed_x_off = _align(route_table_off + 2 * routed_cap * 4) + shared_idx_off = _align( + routed_x_off + routed_cap * hidden_size * payload_itemsize + ) + shared_x_off = _align(shared_idx_off + token_cap * 4) + slot_bytes = _align(shared_x_off + token_cap * hidden_size * payload_itemsize) + return cls( + header_words=header_words, + routed_cap=routed_cap, + token_cap=token_cap, + hidden_size=hidden_size, + payload_itemsize=payload_itemsize, + header_off=header_off, + route_table_off=route_table_off, + routed_x_off=routed_x_off, + shared_idx_off=shared_idx_off, + shared_x_off=shared_x_off, + slot_bytes=slot_bytes, + ) + + +def encode_header( + layout: SlotLayout, + *, + seq: int, + src_role_rank: int, + layer_idx: int, + stage_idx: int, + num_tokens: int, + routed_tokens: int, + shared_tokens: int, + topk: int, + flags: int, + expert_counts: list[int], + echo_seq: int = 0, +) -> torch.Tensor: + """Build the fixed header for one slot as a CPU int32 tensor.""" + expert_per_rank = layout.header_words - HEADER_FIXED_WORDS + if len(expert_counts) != expert_per_rank: + raise ValueError( + f"expert_counts must have {expert_per_rank} entries, " + f"got {len(expert_counts)}", + ) + header = torch.zeros(layout.header_words, dtype=torch.int32) + header[_H_MAGIC] = HEADER_MAGIC + header[_H_VERSION] = HEADER_VERSION + header[_H_SEQ] = seq + header[_H_SRC_ROLE_RANK] = src_role_rank + header[_H_LAYER_IDX] = layer_idx + header[_H_STAGE_IDX] = stage_idx + header[_H_NUM_TOKENS] = num_tokens + header[_H_ROUTED_TOKENS] = routed_tokens + header[_H_SHARED_TOKENS] = shared_tokens + header[_H_TOPK] = topk + header[_H_FLAGS] = flags + header[_H_ECHO_SEQ] = echo_seq + if expert_per_rank: + header[HEADER_FIXED_WORDS:] = torch.tensor(expert_counts, dtype=torch.int32) + return header + + +@dataclass(frozen=True, slots=True) +class SlotHeader: + """Decoded slot header.""" + + seq: int + src_role_rank: int + layer_idx: int + stage_idx: int + num_tokens: int + routed_tokens: int + shared_tokens: int + topk: int + flags: int + echo_seq: int + expert_counts: list[int] + + @property + def is_shutdown(self) -> bool: + return bool(self.flags & FLAG_SHUTDOWN_BIT) + + +def decode_header(header: torch.Tensor) -> SlotHeader: + """Decode a CPU int32 header tensor, validating magic and version.""" + values = header.tolist() + if values[_H_MAGIC] != HEADER_MAGIC: + raise RuntimeError( + f"AFD async GPU header magic mismatch: got {values[_H_MAGIC]:#x}, " + f"expected {HEADER_MAGIC:#x}", + ) + if values[_H_VERSION] != HEADER_VERSION: + raise RuntimeError( + f"AFD async GPU header version {values[_H_VERSION]} is not supported " + f"(expected {HEADER_VERSION})", + ) + return SlotHeader( + seq=values[_H_SEQ], + src_role_rank=values[_H_SRC_ROLE_RANK], + layer_idx=values[_H_LAYER_IDX], + stage_idx=values[_H_STAGE_IDX], + num_tokens=values[_H_NUM_TOKENS], + routed_tokens=values[_H_ROUTED_TOKENS], + shared_tokens=values[_H_SHARED_TOKENS], + topk=values[_H_TOPK], + flags=values[_H_FLAGS], + echo_seq=values[_H_ECHO_SEQ], + expert_counts=values[HEADER_FIXED_WORDS:], + ) + + +@dataclass(slots=True) +class ArrivedSlot: + """One arrival found by ``SymmWindow.poll``.""" + + region: int + ring: int + header: SlotHeader + + +class SymmWindow: + """Symmetric receive window plus the one-sided writes that fill peers'.""" + + def __init__( + self, + *, + num_regions: int, + ring_depth: int, + layout: SlotLayout, + payload_dtype: torch.dtype, + device: torch.device, + group: ProcessGroup, + rank: int, + world_size: int, + ) -> None: + if num_regions <= 0 or ring_depth <= 0: + raise ValueError("num_regions and ring_depth must be positive") + self.num_regions = num_regions + self.ring_depth = ring_depth + self.layout = layout + self.payload_dtype = payload_dtype + self.device = device + self.rank = rank + + self.num_flags = num_regions * ring_depth + self._flag_bytes = _align(self.num_flags * 4) + self.total_bytes = self._flag_bytes + self.num_flags * layout.slot_bytes + + nvshmem_rt.init(group, rank, world_size) + self._base = nvshmem_rt.malloc(self.total_bytes) + # Peer mappings are stable for the life of the allocation, so resolve + # them once instead of per transfer. + self._peer_base = { + pe: (self._base if pe == rank else nvshmem_rt.peer_ptr(self._base, pe)) + for pe in range(world_size) + } + self.local_bytes_view().zero_() + torch.cuda.synchronize() + + # The layout is static, so every window view is built once and then + # sliced. Rebuilding them per transfer meant a __cuda_array_interface__ + # import on every field of every message, which dominated the data path. + self._view_cache: dict[tuple[int, int, int, int], torch.Tensor] = {} + self._flag_cache: dict[tuple[int, int], torch.Tensor] = {} + self._flags_local = nvshmem_rt.tensor_from_ptr( + self._base, + byte_offset=0, + sizes=(self.num_flags,), + dtype=torch.int32, + device=device, + ) + + # Pinned staging keeps header transfers asynchronous; a pageable source + # forces a blocking copy, and there is one header per peer per layer. + # One row per (peer, ring): a single shared row would be overwritten on + # the host by the next peer in the send loop while its own asynchronous + # copy was still in flight, delivering another peer's token counts. + self._header_send = torch.zeros( + (world_size, ring_depth, layout.header_words), + dtype=torch.int32, + ).pin_memory() + self._header_recv = torch.zeros( + layout.header_words, + dtype=torch.int32, + ).pin_memory() + # Host mirror of the local flag array; one D2H per poll refreshes it. + self._flag_host = torch.zeros(self.num_flags, dtype=torch.int32).pin_memory() + self._seen = [FLAG_EMPTY] * self.num_flags + + def local_bytes_view(self) -> torch.Tensor: + return nvshmem_rt.tensor_from_ptr( + self._base, + byte_offset=0, + sizes=(self.total_bytes,), + dtype=torch.uint8, + device=self.device, + ) + + def _slot_byte_off(self, region: int, ring: int) -> int: + return ( + self._flag_bytes + + (region * self.ring_depth + ring) * self.layout.slot_bytes + ) + + def _capacity_view( + self, + peer: int, + region: int, + ring: int, + field_off: int, + ) -> torch.Tensor: + """Return the cached full-capacity view of one slot field.""" + key = (peer, region, ring, field_off) + view = self._view_cache.get(key) + if view is not None: + return view + + layout = self.layout + hidden = layout.hidden_size + if field_off == layout.header_off: + sizes, dtype = (layout.header_words,), torch.int32 + elif field_off == layout.route_table_off: + sizes, dtype = (layout.routed_cap, 2), torch.int32 + elif field_off == layout.routed_x_off: + sizes, dtype = (layout.routed_cap, hidden), self.payload_dtype + elif field_off == layout.shared_idx_off: + sizes, dtype = (layout.token_cap,), torch.int32 + elif field_off == layout.shared_x_off: + sizes, dtype = (layout.token_cap, hidden), self.payload_dtype + else: + raise ValueError(f"unknown slot field offset {field_off}") + + view = nvshmem_rt.tensor_from_ptr( + self._peer_base[peer], + byte_offset=self._slot_byte_off(region, ring) + field_off, + sizes=sizes, + dtype=dtype, + device=self.device, + ) + self._view_cache[key] = view + return view + + def _view( + self, + peer: int, + region: int, + ring: int, + field_off: int, + sizes: tuple[int, ...], + dtype: torch.dtype, + ) -> torch.Tensor: + # Slicing a cached capacity view costs no CUDA calls, unlike importing + # a fresh pointer for every field of every message. + return self._capacity_view(peer, region, ring, field_off)[: sizes[0]] + + def _flag_view(self, peer: int, flag_idx: int) -> torch.Tensor: + key = (peer, flag_idx) + view = self._flag_cache.get(key) + if view is None: + view = nvshmem_rt.tensor_from_ptr( + self._peer_base[peer], + byte_offset=flag_idx * 4, + sizes=(1,), + dtype=torch.int32, + device=self.device, + ) + self._flag_cache[key] = view + return view + + # ------------------------------------------------------------------ + # Send side: every write targets ``peer``'s window, one-sided. + # ------------------------------------------------------------------ + + def write_slot( + self, + *, + peer: int, + region: int, + ring: int, + header: torch.Tensor, + route_table: torch.Tensor | None, + routed_x: torch.Tensor | None, + shared_idx: torch.Tensor | None, + shared_x: torch.Tensor | None, + ) -> None: + """Write one slot into ``peer``'s window, then stamp its flag. + + The flag copy is issued last on the same stream, so a peer that observes + the flag also observes the payload. + """ + layout = self.layout + if routed_x is not None and routed_x.shape[0] > layout.routed_cap: + raise RuntimeError( + f"routed tokens {routed_x.shape[0]} exceed routed_cap " + f"{layout.routed_cap}; raise routed_cap_multiplier", + ) + if shared_x is not None and shared_x.shape[0] > layout.token_cap: + raise RuntimeError( + f"shared tokens {shared_x.shape[0]} exceed token_cap " + f"{layout.token_cap}", + ) + + # Stage through pinned memory so the copy is asynchronous: a pageable + # source would force a blocking transfer, and there is one header per + # peer per layer. + staging = self._header_send[peer][ring] + staging.copy_(header) + self._capacity_view(peer, region, ring, layout.header_off).copy_( + staging, + non_blocking=True, + ) + + if route_table is not None and route_table.numel(): + n = route_table.shape[0] + self._view( + peer, + region, + ring, + layout.route_table_off, + (n, 2), + torch.int32, + ).copy_(route_table, non_blocking=True) + if routed_x is not None and routed_x.numel(): + n = routed_x.shape[0] + self._view( + peer, + region, + ring, + layout.routed_x_off, + (n, layout.hidden_size), + self.payload_dtype, + ).copy_(routed_x, non_blocking=True) + if shared_idx is not None and shared_idx.numel(): + n = shared_idx.shape[0] + self._view( + peer, + region, + ring, + layout.shared_idx_off, + (n,), + torch.int32, + ).copy_(shared_idx, non_blocking=True) + if shared_x is not None and shared_x.numel(): + n = shared_x.shape[0] + self._view( + peer, + region, + ring, + layout.shared_x_off, + (n, layout.hidden_size), + self.payload_dtype, + ).copy_(shared_x, non_blocking=True) + + seq = int(header[_H_SEQ].item()) + flag_idx = region * self.ring_depth + ring + self._flag_view(peer, flag_idx).fill_(seq) + + # ------------------------------------------------------------------ + # Receive side. + # ------------------------------------------------------------------ + + def poll(self) -> ArrivedSlot | None: + """Return the first slot whose flag advanced past what we consumed. + + ponytail: host-side poll, one D2H per call. Correct but it burns a + synchronize per attempt; replace with a device-side ``wait_any`` kernel + spinning on the flag array when the poll shows up in a profile. + """ + self._flag_host.copy_(self._flags_local, non_blocking=False) + host = self._flag_host.tolist() + for idx in range(self.num_flags): + if host[idx] != self._seen[idx]: + self._seen[idx] = host[idx] + region, ring = divmod(idx, self.ring_depth) + return ArrivedSlot( + region=region, + ring=ring, + header=self.read_header(region, ring), + ) + return None + + def read_header(self, region: int, ring: int) -> SlotHeader: + # Reuse the pinned mirror instead of allocating a fresh host tensor on + # every arrival. + self._header_recv.copy_( + self._capacity_view(self.rank, region, ring, self.layout.header_off), + ) + return decode_header(self._header_recv) + + def local_route_table(self, region: int, ring: int, count: int) -> torch.Tensor: + return self._view( + self.rank, + region, + ring, + self.layout.route_table_off, + (count, 2), + torch.int32, + ) + + def local_routed(self, region: int, ring: int, count: int) -> torch.Tensor: + return self._view( + self.rank, + region, + ring, + self.layout.routed_x_off, + (count, self.layout.hidden_size), + self.payload_dtype, + ) + + def local_shared_idx(self, region: int, ring: int, count: int) -> torch.Tensor: + return self._view( + self.rank, + region, + ring, + self.layout.shared_idx_off, + (count,), + torch.int32, + ) + + def local_shared(self, region: int, ring: int, count: int) -> torch.Tensor: + return self._view( + self.rank, + region, + ring, + self.layout.shared_x_off, + (count, self.layout.hidden_size), + self.payload_dtype, + ) + + def close(self) -> None: + # ponytail: the symmetric allocation is left to process teardown. + # nvshmem_free is collective, so freeing here would need both roles to + # shut down in lockstep; add it if windows are ever recreated in-process. + self._peer_base = {} diff --git a/afd_plugin/connectors/npu/async_cam.py b/afd_plugin/connectors/npu/async_cam.py index 00dd128a..57370875 100644 --- a/afd_plugin/connectors/npu/async_cam.py +++ b/afd_plugin/connectors/npu/async_cam.py @@ -45,6 +45,12 @@ coerce_extra_positive_int, coerce_extra_str, ) +from afd_plugin.connectors.async_topology import ( + ASYNC_MOE_REQUEST_SPLIT, + ATTN_RANKS_PER_DP_CONFIG_KEY, + AFDAsyncTopology, + build_async_topology, +) from afd_plugin.connectors.base import ( AFDConnectorBase, ConnectorExtraInfo, @@ -64,9 +70,7 @@ AFD_ASYNC_CAM_GROUP_NAME = "afd_async_cam" CAM_COMM_ID = 0 -ATTN_RANKS_PER_DP_CONFIG_KEY = "attn_ranks_per_dp" ASYNC_MOE_NUM_STAGES = 2 -ASYNC_MOE_REQUEST_SPLIT = "request" ASYNC_MOE_TOKEN_SPLIT = "token" _AFD_ASYNC_EXTRA_CONFIG_FIELDS: Final[frozenset[str]] = frozenset( @@ -190,23 +194,6 @@ class AFDAsyncFFNWorkItem: shared_num_tokens: int -@dataclass(frozen=True, slots=True) -class AFDAsyncTopology: - """Role-local and HCCL-world rank information for one CAM participant.""" - - role: str - role_rank: int - world_rank: int - attn_size: int - ffn_size: int - expert_per_rank: int - - @property - def world_size(self) -> int: - """Return the total number of Attention and FFN ranks.""" - return self.attn_size + self.ffn_size - - class CAMAsyncAFDConnector(AFDConnectorBase): """CAM-backed asynchronous connector for Ascend NPU AFD. @@ -850,56 +837,6 @@ def _log_cam_op_values(op_name: str, label: str, **kwargs: object) -> None: logger.warning("AFD CAM %s %s:\n%s", op_name, label, "\n".join(lines)) -def build_async_topology( - afd_config: AFDConfig, - role_rank: int, - *, - num_routed_experts: int | None = None, -) -> AFDAsyncTopology: - """Validate role-local rank settings and derive the CAM HCCL world rank. - - The world is Attention-first: Attention role rank ``i`` maps to world rank - ``i`` and FFN role rank ``j`` maps to - ``num_attention_ranks + j``. Routed experts are distributed across FFN - ranks using a ceiling division; production model layouts should keep the - routed-expert count divisible by the FFN rank count. - """ - attn_size = afd_config.num_attention_ranks - ffn_size = afd_config.num_ffn_ranks - if attn_size <= 0 or ffn_size <= 0: - raise ValueError("AFD async topology sizes must be positive") - if role_rank < 0: - raise ValueError(f"AFD async role rank must be non-negative, got {role_rank}") - - if afd_config.role == "attention": - if role_rank >= attn_size: - raise ValueError( - "Attention role rank must be within attention size " - f"(rank={role_rank}, size={attn_size})", - ) - world_rank = role_rank - elif afd_config.role == "ffn": - if role_rank >= ffn_size: - raise ValueError( - "FFN role rank must be within FFN size " - f"(rank={role_rank}, size={ffn_size})", - ) - world_rank = attn_size + role_rank - else: - raise ValueError(f"unknown AFD role {afd_config.role!r}") - - expert_count = num_routed_experts or 1 - expert_per_rank = (expert_count + ffn_size - 1) // ffn_size - return AFDAsyncTopology( - role=afd_config.role, - role_rank=role_rank, - world_rank=world_rank, - attn_size=attn_size, - ffn_size=ffn_size, - expert_per_rank=expert_per_rank, - ) - - def _validate_topk_payload( topk_ids: Tensor, topk_weights: Tensor | None, diff --git a/afd_plugin/model_executor/models/deepseek_v2.py b/afd_plugin/model_executor/models/deepseek_v2.py index 90ee2462..5f717f6f 100644 --- a/afd_plugin/model_executor/models/deepseek_v2.py +++ b/afd_plugin/model_executor/models/deepseek_v2.py @@ -21,7 +21,7 @@ from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.models import deepseek_v2 as native -from afd_plugin.config import AFD_ASYNC_CONNECTOR, parse_afd_config +from afd_plugin.config import AFD_ASYNC_CONNECTORS, parse_afd_config from afd_plugin.connectors import ( AFDExpertRoutingSpec, AFDF2ATransferPayload, @@ -512,7 +512,8 @@ def compute_attn_output( topk_weights = None topk_ids = None router_logits = None - # NPU-only: Attention-side gate/topk is implemented in the NPU helper. + # The gate helper delegates expert selection to the connector, so both + # platforms share it despite the module's location. if self.compute_gate_on_attention and self.is_moe_layer: from afd_plugin.model_executor.models.npu import ( deepseek_v2_attention_gate, @@ -568,8 +569,23 @@ def compute_ffn_output( ) return output if self.compute_gate_on_attention: - raise RuntimeError( - "GPU Attention-side gate must call compute_experts_output", + if group_list is None: + # Without a group list the caller is the control-plane path, + # which routes on this side and must use compute_experts_output. + raise RuntimeError( + "GPU Attention-side gate must call compute_experts_output", + ) + # Token-level dispatch: rows arrive pre-routed and grouped by local + # expert, so only the grouped GEMM is left to run here. + from afd_plugin.model_executor.models.gpu import ( + deepseek_v2_attention_gate as gpu_attention_gate, + ) + + return gpu_attention_gate.compute_attention_gate_moe_ffn( + self, + hidden_states=hidden_states, + group_list=group_list, + expand_x_shared=expand_x_shared, ) hidden_states = self.mlp(hidden_states) if ( @@ -703,7 +719,10 @@ def forward( intermediate_tensors: native.IntermediateTensors | None, inputs_embeds: torch.Tensor | None = None, ) -> torch.Tensor | native.IntermediateTensors: - if self.afd_config.connector == AFD_ASYNC_CONNECTOR: + if self.afd_config.connector in AFD_ASYNC_CONNECTORS: + # The schedule below is platform-neutral -- it only drives the model + # and the connector interface -- so both async connectors share it + # despite the module still living under the npu package. from afd_plugin.model_executor.models.npu import ( deepseek_v2_async_cam_forward, ) diff --git a/afd_plugin/model_executor/models/gpu/__init__.py b/afd_plugin/model_executor/models/gpu/__init__.py new file mode 100644 index 00000000..6a15edcf --- /dev/null +++ b/afd_plugin/model_executor/models/gpu/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""GPU-specific AFD model wrappers.""" diff --git a/afd_plugin/model_executor/models/gpu/deepseek_v2_attention_gate.py b/afd_plugin/model_executor/models/gpu/deepseek_v2_attention_gate.py new file mode 100644 index 00000000..75764c81 --- /dev/null +++ b/afd_plugin/model_executor/models/gpu/deepseek_v2_attention_gate.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Attention-side gate helpers for DeepSeek-V2 on CUDA. + +The async GPU connector dispatches tokens that are already routed: one row per +``(token, topk_slot)`` partial, grouped by local expert, with a ``group_list`` +of per-expert counts. The FFN side therefore must not run routing again -- it +only needs the grouped GEMM over its local experts. + +That shape is a ``topk == 1`` problem: give every arriving row its own expert id +and a unit weight, and vLLM's ``fused_experts`` computes exactly the local +expert output. Topk weighting stays on the Attention side, applied during +combine, matching the NPU path. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts + +from afd_plugin.connectors.metadata import AFDF2ATransferPayload + +if TYPE_CHECKING: + from torch import nn + + +def compute_attention_gate_moe_ffn( + layer: nn.Module, + *, + hidden_states: torch.Tensor, + group_list: torch.Tensor, + expand_x_shared: torch.Tensor | None = None, +) -> AFDF2ATransferPayload: + """Run this rank's local experts over pre-routed tokens. + + Args: + layer: The AFD DeepSeek decoder layer owning the MoE module. + hidden_states: ``[num_partials, hidden]`` rows sorted by local expert. + group_list: ``[expert_per_rank]`` per-expert row counts. Must sum to + ``hidden_states.shape[0]``. + expand_x_shared: Optional ``[num_shared, hidden]`` shared-expert rows. + + Returns: + Routed output in the same row order as ``hidden_states``, plus the + shared-expert output when the model has shared experts. + """ + # ``mlp.experts`` is a MoERunner; the weight parameters live on its + # RoutedExperts, and the shared experts hang off the runner rather than + # off the MoE module. + runner = layer.mlp.experts + routed_experts = runner.routed_experts + counts = group_list.to(torch.int64) + num_rows = int(hidden_states.shape[0]) + if int(counts.sum()) != num_rows: + raise ValueError( + f"group_list sums to {int(counts.sum())} but hidden_states has " + f"{num_rows} rows", + ) + + num_local_experts = counts.numel() + if num_rows == 0: + # Routing can leave a peer with nothing -- common in decode, where a + # single token's topk may land entirely on the other FFN rank. + routed_output = hidden_states.new_empty((0, hidden_states.shape[-1])) + shared = runner._shared_experts + return AFDF2ATransferPayload( + routed_output=routed_output, + shared_output=( + shared._layer(expand_x_shared) + if shared is not None + and expand_x_shared is not None + and expand_x_shared.shape[0] > 0 + else None + ), + ) + expert_ids = torch.repeat_interleave( + torch.arange( + num_local_experts, + device=hidden_states.device, + dtype=torch.int32, + ), + counts, + ).unsqueeze(1) + # Unit weights: the real topk weighting happens in the connector's combine. + unit_weights = torch.ones( + (num_rows, 1), + dtype=torch.float32, + device=hidden_states.device, + ) + + routed_output = fused_experts( + hidden_states, + routed_experts.w13_weight, + routed_experts.w2_weight, + unit_weights, + expert_ids, + global_num_experts=num_local_experts, + expert_map=None, + ) + + shared_output = None + shared_experts = runner._shared_experts + if shared_experts is not None and expand_x_shared is not None: + # Call the wrapped MLP rather than SharedExperts.forward: the wrapper is + # a stateful scheduler for the runner's own multi-stream pipeline and + # returns None unless its expected ordering matches. AFD feeds shared + # tokens as their own batch, so that machinery does not apply. + shared_output = shared_experts._layer(expand_x_shared) + + # Mirrors the NPU gate path: scale the routed branch unless fp16, where the + # native model instead scales the shared branch down. + routed_scaling_factor = runner.routed_scaling_factor + if hidden_states.dtype != torch.float16: + routed_output = routed_output * routed_scaling_factor + elif shared_output is not None: + shared_output = shared_output * (1.0 / routed_scaling_factor) + + return AFDF2ATransferPayload( + routed_output=routed_output, + shared_output=shared_output, + ) + + +__all__ = ["compute_attention_gate_moe_ffn"] diff --git a/afd_plugin/v1/worker/attention_model_runner.py b/afd_plugin/v1/worker/attention_model_runner.py index b8ed0d9d..ff6b0e3b 100644 --- a/afd_plugin/v1/worker/attention_model_runner.py +++ b/afd_plugin/v1/worker/attention_model_runner.py @@ -51,6 +51,36 @@ from vllm.v1.core.sched.output import SchedulerOutput +@contextmanager +def _dp_batch_coordination_disabled(disabled: bool): + """Skip vLLM's cross-DP batch agreement for connector-driven runs. + + ``GPUModelRunner._determine_batch_execution_and_padding`` all-reduces the + batch shape across the DP group whenever ``data_parallel_size > 1``. Async + AFD deliberately lets each Attention replica advance on its own, so an idle + replica never joins that collective and a busy one blocks in it forever -- + which is where a 2A2F run hangs before it reaches the first MoE layer. + + Returning the single-rank answer (``num_tokens_across_dp=None``) makes the + upstream function skip its DP-padding branch entirely, exactly as it does + for ``data_parallel_size == 1``. + """ + if not disabled: + yield + return + + original = gpu_model_runner.coordinate_batch_across_dp + + def _single_rank_coordination(*_args: Any, cudagraph_mode: int, **_kwargs: Any): + return False, None, cudagraph_mode + + gpu_model_runner.coordinate_batch_across_dp = _single_rank_coordination + try: + yield + finally: + gpu_model_runner.coordinate_batch_across_dp = original + + class AFDAttentionModelRunner(GPUModelRunner): """Attention model runner that injects AFD metadata into forward context.""" @@ -76,11 +106,9 @@ def __init__( self.afd_config, ) # The connector rendezvous is deferred to the end of ``load_model()`` - # so Attention and FFN weight loading overlap; see that method. - # TODO: Async GPU connector will be supported in the future - assert self.connector.control_plane is not None, ( - "GPU model runner only supports control-plane-driven connectors" - ) + # so Attention and FFN weight loading overlap; see that method. The + # async GPU connector drives FFN work from its own receive loop and so + # has no control plane, which is why there is no assertion here. self._is_warmup = False self._afd_is_graph_capturing = False self._afd_pending_metadata: AFDForwardContextMetadata | None = None @@ -126,9 +154,8 @@ def _send_dp_metadata( dp_metadata: DPMetadata | AFDDPMetadata | None, ubatch_slices: Any, ) -> None: - assert self.connector.control_plane is not None, ( - "_send_dp_metadata needs control plane driven connectors" - ) + if self.connector.control_plane is None: + return if ubatch_slices and len(ubatch_slices) > 1: dp_metadata_list = { @@ -349,25 +376,28 @@ def _determine_batch_execution_and_padding( torch.Tensor | None, CUDAGraphStat | None, ]: - ( - cudagraph_mode, - batch_descriptor, - should_ubatch, - num_tokens_across_dp, - cudagraph_stats, - ) = super()._determine_batch_execution_and_padding( - num_tokens, - num_reqs, - num_scheduled_tokens_np, - max_num_scheduled_tokens, - use_cascade_attn, - allow_microbatching, - force_eager, - force_uniform_decode, - force_has_lora, - force_num_active_loras, - num_encoder_reqs, - ) + with _dp_batch_coordination_disabled( + self.connector.control_plane is None, + ): + ( + cudagraph_mode, + batch_descriptor, + should_ubatch, + num_tokens_across_dp, + cudagraph_stats, + ) = super()._determine_batch_execution_and_padding( + num_tokens, + num_reqs, + num_scheduled_tokens_np, + max_num_scheduled_tokens, + use_cascade_attn, + allow_microbatching, + force_eager, + force_uniform_decode, + force_has_lora, + force_num_active_loras, + num_encoder_reqs, + ) args = ( num_tokens, diff --git a/afd_plugin/v1/worker/ffn_model_runner.py b/afd_plugin/v1/worker/ffn_model_runner.py index a4deb15e..d22b220e 100644 --- a/afd_plugin/v1/worker/ffn_model_runner.py +++ b/afd_plugin/v1/worker/ffn_model_runner.py @@ -30,6 +30,11 @@ AFDControlPayload, AFDDPMetadata, ) +from afd_plugin.connectors.gpu.async_gpu import ( + ConnectorShutdown, + GpuAsyncTransferState, +) +from afd_plugin.connectors.metadata import AFDF2ATransferPayload from afd_plugin.v1.worker.attention_model_runner import ( fail_if_unsupported_ubatching, ) @@ -53,10 +58,12 @@ class GPUFFNModelRunner(LoRAModelRunnerMixin): """FFN model runner for AFD GPU execution. - FFN steps are driven by the connector control plane rather than the vLLM - scheduler. GPU only supports control-plane-driven connectors, so the runner - asserts ``connector.control_plane is not None`` at construction; connectors - without a control plane (``control_plane is None``) are not supported. + FFN steps are driven by the connector rather than the vLLM scheduler, in one + of two ways. Control-plane connectors receive broadcast DP metadata and then + walk every layer in lockstep with the Attention side. Connectors without a + control plane (``control_plane is None``) instead pull one work item at a + time from their receive loop, learning the layer and token counts from the + arriving payload; see ``execute_connector_driven_step``. """ afd_expected_role = "ffn" @@ -80,10 +87,9 @@ def __init__(self, vllm_config: VllmConfig, device: object) -> None: vllm_config, self.afd_config, ) - # TODO: Async GPU connector will be supported in the future - assert self.connector.control_plane is not None, ( - "GPU model runner only supports control-plane-driven connectors" - ) + # A connector without a control plane drives FFN steps from its own + # receive loop instead of from broadcast DP metadata. + self.is_connector_driven = self.connector.control_plane is None self.model: Any | None = None self.model_memory_usage = 0 @@ -237,6 +243,61 @@ def _ffn_forward( self.connector.send_ffn_output(rank_ffn_output, context) return rank_ffn_output + def execute_connector_driven_step(self) -> None: + """Drain whatever the connector has already received, then return. + + Returning on an idle poll rather than blocking forever is what lets the + worker loop observe its shutdown event. The batch size below is only a + drain granularity: successive work items may belong to different layers + of different Attention replicas. + """ + step_afd_gpu_profiler(self.prof) + self._ffn_forward_connector_driven() + + def _ffn_forward_connector_driven( + self, + ) -> torch.Tensor | AFDF2ATransferPayload | None: + stage_idx = 0 + rank_ffn_output = None + connector = self.connector + max_items = max(1, int(self.num_layers)) + + with _ffn_forward_context(self.vllm_config) as forward_context: + for _ in range(max_items): + try: + work_item = connector.recv_ffn_work_item( + stage_idx=stage_idx, + max_num_tokens=self.vllm_config.scheduler_config.max_num_batched_tokens, + ) + except TimeoutError: + # Nothing pending; hand control back so the worker loop can + # check for shutdown. + return rank_ffn_output + except ConnectorShutdown: + raise + + states = work_item.context.states + if not isinstance(states, GpuAsyncTransferState): + raise RuntimeError( + "async GPU FFN work item requires GpuAsyncTransferState", + ) + metadata = work_item.context.metadata + forward_context.dp_metadata = None + forward_context.additional_kwargs["afd_metadata"] = metadata + _set_moe_layer_index(forward_context, work_item.layer_idx) + + rank_ffn_output = self.model.compute_ffn_output( + hidden_states=work_item.hidden_states, + layer_idx=work_item.layer_idx, + group_list=states.group_list, + expand_x_shared=states.expand_x_shared, + ) + rank_ffn_output = connector.send_ffn_work_item_output( + work_item, + rank_ffn_output, + ) + return rank_ffn_output + def _execute_eager_mode( self, hidden_states: torch.Tensor, diff --git a/afd_plugin/v1/worker/ffn_worker.py b/afd_plugin/v1/worker/ffn_worker.py index 1f6ce4d7..69ddca80 100644 --- a/afd_plugin/v1/worker/ffn_worker.py +++ b/afd_plugin/v1/worker/ffn_worker.py @@ -13,6 +13,7 @@ from vllm.v1.worker.gpu_worker import Worker from vllm.v1.worker.worker_base import CompilationTimes +from afd_plugin.connectors.gpu.async_gpu import ConnectorShutdown from afd_plugin.model_executor.models.model_utils import get_afd_model_config from afd_plugin.v1.worker.attention_model_runner import fail_if_unsupported_ubatching from afd_plugin.v1.worker.ffn_model_runner import GPUFFNModelRunner @@ -128,6 +129,13 @@ def ffn_worker_loop() -> None: try: self._run_ffn_server_loop() except Exception as exc: + shutdown_event = self._ffn_shutdown_event + if shutdown_event is not None and shutdown_event.is_set(): + logger.debug( + "AFD FFN receive loop stopped during shutdown", + exc_info=True, + ) + return self._ffn_loop_error = exc logger.exception("AFD FFN worker loop failed") @@ -148,11 +156,17 @@ def _run_ffn_server_loop(self) -> None: while not event.is_set(): if self.model_runner.connector.control_plane is None: - raise NotImplementedError( - "GPU FFN only supports control-plane-driven connectors; " - "connectors without a control plane (control_plane is None) " - "are not supported.", - ) + # Connector-driven: the step returns on an idle poll, so the + # loop gets to re-check the shutdown event. No device-wide + # synchronize here -- it would serialize every receive against + # the previous compute and erase the overlap this path exists + # for; ordering is carried by the connector's own streams. + try: + self.model_runner.execute_connector_driven_step() + except ConnectorShutdown: + logger.info("AFD FFN loop exiting: peer announced shutdown") + return + continue payload = self.model_runner.connector.control_plane.recv_dp_metadata_list() dp_metadata_list = payload.dp_metadata_list diff --git a/afd_plugin/v1/worker/npu/attention_model_runner.py b/afd_plugin/v1/worker/npu/attention_model_runner.py index e5fd8789..b4bc0677 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner.py @@ -74,7 +74,7 @@ stop_afd_npu_profiler, ) from afd_plugin.config import ( - AFD_ASYNC_CONNECTOR, + AFD_ASYNC_NPU_CONNECTOR, AFDConfig, parse_afd_config, ) @@ -143,7 +143,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.afd_config, ) self.afd_async_extra_info = AFDAsyncExtraInfo() - if afd_config.connector == AFD_ASYNC_CONNECTOR: + if afd_config.connector == AFD_ASYNC_NPU_CONNECTOR: connector_extra_info = self.connector.extra_info if not isinstance(connector_extra_info, AFDAsyncExtraInfo): raise TypeError( diff --git a/pyproject.toml b/pyproject.toml index 6ca7cbdc..c3aa7f3a 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/recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/1a1f_eager_async.sh b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/1a1f_eager_async.sh new file mode 100644 index 00000000..b3b74068 --- /dev/null +++ b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/1a1f_eager_async.sh @@ -0,0 +1,128 @@ +#!/usr/bin/env bash +# 1A1F async GPU connector, eager, prefill-only. +# +# Launch under a GPU reservation, which sets CUDA_VISIBLE_DEVICES: +# gpu run --gpu-ids 3,7 -- bash recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/1a1f_eager_async.sh +# +# The two roles are separate vllm serve processes: the AFD process group hosts +# its own TCPStore, which cannot be created under a single torchrun/torchelastic +# launcher. +set -u + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} +LOG_DIR=${LOG_DIR:-.} +mkdir -p "$LOG_DIR" +export VLLM_USE_V2_MODEL_RUNNER=0 +# Single node over NVLink: skip the IB transport probe. +export NVSHMEM_REMOTE_TRANSPORT=${NVSHMEM_REMOTE_TRANSPORT:-none} +# Two servers on one box spawn a lot of threads; the HF tokenizer's rayon pool +# is the first thing to fail when thread creation gets refused. +export TOKENIZERS_PARALLELISM=${TOKENIZERS_PARALLELISM:-false} +export RAYON_NUM_THREADS=${RAYON_NUM_THREADS:-2} +export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4} + +IFS=',' read -r -a DEVICES <<< "${CUDA_VISIBLE_DEVICES:-0,1}" +if [ "${#DEVICES[@]}" -lt 2 ]; then + echo "need 2 visible GPUs, got ${#DEVICES[@]}: ${CUDA_VISIBLE_DEVICES:-unset}" >&2 + exit 1 +fi +ATTN_DEVICES="${DEVICES[0]}" +FFN_DEVICES="${DEVICES[1]}" +echo "attention on ${ATTN_DEVICES}, ffn on ${FFN_DEVICES}" + +# Lower this when sharing a box: vLLM refuses to start if the desired +# fraction exceeds what is actually free. +GPU_MEM_UTIL=${GPU_MEM_UTIL:-0.9} +# Prefill batch size drives whether each MoE call clears the compute-bound +# inflection point, so it is the knob to raise when benchmarking. +MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-512} +MAX_NUM_SEQS=${MAX_NUM_SEQS:-8} +AFD_PORT=${AFD_PORT:-6275} +API_PORT=${API_PORT:-18311} + +CUDA_VISIBLE_DEVICES="$ATTN_DEVICES" uv run vllm serve "$MODEL_PATH" \ + --data-parallel-size 1 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --additional-config '{ + "afd": { + "role": "attention", + "connector": "GpuAsyncAFDConnector", + "async": true, + "compute_gate_on_attention": true, + "host": "127.0.0.1", + "port": '"$AFD_PORT"', + "num_attention_ranks": 1, + "num_ffn_ranks": 1 + } + }' \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \ + --api-server-count 1 \ + --gpu-memory-utilization "$GPU_MEM_UTIL" \ + --enforce-eager \ + --host 127.0.0.1 \ + --port "$API_PORT" \ + --trust-remote-code > "$LOG_DIR/attn.log" 2>&1 & +ATTN_PID=$! + +CUDA_VISIBLE_DEVICES="$FFN_DEVICES" uv run vllm serve "$MODEL_PATH" \ + --data-parallel-size 1 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --additional-config '{ + "afd": { + "role": "ffn", + "connector": "GpuAsyncAFDConnector", + "async": true, + "compute_gate_on_attention": true, + "host": "127.0.0.1", + "port": '"$AFD_PORT"', + "num_attention_ranks": 1, + "num_ffn_ranks": 1 + } + }' \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \ + --api-server-count 1 \ + --gpu-memory-utilization "$GPU_MEM_UTIL" \ + --enforce-eager \ + --host 127.0.0.1 \ + --port "$API_PORT" \ + --trust-remote-code > "$LOG_DIR/ffn.log" 2>&1 & +FFN_PID=$! + +cleanup() { + kill "$ATTN_PID" "$FFN_PID" 2>/dev/null + wait "$ATTN_PID" "$FFN_PID" 2>/dev/null +} +trap cleanup EXIT + +for _ in $(seq 1 "${READY_TIMEOUT:-600}"); do + if curl -sf "http://127.0.0.1:$API_PORT/health" > /dev/null 2>&1; then + echo "server ready on http://127.0.0.1:$API_PORT" + echo + echo "curl -s http://127.0.0.1:$API_PORT/v1/completions \\" + echo " -H 'Content-Type: application/json' \\" + echo " -d '{\"model\":\"$MODEL_PATH\",\"prompt\":\"The capital of France is\",\"max_tokens\":16,\"temperature\":0}'" + echo + if [ -n "${SMOKE:-}" ]; then + curl -s "http://127.0.0.1:$API_PORT/v1/completions" \ + -H 'Content-Type: application/json' \ + -d '{"model":"'"$MODEL_PATH"'","prompt":"The capital of France is", + "max_tokens":16,"temperature":0}' + echo + exit 0 + fi + # Stay up so the servers can take requests; Ctrl-C tears both down. + wait "$ATTN_PID" "$FFN_PID" + exit 0 + fi + if ! kill -0 "$ATTN_PID" 2>/dev/null || ! kill -0 "$FFN_PID" 2>/dev/null; then + echo "a server exited early; see $LOG_DIR/attn.log and $LOG_DIR/ffn.log" >&2 + exit 1 + fi + sleep 1 +done +echo "timed out waiting for the server" >&2 +exit 1 diff --git a/recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/2a2f_eager_async.sh b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/2a2f_eager_async.sh new file mode 100644 index 00000000..01cd83dd --- /dev/null +++ b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/2a2f_eager_async.sh @@ -0,0 +1,129 @@ +#!/usr/bin/env bash +# 2A2F async GPU connector, eager, prefill-only. +# +# Launch under a GPU reservation, which sets CUDA_VISIBLE_DEVICES: +# gpu run --gpus 4 -- bash recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/2a2f_eager_async.sh +# +# The two roles are separate vllm serve processes: the AFD process group hosts +# its own TCPStore, which cannot be created under a single torchrun/torchelastic +# launcher. +set -u + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} +LOG_DIR=${LOG_DIR:-.} +mkdir -p "$LOG_DIR" +export VLLM_USE_V2_MODEL_RUNNER=0 +# Single node over NVLink: skip the IB transport probe. +export NVSHMEM_REMOTE_TRANSPORT=${NVSHMEM_REMOTE_TRANSPORT:-none} +# Two servers on one box spawn a lot of threads; the HF tokenizer's rayon pool +# is the first thing to fail when thread creation gets refused. +export TOKENIZERS_PARALLELISM=${TOKENIZERS_PARALLELISM:-false} +export RAYON_NUM_THREADS=${RAYON_NUM_THREADS:-2} +export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4} + +# Split the reserved devices in half: first two Attention, last two FFN. +IFS=',' read -r -a DEVICES <<< "${CUDA_VISIBLE_DEVICES:-0,1,2,3}" +if [ "${#DEVICES[@]}" -lt 4 ]; then + echo "need 4 visible GPUs, got ${#DEVICES[@]}: ${CUDA_VISIBLE_DEVICES:-unset}" >&2 + exit 1 +fi +ATTN_DEVICES="${DEVICES[0]},${DEVICES[1]}" +FFN_DEVICES="${DEVICES[2]},${DEVICES[3]}" +echo "attention on ${ATTN_DEVICES}, ffn on ${FFN_DEVICES}" + +# Lower this when sharing a box: vLLM refuses to start if the desired +# fraction exceeds what is actually free. +GPU_MEM_UTIL=${GPU_MEM_UTIL:-0.9} +# Prefill batch size drives whether each MoE call clears the compute-bound +# inflection point, so it is the knob to raise when benchmarking. +MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-512} +MAX_NUM_SEQS=${MAX_NUM_SEQS:-8} +AFD_PORT=${AFD_PORT:-6271} +API_PORT=${API_PORT:-18307} + +CUDA_VISIBLE_DEVICES="$ATTN_DEVICES" uv run vllm serve "$MODEL_PATH" \ + --data-parallel-size 2 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --additional-config '{ + "afd": { + "role": "attention", + "connector": "GpuAsyncAFDConnector", + "async": true, + "compute_gate_on_attention": true, + "host": "127.0.0.1", + "port": '"$AFD_PORT"', + "num_attention_ranks": 2, + "num_ffn_ranks": 2 + } + }' \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \ + --api-server-count 1 \ + --gpu-memory-utilization "$GPU_MEM_UTIL" \ + --enforce-eager \ + --host 127.0.0.1 \ + --port "$API_PORT" \ + --trust-remote-code > "$LOG_DIR/attn.log" 2>&1 & +ATTN_PID=$! + +CUDA_VISIBLE_DEVICES="$FFN_DEVICES" uv run vllm serve "$MODEL_PATH" \ + --data-parallel-size 2 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --additional-config '{ + "afd": { + "role": "ffn", + "connector": "GpuAsyncAFDConnector", + "async": true, + "compute_gate_on_attention": true, + "host": "127.0.0.1", + "port": '"$AFD_PORT"', + "num_attention_ranks": 2, + "num_ffn_ranks": 2 + } + }' \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \ + --api-server-count 1 \ + --gpu-memory-utilization "$GPU_MEM_UTIL" \ + --enforce-eager \ + --host 127.0.0.1 \ + --port "$API_PORT" \ + --trust-remote-code > "$LOG_DIR/ffn.log" 2>&1 & +FFN_PID=$! + +cleanup() { + kill "$ATTN_PID" "$FFN_PID" 2>/dev/null + wait "$ATTN_PID" "$FFN_PID" 2>/dev/null +} +trap cleanup EXIT + +for _ in $(seq 1 "${READY_TIMEOUT:-600}"); do + if curl -sf "http://127.0.0.1:$API_PORT/health" > /dev/null 2>&1; then + echo "server ready on http://127.0.0.1:$API_PORT" + echo + echo "curl -s http://127.0.0.1:$API_PORT/v1/completions \\" + echo " -H 'Content-Type: application/json' \\" + echo " -d '{\"model\":\"$MODEL_PATH\",\"prompt\":\"The capital of France is\",\"max_tokens\":16,\"temperature\":0}'" + echo + if [ -n "${SMOKE:-}" ]; then + curl -s "http://127.0.0.1:$API_PORT/v1/completions" \ + -H 'Content-Type: application/json' \ + -d '{"model":"'"$MODEL_PATH"'","prompt":"The capital of France is", + "max_tokens":16,"temperature":0}' + echo + exit 0 + fi + # Stay up so the servers can take requests; Ctrl-C tears both down. + wait "$ATTN_PID" "$FFN_PID" + exit 0 + fi + if ! kill -0 "$ATTN_PID" 2>/dev/null || ! kill -0 "$FFN_PID" 2>/dev/null; then + echo "a server exited early; see $LOG_DIR/attn.log and $LOG_DIR/ffn.log" >&2 + exit 1 + fi + sleep 1 +done +echo "timed out waiting for the server" >&2 +exit 1 diff --git a/tests/e2e/async_gpu_connector_e2e.py b/tests/e2e/async_gpu_connector_e2e.py new file mode 100644 index 00000000..3b9cb2be --- /dev/null +++ b/tests/e2e/async_gpu_connector_e2e.py @@ -0,0 +1,241 @@ +"""End-to-end pass over the async GPU connector's public API, two processes. + +Rank 0 runs the Attention side (``send_attn_output`` / ``recv_ffn_output``), +rank 1 runs the FFN side (``recv_ffn_work_item`` / ``send_ffn_work_item_output``) +with the real grouped-GEMM helper. Everything between the gate and the combined +result is exercised: routing, one-sided dispatch, local expert compute, the +write-back, and the weighted reduction. + +Run with two GPUs:: + + python tests/e2e/async_gpu_connector_e2e.py + +Deliberately *not* launched with torchrun. ``init_afd_process_group`` builds its +own TCPStore on the AFD port, and under torchelastic every rank is forced to +``is_master=False`` (``torch/distributed/rendezvous.py:188``), so no rank hosts +the store and the group never forms. Production launches the two roles as +separate ``vllm serve`` processes, which this mirrors. +""" + +import multiprocessing as mp +import sys +from types import SimpleNamespace + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from afd_plugin.config import AFDConfig +from afd_plugin.connectors.gpu.async_gpu import GpuAsyncAFDConnector +from afd_plugin.connectors.metadata import AFDTransferContext, AFDTransferMetadata +from afd_plugin.model_executor.models.gpu.deepseek_v2_attention_gate import ( + compute_attention_gate_moe_ffn, +) + +NUM_TOKENS = 48 +HIDDEN = 128 +INTERMEDIATE = 256 +TOPK = 4 +NUM_EXPERTS = 8 +NUM_LAYERS = 3 +PORT = 29655 +WORLD_PORT = 29656 +SCALING = 1.7 + + +def build_connector(role: str, local_rank: int) -> GpuAsyncAFDConnector: + vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + hf_config=SimpleNamespace( + hidden_size=HIDDEN, + num_experts_per_tok=TOPK, + n_routed_experts=NUM_EXPERTS, + ), + dtype=torch.bfloat16, + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=NUM_TOKENS), + additional_config={ + "afd": { + "role": role, + "connector": "GpuAsyncAFDConnector", + "async": True, + "compute_gate_on_attention": True, + "num_attention_ranks": 1, + "num_ffn_ranks": 1, + "port": PORT, + "connector_extra_config": {"ring_depth": 1}, + }, + }, + ) + afd_config = AFDConfig( + role=role, + connector="GpuAsyncAFDConnector", + async_dp=True, + compute_gate_on_attention=True, + num_attention_ranks=1, + num_ffn_ranks=1, + host="127.0.0.1", + port=PORT, + ) + connector = GpuAsyncAFDConnector( + rank=local_rank, + local_rank=local_rank, + vllm_config=vllm_config, + afd_config=afd_config, + role_rank=0, + ) + connector.init_afd_connector() + return connector + + +def make_weights(device): + generator = torch.Generator(device="cpu").manual_seed(11) + w13 = ( + torch.randn(NUM_EXPERTS, 2 * INTERMEDIATE, HIDDEN, generator=generator) + / HIDDEN**0.5 + ).to(device, torch.bfloat16) + w2 = ( + torch.randn(NUM_EXPERTS, HIDDEN, INTERMEDIATE, generator=generator) + / INTERMEDIATE**0.5 + ).to(device, torch.bfloat16) + return w13, w2 + + +def make_layer_inputs(layer_idx, device): + gen = torch.Generator(device="cpu").manual_seed(100 + layer_idx) + x = torch.randn(NUM_TOKENS, HIDDEN, generator=gen).to(device, torch.bfloat16) + topk_ids = torch.stack( + [torch.randperm(NUM_EXPERTS, generator=gen)[:TOPK] for _ in range(NUM_TOKENS)], + ).to(device, torch.int32) + topk_weights = torch.rand(NUM_TOKENS, TOPK, generator=gen).to(device) + return x, topk_ids, topk_weights + + +def reference_moe(x, w13, w2, topk_ids, topk_weights): + out = torch.zeros(x.shape[0], HIDDEN, dtype=torch.float32, device=x.device) + for token in range(x.shape[0]): + for slot in range(topk_ids.shape[1]): + expert = int(topk_ids[token, slot]) + hidden = x[token].to(torch.float32) @ w13[expert].to(torch.float32).T + gate, up = hidden.chunk(2, dim=-1) + y = (F.silu(gate) * up) @ w2[expert].to(torch.float32).T + out[token] += float(topk_weights[token, slot]) * y * SCALING + return out + + +def init_world(rank: int) -> None: + """Mimic a single `vllm serve`: a private default group of size 1. + + The connector bootstraps NVSHMEM on the AFD group itself, so the default + group deliberately does *not* span both roles -- that is the topology the + real deployment has. + """ + dist.init_process_group( + "nccl", + init_method=f"tcp://127.0.0.1:{WORLD_PORT + rank}", + world_size=1, + rank=0, + ) + + +def run_attention(rank: int) -> None: + init_world(0) + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + connector = build_connector("attention", rank) + w13, w2 = make_weights(device) + + for layer_idx in range(NUM_LAYERS): + x, topk_ids, topk_weights = make_layer_inputs(layer_idx, device) + context = AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=layer_idx, + stage_idx=0, + seq_len=NUM_TOKENS, + ), + ) + connector.send_attn_output( + x, + context, + topk_ids=topk_ids, + topk_weights=topk_weights, + ) + got = connector.recv_ffn_output(ref_tensor=x, ubatch_idx=0) + expected = reference_moe(x, w13, w2, topk_ids, topk_weights) + torch.testing.assert_close( + got.to(torch.float32), + expected, + rtol=8e-2, + atol=8e-2, + ) + print( + f"[A] layer {layer_idx}: combined output matches reference MoE", flush=True + ) + + print("PASS: async GPU connector end-to-end", flush=True) + connector.close() + + +def run_ffn(rank: int) -> None: + init_world(1) + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + connector = build_connector("ffn", rank) + w13, w2 = make_weights(device) + layer = SimpleNamespace( + mlp=SimpleNamespace( + experts=SimpleNamespace( + routed_experts=SimpleNamespace(w13_weight=w13, w2_weight=w2), + _shared_experts=None, + routed_scaling_factor=SCALING, + ), + ), + ) + + for _ in range(NUM_LAYERS): + while True: + try: + work_item = connector.recv_ffn_work_item( + stage_idx=0, + max_num_tokens=NUM_TOKENS, + ) + break + except TimeoutError: + continue + states = work_item.context.states + payload = compute_attention_gate_moe_ffn( + layer, + hidden_states=work_item.hidden_states, + group_list=states.group_list, + expand_x_shared=None, + ) + connector.send_ffn_work_item_output(work_item, payload) + print( + f"[F] layer {work_item.layer_idx}: served " + f"{work_item.num_tokens} routed tokens", + flush=True, + ) + + connector.close() + + +def main() -> None: + if torch.cuda.device_count() < 2: + raise SystemExit("this test needs two visible GPUs") + mp.set_start_method("spawn", force=True) + procs = [ + mp.Process(target=run_ffn, args=(1,)), + mp.Process(target=run_attention, args=(0,)), + ] + for proc in procs: + proc.start() + failed = False + for proc in procs: + proc.join(timeout=300) + if proc.exitcode != 0: + failed = True + sys.exit(1 if failed else 0) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/async_gpu_moe_equivalence.py b/tests/e2e/async_gpu_moe_equivalence.py new file mode 100644 index 00000000..d5f22e67 --- /dev/null +++ b/tests/e2e/async_gpu_moe_equivalence.py @@ -0,0 +1,121 @@ +"""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. +""" + +import torch +import torch.nn.functional as F +from types import SimpleNamespace + +from afd_plugin.connectors.gpu.async_gpu import plan_dispatch +from afd_plugin.model_executor.models.gpu.deepseek_v2_attention_gate import ( + compute_attention_gate_moe_ffn, +) + +NUM_TOKENS = 32 +HIDDEN = 128 +INTERMEDIATE = 256 +TOPK = 4 +NUM_EXPERTS = 8 +FFN_SIZE = 2 +EXPERT_PER_RANK = NUM_EXPERTS // FFN_SIZE + + +def reference_moe(x, w13, w2, topk_ids, topk_weights): + """Naive per-token MoE: sum over each token's topk experts.""" + out = torch.zeros(x.shape[0], HIDDEN, dtype=torch.float32, device=x.device) + for token in range(x.shape[0]): + for slot in range(topk_ids.shape[1]): + expert = int(topk_ids[token, slot]) + hidden = x[token].to(torch.float32) @ w13[expert].to(torch.float32).T + gate, up = hidden.chunk(2, dim=-1) + y = (F.silu(gate) * up) @ w2[expert].to(torch.float32).T + out[token] += float(topk_weights[token, slot]) * y + return out + + +def main() -> None: + torch.manual_seed(0) + device = torch.device("cuda", 0) + dtype = torch.bfloat16 + + x = torch.randn(NUM_TOKENS, HIDDEN, device=device, dtype=dtype) + w13 = torch.randn( + NUM_EXPERTS, 2 * INTERMEDIATE, HIDDEN, device=device, dtype=dtype, + ) / HIDDEN**0.5 + w2 = torch.randn( + NUM_EXPERTS, HIDDEN, INTERMEDIATE, device=device, dtype=dtype, + ) / INTERMEDIATE**0.5 + topk_ids = torch.stack( + [torch.randperm(NUM_EXPERTS)[:TOPK] for _ in range(NUM_TOKENS)], + ).to(device, torch.int32) + topk_weights = torch.rand(NUM_TOKENS, TOPK, device=device, dtype=torch.float32) + + expected = reference_moe(x, w13, w2, topk_ids, topk_weights) + + # --- what the connector + FFN runner actually do ----------------------- + plan = plan_dispatch( + topk_ids, topk_weights, ffn_size=FFN_SIZE, expert_per_rank=EXPERT_PER_RANK, + ) + # A destination locates its own partials from the two header words the + # sender fills on the device; nothing about the routing is read back here, + # which is what the send path does now too. + starts = plan.segment_start.cpu().tolist() + routed = plan.routed_per_rank.cpu().tolist() + + accumulator = torch.zeros(NUM_TOKENS, HIDDEN, dtype=torch.float32, device=device) + for ffn_rank in range(FFN_SIZE): + base = ffn_rank * EXPERT_PER_RANK + group_list = plan.counts[base : base + EXPERT_PER_RANK] + partials = slice(starts[ffn_rank], starts[ffn_rank] + routed[ffn_rank]) + expand = plan.expand_idx[partials].to(torch.int64) + + # This FFN rank owns experts [base, base + EXPERT_PER_RANK). + layer = SimpleNamespace( + mlp=SimpleNamespace( + experts=SimpleNamespace( + routed_experts=SimpleNamespace( + w13_weight=w13[base : base + EXPERT_PER_RANK].contiguous(), + w2_weight=w2[base : base + EXPERT_PER_RANK].contiguous(), + ), + _shared_experts=None, + routed_scaling_factor=1.0, + ), + ), + ) + # The whole batch crosses the wire; the FFN side gathers one row per + # partial from it before the grouped GEMM, which applies the partial + # weights in its own epilogue. + payload = compute_attention_gate_moe_ffn( + layer, + hidden_states=x.index_select(0, expand), + group_list=group_list, + expand_x_shared=None, + ) + # ...then weights and reduces back to one row per token, in the payload + # dtype, leaving a zero row for tokens this rank held no expert for. + reduced = torch.zeros( + NUM_TOKENS, HIDDEN, dtype=payload.routed_output.dtype, device=device + ) + weighted = payload.routed_output * plan.weights[partials].unsqueeze(1).to( + payload.routed_output.dtype + ) + reduced.index_add_(0, expand, weighted) + accumulator += reduced.to(torch.float32) + + diff = (accumulator - expected).abs() + rel = diff.max() / expected.abs().max() + print(f"max abs diff={diff.max():.4e} max rel={rel:.4e}") + torch.testing.assert_close(accumulator, expected, rtol=6e-2, atol=6e-2) + print("PASS: dispatch -> local experts -> combine matches naive MoE") + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/async_gpu_window_roundtrip.py b/tests/e2e/async_gpu_window_roundtrip.py new file mode 100644 index 00000000..cc71b7c8 --- /dev/null +++ b/tests/e2e/async_gpu_window_roundtrip.py @@ -0,0 +1,179 @@ +"""Two-rank A->F->A round trip over SymmWindow with real routing data. + +Run with two GPUs:: + + torchrun --nproc_per_node=2 tests/e2e/async_gpu_window_roundtrip.py + + +Rank 0 plays one Attention rank, rank 1 one FFN rank holding every expert. +The FFN side runs an identity "expert", so the recombined result must equal the +weighted sum of the inputs -- the same invariant the CPU self-check asserts, +but now carried across the wire. +""" + +import os + +import torch +import torch.distributed as dist + +from afd_plugin.connectors.gpu.async_gpu import plan_dispatch +from afd_plugin.connectors.gpu.symm_window import ( + SlotLayout, + SymmWindow, + encode_header, +) + +NUM_TOKENS = 64 +HIDDEN = 128 +TOPK = 6 +FFN_SIZE = 1 +EXPERT_PER_RANK = 16 +RING_DEPTH = 2 +NUM_LAYERS = 3 + + +def main() -> None: + rank = int(os.environ["RANK"]) + # Device index must come from LOCAL_RANK: under a GPU reservation the + # visible devices are remapped, so the global rank can be out of range. + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("nccl") + + layout = SlotLayout.build( + expert_per_rank=EXPERT_PER_RANK, + partial_cap=NUM_TOKENS * TOPK, + token_cap=NUM_TOKENS, + hidden_size=HIDDEN, + payload_itemsize=2, + ) + window = SymmWindow( + num_regions=1, + ring_depth=RING_DEPTH, + layout=layout, + payload_dtype=torch.bfloat16, + device=device, + group=dist.group.WORLD, + rank=rank, + world_size=2, + ) + if rank == 0: + print(f"slot={layout.slot_bytes / 2**10:.0f}KiB " + f"window={window.total_bytes / 2**10:.0f}KiB", flush=True) + + for layer_idx in range(NUM_LAYERS): + ring = layer_idx % RING_DEPTH + gen = torch.Generator(device="cpu").manual_seed(layer_idx) + hidden = torch.randn(NUM_TOKENS, HIDDEN, generator=gen).to(device, torch.bfloat16) + topk_ids = torch.stack( + [torch.randperm(EXPERT_PER_RANK, generator=gen)[:TOPK] + for _ in range(NUM_TOKENS)], + ).to(device, torch.int32) + topk_weights = torch.rand(NUM_TOKENS, TOPK, generator=gen).to(device) + + # Both ranks build the plan from the same inputs. Rank 0 sends from it; + # rank 1 only uses it as an oracle for what should have arrived. + plan = plan_dispatch( + topk_ids, topk_weights, + ffn_size=FFN_SIZE, expert_per_rank=EXPERT_PER_RANK, + ) + + if rank == 0: + counts_host = plan.counts.cpu().tolist() + # Shared rows are a contiguous range; this rank owns all of them. + shared = slice(0, NUM_TOKENS) + header = encode_header( + layout, seq=layer_idx + 1, src_role_rank=0, layer_idx=layer_idx, + stage_idx=0, num_tokens=NUM_TOKENS, + routed_tokens=NUM_TOKENS * TOPK, + shared_tokens=shared.stop - shared.start, topk=TOPK, flags=0, + expert_counts=counts_host, + segment_start=0, + ) + # Capacity form: the whole batch and every partial go out, and the + # rank-1 side locates its own run from the header. + window.write_slot( + peer=1, region=0, ring=ring, header=header, + expand_idx=plan.expand_idx, + weights=plan.weights, + routed_x=hidden, + shared_x=hidden[shared], + ) + + # Wait for the expert output and reduce it. + while True: + arrived = window.poll() + if arrived is not None: + break + got = arrived.header + assert got.layer_idx == layer_idx, (got.layer_idx, layer_idx) + assert got.src_role_rank == 0, got.src_role_rank + assert got.routed_tokens == NUM_TOKENS * TOPK, got.routed_tokens + + # A reply is a whole batch, so combine adds it without an index. + acc = window.local_routed( + arrived.region, arrived.ring, got.num_tokens, + ).to(torch.float32).clone() + acc[shared] += window.local_shared( + arrived.region, arrived.ring, got.shared_tokens + ).to(torch.float32) + expected = ( + hidden.to(torch.float32) * topk_weights.sum(dim=1, keepdim=True) + + hidden.to(torch.float32) # identity shared expert + ) + torch.testing.assert_close(acc, expected, rtol=2e-2, atol=2e-2) + print(f"layer {layer_idx}: combine matches reference", flush=True) + else: + while True: + arrived = window.poll() + if arrived is not None: + break + got = arrived.header + assert got.layer_idx == layer_idx, (got.layer_idx, layer_idx) + assert sum(got.expert_counts) == got.routed_tokens + shipped = window.local_routed( + arrived.region, arrived.ring, got.num_tokens) + shared_rows = window.local_shared( + arrived.region, arrived.ring, got.shared_tokens) + expand = window.local_expand_idx( + arrived.region, arrived.ring, got.routed_tokens, + got.segment_start).to(torch.int64) + weights = window.local_weights( + arrived.region, arrived.ring, got.routed_tokens, got.segment_start) + # Gathering by the partial indices must reproduce the sender's rows. + expanded = shipped.index_select(0, expand) + torch.testing.assert_close(expanded, hidden.index_select(0, expand)) + + # Identity experts, then the weighted reduce back to token rows. + reduced = torch.zeros( + got.num_tokens, HIDDEN, dtype=torch.float32, device=device) + reduced.index_add_( + 0, expand, expanded.to(torch.float32) * weights.unsqueeze(1)) + + echo = encode_header( + layout, seq=layer_idx + 1, src_role_rank=0, layer_idx=got.layer_idx, + stage_idx=got.stage_idx, num_tokens=got.num_tokens, + routed_tokens=got.routed_tokens, shared_tokens=got.shared_tokens, + topk=TOPK, flags=0, expert_counts=got.expert_counts, + segment_start=got.segment_start, + ) + window.write_slot( + peer=0, region=0, ring=arrived.ring, header=echo, + expand_idx=None, + weights=None, + routed_x=reduced, + shared_x=shared_rows.clone(), + ) + print(f"layer {layer_idx}: dispatch payload verified, echoed back", + flush=True) + + dist.barrier() + if rank == 0: + print("PASS: A->F->A round trip over SymmWindow", flush=True) + window.close() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/unit/config/test_config.py b/tests/unit/config/test_config.py index 7728bd8a..fa77ba54 100644 --- a/tests/unit/config/test_config.py +++ b/tests/unit/config/test_config.py @@ -119,7 +119,7 @@ def test_parse_async_dp_config_from_async_alias(): def test_async_dp_requires_async_connector(): - with pytest.raises(ValueError, match="requires connector='CAMAsyncAFDConnector'"): + with pytest.raises(ValueError, match="AFD async mode requires one of"): parse_afd_config( { "afd": { @@ -131,6 +131,24 @@ def test_async_dp_requires_async_connector(): ) +@pytest.mark.parametrize( + "connector", + ["CAMAsyncAFDConnector", "GpuAsyncAFDConnector"], +) +def test_async_dp_accepts_every_async_connector(connector): + config = parse_afd_config( + { + "afd": { + "connector": connector, + "role": "attention", + "async": True, + }, + }, + ) + assert config.connector == connector + assert config.async_dp + + def test_original_common_afd_field_aliases_are_supported(): raw = { "afd_role": "ffn", diff --git a/tests/unit/connectors/test_async_gpu_connector.py b/tests/unit/connectors/test_async_gpu_connector.py new file mode 100644 index 00000000..bf320309 --- /dev/null +++ b/tests/unit/connectors/test_async_gpu_connector.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Unit tests for the async GPU connector's wire format and routing math.""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("torch") + +import torch # noqa: E402 + +from afd_plugin.connectors.factory import AFDConnectorFactory # noqa: E402 +from afd_plugin.connectors.gpu.async_gpu import ( # noqa: E402 + GpuAsyncAFDConnector, + GpuAsyncExtraInfo, + combine_scatter, + plan_dispatch, +) +from afd_plugin.connectors.gpu.symm_window import ( # noqa: E402 + FLAG_SHUTDOWN_BIT, + HEADER_FIXED_WORDS, + SlotLayout, + decode_header, + encode_header, +) + + +@pytest.fixture +def layout() -> SlotLayout: + return SlotLayout.build( + expert_per_rank=4, + routed_cap=100, + token_cap=32, + hidden_size=8, + payload_itemsize=2, + ) + + +# ---------------------------------------------------------------------- +# Config +# ---------------------------------------------------------------------- + + +def test_connector_is_registered_and_has_no_control_plane(): + connector_cls = AFDConnectorFactory.get_connector_class("GpuAsyncAFDConnector") + assert connector_cls is GpuAsyncAFDConnector + assert connector_cls.control_plane is None + + +def test_ring_depth_defaults_to_the_number_of_live_stages(): + assert GpuAsyncExtraInfo.from_mapping(None).ring_depth == 1 + assert GpuAsyncExtraInfo.from_mapping({"async_moe_ubatching": True}).ring_depth == 2 + assert GpuAsyncExtraInfo.from_mapping({"ring_depth": 4}).ring_depth == 4 + + +def test_unknown_extra_config_field_is_rejected(): + with pytest.raises(ValueError, match="unknown AFD async GPU"): + GpuAsyncExtraInfo.from_mapping({"nope": 1}) + + +def test_routed_cap_multiplier_must_be_positive(): + with pytest.raises(ValueError, match="routed_cap_multiplier"): + GpuAsyncExtraInfo.from_mapping({"routed_cap_multiplier": 0}) + + +# ---------------------------------------------------------------------- +# Slot layout +# ---------------------------------------------------------------------- + + +def test_slot_fields_are_disjoint_and_fit_inside_the_slot(layout: SlotLayout): + assert layout.header_words == HEADER_FIXED_WORDS + 4 + assert layout.route_table_off >= layout.header_off + layout.header_words * 4 + assert layout.routed_x_off >= layout.route_table_off + 2 * 100 * 4 + assert layout.shared_idx_off >= layout.routed_x_off + 100 * 8 * 2 + assert layout.shared_x_off >= layout.shared_idx_off + 32 * 4 + assert layout.slot_bytes >= layout.shared_x_off + 32 * 8 * 2 + + +def test_every_field_offset_is_viewable_as_int32_and_payload(layout: SlotLayout): + # get_buffer takes an element offset, so a byte offset that is not a + # multiple of the element size would silently land on the wrong address. + for offset in ( + layout.header_off, + layout.route_table_off, + layout.routed_x_off, + layout.shared_idx_off, + layout.shared_x_off, + ): + assert offset % 4 == 0 + assert offset % layout.payload_itemsize == 0 + + +# ---------------------------------------------------------------------- +# Header codec +# ---------------------------------------------------------------------- + + +def test_header_round_trip(layout: SlotLayout): + header = encode_header( + layout, + seq=7, + src_role_rank=3, + layer_idx=11, + stage_idx=1, + num_tokens=32, + routed_tokens=90, + shared_tokens=16, + topk=6, + flags=0, + expert_counts=[10, 20, 30, 30], + ) + decoded = decode_header(header) + assert decoded.seq == 7 + assert decoded.src_role_rank == 3 + assert decoded.layer_idx == 11 + assert decoded.stage_idx == 1 + assert decoded.num_tokens == 32 + assert decoded.routed_tokens == 90 + assert decoded.shared_tokens == 16 + assert decoded.topk == 6 + assert decoded.expert_counts == [10, 20, 30, 30] + assert sum(decoded.expert_counts) == decoded.routed_tokens + assert not decoded.is_shutdown + + +def test_shutdown_flag_survives_the_round_trip(layout: SlotLayout): + header = encode_header( + layout, + seq=8, + src_role_rank=0, + layer_idx=0, + stage_idx=0, + num_tokens=0, + routed_tokens=0, + shared_tokens=0, + topk=6, + flags=FLAG_SHUTDOWN_BIT, + expert_counts=[0, 0, 0, 0], + ) + assert decode_header(header).is_shutdown + + +def test_corrupt_magic_is_rejected(layout: SlotLayout): + header = encode_header( + layout, + seq=1, + src_role_rank=0, + layer_idx=0, + stage_idx=0, + num_tokens=1, + routed_tokens=0, + shared_tokens=0, + topk=6, + flags=0, + expert_counts=[0, 0, 0, 0], + ) + header[0] = 0 + with pytest.raises(RuntimeError, match="magic mismatch"): + decode_header(header) + + +def test_expert_counts_length_must_match_the_layout(layout: SlotLayout): + with pytest.raises(ValueError, match="expert_counts"): + encode_header( + layout, + seq=1, + src_role_rank=0, + layer_idx=0, + stage_idx=0, + num_tokens=1, + routed_tokens=0, + shared_tokens=0, + topk=6, + flags=0, + expert_counts=[1, 2], + ) + + +# ---------------------------------------------------------------------- +# Routing +# ---------------------------------------------------------------------- + +_NUM_TOKENS = 7 +_TOPK = 3 +_FFN_SIZE = 2 +_EXPERT_PER_RANK = 4 +_HIDDEN = 5 + + +@pytest.fixture +def routing_inputs(): + generator = torch.Generator().manual_seed(0) + num_experts = _FFN_SIZE * _EXPERT_PER_RANK + topk_ids = torch.stack( + [ + torch.randperm(num_experts, generator=generator)[:_TOPK] + for _ in range(_NUM_TOKENS) + ], + ).to(torch.int32) + hidden_states = torch.randn(_NUM_TOKENS, _HIDDEN, generator=generator) + topk_weights = torch.rand(_NUM_TOKENS, _TOPK, generator=generator) + return topk_ids, hidden_states, topk_weights + + +def test_every_partial_is_routed_exactly_once(routing_inputs): + topk_ids, _, _ = routing_inputs + route_table, counts, _ = plan_dispatch( + topk_ids, + ffn_size=_FFN_SIZE, + expert_per_rank=_EXPERT_PER_RANK, + ) + assert route_table.shape == (_NUM_TOKENS * _TOPK, 2) + assert int(counts.sum()) == _NUM_TOKENS * _TOPK + seen = {(token_idx, slot) for token_idx, slot in route_table.tolist()} + assert len(seen) == _NUM_TOKENS * _TOPK + + +def test_each_destination_segment_is_grouped_by_local_expert(routing_inputs): + topk_ids, _, _ = routing_inputs + route_table, counts, offsets = plan_dispatch( + topk_ids, + ffn_size=_FFN_SIZE, + expert_per_rank=_EXPERT_PER_RANK, + ) + counts_host, offsets_host = counts.tolist(), offsets.tolist() + for ffn_rank in range(_FFN_SIZE): + base = ffn_rank * _EXPERT_PER_RANK + cursor = offsets_host[base] + for local_expert in range(_EXPERT_PER_RANK): + for _ in range(counts_host[base + local_expert]): + token_idx, slot = route_table[cursor].tolist() + assert int(topk_ids[token_idx, slot]) == base + local_expert + cursor += 1 + assert cursor == offsets_host[base] + sum( + counts_host[base : base + _EXPERT_PER_RANK], + ) + + +def test_identity_experts_recombine_to_the_weighted_sum(routing_inputs): + topk_ids, hidden_states, topk_weights = routing_inputs + route_table, counts, offsets = plan_dispatch( + topk_ids, + ffn_size=_FFN_SIZE, + expert_per_rank=_EXPERT_PER_RANK, + ) + counts_host, offsets_host = counts.tolist(), offsets.tolist() + accumulator = torch.zeros(_NUM_TOKENS, _HIDDEN, dtype=torch.float32) + for ffn_rank in range(_FFN_SIZE): + base = ffn_rank * _EXPERT_PER_RANK + start = offsets_host[base] + total = sum(counts_host[base : base + _EXPERT_PER_RANK]) + segment = route_table[start : start + total] + combine_scatter( + accumulator, + routed_out=hidden_states.index_select(0, segment[:, 0].to(torch.int64)), + route_table=segment, + topk_weights=topk_weights, + ) + expected = hidden_states * topk_weights.sum(dim=1, keepdim=True) + torch.testing.assert_close(accumulator, expected.to(torch.float32)) + + +def test_routing_handles_experts_not_divisible_by_ffn_size(): + # expert_per_rank is a ceiling division, so the padded tail must stay empty + # instead of silently absorbing real partials. + ffn_size, expert_per_rank, num_experts = 3, 2, 5 + topk_ids = torch.tensor([[0, 4], [1, 3], [2, 4]], dtype=torch.int32) + _, counts, _ = plan_dispatch( + topk_ids, + ffn_size=ffn_size, + expert_per_rank=expert_per_rank, + ) + assert counts.numel() == ffn_size * expert_per_rank + assert int(counts.sum()) == topk_ids.numel() + assert int(counts[num_experts:].sum()) == 0 + + +def test_routing_can_leave_one_destination_empty(): + """A single decode token's topk can land entirely on one FFN rank. + + The peer that gets nothing must still see a well-formed, empty segment -- + zero-length windows are what crashed a 2A2F decode step. + """ + ffn_size, expert_per_rank = 2, 4 + # Every partial targets experts owned by FFN rank 0. + topk_ids = torch.tensor([[0, 1, 2]], dtype=torch.int32) + route_table, counts, offsets = plan_dispatch( + topk_ids, + ffn_size=ffn_size, + expert_per_rank=expert_per_rank, + ) + counts_host, offsets_host = counts.tolist(), offsets.tolist() + + base_zero = 0 + assert sum(counts_host[base_zero : base_zero + expert_per_rank]) == 3 + base_one = expert_per_rank + empty_total = sum(counts_host[base_one : base_one + expert_per_rank]) + assert empty_total == 0 + empty_segment = route_table[ + offsets_host[base_one] : offsets_host[base_one] + empty_total + ] + assert empty_segment.shape == (0, 2) + + # Combining an empty segment must be a no-op, not an error. + accumulator = torch.zeros(1, 4, dtype=torch.float32) + combine_scatter( + accumulator, + routed_out=torch.zeros(0, 4), + route_table=empty_segment, + topk_weights=torch.ones(1, 3), + ) + assert torch.count_nonzero(accumulator) == 0 diff --git a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py index b650f480..6c6a2cd1 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py @@ -8,7 +8,7 @@ pytest.importorskip("vllm") nn = torch.nn -from afd_plugin.config import AFD_ASYNC_CONNECTOR, AFDConfig # noqa: E402 +from afd_plugin.config import AFD_ASYNC_NPU_CONNECTOR, AFDConfig # noqa: E402 from afd_plugin.model_executor.models import deepseek_v2 as adapter # noqa: E402 @@ -247,7 +247,7 @@ def async_forward(*args): nn.Module.__init__(model) model.afd_config = AFDConfig( role="attention", - connector=AFD_ASYNC_CONNECTOR, + connector=AFD_ASYNC_NPU_CONNECTOR, ) positions = torch.arange(1) diff --git a/tests/unit/model_executor/models/test_forward_context.py b/tests/unit/model_executor/models/test_forward_context.py index bff67d8d..45d52b96 100644 --- a/tests/unit/model_executor/models/test_forward_context.py +++ b/tests/unit/model_executor/models/test_forward_context.py @@ -412,14 +412,19 @@ def test_deepseek_compute_gate_on_attention_selects_backend_boundary(): assert "self.mlp = AFDDeepseekV2RemoteExpertsMoE(" in source assert "self.mlp = GateOnlyRemoteMoE(" in source assert 'prefix=f"{prefix}.mlp"' in source + # The gate/topk helper delegates expert selection to the connector, so both + # platforms share it; only the FFN-side MoE compute stays platform-split. assert ( - "# NPU-only: Attention-side gate/topk is implemented in the NPU helper." + "# The gate helper delegates expert selection to the connector, so both" in source ) assert ( "# NPU-only: gated MoE FFN compute consumes Attention-side topk payloads." in source ) + # CUDA reaches its own grouped-GEMM entry point only once tokens arrive + # pre-routed; without a group list the control-plane path still applies. + assert "gpu_attention_gate.compute_attention_gate_moe_ffn(" in source def test_async_moe_pipeline_preserves_stage_order(monkeypatch): diff --git a/tests/unit/v1/worker/test_attention_model_runner.py b/tests/unit/v1/worker/test_attention_model_runner.py index 94704bd6..f6341458 100644 --- a/tests/unit/v1/worker/test_attention_model_runner.py +++ b/tests/unit/v1/worker/test_attention_model_runner.py @@ -352,13 +352,21 @@ def test_phase5_allows_two_way_ubatching_but_rejects_other_counts(): ) -def _ubatch_runner(uniform_decode, **parallel_overrides): +_DUMMY_CONTROL_PLANE = object() + + +def _ubatch_runner( + uniform_decode, *, control_plane=_DUMMY_CONTROL_PLANE, **parallel_overrides +): runner = object.__new__(AFDAttentionModelRunner) runner.vllm_config = SimpleNamespace( parallel_config=_parallel_config(**parallel_overrides), ) runner.uniform_decode_query_len = 1 runner._is_uniform_decode = lambda **_kwargs: uniform_decode + # The override consults the connector to decide whether cross-DP batch + # coordination applies; a non-None control plane keeps upstream behaviour. + runner.connector = SimpleNamespace(control_plane=control_plane) return runner @@ -1098,3 +1106,34 @@ def test_attention_runner_load_model_initializes_connector_after_weights( expected.append("connector_init") assert events == expected assert connector.is_initialized is True + +def test_connector_driven_runs_skip_cross_dp_batch_coordination(): + """An idle Attention replica never joins the DP all-reduce. + + Async AFD lets each replica advance alone, so a busy replica must not block + in ``coordinate_batch_across_dp`` waiting for one that never steps. + """ + import vllm.v1.worker.gpu_model_runner as gpu_model_runner + + from afd_plugin.v1.worker.attention_model_runner import ( + _dp_batch_coordination_disabled, + ) + + original = gpu_model_runner.coordinate_batch_across_dp + + with _dp_batch_coordination_disabled(True): + assert gpu_model_runner.coordinate_batch_across_dp is not original + result = gpu_model_runner.coordinate_batch_across_dp( + num_tokens_unpadded=8, + parallel_config=None, + allow_microbatching=False, + num_tokens_padded=8, + uniform_decode=True, + cudagraph_mode=0, + ) + # num_tokens_across_dp None makes upstream skip its DP-padding branch. + assert result == (False, None, 0) + assert gpu_model_runner.coordinate_batch_across_dp is original + + with _dp_batch_coordination_disabled(False): + assert gpu_model_runner.coordinate_batch_across_dp is original diff --git a/tests/unit/v1/worker/test_ffn_model_runner.py b/tests/unit/v1/worker/test_ffn_model_runner.py index bf018e01..0743547f 100644 --- a/tests/unit/v1/worker/test_ffn_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_model_runner.py @@ -20,6 +20,7 @@ AFDTransferContext, AFDTransferMetadata, ) +from afd_plugin.connectors.gpu.async_gpu import ConnectorShutdown # noqa: E402 from afd_plugin.v1.worker.cuda_graph import make_ffn_graph_key # noqa: E402 from afd_plugin.v1.worker.ffn_model_runner import ( # noqa: E402 GPUFFNModelRunner, @@ -654,18 +655,45 @@ def test_ffn_worker_reports_zero_compilation_times(): assert compilation_times.encoder == 0.0 -def test_ffn_worker_loop_rejects_connector_without_control_plane(): +def test_ffn_worker_loop_drives_connector_without_control_plane(): worker = object.__new__(AFDFFNWorker) event = threading.Event() + steps = [] + + def execute_connector_driven_step(): + steps.append(1) + # The connector-driven step returns on an idle poll; the loop must come + # back to the shutdown event rather than block forever. + if len(steps) == 3: + event.set() worker._ffn_shutdown_event = event worker.device = SimpleNamespace(type="cpu") worker.model_runner = SimpleNamespace( connector=_ConnectorDrivenFakeConnector(), + execute_connector_driven_step=execute_connector_driven_step, + ) + + worker._run_ffn_server_loop() + + assert len(steps) == 3 + + +def test_ffn_worker_loop_exits_cleanly_when_peer_announces_shutdown(): + worker = object.__new__(AFDFFNWorker) + + def execute_connector_driven_step(): + raise ConnectorShutdown("peer left") + + worker._ffn_shutdown_event = threading.Event() + worker.device = SimpleNamespace(type="cpu") + worker.model_runner = SimpleNamespace( + connector=_ConnectorDrivenFakeConnector(), + execute_connector_driven_step=execute_connector_driven_step, ) - with pytest.raises(NotImplementedError, match="control-plane-driven"): - worker._run_ffn_server_loop() + # A peer shutdown is an ordinary exit, not a loop failure. + worker._run_ffn_server_loop() def test_ffn_worker_loop_logs_unexpected_thread_errors(caplog): From 71892a3239428b934a8e26a30582e76ab23801b2 Mon Sep 17 00:00:00 2001 From: specture724 Date: Thu, 13 Aug 2026 17:18:29 +0800 Subject: [PATCH 02/18] fix: flash_comm_v1_enabled for NPU only Signed-off-by: specture724 --- .../models/npu/deepseek_v2_async_cam_forward.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py b/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py index 383ea56e..5db05a54 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 @@ -173,7 +173,14 @@ def run_attention_gate_afd_forward( topk_weights, topk_ids, router_logits, - use_sequence_parallel=forward_context.flash_comm_v1_enabled, + # Ascend-only field: FlashComm1 leaves each Attention TP rank with a + # disjoint token shard. The CUDA forward context has no such field, + # so its absence means the token dimension is still replicated. + use_sequence_parallel=getattr( + forward_context, + "flash_comm_v1_enabled", + False, + ), ) metadata = AFDTransferMetadata.create_attention_metadata( layer_idx=layer.layer_idx, @@ -222,7 +229,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 " From 69bf40b2704ec1852b7f990667ddc5855877de17 Mon Sep 17 00:00:00 2001 From: specture724 Date: Fri, 14 Aug 2026 09:33:42 +0800 Subject: [PATCH 03/18] perf: gather dispatch rows into the peer window, poll off the compute stream Two transport-path costs on the async GPU connector, both invisible to the numerics. send_attn_output materialized each peer's routed and shared rows with index_select and then copied that staging tensor into the peer's window, so every payload byte was written and re-read locally before it went on the wire. write_slot now takes the row index instead and gathers straight into the destination view. At 2A2F that is ~650 MiB of local bandwidth per forward. SymmWindow.poll and read_header ran their D2H on the compute stream, so a rank could not even look at a flag until its own previous kernel had retired -- the receive was serialized against exactly the compute it exists to overlap with. They now run on a dedicated poll stream. This needs no ordering against local work: flags and headers are written by a peer, never by anything this rank queued. Verified on 4x L20X, DeepSeek-V2-Lite: - tests/unit: 619 passed - async_gpu_window_roundtrip, async_gpu_connector_e2e: pass - async_gpu_moe_equivalence: max rel 4.5367e-03, unchanged - 2A2F greedy probe byte-for-byte identical to the same stack before this change, so the remaining async/sync divergence is the pre-existing bf16 reduction-order difference Co-Authored-By: Claude Opus 5 Signed-off-by: specture724 --- afd_plugin/connectors/gpu/async_gpu.py | 9 +- afd_plugin/connectors/gpu/symm_window.py | 142 +++++++++++++++-------- 2 files changed, 99 insertions(+), 52 deletions(-) diff --git a/afd_plugin/connectors/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py index add5c69b..43904e87 100644 --- a/afd_plugin/connectors/gpu/async_gpu.py +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -538,15 +538,20 @@ def send_attn_output( flags=0, expert_counts=expert_counts, ) + # Rows are gathered straight into the FFN rank's window; handing + # ``write_slot`` the index instead of a gathered tensor saves a + # local write and re-read of every payload byte. window.write_slot( peer=self.attn_size + ffn_rank, region=self.role_rank, ring=ring, header=header, route_table=route_table[segment], - routed_x=hidden_states.index_select(0, token_ids[segment]), + routed_x=hidden_states, + routed_rows=token_ids[segment], shared_idx=shared_idx, - shared_x=hidden_states.index_select(0, shared_idx.to(torch.int64)), + shared_x=hidden_states, + shared_rows=shared_idx.to(torch.int64), ) logger.debug( diff --git a/afd_plugin/connectors/gpu/symm_window.py b/afd_plugin/connectors/gpu/symm_window.py index af920960..8d32a728 100644 --- a/afd_plugin/connectors/gpu/symm_window.py +++ b/afd_plugin/connectors/gpu/symm_window.py @@ -20,6 +20,11 @@ device-to-device copies complete in issue order, so a visible flag implies a complete payload. That holds for NVLink-mapped peer memory; a cross-node transport would need an explicit fence here. + +The receive side reads flags and headers on its own stream. What it reads was +produced by a peer, never by anything this rank queued, so the read needs no +ordering against local compute -- and staying off the compute stream is what +lets an arrival be noticed while the previous kernel is still running. """ from __future__ import annotations @@ -212,6 +217,13 @@ class ArrivedSlot: header: SlotHeader +def _row_count(src: torch.Tensor | None, rows: torch.Tensor | None) -> int: + """Rows a payload field will occupy, gathered or copied whole.""" + if src is None: + return 0 + return int(src.shape[0]) if rows is None else int(rows.numel()) + + class SymmWindow: """Symmetric receive window plus the one-sided writes that fill peers'.""" @@ -280,6 +292,11 @@ def __init__( # Host mirror of the local flag array; one D2H per poll refreshes it. self._flag_host = torch.zeros(self.num_flags, dtype=torch.int32).pin_memory() self._seen = [FLAG_EMPTY] * self.num_flags + # Polls read flags and headers that a *peer* wrote into our window, so + # they depend on nothing this rank has queued. On the compute stream the + # poll would still queue behind our own previous kernel, which serializes + # every receive against the compute it is supposed to overlap with. + self._poll_stream = torch.cuda.Stream(device=device) def local_bytes_view(self) -> torch.Tensor: return nvshmem_rt.tensor_from_ptr( @@ -376,22 +393,30 @@ def write_slot( routed_x: torch.Tensor | None, shared_idx: torch.Tensor | None, shared_x: torch.Tensor | None, + routed_rows: torch.Tensor | None = None, + shared_rows: torch.Tensor | None = None, ) -> None: """Write one slot into ``peer``'s window, then stamp its flag. The flag copy is issued last on the same stream, so a peer that observes the flag also observes the payload. + + ``routed_rows``/``shared_rows`` are row indices into ``routed_x`` / + ``shared_x``: the gather then lands directly in the peer's window. + Materializing the gathered rows locally first would write and re-read + every payload byte for nothing. """ layout = self.layout - if routed_x is not None and routed_x.shape[0] > layout.routed_cap: + routed_count = _row_count(routed_x, routed_rows) + shared_count = _row_count(shared_x, shared_rows) + if routed_count > layout.routed_cap: raise RuntimeError( - f"routed tokens {routed_x.shape[0]} exceed routed_cap " + f"routed tokens {routed_count} exceed routed_cap " f"{layout.routed_cap}; raise routed_cap_multiplier", ) - if shared_x is not None and shared_x.shape[0] > layout.token_cap: + if shared_count > layout.token_cap: raise RuntimeError( - f"shared tokens {shared_x.shape[0]} exceed token_cap " - f"{layout.token_cap}", + f"shared tokens {shared_count} exceed token_cap {layout.token_cap}", ) # Stage through pinned memory so the copy is asynchronous: a pageable @@ -404,46 +429,59 @@ def write_slot( non_blocking=True, ) - if route_table is not None and route_table.numel(): - n = route_table.shape[0] - self._view( - peer, - region, - ring, - layout.route_table_off, - (n, 2), - torch.int32, - ).copy_(route_table, non_blocking=True) - if routed_x is not None and routed_x.numel(): - n = routed_x.shape[0] - self._view( - peer, - region, - ring, - layout.routed_x_off, - (n, layout.hidden_size), - self.payload_dtype, - ).copy_(routed_x, non_blocking=True) - if shared_idx is not None and shared_idx.numel(): - n = shared_idx.shape[0] - self._view( - peer, - region, - ring, - layout.shared_idx_off, - (n,), - torch.int32, - ).copy_(shared_idx, non_blocking=True) - if shared_x is not None and shared_x.numel(): - n = shared_x.shape[0] - self._view( + def write_field( + field_off: int, + trailing_sizes: tuple[int, ...], + dtype: torch.dtype, + src: torch.Tensor | None, + count: int, + rows: torch.Tensor | None = None, + ) -> None: + if src is None or not count: + return + view = self._view( peer, region, ring, - layout.shared_x_off, - (n, layout.hidden_size), - self.payload_dtype, - ).copy_(shared_x, non_blocking=True) + field_off, + (count, *trailing_sizes), + dtype, + ) + if rows is None: + view.copy_(src, non_blocking=True) + else: + torch.index_select(src, 0, rows, out=view) + + write_field( + layout.route_table_off, + (2,), + torch.int32, + route_table, + _row_count(route_table, None), + ) + write_field( + layout.routed_x_off, + (layout.hidden_size,), + self.payload_dtype, + routed_x, + routed_count, + routed_rows, + ) + write_field( + layout.shared_idx_off, + (), + torch.int32, + shared_idx, + _row_count(shared_idx, None), + ) + write_field( + layout.shared_x_off, + (layout.hidden_size,), + self.payload_dtype, + shared_x, + shared_count, + shared_rows, + ) seq = int(header[_H_SEQ].item()) flag_idx = region * self.ring_depth + ring @@ -456,11 +494,13 @@ def write_slot( def poll(self) -> ArrivedSlot | None: """Return the first slot whose flag advanced past what we consumed. - ponytail: host-side poll, one D2H per call. Correct but it burns a - synchronize per attempt; replace with a device-side ``wait_any`` kernel - spinning on the flag array when the poll shows up in a profile. + ponytail: host-side poll, one D2H per call on the poll stream. Correct + but it burns a synchronize per attempt; replace with a device-side + ``wait_any`` kernel spinning on the flag array when the poll shows up in + a profile. """ - self._flag_host.copy_(self._flags_local, non_blocking=False) + with torch.cuda.stream(self._poll_stream): + self._flag_host.copy_(self._flags_local, non_blocking=False) host = self._flag_host.tolist() for idx in range(self.num_flags): if host[idx] != self._seen[idx]: @@ -475,10 +515,12 @@ def poll(self) -> ArrivedSlot | None: def read_header(self, region: int, ring: int) -> SlotHeader: # Reuse the pinned mirror instead of allocating a fresh host tensor on - # every arrival. - self._header_recv.copy_( - self._capacity_view(self.rank, region, ring, self.layout.header_off), - ) + # every arrival, and read it on the poll stream for the same reason the + # flag is read there. + with torch.cuda.stream(self._poll_stream): + self._header_recv.copy_( + self._capacity_view(self.rank, region, ring, self.layout.header_off), + ) return decode_header(self._header_recv) def local_route_table(self, region: int, ring: int, count: int) -> torch.Tensor: From 84bcd6c92731a2eff975215ca0c0a39ecad6e0cf Mon Sep 17 00:00:00 2001 From: specture724 Date: Fri, 14 Aug 2026 09:39:10 +0800 Subject: [PATCH 04/18] perf: drop the per-work-item device synchronizes from the FFN path Each FFN work item forced four full pipeline drains before it could run, which in eager mode also cost the host its run-ahead: - the per-expert group list was rebuilt with a pageable H2D from the decoded header. The counts already sit in device memory as the header's trailing words, so SymmWindow.local_expert_counts hands the grouped GEMM a view of them and the copy disappears. - send_ffn_output echoed that group list back through .cpu().tolist(), reading back numbers the receive had decoded on the host thirty lines earlier. They are kept on GpuAsyncTransferState instead. - the gate validated group_list with int(counts.sum()). - repeat_interleave sized its output by reading the counts back to the host. Passing output_size does both jobs at once: no readback, and it raises if the counts disagree with the row count, which is what the removed check tested. The header's own consistency is now checked where it is free, on the decoded host values in recv_attn_output. What is left on the layer path is the counts D2H in send_attn_output, which is intrinsic to slicing the per-rank segments on the host; the comment there says what removes it. Verified on L20X, DeepSeek-V2-Lite: tests/unit 619 passed, async_gpu_connector_e2e passes, async_gpu_moe_equivalence unchanged at max rel 4.5367e-03. Co-Authored-By: Claude Opus 5 Signed-off-by: specture724 --- afd_plugin/connectors/gpu/async_gpu.py | 52 ++++++++++--------- afd_plugin/connectors/gpu/symm_window.py | 16 ++++++ .../models/gpu/deepseek_v2_attention_gate.py | 11 ++-- 3 files changed, 48 insertions(+), 31 deletions(-) diff --git a/afd_plugin/connectors/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py index 43904e87..2942213a 100644 --- a/afd_plugin/connectors/gpu/async_gpu.py +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -21,8 +21,9 @@ from __future__ import annotations +import time from collections.abc import Mapping -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import timedelta from typing import TYPE_CHECKING, Any, Final @@ -194,6 +195,10 @@ class GpuAsyncTransferState(AFDTransferState): ``region``/``ring`` locate the window slot so ``send_ffn_work_item_output`` can write back to the originating Attention rank and release the slot; ``route_table`` is echoed so the Attention side can scatter the result. + + ``group_list`` is a device view of the arrived header's trailing words; + ``expert_counts_host`` is the same numbers as they were decoded on the host, + kept so the combine header can be built without reading the device back. """ region: int = 0 @@ -206,6 +211,7 @@ class GpuAsyncTransferState(AFDTransferState): routed_tokens: int = 0 shared_tokens: int = 0 group_list: Tensor | None = None + expert_counts_host: list[int] = field(default_factory=list) route_table: Tensor | None = None shared_idx: Tensor | None = None expand_x_shared: Tensor | None = None @@ -496,7 +502,10 @@ def send_attn_output( expert_per_rank=self.expert_per_rank, ) # One D2H per send: offsets are a prefix sum, cheaper to redo on host - # than to fetch a second tensor. + # than to fetch a second tensor. This is the last synchronize left on + # the layer path, and it is intrinsic to slicing the segments on the + # host -- removing it means computing the destination offsets on the + # device and writing full-capacity segments instead. counts_host = counts.cpu().tolist() offsets_host = [0] * len(counts_host) for i in range(1, len(counts_host)): @@ -678,27 +687,30 @@ def recv_attn_output( """ window = self._require_initialized() timeout_ms = int(kwargs.get("timeout_ms", 0)) - deadline = None - if timeout_ms: - import time - - deadline = time.monotonic() + timeout_ms / 1000.0 + deadline = time.monotonic() + timeout_ms / 1000.0 if timeout_ms else None while True: arrived = window.poll() if arrived is not None: break - if deadline is not None: - import time - - if time.monotonic() >= deadline: - raise TimeoutError("AFD async GPU dispatch recv timed out") + if deadline is not None and time.monotonic() >= deadline: + raise TimeoutError("AFD async GPU dispatch recv timed out") header = arrived.header if header.is_shutdown: raise ConnectorShutdown( f"Attention rank {header.src_role_rank} announced shutdown", ) + # The group list is echoed back on the combine header, so the sender's + # own accounting has to agree with 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{header.src_role_rank} " + f"has expert counts summing to {sum(header.expert_counts)} but " + f"declares {header.routed_tokens} routed tokens", + ) states = GpuAsyncTransferState( region=arrived.region, @@ -710,11 +722,8 @@ def recv_attn_output( num_tokens=header.num_tokens, routed_tokens=header.routed_tokens, shared_tokens=header.shared_tokens, - group_list=torch.tensor( - header.expert_counts, - dtype=torch.int64, - device=torch.device("cuda", self.local_rank), - ), + group_list=window.local_expert_counts(arrived.region, arrived.ring), + expert_counts_host=header.expert_counts, route_table=window.local_route_table( arrived.region, arrived.ring, @@ -783,7 +792,7 @@ def send_ffn_output( shared_tokens=states.shared_tokens if shared_output is not None else 0, topk=self.topk, flags=0, - expert_counts=list(header_counts(states)), + expert_counts=states.expert_counts_host, echo_seq=states.seq, ) window.write_slot( @@ -875,13 +884,6 @@ def announce_shutdown(self) -> None: ) -def header_counts(states: GpuAsyncTransferState) -> list[int]: - """Echo the per-expert group list back on the combine header.""" - if states.group_list is None: - return [] - return states.group_list.cpu().tolist() - - __all__ = [ "AFD_ASYNC_GPU_GROUP_NAME", "ConnectorShutdown", diff --git a/afd_plugin/connectors/gpu/symm_window.py b/afd_plugin/connectors/gpu/symm_window.py index 8d32a728..95c573a7 100644 --- a/afd_plugin/connectors/gpu/symm_window.py +++ b/afd_plugin/connectors/gpu/symm_window.py @@ -523,6 +523,22 @@ def read_header(self, region: int, ring: int) -> SlotHeader: ) return decode_header(self._header_recv) + def local_expert_counts(self, region: int, ring: int) -> torch.Tensor: + """Device view of the arrived header's per-expert counts. + + The counts already sit in device memory as the header's trailing words, + so the grouped GEMM can read them straight from the slot. Rebuilding the + tensor from the decoded host list would cost one blocking H2D per work + item. + """ + header = self._capacity_view( + self.rank, + region, + ring, + self.layout.header_off, + ) + return header[HEADER_FIXED_WORDS:] + def local_route_table(self, region: int, ring: int, count: int) -> torch.Tensor: return self._view( self.rank, 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 index 75764c81..abe932e2 100644 --- a/afd_plugin/model_executor/models/gpu/deepseek_v2_attention_gate.py +++ b/afd_plugin/model_executor/models/gpu/deepseek_v2_attention_gate.py @@ -53,12 +53,6 @@ def compute_attention_gate_moe_ffn( routed_experts = runner.routed_experts counts = group_list.to(torch.int64) num_rows = int(hidden_states.shape[0]) - if int(counts.sum()) != num_rows: - raise ValueError( - f"group_list sums to {int(counts.sum())} but hidden_states has " - f"{num_rows} rows", - ) - num_local_experts = counts.numel() if num_rows == 0: # Routing can leave a peer with nothing -- common in decode, where a @@ -75,6 +69,10 @@ def compute_attention_gate_moe_ffn( 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, @@ -82,6 +80,7 @@ def compute_attention_gate_moe_ffn( dtype=torch.int32, ), counts, + output_size=num_rows, ).unsqueeze(1) # Unit weights: the real topk weighting happens in the connector's combine. unit_weights = torch.ones( From 5b20d7cdd23a969ba8312b037fe6dc4f46abbb5d Mon Sep 17 00:00:00 2001 From: specture724 Date: Fri, 14 Aug 2026 10:02:46 +0800 Subject: [PATCH 05/18] perf: ship each token once per destination and reduce on the FFN side Dispatch sent one row per (token, topk slot), so a token whose topk landed several times on the same FFN rank crossed the wire several times. With DeepSeek-V2-Lite's topk=6 over 2 FFN ranks that was measured at 1536 rows per (A,F) pair for a 512-token batch against 512 distinct tokens: 3x the traffic, for rows the receiver could rebuild locally. Dispatch now carries the distinct tokens plus, per partial, a 4-byte row index and its topk weight. The FFN side expands the rows back with a local gather before the grouped GEMM, then weights and sums each token's partials in float32 before replying, so the return trip carries the same distinct rows. Combine is left a plain scatter-add; combine_scatter and the route table are gone from both the wire and the code. Measured on 2A2F, 512-token batches: 512 rows per pair per direction instead of 1536. Against the sync connector's 512, the routed traffic goes from 3.0x to 1.0x, and 1.5x with the shared-expert rows counted. The window shrinks with it. Payload rows are bounded by the batch no matter how skewed the gate is, so the payload is sized by token_cap and only the 4-byte index arrays carry the every-partial-to-one-rank worst case. A 2A2F slot goes from 14.0 MiB to 4.0 MiB, capacity no longer grows with ffn_size**2, and routed_cap_multiplier -- which had no safe value above ffn_size=2 -- is deleted rather than retuned. Reduction order changes: a token's partials are now summed in float32 before the single narrowing to the payload dtype, where each partial used to be narrowed separately and summed on the Attention side. That is one rounding instead of many, so greedy output shifts slightly while getting marginally more accurate. Verified on 4x L20X, DeepSeek-V2-Lite: tests/unit 619 passed, async_gpu_window_roundtrip and async_gpu_connector_e2e pass on the new format, async_gpu_moe_equivalence at max rel 4.5367e-03, and a 2A2F server answers the greedy probe sensibly, matching the sync connector on 3 of 5 prompts. Co-Authored-By: Claude Opus 5 Signed-off-by: specture724 --- afd_plugin/connectors/gpu/async_gpu.py | 296 ++++++++++++------ afd_plugin/connectors/gpu/symm_window.py | 90 ++++-- .../connectors/test_async_gpu_connector.py | 172 ++++++---- 3 files changed, 369 insertions(+), 189 deletions(-) diff --git a/afd_plugin/connectors/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py index 2942213a..d5b45c29 100644 --- a/afd_plugin/connectors/gpu/async_gpu.py +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -67,7 +67,6 @@ { "attn_ranks_per_dp", "ring_depth", - "routed_cap_multiplier", "recv_poll_timeout_ms", "async_moe_ubatching", "async_moe_num_ubatches", @@ -82,15 +81,6 @@ logger = init_logger(f"vllm.{__name__}") -def _coerce_extra_float(value: Any, *, field_name: str) -> float: - if isinstance(value, bool) or not isinstance(value, (int, float, str)): - raise TypeError(f"{field_name} must be a number, got {type(value).__name__}") - result = float(value) - if result <= 0.0: - raise ValueError(f"{field_name} must be positive, got {result}") - return result - - @dataclass(frozen=True) class GpuAsyncExtraInfo(ConnectorExtraInfo): """Typed async GPU connector configuration. @@ -100,10 +90,6 @@ class GpuAsyncExtraInfo(ConnectorExtraInfo): ring_depth: Slots per peer region. Derived from the send-then-recv invariant, not a performance knob: an Attention rank has at most one in-flight request per ``(peer, stage)``, so ``num_stages`` suffices. - routed_cap_multiplier: Headroom over the balanced routed-token estimate. - Real gates are not balanced -- DeepSeek-V2-Lite at 2 FFN ranks was - observed 1.33x over the even split -- so the default leaves room. - The true worst case is ``ffn_size`` (every partial to one rank). recv_poll_timeout_ms: Idle poll timeout on the FFN loop; bounds shutdown response time. async_moe_ubatching: Whether request-boundary async MoE ubatching is used. @@ -113,7 +99,6 @@ class GpuAsyncExtraInfo(ConnectorExtraInfo): attn_ranks_per_dp: int = 1 ring_depth: int = 0 - routed_cap_multiplier: float = 2.0 recv_poll_timeout_ms: int = 50 async_moe_ubatching: bool = False async_moe_num_ubatches: int = 2 @@ -156,10 +141,6 @@ def from_mapping(cls, raw: Mapping[str, Any] | None) -> GpuAsyncExtraInfo: field_name="attn_ranks_per_dp", ), ring_depth=ring_depth, - routed_cap_multiplier=_coerce_extra_float( - raw.get("routed_cap_multiplier", 2.0), - field_name="routed_cap_multiplier", - ), recv_poll_timeout_ms=coerce_extra_positive_int( raw.get("recv_poll_timeout_ms", 50), field_name="recv_poll_timeout_ms", @@ -176,7 +157,6 @@ def to_mapping(self) -> dict[str, Any]: return { "attn_ranks_per_dp": self.attn_ranks_per_dp, "ring_depth": self.ring_depth, - "routed_cap_multiplier": self.routed_cap_multiplier, "recv_poll_timeout_ms": self.recv_poll_timeout_ms, "async_moe_ubatching": self.async_moe_ubatching, "async_moe_num_ubatches": self.async_moe_num_ubatches, @@ -193,12 +173,15 @@ class GpuAsyncTransferState(AFDTransferState): """FFN-side state carried from dispatch recv through combine send. ``region``/``ring`` locate the window slot so ``send_ffn_work_item_output`` - can write back to the originating Attention rank and release the slot; - ``route_table`` is echoed so the Attention side can scatter the result. + 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; ``expert_counts_host`` is the same numbers as they were decoded on the host, kept so the combine header can be built without reading the device back. + + ``expand_idx`` and ``weights`` describe the partials: which shipped row each + one reads on the way in, and what to weight it by when the expert output is + reduced back to one row per shipped token on the way out. """ region: int = 0 @@ -209,10 +192,12 @@ class GpuAsyncTransferState(AFDTransferState): seq: int = 0 num_tokens: int = 0 routed_tokens: int = 0 + uniq_tokens: int = 0 shared_tokens: int = 0 group_list: Tensor | None = None expert_counts_host: list[int] = field(default_factory=list) - route_table: Tensor | None = None + expand_idx: Tensor | None = None + weights: Tensor | None = None shared_idx: Tensor | None = None expand_x_shared: Tensor | None = None @@ -233,58 +218,114 @@ class GpuAsyncFFNWorkItem: @dataclass(slots=True) class _PendingDispatch: - """Attention-side record of one in-flight layer, popped by combine recv.""" + """Attention-side record of one in-flight layer, popped by combine recv. + + ``uniq_ids[r]`` holds the token ids behind the rows shipped to FFN rank + ``r``, in the order they were shipped. The reply comes back one row per + shipped token, already weighted and summed, so combine is a scatter-add at + exactly those ids. + """ context: AFDTransferContext - topk_weights: Tensor + uniq_ids: list[Tensor] num_tokens: int ring: int seq: int expected_ffn: list[int] +@dataclass(frozen=True, slots=True) +class DispatchPlan: + """One layer's routing plan, already in the order every consumer wants. + + Partial-indexed fields (``expand_idx``, ``weights``, length + ``num_tokens * topk``) are sorted by global expert, so a destination's slice + arrives grouped by local expert and feeds the grouped GEMM directly. + Row-indexed fields (``uniq_token_ids``) are grouped by destination rank, so + a destination's payload is also one contiguous slice. + + Attributes: + counts: Partials per global expert, padded to + ``ffn_size * expert_per_rank``. + offsets: Exclusive prefix sum of ``counts``. + uniq_per_rank: Distinct tokens each FFN rank receives. + expand_idx: Per partial, which of its destination's shipped rows it + reads. Already rank-local, so a slice needs no rebasing. + uniq_token_ids: Per shipped row, the token it carries. + weights: Per partial, its topk weight. + """ + + counts: Tensor + offsets: Tensor + uniq_per_rank: Tensor + expand_idx: Tensor + uniq_token_ids: Tensor + weights: Tensor + + def plan_dispatch( topk_ids: Tensor, + topk_weights: Tensor, *, ffn_size: int, expert_per_rank: int, -) -> tuple[Tensor, Tensor, Tensor]: - """Cluster ``(token, topk_slot)`` partials by destination FFN rank. +) -> DispatchPlan: + """Cluster ``(token, topk_slot)`` partials by destination and deduplicate. Sorting by the global expert id groups partials by destination rank and, in - the same pass, by local expert inside each destination -- the two orderings - the receiver needs. Returns ``(route_table, counts, offsets)`` where - ``route_table[i] = (token_idx, topk_slot)`` in that order, ``counts`` holds - per-global-expert partial counts padded to ``ffn_size * expert_per_rank``, - and ``offsets`` is the exclusive prefix sum of ``counts``. + the same pass, by local expert inside each destination. A second sort by + ``(destination, token)`` makes the partials that share a shipped row + adjacent, which is what turns ``topk`` copies of a token into one. + + Every step is a whole-tensor op: the host learns nothing here, so the caller + can fetch ``counts`` and ``uniq_per_rank`` in a single readback. """ - num_slots = topk_ids.shape[1] + num_tokens, num_slots = topk_ids.shape flat = topk_ids.reshape(-1).to(torch.int64) order = torch.argsort(flat, stable=True) counts = torch.bincount(flat, minlength=ffn_size * expert_per_rank) offsets = torch.cumsum(counts, dim=0) - counts - route_table = torch.stack( - (order // num_slots, order % num_slots), - dim=1, - ).to(torch.int32) - return route_table, counts, offsets - -def combine_scatter( - accumulator: Tensor, - *, - routed_out: Tensor, - route_table: Tensor, - topk_weights: Tensor, -) -> None: - """Weighted-scatter one FFN rank's routed output into ``accumulator``.""" - token_idx = route_table[:, 0].to(torch.int64) - slot_idx = route_table[:, 1].to(torch.int64) - weights = topk_weights[token_idx, slot_idx].to(accumulator.dtype) - accumulator.index_add_( - 0, - token_idx, - routed_out.to(accumulator.dtype) * weights.unsqueeze(1), + token_of_partial = order // num_slots + weights = topk_weights.reshape(-1)[order].to(torch.float32) + dest_rank = flat[order] // expert_per_rank + + # Partials sharing a (destination, token) share a shipped row. Sorting by + # that key makes them adjacent, so "is this row new" is one neighbour + # comparison and the row numbering is its running sum. + key = dest_rank * num_tokens + token_of_partial + key_order = torch.argsort(key, stable=True) + sorted_key = key[key_order] + is_new = torch.ones_like(sorted_key, dtype=torch.int32) + is_new[1:] = (sorted_key[1:] != sorted_key[:-1]).to(torch.int32) + row_of_partial = is_new.cumsum(0) - 1 + + uniq_per_rank = torch.zeros( + ffn_size, + dtype=torch.int32, + device=topk_ids.device, + ) + uniq_per_rank.index_add_(0, dest_rank[key_order], is_new) + + # Duplicates write the same token id to the same row, so the scatter is + # well defined despite the repeated indices. + uniq_token_ids = torch.zeros_like(token_of_partial) + uniq_token_ids.scatter_(0, row_of_partial, token_of_partial[key_order]) + + # Back to expert-sorted order, rebased so each destination's slice indexes + # its own payload from zero. + row_base = (torch.cumsum(uniq_per_rank, 0) - uniq_per_rank).to(torch.int32) + expand_idx = torch.empty_like(row_of_partial, dtype=torch.int32) + expand_idx.scatter_(0, key_order, row_of_partial.to(torch.int32)) + expand_idx -= row_base[dest_rank] + + return DispatchPlan( + counts=counts, + offsets=offsets, + uniq_per_rank=uniq_per_rank, + expand_idx=expand_idx, + uniq_token_ids=uniq_token_ids, + weights=weights, ) @@ -334,15 +375,14 @@ def __init__( # region per opposite-role peer. Both roles allocate the larger of the # two so the symmetric allocation matches. self.num_regions = max(self.attn_size, self.ffn_size) - routed_cap = int( - -(-self.max_seq_len * self.topk // self.ffn_size) - * self.extra_info.routed_cap_multiplier, - ) - self.routed_cap = max(1, routed_cap) + # 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) self.layout = SlotLayout.build( expert_per_rank=self.expert_per_rank, - routed_cap=self.routed_cap, + partial_cap=self.partial_cap, token_cap=self.token_cap, hidden_size=self.hidden_size, payload_itemsize=torch.empty(0, dtype=self.payload_dtype).element_size(), @@ -391,14 +431,14 @@ def init_afd_connector(self) -> None: self._free_rings[stage] = list(range(self.ring_depth)) logger.info( "AFD async GPU window ready: role=%s role_rank=%d world_rank=%d/%d " - "regions=%d rings=%d routed_cap=%d slot=%.1fMiB total=%.1fMiB", + "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.routed_cap, + self.partial_cap, self.layout.slot_bytes / 2**20, self.window.total_bytes / 2**20, ) @@ -479,10 +519,17 @@ def send_attn_output( f"hidden_states has {hidden_states.shape[0]} rows but metadata " f"expects {num_tokens}", ) - if tuple(topk_ids.shape) != (num_tokens, self.topk): + expected_shape = (num_tokens, self.topk) + if tuple(topk_ids.shape) != expected_shape: raise ValueError( - f"topk_ids shape must be ({num_tokens}, {self.topk}), " - f"got {tuple(topk_ids.shape)}", + 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 @@ -496,21 +543,27 @@ def send_attn_output( ring = rings.pop(0) self._seq += 1 - route_table, counts, _ = plan_dispatch( + plan = plan_dispatch( topk_ids, + topk_weights, ffn_size=self.ffn_size, expert_per_rank=self.expert_per_rank, ) - # One D2H per send: offsets are a prefix sum, cheaper to redo on host - # than to fetch a second tensor. This is the last synchronize left on - # the layer path, and it is intrinsic to slicing the segments on the - # host -- removing it means computing the destination offsets on the - # device and writing full-capacity segments instead. - counts_host = counts.cpu().tolist() - offsets_host = [0] * len(counts_host) - for i in range(1, len(counts_host)): + # One D2H per send, carrying everything the host needs to slice the + # plan: offsets are a prefix sum, cheaper to redo here than to fetch. + # This is the last synchronize left on the layer path, and it is + # intrinsic to slicing the segments on the host -- removing it means + # computing the destination offsets on the device and writing + # full-capacity segments instead. + num_experts = self.ffn_size * self.expert_per_rank + plan_host = ( + torch.cat((plan.counts.to(torch.int32), plan.uniq_per_rank)).cpu().tolist() + ) + counts_host = plan_host[:num_experts] + uniq_host = plan_host[num_experts:] + offsets_host = [0] * num_experts + for i in range(1, num_experts): offsets_host[i] = offsets_host[i - 1] + counts_host[i - 1] - token_ids = route_table[:, 0].to(torch.int64) # Every FFN rank gets a slot even when routing sends it nothing, and it # replies to every slot, so a reply is expected from all of them. @@ -519,12 +572,21 @@ def send_attn_output( # 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)) + uniq_ids: list[Tensor] = [] + uniq_start = 0 for ffn_rank in range(self.ffn_size): base = ffn_rank * self.expert_per_rank expert_counts = counts_host[base : base + self.expert_per_rank] start = offsets_host[base] routed_tokens = sum(expert_counts) segment = slice(start, start + routed_tokens) + # The shipped rows this rank owns, and the tokens they carry. The + # ids stay here rather than going on the wire: combine scatters the + # reply back to exactly these rows. + uniq_tokens = uniq_host[ffn_rank] + rows = plan.uniq_token_ids[uniq_start : uniq_start + uniq_tokens] + uniq_start += uniq_tokens + uniq_ids.append(rows) # Shared-expert tokens are split round-robin across FFN ranks. shared_idx = torch.arange( @@ -546,6 +608,7 @@ def send_attn_output( topk=self.topk, flags=0, expert_counts=expert_counts, + uniq_tokens=uniq_tokens, ) # Rows are gathered straight into the FFN rank's window; handing # ``write_slot`` the index instead of a gathered tensor saves a @@ -555,9 +618,10 @@ def send_attn_output( region=self.role_rank, ring=ring, header=header, - route_table=route_table[segment], + expand_idx=plan.expand_idx[segment], + weights=plan.weights[segment], routed_x=hidden_states, - routed_rows=token_ids[segment], + routed_rows=rows, shared_idx=shared_idx, shared_x=hidden_states, shared_rows=shared_idx.to(torch.int64), @@ -565,18 +629,20 @@ def send_attn_output( logger.debug( "AFD dispatch sent: A%d layer=%d stage=%d tokens=%d ring=%d " - "awaiting_ffn=%s", + "rows_per_ffn=%s partials=%d awaiting_ffn=%s", self.role_rank, metadata.layer_idx, stage_idx, num_tokens, ring, + uniq_host, + num_tokens * self.topk, expected_ffn, ) self._pending.setdefault(stage_idx, []).append( _PendingDispatch( context=context, - topk_weights=topk_weights, + uniq_ids=uniq_ids, num_tokens=num_tokens, ring=ring, seq=self._seq, @@ -638,20 +704,17 @@ def recv_ffn_output( sorted(outstanding), ) - if header.routed_tokens: - combine_scatter( - accumulator, - routed_out=window.local_routed( - arrived.region, - arrived.ring, - header.routed_tokens, - ), - route_table=window.local_route_table( + if header.uniq_tokens: + # One row per token this rank was sent, already weighted and + # summed over that token's partials on the FFN side. + accumulator.index_add_( + 0, + pending.uniq_ids[header.src_role_rank], + window.local_routed( arrived.region, arrived.ring, - header.routed_tokens, - ), - topk_weights=pending.topk_weights, + header.uniq_tokens, + ).to(accumulator.dtype), ) if header.shared_tokens: accumulator.index_add_( @@ -684,6 +747,10 @@ def recv_attn_output( 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 payload carries each token once, so the rows are expanded back to + one per partial here -- a local gather that replaces the duplicate rows + the sender used to put on the wire. """ window = self._require_initialized() timeout_ms = int(kwargs.get("timeout_ms", 0)) @@ -712,6 +779,11 @@ def recv_attn_output( f"declares {header.routed_tokens} routed tokens", ) + expand_idx = window.local_expand_idx( + arrived.region, + arrived.ring, + header.routed_tokens, + ).to(torch.int64) states = GpuAsyncTransferState( region=arrived.region, ring=arrived.ring, @@ -721,10 +793,12 @@ def recv_attn_output( stage_idx=header.stage_idx, num_tokens=header.num_tokens, routed_tokens=header.routed_tokens, + uniq_tokens=header.uniq_tokens, shared_tokens=header.shared_tokens, group_list=window.local_expert_counts(arrived.region, arrived.ring), expert_counts_host=header.expert_counts, - route_table=window.local_route_table( + expand_idx=expand_idx, + weights=window.local_weights( arrived.region, arrived.ring, header.routed_tokens, @@ -761,8 +835,8 @@ def recv_attn_output( hidden_states=window.local_routed( arrived.region, arrived.ring, - header.routed_tokens, - ), + header.uniq_tokens, + ).index_select(0, expand_idx), context=AFDTransferContext(metadata=metadata, states=states), ) @@ -772,13 +846,33 @@ def send_ffn_output( context: AFDTransferContext, **kwargs: Any, ) -> None: - """Write expert output back to the originating Attention rank.""" + """Reduce expert output to one row per shipped token and write it back. + + Every partial of a token that landed on this rank is weighted and summed + here, so the reply carries the same rows the dispatch did. Doing it on + this side keeps the duplicates off the wire and leaves the Attention + side a plain scatter-add. + """ 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.uniq_tokens, self.hidden_size), + dtype=torch.float32, + device=ffn_output.device, + ) + if states.routed_tokens: + weighted = ffn_output.to(torch.float32) + weighted.mul_(states.weights.unsqueeze(1)) + reduced.index_add_(0, states.expand_idx, weighted) + shared_output: Tensor | None = kwargs.get("shared_output") self._seq += 1 header = encode_header( @@ -793,15 +887,20 @@ def send_ffn_output( topk=self.topk, flags=0, expert_counts=states.expert_counts_host, + uniq_tokens=states.uniq_tokens, echo_seq=states.seq, ) + # ``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, header=header, - route_table=states.route_table, - routed_x=ffn_output, + expand_idx=None, + weights=None, + routed_x=reduced, shared_idx=states.shared_idx if shared_output is not None else None, shared_x=shared_output, ) @@ -877,7 +976,6 @@ def announce_shutdown(self) -> None: region=self.role_rank, ring=0, header=header, - route_table=None, routed_x=None, shared_idx=None, shared_x=None, @@ -887,10 +985,10 @@ def announce_shutdown(self) -> None: __all__ = [ "AFD_ASYNC_GPU_GROUP_NAME", "ConnectorShutdown", + "DispatchPlan", "GpuAsyncAFDConnector", "GpuAsyncExtraInfo", "GpuAsyncFFNWorkItem", "GpuAsyncTransferState", - "combine_scatter", "plan_dispatch", ] diff --git a/afd_plugin/connectors/gpu/symm_window.py b/afd_plugin/connectors/gpu/symm_window.py index 95c573a7..30436c83 100644 --- a/afd_plugin/connectors/gpu/symm_window.py +++ b/afd_plugin/connectors/gpu/symm_window.py @@ -16,6 +16,14 @@ payloads. Dispatch (A -> F) and combine (F -> A) use the same slot layout, so a single spec sizes both directions. +A token is shipped once per destination rank however many of its topk slots +landed there, so the payload rows are *distinct tokens* (``uniq_tokens``, at +most ``token_cap``) while ``expand_idx`` and ``weights`` carry one entry per +*partial* (``routed_tokens``, at most ``partial_cap``). Those two index arrays +are 4 bytes a row against the payload's ``hidden_size`` elements, which is why +they can be sized for the true worst case -- every partial to one rank -- while +the payload is sized by a bound that no routing skew can exceed. + Flag words are written *after* the payload on the same stream. Same-stream device-to-device copies complete in issue order, so a visible flag implies a complete payload. That holds for NVLink-mapped peer memory; a cross-node @@ -41,7 +49,7 @@ # Header words shared by dispatch and combine, followed by expert_counts. HEADER_MAGIC = 0x41464447 # "AFDG" -HEADER_VERSION = 1 +HEADER_VERSION = 2 _H_MAGIC = 0 _H_VERSION = 1 _H_SEQ = 2 @@ -49,12 +57,13 @@ _H_LAYER_IDX = 4 _H_STAGE_IDX = 5 _H_NUM_TOKENS = 6 -_H_ROUTED_TOKENS = 7 +_H_ROUTED_TOKENS = 7 # partials: one per (token, topk slot) landing here _H_SHARED_TOKENS = 8 _H_TOPK = 9 _H_FLAGS = 10 _H_ECHO_SEQ = 11 # combine only: the dispatch seq being answered -HEADER_FIXED_WORDS = 12 +_H_UNIQ_TOKENS = 12 # payload rows: distinct tokens behind those partials +HEADER_FIXED_WORDS = 13 FLAG_EMPTY = 0 FLAG_SHUTDOWN_BIT = 1 << 1 @@ -73,13 +82,14 @@ class SlotLayout: """Byte offsets and element counts for the fields inside one slot.""" header_words: int - routed_cap: int + partial_cap: int token_cap: int hidden_size: int payload_itemsize: int header_off: int - route_table_off: int + expand_idx_off: int + weights_off: int routed_x_off: int shared_idx_off: int shared_x_off: int @@ -90,28 +100,30 @@ def build( cls, *, expert_per_rank: int, - routed_cap: int, + partial_cap: int, token_cap: int, hidden_size: int, payload_itemsize: int, ) -> SlotLayout: header_words = HEADER_FIXED_WORDS + expert_per_rank header_off = 0 - route_table_off = _align(header_off + header_words * 4) - routed_x_off = _align(route_table_off + 2 * routed_cap * 4) + expand_idx_off = _align(header_off + header_words * 4) + weights_off = _align(expand_idx_off + partial_cap * 4) + routed_x_off = _align(weights_off + partial_cap * 4) shared_idx_off = _align( - routed_x_off + routed_cap * hidden_size * payload_itemsize + routed_x_off + token_cap * hidden_size * payload_itemsize ) shared_x_off = _align(shared_idx_off + token_cap * 4) slot_bytes = _align(shared_x_off + token_cap * hidden_size * payload_itemsize) return cls( header_words=header_words, - routed_cap=routed_cap, + partial_cap=partial_cap, token_cap=token_cap, hidden_size=hidden_size, payload_itemsize=payload_itemsize, header_off=header_off, - route_table_off=route_table_off, + expand_idx_off=expand_idx_off, + weights_off=weights_off, routed_x_off=routed_x_off, shared_idx_off=shared_idx_off, shared_x_off=shared_x_off, @@ -132,6 +144,7 @@ def encode_header( topk: int, flags: int, expert_counts: list[int], + uniq_tokens: int = 0, echo_seq: int = 0, ) -> torch.Tensor: """Build the fixed header for one slot as a CPU int32 tensor.""" @@ -154,6 +167,7 @@ def encode_header( header[_H_TOPK] = topk header[_H_FLAGS] = flags header[_H_ECHO_SEQ] = echo_seq + header[_H_UNIQ_TOKENS] = uniq_tokens if expert_per_rank: header[HEADER_FIXED_WORDS:] = torch.tensor(expert_counts, dtype=torch.int32) return header @@ -173,6 +187,7 @@ class SlotHeader: topk: int flags: int echo_seq: int + uniq_tokens: int expert_counts: list[int] @property @@ -204,6 +219,7 @@ def decode_header(header: torch.Tensor) -> SlotHeader: topk=values[_H_TOPK], flags=values[_H_FLAGS], echo_seq=values[_H_ECHO_SEQ], + uniq_tokens=values[_H_UNIQ_TOKENS], expert_counts=values[HEADER_FIXED_WORDS:], ) @@ -330,10 +346,12 @@ def _capacity_view( hidden = layout.hidden_size if field_off == layout.header_off: sizes, dtype = (layout.header_words,), torch.int32 - elif field_off == layout.route_table_off: - sizes, dtype = (layout.routed_cap, 2), torch.int32 + elif field_off == layout.expand_idx_off: + sizes, dtype = (layout.partial_cap,), torch.int32 + elif field_off == layout.weights_off: + sizes, dtype = (layout.partial_cap,), torch.float32 elif field_off == layout.routed_x_off: - sizes, dtype = (layout.routed_cap, hidden), self.payload_dtype + sizes, dtype = (layout.token_cap, hidden), self.payload_dtype elif field_off == layout.shared_idx_off: sizes, dtype = (layout.token_cap,), torch.int32 elif field_off == layout.shared_x_off: @@ -389,7 +407,8 @@ def write_slot( region: int, ring: int, header: torch.Tensor, - route_table: torch.Tensor | None, + expand_idx: torch.Tensor | None, + weights: torch.Tensor | None, routed_x: torch.Tensor | None, shared_idx: torch.Tensor | None, shared_x: torch.Tensor | None, @@ -409,10 +428,14 @@ def write_slot( layout = self.layout routed_count = _row_count(routed_x, routed_rows) shared_count = _row_count(shared_x, shared_rows) - if routed_count > layout.routed_cap: + partial_count = _row_count(expand_idx, None) + if routed_count > layout.token_cap: + raise RuntimeError( + f"payload rows {routed_count} exceed token_cap {layout.token_cap}", + ) + if partial_count > layout.partial_cap: raise RuntimeError( - f"routed tokens {routed_count} exceed routed_cap " - f"{layout.routed_cap}; raise routed_cap_multiplier", + f"partials {partial_count} exceed partial_cap {layout.partial_cap}", ) if shared_count > layout.token_cap: raise RuntimeError( @@ -453,11 +476,18 @@ def write_field( torch.index_select(src, 0, rows, out=view) write_field( - layout.route_table_off, - (2,), + layout.expand_idx_off, + (), torch.int32, - route_table, - _row_count(route_table, None), + expand_idx, + partial_count, + ) + write_field( + layout.weights_off, + (), + torch.float32, + weights, + _row_count(weights, None), ) write_field( layout.routed_x_off, @@ -539,16 +569,26 @@ def local_expert_counts(self, region: int, ring: int) -> torch.Tensor: ) return header[HEADER_FIXED_WORDS:] - def local_route_table(self, region: int, ring: int, count: int) -> torch.Tensor: + def local_expand_idx(self, region: int, ring: int, count: int) -> torch.Tensor: return self._view( self.rank, region, ring, - self.layout.route_table_off, - (count, 2), + self.layout.expand_idx_off, + (count,), torch.int32, ) + def local_weights(self, region: int, ring: int, count: int) -> torch.Tensor: + return self._view( + self.rank, + region, + ring, + self.layout.weights_off, + (count,), + torch.float32, + ) + def local_routed(self, region: int, ring: int, count: int) -> torch.Tensor: return self._view( self.rank, diff --git a/tests/unit/connectors/test_async_gpu_connector.py b/tests/unit/connectors/test_async_gpu_connector.py index bf320309..7d797473 100644 --- a/tests/unit/connectors/test_async_gpu_connector.py +++ b/tests/unit/connectors/test_async_gpu_connector.py @@ -14,7 +14,6 @@ from afd_plugin.connectors.gpu.async_gpu import ( # noqa: E402 GpuAsyncAFDConnector, GpuAsyncExtraInfo, - combine_scatter, plan_dispatch, ) from afd_plugin.connectors.gpu.symm_window import ( # noqa: E402 @@ -30,7 +29,7 @@ def layout() -> SlotLayout: return SlotLayout.build( expert_per_rank=4, - routed_cap=100, + partial_cap=100, token_cap=32, hidden_size=8, payload_itemsize=2, @@ -59,11 +58,6 @@ def test_unknown_extra_config_field_is_rejected(): GpuAsyncExtraInfo.from_mapping({"nope": 1}) -def test_routed_cap_multiplier_must_be_positive(): - with pytest.raises(ValueError, match="routed_cap_multiplier"): - GpuAsyncExtraInfo.from_mapping({"routed_cap_multiplier": 0}) - - # ---------------------------------------------------------------------- # Slot layout # ---------------------------------------------------------------------- @@ -71,9 +65,11 @@ def test_routed_cap_multiplier_must_be_positive(): def test_slot_fields_are_disjoint_and_fit_inside_the_slot(layout: SlotLayout): assert layout.header_words == HEADER_FIXED_WORDS + 4 - assert layout.route_table_off >= layout.header_off + layout.header_words * 4 - assert layout.routed_x_off >= layout.route_table_off + 2 * 100 * 4 - assert layout.shared_idx_off >= layout.routed_x_off + 100 * 8 * 2 + assert layout.expand_idx_off >= layout.header_off + layout.header_words * 4 + assert layout.weights_off >= layout.expand_idx_off + 100 * 4 + assert layout.routed_x_off >= layout.weights_off + 100 * 4 + # The payload is sized by distinct tokens, not by partials. + assert layout.shared_idx_off >= layout.routed_x_off + 32 * 8 * 2 assert layout.shared_x_off >= layout.shared_idx_off + 32 * 4 assert layout.slot_bytes >= layout.shared_x_off + 32 * 8 * 2 @@ -83,7 +79,8 @@ def test_every_field_offset_is_viewable_as_int32_and_payload(layout: SlotLayout) # multiple of the element size would silently land on the wrong address. for offset in ( layout.header_off, - layout.route_table_off, + layout.expand_idx_off, + layout.weights_off, layout.routed_x_off, layout.shared_idx_off, layout.shared_x_off, @@ -110,6 +107,7 @@ def test_header_round_trip(layout: SlotLayout): topk=6, flags=0, expert_counts=[10, 20, 30, 30], + uniq_tokens=25, ) decoded = decode_header(header) assert decoded.seq == 7 @@ -121,6 +119,7 @@ def test_header_round_trip(layout: SlotLayout): assert decoded.shared_tokens == 16 assert decoded.topk == 6 assert decoded.expert_counts == [10, 20, 30, 30] + assert decoded.uniq_tokens == 25 assert sum(decoded.expert_counts) == decoded.routed_tokens assert not decoded.is_shutdown @@ -204,60 +203,107 @@ def routing_inputs(): return topk_ids, hidden_states, topk_weights +def _destination_slices(plan, ffn_size, expert_per_rank): + """Walk the plan the way send_attn_output does: one slice per destination.""" + counts, offsets = plan.counts.tolist(), plan.offsets.tolist() + uniq = plan.uniq_per_rank.tolist() + uniq_start = 0 + for ffn_rank in range(ffn_size): + base = ffn_rank * expert_per_rank + total = sum(counts[base : base + expert_per_rank]) + partials = slice(offsets[base], offsets[base] + total) + rows = plan.uniq_token_ids[uniq_start : uniq_start + uniq[ffn_rank]] + uniq_start += uniq[ffn_rank] + yield ffn_rank, partials, rows + + def test_every_partial_is_routed_exactly_once(routing_inputs): - topk_ids, _, _ = routing_inputs - route_table, counts, _ = plan_dispatch( + topk_ids, _, topk_weights = routing_inputs + plan = plan_dispatch( topk_ids, + topk_weights, ffn_size=_FFN_SIZE, expert_per_rank=_EXPERT_PER_RANK, ) - assert route_table.shape == (_NUM_TOKENS * _TOPK, 2) - assert int(counts.sum()) == _NUM_TOKENS * _TOPK - seen = {(token_idx, slot) for token_idx, slot in route_table.tolist()} - assert len(seen) == _NUM_TOKENS * _TOPK + 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.uniq_per_rank.sum()) <= _NUM_TOKENS * _TOPK def test_each_destination_segment_is_grouped_by_local_expert(routing_inputs): - topk_ids, _, _ = routing_inputs - route_table, counts, offsets = plan_dispatch( + topk_ids, _, topk_weights = routing_inputs + plan = plan_dispatch( topk_ids, + topk_weights, ffn_size=_FFN_SIZE, expert_per_rank=_EXPERT_PER_RANK, ) - counts_host, offsets_host = counts.tolist(), offsets.tolist() - for ffn_rank in range(_FFN_SIZE): + counts, offsets = plan.counts.tolist(), plan.offsets.tolist() + for ffn_rank, partials, rows in _destination_slices( + plan, + _FFN_SIZE, + _EXPERT_PER_RANK, + ): base = ffn_rank * _EXPERT_PER_RANK - cursor = offsets_host[base] + # The token behind each partial, in the order the receiver sees them. + tokens = rows.index_select(0, plan.expand_idx[partials].to(torch.int64)) + cursor = 0 for local_expert in range(_EXPERT_PER_RANK): - for _ in range(counts_host[base + local_expert]): - token_idx, slot = route_table[cursor].tolist() - assert int(topk_ids[token_idx, slot]) == base + local_expert + 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 == offsets_host[base] + sum( - counts_host[base : base + _EXPERT_PER_RANK], - ) + assert cursor == partials.stop - partials.start + assert offsets[base] == partials.start + + +def test_each_token_is_shipped_once_per_destination(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, + ) + for ffn_rank, partials, rows in _destination_slices( + plan, + _FFN_SIZE, + _EXPERT_PER_RANK, + ): + shipped = rows.tolist() + assert len(set(shipped)) == len(shipped), "a token was shipped twice" + 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(shipped) == expected + # Every partial must land inside its destination's own payload. + expand = plan.expand_idx[partials] + assert int(expand.min()) >= 0 + assert int(expand.max()) < len(shipped) def test_identity_experts_recombine_to_the_weighted_sum(routing_inputs): + """The full chain: ship distinct rows, expand, weight, reduce, scatter.""" topk_ids, hidden_states, topk_weights = routing_inputs - route_table, counts, offsets = plan_dispatch( + plan = plan_dispatch( topk_ids, + topk_weights, ffn_size=_FFN_SIZE, expert_per_rank=_EXPERT_PER_RANK, ) - counts_host, offsets_host = counts.tolist(), offsets.tolist() accumulator = torch.zeros(_NUM_TOKENS, _HIDDEN, dtype=torch.float32) - for ffn_rank in range(_FFN_SIZE): - base = ffn_rank * _EXPERT_PER_RANK - start = offsets_host[base] - total = sum(counts_host[base : base + _EXPERT_PER_RANK]) - segment = route_table[start : start + total] - combine_scatter( - accumulator, - routed_out=hidden_states.index_select(0, segment[:, 0].to(torch.int64)), - route_table=segment, - topk_weights=topk_weights, - ) + for _, partials, rows in _destination_slices(plan, _FFN_SIZE, _EXPERT_PER_RANK): + expand = plan.expand_idx[partials].to(torch.int64) + # What crosses the wire, and what the FFN side does with it. + expanded = hidden_states.index_select(0, rows).index_select(0, expand) + reduced = torch.zeros(rows.numel(), _HIDDEN, dtype=torch.float32) + reduced.index_add_(0, expand, expanded * plan.weights[partials].unsqueeze(1)) + accumulator.index_add_(0, rows, reduced) expected = hidden_states * topk_weights.sum(dim=1, keepdim=True) torch.testing.assert_close(accumulator, expected.to(torch.float32)) @@ -267,14 +313,15 @@ def test_routing_handles_experts_not_divisible_by_ffn_size(): # instead of silently absorbing real partials. ffn_size, expert_per_rank, num_experts = 3, 2, 5 topk_ids = torch.tensor([[0, 4], [1, 3], [2, 4]], dtype=torch.int32) - _, counts, _ = plan_dispatch( + plan = plan_dispatch( topk_ids, + torch.ones(3, 2), ffn_size=ffn_size, expert_per_rank=expert_per_rank, ) - assert counts.numel() == ffn_size * expert_per_rank - assert int(counts.sum()) == topk_ids.numel() - assert int(counts[num_experts:].sum()) == 0 + 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(): @@ -286,29 +333,24 @@ def test_routing_can_leave_one_destination_empty(): ffn_size, expert_per_rank = 2, 4 # Every partial targets experts owned by FFN rank 0. topk_ids = torch.tensor([[0, 1, 2]], dtype=torch.int32) - route_table, counts, offsets = plan_dispatch( + plan = plan_dispatch( topk_ids, + torch.ones(1, 3), ffn_size=ffn_size, expert_per_rank=expert_per_rank, ) - counts_host, offsets_host = counts.tolist(), offsets.tolist() - - base_zero = 0 - assert sum(counts_host[base_zero : base_zero + expert_per_rank]) == 3 - base_one = expert_per_rank - empty_total = sum(counts_host[base_one : base_one + expert_per_rank]) - assert empty_total == 0 - empty_segment = route_table[ - offsets_host[base_one] : offsets_host[base_one] + empty_total - ] - assert empty_segment.shape == (0, 2) - - # Combining an empty segment must be a no-op, not an error. + slices = list(_destination_slices(plan, ffn_size, expert_per_rank)) + + _, busy_partials, busy_rows = slices[0] + assert busy_partials.stop - busy_partials.start == 3 + # One token, three of its partials: it is shipped once, read three times. + assert busy_rows.tolist() == [0] + + _, empty_partials, empty_rows = slices[1] + assert empty_partials.stop - empty_partials.start == 0 + assert empty_rows.numel() == 0 + + # Reducing an empty destination must be a no-op, not an error. accumulator = torch.zeros(1, 4, dtype=torch.float32) - combine_scatter( - accumulator, - routed_out=torch.zeros(0, 4), - route_table=empty_segment, - topk_weights=torch.ones(1, 3), - ) + accumulator.index_add_(0, empty_rows, torch.zeros(0, 4)) assert torch.count_nonzero(accumulator) == 0 From f2c9a08f14a970fb34916af27d89168a2df72d78 Mon Sep 17 00:00:00 2001 From: specture724 Date: Fri, 14 Aug 2026 11:12:31 +0800 Subject: [PATCH 06/18] refactor: fold the two receive spin loops into SymmWindow.wait Both roles spun on poll() with their own copy of the loop. One helper now owns it, and its docstring carries the measurement that says to leave it hot. Profiling 2A2F under load showed the poll costing more GPU copy-engine time than the expert GEMM: ~25k D2H polls a second per rank, 42.8 ms of Memcpy DtoH against 1.8 ms of fused_moe_kernel in a 1 s window on an FFN rank, which sits 95.4% idle. Backing off looked like the fix. It is not: every layer is a serialized A->F->A round trip, so detection latency is paid twice per layer on the critical path, and on the FFN side it also delays picking up the next rank's dispatch. Sleeping between attempts (4 hot tries, then 50us doubling to 1 ms) made mean TTFT worse at every rate -- 346 -> 431 ms at 32 rps, 1869 -> 3906 ms at 64 rps -- so the backoff is not here, only the note saying why. No behaviour change: reverting to the hot spin reproduces the baseline sweep (335/445/1899 ms at 32/48/64 rps against 346/425/1869 before the experiment). Co-Authored-By: Claude Opus 5 Signed-off-by: specture724 --- afd_plugin/connectors/gpu/async_gpu.py | 14 +++------- afd_plugin/connectors/gpu/symm_window.py | 35 +++++++++++++++++++++--- 2 files changed, 35 insertions(+), 14 deletions(-) diff --git a/afd_plugin/connectors/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py index d5b45c29..ec7ebdb6 100644 --- a/afd_plugin/connectors/gpu/async_gpu.py +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -21,7 +21,6 @@ from __future__ import annotations -import time from collections.abc import Mapping from dataclasses import dataclass, field from datetime import timedelta @@ -674,7 +673,7 @@ def recv_ffn_output( ) outstanding = set(pending.expected_ffn) while outstanding: - arrived = window.poll() + arrived = window.wait() if arrived is None: continue header = arrived.header @@ -754,14 +753,9 @@ def recv_attn_output( """ window = self._require_initialized() timeout_ms = int(kwargs.get("timeout_ms", 0)) - deadline = time.monotonic() + timeout_ms / 1000.0 if timeout_ms else None - - while True: - arrived = window.poll() - if arrived is not None: - break - if deadline is not None and time.monotonic() >= deadline: - raise TimeoutError("AFD async GPU dispatch recv timed out") + 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 if header.is_shutdown: diff --git a/afd_plugin/connectors/gpu/symm_window.py b/afd_plugin/connectors/gpu/symm_window.py index 30436c83..f115c200 100644 --- a/afd_plugin/connectors/gpu/symm_window.py +++ b/afd_plugin/connectors/gpu/symm_window.py @@ -37,6 +37,7 @@ from __future__ import annotations +import time from dataclasses import dataclass from typing import TYPE_CHECKING @@ -521,13 +522,39 @@ def write_field( # Receive side. # ------------------------------------------------------------------ + def wait(self, *, timeout_s: float | None = None) -> ArrivedSlot | None: + """Spin until a slot arrives, or ``timeout_s`` passes. + + The spin is hot on purpose. A poll is a D2H plus a stream synchronize, + about 40us, and a profile shows this loop issuing ~25k of them a second + per rank -- more copy-engine time than the expert GEMM itself. Backing + off looks like the obvious fix and is not: every layer is a serialized + A->F->A round trip, so detection latency lands directly on the critical + path twice per layer, and on the FFN side it also delays picking up the + next rank's dispatch. Measured on 2A2F, sleeping between attempts (4 hot + tries, then 50us doubling to 1ms) made mean TTFT worse at every rate: + 346 -> 431 ms at 32 rps, 1869 -> 3906 ms at 64 rps. + + ponytail: hot spin, and it costs real GPU copy-engine time. The fix is + not to poll less often but to stop polling from the host: a device-side + ``wait_any`` kernel spinning on the flag array would notice an arrival + in ~1us and let the host block on one synchronize. The larger win is to + take the round trip off the critical path entirely (ubatching), after + which detection latency stops being paid per layer. + """ + deadline = None if timeout_s is None else time.monotonic() + timeout_s + while True: + arrived = self.poll() + if arrived is not None: + return arrived + if deadline is not None and time.monotonic() >= deadline: + return None + def poll(self) -> ArrivedSlot | None: """Return the first slot whose flag advanced past what we consumed. - ponytail: host-side poll, one D2H per call on the poll stream. Correct - but it burns a synchronize per attempt; replace with a device-side - ``wait_any`` kernel spinning on the flag array when the poll shows up in - a profile. + One D2H per call, on the poll stream. Callers should use ``wait`` + instead of spinning on this. """ with torch.cuda.stream(self._poll_stream): self._flag_host.copy_(self._flags_local, non_blocking=False) From 1f38e071aaa23715c332e22b6d8a87bfef08a7da Mon Sep 17 00:00:00 2001 From: specture724 Date: Fri, 14 Aug 2026 13:10:52 +0800 Subject: [PATCH 07/18] feat: run async MoE ubatching on GPU The two-stage schedule was written for CAM and never ran on CUDA: the switch existed, the model-side schedule was already platform-neutral, but nothing on the GPU side planned the stages. It does now, and four things had to be fixed before a request survived the path. Stage planning reuses plan_async_moe_stages and vLLM's own per-ubatch metadata build, so the split is one extra _build_attention_metadata call over the same inputs. Batches with fewer than two requests, dummy runs and graph capture fall back to the unstaged schedule -- a synthetic batch has no peer waiting for its stages. The four fixes: - vLLM sizes the per-ubatch metadata builder pool from its own ubatching switch, which AFD leaves off, so only one builder existed and building the second stage asserted. initialize_metadata_builders now asks for one per stage, plus one for the full batch: the dense prefix keeps running whole, so its metadata is live at the same time as both stages' and cannot share a builder with them. Sharing one showed up as an illegal memory access inside attention. - Each stage was handed the same ring list, so both stages of a layer claimed the same window slot and the second dispatch overwrote the first's payload and flag. Rings are now partitioned across stages. This was a latent bug in the connector, not in the schedule. - The stage forward context kept the full batch's slot_mapping, so MLA wrote each stage's rows into the whole batch's KV-cache slots. - Restoring the stage outputs splits a concatenation along the last dimension, and CUDA's fused_add_rms_norm aborts on the resulting non-contiguous views. Verified on 4x L20X, DeepSeek-V2-Lite 2A2F: tests/unit 619 passed, and the greedy probe is character-for-character identical to the unstaged schedule on all five prompts. It is slower, so it stays off by default and lives in its own recipe. Mean TTFT against the unstaged schedule: 144 vs 144 ms at 8 rps, 224 vs 185 at 16, 546 vs 346 at 32, 798 vs 426 at 48, no failures at any rate. Halving the batch halves the work per stage but doubles the per-layer fixed cost -- two dispatches, two waits and two host synchronizes instead of one each -- and at these batch sizes that costs more than the overlap returns. Co-Authored-By: Claude Opus 5 Signed-off-by: specture724 --- afd_plugin/connectors/gpu/async_gpu.py | 31 ++- .../npu/deepseek_v2_async_cam_forward.py | 15 +- .../v1/worker/attention_model_runner.py | 232 +++++++++++++++++- .../2a2f_eager_async_ubatch.sh | 136 ++++++++++ 4 files changed, 407 insertions(+), 7 deletions(-) create mode 100644 recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/2a2f_eager_async_ubatch.sh diff --git a/afd_plugin/connectors/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py index ec7ebdb6..2438e763 100644 --- a/afd_plugin/connectors/gpu/async_gpu.py +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -370,6 +370,17 @@ def __init__( self.is_attention = afd_config.role == "attention" self.ring_depth = self.extra_info.ring_depth + self.num_stages = ( + max(1, self.extra_info.async_moe_num_ubatches) + if self.extra_info.async_moe_ubatching + else 1 + ) + 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. @@ -397,6 +408,17 @@ def __init__( 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. @@ -426,8 +448,8 @@ def init_afd_connector(self) -> None: rank=self.world_rank, world_size=self.topology.world_size, ) - for stage in range(max(1, self.extra_info.async_moe_num_ubatches)): - self._free_rings[stage] = list(range(self.ring_depth)) + 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", @@ -532,7 +554,10 @@ def send_attn_output( ) stage_idx = metadata.stage_idx - rings = self._free_rings.setdefault(stage_idx, list(range(self.ring_depth))) + 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 " 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 5db05a54..e2e9844a 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 @@ -338,6 +338,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, @@ -514,7 +524,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/attention_model_runner.py b/afd_plugin/v1/worker/attention_model_runner.py index ff6b0e3b..598ed587 100644 --- a/afd_plugin/v1/worker/attention_model_runner.py +++ b/afd_plugin/v1/worker/attention_model_runner.py @@ -23,6 +23,7 @@ from vllm.v1.worker.gpu_model_runner import GPUModelRunner, PerLayerAttnMetadata from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper from vllm.v1.worker.ubatch_utils import ( + UBatchSlice, UBatchSlices, check_ubatch_thresholds, is_last_ubatch_empty, @@ -40,7 +41,14 @@ AFDDPMetadata, AFDForwardContextMetadata, ) +from afd_plugin.connectors.async_topology import ASYNC_MOE_REQUEST_SPLIT +from afd_plugin.connectors.gpu.async_gpu import GpuAsyncExtraInfo from afd_plugin.model_executor.models.forward_context import use_afd_metadata_provider +from afd_plugin.model_executor.models.npu.async_cam_layout import ( + ASYNC_MOE_UBATCH_METADATA_KEY, + AsyncMoeUbatchMetadata, +) +from afd_plugin.model_executor.npu.async_cam_ubatching import plan_async_moe_stages from afd_plugin.v1.worker.cuda_graph import validate_cuda_graph_mode from afd_plugin.v1.worker.ubatch_wrapper import ( AFDUBatchWrapper, @@ -49,6 +57,11 @@ if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput + from vllm.v1.kv_cache_interface import KVCacheConfig + +# The async MoE schedule interleaves exactly two stages: while the FFN ranks +# hold one stage's tokens, Attention computes the other's. +_ASYNC_MOE_STAGE_COUNT = 2 @contextmanager @@ -86,6 +99,12 @@ class AFDAttentionModelRunner(GPUModelRunner): afd_expected_role = "attention" + # Class-level defaults so a runner built without ``__init__`` -- which is how + # the unit tests exercise the metadata plumbing -- still reads as "no async + # ubatching" instead of raising on a missing attribute. + _afd_async_extra_info: GpuAsyncExtraInfo | None = None + _afd_async_moe_ubatch_metadata: AsyncMoeUbatchMetadata | None = None + def __init__( self, vllm_config: VllmConfig, @@ -114,6 +133,19 @@ def __init__( self._afd_pending_metadata: AFDForwardContextMetadata | None = None self._afd_suppress_metadata_send = False self._afd_transaction_counter = 0 + # Async MoE ubatching is a property of the async connector, so only that + # connector's typed config carries the switch. + extra_info = self.connector.extra_info + self._afd_async_extra_info = ( + extra_info if isinstance(extra_info, GpuAsyncExtraInfo) else None + ) + self._afd_async_moe_ubatch_metadata = None + if self._afd_async_extra_info is not None: + fail_if_unsupported_async_moe_ubatching( + self.vllm_config, + self._afd_async_extra_info, + self.afd_config, + ) self.prof = create_afd_gpu_profiler("attention") @staticmethod @@ -285,6 +317,10 @@ def _install_afd_metadata_on_forward_context( forward_context.additional_kwargs["afd_metadata"] = ( self._afd_pending_metadata ) + if self._afd_async_moe_ubatch_metadata is not None: + forward_context.additional_kwargs[ASYNC_MOE_UBATCH_METADATA_KEY] = ( + self._afd_async_moe_ubatch_metadata + ) if bool(getattr(self, "_afd_suppress_metadata_send", False)): return dp_metadata = forward_context.dp_metadata @@ -294,6 +330,117 @@ def _install_afd_metadata_on_forward_context( dp_metadata = self._build_capture_dp_metadata(padded_graph_tokens) self._send_dp_metadata(dp_metadata, ubatch_slices) + # Upstream source: vLLM v0.26.0, + # vllm/v1/worker/gpu_model_runner.py GPUModelRunner.initialize_metadata_builders + # Patch reason: upstream sizes the per-ubatch metadata builder pool from + # vLLM's own ubatching switch. AFD async MoE ubatching splits the batch + # itself and leaves that switch off, so only one builder exists and building + # the second stage's metadata asserts. + # Patch functionality: ask for one builder per AFD MoE stage instead. + # Signature: matches upstream; no added parameters. + def initialize_metadata_builders( + self, + kv_cache_config: KVCacheConfig, + kernel_block_sizes: list[int], + ) -> None: + for kv_cache_group_id in range(len(kv_cache_config.kv_cache_groups)): + for attn_group in self.attn_groups[kv_cache_group_id]: + attn_group.create_metadata_builders( + self.vllm_config, + self.device, + kernel_block_sizes[kv_cache_group_id] + if kv_cache_group_id < len(kernel_block_sizes) + else None, + # ### PATCH START: one builder per AFD MoE stage + num_metadata_builders=self._num_afd_metadata_builders(), + # ### PATCH END: one builder per AFD MoE stage + ) + # Calculate reorder batch threshold (if needed) + # Note (tdoublep): do this *after* constructing builders, + # because some of them change the threshold at init time. + self.calculate_reorder_batch_threshold() + + # Initialize drafter attention backend + if self.speculative_config and ( + self.speculative_config.use_eagle() + or self.speculative_config.uses_draft_model() + ): + self.drafter.initialize_attn_backend(kv_cache_config, kernel_block_sizes) + + def _num_afd_metadata_builders(self) -> int: + """Builders needed, one per set of metadata that is live at once. + + Async MoE ubatching keeps the full-batch metadata alive for the dense + prefix while both stages' metadata drive the MoE layers, so it needs a + builder for the full batch plus one per stage. Sharing builder 0 between + the full batch and stage 0 lets the second build hand back reusable + runtime buffers sized for the first, which reads as an illegal memory + access inside attention. + """ + if self.parallel_config.use_ubatching: + return int(self.parallel_config.num_ubatches) + extra_info = self._afd_async_extra_info + if extra_info is not None and extra_info.async_moe_ubatching: + return 1 + _ASYNC_MOE_STAGE_COUNT + return 1 + + @contextmanager + def _afd_stage_metadata_builders(self): + """Hide builder 0 so a stage build cannot touch the full batch's.""" + saved: list[tuple[Any, list[AttentionMetadataBuilder]]] = [] + for kv_cache_group in self.attn_groups: + for attention_group in kv_cache_group: + saved.append((attention_group, attention_group.metadata_builders)) + attention_group.metadata_builders = attention_group.metadata_builders[ + 1: + ] + try: + yield + finally: + for attention_group, builders in saved: + attention_group.metadata_builders = builders + + def _plan_async_moe_stages( + self, + *, + num_reqs: int, + ubatch_slices: UBatchSlices | None, + for_cudagraph_capture: bool, + num_scheduled_tokens: dict[str, int] | None, + ) -> tuple[Any, ...] | None: + """Split this batch into two MoE stages, or ``None`` to run it whole. + + Splitting is what gives the FFN ranks something to chew on while + Attention computes: one stage is in flight across the wire while the + other runs locally. It needs at least two requests to split at a request + boundary, so a single-request batch falls back to the unstaged schedule. + + Capture and dummy runs are excluded: the staged schedule dispatches to + peers that are not expecting a synthetic batch. + """ + extra_info = self._afd_async_extra_info + if extra_info is None or not extra_info.async_moe_ubatching: + return None + if for_cudagraph_capture or self._afd_is_graph_capturing: + return None + if ubatch_slices is not None or num_scheduled_tokens is None: + # vLLM is already ubatching this batch; two splitters would fight. + return None + if num_reqs < _ASYNC_MOE_STAGE_COUNT: + return None + # Request order in the batch, which is the order the token dimension is + # laid out in -- the dict is keyed by request id and is not that order. + req_ids = self.input_batch.req_ids[:num_reqs] + scheduled = [int(num_scheduled_tokens[req_id]) for req_id in req_ids] + if any(token_count <= 0 for token_count in scheduled): + return None + return plan_async_moe_stages( + scheduled, + split=extra_info.async_moe_split, + use_sequence_parallel=False, + tensor_parallel_size=self.vllm_config.parallel_config.tensor_parallel_size, + ) + # Patch reason: AFD stages connector metadata before native Attention # metadata construction. In addition, vLLM v0.26 caches Attention metadata # without the ubatch id, so update-capable backends can reuse ubatch 0 @@ -327,17 +474,28 @@ def _build_attention_metadata( ubatch_slices, int(num_tokens), ) + stages = self._plan_async_moe_stages( + num_reqs=num_reqs, + ubatch_slices=ubatch_slices, + for_cudagraph_capture=for_cudagraph_capture, + num_scheduled_tokens=num_scheduled_tokens, + ) + self._afd_async_moe_ubatch_metadata = None + num_metadata_ubatches = max( + len(ubatch_slices) if ubatch_slices is not None else 1, + len(stages) if stages is not None else 1, + ) disabled_metadata_builders: dict[int, AttentionMetadataBuilder] = {} try: - if ubatch_slices is not None and len(ubatch_slices) > 1: + if num_metadata_ubatches > 1: for kv_cache_group in self.attn_groups: for attention_group in kv_cache_group: - for ubatch_idx in range(len(ubatch_slices)): + for ubatch_idx in range(num_metadata_ubatches): builder = attention_group.get_metadata_builder(ubatch_idx) if builder.supports_update_block_table: disabled_metadata_builders[id(builder)] = builder builder.supports_update_block_table = False - return super()._build_attention_metadata( + full_metadata = super()._build_attention_metadata( num_tokens, num_reqs, max_query_len, @@ -351,6 +509,35 @@ def _build_attention_metadata( cascade_attn_prefix_lens, slot_mappings, ) + if stages is not None: + # A second build over the same inputs, sliced per stage, on the + # stages' own builders. The dense prefix still runs on the whole + # batch and needs the full metadata, so both live at once. + with self._afd_stage_metadata_builders(): + stage_metadata, _ = super()._build_attention_metadata( + num_tokens, + num_reqs, + max_query_len, + num_tokens_padded, + num_reqs_padded, + [ + UBatchSlice(stage.request_slice, stage.token_slice) + for stage in stages + ], + logits_indices, + use_spec_decode, + for_cudagraph_capture, + num_scheduled_tokens, + cascade_attn_prefix_lens, + slot_mappings, + ) + self._afd_async_moe_ubatch_metadata = AsyncMoeUbatchMetadata( + attn_metadata=stage_metadata, + stages=stages, + parent_input_tokens=int(num_tokens_padded or num_tokens), + use_sequence_parallel=False, + ) + return full_metadata finally: for builder in disabled_metadata_builders.values(): builder.supports_update_block_table = True @@ -539,6 +726,10 @@ def _dummy_run( """ previous_metadata = self._afd_pending_metadata + previous_stage_metadata = self._afd_async_moe_ubatch_metadata + # A dummy batch is synthetic: nobody on the FFN side is waiting for its + # stages, so it always runs the unstaged schedule. + self._afd_async_moe_ubatch_metadata = None previous_is_graph_capturing = getattr( self, "_afd_is_graph_capturing", @@ -564,6 +755,7 @@ def _dummy_run( finally: self._afd_is_graph_capturing = previous_is_graph_capturing self._afd_pending_metadata = previous_metadata + self._afd_async_moe_ubatch_metadata = previous_stage_metadata # Patch reason: native capture does not publish AFD warmup/capture metadata. # Patch functionality: preserve the upstream warmup/capture flow while @@ -697,6 +889,40 @@ def fail_if_unsupported_ubatching(vllm_config: VllmConfig) -> None: fail_if_ubatching_enabled = fail_if_unsupported_ubatching +def fail_if_unsupported_async_moe_ubatching( + vllm_config: VllmConfig, + extra_info: GpuAsyncExtraInfo, + afd_config: AFDConfig, +) -> None: + """Reject async MoE ubatching settings the GPU schedule cannot honour. + + Mirrors the NPU checks: the schedule interleaves exactly two stages, needs + the Attention-side gate to have the routing payloads to dispatch, and splits + the batch itself, so vLLM's own ubatching must stay off. + """ + if not extra_info.async_moe_ubatching: + return + if extra_info.async_moe_num_ubatches != _ASYNC_MOE_STAGE_COUNT: + raise RuntimeError( + "async_moe_ubatching currently supports exactly two stages; got " + f"async_moe_num_ubatches={extra_info.async_moe_num_ubatches}", + ) + if extra_info.async_moe_split != ASYNC_MOE_REQUEST_SPLIT: + raise RuntimeError( + "async_moe_ubatching only supports request-boundary splits; got " + f"async_moe_split={extra_info.async_moe_split!r}", + ) + if not afd_config.compute_gate_on_attention: + raise RuntimeError( + "async_moe_ubatching requires compute_gate_on_attention=true", + ) + if bool(vllm_config.parallel_config.use_ubatching): + raise RuntimeError( + "async_moe_ubatching splits the batch itself; vLLM's own ubatching " + "must stay disabled", + ) + + def fail_if_cuda_graph_enabled(vllm_config: VllmConfig) -> None: validate_cuda_graph_mode(vllm_config) diff --git a/recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/2a2f_eager_async_ubatch.sh b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/2a2f_eager_async_ubatch.sh new file mode 100644 index 00000000..0720dd93 --- /dev/null +++ b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/2a2f_eager_async_ubatch.sh @@ -0,0 +1,136 @@ +#!/usr/bin/env bash +# 2A2F async GPU connector with async MoE ubatching, eager, prefill-only. +# +# The batch is split at a request boundary into two stages that leapfrog: while +# the FFN ranks hold one stage, Attention computes the other. Without it every +# layer is a serialized A->F->A round trip and both roles spend most of their +# time waiting for each other. +# +# Launch under a GPU reservation, which sets CUDA_VISIBLE_DEVICES: +# gpu run --gpus 4 -- bash recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/2a2f_eager_async_ubatch.sh +# +# The two roles are separate vllm serve processes: the AFD process group hosts +# its own TCPStore, which cannot be created under a single torchrun/torchelastic +# launcher. +set -u + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} +LOG_DIR=${LOG_DIR:-.} +mkdir -p "$LOG_DIR" +export VLLM_USE_V2_MODEL_RUNNER=0 +# Single node over NVLink: skip the IB transport probe. +export NVSHMEM_REMOTE_TRANSPORT=${NVSHMEM_REMOTE_TRANSPORT:-none} +# Two servers on one box spawn a lot of threads; the HF tokenizer's rayon pool +# is the first thing to fail when thread creation gets refused. +export TOKENIZERS_PARALLELISM=${TOKENIZERS_PARALLELISM:-false} +export RAYON_NUM_THREADS=${RAYON_NUM_THREADS:-2} +export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4} + +# Split the reserved devices in half: first two Attention, last two FFN. +IFS=',' read -r -a DEVICES <<< "${CUDA_VISIBLE_DEVICES:-0,1,2,3}" +if [ "${#DEVICES[@]}" -lt 4 ]; then + echo "need 4 visible GPUs, got ${#DEVICES[@]}: ${CUDA_VISIBLE_DEVICES:-unset}" >&2 + exit 1 +fi +ATTN_DEVICES="${DEVICES[0]},${DEVICES[1]}" +FFN_DEVICES="${DEVICES[2]},${DEVICES[3]}" +echo "attention on ${ATTN_DEVICES}, ffn on ${FFN_DEVICES}" + +# Lower this when sharing a box: vLLM refuses to start if the desired +# fraction exceeds what is actually free. +GPU_MEM_UTIL=${GPU_MEM_UTIL:-0.9} +# Prefill batch size drives whether each MoE call clears the compute-bound +# inflection point, so it is the knob to raise when benchmarking. +MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-512} +MAX_NUM_SEQS=${MAX_NUM_SEQS:-8} +AFD_PORT=${AFD_PORT:-6301} +API_PORT=${API_PORT:-18337} + +CUDA_VISIBLE_DEVICES="$ATTN_DEVICES" uv run vllm serve "$MODEL_PATH" \ + --data-parallel-size 2 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --additional-config '{ + "afd": { + "role": "attention", + "connector": "GpuAsyncAFDConnector", + "async": true, + "compute_gate_on_attention": true, + "host": "127.0.0.1", + "port": '"$AFD_PORT"', + "num_attention_ranks": 2, + "num_ffn_ranks": 2, + "connector_extra_config": {"async_moe_ubatching": true} + } + }' \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \ + --api-server-count 1 \ + --gpu-memory-utilization "$GPU_MEM_UTIL" \ + --enforce-eager \ + --host 127.0.0.1 \ + --port "$API_PORT" \ + --trust-remote-code > "$LOG_DIR/attn.log" 2>&1 & +ATTN_PID=$! + +CUDA_VISIBLE_DEVICES="$FFN_DEVICES" uv run vllm serve "$MODEL_PATH" \ + --data-parallel-size 2 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --additional-config '{ + "afd": { + "role": "ffn", + "connector": "GpuAsyncAFDConnector", + "async": true, + "compute_gate_on_attention": true, + "host": "127.0.0.1", + "port": '"$AFD_PORT"', + "num_attention_ranks": 2, + "num_ffn_ranks": 2, + "connector_extra_config": {"async_moe_ubatching": true} + } + }' \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \ + --api-server-count 1 \ + --gpu-memory-utilization "$GPU_MEM_UTIL" \ + --enforce-eager \ + --host 127.0.0.1 \ + --port "$API_PORT" \ + --trust-remote-code > "$LOG_DIR/ffn.log" 2>&1 & +FFN_PID=$! + +cleanup() { + kill "$ATTN_PID" "$FFN_PID" 2>/dev/null + wait "$ATTN_PID" "$FFN_PID" 2>/dev/null +} +trap cleanup EXIT + +for _ in $(seq 1 "${READY_TIMEOUT:-600}"); do + if curl -sf "http://127.0.0.1:$API_PORT/health" > /dev/null 2>&1; then + echo "server ready on http://127.0.0.1:$API_PORT" + echo + echo "curl -s http://127.0.0.1:$API_PORT/v1/completions \\" + echo " -H 'Content-Type: application/json' \\" + echo " -d '{\"model\":\"$MODEL_PATH\",\"prompt\":\"The capital of France is\",\"max_tokens\":16,\"temperature\":0}'" + echo + if [ -n "${SMOKE:-}" ]; then + curl -s "http://127.0.0.1:$API_PORT/v1/completions" \ + -H 'Content-Type: application/json' \ + -d '{"model":"'"$MODEL_PATH"'","prompt":"The capital of France is", + "max_tokens":16,"temperature":0}' + echo + exit 0 + fi + # Stay up so the servers can take requests; Ctrl-C tears both down. + wait "$ATTN_PID" "$FFN_PID" + exit 0 + fi + if ! kill -0 "$ATTN_PID" 2>/dev/null || ! kill -0 "$FFN_PID" 2>/dev/null; then + echo "a server exited early; see $LOG_DIR/attn.log and $LOG_DIR/ffn.log" >&2 + exit 1 + fi + sleep 1 +done +echo "timed out waiting for the server" >&2 +exit 1 From 741db0de68aba97ec4cfee64939a4e3e01301e00 Mon Sep 17 00:00:00 2001 From: specture724 Date: Fri, 14 Aug 2026 16:13:20 +0800 Subject: [PATCH 08/18] perf: weight and reduce the expert output in the payload dtype A profile of a 6A2F FFN rank under load put the scatter machinery at 2.2 s of a 6.6 s busy window against 2.7 s for fused_moe_kernel itself: the reduce from partials back to one row per shipped token widened the expert output to float32, scaled it, and scattered it, which is three passes over [partials, hidden] with the last two moving twice the bytes they need to. The Attention-side combine did the same thing, widening every arriving block and narrowing the result on the way out; its scatter and gather together cost 60 ms against FlashAttention's 24 ms in the same window. Both now stay in the payload dtype, so the reduce is a multiply and a scatter. What float32 was protecting is a sum of at most topk terms on the FFN side and one term per FFN rank plus the shared one on the Attention side, and the result crosses the wire narrowed either way. Separately, the gate scaled its output by routed_scaling_factor in a pass over [num_partials, hidden]. fused_experts multiplies every row by its topk weight in an epilogue it already runs, so the factor goes in there instead and the pass disappears. The topk weighting itself deliberately stays in the connector: an earlier version of this change handed the per-partial weights to the gate and dropped the multiply from the connector, which left send_ffn_output silently depending on its caller having applied them -- the end-to-end test's own FFN loop had not, and the failure mode was a wrong answer rather than an error. Numerics: the equivalence test moves from 4.54e-3 to 6.60e-3 max relative error, about one bf16 ulp and still an order inside its tolerance. Verified on L20X: tests/unit 619 passed, async_gpu_connector_e2e passes over the wire. The end-to-end effect is not measured yet -- every GPU on the box is busy. Co-Authored-By: Claude Opus 5 Signed-off-by: specture724 --- afd_plugin/connectors/gpu/async_gpu.py | 24 ++++++++++++----- .../models/gpu/deepseek_v2_attention_gate.py | 27 ++++++++++--------- 2 files changed, 32 insertions(+), 19 deletions(-) diff --git a/afd_plugin/connectors/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py index 2438e763..1792241b 100644 --- a/afd_plugin/connectors/gpu/async_gpu.py +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -691,9 +691,14 @@ def recv_ffn_output( ) 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=torch.float32, + dtype=self.payload_dtype, device=ref_tensor.device, ) outstanding = set(pending.expected_ffn) @@ -738,7 +743,7 @@ def recv_ffn_output( arrived.region, arrived.ring, header.uniq_tokens, - ).to(accumulator.dtype), + ), ) if header.shared_tokens: accumulator.index_add_( @@ -752,11 +757,11 @@ def recv_ffn_output( arrived.region, arrived.ring, header.shared_tokens, - ).to(accumulator.dtype), + ), ) self._free_rings.setdefault(ubatch_idx, []).append(pending.ring) - return accumulator.to(ref_tensor.dtype) + return accumulator # ================================================================== # FFN-side data path @@ -871,6 +876,12 @@ def send_ffn_output( here, so the reply carries the same rows the dispatch did. Doing it on this side keeps the duplicates off the wire and leaves the Attention side a plain scatter-add. + + 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 @@ -884,12 +895,11 @@ def send_ffn_output( ) reduced = torch.zeros( (states.uniq_tokens, self.hidden_size), - dtype=torch.float32, + dtype=self.payload_dtype, device=ffn_output.device, ) if states.routed_tokens: - weighted = ffn_output.to(torch.float32) - weighted.mul_(states.weights.unsqueeze(1)) + 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") 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 index abe932e2..a26321d9 100644 --- a/afd_plugin/model_executor/models/gpu/deepseek_v2_attention_gate.py +++ b/afd_plugin/model_executor/models/gpu/deepseek_v2_attention_gate.py @@ -8,9 +8,8 @@ only needs the grouped GEMM over its local experts. That shape is a ``topk == 1`` problem: give every arriving row its own expert id -and a unit weight, and vLLM's ``fused_experts`` computes exactly the local -expert output. Topk weighting stays on the Attention side, applied during -combine, matching the NPU path. +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 @@ -82,9 +81,18 @@ def compute_attention_gate_moe_ffn( counts, output_size=num_rows, ).unsqueeze(1) - # Unit weights: the real topk weighting happens in the connector's combine. - unit_weights = torch.ones( + # 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, ) @@ -93,7 +101,7 @@ def compute_attention_gate_moe_ffn( hidden_states, routed_experts.w13_weight, routed_experts.w2_weight, - unit_weights, + row_weights, expert_ids, global_num_experts=num_local_experts, expert_map=None, @@ -108,12 +116,7 @@ def compute_attention_gate_moe_ffn( # tokens as their own batch, so that machinery does not apply. shared_output = shared_experts._layer(expand_x_shared) - # Mirrors the NPU gate path: scale the routed branch unless fp16, where the - # native model instead scales the shared branch down. - routed_scaling_factor = runner.routed_scaling_factor - if hidden_states.dtype != torch.float16: - routed_output = routed_output * routed_scaling_factor - elif shared_output is not None: + if scale_shared_instead and shared_output is not None: shared_output = shared_output * (1.0 / routed_scaling_factor) return AFDF2ATransferPayload( From a3421e41075fd282a81beecfa73ac417ea28b328 Mon Sep 17 00:00:00 2001 From: specture724 Date: Fri, 14 Aug 2026 16:44:05 +0800 Subject: [PATCH 09/18] perf: wait for the expert reply on the stream, not on the host Combine polled for its replies: copy the flag word down, look at it, try again. That cost a synchronize per attempt, but the real damage was that it left the GPU with nothing queued behind the wait. A profile of a 2A2F attention rank found the host running 2.3us ahead of the GPU at the median, with 76.9% of kernel launches executing within 5us of being issued -- the device was idle waiting to be fed one kernel at a time, because every layer blocked the host twice. cuStreamWaitValue32 moves the wait onto the stream, so the host enqueues "block until this word reaches N" and goes straight on to the next layer. For that to work combine has to know, before anything arrives, which slot each reply lands in, how many rows it carries, and what value its flag will take: - the slot is the region the FFN rank owns, on the ring the dispatch used; - the row counts are the ones the dispatch shipped, so they are recorded in the pending record instead of read back out of the reply header; - a reply now stamps the flag with the dispatch sequence it answers rather than the sender's own counter, which is a number the waiting rank already holds. The comparison is GEQ, and dispatch sequence numbers only ever increase, so a stale reply in the slot cannot satisfy it. The per-arrival header checks go with the poll, and the stream wait subsumes them: it blocks on one specific slot reaching one specific sequence, where the poll took whatever had landed and had to check afterwards that it was the thing it wanted. The shutdown branch went too; announce_shutdown has no callers. Measured on 4x L20X, DeepSeek-V2-Lite 2A2F throughout. Mean TTFT against the previous commit: 144 -> 129 ms at 8 rps, 185 -> 163 at 16, 346 -> 261 at 32, 426 -> 347 at 48, 1869 -> 951 at 64, with achieved throughput at offered 64 going 55.6 -> 56.6 rps. Against the synchronous connector async now wins by 21% at 8 rps and 25% at 16, and the gap at 64 rps closes from 4.8x to 2.5x. On the same attention rank at rate 32, queue lag moves from a 13.2us p90 to 660us and launches into an empty queue from 76.9% to 59.5% -- what is left is the counts readback that still blocks the host once per dispatch. async_gpu_connector_e2e passes over the wire, tests/unit 619 passed, and the greedy probe still answers sensibly. Co-Authored-By: Claude Opus 5 --- afd_plugin/connectors/gpu/async_gpu.py | 106 +++++++++---------- afd_plugin/connectors/gpu/cuda_rt.py | 125 +++++++++++++++++++++++ afd_plugin/connectors/gpu/symm_window.py | 27 ++++- tests/e2e/async_gpu_connector_e2e.py | 2 + 4 files changed, 206 insertions(+), 54 deletions(-) create mode 100644 afd_plugin/connectors/gpu/cuda_rt.py diff --git a/afd_plugin/connectors/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py index 1792241b..889cb865 100644 --- a/afd_plugin/connectors/gpu/async_gpu.py +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -223,10 +223,16 @@ class _PendingDispatch: ``r``, in the order they were shipped. The reply comes back one row per shipped token, already weighted and summed, so combine is a scatter-add at exactly those ids. + + ``shared_ids[r]`` is the same thing for the shared-expert rows. Both are + 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 uniq_ids: list[Tensor] + shared_ids: list[Tensor] num_tokens: int ring: int seq: int @@ -354,6 +360,9 @@ def __init__( 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 @@ -597,6 +606,7 @@ def send_attn_output( # round-robin shared slice can come out empty for one rank. expected_ffn = list(range(self.ffn_size)) uniq_ids: list[Tensor] = [] + shared_ids: list[Tensor] = [] uniq_start = 0 for ffn_rank in range(self.ffn_size): base = ffn_rank * self.expert_per_rank @@ -620,6 +630,8 @@ def send_attn_output( device=hidden_states.device, dtype=torch.int32, ) + shared_rows = shared_idx.to(torch.int64) + shared_ids.append(shared_rows) header = encode_header( self.layout, seq=self._seq, @@ -648,7 +660,7 @@ def send_attn_output( routed_rows=rows, shared_idx=shared_idx, shared_x=hidden_states, - shared_rows=shared_idx.to(torch.int64), + shared_rows=shared_rows, ) logger.debug( @@ -667,6 +679,7 @@ def send_attn_output( _PendingDispatch( context=context, uniq_ids=uniq_ids, + shared_ids=shared_ids, num_tokens=num_tokens, ring=ring, seq=self._seq, @@ -680,7 +693,21 @@ def recv_ffn_output( ubatch_idx: int = 0, **kwargs: Any, ) -> Tensor: - """Wait for this layer's expert output and reduce it back to ``[B, H]``.""" + """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 + reaching one specific sequence number, 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) @@ -701,64 +728,35 @@ def recv_ffn_output( dtype=self.payload_dtype, device=ref_tensor.device, ) - outstanding = set(pending.expected_ffn) - while outstanding: - arrived = window.wait() - if arrived is None: - continue - header = arrived.header - if header.is_shutdown: - raise ConnectorShutdown( - f"FFN rank {header.src_role_rank} announced shutdown", - ) - if header.src_role_rank not in outstanding: - raise RuntimeError( - "AFD async GPU combine received an unexpected FFN rank " - f"{header.src_role_rank}; expected one of {sorted(outstanding)}", - ) - if header.echo_seq != pending.seq: - raise RuntimeError( - "AFD async GPU combine answered dispatch seq " - f"{header.echo_seq} while waiting on {pending.seq} " - f"(F{header.src_role_rank}, layer {header.layer_idx}); the " - "pending FIFO and the wire have diverged", - ) - outstanding.discard(header.src_role_rank) - logger.debug( - "AFD combine recv: A%d <- F%d layer=%d routed=%d still_waiting=%s", - self.role_rank, - header.src_role_rank, - header.layer_idx, - header.routed_tokens, - sorted(outstanding), - ) - - if header.uniq_tokens: + for ffn_rank in pending.expected_ffn: + # An FFN rank replies into the region it owns, on the ring the + # dispatch used. + window.stream_wait(ffn_rank, pending.ring, pending.seq) + uniq_ids = pending.uniq_ids[ffn_rank] + if uniq_ids.numel(): # One row per token this rank was sent, already weighted and # summed over that token's partials on the FFN side. accumulator.index_add_( 0, - pending.uniq_ids[header.src_role_rank], - window.local_routed( - arrived.region, - arrived.ring, - header.uniq_tokens, - ), + uniq_ids, + window.local_routed(ffn_rank, pending.ring, uniq_ids.numel()), ) - if header.shared_tokens: + shared_ids = pending.shared_ids[ffn_rank] + if self.has_shared_experts and shared_ids.numel(): accumulator.index_add_( 0, - window.local_shared_idx( - arrived.region, - arrived.ring, - header.shared_tokens, - ).to(torch.int64), - window.local_shared( - arrived.region, - arrived.ring, - header.shared_tokens, - ), + shared_ids, + window.local_shared(ffn_rank, pending.ring, shared_ids.numel()), ) + logger.debug( + "AFD combine queued: A%d layer=%d stage=%d ring=%d seq=%d from=%s", + self.role_rank, + pending.context.metadata.layer_idx, + ubatch_idx, + pending.ring, + pending.seq, + pending.expected_ffn, + ) self._free_rings.setdefault(ubatch_idx, []).append(pending.ring) return accumulator @@ -932,6 +930,10 @@ def send_ffn_output( routed_x=reduced, shared_idx=states.shared_idx if shared_output is not None else None, shared_x=shared_output, + # Stamp the dispatch this answers, not our own counter: the + # Attention rank knows that number already and can hand it to a + # stream wait before the reply exists. + flag_value=states.seq, ) # ================================================================== diff --git a/afd_plugin/connectors/gpu/cuda_rt.py b/afd_plugin/connectors/gpu/cuda_rt.py new file mode 100644 index 00000000..0c3ab240 --- /dev/null +++ b/afd_plugin/connectors/gpu/cuda_rt.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Stream memory operations, for waiting on a peer's flag without the host. + +A receive is a wait on one 32-bit word that a peer writes. Doing that wait on +the host -- copy the flag down, look at it, try again -- costs a synchronize per +attempt and, far worse, stops the host from queueing anything behind the wait. +A profile of the async connector showed 86% of kernel launches starting within +5us of being issued: the GPU was idle waiting to be fed one kernel at a time, +because every layer blocked the host twice. + +``cuStreamWaitValue32`` moves the wait onto the stream. The host enqueues "wait +until this word reaches N" and keeps going, so the work behind the wait is +already queued when the flag arrives. It is a driver API with no runtime API +equivalent and no PyTorch binding, hence ctypes -- the same approach +``nvshmem_rt`` already takes for the symmetric allocator. +""" + +from __future__ import annotations + +import ctypes +from typing import Final + +# CUstreamWaitValue_flags. GEQ is a cyclic comparison, so a monotonically +# increasing sequence number keeps working across 32-bit wraparound. +CU_STREAM_WAIT_VALUE_GEQ: Final[int] = 0x0 +CU_STREAM_WAIT_VALUE_EQ: Final[int] = 0x1 +# CUdevice_attribute: stream memory ops must be supported by the device. +_CU_DEVICE_ATTRIBUTE_CAN_USE_STREAM_MEM_OPS: Final[int] = 74 + +_lib: ctypes.CDLL | None = None +_wait_value32 = None +_checked_devices: set[int] = set() + + +def _load() -> ctypes.CDLL: + global _lib, _wait_value32 + if _lib is not None: + return _lib + + lib = ctypes.CDLL("libcuda.so.1") + # CUDA 11.7 renamed the entry point; the unsuffixed symbol still exists on + # some builds, so take whichever this driver exports. + for symbol in ("cuStreamWaitValue32_v2", "cuStreamWaitValue32"): + fn = getattr(lib, symbol, None) + if fn is not None: + fn.argtypes = [ + ctypes.c_void_p, # CUstream + ctypes.c_ulonglong, # CUdeviceptr + ctypes.c_uint32, # value + ctypes.c_uint32, # flags + ] + fn.restype = ctypes.c_int + _wait_value32 = fn + break + else: + raise RuntimeError( + "libcuda.so.1 exports no cuStreamWaitValue32; this driver cannot " + "wait on a flag from a stream", + ) + lib.cuDeviceGetAttribute.argtypes = [ + ctypes.POINTER(ctypes.c_int), + ctypes.c_int, + ctypes.c_int, + ] + lib.cuDeviceGetAttribute.restype = ctypes.c_int + _lib = lib + return lib + + +def require_stream_mem_ops(device_index: int) -> None: + """Fail loudly at setup if the device cannot wait on memory from a stream.""" + if device_index in _checked_devices: + return + lib = _load() + supported = ctypes.c_int(0) + status = lib.cuDeviceGetAttribute( + ctypes.byref(supported), + _CU_DEVICE_ATTRIBUTE_CAN_USE_STREAM_MEM_OPS, + device_index, + ) + if status != 0: + raise RuntimeError( + f"cuDeviceGetAttribute failed with status {status} while checking " + "stream memory op support", + ) + if not supported.value: + raise RuntimeError( + f"CUDA device {device_index} does not support stream memory " + "operations, which the async GPU connector needs to wait on a " + "peer's flag without blocking the host", + ) + _checked_devices.add(device_index) + + +def stream_wait_value32( + stream: int, + device_ptr: int, + value: int, + *, + flags: int = CU_STREAM_WAIT_VALUE_GEQ, +) -> None: + """Enqueue "block this stream until ``*device_ptr`` reaches ``value``".""" + if _wait_value32 is None: + _load() + assert _wait_value32 is not None + status = _wait_value32( + ctypes.c_void_p(stream), + ctypes.c_ulonglong(device_ptr), + ctypes.c_uint32(value & 0xFFFFFFFF), + ctypes.c_uint32(flags), + ) + if status != 0: + raise RuntimeError( + f"cuStreamWaitValue32 failed with status {status} " + f"(ptr={device_ptr:#x}, value={value})", + ) + + +__all__ = [ + "CU_STREAM_WAIT_VALUE_EQ", + "CU_STREAM_WAIT_VALUE_GEQ", + "require_stream_mem_ops", + "stream_wait_value32", +] diff --git a/afd_plugin/connectors/gpu/symm_window.py b/afd_plugin/connectors/gpu/symm_window.py index f115c200..df30d858 100644 --- a/afd_plugin/connectors/gpu/symm_window.py +++ b/afd_plugin/connectors/gpu/symm_window.py @@ -43,7 +43,7 @@ import torch -from afd_plugin.connectors.gpu import nvshmem_rt +from afd_plugin.connectors.gpu import cuda_rt, nvshmem_rt if TYPE_CHECKING: from torch.distributed.distributed_c10d import ProcessGroup @@ -415,6 +415,7 @@ def write_slot( shared_x: torch.Tensor | None, routed_rows: torch.Tensor | None = None, shared_rows: torch.Tensor | None = None, + flag_value: int | None = None, ) -> None: """Write one slot into ``peer``'s window, then stamp its flag. @@ -425,6 +426,12 @@ def write_slot( ``shared_x``: the gather then lands directly in the peer's window. Materializing the gathered rows locally first would write and re-read every payload byte for nothing. + + ``flag_value`` overrides what the flag is stamped with, which defaults + to this message's own sequence number. A reply stamps the sequence it + answers instead, so the rank waiting for it knows the value to wait for + before it arrives -- that is what lets the wait happen on a stream + rather than on the host. """ layout = self.layout routed_count = _row_count(routed_x, routed_rows) @@ -514,7 +521,7 @@ def write_field( shared_rows, ) - seq = int(header[_H_SEQ].item()) + seq = int(header[_H_SEQ].item()) if flag_value is None else flag_value flag_idx = region * self.ring_depth + ring self._flag_view(peer, flag_idx).fill_(seq) @@ -550,6 +557,22 @@ def wait(self, *, timeout_s: float | None = None) -> ArrivedSlot | None: if deadline is not None and time.monotonic() >= deadline: return None + def stream_wait(self, region: int, ring: int, value: int) -> None: + """Block the current stream until this slot's flag reaches ``value``. + + The host returns immediately, so whatever it queues next -- the combine + that consumes this slot, the next layer -- is already on the stream when + the peer's write lands. Nothing here tells the host that the data + arrived, so the caller must already know the shapes it is going to read. + """ + cuda_rt.require_stream_mem_ops(self.device.index) + flag_idx = region * self.ring_depth + ring + cuda_rt.stream_wait_value32( + torch.cuda.current_stream(self.device).cuda_stream, + self._base + flag_idx * 4, + value, + ) + def poll(self) -> ArrivedSlot | None: """Return the first slot whose flag advanced past what we consumed. diff --git a/tests/e2e/async_gpu_connector_e2e.py b/tests/e2e/async_gpu_connector_e2e.py index 3b9cb2be..85b23c46 100644 --- a/tests/e2e/async_gpu_connector_e2e.py +++ b/tests/e2e/async_gpu_connector_e2e.py @@ -50,6 +50,8 @@ def build_connector(role: str, local_rank: int) -> GpuAsyncAFDConnector: 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, ), From 32aeacea070e51a2f35630e6709695a03772be40 Mon Sep 17 00:00:00 2001 From: specture724 Date: Fri, 14 Aug 2026 17:17:54 +0800 Subject: [PATCH 10/18] perf: resolve the Ascend DBO yield once instead of per MoE layer The yield hook imported afd_plugin.v1.worker.npu.ubatching inside the custom op's body, which runs once per MoE layer. On a CUDA build that module pulls in torch_npu and cannot import, and Python does not cache a failed import: every call re-walked the import machinery. A profile of a 2A2F Attention rank put vllm::manual_dbo_yield at 833us per call at the median with nothing nested inside it -- 258 ms of a 1403 ms window, the largest single host cost on the layer path, all of it spent failing to import a module that can never be there. The import is now resolved once at module import, so the op does what it says: check whether DBO is on, and return. Measured on 4x L20X, DeepSeek-V2-Lite 2A2F, mean TTFT against the previous commit: 129 -> 114 ms at 8 rps, 163 -> 147 at 16, 261 -> 228 at 32, 951 -> 826 at 64. Against the synchronous connector, async is now ahead at 8 rps (114 vs 143) and 16 rps (147 vs 155). The 48 rps point moved the wrong way in this run, 347 -> 624 ms with a p50 of 383 against a mean of 624, which reads as one slow stretch inside a 25 s window rather than a change in the steady state; the box is shared and that point wants re-running. The unit tests patched sys.modules to exercise the per-call import, so they now patch the resolved names instead. Same behaviour under test. Co-Authored-By: Claude Opus 5 Signed-off-by: specture724 --- afd_plugin/v1/worker/dbo.py | 36 +++++++++++++++----------- tests/unit/v1/worker/test_dbo.py | 44 ++++++++------------------------ 2 files changed, 31 insertions(+), 49 deletions(-) diff --git a/afd_plugin/v1/worker/dbo.py b/afd_plugin/v1/worker/dbo.py index d2b9bf71..19bde5e6 100644 --- a/afd_plugin/v1/worker/dbo.py +++ b/afd_plugin/v1/worker/dbo.py @@ -6,6 +6,23 @@ from vllm.utils.torch_utils import direct_register_custom_op from vllm.v1.worker.ubatching import 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. +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 @@ -50,23 +67,12 @@ 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(): diff --git a/tests/unit/v1/worker/test_dbo.py b/tests/unit/v1/worker/test_dbo.py index 6c6b972f..9b7c263f 100644 --- a/tests/unit/v1/worker/test_dbo.py +++ b/tests/unit/v1/worker/test_dbo.py @@ -1,8 +1,6 @@ from __future__ import annotations import builtins -import sys -from types import SimpleNamespace import pytest @@ -84,22 +82,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,22 +98,8 @@ 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")) From cbb8a24ef7a11867c736711a8892ff8446467c5d Mon Sep 17 00:00:00 2001 From: specture724 Date: Fri, 14 Aug 2026 17:39:04 +0800 Subject: [PATCH 11/18] perf: cut the per-layer host work on the dispatch path With the poll gone, the host stopped being blocked and started being busy: a profile of an Attention rank attributed only 7.8 ms of a 1403 ms window to synchronizes, while roughly half the window was plain Python between recorded ops. Three things on the dispatch path, all of them per peer per MoE layer: - encode_header assigned fourteen fields into a fresh tensor one setitem at a time. It builds a numpy array and wraps it instead, which shares the buffer. - the caller summed each rank's expert counts and prefix-summed all of them in Python to find segment bounds, a 64-iteration loop per layer for the counts this model has. plan_dispatch returns the per-rank totals and starts, which ride back in the readback that already happens. - shared-expert tokens were split round robin, so each peer needed an arange, a gather through it, and a whole slot field to carry the index. A contiguous split gives the same balance -- the shared expert treats every token alike -- and makes a rank's slice a view: no index built, none gathered through, none sent. The slot field is gone with it, and combine adds a slice instead of scattering. cudaMemcpyAsync was costing 25us of host time per call, 4000 calls in that window, so dropping one of the five copies per peer per layer is worth more than its bytes suggest. Measured on 4x L20X, DeepSeek-V2-Lite 2A2F, SECS=25, mean TTFT against the previous commit: 116 -> 106 ms at 8 rps, 146 -> 129 at 16, 245 -> 205 at 32, 355 -> 308 at 48, 870 -> 727 at 64. Against the synchronous connector async is now ahead by 25% at 8 rps and 17% at 16, within 11% at 32, and achieved throughput at offered 64 is 58.4 rps against 59.8. The shared split is a wire change but not a numerical one: a token's shared output is the same function wherever it runs. async_gpu_window_roundtrip and async_gpu_connector_e2e pass, tests/unit 619 passed. Co-Authored-By: Claude Opus 5 Signed-off-by: specture724 --- afd_plugin/connectors/gpu/async_gpu.py | 95 +++++++++++-------- afd_plugin/connectors/gpu/symm_window.py | 74 ++++++--------- .../connectors/test_async_gpu_connector.py | 7 +- 3 files changed, 84 insertions(+), 92 deletions(-) diff --git a/afd_plugin/connectors/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py index 889cb865..8c21c837 100644 --- a/afd_plugin/connectors/gpu/async_gpu.py +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -197,7 +197,6 @@ class GpuAsyncTransferState(AFDTransferState): expert_counts_host: list[int] = field(default_factory=list) expand_idx: Tensor | None = None weights: Tensor | None = None - shared_idx: Tensor | None = None expand_x_shared: Tensor | None = None @@ -224,15 +223,16 @@ class _PendingDispatch: shipped token, already weighted and summed, so combine is a scatter-add at exactly those ids. - ``shared_ids[r]`` is the same thing for the shared-expert rows. Both are - 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. + ``shared_slices[r]`` is the contiguous token range whose shared-expert + output that rank returns. Both are 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 uniq_ids: list[Tensor] - shared_ids: list[Tensor] + shared_slices: list[slice] num_tokens: int ring: int seq: int @@ -254,6 +254,10 @@ class DispatchPlan: ``ffn_size * expert_per_rank``. offsets: Exclusive prefix sum of ``counts``. uniq_per_rank: Distinct tokens each FFN rank receives. + routed_per_rank: Partials each FFN rank receives. + segment_start: Where each FFN rank's partials begin. Summing counts and + prefix-summing them on the host cost a Python loop per layer; they + come back in the same readback instead. expand_idx: Per partial, which of its destination's shipped rows it reads. Already rank-local, so a slice needs no rebasing. uniq_token_ids: Per shipped row, the token it carries. @@ -263,6 +267,8 @@ class DispatchPlan: counts: Tensor offsets: Tensor uniq_per_rank: Tensor + routed_per_rank: Tensor + segment_start: Tensor expand_idx: Tensor uniq_token_ids: Tensor weights: Tensor @@ -324,10 +330,13 @@ def plan_dispatch( expand_idx.scatter_(0, key_order, row_of_partial.to(torch.int32)) expand_idx -= row_base[dest_rank] + by_rank = counts.view(ffn_size, expert_per_rank) return DispatchPlan( counts=counts, offsets=offsets, uniq_per_rank=uniq_per_rank, + routed_per_rank=by_rank.sum(1).to(torch.int32), + segment_start=offsets.view(ffn_size, expert_per_rank)[:, 0].to(torch.int32), expand_idx=expand_idx, uniq_token_ids=uniq_token_ids, weights=weights, @@ -590,13 +599,22 @@ def send_attn_output( # full-capacity segments instead. num_experts = self.ffn_size * self.expert_per_rank plan_host = ( - torch.cat((plan.counts.to(torch.int32), plan.uniq_per_rank)).cpu().tolist() + torch.cat( + ( + plan.counts.to(torch.int32), + plan.uniq_per_rank, + plan.routed_per_rank, + plan.segment_start, + ), + ) + .cpu() + .tolist() ) counts_host = plan_host[:num_experts] - uniq_host = plan_host[num_experts:] - offsets_host = [0] * num_experts - for i in range(1, num_experts): - offsets_host[i] = offsets_host[i - 1] + counts_host[i - 1] + rank_fields = plan_host[num_experts:] + uniq_host = rank_fields[: self.ffn_size] + routed_host = rank_fields[self.ffn_size : 2 * self.ffn_size] + start_host = rank_fields[2 * self.ffn_size :] # 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. @@ -606,14 +624,16 @@ def send_attn_output( # round-robin shared slice can come out empty for one rank. expected_ffn = list(range(self.ffn_size)) uniq_ids: list[Tensor] = [] - shared_ids: list[Tensor] = [] + shared_slices: list[slice] = [] uniq_start = 0 for ffn_rank in range(self.ffn_size): base = ffn_rank * self.expert_per_rank expert_counts = counts_host[base : base + self.expert_per_rank] - start = offsets_host[base] - routed_tokens = sum(expert_counts) - segment = slice(start, start + routed_tokens) + routed_tokens = routed_host[ffn_rank] + segment = slice( + start_host[ffn_rank], + start_host[ffn_rank] + routed_tokens, + ) # The shipped rows this rank owns, and the tokens they carry. The # ids stay here rather than going on the wire: combine scatters the # reply back to exactly these rows. @@ -622,16 +642,14 @@ def send_attn_output( uniq_start += uniq_tokens uniq_ids.append(rows) - # Shared-expert tokens are split round-robin across FFN ranks. - shared_idx = torch.arange( - ffn_rank, - num_tokens, - self.ffn_size, - device=hidden_states.device, - dtype=torch.int32, - ) - shared_rows = shared_idx.to(torch.int64) - shared_ids.append(shared_rows) + # Shared-expert tokens are split into 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. Round-robin needed an + # arange, a gather and a whole slot field per peer per layer to + # achieve the same balance. + shared_start = ffn_rank * num_tokens // self.ffn_size + shared_stop = (ffn_rank + 1) * num_tokens // self.ffn_size + shared_slices.append(slice(shared_start, shared_stop)) header = encode_header( self.layout, seq=self._seq, @@ -640,7 +658,7 @@ def send_attn_output( stage_idx=stage_idx, num_tokens=num_tokens, routed_tokens=routed_tokens, - shared_tokens=int(shared_idx.numel()), + shared_tokens=shared_stop - shared_start, topk=self.topk, flags=0, expert_counts=expert_counts, @@ -658,9 +676,7 @@ def send_attn_output( weights=plan.weights[segment], routed_x=hidden_states, routed_rows=rows, - shared_idx=shared_idx, - shared_x=hidden_states, - shared_rows=shared_rows, + shared_x=hidden_states[shared_start:shared_stop], ) logger.debug( @@ -679,7 +695,7 @@ def send_attn_output( _PendingDispatch( context=context, uniq_ids=uniq_ids, - shared_ids=shared_ids, + shared_slices=shared_slices, num_tokens=num_tokens, ring=ring, seq=self._seq, @@ -741,12 +757,13 @@ def recv_ffn_output( uniq_ids, window.local_routed(ffn_rank, pending.ring, uniq_ids.numel()), ) - shared_ids = pending.shared_ids[ffn_rank] - if self.has_shared_experts and shared_ids.numel(): - accumulator.index_add_( - 0, - shared_ids, - window.local_shared(ffn_rank, pending.ring, shared_ids.numel()), + 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, ) logger.debug( "AFD combine queued: A%d layer=%d stage=%d ring=%d seq=%d from=%s", @@ -825,11 +842,6 @@ def recv_attn_output( arrived.ring, header.routed_tokens, ), - shared_idx=window.local_shared_idx( - arrived.region, - arrived.ring, - header.shared_tokens, - ), expand_x_shared=window.local_shared( arrived.region, arrived.ring, @@ -928,7 +940,6 @@ def send_ffn_output( expand_idx=None, weights=None, routed_x=reduced, - shared_idx=states.shared_idx if shared_output is not None else None, shared_x=shared_output, # Stamp the dispatch this answers, not our own counter: the # Attention rank knows that number already and can hand it to a diff --git a/afd_plugin/connectors/gpu/symm_window.py b/afd_plugin/connectors/gpu/symm_window.py index df30d858..139a1fb8 100644 --- a/afd_plugin/connectors/gpu/symm_window.py +++ b/afd_plugin/connectors/gpu/symm_window.py @@ -14,7 +14,9 @@ ``slot(region, ring)`` holds a fixed header followed by the routed/shared payloads. Dispatch (A -> F) and combine (F -> A) use the same slot layout, so a -single spec sizes both directions. +single spec sizes both directions. Shared-expert tokens are a contiguous range +of the batch, so the slot carries only their count -- the range is implied by +which FFN rank the slot belongs to. A token is shipped once per destination rank however many of its topk slots landed there, so the payload rows are *distinct tokens* (``uniq_tokens``, at @@ -41,6 +43,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING +import numpy as np import torch from afd_plugin.connectors.gpu import cuda_rt, nvshmem_rt @@ -92,7 +95,6 @@ class SlotLayout: expand_idx_off: int weights_off: int routed_x_off: int - shared_idx_off: int shared_x_off: int slot_bytes: int @@ -111,10 +113,7 @@ def build( expand_idx_off = _align(header_off + header_words * 4) weights_off = _align(expand_idx_off + partial_cap * 4) routed_x_off = _align(weights_off + partial_cap * 4) - shared_idx_off = _align( - routed_x_off + token_cap * hidden_size * payload_itemsize - ) - shared_x_off = _align(shared_idx_off + token_cap * 4) + shared_x_off = _align(routed_x_off + token_cap * hidden_size * payload_itemsize) slot_bytes = _align(shared_x_off + token_cap * hidden_size * payload_itemsize) return cls( header_words=header_words, @@ -126,7 +125,6 @@ def build( expand_idx_off=expand_idx_off, weights_off=weights_off, routed_x_off=routed_x_off, - shared_idx_off=shared_idx_off, shared_x_off=shared_x_off, slot_bytes=slot_bytes, ) @@ -155,23 +153,29 @@ def encode_header( f"expert_counts must have {expert_per_rank} entries, " f"got {len(expert_counts)}", ) - header = torch.zeros(layout.header_words, dtype=torch.int32) - header[_H_MAGIC] = HEADER_MAGIC - header[_H_VERSION] = HEADER_VERSION - header[_H_SEQ] = seq - header[_H_SRC_ROLE_RANK] = src_role_rank - header[_H_LAYER_IDX] = layer_idx - header[_H_STAGE_IDX] = stage_idx - header[_H_NUM_TOKENS] = num_tokens - header[_H_ROUTED_TOKENS] = routed_tokens - header[_H_SHARED_TOKENS] = shared_tokens - header[_H_TOPK] = topk - header[_H_FLAGS] = flags - header[_H_ECHO_SEQ] = echo_seq - header[_H_UNIQ_TOKENS] = uniq_tokens + # Built with numpy and wrapped, not assigned field by field into a tensor: + # this runs once per peer per MoE layer, and fourteen tensor setitems there + # cost more host time than the transfer they describe. torch.from_numpy + # shares the buffer, so the wrap is free. + header = np.empty(layout.header_words, dtype=np.int32) + header[:HEADER_FIXED_WORDS] = ( + HEADER_MAGIC, + HEADER_VERSION, + seq, + src_role_rank, + layer_idx, + stage_idx, + num_tokens, + routed_tokens, + shared_tokens, + topk, + flags, + echo_seq, + uniq_tokens, + ) if expert_per_rank: - header[HEADER_FIXED_WORDS:] = torch.tensor(expert_counts, dtype=torch.int32) - return header + header[HEADER_FIXED_WORDS:] = expert_counts + return torch.from_numpy(header) @dataclass(frozen=True, slots=True) @@ -351,11 +355,7 @@ def _capacity_view( sizes, dtype = (layout.partial_cap,), torch.int32 elif field_off == layout.weights_off: sizes, dtype = (layout.partial_cap,), torch.float32 - elif field_off == layout.routed_x_off: - sizes, dtype = (layout.token_cap, hidden), self.payload_dtype - elif field_off == layout.shared_idx_off: - sizes, dtype = (layout.token_cap,), torch.int32 - elif field_off == layout.shared_x_off: + elif field_off in (layout.routed_x_off, layout.shared_x_off): sizes, dtype = (layout.token_cap, hidden), self.payload_dtype else: raise ValueError(f"unknown slot field offset {field_off}") @@ -411,7 +411,6 @@ def write_slot( expand_idx: torch.Tensor | None, weights: torch.Tensor | None, routed_x: torch.Tensor | None, - shared_idx: torch.Tensor | None, shared_x: torch.Tensor | None, routed_rows: torch.Tensor | None = None, shared_rows: torch.Tensor | None = None, @@ -505,13 +504,6 @@ def write_field( routed_count, routed_rows, ) - write_field( - layout.shared_idx_off, - (), - torch.int32, - shared_idx, - _row_count(shared_idx, None), - ) write_field( layout.shared_x_off, (layout.hidden_size,), @@ -649,16 +641,6 @@ def local_routed(self, region: int, ring: int, count: int) -> torch.Tensor: self.payload_dtype, ) - def local_shared_idx(self, region: int, ring: int, count: int) -> torch.Tensor: - return self._view( - self.rank, - region, - ring, - self.layout.shared_idx_off, - (count,), - torch.int32, - ) - def local_shared(self, region: int, ring: int, count: int) -> torch.Tensor: return self._view( self.rank, diff --git a/tests/unit/connectors/test_async_gpu_connector.py b/tests/unit/connectors/test_async_gpu_connector.py index 7d797473..c71bde7f 100644 --- a/tests/unit/connectors/test_async_gpu_connector.py +++ b/tests/unit/connectors/test_async_gpu_connector.py @@ -68,9 +68,9 @@ def test_slot_fields_are_disjoint_and_fit_inside_the_slot(layout: SlotLayout): assert layout.expand_idx_off >= layout.header_off + layout.header_words * 4 assert layout.weights_off >= layout.expand_idx_off + 100 * 4 assert layout.routed_x_off >= layout.weights_off + 100 * 4 - # The payload is sized by distinct tokens, not by partials. - assert layout.shared_idx_off >= layout.routed_x_off + 32 * 8 * 2 - assert layout.shared_x_off >= layout.shared_idx_off + 32 * 4 + # The payload is sized by distinct tokens, not by partials, and the shared + # rows are a contiguous range so no index rides along with them. + assert layout.shared_x_off >= layout.routed_x_off + 32 * 8 * 2 assert layout.slot_bytes >= layout.shared_x_off + 32 * 8 * 2 @@ -82,7 +82,6 @@ def test_every_field_offset_is_viewable_as_int32_and_payload(layout: SlotLayout) layout.expand_idx_off, layout.weights_off, layout.routed_x_off, - layout.shared_idx_off, layout.shared_x_off, ): assert offset % 4 == 0 From 3f8c6b01b87129fcb1a18a25f761fbdd7479ad74 Mon Sep 17 00:00:00 2001 From: specture724 Date: Fri, 14 Aug 2026 17:53:31 +0800 Subject: [PATCH 12/18] perf: read the dispatch plan back into pinned memory The readback used .cpu(), which allocates a pageable destination. A pageable device-to-host copy cannot be handed straight to the DMA engine, so the driver stages it through its own buffer and blocks inside the memcpy: a profile of an Attention rank measured 466us of host time per call over 468 calls, 218 ms of a 1163 ms window, against 8us for the device-to-device peer writes and 16us for the same direction into pinned memory. Copying into a pre-allocated pinned buffer instead makes it an ordinary transfer. The buffer is sized once from the topology, and the concatenation writes into a matching device buffer rather than allocating one per layer. This is the readback I said earlier was worth about 2 ms. That was wrong: I had measured its cudaStreamSynchronize and missed that a pageable copy does its waiting inside cudaMemcpyAsync, where it does not show up as a synchronize at all. Measured on 4x L20X, DeepSeek-V2-Lite 2A2F, SECS=25, mean TTFT against the previous commit: 106 -> 101 ms at 8 rps, 129 -> 125 at 16, 205 -> 191 at 32, 308 -> 298 at 48, 727 -> 724 at 64. Against the synchronous connector that is 29% faster at 8 rps, 20% at 16, and within 4% at 32. Co-Authored-By: Claude Opus 5 Signed-off-by: specture724 --- afd_plugin/connectors/gpu/async_gpu.py | 37 ++++++++++++++++++-------- 1 file changed, 26 insertions(+), 11 deletions(-) diff --git a/afd_plugin/connectors/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py index 8c21c837..5a7c1653 100644 --- a/afd_plugin/connectors/gpu/async_gpu.py +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -416,6 +416,12 @@ def __init__( payload_itemsize=torch.empty(0, dtype=self.payload_dtype).element_size(), ) + # One readback per dispatch: per-expert counts followed by the three + # per-rank vectors combine and the headers need. + plan_words = self.ffn_size * (self.expert_per_rank + 3) + self._plan_host = torch.empty(plan_words, dtype=torch.int32).pin_memory() + self._plan_device: Tensor | None = None + self.pg: ProcessGroup | None = None self.window: SymmWindow | None = None self._seq = 0 @@ -456,6 +462,11 @@ def init_afd_connector(self) -> None: timeout=timedelta(minutes=30), ) device = torch.device("cuda", self.local_rank) + self._plan_device = torch.empty( + self._plan_host.numel(), + dtype=torch.int32, + device=device, + ) self.window = SymmWindow( num_regions=self.num_regions, ring_depth=self.ring_depth, @@ -598,18 +609,22 @@ def send_attn_output( # computing the destination offsets on the device and writing # full-capacity segments instead. num_experts = self.ffn_size * self.expert_per_rank - plan_host = ( - torch.cat( - ( - plan.counts.to(torch.int32), - plan.uniq_per_rank, - plan.routed_per_rank, - plan.segment_start, - ), - ) - .cpu() - .tolist() + # Land the readback in pinned memory. ``.cpu()`` allocates a pageable + # destination, and a pageable device-to-host copy has to stage through + # a driver buffer: a profile measured 466us of host time per call + # against 16us for the same copy into pinned memory, which made this one + # readback a fifth of the Attention rank's host time. + torch.cat( + ( + plan.counts.to(torch.int32), + plan.uniq_per_rank, + plan.routed_per_rank, + plan.segment_start, + ), + out=self._plan_device, ) + self._plan_host.copy_(self._plan_device) + plan_host = self._plan_host.tolist() counts_host = plan_host[:num_experts] rank_fields = plan_host[num_experts:] uniq_host = rank_fields[: self.ffn_size] From a0947607e20b59db2ff49405ed7237e240130b23 Mon Sep 17 00:00:00 2001 From: specture724 Date: Fri, 14 Aug 2026 18:07:46 +0800 Subject: [PATCH 13/18] perf: find the expert boundaries with searchsorted, not bincount torch.bincount sizes its output from the maximum value in the data, so it reads that maximum back to the host -- into pageable memory, which blocks inside the copy. A profile of an Attention rank at 48 rps measured 783us of host time per call over 260 calls, 203 ms of a 1170 ms window: the largest single host cost, inside a call whose output size we already knew. The partials are already sorted by global expert id for the dispatch, so the boundaries are one searchsorted over that sorted array against the expert range. It needs no readback, and the lower boundaries are the offsets that the cumulative sum used to produce, so a pass goes with it. Measured on 4x L20X, DeepSeek-V2-Lite 2A2F, SECS=25, mean TTFT against the previous commit: 101 -> 95 ms at 8 rps, 125 -> 118 at 16, 191 -> 178 at 32, 298 -> 276 at 48, 724 -> 517 at 64. Achieved throughput at offered 64 goes 58.4 -> 59.4 rps against the synchronous connector's 59.8. Async is now ahead of synchronous at 8, 16 and 32 rps (95 vs 143, 118 vs 155, 178 vs 185) and behind at 48 and 64. Co-Authored-By: Claude Opus 5 Signed-off-by: specture724 --- afd_plugin/connectors/gpu/async_gpu.py | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/afd_plugin/connectors/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py index 5a7c1653..db159cd3 100644 --- a/afd_plugin/connectors/gpu/async_gpu.py +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -294,12 +294,29 @@ def plan_dispatch( num_tokens, num_slots = topk_ids.shape flat = topk_ids.reshape(-1).to(torch.int64) order = torch.argsort(flat, stable=True) - counts = torch.bincount(flat, minlength=ffn_size * expert_per_rank) - offsets = torch.cumsum(counts, dim=0) - counts + sorted_experts = flat[order] + + # 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. + bounds = torch.searchsorted( + sorted_experts, + torch.arange( + ffn_size * expert_per_rank + 1, + device=flat.device, + dtype=flat.dtype, + ), + ) + offsets = bounds[:-1] + counts = bounds[1:] - offsets token_of_partial = order // num_slots weights = topk_weights.reshape(-1)[order].to(torch.float32) - dest_rank = flat[order] // expert_per_rank + dest_rank = sorted_experts // expert_per_rank # Partials sharing a (destination, token) share a shipped row. Sorting by # that key makes them adjacent, so "is this row new" is one neighbour From ce2a6acc1dc1c1eb1fde3492e6494e672cc7c122 Mon Sep 17 00:00:00 2001 From: specture724 Date: Mon, 17 Aug 2026 11:24:54 +0800 Subject: [PATCH 14/18] test: give the lifecycle connector fake an extra_info Rebasing onto upstream/main brought together upstream's connector-lifecycle test and this branch's constructor, which reads connector.extra_info to decide whether async MoE ubatching is configured. Every real connector has one from AFDConnectorBase; the fake did not, so construction raised. Co-Authored-By: Claude Opus 5 --- tests/unit/v1/worker/test_attention_model_runner.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unit/v1/worker/test_attention_model_runner.py b/tests/unit/v1/worker/test_attention_model_runner.py index f6341458..9c18cbfb 100644 --- a/tests/unit/v1/worker/test_attention_model_runner.py +++ b/tests/unit/v1/worker/test_attention_model_runner.py @@ -1005,6 +1005,9 @@ def __init__(self, events): self.events = events self.control_plane = object() self._initialized = False + # Every real connector carries one; the runner reads it to decide + # whether async MoE ubatching is configured. + self.extra_info = None @property def is_initialized(self): @@ -1107,6 +1110,7 @@ def test_attention_runner_load_model_initializes_connector_after_weights( assert events == expected assert connector.is_initialized is True + def test_connector_driven_runs_skip_cross_dp_batch_coordination(): """An idle Attention replica never joins the DP all-reduce. From 6d26165255d51a5281ad05f54c0742005379398d Mon Sep 17 00:00:00 2001 From: specture724 Date: Mon, 17 Aug 2026 21:00:19 +0800 Subject: [PATCH 15/18] perf: build dispatch headers on the device and write slots at capacity The send path read the routing back to the host once per MoE layer to size each destination's slice. A profile put that synchronize at 392us a call -- not the copy, but the host waiting for everything queued ahead of it -- and it capped run-ahead at one layer, turning the forward into a sum of per-layer maxima instead of a max of sums. Write every slot at capacity instead: the payload carries the whole batch and the index arrays carry every partial, with a new segment_start header word telling a destination which run is its own. The header's routing tail is filled straight from the plan on the device, so nothing is read back. The bytes are nearly free -- with topk slots over ffn_size destinations a token misses a given destination only (1 - 1/ffn_size) ** topk of the time, 1.6% at 2A2F -- and deduplication, its presence grid and the combine scatter all go away with it. Signed-off-by: specture724 --- afd_plugin/connectors/gpu/async_gpu.py | 313 ++++++++---------- afd_plugin/connectors/gpu/symm_window.py | 184 +++++++--- .../connectors/test_async_gpu_connector.py | 84 ++--- 3 files changed, 315 insertions(+), 266 deletions(-) diff --git a/afd_plugin/connectors/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py index db159cd3..9d8dd06e 100644 --- a/afd_plugin/connectors/gpu/async_gpu.py +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -43,9 +43,14 @@ 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, + HEADER_HOST_WORDS, SlotLayout, SymmWindow, encode_header, + encode_header_host_words, ) from afd_plugin.connectors.metadata import ( AFDA2FTransferPayload, @@ -178,9 +183,9 @@ class GpuAsyncTransferState(AFDTransferState): ``expert_counts_host`` is the same numbers as they were decoded on the host, kept so the combine header can be built without reading the device back. - ``expand_idx`` and ``weights`` describe the partials: which shipped row each - one reads on the way in, and what to weight it by when the expert output is - reduced back to one row per shipped token on the way out. + ``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 @@ -191,7 +196,6 @@ class GpuAsyncTransferState(AFDTransferState): seq: int = 0 num_tokens: int = 0 routed_tokens: int = 0 - uniq_tokens: int = 0 shared_tokens: int = 0 group_list: Tensor | None = None expert_counts_host: list[int] = field(default_factory=list) @@ -218,20 +222,17 @@ class GpuAsyncFFNWorkItem: class _PendingDispatch: """Attention-side record of one in-flight layer, popped by combine recv. - ``uniq_ids[r]`` holds the token ids behind the rows shipped to FFN rank - ``r``, in the order they were shipped. The reply comes back one row per - shipped token, already weighted and summed, so combine is a scatter-add at - exactly those ids. + 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. Both are 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. + 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 - uniq_ids: list[Tensor] shared_slices: list[slice] num_tokens: int ring: int @@ -243,34 +244,28 @@ class _PendingDispatch: class DispatchPlan: """One layer's routing plan, already in the order every consumer wants. - Partial-indexed fields (``expand_idx``, ``weights``, length - ``num_tokens * topk``) are sorted by global expert, so a destination's slice - arrives grouped by local expert and feeds the grouped GEMM directly. - Row-indexed fields (``uniq_token_ids``) are grouped by destination rank, so - a destination's payload is also one contiguous slice. + 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``. - offsets: Exclusive prefix sum of ``counts``. - uniq_per_rank: Distinct tokens each FFN rank receives. + ``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 partials begin. Summing counts and - prefix-summing them on the host cost a Python loop per layer; they - come back in the same readback instead. - expand_idx: Per partial, which of its destination's shipped rows it - reads. Already rank-local, so a slice needs no rebasing. - uniq_token_ids: Per shipped row, the token it carries. + 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 - offsets: Tensor - uniq_per_rank: Tensor routed_per_rank: Tensor segment_start: Tensor expand_idx: Tensor - uniq_token_ids: Tensor weights: Tensor @@ -281,17 +276,17 @@ def plan_dispatch( ffn_size: int, expert_per_rank: int, ) -> DispatchPlan: - """Cluster ``(token, topk_slot)`` partials by destination and deduplicate. + """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. A second sort by - ``(destination, token)`` makes the partials that share a shipped row - adjacent, which is what turns ``topk`` copies of a token into one. + 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: the host learns nothing here, so the caller - can fetch ``counts`` and ``uniq_per_rank`` in a single readback. + 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_tokens, num_slots = topk_ids.shape + num_slots = topk_ids.shape[1] flat = topk_ids.reshape(-1).to(torch.int64) order = torch.argsort(flat, stable=True) sorted_experts = flat[order] @@ -314,48 +309,18 @@ def plan_dispatch( offsets = bounds[:-1] counts = bounds[1:] - offsets - token_of_partial = order // num_slots + # 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 = (order // num_slots).to(torch.int32) weights = topk_weights.reshape(-1)[order].to(torch.float32) - dest_rank = sorted_experts // expert_per_rank - - # Partials sharing a (destination, token) share a shipped row. Sorting by - # that key makes them adjacent, so "is this row new" is one neighbour - # comparison and the row numbering is its running sum. - key = dest_rank * num_tokens + token_of_partial - key_order = torch.argsort(key, stable=True) - sorted_key = key[key_order] - is_new = torch.ones_like(sorted_key, dtype=torch.int32) - is_new[1:] = (sorted_key[1:] != sorted_key[:-1]).to(torch.int32) - row_of_partial = is_new.cumsum(0) - 1 - - uniq_per_rank = torch.zeros( - ffn_size, - dtype=torch.int32, - device=topk_ids.device, - ) - uniq_per_rank.index_add_(0, dest_rank[key_order], is_new) - - # Duplicates write the same token id to the same row, so the scatter is - # well defined despite the repeated indices. - uniq_token_ids = torch.zeros_like(token_of_partial) - uniq_token_ids.scatter_(0, row_of_partial, token_of_partial[key_order]) - # Back to expert-sorted order, rebased so each destination's slice indexes - # its own payload from zero. - row_base = (torch.cumsum(uniq_per_rank, 0) - uniq_per_rank).to(torch.int32) - expand_idx = torch.empty_like(row_of_partial, dtype=torch.int32) - expand_idx.scatter_(0, key_order, row_of_partial.to(torch.int32)) - expand_idx -= row_base[dest_rank] - - by_rank = counts.view(ffn_size, expert_per_rank) + # 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, - offsets=offsets, - uniq_per_rank=uniq_per_rank, - routed_per_rank=by_rank.sum(1).to(torch.int32), - segment_start=offsets.view(ffn_size, expert_per_rank)[:, 0].to(torch.int32), + 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, - uniq_token_ids=uniq_token_ids, weights=weights, ) @@ -433,11 +398,7 @@ def __init__( payload_itemsize=torch.empty(0, dtype=self.payload_dtype).element_size(), ) - # One readback per dispatch: per-expert counts followed by the three - # per-rank vectors combine and the headers need. - plan_words = self.ffn_size * (self.expert_per_rank + 3) - self._plan_host = torch.empty(plan_words, dtype=torch.int32).pin_memory() - self._plan_device: Tensor | None = None + self._header_device: Tensor | None = None self.pg: ProcessGroup | None = None self.window: SymmWindow | None = None @@ -479,8 +440,8 @@ def init_afd_connector(self) -> None: timeout=timedelta(minutes=30), ) device = torch.device("cuda", self.local_rank) - self._plan_device = torch.empty( - self._plan_host.numel(), + self._header_device = torch.empty( + (self.ffn_size, self.layout.header_words), dtype=torch.int32, device=device, ) @@ -555,6 +516,66 @@ def _require_initialized(self) -> SymmWindow: raise RuntimeError("AFD async GPU connector is not initialized") return self.window + def _headers_for( + self, + *, + seq: int, + layer_idx: int, + stage_idx: int, + num_tokens: int, + plan: DispatchPlan, + ) -> Tensor: + """Assemble one dispatch header per FFN rank, on the device. + + The host fills the prefix it knows and the plan fills the routing tail + without ever leaving the device. Reading the routing back to encode it + here 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. + """ + assert self._header_device is not None + # A fresh staging buffer per dispatch, from the pinned caching allocator, + # dropped as soon as the copy is queued. The allocator will not hand the + # block out again until that copy has run, which is exactly the guarantee + # an asynchronous copy needs and the reason there is no event here. + # + # Reusing one buffer is the obvious thing and it is wrong: with nothing + # blocking on the layer path the host runs several layers ahead, so it + # rewrites the buffer under the in-flight copy. The header then carries a + # later layer's sequence number, the FFN rank echoes it onto the reply + # flag, and because the stream wait is a ``GEQ`` compare the Attention + # rank leaves it early and overwrites a slot still being read. A stress + # test of that pattern corrupted 298 copies out of 300. + staging = torch.empty( + (self.ffn_size, HEADER_HOST_WORDS), + dtype=torch.int32, + pin_memory=True, + ) + staging_words = staging.numpy() + for ffn_rank in range(self.ffn_size): + shared_start = ffn_rank * num_tokens // self.ffn_size + shared_stop = (ffn_rank + 1) * num_tokens // self.ffn_size + staging_words[ffn_rank] = encode_header_host_words( + seq=seq, + src_role_rank=self.role_rank, + layer_idx=layer_idx, + stage_idx=stage_idx, + num_tokens=num_tokens, + shared_tokens=shared_stop - shared_start, + topk=self.topk, + flags=0, + ) + headers = self._header_device + headers[:, :HEADER_HOST_WORDS].copy_(staging, non_blocking=True) + 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 # ================================================================== @@ -619,35 +640,6 @@ def send_attn_output( ffn_size=self.ffn_size, expert_per_rank=self.expert_per_rank, ) - # One D2H per send, carrying everything the host needs to slice the - # plan: offsets are a prefix sum, cheaper to redo here than to fetch. - # This is the last synchronize left on the layer path, and it is - # intrinsic to slicing the segments on the host -- removing it means - # computing the destination offsets on the device and writing - # full-capacity segments instead. - num_experts = self.ffn_size * self.expert_per_rank - # Land the readback in pinned memory. ``.cpu()`` allocates a pageable - # destination, and a pageable device-to-host copy has to stage through - # a driver buffer: a profile measured 466us of host time per call - # against 16us for the same copy into pinned memory, which made this one - # readback a fifth of the Attention rank's host time. - torch.cat( - ( - plan.counts.to(torch.int32), - plan.uniq_per_rank, - plan.routed_per_rank, - plan.segment_start, - ), - out=self._plan_device, - ) - self._plan_host.copy_(self._plan_device) - plan_host = self._plan_host.tolist() - counts_host = plan_host[:num_experts] - rank_fields = plan_host[num_experts:] - uniq_host = rank_fields[: self.ffn_size] - routed_host = rank_fields[self.ffn_size : 2 * self.ffn_size] - start_host = rank_fields[2 * self.ffn_size :] - # 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 @@ -655,25 +647,15 @@ def send_attn_output( # 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)) - uniq_ids: list[Tensor] = [] shared_slices: list[slice] = [] - uniq_start = 0 + headers = self._headers_for( + seq=self._seq, + layer_idx=metadata.layer_idx, + stage_idx=stage_idx, + num_tokens=num_tokens, + plan=plan, + ) for ffn_rank in range(self.ffn_size): - base = ffn_rank * self.expert_per_rank - expert_counts = counts_host[base : base + self.expert_per_rank] - routed_tokens = routed_host[ffn_rank] - segment = slice( - start_host[ffn_rank], - start_host[ffn_rank] + routed_tokens, - ) - # The shipped rows this rank owns, and the tokens they carry. The - # ids stay here rather than going on the wire: combine scatters the - # reply back to exactly these rows. - uniq_tokens = uniq_host[ffn_rank] - rows = plan.uniq_token_ids[uniq_start : uniq_start + uniq_tokens] - uniq_start += uniq_tokens - uniq_ids.append(rows) - # Shared-expert tokens are split into 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. Round-robin needed an @@ -682,51 +664,37 @@ def send_attn_output( shared_start = ffn_rank * num_tokens // self.ffn_size shared_stop = (ffn_rank + 1) * num_tokens // self.ffn_size shared_slices.append(slice(shared_start, shared_stop)) - header = encode_header( - self.layout, - seq=self._seq, - src_role_rank=self.role_rank, - layer_idx=metadata.layer_idx, - stage_idx=stage_idx, - num_tokens=num_tokens, - routed_tokens=routed_tokens, - shared_tokens=shared_stop - shared_start, - topk=self.topk, - flags=0, - expert_counts=expert_counts, - uniq_tokens=uniq_tokens, - ) - # Rows are gathered straight into the FFN rank's window; handing - # ``write_slot`` the index instead of a gathered tensor saves a - # local write and re-read of every payload byte. + # 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=header, - expand_idx=plan.expand_idx[segment], - weights=plan.weights[segment], + header=headers[ffn_rank], + expand_idx=plan.expand_idx, + weights=plan.weights, routed_x=hidden_states, - routed_rows=rows, shared_x=hidden_states[shared_start:shared_stop], + flag_value=self._seq, ) logger.debug( "AFD dispatch sent: A%d layer=%d stage=%d tokens=%d ring=%d " - "rows_per_ffn=%s partials=%d awaiting_ffn=%s", + "partials=%d awaiting_ffn=%s", self.role_rank, metadata.layer_idx, stage_idx, num_tokens, ring, - uniq_host, num_tokens * self.topk, expected_ffn, ) self._pending.setdefault(stage_idx, []).append( _PendingDispatch( context=context, - uniq_ids=uniq_ids, shared_slices=shared_slices, num_tokens=num_tokens, ring=ring, @@ -780,15 +748,15 @@ def recv_ffn_output( # An FFN rank replies into the region it owns, on the ring the # dispatch used. window.stream_wait(ffn_rank, pending.ring, pending.seq) - uniq_ids = pending.uniq_ids[ffn_rank] - if uniq_ids.numel(): - # One row per token this rank was sent, already weighted and - # summed over that token's partials on the FFN side. - accumulator.index_add_( - 0, - uniq_ids, - window.local_routed(ffn_rank, pending.ring, uniq_ids.numel()), - ) + # 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: @@ -824,9 +792,11 @@ def recv_attn_output( 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 payload carries each token once, so the rows are expanded back to - one per partial here -- a local gather that replaces the duplicate rows - the sender used to put on the wire. + 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. """ window = self._require_initialized() timeout_ms = int(kwargs.get("timeout_ms", 0)) @@ -854,6 +824,7 @@ def recv_attn_output( arrived.region, arrived.ring, header.routed_tokens, + header.segment_start, ).to(torch.int64) states = GpuAsyncTransferState( region=arrived.region, @@ -864,7 +835,6 @@ def recv_attn_output( stage_idx=header.stage_idx, num_tokens=header.num_tokens, routed_tokens=header.routed_tokens, - uniq_tokens=header.uniq_tokens, shared_tokens=header.shared_tokens, group_list=window.local_expert_counts(arrived.region, arrived.ring), expert_counts_host=header.expert_counts, @@ -873,6 +843,7 @@ def recv_attn_output( arrived.region, arrived.ring, header.routed_tokens, + header.segment_start, ), expand_x_shared=window.local_shared( arrived.region, @@ -901,7 +872,7 @@ def recv_attn_output( hidden_states=window.local_routed( arrived.region, arrived.ring, - header.uniq_tokens, + header.num_tokens, ).index_select(0, expand_idx), context=AFDTransferContext(metadata=metadata, states=states), ) @@ -912,12 +883,13 @@ def send_ffn_output( context: AFDTransferContext, **kwargs: Any, ) -> None: - """Reduce expert output to one row per shipped token and write it back. + """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, so the reply carries the same rows the dispatch did. Doing it on - this side keeps the duplicates off the wire and leaves the Attention - side a plain scatter-add. + 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 @@ -936,7 +908,7 @@ def send_ffn_output( "AFD async GPU send_ffn_output requires the dispatch expansion", ) reduced = torch.zeros( - (states.uniq_tokens, self.hidden_size), + (states.num_tokens, self.hidden_size), dtype=self.payload_dtype, device=ffn_output.device, ) @@ -958,7 +930,6 @@ def send_ffn_output( topk=self.topk, flags=0, expert_counts=states.expert_counts_host, - uniq_tokens=states.uniq_tokens, echo_seq=states.seq, ) # ``reduced`` is float32 and the slot is the payload dtype; the copy diff --git a/afd_plugin/connectors/gpu/symm_window.py b/afd_plugin/connectors/gpu/symm_window.py index 139a1fb8..44e48997 100644 --- a/afd_plugin/connectors/gpu/symm_window.py +++ b/afd_plugin/connectors/gpu/symm_window.py @@ -18,13 +18,19 @@ of the batch, so the slot carries only their count -- the range is implied by which FFN rank the slot belongs to. -A token is shipped once per destination rank however many of its topk slots -landed there, so the payload rows are *distinct tokens* (``uniq_tokens``, at -most ``token_cap``) while ``expand_idx`` and ``weights`` carry one entry per -*partial* (``routed_tokens``, at most ``partial_cap``). Those two index arrays -are 4 bytes a row against the payload's ``hidden_size`` elements, which is why -they can be sized for the true worst case -- every partial to one rank -- while -the payload is sized by a bound that no routing skew can exceed. +Every slot is written at capacity: the payload carries the whole batch +(``num_tokens`` rows) and the ``expand_idx``/``weights`` arrays carry every +partial, with ``segment_start``/``routed_tokens`` in the header telling a +destination which run of partials is its own. Sizing the writes by the routing +instead would save a fraction of the bytes and cost a device-to-host readback +per layer to learn the sizes, which measured far more than the bytes are worth: +with ``topk`` slots spread over ``ffn_size`` destinations a token misses a given +destination only ``(1 - 1/ffn_size) ** topk`` of the time, so at 2A2F and +``topk=6`` the capacity payload is about 1.6% larger than the exact one. + +``expand_idx`` and ``weights`` are 4 bytes a row against the payload's +``hidden_size`` elements, so shipping them whole to every peer costs well under +a percent of the slot. Flag words are written *after* the payload on the same stream. Same-stream device-to-device copies complete in issue order, so a visible flag implies a @@ -53,7 +59,7 @@ # Header words shared by dispatch and combine, followed by expert_counts. HEADER_MAGIC = 0x41464447 # "AFDG" -HEADER_VERSION = 2 +HEADER_VERSION = 3 _H_MAGIC = 0 _H_VERSION = 1 _H_SEQ = 2 @@ -61,12 +67,17 @@ _H_LAYER_IDX = 4 _H_STAGE_IDX = 5 _H_NUM_TOKENS = 6 -_H_ROUTED_TOKENS = 7 # partials: one per (token, topk slot) landing here -_H_SHARED_TOKENS = 8 -_H_TOPK = 9 -_H_FLAGS = 10 -_H_ECHO_SEQ = 11 # combine only: the dispatch seq being answered -_H_UNIQ_TOKENS = 12 # payload rows: distinct tokens behind those partials +_H_SHARED_TOKENS = 7 +_H_TOPK = 8 +_H_FLAGS = 9 +_H_ECHO_SEQ = 10 # combine only: the dispatch seq being answered +# Everything from here on is routing, which only the device knows. Keeping those +# words in one contiguous tail is what lets a sender fill the host-known prefix +# with a single copy and the rest straight from the plan, so no dispatch ever +# reads the routing back to the host. +HEADER_HOST_WORDS = 11 +H_ROUTED_TOKENS = 11 # partials: one per (token, topk slot) landing here +H_SEGMENT_START = 12 # where those partials begin in the shipped index arrays HEADER_FIXED_WORDS = 13 FLAG_EMPTY = 0 @@ -130,6 +141,44 @@ def build( ) +def encode_header_host_words( + *, + seq: int, + src_role_rank: int, + layer_idx: int, + stage_idx: int, + num_tokens: int, + shared_tokens: int, + topk: int, + flags: int, + echo_seq: int = 0, +) -> np.ndarray: + """Build the host-known prefix of one slot header as int32. + + The routing tail -- ``routed_tokens``, ``segment_start`` and the per-expert + counts -- is not here: a dispatch fills it straight from the plan on the + device, which is what keeps the send path free of a readback. + """ + # Built with numpy and assigned as one tuple, not field by field into a + # tensor: this runs once per peer per MoE layer, and a dozen tensor setitems + # there cost more host time than the transfer they describe. + words = np.empty(HEADER_HOST_WORDS, dtype=np.int32) + words[:] = ( + HEADER_MAGIC, + HEADER_VERSION, + seq, + src_role_rank, + layer_idx, + stage_idx, + num_tokens, + shared_tokens, + topk, + flags, + echo_seq, + ) + return words + + def encode_header( layout: SlotLayout, *, @@ -143,38 +192,38 @@ def encode_header( topk: int, flags: int, expert_counts: list[int], - uniq_tokens: int = 0, + segment_start: int = 0, echo_seq: int = 0, ) -> torch.Tensor: - """Build the fixed header for one slot as a CPU int32 tensor.""" + """Build a whole slot header as a CPU int32 tensor. + + Used where the routing is already on the host: the FFN reply, which decoded + it from the dispatch it answers, and the shutdown announcement, which has no + routing at all. + """ expert_per_rank = layout.header_words - HEADER_FIXED_WORDS if len(expert_counts) != expert_per_rank: raise ValueError( f"expert_counts must have {expert_per_rank} entries, " f"got {len(expert_counts)}", ) - # Built with numpy and wrapped, not assigned field by field into a tensor: - # this runs once per peer per MoE layer, and fourteen tensor setitems there - # cost more host time than the transfer they describe. torch.from_numpy - # shares the buffer, so the wrap is free. header = np.empty(layout.header_words, dtype=np.int32) - header[:HEADER_FIXED_WORDS] = ( - HEADER_MAGIC, - HEADER_VERSION, - seq, - src_role_rank, - layer_idx, - stage_idx, - num_tokens, - routed_tokens, - shared_tokens, - topk, - flags, - echo_seq, - uniq_tokens, + header[:HEADER_HOST_WORDS] = encode_header_host_words( + seq=seq, + src_role_rank=src_role_rank, + layer_idx=layer_idx, + stage_idx=stage_idx, + num_tokens=num_tokens, + shared_tokens=shared_tokens, + topk=topk, + flags=flags, + echo_seq=echo_seq, ) + header[H_ROUTED_TOKENS] = routed_tokens + header[H_SEGMENT_START] = segment_start if expert_per_rank: header[HEADER_FIXED_WORDS:] = expert_counts + # torch.from_numpy shares the buffer, so the wrap is free. return torch.from_numpy(header) @@ -192,7 +241,7 @@ class SlotHeader: topk: int flags: int echo_seq: int - uniq_tokens: int + segment_start: int expert_counts: list[int] @property @@ -219,12 +268,12 @@ def decode_header(header: torch.Tensor) -> SlotHeader: layer_idx=values[_H_LAYER_IDX], stage_idx=values[_H_STAGE_IDX], num_tokens=values[_H_NUM_TOKENS], - routed_tokens=values[_H_ROUTED_TOKENS], + routed_tokens=values[H_ROUTED_TOKENS], shared_tokens=values[_H_SHARED_TOKENS], topk=values[_H_TOPK], flags=values[_H_FLAGS], echo_seq=values[_H_ECHO_SEQ], - uniq_tokens=values[_H_UNIQ_TOKENS], + segment_start=values[H_SEGMENT_START], expert_counts=values[HEADER_FIXED_WORDS:], ) @@ -431,6 +480,10 @@ def write_slot( answers instead, so the rank waiting for it knows the value to wait for before it arrives -- that is what lets the wait happen on a stream rather than on the host. + + ``header`` may live on the device, which is how a dispatch ships routing + it never read back; such a header has no readable sequence number, so + ``flag_value`` is then required. """ layout = self.layout routed_count = _row_count(routed_x, routed_rows) @@ -449,15 +502,16 @@ def write_slot( f"shared tokens {shared_count} exceed token_cap {layout.token_cap}", ) - # Stage through pinned memory so the copy is asynchronous: a pageable - # source would force a blocking transfer, and there is one header per - # peer per layer. - staging = self._header_send[peer][ring] - staging.copy_(header) - self._capacity_view(peer, region, ring, layout.header_off).copy_( - staging, - non_blocking=True, - ) + header_view = self._capacity_view(peer, region, ring, layout.header_off) + if header.is_cuda: + header_view.copy_(header, non_blocking=True) + else: + # Stage through pinned memory so the copy is asynchronous: a + # pageable source would force a blocking transfer, and there is one + # header per peer per layer. + staging = self._header_send[peer][ring] + staging.copy_(header) + header_view.copy_(staging, non_blocking=True) def write_field( field_off: int, @@ -513,7 +567,14 @@ def write_field( shared_rows, ) - seq = int(header[_H_SEQ].item()) if flag_value is None else flag_value + if flag_value is None: + if header.is_cuda: + raise ValueError( + "write_slot needs flag_value for a device-resident header: " + "reading its sequence number back would synchronize", + ) + flag_value = int(header[_H_SEQ].item()) + seq = flag_value flag_idx = region * self.ring_depth + ring self._flag_view(peer, flag_idx).fill_(seq) @@ -611,25 +672,42 @@ def local_expert_counts(self, region: int, ring: int) -> torch.Tensor: ) return header[HEADER_FIXED_WORDS:] - def local_expand_idx(self, region: int, ring: int, count: int) -> torch.Tensor: + def local_expand_idx( + self, + region: int, + ring: int, + count: int, + start: int = 0, + ) -> torch.Tensor: + """Device view of this destination's run of partial indices. + + Senders ship the whole array, so a destination's own partials start at + the ``segment_start`` its header carries. + """ return self._view( self.rank, region, ring, self.layout.expand_idx_off, - (count,), + (start + count,), torch.int32, - ) + )[start:] - def local_weights(self, region: int, ring: int, count: int) -> torch.Tensor: + def local_weights( + self, + region: int, + ring: int, + count: int, + start: int = 0, + ) -> torch.Tensor: return self._view( self.rank, region, ring, self.layout.weights_off, - (count,), + (start + count,), torch.float32, - ) + )[start:] def local_routed(self, region: int, ring: int, count: int) -> torch.Tensor: return self._view( diff --git a/tests/unit/connectors/test_async_gpu_connector.py b/tests/unit/connectors/test_async_gpu_connector.py index c71bde7f..faaed020 100644 --- a/tests/unit/connectors/test_async_gpu_connector.py +++ b/tests/unit/connectors/test_async_gpu_connector.py @@ -106,7 +106,7 @@ def test_header_round_trip(layout: SlotLayout): topk=6, flags=0, expert_counts=[10, 20, 30, 30], - uniq_tokens=25, + segment_start=25, ) decoded = decode_header(header) assert decoded.seq == 7 @@ -118,7 +118,7 @@ def test_header_round_trip(layout: SlotLayout): assert decoded.shared_tokens == 16 assert decoded.topk == 6 assert decoded.expert_counts == [10, 20, 30, 30] - assert decoded.uniq_tokens == 25 + assert decoded.segment_start == 25 assert sum(decoded.expert_counts) == decoded.routed_tokens assert not decoded.is_shutdown @@ -203,17 +203,16 @@ def routing_inputs(): def _destination_slices(plan, ffn_size, expert_per_rank): - """Walk the plan the way send_attn_output does: one slice per destination.""" - counts, offsets = plan.counts.tolist(), plan.offsets.tolist() - uniq = plan.uniq_per_rank.tolist() - uniq_start = 0 + """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): - base = ffn_rank * expert_per_rank - total = sum(counts[base : base + expert_per_rank]) - partials = slice(offsets[base], offsets[base] + total) - rows = plan.uniq_token_ids[uniq_start : uniq_start + uniq[ffn_rank]] - uniq_start += uniq[ffn_rank] - yield ffn_rank, partials, rows + yield ffn_rank, slice(starts[ffn_rank], starts[ffn_rank] + routed[ffn_rank]) def test_every_partial_is_routed_exactly_once(routing_inputs): @@ -227,7 +226,14 @@ def test_every_partial_is_routed_exactly_once(routing_inputs): 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.uniq_per_rank.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): @@ -238,15 +244,15 @@ def test_each_destination_segment_is_grouped_by_local_expert(routing_inputs): ffn_size=_FFN_SIZE, expert_per_rank=_EXPERT_PER_RANK, ) - counts, offsets = plan.counts.tolist(), plan.offsets.tolist() - for ffn_rank, partials, rows in _destination_slices( + 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 = rows.index_select(0, plan.expand_idx[partials].to(torch.int64)) + tokens = plan.expand_idx[partials] cursor = 0 for local_expert in range(_EXPERT_PER_RANK): for _ in range(counts[base + local_expert]): @@ -254,10 +260,9 @@ def test_each_destination_segment_is_grouped_by_local_expert(routing_inputs): assert base + local_expert in topk_ids[token_idx].tolist() cursor += 1 assert cursor == partials.stop - partials.start - assert offsets[base] == partials.start -def test_each_token_is_shipped_once_per_destination(routing_inputs): +def test_every_partial_names_a_token_of_this_batch(routing_inputs): topk_ids, _, topk_weights = routing_inputs plan = plan_dispatch( topk_ids, @@ -265,13 +270,11 @@ def test_each_token_is_shipped_once_per_destination(routing_inputs): ffn_size=_FFN_SIZE, expert_per_rank=_EXPERT_PER_RANK, ) - for ffn_rank, partials, rows in _destination_slices( - plan, - _FFN_SIZE, - _EXPERT_PER_RANK, - ): - shipped = rows.tolist() - assert len(set(shipped)) == len(shipped), "a token was shipped twice" + # 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 @@ -279,15 +282,11 @@ def test_each_token_is_shipped_once_per_destination(routing_inputs): for expert in topk_ids[token].tolist() if base <= expert < base + _EXPERT_PER_RANK } - assert set(shipped) == expected - # Every partial must land inside its destination's own payload. - expand = plan.expand_idx[partials] - assert int(expand.min()) >= 0 - assert int(expand.max()) < len(shipped) + assert set(plan.expand_idx[partials].tolist()) == expected def test_identity_experts_recombine_to_the_weighted_sum(routing_inputs): - """The full chain: ship distinct rows, expand, weight, reduce, scatter.""" + """The full chain: ship the batch, expand, weight, reduce, add.""" topk_ids, hidden_states, topk_weights = routing_inputs plan = plan_dispatch( topk_ids, @@ -296,13 +295,14 @@ def test_identity_experts_recombine_to_the_weighted_sum(routing_inputs): expert_per_rank=_EXPERT_PER_RANK, ) accumulator = torch.zeros(_NUM_TOKENS, _HIDDEN, dtype=torch.float32) - for _, partials, rows in _destination_slices(plan, _FFN_SIZE, _EXPERT_PER_RANK): + for _, partials in _destination_slices(plan, _FFN_SIZE, _EXPERT_PER_RANK): expand = plan.expand_idx[partials].to(torch.int64) - # What crosses the wire, and what the FFN side does with it. - expanded = hidden_states.index_select(0, rows).index_select(0, expand) - reduced = torch.zeros(rows.numel(), _HIDDEN, dtype=torch.float32) + # 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)) - accumulator.index_add_(0, rows, reduced) + # 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)) @@ -340,16 +340,16 @@ def test_routing_can_leave_one_destination_empty(): ) slices = list(_destination_slices(plan, ffn_size, expert_per_rank)) - _, busy_partials, busy_rows = slices[0] + _, busy_partials = slices[0] assert busy_partials.stop - busy_partials.start == 3 - # One token, three of its partials: it is shipped once, read three times. - assert busy_rows.tolist() == [0] + # 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, empty_rows = slices[1] + _, empty_partials = slices[1] assert empty_partials.stop - empty_partials.start == 0 - assert empty_rows.numel() == 0 # Reducing an empty destination must be a no-op, not an error. accumulator = torch.zeros(1, 4, dtype=torch.float32) - accumulator.index_add_(0, empty_rows, torch.zeros(0, 4)) + empty = plan.expand_idx[empty_partials].to(torch.int64) + accumulator.index_add_(0, empty, torch.zeros(0, 4)) assert torch.count_nonzero(accumulator) == 0 From 8a0e7b90ad0363118ab25e20898a72ed4c48e335 Mon Sep 17 00:00:00 2001 From: specture724 Date: Wed, 26 Aug 2026 14:12:32 +0800 Subject: [PATCH 16/18] style: reformat the window roundtrip smoke test This file predates the ruff format hook and never satisfied it: a bare `pre-commit run` on it fails both ruff-check (E501 at the torch.randn line) and ruff-format. Touching it for anything else drags the whole reflow into that commit, so it goes in on its own here. No behaviour change -- the parsed AST is byte-identical before and after. Co-Authored-By: Claude Opus 5 Signed-off-by: specture724 --- tests/e2e/async_gpu_window_roundtrip.py | 98 +++++++++++++++++-------- 1 file changed, 68 insertions(+), 30 deletions(-) diff --git a/tests/e2e/async_gpu_window_roundtrip.py b/tests/e2e/async_gpu_window_roundtrip.py index cc71b7c8..529267d2 100644 --- a/tests/e2e/async_gpu_window_roundtrip.py +++ b/tests/e2e/async_gpu_window_roundtrip.py @@ -59,24 +59,33 @@ def main() -> None: world_size=2, ) if rank == 0: - print(f"slot={layout.slot_bytes / 2**10:.0f}KiB " - f"window={window.total_bytes / 2**10:.0f}KiB", flush=True) + print( + f"slot={layout.slot_bytes / 2**10:.0f}KiB " + f"window={window.total_bytes / 2**10:.0f}KiB", + flush=True, + ) for layer_idx in range(NUM_LAYERS): ring = layer_idx % RING_DEPTH gen = torch.Generator(device="cpu").manual_seed(layer_idx) - hidden = torch.randn(NUM_TOKENS, HIDDEN, generator=gen).to(device, torch.bfloat16) + hidden = torch.randn(NUM_TOKENS, HIDDEN, generator=gen).to( + device, torch.bfloat16 + ) topk_ids = torch.stack( - [torch.randperm(EXPERT_PER_RANK, generator=gen)[:TOPK] - for _ in range(NUM_TOKENS)], + [ + torch.randperm(EXPERT_PER_RANK, generator=gen)[:TOPK] + for _ in range(NUM_TOKENS) + ], ).to(device, torch.int32) topk_weights = torch.rand(NUM_TOKENS, TOPK, generator=gen).to(device) # Both ranks build the plan from the same inputs. Rank 0 sends from it; # rank 1 only uses it as an oracle for what should have arrived. plan = plan_dispatch( - topk_ids, topk_weights, - ffn_size=FFN_SIZE, expert_per_rank=EXPERT_PER_RANK, + topk_ids, + topk_weights, + ffn_size=FFN_SIZE, + expert_per_rank=EXPERT_PER_RANK, ) if rank == 0: @@ -84,17 +93,26 @@ def main() -> None: # Shared rows are a contiguous range; this rank owns all of them. shared = slice(0, NUM_TOKENS) header = encode_header( - layout, seq=layer_idx + 1, src_role_rank=0, layer_idx=layer_idx, - stage_idx=0, num_tokens=NUM_TOKENS, + layout, + seq=layer_idx + 1, + src_role_rank=0, + layer_idx=layer_idx, + stage_idx=0, + num_tokens=NUM_TOKENS, routed_tokens=NUM_TOKENS * TOPK, - shared_tokens=shared.stop - shared.start, topk=TOPK, flags=0, + shared_tokens=shared.stop - shared.start, + topk=TOPK, + flags=0, expert_counts=counts_host, segment_start=0, ) # Capacity form: the whole batch and every partial go out, and the # rank-1 side locates its own run from the header. window.write_slot( - peer=1, region=0, ring=ring, header=header, + peer=1, + region=0, + ring=ring, + header=header, expand_idx=plan.expand_idx, weights=plan.weights, routed_x=hidden, @@ -112,15 +130,21 @@ def main() -> None: assert got.routed_tokens == NUM_TOKENS * TOPK, got.routed_tokens # A reply is a whole batch, so combine adds it without an index. - acc = window.local_routed( - arrived.region, arrived.ring, got.num_tokens, - ).to(torch.float32).clone() + acc = ( + window.local_routed( + arrived.region, + arrived.ring, + got.num_tokens, + ) + .to(torch.float32) + .clone() + ) acc[shared] += window.local_shared( arrived.region, arrived.ring, got.shared_tokens ).to(torch.float32) expected = ( hidden.to(torch.float32) * topk_weights.sum(dim=1, keepdim=True) - + hidden.to(torch.float32) # identity shared expert + + hidden.to(torch.float32) # identity shared expert ) torch.testing.assert_close(acc, expected, rtol=2e-2, atol=2e-2) print(f"layer {layer_idx}: combine matches reference", flush=True) @@ -132,41 +156,55 @@ def main() -> None: got = arrived.header assert got.layer_idx == layer_idx, (got.layer_idx, layer_idx) assert sum(got.expert_counts) == got.routed_tokens - shipped = window.local_routed( - arrived.region, arrived.ring, got.num_tokens) + shipped = window.local_routed(arrived.region, arrived.ring, got.num_tokens) shared_rows = window.local_shared( - arrived.region, arrived.ring, got.shared_tokens) + arrived.region, arrived.ring, got.shared_tokens + ) expand = window.local_expand_idx( - arrived.region, arrived.ring, got.routed_tokens, - got.segment_start).to(torch.int64) + arrived.region, arrived.ring, got.routed_tokens, got.segment_start + ).to(torch.int64) weights = window.local_weights( - arrived.region, arrived.ring, got.routed_tokens, got.segment_start) + arrived.region, arrived.ring, got.routed_tokens, got.segment_start + ) # Gathering by the partial indices must reproduce the sender's rows. expanded = shipped.index_select(0, expand) torch.testing.assert_close(expanded, hidden.index_select(0, expand)) # Identity experts, then the weighted reduce back to token rows. reduced = torch.zeros( - got.num_tokens, HIDDEN, dtype=torch.float32, device=device) + got.num_tokens, HIDDEN, dtype=torch.float32, device=device + ) reduced.index_add_( - 0, expand, expanded.to(torch.float32) * weights.unsqueeze(1)) + 0, expand, expanded.to(torch.float32) * weights.unsqueeze(1) + ) echo = encode_header( - layout, seq=layer_idx + 1, src_role_rank=0, layer_idx=got.layer_idx, - stage_idx=got.stage_idx, num_tokens=got.num_tokens, - routed_tokens=got.routed_tokens, shared_tokens=got.shared_tokens, - topk=TOPK, flags=0, expert_counts=got.expert_counts, + layout, + seq=layer_idx + 1, + src_role_rank=0, + layer_idx=got.layer_idx, + stage_idx=got.stage_idx, + num_tokens=got.num_tokens, + routed_tokens=got.routed_tokens, + shared_tokens=got.shared_tokens, + topk=TOPK, + flags=0, + expert_counts=got.expert_counts, segment_start=got.segment_start, ) window.write_slot( - peer=0, region=0, ring=arrived.ring, header=echo, + peer=0, + region=0, + ring=arrived.ring, + header=echo, expand_idx=None, weights=None, routed_x=reduced, shared_x=shared_rows.clone(), ) - print(f"layer {layer_idx}: dispatch payload verified, echoed back", - flush=True) + print( + f"layer {layer_idx}: dispatch payload verified, echoed back", flush=True + ) dist.barrier() if rank == 0: From 48ac7ecb4c2523961799c1ece0b2d169f1d31d8d Mon Sep 17 00:00:00 2001 From: specture724 Date: Wed, 26 Aug 2026 14:13:01 +0800 Subject: [PATCH 17/18] perf: size the shared field by the per-rank split, not by the batch The slot reserved `token_cap` rows for `shared_x`, the same as `routed_x`, but nothing can ever fill them. Shared-expert tokens are split across the FFN ranks as contiguous chunks, so a slot holds at most `ceil(token_cap / ffn_size)` of them, in both directions: the dispatch ships one rank's chunk, and the reply answers that same chunk. A model with no shared experts does not need the field at all -- and was still being sent one every layer to every peer, because the send path never consulted `has_shared_experts`. Sizing the field by `shared_cap` cuts the payload half of every slot by `(1 - 1/ffn_size)`. Measured on a 2A2F stack (DeepSeek-V2-Lite, 4x L20X, max-num-batched-tokens 8192), the window logged at startup goes from 64.4 MiB per slot / 128.8 MiB total to 48.4 / 96.8, a 25% saving that comes straight back to the KV cache budget. It grows with the FFN rank count: 41% at 2A6F, and 50% for a model with no shared experts, which now ships no shared rows at all. The split arithmetic lived in two places -- `_headers_for`, which puts `shared_tokens` on the wire, and `send_attn_output`, which writes the rows it describes. Those two have to agree or the receiver reads a count that does not match the payload, so they now come from one `_shared_slice` rather than from two copies that can drift. Verified beyond the unit suite, because neither GPU e2e exercises this path: both run 1A1F, where the split is the whole batch and `shared_cap` collapses back to `token_cap`. A 4-GPU 2A2F run with real inference is what covers `shared_cap != token_cap`; greedy output stays coherent, a 90k-token prompt chunk-prefills without tripping the new bound, and no capacity error appears in either log. Co-Authored-By: Claude Opus 5 Signed-off-by: specture724 --- afd_plugin/connectors/gpu/async_gpu.py | 47 ++++++++++++++----- afd_plugin/connectors/gpu/symm_window.py | 20 ++++++-- tests/e2e/async_gpu_window_roundtrip.py | 3 ++ .../connectors/test_async_gpu_connector.py | 45 +++++++++++++++++- 4 files changed, 98 insertions(+), 17 deletions(-) diff --git a/afd_plugin/connectors/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py index 9d8dd06e..4d875029 100644 --- a/afd_plugin/connectors/gpu/async_gpu.py +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -390,10 +390,18 @@ def __init__( # 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(), ) @@ -516,6 +524,25 @@ def _require_initialized(self) -> SymmWindow: 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( self, *, @@ -554,15 +581,14 @@ def _headers_for( ) staging_words = staging.numpy() for ffn_rank in range(self.ffn_size): - shared_start = ffn_rank * num_tokens // self.ffn_size - shared_stop = (ffn_rank + 1) * num_tokens // self.ffn_size + shared = self._shared_slice(ffn_rank, num_tokens) staging_words[ffn_rank] = encode_header_host_words( seq=seq, src_role_rank=self.role_rank, layer_idx=layer_idx, stage_idx=stage_idx, num_tokens=num_tokens, - shared_tokens=shared_stop - shared_start, + shared_tokens=shared.stop - shared.start, topk=self.topk, flags=0, ) @@ -656,14 +682,11 @@ def send_attn_output( plan=plan, ) for ffn_rank in range(self.ffn_size): - # Shared-expert tokens are split into 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. Round-robin needed an - # arange, a gather and a whole slot field per peer per layer to - # achieve the same balance. - shared_start = ffn_rank * num_tokens // self.ffn_size - shared_stop = (ffn_rank + 1) * num_tokens // self.ffn_size - shared_slices.append(slice(shared_start, shared_stop)) + # 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 @@ -677,7 +700,7 @@ def send_attn_output( expand_idx=plan.expand_idx, weights=plan.weights, routed_x=hidden_states, - shared_x=hidden_states[shared_start:shared_stop], + shared_x=hidden_states[shared], flag_value=self._seq, ) diff --git a/afd_plugin/connectors/gpu/symm_window.py b/afd_plugin/connectors/gpu/symm_window.py index 44e48997..5a4330a9 100644 --- a/afd_plugin/connectors/gpu/symm_window.py +++ b/afd_plugin/connectors/gpu/symm_window.py @@ -32,6 +32,13 @@ ``hidden_size`` elements, so shipping them whole to every peer costs well under a percent of the slot. +``shared_x`` is the exception, and is sized ``shared_cap`` rather than +``token_cap``: shared-expert tokens are split across the FFN ranks, so a slot +never holds more than ``ceil(token_cap / ffn_size)`` of them in either +direction, and a model with no shared experts does not need the field at all. +Sizing it like ``routed_x`` reserved a second whole-batch payload per slot that +nothing could ever fill. + Flag words are written *after* the payload on the same stream. Same-stream device-to-device copies complete in issue order, so a visible flag implies a complete payload. That holds for NVLink-mapped peer memory; a cross-node @@ -99,6 +106,7 @@ class SlotLayout: header_words: int partial_cap: int token_cap: int + shared_cap: int hidden_size: int payload_itemsize: int @@ -116,6 +124,7 @@ def build( expert_per_rank: int, partial_cap: int, token_cap: int, + shared_cap: int, hidden_size: int, payload_itemsize: int, ) -> SlotLayout: @@ -125,11 +134,12 @@ def build( weights_off = _align(expand_idx_off + partial_cap * 4) routed_x_off = _align(weights_off + partial_cap * 4) shared_x_off = _align(routed_x_off + token_cap * hidden_size * payload_itemsize) - slot_bytes = _align(shared_x_off + token_cap * hidden_size * payload_itemsize) + slot_bytes = _align(shared_x_off + shared_cap * hidden_size * payload_itemsize) return cls( header_words=header_words, partial_cap=partial_cap, token_cap=token_cap, + shared_cap=shared_cap, hidden_size=hidden_size, payload_itemsize=payload_itemsize, header_off=header_off, @@ -404,8 +414,10 @@ def _capacity_view( sizes, dtype = (layout.partial_cap,), torch.int32 elif field_off == layout.weights_off: sizes, dtype = (layout.partial_cap,), torch.float32 - elif field_off in (layout.routed_x_off, layout.shared_x_off): + elif field_off == layout.routed_x_off: sizes, dtype = (layout.token_cap, hidden), self.payload_dtype + elif field_off == layout.shared_x_off: + sizes, dtype = (layout.shared_cap, hidden), self.payload_dtype else: raise ValueError(f"unknown slot field offset {field_off}") @@ -497,9 +509,9 @@ def write_slot( raise RuntimeError( f"partials {partial_count} exceed partial_cap {layout.partial_cap}", ) - if shared_count > layout.token_cap: + if shared_count > layout.shared_cap: raise RuntimeError( - f"shared tokens {shared_count} exceed token_cap {layout.token_cap}", + f"shared tokens {shared_count} exceed shared_cap {layout.shared_cap}", ) header_view = self._capacity_view(peer, region, ring, layout.header_off) diff --git a/tests/e2e/async_gpu_window_roundtrip.py b/tests/e2e/async_gpu_window_roundtrip.py index 529267d2..c6d96eb6 100644 --- a/tests/e2e/async_gpu_window_roundtrip.py +++ b/tests/e2e/async_gpu_window_roundtrip.py @@ -45,6 +45,9 @@ def main() -> None: expert_per_rank=EXPERT_PER_RANK, partial_cap=NUM_TOKENS * TOPK, token_cap=NUM_TOKENS, + # This smoke test ships every token as shared, and FFN_SIZE is 1, so the + # per-rank split is the whole batch. + shared_cap=-(-NUM_TOKENS // FFN_SIZE), hidden_size=HIDDEN, payload_itemsize=2, ) diff --git a/tests/unit/connectors/test_async_gpu_connector.py b/tests/unit/connectors/test_async_gpu_connector.py index faaed020..b19fb581 100644 --- a/tests/unit/connectors/test_async_gpu_connector.py +++ b/tests/unit/connectors/test_async_gpu_connector.py @@ -4,6 +4,9 @@ from __future__ import annotations +from itertools import pairwise +from types import SimpleNamespace + import pytest pytest.importorskip("torch") @@ -27,10 +30,14 @@ @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, ) @@ -71,7 +78,43 @@ def test_slot_fields_are_disjoint_and_fit_inside_the_slot(layout: SlotLayout): # The payload is sized by distinct tokens, not by partials, and the shared # rows are a contiguous range so no index rides along with them. assert layout.shared_x_off >= layout.routed_x_off + 32 * 8 * 2 - assert layout.slot_bytes >= layout.shared_x_off + 32 * 8 * 2 + # The shared field is sized by the per-rank split, not by the batch: a slot + # can never hold a whole batch of shared rows, and reserving room for one + # doubled the payload half of every slot. + assert layout.slot_bytes >= layout.shared_x_off + 8 * 8 * 2 + assert layout.slot_bytes < layout.shared_x_off + 32 * 8 * 2 + + +@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_every_field_offset_is_viewable_as_int32_and_payload(layout: SlotLayout): From c26c3daf8593e6c760370064c2bfd1696820e894 Mon Sep 17 00:00:00 2001 From: specture724 Date: Wed, 26 Aug 2026 14:16:13 +0800 Subject: [PATCH 18/18] fix: make the shutdown broadcast callable `announce_shutdown` passed `shared_idx=`, which `write_slot` has no parameter for, and omitted `expand_idx` and `weights`, which it requires. Any call raised TypeError before reaching the window. Nothing calls it -- the FFN loop leaves on the `recv_poll_timeout_ms` timeout instead, and the receive half is fully wired (`recv_attn_output` raises ConnectorShutdown on the header bit) -- so the wire protocol's shutdown path has never actually run. The signature drifted when the slot layout replaced a per-destination shared index with a contiguous range; every live call site was updated and this one, having no caller, was not. Since no runtime path exercises it, the regression guard is a signature bind: the test captures the arguments announce_shutdown really sends and binds them against `SymmWindow.write_slot`, which fails on both a stale keyword and a missing required one. Restoring the old call makes it fail with the TypeError it was hiding. Co-Authored-By: Claude Opus 5 Signed-off-by: specture724 --- afd_plugin/connectors/gpu/async_gpu.py | 5 +++- .../connectors/test_async_gpu_connector.py | 30 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/afd_plugin/connectors/gpu/async_gpu.py b/afd_plugin/connectors/gpu/async_gpu.py index 4d875029..6e59e99a 100644 --- a/afd_plugin/connectors/gpu/async_gpu.py +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -1039,13 +1039,16 @@ def announce_shutdown(self) -> None: 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_idx=None, shared_x=None, ) diff --git a/tests/unit/connectors/test_async_gpu_connector.py b/tests/unit/connectors/test_async_gpu_connector.py index b19fb581..b9e2743e 100644 --- a/tests/unit/connectors/test_async_gpu_connector.py +++ b/tests/unit/connectors/test_async_gpu_connector.py @@ -4,6 +4,7 @@ from __future__ import annotations +import inspect from itertools import pairwise from types import SimpleNamespace @@ -23,6 +24,7 @@ FLAG_SHUTDOWN_BIT, HEADER_FIXED_WORDS, SlotLayout, + SymmWindow, decode_header, encode_header, ) @@ -117,6 +119,34 @@ def test_shared_split_is_empty_without_shared_experts(): 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, + _seq=7, + 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 + + def test_every_field_offset_is_viewable_as_int32_and_payload(layout: SlotLayout): # get_buffer takes an element offset, so a byte offset that is not a # multiple of the element size would silently land on the wrong address.