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..6e59e99a --- /dev/null +++ b/afd_plugin/connectors/gpu/async_gpu.py @@ -0,0 +1,1065 @@ +# 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, field +from datetime import timedelta +from typing import TYPE_CHECKING, Any, Final + +import torch +from torch import Tensor +from vllm.logger import init_logger + +from afd_plugin.config import AFDConfig +from afd_plugin.config_utils import ( + coerce_extra_bool, + coerce_extra_positive_int, + coerce_extra_str, +) +from afd_plugin.connectors.async_topology import ( + ASYNC_MOE_REQUEST_SPLIT, + build_async_topology, +) +from afd_plugin.connectors.base import AFDConnectorBase, ConnectorExtraInfo +from afd_plugin.connectors.gpu.symm_window import ( + FLAG_SHUTDOWN_BIT, + H_ROUTED_TOKENS, + H_SEGMENT_START, + HEADER_FIXED_WORDS, + HEADER_HOST_WORDS, + SlotLayout, + SymmWindow, + encode_header, + encode_header_host_words, +) +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", + "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__}") + + +@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. + 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 + 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, + recv_poll_timeout_ms=coerce_extra_positive_int( + raw.get("recv_poll_timeout_ms", 50), + field_name="recv_poll_timeout_ms", + ), + async_moe_ubatching=ubatching, + async_moe_num_ubatches=num_ubatches, + async_moe_split=coerce_extra_str( + raw.get("async_moe_split", ASYNC_MOE_REQUEST_SPLIT), + field_name="async_moe_split", + ), + ) + + def to_mapping(self) -> dict[str, Any]: + return { + "attn_ranks_per_dp": self.attn_ranks_per_dp, + "ring_depth": self.ring_depth, + "recv_poll_timeout_ms": self.recv_poll_timeout_ms, + "async_moe_ubatching": self.async_moe_ubatching, + "async_moe_num_ubatches": self.async_moe_num_ubatches, + "async_moe_split": self.async_moe_split, + } + + +class ConnectorShutdown(RuntimeError): # noqa: N818 + """Raised on the FFN loop when a peer announced shutdown.""" + + +@dataclass(slots=True) +class GpuAsyncTransferState(AFDTransferState): + """FFN-side state carried from dispatch recv through combine send. + + ``region``/``ring`` locate the window slot so ``send_ffn_work_item_output`` + can write back to the originating Attention rank and release the slot. + + ``group_list`` is a device view of the arrived header's trailing words; + ``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 this rank's own run of partials: + which token each one reads on the way in, and what to weight it by when the + expert output is reduced back to one row per token on the way out. + """ + + region: int = 0 + ring: int = 0 + src_role_rank: int = 0 + layer_idx: int = 0 + stage_idx: int = 0 + seq: int = 0 + num_tokens: int = 0 + routed_tokens: int = 0 + shared_tokens: int = 0 + group_list: Tensor | None = None + expert_counts_host: list[int] = field(default_factory=list) + expand_idx: Tensor | None = None + weights: Tensor | None = None + expand_x_shared: Tensor | None = None + + +@dataclass(slots=True) +class GpuAsyncFFNWorkItem: + """Normalized FFN-side work item produced by a window arrival.""" + + hidden_states: Tensor + context: AFDTransferContext + recv_output: AFDA2FTransferPayload + layer_idx: int + stage_idx: int + num_tokens: int + total_num_tokens: int + shared_num_tokens: int + + +@dataclass(slots=True) +class _PendingDispatch: + """Attention-side record of one in-flight layer, popped by combine recv. + + Every reply carries a full batch of rows, already weighted and summed over + that rank's experts, so combine is a plain add at matching row positions and + needs no index from the dispatch. + + ``shared_slices[r]`` is the contiguous token range whose shared-expert + output that rank returns. It is recorded here because combine must know the + shape of a reply *before* it arrives: that is what lets the wait happen on a + stream instead of on the host, which cannot then be told what turned up. + """ + + context: AFDTransferContext + shared_slices: list[slice] + num_tokens: int + ring: int + 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. + + Every field is indexed by partial -- one entry per ``(token, topk slot)`` -- + and sorted by global expert, so a destination's partials are one contiguous + run, grouped by local expert inside it, which feeds the grouped GEMM + directly. + + Nothing here is ever read back to the host. The three per-destination + vectors go into the slot headers on the device, and the arrays are shipped + whole, so the sender never needs to know how the routing came out. + + Attributes: + counts: Partials per global expert, padded to + ``ffn_size * expert_per_rank``. These are the header's group list. + routed_per_rank: Partials each FFN rank receives. + segment_start: Where each FFN rank's run of partials begins. + expand_idx: Per partial, the token it carries. + weights: Per partial, its topk weight. + """ + + counts: Tensor + routed_per_rank: Tensor + segment_start: Tensor + expand_idx: Tensor + weights: Tensor + + +def plan_dispatch( + topk_ids: Tensor, + topk_weights: Tensor, + *, + ffn_size: int, + expert_per_rank: int, +) -> DispatchPlan: + """Cluster ``(token, topk_slot)`` partials by destination expert. + + Sorting by the global expert id groups partials by destination rank and, in + the same pass, by local expert inside each destination, which is every + grouping any consumer needs. One sort and a bounds lookup is the whole plan. + + Every step is a whole-tensor op and the host learns nothing here, by + design: a readback of the routing would block the send path behind the + device once per MoE layer. + """ + num_slots = topk_ids.shape[1] + flat = topk_ids.reshape(-1).to(torch.int64) + order = torch.argsort(flat, stable=True) + 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 + + # 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) + + # A rank owns a contiguous block of experts, so its partials begin where its + # first expert's do and run to the end of its last. + return DispatchPlan( + counts=counts, + routed_per_rank=counts.view(ffn_size, expert_per_rank).sum(1), + segment_start=offsets.view(ffn_size, expert_per_rank)[:, 0], + expand_idx=expand_idx, + weights=weights, + ) + + +class GpuAsyncAFDConnector(AFDConnectorBase): + """NVSHMEM symmetric-window asynchronous connector for CUDA AFD.""" + + control_plane = None + + @classmethod + def parse_extra_config( + cls, + raw: Mapping[str, Any] | None, + ) -> GpuAsyncExtraInfo: + return GpuAsyncExtraInfo.from_mapping(raw) + + def __init__( + self, + rank: int, + local_rank: int, + vllm_config: VllmConfig, + afd_config: AFDConfig, + role_rank: int, + ) -> None: + super().__init__(rank, local_rank, vllm_config, afd_config, role_rank) + self._initialized = False + hf_config = vllm_config.model_config.hf_config + self.hidden_size = hf_config.hidden_size + self.topk = hf_config.num_experts_per_tok + self.num_routed_experts = hf_config.n_routed_experts + # Combine has to know whether a reply carries shared-expert rows before + # it arrives, and only the model config says so. + self.has_shared_experts = bool(hf_config.n_shared_experts) + self.payload_dtype = vllm_config.model_config.dtype + self.max_seq_len = vllm_config.scheduler_config.max_num_batched_tokens + self.tp_size = self.extra_info.attn_ranks_per_dp + + self.topology = build_async_topology( + afd_config, + role_rank, + num_routed_experts=self.num_routed_experts, + ) + self.world_rank = self.topology.world_rank + self.attn_size = self.topology.attn_size + self.ffn_size = self.topology.ffn_size + self.expert_per_rank = self.topology.expert_per_rank + self.is_attention = afd_config.role == "attention" + + 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. + self.num_regions = max(self.attn_size, self.ffn_size) + # Payload rows are distinct tokens, so a batch is their bound no matter + # how skewed the gate is. Only the 4-byte-per-partial index arrays need + # the every-partial-to-one-rank worst case. + self.token_cap = max(1, self.max_seq_len) + self.partial_cap = max(1, self.max_seq_len * self.topk) + # Shared-expert rows are split across the FFN ranks, so a slot holds a + # fraction of the batch, not all of it -- and none at all when the model + # has no shared experts. Both roles derive this from the same config, so + # the symmetric allocation still matches. + self.shared_cap = ( + -(-self.token_cap // self.ffn_size) if self.has_shared_experts else 0 + ) + self.layout = SlotLayout.build( + expert_per_rank=self.expert_per_rank, + partial_cap=self.partial_cap, + token_cap=self.token_cap, + shared_cap=self.shared_cap, + hidden_size=self.hidden_size, + payload_itemsize=torch.empty(0, dtype=self.payload_dtype).element_size(), + ) + + self._header_device: Tensor | None = None + + 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 _rings_for_stage(self, stage_idx: int) -> list[int]: + """Ring slots this stage owns. + + A ring names a window slot, so stages must not share one: two stages of + the same layer are in flight at the same time, and handing both the same + slot means the second dispatch overwrites the first's payload and flag, + after which the reply the first is waiting for never arrives. + """ + first = stage_idx % self.num_stages + return list(range(first, self.ring_depth, self.num_stages)) + + def init_afd_connector(self) -> None: + """Collectively create the AFD world group and the symmetric window. + + All Attention and FFN ranks must call this with identical rendezvous and + topology settings; the window allocation is symmetric, so a mismatched + size fails here rather than corrupting a later transfer. + """ + if self._initialized: + return + + self.pg = init_afd_process_group( + backend="nccl", + init_method=f"tcp://{self.afd_config.host}:{self.afd_config.port}", + world_size=self.topology.world_size, + rank=self.world_rank, + group_name=AFD_ASYNC_GPU_GROUP_NAME, + timeout=timedelta(minutes=30), + ) + device = torch.device("cuda", self.local_rank) + self._header_device = torch.empty( + (self.ffn_size, self.layout.header_words), + dtype=torch.int32, + device=device, + ) + self.window = SymmWindow( + num_regions=self.num_regions, + ring_depth=self.ring_depth, + layout=self.layout, + payload_dtype=self.payload_dtype, + device=device, + group=self.pg, + rank=self.world_rank, + world_size=self.topology.world_size, + ) + for stage in range(self.num_stages): + self._free_rings[stage] = self._rings_for_stage(stage) + logger.info( + "AFD async GPU window ready: role=%s role_rank=%d world_rank=%d/%d " + "regions=%d rings=%d partial_cap=%d slot=%.1fMiB total=%.1fMiB", + self.afd_config.role, + self.role_rank, + self.world_rank, + self.topology.world_size, + self.num_regions, + self.ring_depth, + self.partial_cap, + self.layout.slot_bytes / 2**20, + self.window.total_bytes / 2**20, + ) + self._initialized = True + + def close(self) -> None: + if self.window is not None: + self.window.close() + self.window = None + if self.pg is not None: + import torch.distributed as dist + + dist.destroy_process_group(self.pg) + self.pg = None + self._pending.clear() + self._free_rings.clear() + self._initialized = False + + def select_experts(self, **kwargs: Any) -> tuple[Tensor, Tensor]: + """Run vLLM's grouped top-k on the Attention side. + + ``compute_gate_topk`` delegates expert selection to the connector so the + CAM and CUDA paths can share one gate; this is the CUDA half. + """ + from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import ( + grouped_topk, + ) + + if kwargs.get("mix_placement"): + raise RuntimeError( + "AFD async GPU connector does not support mix_placement", + ) + return grouped_topk( + hidden_states=kwargs["hidden_states"], + gating_output=kwargs["router_logits"], + topk=kwargs["top_k"], + renormalize=kwargs["renormalize"], + num_expert_group=kwargs.get("num_expert_group", 0), + topk_group=kwargs.get("topk_group", 0), + scoring_func=kwargs.get("scoring_func", "softmax"), + routed_scaling_factor=kwargs.get("routed_scaling_factor", 1.0), + e_score_correction_bias=kwargs.get("e_score_correction_bias"), + ) + + def _require_initialized(self) -> SymmWindow: + if not self._initialized or self.window is None: + raise RuntimeError("AFD async GPU connector is not initialized") + return self.window + + def _shared_slice(self, ffn_rank: int, num_tokens: int) -> slice: + """Token range whose shared-expert output ``ffn_rank`` owns. + + Contiguous chunks, so a rank's slice is a view: no index to build, none + to gather through, and none to put on the wire. Empty when the model has + no shared experts, which is what keeps the payload off the wire and the + field out of the slot. + + The header's ``shared_tokens`` and the rows actually written have to + agree, and they are produced by different callers, so both come from + here rather than from two copies of the arithmetic. + """ + if not self.has_shared_experts: + return slice(0, 0) + return slice( + ffn_rank * num_tokens // self.ffn_size, + (ffn_rank + 1) * num_tokens // self.ffn_size, + ) + + def _headers_for( + 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 = 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, + 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 + # ================================================================== + + 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}", + ) + expected_shape = (num_tokens, self.topk) + if tuple(topk_ids.shape) != expected_shape: + raise ValueError( + f"topk_ids shape must be {expected_shape}, got {tuple(topk_ids.shape)}", + ) + # The weights are flattened alongside the ids to give each partial its + # own weight, so a mismatched shape would silently misalign them. + if tuple(topk_weights.shape) != expected_shape: + raise ValueError( + f"topk_weights shape must be {expected_shape}, " + f"got {tuple(topk_weights.shape)}", + ) + + stage_idx = metadata.stage_idx + rings = self._free_rings.setdefault( + stage_idx, + self._rings_for_stage(stage_idx), + ) + if not rings: + raise RuntimeError( + f"AFD async GPU ring exhausted on stage {stage_idx}; the " + "send-then-recv invariant was violated or the topology config " + "does not match the actual peer count", + ) + ring = rings.pop(0) + self._seq += 1 + + plan = plan_dispatch( + topk_ids, + topk_weights, + ffn_size=self.ffn_size, + expert_per_rank=self.expert_per_rank, + ) + # Every FFN rank gets a slot even when routing sends it nothing, and it + # replies to every slot, so a reply is expected from all of them. + # Expecting only the ranks that received data leaves the empty rank's + # reply unmatched and its ring slot never released -- which is what a + # single-token decode hits, since both the routed segment and the + # round-robin shared slice can come out empty for one rank. + expected_ffn = list(range(self.ffn_size)) + shared_slices: list[slice] = [] + headers = self._headers_for( + 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): + # Round-robin needed an arange, a gather and a whole slot field per + # peer per layer to achieve the balance this contiguous split gets + # from a view. + shared = self._shared_slice(ffn_rank, num_tokens) + shared_slices.append(shared) + # Everything but the shared slice goes out whole. The index arrays + # cost a fraction of a percent of the slot, and the payload rows a + # destination does not need are the few tokens none of whose topk + # slots landed on it -- 1.6% of them at 2A2F. Sizing either to the + # routing is what used to make the host wait for the device here. + window.write_slot( + peer=self.attn_size + ffn_rank, + region=self.role_rank, + ring=ring, + header=headers[ffn_rank], + expand_idx=plan.expand_idx, + weights=plan.weights, + routed_x=hidden_states, + shared_x=hidden_states[shared], + flag_value=self._seq, + ) + + logger.debug( + "AFD dispatch sent: A%d layer=%d stage=%d tokens=%d ring=%d " + "partials=%d awaiting_ffn=%s", + self.role_rank, + metadata.layer_idx, + stage_idx, + num_tokens, + ring, + num_tokens * self.topk, + expected_ffn, + ) + self._pending.setdefault(stage_idx, []).append( + _PendingDispatch( + context=context, + shared_slices=shared_slices, + num_tokens=num_tokens, + ring=ring, + seq=self._seq, + expected_ffn=expected_ffn, + ), + ) + + def recv_ffn_output( + self, + ref_tensor: Tensor, + ubatch_idx: int = 0, + **kwargs: Any, + ) -> Tensor: + """Queue this layer's combine and return, without waiting for the data. + + The waiting happens on the stream: every reply is answered into a known + slot, carries a known number of rows, and stamps the dispatch sequence + it answers, so the whole combine can be enqueued before any of it has + arrived. The host goes straight on to the next layer, which is the point + -- polling for the reply here left the GPU with nothing queued behind + the wait, and a profile found 86% of kernel launches executing within + 5us of being issued because of it. + + Nothing reports what turned up, so there is no per-arrival header check + any more. The stream wait replaces it: it blocks on one specific slot + 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) + if not queue: + raise RuntimeError( + f"AFD async GPU recv_ffn_output has no pending dispatch on " + f"stage {ubatch_idx}", + ) + pending = queue.pop(0) + + # Accumulate in the payload dtype. Each row takes at most one + # contribution per FFN rank plus its shared one -- the topk sum already + # happened on the FFN side -- so there is little left for a wider + # accumulator to protect, and float32 cost a widening pass over every + # arriving block plus a narrowing one on the way out. + accumulator = torch.zeros( + (pending.num_tokens, self.hidden_size), + dtype=self.payload_dtype, + device=ref_tensor.device, + ) + for ffn_rank in pending.expected_ffn: + # An FFN rank replies into the region it owns, on the ring the + # dispatch used. + window.stream_wait(ffn_rank, pending.ring, pending.seq) + # A reply is a whole batch, already weighted and summed over that + # rank's experts, with a zero row wherever the rank held none of a + # token's experts. Row i answers token i, so this is a plain add: + # the scatter it replaces was the second largest kernel on the rank. + accumulator += window.local_routed( + ffn_rank, + pending.ring, + pending.num_tokens, + ) + shared = pending.shared_slices[ffn_rank] + shared_tokens = shared.stop - shared.start + if self.has_shared_experts and shared_tokens: + accumulator[shared] += window.local_shared( + ffn_rank, + pending.ring, + shared_tokens, + ) + 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 + + # ================================================================== + # 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. + + 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)) + 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: + 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", + ) + + expand_idx = window.local_expand_idx( + arrived.region, + arrived.ring, + header.routed_tokens, + header.segment_start, + ).to(torch.int64) + states = GpuAsyncTransferState( + region=arrived.region, + ring=arrived.ring, + 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=window.local_expert_counts(arrived.region, arrived.ring), + expert_counts_host=header.expert_counts, + expand_idx=expand_idx, + weights=window.local_weights( + arrived.region, + arrived.ring, + header.routed_tokens, + header.segment_start, + ), + expand_x_shared=window.local_shared( + arrived.region, + arrived.ring, + 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.num_tokens, + ).index_select(0, expand_idx), + context=AFDTransferContext(metadata=metadata, states=states), + ) + + def send_ffn_output( + self, + ffn_output: Tensor, + context: AFDTransferContext, + **kwargs: Any, + ) -> None: + """Reduce expert output back to one row per token and write it back. + + Every partial of a token that landed on this rank is weighted and summed + here, into the token's own row of a full batch. Tokens this rank held no + expert for keep their zero row, which is what lets the Attention side + add replies together without an index. Doing the reduction on this side + keeps the duplicates off the wire. + + Both the weighting and the sum stay in the payload dtype. Widening to + float32 first cost two extra passes over ``[partials, hidden]`` and made + the scatter move twice the bytes, which a profile showed costing almost + as much GPU time as the expert GEMM itself -- to protect a sum of at + most ``topk`` terms whose result goes on the wire narrowed anyway. + """ + window = self._require_initialized() + states = context.states + if not isinstance(states, GpuAsyncTransferState): + raise RuntimeError( + "AFD async GPU send_ffn_output requires GpuAsyncTransferState", + ) + if states.expand_idx is None or states.weights is None: + raise RuntimeError( + "AFD async GPU send_ffn_output requires the dispatch expansion", + ) + reduced = torch.zeros( + (states.num_tokens, self.hidden_size), + dtype=self.payload_dtype, + device=ffn_output.device, + ) + if states.routed_tokens: + weighted = ffn_output * states.weights.unsqueeze(1).to(ffn_output.dtype) + reduced.index_add_(0, states.expand_idx, weighted) + + shared_output: Tensor | None = kwargs.get("shared_output") + 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=states.expert_counts_host, + 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, + expand_idx=None, + weights=None, + routed_x=reduced, + 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, + ) + + # ================================================================== + # 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: + # Header only: the shutdown bit is the whole message, so every + # payload field is empty and write_slot skips it. + window.write_slot( + peer=peer, + region=self.role_rank, + ring=0, + header=header, + expand_idx=None, + weights=None, + routed_x=None, + shared_x=None, + ) + + +__all__ = [ + "AFD_ASYNC_GPU_GROUP_NAME", + "ConnectorShutdown", + "DispatchPlan", + "GpuAsyncAFDConnector", + "GpuAsyncExtraInfo", + "GpuAsyncFFNWorkItem", + "GpuAsyncTransferState", + "plan_dispatch", +] 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/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..5a4330a9 --- /dev/null +++ b/afd_plugin/connectors/gpu/symm_window.py @@ -0,0 +1,748 @@ +# 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. 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. + +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. + +``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 +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 + +import time +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 + +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 = 3 +_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_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 +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 + partial_cap: int + token_cap: int + shared_cap: int + hidden_size: int + payload_itemsize: int + + header_off: int + expand_idx_off: int + weights_off: int + routed_x_off: int + shared_x_off: int + slot_bytes: int + + @classmethod + def build( + cls, + *, + expert_per_rank: int, + partial_cap: int, + token_cap: int, + shared_cap: int, + hidden_size: int, + payload_itemsize: int, + ) -> SlotLayout: + header_words = HEADER_FIXED_WORDS + expert_per_rank + header_off = 0 + 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_x_off = _align(routed_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, + expand_idx_off=expand_idx_off, + weights_off=weights_off, + routed_x_off=routed_x_off, + shared_x_off=shared_x_off, + slot_bytes=slot_bytes, + ) + + +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, + *, + 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], + segment_start: int = 0, + echo_seq: int = 0, +) -> torch.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)}", + ) + header = np.empty(layout.header_words, dtype=np.int32) + 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) + + +@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 + segment_start: 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], + segment_start=values[H_SEGMENT_START], + expert_counts=values[HEADER_FIXED_WORDS:], + ) + + +@dataclass(slots=True) +class ArrivedSlot: + """One arrival found by ``SymmWindow.poll``.""" + + region: int + ring: int + 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'.""" + + 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 + # 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( + 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.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.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}") + + 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, + expand_idx: torch.Tensor | None, + weights: torch.Tensor | None, + routed_x: torch.Tensor | None, + 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. + + 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. + + ``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. + + ``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) + shared_count = _row_count(shared_x, shared_rows) + 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"partials {partial_count} exceed partial_cap {layout.partial_cap}", + ) + if shared_count > layout.shared_cap: + raise RuntimeError( + f"shared tokens {shared_count} exceed shared_cap {layout.shared_cap}", + ) + + 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, + 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, + 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.expand_idx_off, + (), + torch.int32, + expand_idx, + partial_count, + ) + write_field( + layout.weights_off, + (), + torch.float32, + weights, + _row_count(weights, None), + ) + write_field( + layout.routed_x_off, + (layout.hidden_size,), + self.payload_dtype, + routed_x, + routed_count, + routed_rows, + ) + write_field( + layout.shared_x_off, + (layout.hidden_size,), + self.payload_dtype, + shared_x, + shared_count, + shared_rows, + ) + + 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) + + # ------------------------------------------------------------------ + # 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 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. + + 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) + 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, 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_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_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, + (start + count,), + torch.int32, + )[start:] + + 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, + (start + count,), + torch.float32, + )[start:] + + 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(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..a26321d9 --- /dev/null +++ b/afd_plugin/model_executor/models/gpu/deepseek_v2_attention_gate.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Attention-side gate helpers for DeepSeek-V2 on CUDA. + +The async GPU connector dispatches tokens that are already routed: one row per +``(token, topk_slot)`` partial, grouped by local expert, with a ``group_list`` +of per-expert counts. The FFN side therefore must not run routing again -- it +only needs the grouped GEMM over its local experts. + +That shape is a ``topk == 1`` problem: give every arriving row its own expert id +and its own weight, and vLLM's ``fused_experts`` computes exactly the local +expert output, already scaled, in the epilogue it runs anyway. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts + +from afd_plugin.connectors.metadata import AFDF2ATransferPayload + +if TYPE_CHECKING: + from torch import nn + + +def compute_attention_gate_moe_ffn( + layer: nn.Module, + *, + hidden_states: torch.Tensor, + group_list: torch.Tensor, + expand_x_shared: torch.Tensor | None = None, +) -> AFDF2ATransferPayload: + """Run this rank's local experts over pre-routed tokens. + + Args: + layer: The AFD DeepSeek decoder layer owning the MoE module. + hidden_states: ``[num_partials, hidden]`` rows sorted by local expert. + group_list: ``[expert_per_rank]`` per-expert row counts. Must sum to + ``hidden_states.shape[0]``. + expand_x_shared: Optional ``[num_shared, hidden]`` shared-expert rows. + + Returns: + Routed output in the same row order as ``hidden_states``, plus the + shared-expert output when the model has shared experts. + """ + # ``mlp.experts`` is a MoERunner; the weight parameters live on its + # RoutedExperts, and the shared experts hang off the runner rather than + # off the MoE module. + runner = layer.mlp.experts + routed_experts = runner.routed_experts + counts = group_list.to(torch.int64) + num_rows = int(hidden_states.shape[0]) + num_local_experts = counts.numel() + if num_rows == 0: + # Routing can leave a peer with nothing -- common in decode, where a + # single token's topk may land entirely on the other FFN rank. + routed_output = hidden_states.new_empty((0, hidden_states.shape[-1])) + shared = runner._shared_experts + return AFDF2ATransferPayload( + routed_output=routed_output, + shared_output=( + shared._layer(expand_x_shared) + if shared is not None + and expand_x_shared is not None + and expand_x_shared.shape[0] > 0 + else None + ), + ) + # ``output_size`` keeps this off the device: without it repeat_interleave + # reads the counts back to the host to size its output, which is a + # synchronize on every work item. It also enforces what the removed + # ``counts.sum() == num_rows`` check used to, raising if they disagree. + expert_ids = torch.repeat_interleave( + torch.arange( + num_local_experts, + device=hidden_states.device, + dtype=torch.int32, + ), + counts, + output_size=num_rows, + ).unsqueeze(1) + # Mirrors the NPU gate path: scale the routed branch unless fp16, where the + # native model instead scales the shared branch down. ``fused_experts`` + # multiplies every row by its topk weight in an epilogue it runs anyway, so + # feeding the factor in there costs nothing, where scaling its output cost a + # full pass over ``[num_partials, hidden]``. The topk weighting itself is + # not applied here -- the connector owns it, and applies it when it reduces + # the partials back to one row per token. + routed_scaling_factor = runner.routed_scaling_factor + scale_shared_instead = hidden_states.dtype == torch.float16 + row_weights = torch.full( + (num_rows, 1), + 1.0 if scale_shared_instead else routed_scaling_factor, + dtype=torch.float32, + device=hidden_states.device, + ) + + routed_output = fused_experts( + hidden_states, + routed_experts.w13_weight, + routed_experts.w2_weight, + row_weights, + expert_ids, + global_num_experts=num_local_experts, + expert_map=None, + ) + + shared_output = None + shared_experts = runner._shared_experts + 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) + + if scale_shared_instead and shared_output is not None: + shared_output = shared_output * (1.0 / routed_scaling_factor) + + return AFDF2ATransferPayload( + routed_output=routed_output, + shared_output=shared_output, + ) + + +__all__ = ["compute_attention_gate_moe_ffn"] diff --git a/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py b/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py index 383ea56e..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 @@ -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 " @@ -327,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, @@ -503,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 b8ed0d9d..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,41 @@ 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 +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): @@ -56,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, @@ -76,16 +125,27 @@ 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 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 @@ -126,9 +186,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 = { @@ -258,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 @@ -267,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 @@ -300,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, @@ -324,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 @@ -349,25 +563,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, @@ -509,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", @@ -534,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 @@ -667,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/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/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/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 diff --git a/tests/e2e/async_gpu_connector_e2e.py b/tests/e2e/async_gpu_connector_e2e.py new file mode 100644 index 00000000..85b23c46 --- /dev/null +++ b/tests/e2e/async_gpu_connector_e2e.py @@ -0,0 +1,243 @@ +"""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, + # This test's FFN loop runs routed experts only. + n_shared_experts=0, + ), + dtype=torch.bfloat16, + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=NUM_TOKENS), + additional_config={ + "afd": { + "role": role, + "connector": "GpuAsyncAFDConnector", + "async": True, + "compute_gate_on_attention": True, + "num_attention_ranks": 1, + "num_ffn_ranks": 1, + "port": PORT, + "connector_extra_config": {"ring_depth": 1}, + }, + }, + ) + afd_config = AFDConfig( + role=role, + connector="GpuAsyncAFDConnector", + async_dp=True, + compute_gate_on_attention=True, + num_attention_ranks=1, + num_ffn_ranks=1, + host="127.0.0.1", + port=PORT, + ) + connector = GpuAsyncAFDConnector( + rank=local_rank, + local_rank=local_rank, + vllm_config=vllm_config, + afd_config=afd_config, + role_rank=0, + ) + connector.init_afd_connector() + return connector + + +def make_weights(device): + generator = torch.Generator(device="cpu").manual_seed(11) + w13 = ( + torch.randn(NUM_EXPERTS, 2 * INTERMEDIATE, HIDDEN, generator=generator) + / HIDDEN**0.5 + ).to(device, torch.bfloat16) + w2 = ( + torch.randn(NUM_EXPERTS, HIDDEN, INTERMEDIATE, generator=generator) + / INTERMEDIATE**0.5 + ).to(device, torch.bfloat16) + return w13, w2 + + +def make_layer_inputs(layer_idx, device): + gen = torch.Generator(device="cpu").manual_seed(100 + layer_idx) + x = torch.randn(NUM_TOKENS, HIDDEN, generator=gen).to(device, torch.bfloat16) + topk_ids = torch.stack( + [torch.randperm(NUM_EXPERTS, generator=gen)[:TOPK] for _ in range(NUM_TOKENS)], + ).to(device, torch.int32) + topk_weights = torch.rand(NUM_TOKENS, TOPK, generator=gen).to(device) + return x, topk_ids, topk_weights + + +def reference_moe(x, w13, w2, topk_ids, topk_weights): + out = torch.zeros(x.shape[0], HIDDEN, dtype=torch.float32, device=x.device) + for token in range(x.shape[0]): + for slot in range(topk_ids.shape[1]): + expert = int(topk_ids[token, slot]) + hidden = x[token].to(torch.float32) @ w13[expert].to(torch.float32).T + gate, up = hidden.chunk(2, dim=-1) + y = (F.silu(gate) * up) @ w2[expert].to(torch.float32).T + out[token] += float(topk_weights[token, slot]) * y * SCALING + return out + + +def init_world(rank: int) -> None: + """Mimic a single `vllm serve`: a private default group of size 1. + + The connector bootstraps NVSHMEM on the AFD group itself, so the default + group deliberately does *not* span both roles -- that is the topology the + real deployment has. + """ + dist.init_process_group( + "nccl", + init_method=f"tcp://127.0.0.1:{WORLD_PORT + rank}", + world_size=1, + rank=0, + ) + + +def run_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..c6d96eb6 --- /dev/null +++ b/tests/e2e/async_gpu_window_roundtrip.py @@ -0,0 +1,220 @@ +"""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, + # 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, + ) + 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..b9e2743e --- /dev/null +++ b/tests/unit/connectors/test_async_gpu_connector.py @@ -0,0 +1,428 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Unit tests for the async GPU connector's wire format and routing math.""" + +from __future__ import annotations + +import inspect +from itertools import pairwise +from types import SimpleNamespace + +import pytest + +pytest.importorskip("torch") + +import torch # noqa: E402 + +from afd_plugin.connectors.factory import AFDConnectorFactory # noqa: E402 +from afd_plugin.connectors.gpu.async_gpu import ( # noqa: E402 + GpuAsyncAFDConnector, + GpuAsyncExtraInfo, + plan_dispatch, +) +from afd_plugin.connectors.gpu.symm_window import ( # noqa: E402 + FLAG_SHUTDOWN_BIT, + HEADER_FIXED_WORDS, + SlotLayout, + SymmWindow, + decode_header, + encode_header, +) + + +@pytest.fixture +def layout() -> SlotLayout: + # shared_cap is deliberately not token_cap: the shared field is sized by the + # per-rank split, so a layout that quietly reused token_cap for it would + # otherwise still satisfy every assertion below. + return SlotLayout.build( + expert_per_rank=4, + partial_cap=100, + token_cap=32, + shared_cap=8, + hidden_size=8, + payload_itemsize=2, + ) + + +# ---------------------------------------------------------------------- +# Config +# ---------------------------------------------------------------------- + + +def test_connector_is_registered_and_has_no_control_plane(): + connector_cls = AFDConnectorFactory.get_connector_class("GpuAsyncAFDConnector") + assert connector_cls is GpuAsyncAFDConnector + assert connector_cls.control_plane is None + + +def test_ring_depth_defaults_to_the_number_of_live_stages(): + assert GpuAsyncExtraInfo.from_mapping(None).ring_depth == 1 + assert GpuAsyncExtraInfo.from_mapping({"async_moe_ubatching": True}).ring_depth == 2 + assert GpuAsyncExtraInfo.from_mapping({"ring_depth": 4}).ring_depth == 4 + + +def test_unknown_extra_config_field_is_rejected(): + with pytest.raises(ValueError, match="unknown AFD async GPU"): + GpuAsyncExtraInfo.from_mapping({"nope": 1}) + + +# ---------------------------------------------------------------------- +# 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.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, 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 + # 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_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. + for offset in ( + layout.header_off, + layout.expand_idx_off, + layout.weights_off, + layout.routed_x_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], + segment_start=25, + ) + 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 decoded.segment_start == 25 + 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 _destination_slices(plan, ffn_size, expert_per_rank): + """Walk the plan the way an FFN rank does: one run of partials each. + + A destination is told where its run starts and how long it is, and reads the + index arrays the sender shipped whole -- there is no per-destination slicing + on the send side any more, which is what removed the readback. + """ + routed = plan.routed_per_rank.tolist() + starts = plan.segment_start.tolist() + for ffn_rank in range(ffn_size): + yield ffn_rank, slice(starts[ffn_rank], starts[ffn_rank] + routed[ffn_rank]) + + +def test_every_partial_is_routed_exactly_once(routing_inputs): + topk_ids, _, topk_weights = routing_inputs + plan = plan_dispatch( + topk_ids, + topk_weights, + ffn_size=_FFN_SIZE, + expert_per_rank=_EXPERT_PER_RANK, + ) + assert plan.expand_idx.shape == (_NUM_TOKENS * _TOPK,) + assert plan.weights.shape == (_NUM_TOKENS * _TOPK,) + assert int(plan.counts.sum()) == _NUM_TOKENS * _TOPK + assert int(plan.routed_per_rank.sum()) == _NUM_TOKENS * _TOPK + # The runs must tile the array end to end, or a partial is read twice or not + # at all: they are the only thing a destination gets to locate itself by. + cursor = 0 + for _, partials in _destination_slices(plan, _FFN_SIZE, _EXPERT_PER_RANK): + assert partials.start == cursor + cursor = partials.stop + assert cursor == _NUM_TOKENS * _TOPK + + +def test_each_destination_segment_is_grouped_by_local_expert(routing_inputs): + topk_ids, _, topk_weights = routing_inputs + plan = plan_dispatch( + topk_ids, + topk_weights, + ffn_size=_FFN_SIZE, + expert_per_rank=_EXPERT_PER_RANK, + ) + counts = plan.counts.tolist() + for ffn_rank, partials in _destination_slices( + plan, + _FFN_SIZE, + _EXPERT_PER_RANK, + ): + base = ffn_rank * _EXPERT_PER_RANK + # The token behind each partial, in the order the receiver sees them. + tokens = plan.expand_idx[partials] + cursor = 0 + for local_expert in range(_EXPERT_PER_RANK): + for _ in range(counts[base + local_expert]): + token_idx = int(tokens[cursor]) + assert base + local_expert in topk_ids[token_idx].tolist() + cursor += 1 + assert cursor == partials.stop - partials.start + + +def test_every_partial_names_a_token_of_this_batch(routing_inputs): + topk_ids, _, topk_weights = routing_inputs + plan = plan_dispatch( + topk_ids, + topk_weights, + ffn_size=_FFN_SIZE, + expert_per_rank=_EXPERT_PER_RANK, + ) + # Destinations read the whole batch out of the slot, so an index is only in + # range if it names a row of it. + assert int(plan.expand_idx.min()) >= 0 + assert int(plan.expand_idx.max()) < _NUM_TOKENS + for ffn_rank, partials in _destination_slices(plan, _FFN_SIZE, _EXPERT_PER_RANK): + base = ffn_rank * _EXPERT_PER_RANK + expected = { + token + for token in range(_NUM_TOKENS) + for expert in topk_ids[token].tolist() + if base <= expert < base + _EXPERT_PER_RANK + } + assert set(plan.expand_idx[partials].tolist()) == expected + + +def test_identity_experts_recombine_to_the_weighted_sum(routing_inputs): + """The full chain: ship the batch, expand, weight, reduce, add.""" + topk_ids, hidden_states, topk_weights = routing_inputs + plan = plan_dispatch( + topk_ids, + topk_weights, + ffn_size=_FFN_SIZE, + expert_per_rank=_EXPERT_PER_RANK, + ) + accumulator = torch.zeros(_NUM_TOKENS, _HIDDEN, dtype=torch.float32) + for _, partials in _destination_slices(plan, _FFN_SIZE, _EXPERT_PER_RANK): + expand = plan.expand_idx[partials].to(torch.int64) + # What the FFN side does with the batch it was sent. + expanded = hidden_states.index_select(0, expand) + reduced = torch.zeros(_NUM_TOKENS, _HIDDEN, dtype=torch.float32) + reduced.index_add_(0, expand, expanded * plan.weights[partials].unsqueeze(1)) + # A reply is a whole batch, so combine adds it without an index. + accumulator += reduced + expected = hidden_states * topk_weights.sum(dim=1, keepdim=True) + torch.testing.assert_close(accumulator, expected.to(torch.float32)) + + +def test_routing_handles_experts_not_divisible_by_ffn_size(): + # expert_per_rank is a ceiling division, so the padded tail must stay empty + # instead of silently absorbing real partials. + ffn_size, expert_per_rank, num_experts = 3, 2, 5 + topk_ids = torch.tensor([[0, 4], [1, 3], [2, 4]], dtype=torch.int32) + plan = plan_dispatch( + topk_ids, + torch.ones(3, 2), + ffn_size=ffn_size, + expert_per_rank=expert_per_rank, + ) + assert plan.counts.numel() == ffn_size * expert_per_rank + assert int(plan.counts.sum()) == topk_ids.numel() + assert int(plan.counts[num_experts:].sum()) == 0 + + +def test_routing_can_leave_one_destination_empty(): + """A single decode token's topk can land entirely on one FFN rank. + + The peer that gets nothing must still see a well-formed, empty segment -- + zero-length windows are what crashed a 2A2F decode step. + """ + ffn_size, expert_per_rank = 2, 4 + # Every partial targets experts owned by FFN rank 0. + topk_ids = torch.tensor([[0, 1, 2]], dtype=torch.int32) + plan = plan_dispatch( + topk_ids, + torch.ones(1, 3), + ffn_size=ffn_size, + expert_per_rank=expert_per_rank, + ) + slices = list(_destination_slices(plan, ffn_size, expert_per_rank)) + + _, busy_partials = slices[0] + assert busy_partials.stop - busy_partials.start == 3 + # One token, three of its partials: it is sent once, read three times. + assert plan.expand_idx[busy_partials].tolist() == [0, 0, 0] + + _, empty_partials = slices[1] + assert empty_partials.stop - empty_partials.start == 0 + + # Reducing an empty destination must be a no-op, not an error. + accumulator = torch.zeros(1, 4, dtype=torch.float32) + empty = plan.expand_idx[empty_partials].to(torch.int64) + accumulator.index_add_(0, empty, torch.zeros(0, 4)) + assert torch.count_nonzero(accumulator) == 0 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..9c18cbfb 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 @@ -997,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): @@ -1098,3 +1109,35 @@ 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_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")) 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):