From b2b5053db53c13db5c41bdb693ed0acbfddfd67b Mon Sep 17 00:00:00 2001 From: swjeong9 Date: Fri, 4 Sep 2026 16:02:01 +0900 Subject: [PATCH] [Feature] Support arbitrary attention-to-FFN rank ratios in P2pNcclAFDConnector `validate_p2p_topology` rejected every layout except `A >= F` with `A % F == 0`, so `A < F` and non-divisible ratios could not be expressed. Generalize the subgroup partition to a balanced block distribution over `G = min(A, F)` groups: attention rank `a` joins group `a * G // A` and FFN rank `f` joins group `f * G // F`. Both maps are surjective, so no group is left without a peer. The previously supported layouts are the `G == F` case of this formula and produce the same groups, so their rank mapping and execution path are unchanged. On the transport side a group is always one member of the smaller role plus several of the larger one. For `A >= F` this is the existing gather. For `A < F` the attention rank splits its tokens across its FFN members and concatenates the results back. Split sizes come from `split_send_sizes`, a closed form that avoids branching on the remainder so the arithmetic survives `torch.compile` tracing with a symbolic token dimension. Slices are passed to the custom ops as (start, size) pairs rather than as views, because a view with a non-zero storage offset reached the op at offset 0 in generated inductor code. An FFN rank whose share is zero still runs a one-token dummy batch so it can take part in the expert all-to-all inside the FFN role. This also fixes subgroup rendezvous for multi-node deployments. Every subgroup used the configured `afd.host` for its `StatelessProcessGroup`, but the socket is opened by that group's first member, so a second subgroup whose first member lives on another node could never be reached. Each rank now exchanges its own address once over the existing AFD group, and each subgroup dials its first member. On a single host the resolved address equals the configured host, leaving existing recipes unchanged. The GPU results for this change were all collected across several nodes and are not reproducible without it. vLLM sources are untouched and no compat patch is added. Signed-off-by: swjeong9 --- afd_plugin/connectors/gpu/p2p.py | 484 +++++++++++----- afd_plugin/distributed/__init__.py | 2 + afd_plugin/distributed/topology.py | 76 ++- afd_plugin/v1/worker/ffn_metadata.py | 52 +- .../unit/connectors/test_p2p_compiled_send.py | 198 +++++++ tests/unit/connectors/test_p2p_connector.py | 54 +- .../test_p2p_generalized_topology.py | 529 ++++++++++++++++++ .../connectors/test_p2p_multinode_address.py | 44 ++ .../distributed/test_topology_partition.py | 222 ++++++++ .../distributed/test_topology_snapshot.py | 299 ++++++++++ tests/unit/v1/worker/test_ffn_metadata.py | 48 ++ 11 files changed, 1817 insertions(+), 191 deletions(-) create mode 100644 tests/unit/connectors/test_p2p_compiled_send.py create mode 100644 tests/unit/connectors/test_p2p_generalized_topology.py create mode 100644 tests/unit/connectors/test_p2p_multinode_address.py create mode 100644 tests/unit/distributed/test_topology_partition.py create mode 100644 tests/unit/distributed/test_topology_snapshot.py create mode 100644 tests/unit/v1/worker/test_ffn_metadata.py diff --git a/afd_plugin/connectors/gpu/p2p.py b/afd_plugin/connectors/gpu/p2p.py index 5534859d..892b483d 100644 --- a/afd_plugin/connectors/gpu/p2p.py +++ b/afd_plugin/connectors/gpu/p2p.py @@ -11,16 +11,21 @@ Topology: The connector creates one AFD NCCL world ordered as ``[F0, F1, ..., A0, A1, ...]``: FFN ranks first, followed by Attention - ranks. Each FFN rank owns one subgroup containing itself and one or more - consecutive Attention ranks, which requires:: - - num_attention_ranks >= num_ffn_ranks - num_attention_ranks % num_ffn_ranks == 0 - - Attention sends hidden states to its mapped FFN rank; the FFN rank - concatenates inputs from its Attention peers, runs FFN work, splits the - output by the recorded sequence lengths, and sends each slice back to the - originating Attention rank. + ranks. The world is partitioned into ``G = min(A, F)`` subgroups by a + balanced block distribution: Attention ``a`` joins subgroup ``a*G//A`` + and FFN ``f`` joins subgroup ``f*G//F``. Every positive ``(A, F)`` pair + is valid, and the role with fewer ranks contributes exactly one member + per subgroup: + + - ``A >= F``: one FFN rank and one or more Attention ranks. Attention + sends its hidden states whole; the FFN rank concatenates the inputs + from its Attention peers, runs FFN work, splits the output by the + recorded sequence lengths, and sends each slice back to its origin. + - ``A < F``: one Attention rank and several FFN ranks. Attention splits + its tokens across the FFN members with ``split_send_sizes``; each FFN + returns its share, and Attention reassembles them in member order. An + FFN that receives no tokens runs its forward on a one-token dummy + batch and returns a zero-size tensor. Control and data planes: DP metadata handling is a pluggable control plane, not part of the @@ -61,6 +66,7 @@ from __future__ import annotations +import socket from collections.abc import Callable, Mapping from datetime import timedelta from typing import TYPE_CHECKING, Any, NamedTuple @@ -92,6 +98,7 @@ DefaultProcessGroupSwitcher, build_rank_mapping, init_afd_process_group, + split_send_sizes, ) if TYPE_CHECKING: @@ -115,9 +122,11 @@ class P2pNcclAFDConnector(AFDConnectorBase): """NCCL-backed Attention <-> FFN connector for CUDA deployments. The P2P topology places FFN ranks before Attention ranks in the AFD world - (``[F0, F1, ..., A0, A1, ...]``), and each FFN rank owns a subgroup with - one or more consecutive Attention ranks. Within a subgroup, the FFN rank - is subgroup rank ``0`` and its Attention peers occupy ranks ``1..ratio``. + (``[F0, F1, ..., A0, A1, ...]``) and partitions it into ``min(A, F)`` + subgroups (see the module docstring). Within a subgroup the FFN members + occupy ranks ``0..subgroup_ffn_count-1`` and the Attention members + follow; historically that is FFN rank ``0`` plus Attention ranks + ``1..ratio``. Hidden states move through per-subgroup ``PyNcclCommunicator`` instances on the current CUDA stream; a separate NCCL process group distributes the @@ -181,6 +190,9 @@ def __init__( self.min_size = self.mapping.min_size self.ratio = self.mapping.ratio self.group_size = len(self.mapping.subgroup_ranks) + self.subgroup_ffn_count = sum( + 1 for rank in self.mapping.subgroup_ranks if rank < self.ffn_size + ) self.dst_list = list(self.mapping.dp_metadata_destinations) text_config = vllm_config.model_config.hf_text_config self.num_hidden_layers = (text_config.num_hidden_layers,) @@ -193,6 +205,8 @@ def __init__( tuple[int, int], _TensorMetadata, ] = {} + # Stages of this step whose incoming shares are all zero. + self._dummy_stages: set[int] = set() self._recv_attn_buffers: dict[ tuple[int, int, tuple[int, ...]], torch.Tensor, @@ -242,9 +256,10 @@ def init_afd_connector(self) -> None: 1. Joins the AFD world process group (FFN ranks first, then Attention ranks) rendezvoused at ``tcp://host:port``. 2. Creates this rank's subgroup ``StatelessProcessGroup`` on - ``port + subgroup_index + 1`` and two ``PyNcclCommunicator`` - instances over it (Attention-to-FFN and FFN-to-Attention), each - registered for use by the P2P custom ops. + ``port + subgroup_index + 1``, at the address of the subgroup's + first member, and two ``PyNcclCommunicator`` instances over it + (Attention-to-FFN and FFN-to-Attention), each registered for use + by the P2P custom ops. 3. On ranks that participate in the DP metadata control plane, joins the ``p2p`` process group that ``control_plane`` uses to distribute per-stage token counts. @@ -265,10 +280,20 @@ def init_afd_connector(self) -> None: timeout=timedelta(minutes=2), ) + # A subgroup's group is served by its first member, which can only + # bind its own node's address. + addresses = _gather_rank_addresses( + afd_pg, + host=self.afd_config.host, + port=self.afd_config.port, + world_size=self.ffn_size + self.attn_size, + ) + subgroup_host = addresses[self.mapping.subgroup_ranks[0]] + with DefaultProcessGroupSwitcher(_get_default_group(), afd_pg): base_port = self.afd_config.port self.a2e_group = StatelessProcessGroup.create( - host=self.afd_config.host, + host=subgroup_host, port=base_port + self.mapping.subgroup_index + 1, rank=self.mapping.rank_in_subgroup, world_size=len(self.mapping.subgroup_ranks), @@ -308,7 +333,11 @@ def send_attn_output( context: AFDTransferContext, **kwargs: Any, ) -> None: - """Send Attention hidden states to this rank's mapped FFN rank. + """Send Attention hidden states to this rank's subgroup FFN member(s). + + Each subgroup FFN member receives one contiguous share of this + rank's tokens, sized by ``split_send_sizes``. A single FFN member + receives the whole tensor. Args: hidden_states: CUDA tensor of shape ``(num_tokens, hidden_size)`` @@ -350,26 +379,27 @@ def send_attn_output( ) if input_ids.dtype != _INPUT_IDS_DTYPE: raise ValueError("P2P AFD input_ids must use torch.int32") - self._send_hidden_states( - hidden_states, - 0, - self.a2e_group, - self.a2e_comm_id, - ) - if router_logits is not None: - self._send_hidden_states( - router_logits, - 0, - self.a2e_group, - self.a2e_comm_id, - ) - if input_ids is not None: - self._send_hidden_states( - input_ids, - 0, - self.a2e_group, - self.a2e_comm_id, - ) + if self.subgroup_ffn_count == 1: + self._send_hidden_states(hidden_states, 0, self.a2e_group, self.a2e_comm_id) + if router_logits is not None: + self._send_hidden_states( + router_logits, 0, self.a2e_group, self.a2e_comm_id + ) + if input_ids is not None: + self._send_hidden_states(input_ids, 0, self.a2e_group, self.a2e_comm_id) + return + # Traced by torch.compile: pass the bounds to the op rather than + # slicing here, and send every share (see _register_p2p_custom_ops). + sizes = split_send_sizes(hidden_states.shape[0], self.subgroup_ffn_count) + start = 0 + for dst, size in enumerate(sizes): + # Wire order per destination, matching recv_attn_output. + self._send_slice(hidden_states, start, size, dst) + if router_logits is not None: + self._send_slice(router_logits, start, size, dst) + if input_ids is not None: + self._send_slice(input_ids, start, size, dst) + start = start + size def recv_ffn_output( self, @@ -377,13 +407,17 @@ def recv_ffn_output( ubatch_idx: int = 0, **kwargs: Any, ) -> torch.Tensor: - """Receive this rank's FFN output slice on the Attention side. + """Receive this rank's FFN output on the Attention side. + + Receives one share per subgroup FFN member, in the sizes and member + order ``send_attn_output`` used, and assembles them into an output + shaped like ``ref_tensor``, restoring the original token order. Args: - ref_tensor: Preallocated CUDA tensor to receive into. Used as a - stable buffer for CUDA graph capture and returned directly - when the subgroup has a single rank and no wire transfer - occurs. + ref_tensor: Preallocated CUDA tensor of shape + ``(num_tokens, hidden_size)``. A single FFN member receives + into it directly, keeping the allocation stable for CUDA + graph capture; several members use its shape and device. ubatch_idx: Stage/microbatch index. Defaults to ``0``. **kwargs: Unused; accepted for interface compatibility. @@ -394,17 +428,26 @@ def recv_ffn_output( RuntimeError: If the connector is not initialized, or no receive is performed for a single-rank subgroup. """ - output = self._recv_hidden_states( - 0, - self.e2a_group, - self.e2a_comm_id, - self.tensor_metadata_list[ubatch_idx], - ref_tensor=ref_tensor, - ) - if output is None: - raise RuntimeError( - "P2P recv_ffn_output requires ref_tensor when no receive is performed", + if self.subgroup_ffn_count == 1: + output = self._recv_hidden_states( + 0, + self.e2a_group, + self.e2a_comm_id, + self.tensor_metadata_list[ubatch_idx], + ref_tensor=ref_tensor, ) + if output is None: + raise RuntimeError( + "P2P recv_ffn_output requires ref_tensor when no receive " + "is performed", + ) + return output + sizes = split_send_sizes(ref_tensor.shape[0], self.subgroup_ffn_count) + output = torch.empty_like(ref_tensor) + start = 0 + for src, size in enumerate(sizes): + self._recv_slice(ref_tensor, output, start, size, src) + start = start + size return output def recv_attn_output( @@ -414,12 +457,14 @@ def recv_attn_output( ) -> AFDA2FTransferPayload: """Receive and concatenate Attention hidden states on the FFN rank. - Receives one tensor from every Attention peer in this rank's - subgroup (subgroup ranks ``1..ratio``), concatenates them along the - token dimension, and records each peer's sequence length in the - returned metadata so ``send_ffn_output`` can split the FFN output - back per peer. When CUDA graphs are enabled, receives reuse the - buffers preallocated by ``control_plane.update_state_from_dp_metadata``. + Receives one tensor from every Attention member of this rank's + subgroup, concatenates them along the token dimension, and records + each peer's sequence length in the returned metadata so + ``send_ffn_output`` can split the FFN output back per peer. + Zero-size shares add no rows; when every share is zero a one-token + dummy batch is returned instead. When CUDA graphs are enabled, + receives reuse the buffers preallocated by + ``control_plane.update_state_from_dp_metadata``. Args: ubatch_idx: Stage/microbatch index to receive. Defaults to ``0``. @@ -432,8 +477,7 @@ def recv_attn_output( lengths. Raises: - RuntimeError: If the connector is not initialized or the subgroup - has no Attention peers. + RuntimeError: If the connector is not initialized. """ routing_spec: AFDExpertRoutingSpec | None = kwargs.get("routing_spec") recv_input_ids: bool = kwargs.get("recv_input_ids", False) @@ -441,25 +485,27 @@ def recv_attn_output( router_logits_list: list[torch.Tensor] = [] input_ids_list: list[torch.Tensor] = [] - for src in range(1, self.group_size): + for src in range(self.subgroup_ffn_count, self.group_size): tensor_metadata = self._recv_attn_tensor_metadata_list.get( (ubatch_idx, src), self.tensor_metadata_list[ubatch_idx], ) + # The sender posts every share; a zero-size one adds no rows. + share_is_empty = tensor_metadata.size[0] == 0 ref_tensor = None if not self.vllm_config.model_config.enforce_eager: ref_tensor = self._recv_attn_buffers.get( (ubatch_idx, src, tuple(tensor_metadata.size)), ) - hidden_states_list.append( - self._recv_hidden_states( - src, - self.a2e_group, - self.a2e_comm_id, - tensor_metadata, - ref_tensor=ref_tensor, - ), + received = self._recv_hidden_states( + src, + self.a2e_group, + self.a2e_comm_id, + tensor_metadata, + ref_tensor=ref_tensor, ) + if not share_is_empty: + hidden_states_list.append(received) if routing_spec is not None: router_metadata = _TensorMetadata( device=tensor_metadata.device, @@ -471,14 +517,14 @@ def recv_attn_output( ], ), ) - router_logits_list.append( - self._recv_hidden_states( - src, - self.a2e_group, - self.a2e_comm_id, - router_metadata, - ), + received = self._recv_hidden_states( + src, + self.a2e_group, + self.a2e_comm_id, + router_metadata, ) + if not share_is_empty: + router_logits_list.append(received) if recv_input_ids: input_ids_metadata = _TensorMetadata( device=tensor_metadata.device, @@ -490,18 +536,52 @@ def recv_attn_output( input_ids_ref = self._recv_attn_input_ids_buffers.get( (ubatch_idx, src, tuple(input_ids_metadata.size)), ) - input_ids_list.append( - self._recv_hidden_states( - src, - self.a2e_group, - self.a2e_comm_id, - input_ids_metadata, - ref_tensor=input_ids_ref, - ), + received = self._recv_hidden_states( + src, + self.a2e_group, + self.a2e_comm_id, + input_ids_metadata, + ref_tensor=input_ids_ref, ) + if not share_is_empty: + input_ids_list.append(received) if not hidden_states_list: - raise RuntimeError("P2P FFN rank has no Attention peers") + # Nothing arrived: run the forward on a dummy batch, as vLLM does + # for an idle DP rank. AFDTransferMetadata rejects zero lengths. + stage_metadata = self.tensor_metadata_list[ubatch_idx] + return AFDA2FTransferPayload( + hidden_states=torch.zeros( + (1, self.hidden_size), + dtype=stage_metadata.dtype, + device=stage_metadata.device, + ), + context=AFDTransferContext( + metadata=AFDTransferMetadata.create_ffn_metadata( + layer_idx=0, + stage_idx=ubatch_idx, + seq_lens=[1], + ), + ), + router_logits=( + None + if routing_spec is None + else torch.zeros( + (1, routing_spec.router_logits_width), + dtype=routing_spec.router_logits_dtype, + device=stage_metadata.device, + ) + ), + input_ids=( + None + if not recv_input_ids + else torch.zeros( + (1,), + dtype=_INPUT_IDS_DTYPE, + device=stage_metadata.device, + ) + ), + ) hidden_states = ( torch.cat(hidden_states_list, dim=0) if len(hidden_states_list) > 1 @@ -544,11 +624,11 @@ def send_ffn_output( ) -> None: """Split the FFN output and send each slice back to its Attention peer. - With a one-to-one mapping (``ratio == 1``) the whole tensor is sent to - the single Attention peer. Otherwise the output is split along the - token dimension by ``metadata.seq_lens`` — falling back to an even - split when the recorded lengths do not cover every peer — and each - slice is sent to the Attention rank that originally produced it. + With a single Attention member the whole tensor is sent to it, or a + zero-size tensor when nothing arrived this step. Otherwise the output + is split along the token dimension by ``metadata.seq_lens``, falling + back to an even split when the recorded lengths do not cover every + peer, and each slice is sent to the Attention rank that produced it. Args: ffn_output: CUDA tensor of shape ``(num_tokens, hidden_size)`` @@ -566,6 +646,14 @@ def send_ffn_output( RuntimeError: If the connector is not initialized. """ metadata = context.metadata + if metadata.stage_idx in self._dummy_stages: + self._send_hidden_states( + ffn_output.narrow(0, 0, 0), + self.subgroup_ffn_count, + self.e2a_group, + self.e2a_comm_id, + ) + return if not torch.compiler.is_compiling() and not metadata.validate_tensor_shape( tuple(ffn_output.shape), ): @@ -573,7 +661,12 @@ def send_ffn_output( f"ffn_output shape {ffn_output.shape!r} does not match metadata", ) if self.ratio == 1: - self._send_hidden_states(ffn_output, 1, self.e2a_group, self.e2a_comm_id) + self._send_hidden_states( + ffn_output, + self.subgroup_ffn_count, + self.e2a_group, + self.e2a_comm_id, + ) return split_sizes = metadata.seq_lens @@ -589,7 +682,7 @@ def send_ffn_output( start = 0 for dst, token_count in zip( - range(1, self.group_size), + range(self.subgroup_ffn_count, self.group_size), split_sizes, strict=False, ): @@ -631,6 +724,40 @@ def _send_hidden_states( ) return + def _send_slice( + self, + base: torch.Tensor, + start: int, + size: int, + dst: int, + ) -> None: + """Send rows ``[start, start + size)`` of ``base`` to subgroup rank ``dst``.""" + if self.a2e_group is None or self.a2e_comm_id is None: + raise RuntimeError("P2P connector is not initialized") + if dst >= self.a2e_group.world_size: + raise ValueError(f"invalid P2P destination rank {dst}") + torch.ops.vllm.afd_p2p_send_slice(base, start, size, dst, self.a2e_comm_id) + + def _recv_slice( + self, + base: torch.Tensor, + out: torch.Tensor, + start: int, + size: int, + src: int, + ) -> None: + """Receive ``size`` rows from subgroup rank ``src`` into ``out``. + + ``base`` is only read, to order this receive after the sends that + declare it mutated; without that dependency the compiler hoisted the + receives above the sends and deadlocked. + """ + if self.e2a_group is None or self.e2a_comm_id is None: + raise RuntimeError("P2P connector is not initialized") + if src >= self.e2a_group.world_size: + raise ValueError(f"invalid P2P source rank {src}") + torch.ops.vllm.afd_p2p_recv_slice(base, out, start, size, src, self.e2a_comm_id) + def _recv_hidden_states( self, src: int, @@ -729,34 +856,30 @@ def update_state_from_dp_metadata( connector.is_warmup = payload.is_warmup connector.tensor_metadata_list = {} connector._recv_attn_tensor_metadata_list = {} + connector._dummy_stages = set() device = torch.device(f"cuda:{connector.local_rank}") dtype = connector.vllm_config.model_config.dtype for stage_idx, dp_metadata in payload.dp_metadata_list.items(): stage_idx = stage_idx if connector.afd_config.role == "ffn": peer_metadata: list[_TensorMetadata] = [] - for src_rank in range(1, connector.group_size): - if src_rank <= 0 or src_rank >= connector.group_size: - raise ValueError(f"invalid Attention subgroup rank {src_rank}") + ffn_member_count = connector.subgroup_ffn_count + for src_rank in range(ffn_member_count, connector.group_size): attention_rank = ( - connector.mapping.subgroup_index * connector.ratio - + src_rank - - 1 + connector.mapping.subgroup_ranks[src_rank] - connector.ffn_size ) - + token_count = _num_tokens_for_attention_rank( + dp_metadata, + attention_rank=attention_rank, + attention_size=connector.attn_size, + ) + share = split_send_sizes(token_count, ffn_member_count)[ + connector.mapping.rank_in_subgroup + ] tensor_metadata = _TensorMetadata( device, dtype, - torch.Size( - [ - _num_tokens_for_attention_rank( - dp_metadata, - attention_rank=attention_rank, - attention_size=connector.attn_size, - ), - connector.hidden_size, - ], - ), + torch.Size([share, connector.hidden_size]), ) connector._recv_attn_tensor_metadata_list[(stage_idx, src_rank)] = ( tensor_metadata @@ -765,6 +888,8 @@ def update_state_from_dp_metadata( num_tokens = sum( tensor_metadata.size[0] for tensor_metadata in peer_metadata ) + if num_tokens == 0: + connector._dummy_stages.add(stage_idx) else: num_tokens = _num_tokens_for_attention_rank( dp_metadata, @@ -902,6 +1027,58 @@ def _register_comm(communicator: PyNcclCommunicator) -> int: return comm_id +def _local_address_toward(host: str, port: int) -> str: + """Address of the local interface that routes to ``host:port``. + + Connecting a UDP socket sends nothing; it only makes the kernel pick the + outgoing interface, whose address is then read back. On a single host + this is ``host`` itself (for example ``127.0.0.1``). + """ + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as probe: + probe.connect((host, port)) + return probe.getsockname()[0] + + +def _gather_rank_addresses( + group: ProcessGroup, + *, + host: str, + port: int, + world_size: int, +) -> list[str]: + """Every rank's routable address, indexed by AFD world rank.""" + gathered: list[str | None] = [None] * world_size + torch.distributed.all_gather_object( + gathered, + _local_address_toward(host, port), + group=group, + ) + return [address or host for address in gathered] + + +def _register_op_once( + op_name: str, + op_func: Callable[..., None], + mutates_args: list[str], + fake_impl: Callable[..., None], +) -> None: + """Register ``op_name`` under ``torch.ops.vllm`` unless it already exists. + + A prior import of this module (for example after an importlib reload) can + leave the op defined in torch's process-global registry while + ``_AFD_CUSTOM_OPS_REGISTERED`` is back to ``False``. Reuse the existing op + instead of re-defining it, which torch rejects. + """ + if hasattr(torch.ops.vllm, op_name): + return + direct_register_custom_op( + op_name=op_name, + op_func=op_func, + mutates_args=mutates_args, + fake_impl=fake_impl, + ) + + def _register_p2p_custom_ops() -> None: """Register the AFD P2P send and receive custom ops. @@ -909,6 +1086,11 @@ def _register_p2p_custom_ops() -> None: fake implementations keeps the transfers traceable by ``torch.compile`` and capturable by CUDA graphs. Registration happens once per process; ops already registered by another connector instance are reused. + + The ``_slice`` variants take a whole tensor plus ``(start, size)`` and + slice inside the op, because a view at a non-zero offset passed to a + mutating custom op reached the op with offset 0 under ``torch.compile``. + Only subgroups with several FFN members call them. """ global _AFD_CUSTOM_OPS_REGISTERED @@ -958,37 +1140,83 @@ def afd_p2p_recv_fake( ) -> None: pass - def register_one( - op_name: str, - op_func: Callable[..., None], - mutates_args: list[str], - fake_impl: Callable[..., None], + def afd_p2p_send_slice_impl( + base: torch.Tensor, + start: int, + size: int, + dst: int, + comm_id: int, ) -> None: - # A prior import of this module (for example after an importlib reload) - # can leave the op defined in torch's process-global registry while - # this module's _AFD_CUSTOM_OPS_REGISTERED flag is back to False. Reuse - # the existing op instead of re-defining it, which torch rejects. - if hasattr(torch.ops.vllm, op_name): - return - direct_register_custom_op( - op_name=op_name, - op_func=op_func, - mutates_args=mutates_args, - fake_impl=fake_impl, + communicator = _AFD_COMMUNICATORS.get(comm_id) + if communicator is None: + raise RuntimeError(f"AFD communicator id {comm_id} is not registered") + communicator.send( + base.narrow(0, start, size).contiguous(), + dst, + stream=torch.cuda.current_stream(base.device), + ) + return None + + def afd_p2p_send_slice_fake( + base: torch.Tensor, + start: int, + size: int, + dst: int, + comm_id: int, + ) -> None: + pass + + def afd_p2p_recv_slice_impl( + base: torch.Tensor, + out: torch.Tensor, + start: int, + size: int, + src: int, + comm_id: int, + ) -> None: + communicator = _AFD_COMMUNICATORS.get(comm_id) + if communicator is None: + raise RuntimeError(f"AFD communicator id {comm_id} is not registered") + communicator.recv( + out.narrow(0, start, size), + src, + stream=torch.cuda.current_stream(out.device), ) - register_one( + def afd_p2p_recv_slice_fake( + base: torch.Tensor, + out: torch.Tensor, + start: int, + size: int, + src: int, + comm_id: int, + ) -> None: + pass + + _register_op_once( op_name="afd_p2p_send", op_func=afd_p2p_send_impl, mutates_args=["tensor"], fake_impl=afd_p2p_send_fake, ) - register_one( + _register_op_once( op_name="afd_p2p_recv", op_func=afd_p2p_recv_impl, mutates_args=["out"], fake_impl=afd_p2p_recv_fake, ) + _register_op_once( + op_name="afd_p2p_send_slice", + op_func=afd_p2p_send_slice_impl, + mutates_args=["base"], + fake_impl=afd_p2p_send_slice_fake, + ) + _register_op_once( + op_name="afd_p2p_recv_slice", + op_func=afd_p2p_recv_slice_impl, + mutates_args=["out"], + fake_impl=afd_p2p_recv_slice_fake, + ) _AFD_CUSTOM_OPS_REGISTERED = True diff --git a/afd_plugin/distributed/__init__.py b/afd_plugin/distributed/__init__.py index 8769568e..a06634c6 100644 --- a/afd_plugin/distributed/__init__.py +++ b/afd_plugin/distributed/__init__.py @@ -6,6 +6,7 @@ AFDRankMapping, build_rank_mapping, resolve_role_rank, + split_send_sizes, topology_from_config, validate_p2p_topology, ) @@ -32,6 +33,7 @@ def __getattr__(name: str): "create_hccl_process_group_options", "init_afd_process_group", "resolve_role_rank", + "split_send_sizes", "topology_from_config", "validate_p2p_topology", ] diff --git a/afd_plugin/distributed/topology.py b/afd_plugin/distributed/topology.py index 55732b26..3bef6efa 100644 --- a/afd_plugin/distributed/topology.py +++ b/afd_plugin/distributed/topology.py @@ -18,8 +18,11 @@ class AFDRankMapping: """Rank mapping for the P2P connector. The P2P world always places FFN ranks first, followed by Attention ranks: - ``[F0, F1, ..., A0, A1, ...]``. Each FFN rank owns one subgroup containing - itself at subgroup rank 0 and one or more consecutive Attention ranks. + ``[F0, F1, ..., A0, A1, ...]``. The world is partitioned into + ``min(A, F)`` subgroups, each holding at least one rank of both roles; + ``subgroup_ranks`` lists a subgroup's FFN members first, then its + Attention members, in world order. ``ratio`` is the number of Attention + members in this rank's subgroup. """ role: str @@ -52,16 +55,10 @@ def topology_from_config(config: AFDConfig) -> tuple[int, int]: def validate_p2p_topology(config: AFDConfig) -> None: attention_size, ffn_size = topology_from_config(config) - if attention_size < ffn_size: + if attention_size <= 0 or ffn_size <= 0: raise ValueError( - "P2pNcclAFDConnector currently requires num_attention_ranks >= " - f"num_ffn_ranks, got {attention_size} < {ffn_size}", - ) - if attention_size % ffn_size != 0: - raise ValueError( - "P2pNcclAFDConnector currently requires num_attention_ranks to be a " - "multiple of num_ffn_ranks, got " - f"{attention_size} and {ffn_size}", + "P2P topology requires positive rank counts, got " + f"A={attention_size}, F={ffn_size}", ) @@ -112,6 +109,30 @@ def resolve_role_rank(vllm_config: VllmConfig, config: AFDConfig) -> int: return role_rank +def split_send_sizes(token_count: int, parts: int) -> tuple[int, ...]: + """Split one rank's token count into the sizes it sends to each peer. + + The sender slices its tensor with the whole tuple and each receiver + reads the element at its own position, so both sides reach identical + sizes without extra communication. The rule matches + ``torch.tensor_split``: sizes differ by at most one and the + ``token_count % parts`` leading shares carry the extra token. A share is + zero when ``token_count < parts``. + + Share ``i`` is the closed form ``(token_count - i + parts - 1) // parts`` + rather than a branch on the remainder, so that ``token_count`` may also + be a ``SymInt``: under ``torch.compile`` a Python branch would freeze at + its trace-time outcome and split every other token count wrongly. + """ + if parts <= 0: + raise ValueError(f"parts must be positive, got {parts}") + if token_count < 0: + raise ValueError(f"token_count must be >= 0, got {token_count}") + return tuple( + (token_count - position + parts - 1) // parts for position in range(parts) + ) + + def build_rank_mapping( config: AFDConfig, role_rank: int, @@ -130,7 +151,6 @@ def build_rank_mapping( f"(rank={role_rank}, size={attention_size})", ) world_rank = ffn_size + role_rank - subgroup_index = role_rank // (attention_size // ffn_size) elif config.role == "ffn": if role_rank >= ffn_size: raise ValueError( @@ -138,20 +158,33 @@ def build_rank_mapping( f"(rank={role_rank}, size={ffn_size})", ) world_rank = role_rank - subgroup_index = role_rank else: raise ValueError(f"unknown AFD role {config.role!r}") - ratio = attention_size // ffn_size + # Balanced block distribution: i * G // N maps N ranks onto G contiguous + # blocks differing in size by at most one, and is onto for both roles. min_size = min(ffn_size, attention_size) - ffn_ranks = list(range(ffn_size)) - attention_ranks = list(range(ffn_size, ffn_size + attention_size)) - subgroup_ranks = tuple( - [ffn_ranks[subgroup_index]] - + [attention_ranks[subgroup_index * ratio + offset] for offset in range(ratio)], - ) + if config.role == "attention": + subgroup_index = role_rank * min_size // attention_size + else: + subgroup_index = role_rank * min_size // ffn_size + subgroup_ffn_ranks = [ + ffn_rank + for ffn_rank in range(ffn_size) + if ffn_rank * min_size // ffn_size == subgroup_index + ] + subgroup_attention_ranks = [ + ffn_size + attention_rank + for attention_rank in range(attention_size) + if attention_rank * min_size // attention_size == subgroup_index + ] + subgroup_ranks = tuple(subgroup_ffn_ranks + subgroup_attention_ranks) rank_in_subgroup = subgroup_ranks.index(world_rank) - p2p_rank = role_rank + min_size if config.role == "attention" else role_rank + ratio = len(subgroup_attention_ranks) + + # FFN ranks take p2p_rank 0..F-1, so Attention starts at ffn_size, the + # offset the receiver side assumes. + p2p_rank = role_rank + ffn_size if config.role == "attention" else role_rank destinations: list[int] = [] if ffn_size <= world_rank < ffn_size + min_size: @@ -181,6 +214,7 @@ def build_rank_mapping( "AFDRankMapping", "build_rank_mapping", "resolve_role_rank", + "split_send_sizes", "topology_from_config", "validate_p2p_topology", ] diff --git a/afd_plugin/v1/worker/ffn_metadata.py b/afd_plugin/v1/worker/ffn_metadata.py index e2e15ea9..995c4314 100644 --- a/afd_plugin/v1/worker/ffn_metadata.py +++ b/afd_plugin/v1/worker/ffn_metadata.py @@ -4,6 +4,8 @@ from __future__ import annotations +from afd_plugin.distributed import split_send_sizes + def aggregate_ffn_token_counts( attention_counts: tuple[int, ...], @@ -12,17 +14,20 @@ def aggregate_ffn_token_counts( ffn_size: int, fallback: int = 1, ) -> tuple[int, ...]: - """Aggregate consecutive Attention-rank counts for each FFN rank. + """Token count each FFN rank receives for a stage, from Attention counts. - For example, ``4A2F`` counts ``(0, 4, 5, 6)`` become ``(5, 11)`` because - every zero-token Attention peer contributes one placeholder row. Missing - peers use the same per-peer fallback, so empty counts become ``(2, 2)``. + An FFN rank sums the shares its subgroup's Attention members send it, + splitting with ``split_send_sizes`` as the connector does. For example, + ``4A2F`` counts ``(0, 4, 5, 6)`` become ``(5, 11)`` because every + zero-token Attention peer contributes one placeholder row, and ``1A2F`` + count ``(5,)`` becomes ``(3, 2)``. No count is below one, since an FFN + rank that receives nothing runs a one-token dummy batch. Missing peers + use the same per-peer fallback. """ fallback_count = max(1, int(fallback)) - fallback_counts = tuple(fallback_count for _ in range(max(0, ffn_size))) - if ffn_size <= 0 or attention_size < ffn_size or attention_size % ffn_size != 0: - return fallback_counts + if ffn_size <= 0 or attention_size <= 0: + return tuple(fallback_count for _ in range(max(0, ffn_size))) expanded_counts = attention_counts if ( @@ -36,19 +41,28 @@ def aggregate_ffn_token_counts( for rank in range(attention_size) ) - group_size = attention_size // ffn_size - return tuple( - sum( - max(1, int(expanded_counts[attention_rank])) - if attention_rank < len(expanded_counts) - else fallback_count - for attention_rank in range( - ffn_rank * group_size, - (ffn_rank + 1) * group_size, - ) + def peer_count(attention_rank: int) -> int: + if attention_rank < len(expanded_counts): + return max(1, int(expanded_counts[attention_rank])) + return fallback_count + + group_count = min(attention_size, ffn_size) + ffn_counts = [] + for ffn_rank in range(ffn_size): + group = ffn_rank * group_count // ffn_size + ffn_members = [ + member + for member in range(ffn_size) + if member * group_count // ffn_size == group + ] + position = ffn_members.index(ffn_rank) + total = sum( + split_send_sizes(peer_count(attention_rank), len(ffn_members))[position] + for attention_rank in range(attention_size) + if attention_rank * group_count // attention_size == group ) - for ffn_rank in range(ffn_size) - ) + ffn_counts.append(max(1, total)) + return tuple(ffn_counts) def project_ffn_token_counts_to_dp( diff --git a/tests/unit/connectors/test_p2p_compiled_send.py b/tests/unit/connectors/test_p2p_compiled_send.py new file mode 100644 index 00000000..7526b5ae --- /dev/null +++ b/tests/unit/connectors/test_p2p_compiled_send.py @@ -0,0 +1,198 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""The Attention-side transfer traced once must be right for every token count. + +vLLM compiles the model forward once with a symbolic token dimension and +drops all guards (``TorchCompileWithNoGuardsWrapper``), then reuses that +trace for every later batch. ``send_attn_output`` and ``recv_ffn_output`` +run inside that trace, so their share arithmetic must not depend on +anything frozen at trace time: no Python-int sizes from the metadata and no +branch on the remainder (a ``divmod``-based split traced at 64 tokens would +keep sending (n//2, n//2) — two rows for five tokens). + +These tests compile the connector methods with the same options vLLM uses, +trace at 64 tokens, and then call the compiled code with other token counts +through custom ops that record what actually went on the wire. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + import torch +else: + torch = pytest.importorskip("torch") +pytest.importorskip("vllm") + +from afd_plugin.config import AFDConfig # noqa: E402 +from afd_plugin.connectors import ( # noqa: E402 + AFDConnectorFactory, + AFDControlPayload, + AFDDPMetadata, + AFDTransferContext, + AFDTransferMetadata, +) +from afd_plugin.distributed import split_send_sizes # noqa: E402 + +_HIDDEN = 8 +_SENT: list[tuple[int, int]] = [] +_RECEIVED: list[tuple[int, int]] = [] + + +# Custom ops with fake implementations stand in for the NCCL ops so the +# transfer survives tracing and the recorded sizes are runtime values. Like +# the real slice ops they take whole tensors plus (start, size). +@torch.library.custom_op("afd_test::record_send_slice", mutates_args=("base",)) +def _record_send_slice(base: torch.Tensor, start: int, size: int, dst: int) -> None: + _SENT.append((dst, int(size))) + + +@_record_send_slice.register_fake +def _(base, start, size, dst): + return None + + +@torch.library.custom_op("afd_test::record_recv_slice", mutates_args=("out",)) +def _record_recv_slice( + base: torch.Tensor, out: torch.Tensor, start: int, size: int, src: int +) -> None: + _RECEIVED.append((src, int(size))) + out.narrow(0, start, size).fill_(float(src)) + + +@_record_recv_slice.register_fake +def _(base, out, start, size, src): + return None + + +@torch.library.custom_op("afd_test::record_send", mutates_args=("tensor",)) +def _record_send(tensor: torch.Tensor, dst: int) -> None: + _SENT.append((dst, int(tensor.shape[0]))) + + +@_record_send.register_fake +def _(tensor, dst): + return None + + +def _attention_connector(attention, ffn): + text_config = SimpleNamespace(hidden_size=_HIDDEN, num_hidden_layers=2) + vllm_config = SimpleNamespace( + additional_config={}, + model_config=SimpleNamespace( + dtype=torch.float32, + enforce_eager=False, + hf_config=text_config, + hf_text_config=text_config, + ), + parallel_config=SimpleNamespace( + data_parallel_size=attention, + data_parallel_rank=0, + prefill_context_parallel_size=1, + tensor_parallel_size=1, + ), + ) + connector = AFDConnectorFactory.create_connector( + 0, + 0, + vllm_config, + AFDConfig( + role="attention", + connector="P2pNcclAFDConnector", + num_attention_ranks=attention, + num_ffn_ranks=ffn, + ), + ) + connector.control_plane.update_state_from_dp_metadata( + AFDControlPayload( + dp_metadata_list={0: AFDDPMetadata([64] * attention)}, + is_graph_capturing=False, + is_warmup=False, + ), + ) + connector._send_hidden_states = lambda hidden, dst, group, comm_id: ( + torch.ops.afd_test.record_send(hidden, dst) + ) + connector._send_slice = lambda base, start, size, dst: ( + torch.ops.afd_test.record_send_slice(base, start, size, dst) + ) + connector._recv_slice = lambda base, out, start, size, src: ( + torch.ops.afd_test.record_recv_slice(base, out, start, size, src) + ) + return connector + + +def _compile_like_vllm(fn): + return torch.compile( + fn, + fullgraph=True, + dynamic=False, + backend="aot_eager", + options={"guard_filter_fn": torch.compiler.skip_all_guards_unsafe}, + ) + + +def _context(seq_len): + return AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=0, + stage_idx=0, + seq_len=seq_len, + ), + ) + + +@pytest.mark.parametrize("ffn", [2, 4]) +def test_compiled_send_splits_every_token_count_like_the_receivers_expect(ffn): + torch._dynamo.reset() + connector = _attention_connector(1, ffn) + context = _context(64) + compiled = _compile_like_vllm( + lambda hidden: connector.send_attn_output(hidden, context), + ) + traced = torch.zeros((64, _HIDDEN)) + torch._dynamo.mark_dynamic(traced, 0) + compiled(traced) + for tokens in (64, 5, 7, 1, 128, 0): + _SENT.clear() + compiled(torch.zeros((tokens, _HIDDEN))) + assert list(enumerate(split_send_sizes(tokens, ffn))) == _SENT, tokens + + +def test_compiled_receive_fills_each_share_in_member_order(): + torch._dynamo.reset() + connector = _attention_connector(1, 2) + compiled = _compile_like_vllm( + lambda buffer: connector.recv_ffn_output(ref_tensor=buffer, ubatch_idx=0), + ) + traced = torch.zeros((64, _HIDDEN)) + torch._dynamo.mark_dynamic(traced, 0) + compiled(traced) + for tokens in (5, 7, 1): + _RECEIVED.clear() + output = compiled(torch.full((tokens, _HIDDEN), -1.0)) + first, second = split_send_sizes(tokens, 2) + assert [(0, first), (1, second)] == _RECEIVED, tokens + assert output.shape == (tokens, _HIDDEN) + assert torch.equal(output[:first], torch.zeros((first, _HIDDEN))) + assert torch.equal(output[first:], torch.ones((second, _HIDDEN))) + + +def test_single_member_compiled_send_is_one_whole_send(): + torch._dynamo.reset() + connector = _attention_connector(2, 1) + context = _context(64) + compiled = _compile_like_vllm( + lambda hidden: connector.send_attn_output(hidden, context), + ) + traced = torch.zeros((64, _HIDDEN)) + torch._dynamo.mark_dynamic(traced, 0) + compiled(traced) + for tokens in (5, 1, 128): + _SENT.clear() + compiled(torch.zeros((tokens, _HIDDEN))) + assert [(0, tokens)] == _SENT diff --git a/tests/unit/connectors/test_p2p_connector.py b/tests/unit/connectors/test_p2p_connector.py index 13250390..dab0b73b 100644 --- a/tests/unit/connectors/test_p2p_connector.py +++ b/tests/unit/connectors/test_p2p_connector.py @@ -419,30 +419,33 @@ def test_p2p_tensor_metadata_clamps_idle_attention_rank_to_dummy_token(): assert attention_connector.tensor_metadata_list[0].size == torch.Size([1, 16]) +# The A >= F and divisibility constraints were removed with the unified +# subgroup partition; the shapes this test used to reject now parse cleanly, +# and only non-positive rank counts remain invalid. @pytest.mark.parametrize( - ("raw", "message"), - [ - ( - { - "connector": "P2pNcclAFDConnector", - "num_attention_ranks": 1, - "num_ffn_ranks": 2, - }, - "num_attention_ranks >= num_ffn_ranks", - ), - ( + ("attention", "ffn"), + [(1, 2), (3, 2)], +) +def test_p2p_topology_accepts_previously_constrained_shapes(attention, ffn): + config = afd_config_from_mapping( + { + "connector": "P2pNcclAFDConnector", + "num_attention_ranks": attention, + "num_ffn_ranks": ffn, + }, + ) + assert (config.num_attention_ranks, config.num_ffn_ranks) == (attention, ffn) + + +def test_p2p_topology_rejects_non_positive_rank_counts(): + with pytest.raises(ValueError, match="positive"): + afd_config_from_mapping( { "connector": "P2pNcclAFDConnector", - "num_attention_ranks": 3, + "num_attention_ranks": 0, "num_ffn_ranks": 2, }, - "multiple of num_ffn_ranks", - ), - ], -) -def test_p2p_topology_validation_errors_are_clear(raw, message): - with pytest.raises(ValueError, match=message): - afd_config_from_mapping(raw) + ) def test_p2p_module_exports_connector_class(): @@ -542,11 +545,16 @@ def direct_register_custom_op(**kwargs): assert [call["op_name"] for call in calls] == [ "afd_p2p_send", "afd_p2p_recv", + "afd_p2p_send_slice", + "afd_p2p_recv_slice", + ] + assert [call["mutates_args"] for call in calls] == [ + ["tensor"], + ["out"], + ["base"], + ["out"], ] - assert calls[0]["mutates_args"] == ["tensor"] - assert calls[1]["mutates_args"] == ["out"] - assert callable(calls[0]["fake_impl"]) - assert callable(calls[1]["fake_impl"]) + assert all(callable(call["fake_impl"]) for call in calls) def test_p2p_hidden_state_send_uses_registered_custom_op(monkeypatch): diff --git a/tests/unit/connectors/test_p2p_generalized_topology.py b/tests/unit/connectors/test_p2p_generalized_topology.py new file mode 100644 index 00000000..486a1729 --- /dev/null +++ b/tests/unit/connectors/test_p2p_generalized_topology.py @@ -0,0 +1,529 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Transport behavior of ``P2pNcclAFDConnector`` on generalized topologies. + +The subgroup partition is now defined for any positive (A, F). The +existing connector tests cover the historical A >= F layouts; this file +covers the new shapes — A < F (one attention splitting across several FFN +members) and non-divisible A > F (uneven subgroups) — plus the historical +4A2F layout run through the same checks as a regression anchor. + +All communication is mocked at the ``_send_hidden_states`` / +``_recv_hidden_states`` seam, the same one the existing tests use, so the +connector's own logic (who sends how many rows to whom, in which order) +runs unchanged on the CPU. The round-trip tests connect all ranks of a +topology through a fake wire (a dict of queues keyed by subgroup and +member ranks) so that sizes, addressing, and ordering are checked end to +end exactly as the real NCCL exchange would pair them. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + import torch +else: + torch = pytest.importorskip("torch") +pytest.importorskip("vllm") + +from afd_plugin.config import AFDConfig # noqa: E402 +from afd_plugin.connectors import ( # noqa: E402 + AFDConnectorFactory, + AFDControlPayload, + AFDDPMetadata, + AFDTransferContext, + AFDTransferMetadata, +) +from afd_plugin.connectors.gpu.p2p import _TensorMetadata # noqa: E402 + +_HIDDEN = 16 + + +def _fake_vllm_config(*, dp_size=1, dp_rank=0, enforce_eager=True): + text_config = SimpleNamespace(hidden_size=_HIDDEN, num_hidden_layers=2) + return SimpleNamespace( + additional_config={}, + model_config=SimpleNamespace( + dtype=torch.float32, + enforce_eager=enforce_eager, + hf_config=text_config, + hf_text_config=text_config, + ), + parallel_config=SimpleNamespace( + data_parallel_size=dp_size, + data_parallel_rank=dp_rank, + prefill_context_parallel_size=1, + tensor_parallel_size=1, + ), + ) + + +def _connector(role, role_rank, attention, ffn, **config_kwargs): + # The factory derives the role rank from the vLLM DP coordinates, so the + # fake config carries this connector's own DP size/rank. + role_size = attention if role == "attention" else ffn + return AFDConnectorFactory.create_connector( + role_rank, + 0, + _fake_vllm_config(dp_size=role_size, dp_rank=role_rank, **config_kwargs), + AFDConfig( + role=role, + connector="P2pNcclAFDConnector", + num_attention_ranks=attention, + num_ffn_ranks=ffn, + ), + ) + + +def _all_connectors(attention, ffn): + attentions = [_connector("attention", a, attention, ffn) for a in range(attention)] + ffns = [_connector("ffn", f, attention, ffn) for f in range(ffn)] + return attentions, ffns + + +def _payload(token_counts): + return AFDControlPayload( + dp_metadata_list={0: AFDDPMetadata(token_counts)}, + is_graph_capturing=False, + is_warmup=False, + ) + + +def _apply(connectors, token_counts): + for connector in connectors: + connector.control_plane.update_state_from_dp_metadata(_payload(token_counts)) + # The dummy batch allocates on the stage metadata's device; point + # it at the CPU so the tests can run without CUDA. + old = connector.tensor_metadata_list[0] + connector.tensor_metadata_list[0] = _TensorMetadata( + torch.device("cpu"), + old.dtype, + old.size, + ) + + +def _attention_context(seq_len): + return AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=0, + stage_idx=0, + seq_len=seq_len, + ), + ) + + +def _ffn_context(seq_lens): + return AFDTransferContext( + metadata=AFDTransferMetadata.create_ffn_metadata( + layer_idx=0, + stage_idx=0, + seq_lens=seq_lens, + ), + ) + + +def _record_sends(monkeypatch, connector): + """Replace both send seams with a recorder of (rows, dst, tensor). + + Single-member subgroups send whole tensors through ``_send_hidden_states``; + multi-member subgroups send (base, start, size) slices through + ``_send_slice`` — recorded as the sliced rows. + """ + sent = [] + monkeypatch.setattr( + connector, + "_send_hidden_states", + lambda hidden, dst, group, comm_id: sent.append((hidden.shape[0], dst, hidden)), + ) + monkeypatch.setattr( + connector, + "_send_slice", + lambda base, start, size, dst: sent.append( + (size, dst, base.narrow(0, start, size)) + ), + ) + return sent + + +def _fake_recvs(monkeypatch, connector, fill=None): + """Replace the receive seam; returns the (src, rows) call log. + + Received tensors are filled with ``fill(src)`` (default: the source + rank) so callers can tell which source each row came from. + """ + received = [] + + def fake_recv(src, group, comm_id, tensor_metadata, *, ref_tensor=None): + received.append((src, tensor_metadata.size[0])) + value = float(src) if fill is None else fill(src) + tensor = torch.full( + tuple(tensor_metadata.size), + value, + dtype=tensor_metadata.dtype, + ) + if ref_tensor is not None: + # The real op receives in place into the caller's buffer. + ref_tensor.copy_(tensor) + return ref_tensor + return tensor + + def fake_recv_slice(base, out, start, size, src): + received.append((src, size)) + value = float(src) if fill is None else fill(src) + out.narrow(0, start, size).fill_(value) + + monkeypatch.setattr(connector, "_recv_hidden_states", fake_recv) + monkeypatch.setattr(connector, "_recv_slice", fake_recv_slice) + return received + + +# -------------------------------------------------------------------------- +# 1. Construction: every topology is accepted in eager and graph mode alike +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("role", ["attention", "ffn"]) +@pytest.mark.parametrize("enforce_eager", [True, False]) +def test_multi_ffn_subgroup_is_accepted_in_both_modes(role, enforce_eager): + connector = _connector(role, 0, 1, 2, enforce_eager=enforce_eager) + assert connector.subgroup_ffn_count == 2 + + +def test_registration_defines_the_four_transfer_ops(): + """Every topology registers the same ops, so the schemas are pinned here. + + ``base`` is declared mutated on the send and only read by the receive, + which is what keeps receives ordered after sends in a compiled graph. + """ + from afd_plugin.connectors.gpu import p2p + + p2p._register_p2p_custom_ops() + schemas = { + name: str(getattr(torch.ops.vllm, name).default._schema) + for name in ( + "afd_p2p_send", + "afd_p2p_recv", + "afd_p2p_send_slice", + "afd_p2p_recv_slice", + ) + } + assert schemas["afd_p2p_send"] == ( + "vllm::afd_p2p_send(Tensor(a0!) tensor, SymInt dst, SymInt comm_id) -> ()" + ) + assert schemas["afd_p2p_recv"] == ( + "vllm::afd_p2p_recv(Tensor(a0!) out, SymInt src, SymInt comm_id) -> ()" + ) + assert schemas["afd_p2p_send_slice"] == ( + "vllm::afd_p2p_send_slice(Tensor(a0!) base, SymInt start, SymInt size, " + "SymInt dst, SymInt comm_id) -> ()" + ) + assert schemas["afd_p2p_recv_slice"] == ( + "vllm::afd_p2p_recv_slice(Tensor base, Tensor(a1!) out, SymInt start, " + "SymInt size, SymInt src, SymInt comm_id) -> ()" + ) + + +def test_historical_layout_still_allows_graph_mode(): + connector = _connector("attention", 1, 4, 2, enforce_eager=False) + assert connector.subgroup_ffn_count == 1 + + +# -------------------------------------------------------------------------- +# 2. Size pre-computation (update_state_from_dp_metadata) +# -------------------------------------------------------------------------- + + +def test_sizes_1a2f_split_across_two_ffn_members(): + # Group (F0, F1, A0): A0 sits at group rank 2 and splits 5 tokens 3/2. + attentions, ffns = _all_connectors(1, 2) + _apply(attentions + ffns, [5]) + assert attentions[0].tensor_metadata_list[0].size == torch.Size([5, _HIDDEN]) + expected = ffns[0]._recv_attn_tensor_metadata_list[(0, 2)].size + assert expected == torch.Size([3, _HIDDEN]) + expected = ffns[1]._recv_attn_tensor_metadata_list[(0, 2)].size + assert expected == torch.Size([2, _HIDDEN]) + assert ffns[0].tensor_metadata_list[0].size == torch.Size([3, _HIDDEN]) + assert ffns[1].tensor_metadata_list[0].size == torch.Size([2, _HIDDEN]) + assert all(not ffn._dummy_stages for ffn in ffns) + + +def test_sizes_1a4f_with_fewer_tokens_than_members_marks_dummy_stages(): + # 2 tokens over 4 FFN members: shares (1, 1, 0, 0). F2 and F3 receive + # nothing this step and are marked for the dummy batch. + attentions, ffns = _all_connectors(1, 4) + _apply(attentions + ffns, [2]) + shares = [ffn._recv_attn_tensor_metadata_list[(0, 4)].size[0] for ffn in ffns] + assert shares == [1, 1, 0, 0] + assert [sorted(ffn._dummy_stages) for ffn in ffns] == [[], [], [0], [0]] + + +def test_sizes_6a4f_uneven_subgroups(): + # Subgroups: {F0,A0,A1} {F1,A2} {F2,A3,A4} {F3,A5}; token count = rank+1. + attentions, ffns = _all_connectors(6, 4) + _apply(attentions + ffns, [1, 2, 3, 4, 5, 6]) + + def peers(ffn): + return sorted( + (src, md.size[0]) + for (_, src), md in ffn._recv_attn_tensor_metadata_list.items() + ) + + assert peers(ffns[0]) == [(1, 1), (2, 2)] + assert peers(ffns[1]) == [(1, 3)] + assert peers(ffns[2]) == [(1, 4), (2, 5)] + assert peers(ffns[3]) == [(1, 6)] + assert [ffn.tensor_metadata_list[0].size[0] for ffn in ffns] == [3, 3, 9, 6] + + +def test_sizes_4a2f_match_historical_layout(): + attentions, ffns = _all_connectors(4, 2) + _apply(attentions + ffns, [3, 5, 7, 11]) + assert ffns[0]._recv_attn_tensor_metadata_list[(0, 1)].size[0] == 3 + assert ffns[0]._recv_attn_tensor_metadata_list[(0, 2)].size[0] == 5 + assert ffns[1]._recv_attn_tensor_metadata_list[(0, 1)].size[0] == 7 + assert ffns[1]._recv_attn_tensor_metadata_list[(0, 2)].size[0] == 11 + + +# -------------------------------------------------------------------------- +# 3. Attention send +# -------------------------------------------------------------------------- + + +def test_attention_sends_slices_to_ffn_members_in_order(monkeypatch): + attention = _connector("attention", 0, 1, 2) + _apply([attention], [5]) + sent = _record_sends(monkeypatch, attention) + hidden = torch.arange(5 * _HIDDEN, dtype=torch.float32).reshape(5, _HIDDEN) + attention.send_attn_output(hidden, _attention_context(5)) + assert [(rows, dst) for rows, dst, _ in sent] == [(3, 0), (2, 1)] + assert torch.equal(sent[0][2], hidden[:3]) + assert torch.equal(sent[1][2], hidden[3:]) + + +def test_attention_sends_every_share_including_zero_size_ones(monkeypatch): + # 2 tokens over 4 members: shares (1, 1, 0, 0). The zero-size shares are + # sent too — the send path has no branch on the token count so that it + # traces correctly under torch.compile (see test_p2p_compiled_send). + attention = _connector("attention", 0, 1, 4) + _apply([attention], [2]) + sent = _record_sends(monkeypatch, attention) + attention.send_attn_output(torch.zeros((2, _HIDDEN)), _attention_context(2)) + assert [(rows, dst) for rows, dst, _ in sent] == [(1, 0), (1, 1), (0, 2), (0, 3)] + + +def test_attention_with_single_ffn_member_sends_the_whole_tensor(monkeypatch): + # One FFN member: the upstream path — a single send of the tensor + # object itself (no slicing, no view). + attention = _connector("attention", 1, 4, 2) + _apply([attention], [3, 5, 7, 11]) + sent = _record_sends(monkeypatch, attention) + hidden = torch.zeros((5, _HIDDEN)) + attention.send_attn_output(hidden, _attention_context(5)) + assert len(sent) == 1 + assert sent[0][1] == 0 + assert sent[0][2] is hidden + + +# -------------------------------------------------------------------------- +# 4. FFN receive (incl. the dummy batch) +# -------------------------------------------------------------------------- + + +def test_ffn_receives_its_share_from_the_attention_member(monkeypatch): + ffn1 = _connector("ffn", 1, 1, 2) + _apply([ffn1], [5]) + received = _fake_recvs(monkeypatch, ffn1) + payload = ffn1.recv_attn_output(ubatch_idx=0) + assert received == [(2, 2)] # attention member at group rank 2, 2 rows + assert payload.hidden_states.shape == (2, _HIDDEN) + assert payload.context.metadata.seq_lens == [2] + + +def test_ffn_with_nothing_to_receive_runs_on_a_dummy_batch(monkeypatch): + ffn2 = _connector("ffn", 2, 1, 4) + _apply([ffn2], [2]) + received = _fake_recvs(monkeypatch, ffn2) + payload = ffn2.recv_attn_output(ubatch_idx=0) + # The zero-size share is still received (the attention member sent it). + assert received == [(4, 0)] + assert payload.hidden_states.shape == (1, _HIDDEN) + assert torch.equal(payload.hidden_states, torch.zeros((1, _HIDDEN))) + assert payload.context.metadata.seq_lens == [1] + assert 0 in ffn2._dummy_stages + + +# -------------------------------------------------------------------------- +# 5. FFN return send +# -------------------------------------------------------------------------- + + +def test_ffn_returns_whole_output_to_its_attention_member(monkeypatch): + ffn1 = _connector("ffn", 1, 1, 2) + _apply([ffn1], [5]) + sent = _record_sends(monkeypatch, ffn1) + ffn1.send_ffn_output(torch.zeros((2, _HIDDEN)), _ffn_context([2])) + assert [(rows, dst) for rows, dst, _ in sent] == [(2, 2)] + + +def test_ffn_on_a_dummy_stage_returns_a_zero_size_tensor(monkeypatch): + # The dummy result is discarded; a zero-size tensor goes back to the + # attention member (group rank 4) to match its posted receive. + ffn2 = _connector("ffn", 2, 1, 4) + _apply([ffn2], [2]) + _fake_recvs(monkeypatch, ffn2) + payload = ffn2.recv_attn_output(ubatch_idx=0) + sent = _record_sends(monkeypatch, ffn2) + ffn2.send_ffn_output(payload.hidden_states, payload.context) + assert [(rows, dst) for rows, dst, _ in sent] == [(0, 4)] + + +def test_ffn_with_several_attention_members_splits_by_seq_lens(monkeypatch): + ffn0 = _connector("ffn", 0, 6, 4) # subgroup {F0, A0, A1} + _apply([ffn0], [1, 2, 3, 4, 5, 6]) + sent = _record_sends(monkeypatch, ffn0) + ffn0.send_ffn_output(torch.zeros((3, _HIDDEN)), _ffn_context([1, 2])) + assert [(rows, dst) for rows, dst, _ in sent] == [(1, 1), (2, 2)] + + +# -------------------------------------------------------------------------- +# 6. Attention return receive +# -------------------------------------------------------------------------- + + +def test_attention_reassembles_returns_in_member_order(monkeypatch): + attention = _connector("attention", 0, 1, 2) + _apply([attention], [5]) + received = _fake_recvs(monkeypatch, attention) + buffer = torch.full((5, _HIDDEN), -1.0) + output = attention.recv_ffn_output(ref_tensor=buffer, ubatch_idx=0) + assert received == [(0, 3), (1, 2)] + assert output.shape == (5, _HIDDEN) + assert torch.equal(output[:3], torch.zeros((3, _HIDDEN))) + assert torch.equal(output[3:], torch.ones((2, _HIDDEN))) + + +# -------------------------------------------------------------------------- +# 7 + 8. Round trip over a fake wire +# -------------------------------------------------------------------------- + + +def _wire(monkeypatch, connectors): + """Connect every connector through queues keyed by subgroup and ranks. + + A send from member ``s`` to member ``d`` of subgroup ``g`` lands on + ``wire[(g, s, d)]``; a receive on member ``d`` from ``s`` pops the same + queue and checks the tensor has the size the receiver expected. + """ + wire: dict[tuple[int, int, int], list[torch.Tensor]] = {} + for connector in connectors: + group = connector.mapping.subgroup_index + me = connector.mapping.rank_in_subgroup + + def fake_send(hidden, dst, process_group, comm_id, *, group=group, me=me): + wire.setdefault((group, me, dst), []).append(hidden.clone()) + + def fake_recv( + src, + process_group, + comm_id, + tensor_metadata, + *, + ref_tensor=None, + group=group, + me=me, + ): + queue = wire.get((group, src, me)) + assert queue, f"nothing on the wire from {src} to {me} in subgroup {group}" + tensor = queue.pop(0) + assert tuple(tensor.shape) == tuple(tensor_metadata.size) + if ref_tensor is not None: + ref_tensor.copy_(tensor) + return ref_tensor + return tensor + + def fake_send_slice(base, start, size, dst, *, group=group, me=me): + wire.setdefault((group, me, dst), []).append( + base.narrow(0, start, size).clone() + ) + + def fake_recv_slice(base, out, start, size, src, *, group=group, me=me): + queue = wire.get((group, src, me)) + assert queue, f"nothing on the wire from {src} to {me} in subgroup {group}" + tensor = queue.pop(0) + assert tuple(tensor.shape) == (size, out.shape[1]) + out.narrow(0, start, size).copy_(tensor) + + monkeypatch.setattr(connector, "_send_hidden_states", fake_send) + monkeypatch.setattr(connector, "_recv_hidden_states", fake_recv) + monkeypatch.setattr(connector, "_send_slice", fake_send_slice) + monkeypatch.setattr(connector, "_recv_slice", fake_recv_slice) + return wire + + +@pytest.mark.parametrize( + ("attention", "ffn", "tokens"), + [ + (1, 2, [5]), # A < F: one attention splits 3/2 + (2, 3, [5, 3]), # A < F: one subgroup splits, one does not + (6, 4, [1, 2, 3, 4, 5, 6]), # non-divisible A > F: uneven subgroups + (4, 2, [3, 5, 7, 11]), # historical divisible layout (regression) + (1, 4, [2]), # fewer tokens than members: dummy stages + ], +) +def test_round_trip_preserves_every_token(monkeypatch, attention, ffn, tokens): + attentions, ffns = _all_connectors(attention, ffn) + _apply(attentions + ffns, tokens) + wire = _wire(monkeypatch, attentions + ffns) + + # Distinct values per attention rank and per token so any misrouting or + # reordering shows up in the final comparison. + inputs = [ + torch.arange(t * _HIDDEN, dtype=torch.float32).reshape(t, _HIDDEN) + 1000 * a + for a, t in enumerate(tokens) + ] + for connector, hidden in zip(attentions, inputs, strict=True): + connector.send_attn_output(hidden, _attention_context(hidden.shape[0])) + + # Every FFN rank receives what arrived (or a dummy batch), "computes" + # by adding one, and returns the result. + for connector in ffns: + payload = connector.recv_attn_output(ubatch_idx=0) + connector.send_ffn_output(payload.hidden_states + 1, payload.context) + + for connector, hidden in zip(attentions, inputs, strict=True): + output = connector.recv_ffn_output( + ref_tensor=torch.empty_like(hidden), + ubatch_idx=0, + ) + assert torch.equal(output, hidden + 1) + + # Nothing sent but unreceived, nothing expected but unsent. + assert all(not queue for queue in wire.values()) + + +def test_round_trip_zero_shares_travel_as_zero_size_tensors(monkeypatch): + # 1A4F with 2 tokens: F2 and F3 get zero-size shares. Both directions + # still post the transfer (as zero-row tensors), so the pairs appear on + # the wire, the dummy members return zero rows, and nothing is left over. + attentions, ffns = _all_connectors(1, 4) + _apply(attentions + ffns, [2]) + wire = _wire(monkeypatch, attentions + ffns) + hidden = torch.arange(2 * _HIDDEN, dtype=torch.float32).reshape(2, _HIDDEN) + attentions[0].send_attn_output(hidden, _attention_context(2)) + assert [t.shape[0] for t in wire[(0, 4, 2)]] == [0] + assert [t.shape[0] for t in wire[(0, 4, 3)]] == [0] + for connector in ffns: + payload = connector.recv_attn_output(ubatch_idx=0) + connector.send_ffn_output(payload.hidden_states + 1, payload.context) + assert [t.shape[0] for t in wire[(0, 2, 4)]] == [0] + assert [t.shape[0] for t in wire[(0, 3, 4)]] == [0] + output = attentions[0].recv_ffn_output( + ref_tensor=torch.empty_like(hidden), + ubatch_idx=0, + ) + assert torch.equal(output, hidden + 1) + assert all(not queue for queue in wire.values()) diff --git a/tests/unit/connectors/test_p2p_multinode_address.py b/tests/unit/connectors/test_p2p_multinode_address.py new file mode 100644 index 00000000..3aca84c1 --- /dev/null +++ b/tests/unit/connectors/test_p2p_multinode_address.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Subgroup rendezvous address selection for multi-node P2P deployments. + +Each subgroup's StatelessProcessGroup is served by its first member binding +host:port, so members on other nodes must dial that member's own address. +The connector derives every rank's address from the interface that routes +to the configured rendezvous host. On a single host that address equals the +configured one, which keeps the historical single-host layouts unchanged. +""" + +from __future__ import annotations + +import socket + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("vllm") + +from afd_plugin.connectors.gpu.p2p import _local_address_toward # noqa: E402 + + +def test_loopback_host_yields_loopback_address(): + # The recipes use host=127.0.0.1: every rank must keep binding/dialing + # loopback exactly as before the multi-node change. + assert _local_address_toward("127.0.0.1", 6239) == "127.0.0.1" + + +def test_own_lan_address_yields_itself(): + # host = this machine's own routable address (single host, LAN IP in the + # config): the probe must return that same address. + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as probe: + probe.connect(("192.0.2.1", 9)) # TEST-NET-3, nothing is sent + own = probe.getsockname()[0] + if own.startswith("127."): + pytest.skip("no routable interface on this machine") + assert _local_address_toward(own, 6239) == own + + +def test_probe_sends_nothing_and_needs_no_listener(): + # A closed port on loopback still resolves (UDP connect only sets the + # route), so init does not depend on the master being up yet. + assert _local_address_toward("127.0.0.1", 1) == "127.0.0.1" diff --git a/tests/unit/distributed/test_topology_partition.py b/tests/unit/distributed/test_topology_partition.py new file mode 100644 index 00000000..8b80eb38 --- /dev/null +++ b/tests/unit/distributed/test_topology_partition.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Unified subgroup partition tests for arbitrary positive ``(A, F)``. + +``build_rank_mapping`` partitions the world into ``min(A, F)`` subgroups +(attention ``a`` -> subgroup ``a*G//A``, FFN ``f`` -> subgroup ``f*G//F``). +The historical ``A >= F`` divisible layout is the ``G == F`` special case, +pinned byte-for-byte by ``test_topology_snapshot.py``; this file covers what +the snapshots cannot: the structural properties that must hold on every +positive pair. + +Also guards the DP metadata channel numbering: attention ``p2p_rank`` uses +the ``ffn_size`` offset that the receiver side assumes (``p2p.py``: +``src = p2p_rank % min_size + ffn_size``); under ``A >= F`` it equals the +historical ``min_size`` offset, and under ``A < F`` it is what keeps the +numbering collision-free. +""" + +from __future__ import annotations + +import pytest + +from afd_plugin.config import AFDConfig +from afd_plugin.distributed.topology import ( + build_rank_mapping, + split_send_sizes, + validate_p2p_topology, +) + +# Every positive pair with A, F in 1..6 — divisible, non-divisible, and +# A < F shapes (36 combinations). +_FULL_GRID = [(a, f) for a in range(1, 7) for f in range(1, 7)] + + +def _config(role: str, attention: int, ffn: int) -> AFDConfig: + return AFDConfig( + role=role, + num_attention_ranks=attention, + num_ffn_ranks=ffn, + ) + + +def _mapping(role, attention, ffn, role_rank): + return build_rank_mapping(_config(role, attention, ffn), role_rank) + + +def _all_mappings(attention, ffn): + for role, size in (("ffn", ffn), ("attention", attention)): + for role_rank in range(size): + yield _mapping(role, attention, ffn, role_rank) + + +def _receiver_expected_source_rank(ffn_rank, attention, ffn): + """The p2p_rank the FFN receiver polls for its metadata sender. + + Mirrors the receiver convention in ``connectors/gpu/p2p.py`` + (``recv_dp_metadata_list``): ``src = p2p_rank % min_size + ffn_size`` + with ``p2p_rank == ffn_rank`` on FFN ranks. Kept in sync by review; the + connector cannot be imported here without torch/vLLM. + """ + return ffn_rank % min(attention, ffn) + ffn + + +# -------------------------------------------------------------------------- +# Validation — positive rank counts only +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize(("attention", "ffn"), [(1, 2), (3, 2), (2, 4), (6, 4)]) +def test_validation_accepts_any_positive_pair(attention, ffn): + # The historical A >= F and divisibility constraints are removed; the + # previously rejected shapes validate cleanly. + validate_p2p_topology(_config("attention", attention, ffn)) + + +@pytest.mark.parametrize(("attention", "ffn"), [(0, 2), (2, 0)]) +def test_validation_rejects_non_positive_pair(attention, ffn): + with pytest.raises(ValueError, match="positive rank counts"): + validate_p2p_topology(_config("attention", attention, ffn)) + + +# -------------------------------------------------------------------------- +# Partition structure on the full grid +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize(("attention", "ffn"), _FULL_GRID) +def test_partition_properties(attention, ffn): + group_count = min(attention, ffn) + groups: dict[int, tuple[int, ...]] = {} + + for mapping in _all_mappings(attention, ffn): + # Index formulas and membership consistency: every member computes + # the identical member tuple, and sits at its claimed position. + if mapping.role == "ffn": + assert mapping.subgroup_index == mapping.role_rank * group_count // ffn + else: + assert ( + mapping.subgroup_index == mapping.role_rank * group_count // attention + ) + prior = groups.setdefault(mapping.subgroup_index, mapping.subgroup_ranks) + assert prior == mapping.subgroup_ranks + assert mapping.subgroup_ranks[mapping.rank_in_subgroup] == mapping.world_rank + # ratio is this subgroup's attention member count. + attention_members = [rank for rank in mapping.subgroup_ranks if rank >= ffn] + assert mapping.ratio == len(attention_members) + + # Exactly min(A, F) subgroups, each holding at least one rank of both + # roles, members FFN-first in ascending world order, and together they + # partition the world with no overlap. + assert sorted(groups) == list(range(group_count)) + seen: list[int] = [] + for members in groups.values(): + ffn_members = [rank for rank in members if rank < ffn] + attention_members = [rank for rank in members if rank >= ffn] + assert ffn_members and attention_members + assert list(members) == sorted(ffn_members) + sorted(attention_members) + seen.extend(members) + assert sorted(seen) == list(range(ffn + attention)) + + +def test_example_2a3f_partition(): + # The shape that motivated the unified partition: 2 attention, 3 FFN + # (world order [F0, F1, F2, A0, A1]) splits into {F0, F1, A0} and + # {F2, A1}. + members = { + mapping.subgroup_index: mapping.subgroup_ranks + for mapping in _all_mappings(2, 3) + } + assert members == {0: (0, 1, 3), 1: (2, 4)} + + +@pytest.mark.parametrize( + ("attention", "ffn", "expected"), + [ + (4, 2, {0: (0, 2, 3), 1: (1, 4, 5)}), # historical divisible layout + (3, 2, {0: (0, 2, 3), 1: (1, 4)}), # non-divisible A > F + (1, 2, {0: (0, 1, 2)}), # capacity-bound target shape + ], +) +def test_partition_literal_examples(attention, ffn, expected): + members = { + mapping.subgroup_index: mapping.subgroup_ranks + for mapping in _all_mappings(attention, ffn) + } + assert members == expected + + +# -------------------------------------------------------------------------- +# DP metadata channel numbering +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize(("attention", "ffn"), _FULL_GRID) +def test_metadata_p2p_ranks_are_unique_and_contiguous(attention, ffn): + ranks = sorted( + mapping.p2p_rank + for mapping in _all_mappings(attention, ffn) + if mapping.participates_in_dp_metadata_group + ) + assert ranks == list(range(ffn)) + [ + ffn + attention_rank for attention_rank in range(min(attention, ffn)) + ] + + +@pytest.mark.parametrize(("attention", "ffn"), _FULL_GRID) +def test_sender_p2p_rank_matches_receiver_formula(attention, ffn): + min_size = min(attention, ffn) + for ffn_rank in range(ffn): + representative = ffn_rank % min_size + sender = _mapping("attention", attention, ffn, representative) + assert sender.p2p_rank == _receiver_expected_source_rank( + ffn_rank, + attention, + ffn, + ) + + +@pytest.mark.parametrize(("attention", "ffn"), _FULL_GRID) +def test_dp_metadata_destinations_cover_every_ffn_once(attention, ffn): + destinations: list[int] = [] + for role_rank in range(attention): + mapping = _mapping("attention", attention, ffn, role_rank) + destinations.extend(mapping.dp_metadata_destinations) + assert sorted(destinations) == list(range(ffn)) + + +# -------------------------------------------------------------------------- +# split_send_sizes — the per-peer size arithmetic the transport shares +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("token_count", "parts", "expected"), + [ + (5, 2, (3, 2)), # remainder rides on the leading shares + (6, 3, (2, 2, 2)), # even split + (7, 1, (7,)), # parts=1 = whole count (every A >= F subgroup) + (2, 4, (1, 1, 0, 0)), # fewer tokens than peers -> trailing zeros + (0, 3, (0, 0, 0)), # nothing to send + ], +) +def test_split_send_sizes_literal_cases(token_count, parts, expected): + assert split_send_sizes(token_count, parts) == expected + + +@pytest.mark.parametrize("token_count", range(9)) +@pytest.mark.parametrize("parts", range(1, 5)) +def test_split_send_sizes_properties(token_count, parts): + sizes = split_send_sizes(token_count, parts) + # Conservation, near-evenness, and deterministic front-loading — the + # contract that lets sender and receivers agree without communication. + assert sum(sizes) == token_count + assert len(sizes) == parts + assert max(sizes) - min(sizes) <= 1 + assert list(sizes) == sorted(sizes, reverse=True) + + +@pytest.mark.parametrize(("token_count", "parts"), [(1, 0), (1, -1), (-1, 2)]) +def test_split_send_sizes_rejects_invalid_inputs(token_count, parts): + with pytest.raises(ValueError): + split_send_sizes(token_count, parts) diff --git a/tests/unit/distributed/test_topology_snapshot.py b/tests/unit/distributed/test_topology_snapshot.py new file mode 100644 index 00000000..aff1fdd9 --- /dev/null +++ b/tests/unit/distributed/test_topology_snapshot.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Characterization snapshots of the P2P rank mapping. + +These tests freeze the *current* ``build_rank_mapping`` behavior for legal +``A >= F`` topologies so that the M2N generalization work can prove, at every +step, that existing gather-mode mappings remain byte-for-byte unchanged. + +Two layers of protection: + +1. Literal golden snapshots for representative topologies (1A1F, 2A2F, 4A2F). +2. Structural invariants over a wider grid of legal ``(A, F)`` pairs. +""" + +from __future__ import annotations + +import dataclasses + +import pytest + +from afd_plugin.config import AFDConfig +from afd_plugin.distributed.topology import ( + build_rank_mapping, +) + + +def _config(role: str, attention: int, ffn: int) -> AFDConfig: + return AFDConfig( + role=role, + num_attention_ranks=attention, + num_ffn_ranks=ffn, + ) + + +def _mapping_dict(role: str, attention: int, ffn: int, role_rank: int) -> dict: + mapping = build_rank_mapping(_config(role, attention, ffn), role_rank) + return dataclasses.asdict(mapping) + + +# Captured from main @ 603c111 (pre-M2N). Do not regenerate to make a failing +# test pass: a diff here means gather-mode behavior changed. +_GOLDEN: dict[tuple[int, int], dict[tuple[str, int], dict]] = { + (1, 1): { + ("ffn", 0): { + "role": "ffn", + "role_rank": 0, + "world_rank": 0, + "p2p_rank": 0, + "attention_size": 1, + "ffn_size": 1, + "min_size": 1, + "ratio": 1, + "subgroup_index": 0, + "rank_in_subgroup": 0, + "subgroup_ranks": (0, 1), + "dp_metadata_destinations": (), + }, + ("attention", 0): { + "role": "attention", + "role_rank": 0, + "world_rank": 1, + "p2p_rank": 1, + "attention_size": 1, + "ffn_size": 1, + "min_size": 1, + "ratio": 1, + "subgroup_index": 0, + "rank_in_subgroup": 1, + "subgroup_ranks": (0, 1), + "dp_metadata_destinations": (0,), + }, + }, + (2, 2): { + ("ffn", 0): { + "role": "ffn", + "role_rank": 0, + "world_rank": 0, + "p2p_rank": 0, + "attention_size": 2, + "ffn_size": 2, + "min_size": 2, + "ratio": 1, + "subgroup_index": 0, + "rank_in_subgroup": 0, + "subgroup_ranks": (0, 2), + "dp_metadata_destinations": (), + }, + ("ffn", 1): { + "role": "ffn", + "role_rank": 1, + "world_rank": 1, + "p2p_rank": 1, + "attention_size": 2, + "ffn_size": 2, + "min_size": 2, + "ratio": 1, + "subgroup_index": 1, + "rank_in_subgroup": 0, + "subgroup_ranks": (1, 3), + "dp_metadata_destinations": (), + }, + ("attention", 0): { + "role": "attention", + "role_rank": 0, + "world_rank": 2, + "p2p_rank": 2, + "attention_size": 2, + "ffn_size": 2, + "min_size": 2, + "ratio": 1, + "subgroup_index": 0, + "rank_in_subgroup": 1, + "subgroup_ranks": (0, 2), + "dp_metadata_destinations": (0,), + }, + ("attention", 1): { + "role": "attention", + "role_rank": 1, + "world_rank": 3, + "p2p_rank": 3, + "attention_size": 2, + "ffn_size": 2, + "min_size": 2, + "ratio": 1, + "subgroup_index": 1, + "rank_in_subgroup": 1, + "subgroup_ranks": (1, 3), + "dp_metadata_destinations": (1,), + }, + }, + (4, 2): { + ("ffn", 0): { + "role": "ffn", + "role_rank": 0, + "world_rank": 0, + "p2p_rank": 0, + "attention_size": 4, + "ffn_size": 2, + "min_size": 2, + "ratio": 2, + "subgroup_index": 0, + "rank_in_subgroup": 0, + "subgroup_ranks": (0, 2, 3), + "dp_metadata_destinations": (), + }, + ("ffn", 1): { + "role": "ffn", + "role_rank": 1, + "world_rank": 1, + "p2p_rank": 1, + "attention_size": 4, + "ffn_size": 2, + "min_size": 2, + "ratio": 2, + "subgroup_index": 1, + "rank_in_subgroup": 0, + "subgroup_ranks": (1, 4, 5), + "dp_metadata_destinations": (), + }, + ("attention", 0): { + "role": "attention", + "role_rank": 0, + "world_rank": 2, + "p2p_rank": 2, + "attention_size": 4, + "ffn_size": 2, + "min_size": 2, + "ratio": 2, + "subgroup_index": 0, + "rank_in_subgroup": 1, + "subgroup_ranks": (0, 2, 3), + "dp_metadata_destinations": (0,), + }, + ("attention", 1): { + "role": "attention", + "role_rank": 1, + "world_rank": 3, + "p2p_rank": 3, + "attention_size": 4, + "ffn_size": 2, + "min_size": 2, + "ratio": 2, + "subgroup_index": 0, + "rank_in_subgroup": 2, + "subgroup_ranks": (0, 2, 3), + "dp_metadata_destinations": (1,), + }, + ("attention", 2): { + "role": "attention", + "role_rank": 2, + "world_rank": 4, + "p2p_rank": 4, + "attention_size": 4, + "ffn_size": 2, + "min_size": 2, + "ratio": 2, + "subgroup_index": 1, + "rank_in_subgroup": 1, + "subgroup_ranks": (1, 4, 5), + "dp_metadata_destinations": (), + }, + ("attention", 3): { + "role": "attention", + "role_rank": 3, + "world_rank": 5, + "p2p_rank": 5, + "attention_size": 4, + "ffn_size": 2, + "min_size": 2, + "ratio": 2, + "subgroup_index": 1, + "rank_in_subgroup": 2, + "subgroup_ranks": (1, 4, 5), + "dp_metadata_destinations": (), + }, + }, +} + +_GOLDEN_CASES = [ + (attention, ffn, role, role_rank) + for (attention, ffn), mappings in _GOLDEN.items() + for (role, role_rank) in mappings +] + +_LEGAL_GRID = [ + (1, 1), + (2, 1), + (2, 2), + (3, 1), + (4, 1), + (4, 2), + (4, 4), + (6, 2), + (6, 3), + (8, 2), + (8, 4), +] + + +@pytest.mark.parametrize( + ("attention", "ffn", "role", "role_rank"), + _GOLDEN_CASES, +) +def test_rank_mapping_matches_golden_snapshot(attention, ffn, role, role_rank): + assert ( + _mapping_dict(role, attention, ffn, role_rank) + == _GOLDEN[(attention, ffn)][(role, role_rank)] + ) + + +@pytest.mark.parametrize(("attention", "ffn"), _LEGAL_GRID) +def test_rank_mapping_invariants(attention, ffn): + ratio = attention // ffn + min_size = min(attention, ffn) + + seen_subgroups: dict[int, tuple[int, ...]] = {} + dp_destination_union: list[int] = [] + + for role, size in (("ffn", ffn), ("attention", attention)): + for role_rank in range(size): + mapping = build_rank_mapping(_config(role, attention, ffn), role_rank) + + if role == "ffn": + assert mapping.world_rank == role_rank + assert mapping.p2p_rank == role_rank + assert mapping.subgroup_index == role_rank + else: + assert mapping.world_rank == ffn + role_rank + assert mapping.p2p_rank == role_rank + min_size + assert mapping.subgroup_index == role_rank // ratio + + assert mapping.ratio == ratio + assert mapping.min_size == min_size + assert len(mapping.subgroup_ranks) == 1 + ratio + assert mapping.subgroup_ranks[0] == mapping.subgroup_index + assert ( + mapping.subgroup_ranks[mapping.rank_in_subgroup] == mapping.world_rank + ) + + prior = seen_subgroups.setdefault( + mapping.subgroup_index, + mapping.subgroup_ranks, + ) + assert prior == mapping.subgroup_ranks + + dp_destination_union.extend(mapping.dp_metadata_destinations) + + # The DP metadata senders must cover every FFN rank exactly once. + assert sorted(dp_destination_union) == list(range(ffn)) + # Subgroups must partition the attention ranks. + attention_members = sorted( + rank for ranks in seen_subgroups.values() for rank in ranks[1:] + ) + assert attention_members == list(range(ffn, ffn + attention)) + + +# The historical constraint-characterization test (A < F and non-divisible +# pairs raising ValueError) was removed together with the constraint itself; +# test_topology_partition.py now pins the relaxed validation behavior. diff --git a/tests/unit/v1/worker/test_ffn_metadata.py b/tests/unit/v1/worker/test_ffn_metadata.py new file mode 100644 index 00000000..707e07ca --- /dev/null +++ b/tests/unit/v1/worker/test_ffn_metadata.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""FFN-side token counts feed the EP dispatch sizes; they must follow the +subgroup partition exactly as the connector's receive sizes do.""" + +from __future__ import annotations + +import pytest + +from afd_plugin.v1.worker.ffn_metadata import ( + aggregate_ffn_token_counts, + project_ffn_token_counts_to_dp, +) + + +@pytest.mark.parametrize( + ("attention_counts", "attention", "ffn", "expected"), + [ + # Historical A >= F layouts: unchanged (docstring example and 1:1). + ((0, 4, 5, 6), 4, 2, (5, 11)), + ((3, 5), 2, 2, (3, 5)), + ((), 4, 2, (2, 2)), + # TP-expanded counts (one count per DP rank, two TP ranks each). + ((4, 6), 4, 2, (8, 12)), + # A < F: one attention splits across its FFN members. + ((5,), 1, 2, (3, 2)), + ((5, 3), 2, 4, (3, 2, 2, 1)), + # Fewer tokens than members: zero shares become one-token dummies. + ((2,), 1, 4, (1, 1, 1, 1)), + # Non-divisible A > F: uneven subgroups {A0,A1}{A2}. + ((1, 2, 3), 3, 2, (3, 3)), + ((1, 2, 3, 4, 5, 6), 6, 4, (3, 3, 9, 6)), + ], +) +def test_aggregate_follows_the_subgroup_partition( + attention_counts, attention, ffn, expected +): + assert ( + aggregate_ffn_token_counts( + attention_counts, attention_size=attention, ffn_size=ffn + ) + == expected + ) + + +def test_project_collapses_tp_ranks_to_dp_ranks(): + assert project_ffn_token_counts_to_dp((8, 8, 13, 13), dp_size=2) == (8, 13) + assert project_ffn_token_counts_to_dp((3, 2), dp_size=2) == (3, 2)