[Feature] Relax the divisibility constraint in P2pNcclAFDConnector

README.md

@@ -54,7 +54,7 @@ See the [recipe index](recipe/README.md) for deployment and benchmark examples.
 
 | Connector | Platform | Recommend Stage | Sync or Async | Graph Support | Notes |
 | --- | --- | --- | --- | --- | --- |
-| `P2pNcclAFDConnector` | CUDA | Decode | Sync | `FULL_DECODE_ONLY` CUDA graph | FFN ranks are ordered before Attention ranks. `num_attention_ranks` must be greater than or equal to `num_ffn_ranks` and divisible by it. See the [DeepSeek V2 Lite recipe](recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/README.md). |
+| `P2pNcclAFDConnector` | CUDA | Decode | Sync | `FULL_DECODE_ONLY` CUDA graph | FFN ranks are ordered before Attention ranks. `num_attention_ranks` must be greater than or equal to `num_ffn_ranks`; the Attention ranks are spread over the FFN ranks in blocks that differ in size by at most one. See the [DeepSeek V2 Lite recipe](recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/README.md). |
 | `CAMP2pAFDConnector` | Ascend NPU | Decode | Sync | `FULL_DECODE_ONLY` ACL graph | Uses HCCL/CAMP2P custom ops. Ascend ops build by default on NPU platforms. See the [synchronous DeepSeek V3.2 recipe](recipe/npu/CAMP2pAFDConnector/deepseek_v3_2/README.md). |
 | `CAMAsyncAFDConnector` | Ascend NPU | Prefill / decode | Async | Not supported | Experimental v0.26 DP+TP/SP path with AFD-managed two-stage MoE ubatching; native DBO and PCP are unsupported. Post-fix DeepSeek-V3.2 DP2TP8+EP16 token split reached `0.9522` strict match on the complete GSM8K evaluation. The [legacy PCP8 recipe](recipe/npu/CAMAsyncAFDConnector/deepseek_v3_2/README.md) requires `release/v0.19.1rc1`. |
 

afd_plugin/connectors/gpu/p2p.py

@@ -15,7 +15,9 @@
     consecutive Attention ranks, which requires::
 
         num_attention_ranks >= num_ffn_ranks
-        num_attention_ranks % num_ffn_ranks == 0
+
+    The Attention ranks are spread over the FFN ranks in blocks that differ
+    in size by at most one, so the counts need not divide evenly.
 
     Attention sends hidden states to its mapped FFN rank; the FFN rank
     concatenates inputs from its Attention peers, runs FFN work, splits the
@@ -184,7 +186,7 @@ def __init__(
         self.attn_size = self.mapping.attention_size
         self.ffn_size = self.mapping.ffn_size
         self.min_size = self.mapping.min_size
-        self.ratio = self.mapping.ratio
+        self.ratio = len(self.mapping.subgroup_ranks) - 1
         self.group_size = len(self.mapping.subgroup_ranks)
         self.dst_list = list(self.mapping.dp_metadata_destinations)
         text_config = vllm_config.model_config.hf_text_config
@@ -751,10 +753,10 @@ def update_state_from_dp_metadata(
                 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}")
+                    # Subgroups need not be the same size, so read the peer
+                    # from the roster instead of assuming a uniform ratio.
                     attention_rank = (
-                        connector.mapping.subgroup_index * connector.ratio
-                        + src_rank
-                        - 1
+                        connector.mapping.subgroup_ranks[src_rank] - connector.ffn_size
                     )
 
                     tensor_metadata = _TensorMetadata(

afd_plugin/distributed/__init__.py

@@ -6,6 +6,7 @@
     AFDRankMapping,
     build_rank_mapping,
     resolve_role_rank,
+    subgroup_attention_block,
     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",
+    "subgroup_attention_block",
     "topology_from_config",
     "validate_p2p_topology",
 ]

afd_plugin/distributed/topology.py

@@ -29,7 +29,6 @@ class AFDRankMapping:
     attention_size: int
     ffn_size: int
     min_size: int
-    ratio: int
     subgroup_index: int
     rank_in_subgroup: int
     subgroup_ranks: tuple[int, ...]
@@ -57,12 +56,6 @@ def validate_p2p_topology(config: AFDConfig) -> None:
             "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}",
-        )
 
 
 def resolve_role_rank(vllm_config: VllmConfig, config: AFDConfig) -> int:
@@ -112,6 +105,16 @@ def resolve_role_rank(vllm_config: VllmConfig, config: AFDConfig) -> int:
     return role_rank
 
 
+def subgroup_attention_block(
+    subgroup_index: int, attention_size: int, ffn_size: int
+) -> range:
+    """Attention role ranks owned by one FFN rank: contiguous, sizes within one."""
+    return range(
+        (subgroup_index * attention_size + ffn_size - 1) // ffn_size,
+        ((subgroup_index + 1) * attention_size + ffn_size - 1) // ffn_size,
+    )
+
+
 def build_rank_mapping(
     config: AFDConfig,
     role_rank: int,
@@ -130,7 +133,7 @@ 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)
+        subgroup_index = role_rank * ffn_size // attention_size
     elif config.role == "ffn":
         if role_rank >= ffn_size:
             raise ValueError(
@@ -142,14 +145,19 @@ def build_rank_mapping(
     else:
         raise ValueError(f"unknown AFD role {config.role!r}")
 
-    ratio = attention_size // ffn_size
+    # Balanced block distribution: Attention rank ``a`` joins the subgroup of
+    # FFN rank ``a * F // A``, and ``subgroup_attention_block`` is the reverse
+    # lookup that lists a subgroup's Attention ranks. The blocks are contiguous and
+    # differ in size by at most one, so an integral ratio is no longer
+    # required; when A is a multiple of F this is the same grouping as before.
     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)],
-    )
+    subgroup_attention_ranks = [
+        ffn_size + attention_rank
+        for attention_rank in subgroup_attention_block(
+            subgroup_index, attention_size, ffn_size
+        )
+    ]
+    subgroup_ranks = tuple([subgroup_index] + subgroup_attention_ranks)
     rank_in_subgroup = subgroup_ranks.index(world_rank)
     p2p_rank = role_rank + min_size if config.role == "attention" else role_rank
 
@@ -169,7 +177,6 @@ def build_rank_mapping(
         attention_size=attention_size,
         ffn_size=ffn_size,
         min_size=min_size,
-        ratio=ratio,
         subgroup_index=subgroup_index,
         rank_in_subgroup=rank_in_subgroup,
         subgroup_ranks=subgroup_ranks,

afd_plugin/v1/worker/ffn_metadata.py

@@ -4,6 +4,8 @@
 
 from __future__ import annotations
 
+from afd_plugin.distributed import subgroup_attention_block
+
 
 def aggregate_ffn_token_counts(
     attention_counts: tuple[int, ...],
@@ -14,14 +16,17 @@ def aggregate_ffn_token_counts(
 ) -> tuple[int, ...]:
     """Aggregate consecutive Attention-rank counts for each FFN rank.
 
-    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)``.
+    The blocks are the ones the rank mapping builds: contiguous and differing
+    in size by at most one, so ``3A2F`` groups ``{A0, A1}`` onto ``F0`` and
+    ``{A2}`` onto ``F1``. 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)``.
     """
 
     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:
+    if ffn_size <= 0 or attention_size < ffn_size:
         return fallback_counts
 
     expanded_counts = attention_counts
@@ -36,15 +41,14 @@ def aggregate_ffn_token_counts(
             for rank in range(attention_size)
         )
 
-    group_size = attention_size // ffn_size
+    # Each FFN rank sums the block that ``build_rank_mapping`` assigns to it.
     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,
+            for attention_rank in subgroup_attention_block(
+                ffn_rank, attention_size, ffn_size
             )
         )
         for ffn_rank in range(ffn_size)

docs/design/module/connector_contracts.md

@@ -110,8 +110,9 @@ synchronous NPU runtime requires both common and connector-local values to be
 | `CAMAsyncAFDConnector` | Ascend | Attention ranks, then FFN ranks | `None`; routing/token metadata travels with CAM dispatch payloads | `connector.control_plane is None` |
 
 The CUDA P2P mapping requires
-`num_attention_ranks >= num_ffn_ranks` and an integral A/F ratio. Each FFN rank
-owns a subgroup containing itself and consecutive Attention peers. CAMP2P also
+`num_attention_ranks >= num_ffn_ranks`. Each FFN rank owns a subgroup containing
+itself and consecutive Attention peers; the peers are distributed in blocks that
+differ in size by at most one, so an integral A/F ratio is not required. CAMP2P also
 requires at least as many Attention ranks as FFN ranks; its control and HCCL
 groups remain connector-owned. CAM async maps role ranks directly into an
 Attention-first world and distributes routed experts across FFN ranks.
@@ -258,7 +259,7 @@ metadata is present.
 
 | Connector | Connector-owned resources | Topology constraints and mapping |
 | --- | --- | --- |
-| CUDA P2P | AFD process group, PyNccl data communicators, separate NCCL metadata group, compiled custom-op communicator registry, and graph-oriented receive buffers/state. | Requires `A >= F` and `A % F == 0`. One FFN rank is grouped with a consecutive block of `A/F` Attention ranks in FFN-first ordering. |
+| CUDA P2P | AFD process group, PyNccl data communicators, separate NCCL metadata group, compiled custom-op communicator registry, and graph-oriented receive buffers/state. | Requires `A >= F`. One FFN rank is grouped with a consecutive block of Attention ranks in FFN-first ordering; the blocks differ in size by at most one, and hold `A/F` ranks each when `F` divides `A`. |
 | Ascend CAMP2P | AFD process group, one HCCL communication group per ubatch, FFN HCCL state, Gloo metadata group, custom-op state and transfer handles. | FFN-first ordering and `A >= F`; group construction derives each FFN/Attention mapping. |
 | Ascend CAM async | Attention-first HCCL group, external CAM operator state, per-stage pending Attention payload queues, and connector work-item state. | Role ranks map into a combined Attention-first world; CAM tensor metadata determines actual layer and routed/shared token counts. |
 

docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md

@@ -14,7 +14,7 @@ It supports both prefill and decode which all support eager mode. CUDA graph sup
 
 ## How it works
 
-Throughout this section, let `A = num_attention_ranks`, `F = num_ffn_ranks`, and `ratio = A / F`. The topology rules (`A >= F`, `A % F == 0`) guarantee `ratio` is a whole number and make `min_size = min(A, F) = F`. One physical process sits at up to three different rank numbers — an AFD world rank, a subgroup rank, and a control-plane (`p2p`) rank — all derived deterministically from the role, role rank, and topology counts.
+Throughout this section, let `A = num_attention_ranks` and `F = num_ffn_ranks`. The topology rule (`A >= F`) makes `min_size = min(A, F) = F`. Each subgroup's `ratio` is the number of Attention peers it holds; when `F` divides `A` every subgroup holds `A / F` of them, and otherwise the counts differ by one. One physical process sits at up to three different rank numbers — an AFD world rank, a subgroup rank, and a control-plane (`p2p`) rank — all derived deterministically from the role, role rank, and topology counts.
 
 ### AFD world (shared rendezvous)
 
@@ -29,7 +29,7 @@ FFN role rank `i` gets world rank `i`; Attention role rank `j` gets world rank `
 
 ### Data plane: one subgroup per FFN rank
 
-Each FFN rank `k` owns subgroup `k`, containing itself plus its `ratio` consecutive Attention peers `A(k*ratio) .. A(k*ratio + ratio - 1)`. Inside a subgroup the FFN rank is always subgroup rank `0` and the Attention peers occupy subgroup ranks `1..ratio`.
+Each FFN rank `k` owns subgroup `k`, containing itself plus the consecutive Attention peers assigned to it: Attention rank `a` joins subgroup `a * F // A`, which is `a // ratio` whenever `F` divides `A`. Inside a subgroup the FFN rank is always subgroup rank `0` and the Attention peers occupy subgroup ranks `1..ratio`.
 
 Each subgroup is its own process group, rendezvoused on the AFD world's store under a `PrefixStore` of its own so no extra port is needed, carrying two NCCL communicators: Attention-to-FFN for hidden states and FFN-to-Attention for FFN outputs. Per layer/stage, each Attention peer sends its hidden states to subgroup rank `0`; the FFN rank receives from ranks `1..ratio` in order, concatenates along the token dimension, runs FFN work, splits the output by the recorded sequence lengths, and sends each slice back to the originating Attention rank. The data path uses vLLM `PyNcclCommunicator.send()` / `recv()` on the current CUDA stream.
 
@@ -105,10 +105,9 @@ role rank.
 
 ```text
 num_attention_ranks >= num_ffn_ranks
-num_attention_ranks % num_ffn_ranks == 0
 ```
 
-Therefore, every FFN rank maps to the same integer number of consecutive Attention ranks.
+The Attention ranks are then spread over the FFN ranks in consecutive blocks that differ in size by at most one: Attention rank `a` joins the subgroup of FFN rank `a * F // A`. When `F` divides `A` every FFN rank maps to the same number of Attention ranks, as before.
 
 Examples:
 
@@ -118,7 +117,7 @@ Examples:
 | `2A2F` | Yes | `F0 <-> A0`, `F1 <-> A1` |
 | `4A2F` | Yes | `F0 <-> A0,A1`, `F1 <-> A2,A3` |
 | `1A2F` | No | Attention rank count is smaller than FFN rank count. |
-| `3A2F` | No | Attention rank count is not divisible by FFN rank count. |
+| `3A2F` | Yes | `F0 <-> A0,A1`, `F1 <-> A2` |
 
 ## Minimal launch shape
 

tests/unit/connectors/test_p2p_connector.py

@@ -184,7 +184,6 @@ def test_p2p_topology_supports_equal_and_integer_multiple_attention_counts(
         role_rank,
     )
 
-    assert mapping.ratio == attention_size // ffn_size
     assert mapping.subgroup_ranks == subgroup_ranks
     assert mapping.dp_metadata_destinations == dsts
 
@@ -343,6 +342,9 @@ def record_ffn_send(
         (4, 2, 0, [3, 5, 7, 11], [3, 5]),
         (4, 2, 1, [3, 5, 7, 0], [7, 1]),
         (6, 3, 2, [2, 3, 5, 7, 11, 13], [11, 13]),
+        # A % F != 0: subgroup 0 holds A0 and A1, subgroup 1 holds A2 alone.
+        (3, 2, 0, [3, 5, 7], [3, 5]),
+        (3, 2, 1, [3, 5, 7], [7]),
     ],
 )
 def test_p2p_ffn_metadata_tracks_each_attention_peer_in_xayf(
@@ -440,14 +442,6 @@ def test_p2p_tensor_metadata_clamps_idle_attention_rank_to_dummy_token():
             },
             "num_attention_ranks >= num_ffn_ranks",
         ),
-        (
-            {
-                "connector": "P2pNcclAFDConnector",
-                "num_attention_ranks": 3,
-                "num_ffn_ranks": 2,
-            },
-            "multiple of num_ffn_ranks",
-        ),
     ],
 )
 def test_p2p_topology_validation_errors_are_clear(raw, message):

tests/unit/distributed/test_topology_partition.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project
+"""Subgroup partition tests for the P2P rank mapping.
+
+``build_rank_mapping`` spreads the Attention ranks over the FFN ranks in
+contiguous blocks that differ in size by at most one (attention ``a`` joins
+the subgroup of FFN rank ``a * F // A``). When ``F`` divides ``A`` this is the
+grouping the connector has always built, which
+``test_divisible_layouts_keep_the_historical_grouping`` pins; the other tests
+cover every ``A >= F`` pair, including the layouts that grouping could not
+express.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from afd_plugin.config import AFDConfig
+from afd_plugin.distributed.topology import build_rank_mapping
+
+# Every A >= F pair with A, F in 1..8, divisible and not.
+_GRID = [(a, f) for a in range(1, 9) for f in range(1, a + 1)]
+
+
+def _config(role: str, attention: int, ffn: int) -> AFDConfig:
+    return AFDConfig(
+        role=role,
+        connector="P2pNcclAFDConnector",
+        num_attention_ranks=attention,
+        num_ffn_ranks=ffn,
+    )
+
+
+def _mapping(role: str, role_rank: int, attention: int, ffn: int):
+    return build_rank_mapping(_config(role, attention, ffn), role_rank)
+
+
+@pytest.mark.parametrize(("attention", "ffn"), _GRID)
+def test_every_attention_rank_belongs_to_exactly_one_subgroup(attention, ffn):
+    subgroups = {
+        ffn_rank: _mapping("ffn", ffn_rank, attention, ffn).subgroup_ranks
+        for ffn_rank in range(ffn)
+    }
+
+    # The FFN rank leads its own subgroup, and the Attention members partition
+    # the Attention world in world order.
+    assert [ranks[0] for ranks in subgroups.values()] == list(range(ffn))
+    peers = [rank for ranks in subgroups.values() for rank in ranks[1:]]
+    assert peers == list(range(ffn, ffn + attention))
+
+    # Block sizes differ by at most one, and none is empty.
+    sizes = [len(ranks) - 1 for ranks in subgroups.values()]
+    assert min(sizes) >= 1
+    assert max(sizes) - min(sizes) <= 1
+
+
+@pytest.mark.parametrize(("attention", "ffn"), _GRID)
+def test_both_roles_agree_on_the_subgroup_they_share(attention, ffn):
+    for attention_rank in range(attention):
+        mapping = _mapping("attention", attention_rank, attention, ffn)
+        owner = _mapping("ffn", mapping.subgroup_index, attention, ffn)
+        assert mapping.subgroup_ranks == owner.subgroup_ranks
+        assert mapping.subgroup_ranks[mapping.rank_in_subgroup] == mapping.world_rank
+
+
+@pytest.mark.parametrize(("attention", "ffn"), _GRID)
+def test_world_and_p2p_ranks_follow_the_role_layout(attention, ffn):
+    min_size = min(attention, ffn)
+    dp_destinations: list[int] = []
+    for role, size in (("ffn", ffn), ("attention", attention)):
+        for role_rank in range(size):
+            mapping = _mapping(role, role_rank, attention, ffn)
+            assert mapping.min_size == min_size
+            if role == "ffn":
+                assert mapping.world_rank == role_rank
+                assert mapping.p2p_rank == role_rank
+            else:
+                assert mapping.world_rank == ffn + role_rank
+                assert mapping.p2p_rank == role_rank + min_size
+            dp_destinations.extend(mapping.dp_metadata_destinations)
+
+    # Only Attention ranks send DP metadata, and each FFN rank receives it
+    # from exactly one of them.
+    assert sorted(dp_destinations) == list(range(ffn))
+
+
+@pytest.mark.parametrize(
+    ("attention", "ffn", "expected"),
+    [
+        (3, 2, [(0, 2, 3), (1, 4)]),
+        (5, 2, [(0, 2, 3, 4), (1, 5, 6)]),
+        (5, 3, [(0, 3, 4), (1, 5, 6), (2, 7)]),
+        (6, 4, [(0, 4, 5), (1, 6), (2, 7, 8), (3, 9)]),
+    ],
+)
+def test_partition_literal_examples(attention, ffn, expected):
+    assert [
+        _mapping("ffn", ffn_rank, attention, ffn).subgroup_ranks
+        for ffn_rank in range(ffn)
+    ] == expected
+
+
+@pytest.mark.parametrize(("attention", "ffn"), [(a, f) for a, f in _GRID if a % f == 0])
+def test_divisible_layouts_keep_the_historical_grouping(attention, ffn):
+    ratio = attention // ffn
+    for ffn_rank in range(ffn):
+        mapping = _mapping("ffn", ffn_rank, attention, ffn)
+        assert len(mapping.subgroup_ranks) - 1 == ratio
+        assert mapping.subgroup_ranks == (
+            ffn_rank,
+            *(ffn + ffn_rank * ratio + offset for offset in range(ratio)),
+        )