[Feature] Support arbitrary attention-to-FFN rank ratios in P2pNcclAFDConnector

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",
 ]

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,28 +151,40 @@ 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(
                 "FFN role rank must be within FFN size "
                 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",
 ]

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(