[Feature] Relax the divisibility constraint in P2pNcclAFDConnector

[swjeong9]

Purpose

This relaxes the A % F == 0 constraint in P2pNcclAFDConnector. The subgroup partition changes from a // (A // F) to a * F // A.

Scope

In scope

• distributed/topology.py: drop the divisibility check, replace the subgroup partition, read ratio from the subgroup roster
• v1/worker/ffn_metadata.py: aggregate per-FFN token counts with the same partition
• connectors/gpu/p2p.py: look up the control-plane peer in the subgroup roster instead of computing it from a uniform ratio

Out of scope

• A < F support
• The multi-node rendezvous fix. We will open a separate issue for it.

Test Plan

test_topology_snapshot.py pins the rank mapping of the existing topologies, and test_topology_partition.py runs every combination of A and F from 1 to 8 to check the properties of the partition. test_ffn_metadata.py and test_p2p_connector.py gained 3A2F cases.

The GPU runs are on AWS EKS with two g6.12xlarge nodes (4x L4 24GB each), Attention workers on one node and FFN workers on the other. Topologies are 2A2F, 3A2F and 4A2F.

A=3; F=2                            # 2A2F, 3A2F, 4A2F
MODEL=deepseek-ai/DeepSeek-V2-Lite

COMMON=(--served-model-name deepseek-v2-lite --trust-remote-code --enable-expert-parallel
        --max-model-len 4096 --max-num-seqs 128 --max-num-batched-tokens 128
        --gpu-memory-utilization 0.9 --host 0.0.0.0 --data-parallel-size-local 1)
GRAPH=(--max-cudagraph-capture-size 128
       --compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_capture_sizes":[16,32,48,64,80,96,112,128]}')
# for eager mode, use --enforce-eager in place of GRAPH
afd() { printf '{"afd":{"role":"%s","connector":"P2pNcclAFDConnector","host":"%s","port":6269,"num_attention_ranks":%d,"num_ffn_ranks":%d}}' "$1" "$FFN0_IP" "$A" "$F"; }

# Attention node, rank r = 0..A-1
vllm serve $MODEL "${COMMON[@]}" "${GRAPH[@]}" \
  --port 18305 --data-parallel-size $A \
  --data-parallel-address $ATTN0_IP --data-parallel-rpc-port 13345 \
  --additional-config "$(afd attention)"

# FFN node, rank r = 0..F-1
vllm serve $MODEL "${COMMON[@]}" "${GRAPH[@]}" \
  --port 18400 --data-parallel-size $F \
  --data-parallel-address $FFN0_IP --data-parallel-rpc-port 13346 \
  --additional-config "$(afd ffn)"

# both roles: for r > 0, add --headless --data-parallel-start-rank $r

vllm bench serve --backend vllm --host $ATTN0_IP --port 18305 \
  --model deepseek-v2-lite --tokenizer $MODEL --trust-remote-code \
  --dataset-name random --random-input-len 1 --random-output-len 256 --ignore-eos \
  --request-rate inf --max-concurrency $((64 * A)) --num-prompts $((128 * A)) \
  --num-warmups 8 --percentile-metrics ttft,tpot,itl,e2el --metric-percentiles 50,90,99

Test Result

Times are in ms, throughput in tok/s.

Topology	Mode	Build	TTFT P50	TTFT P99	ITL P50	ITL P99	Throughput
2A2F	eager	main	264	980	107.4	112.9	1155.0
2A2F	eager	this PR	827	1082	108.9	114.9	1121.2
2A2F	graph	main	291	372	108.6	113.7	1169.7
2A2F	graph	this PR	314	467	106.9	110.9	1187.3
4A2F	eager	main	989	1070	122.5	131.2	2016.5
4A2F	eager	this PR	787	983	121.7	133.5	2055.8
4A2F	graph	main	384	525	124.5	132.4	2036.7
4A2F	graph	this PR	372	505	125.0	132.7	2026.6
3A2F	eager	main	ValueError
3A2F	eager	this PR	837	867	116.1	123.0	1601.2
3A2F	graph	main	ValueError
3A2F	graph	this PR	314	413	118.6	124.7	1620.1

3A2F on main:

ValueError: P2pNcclAFDConnector currently requires num_attention_ranks
to be a multiple of num_ffn_ranks, got 3 and 2

For the graph replay check you asked for, we re-ran 2A2F and 3A2F with VLLM_LOGGING_LEVEL=DEBUG.

Capturing CUDA graphs (decode, FULL): 100%|██████████| 8/8

Running batch with cudagraph_mode: FULL, batch_descriptor: BatchDescriptor(
  num_tokens=80, num_reqs=80, uniform=True, ...), num_tokens_across_dp: tensor([80, 80, 80])

On both topologies all 776 decode steps on Attention rank 0 report cudagraph_mode: FULL, with no NONE. On 3A2F the batch lands on 16, 64 and 80, and all three are served by a captured graph.

Docs Impact

README.md, docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md and docs/design/module/connector_contracts.md drop the divisibility rule and gain a 3A2F row in the example table.

---

Essential PR Checklist

• Purpose is clear and linked to public context when possible.
• Scope is bounded.
• Compatibility with vLLM v0.26.0 is considered.
• No changes are made to the vLLM source checkout.
• Plugin-owned classes or explicit dotted class paths are preferred over monkey patches.
• Any compat shim or monkey patch is isolated, idempotent, version-guarded, documented, and tested. (n/a - none added)
• Imports remain CPU-safe; CUDA-heavy work is delayed or GPU-gated.
• Validation evidence is included, including skipped GPU tests when applicable.
• Documentation impact is stated.

[jiangkuaixue123]

cc @specture724

[jiangkuaixue123]

Thanks for sharing the implementation and measurements. I am not sure eager-mode AFD performance with P2pNcclAFDConnector is particularly meaningful for its intended use case: the decode stage of a prefill/decode-disaggregated deployment, with CUDA graphs in FULL_DECODE_ONLY mode.

So far, we have not identified a concrete use case for A < F. Our typical deployments use A > F to improve FFN compute utilization, which is why we have not implemented A < F support.

Without a concrete application scenario and performance evidence for it under the intended execution mode, we would be reluctant to merge the additional software complexity and maintenance burden. If you have a use case that benefits from A < F, along with corresponding performance results, please share them—we would be happy to revisit this with that evidence.

[swjeong9]

Thank you for the thoughtful review. The changes in this PR are of three different kinds, so we have laid them out one by one below.

On the measurements

To clarify the measurement setup, since the PR description may not have made this clear: we measured in both eager and graph mode (FULL_DECODE_ONLY + DBO), and the tables have a row for each. The purpose was to verify that, on an existing layout (2A2F), this change does not alter the output and introduces no performance regression.

• Output: identical probabilities at all 4096 positions, in both modes
• Time per decode token (TPOT): eager 75.1 → 74.6 ms, graph+DBO 112.8 → 115.0 ms
  • The differences are marginal, so we conclude there is no regression.

The three changes in this PR

1. Allowing A < F

We fully agree with your point about the intended regime: since AFD batches tokens from several attention ranks to fill the FFN, A >= F is the normal case.

We are a research group working on optimizing LLM serving systems on heterogeneous GPU clusters. We expected that if attention is placed on GPUs with high memory bandwidth and capacity, and the FFN on GPUs with better compute efficiency, a layout with fewer attention ranks than FFN ranks could arise. Our reference point was MegaScale-Infer, which introduced attention-FFN disaggregation:

• It identifies heterogeneous deployment as one of the two key advantages of disaggregation (§2.4)
  • For example, H20 for attention and L40S for the experts (§4.3)
• The number of attention nodes is not a fixed ratio; it is set by profiling the per-micro-batch compute time of attention and of the experts and balancing the two (§4.2)
• In the deployment-plan experiment of §7.4, they vary the attention node count from 1 to 16 against 16 expert nodes, and the point they report as optimal is 8 attention : 16 expert (Fig. 15).

That said, we recognize this is a special case specific to our research setting, not a requirement for typical deployments.

2. Relaxing A % F == 0

This one is separate from A < F; it sits inside the A >= F regime you described.

• Currently, with F=2, A can be 2, 4, 6 but not 3, 5, 7. The ratio that balances attention and FFN utilization is not always an integer multiple.
• As far as we can tell, the constraint is not a requirement of the communication structure; it appears to guard subgroup_index = role_rank // (A // F), which for A % F != 0 assigns the last attention ranks an out-of-range subgroup index and raises IndexError.
• With the partition written as a * G // A, groups may differ in size, the condition is no longer needed, and the existing layouts produce the same groups as before.

3. Multi-node deployment

Currently P2pNcclAFDConnector only works when every rank is on the same host.

• Every subgroup's StatelessProcessGroup is created against the single afd_config.host, so with two or more FFN ranks on different nodes, the first member of the second subgroup cannot bind that address and initialization hangs.
• The recipes all use 127.0.0.1, and line 178 of the user guide records cross-node use as "not established by the current recipes".
• We reproduced this on 2A2F across four nodes and addressed it by having each rank exchange its own address once, then creating each subgroup at its first member's address. On a single host the resolved address equals the configured one, so existing recipes behave the same.

Two questions

Our understanding is that your concern centers on item 1. Most of this PR's complexity is also in item 1 (the slice ops, the dummy batch for a rank that receives zero tokens, and their handling under torch.compile); items 2 and 3 have none of that.

So we would like to ask:

• Do the same concerns apply to items 2 and 3? Both sit inside the A >= F regime, and item 3 also affects the currently supported 4A2F once its FFN ranks are placed on different nodes.
• If the concern is specific to item 1, we would trim this PR to items 2 and 3 with the A >= F check left in place. Item 1 would stay in the RFC, and we would bring it back if a concrete use case comes up.

We would appreciate your guidance on this, and will gladly reshape the PR whichever way you prefer. Thank you again for your time.

[jiangkuaixue123]

Thanks for the clarification. I agree with item 2—relaxing the divisibility constraint while keeping A >= F is something we can support.

For item 3, could you open a separate bug issue with a reproducer and the deployment details? We would like to reproduce and investigate it on our side as well.

Regarding the eager versus FULL_DECODE_ONLY + DBO measurements, could you verify that FULL_DECODE_ONLY is actually hitting CUDA graph replay? The reported numbers make me suspect that graph replay may not be taking place and execution may be falling back to eager mode. Please check the actual execution path to confirm.

[swjeong9]

Thank you for going through this carefully, and in particular for the question on graph replay, which pointed at a gap in how we measured. Would the following be acceptable for the three items?

Item 2 (relaxing the divisibility constraint): we would narrow this PR to that change alone. The A >= F check is retained and only A % F == 0 is removed. As requested, we will confirm that FULL_DECODE_ONLY reaches graph replay, and attach the profiler traces and capture logs that establish it.

Item 3 (multi-node rendezvous): we will open a separate bug issue with the deployment details and a reproduction that varies only whether Attention and FFN are co-located on one host. If you confirm it and consider it worth addressing, please indicate the approach you would prefer and we will submit it as a separate PR.

Item 1 (A < F): as you note, it is not a requirement for typical deployments, so we would drop it from this PR. Should it become useful later, we would be glad to prepare it separately.

Please let us know if you would prefer a different arrangement, and we will proceed accordingly.

[jiangkuaixue123]

Thank you for going through this carefully, and in particular for the question on graph replay, which pointed at a gap in how we measured. Would the following be acceptable for the three items?

Item 2 (relaxing the divisibility constraint): we would narrow this PR to that change alone. The A >= F check is retained and only A % F == 0 is removed. As requested, we will confirm that FULL_DECODE_ONLY reaches graph replay, and attach the profiler traces and capture logs that establish it.

Item 3 (multi-node rendezvous): we will open a separate bug issue with the deployment details and a reproduction that varies only whether Attention and FFN are co-located on one host. If you confirm it and consider it worth addressing, please indicate the approach you would prefer and we will submit it as a separate PR.

Item 1 (A < F): as you note, it is not a requirement for typical deployments, so we would drop it from this PR. Should it become useful later, we would be glad to prepare it separately.

Please let us know if you would prefer a different arrangement, and we will proceed accordingly.

LGTM

[specture724]

Thank @swjeong9 for the work.

Item 3 (multi-node rendezvous):

We prefer a separate issue & PR for this. And for the design, subgroups only need the stateless process group to make sure every rank in the group get the ncclUniqueId of rank0. And now building stateless process group makes multi-node run fail. I'd like to re-design the subgroup logic:
Since afd_world has already built a global rendezvous store, we can reuse it for subgroups.

# afd_process_group.py: expose the rendezvous_store
def afd_rendezvous_store(
    init_method: str, rank: int, world_size: int, timeout: timedelta
) -> Store:
    store, _, _ = next(rendezvous(init_method, rank, world_size, timeout=timeout))
    store.set_timeout(timeout)
    return store

#p2p.py: replace _gather_rank_addresses + StatelessProcessGroup.create
subgroup_store = PrefixStore(
    f"afd_subgroup_{self.mapping.subgroup_index}", root_store
)
self.a2e_group = StatelessProcessGroup(
    rank=self.mapping.rank_in_subgroup,
    world_size=len(self.mapping.subgroup_ranks),
    store=subgroup_store,
)

[specture724]

Item 1 & 2 LGTM

[swjeong9]

As we discussed, this PR is now narrowed to the change that removes the A % F == 0 constraint, and the title has been updated to match.

We took another look at the graph side, as you suggested. We confirmed with DEBUG logs that replay does take place, then turned DEBUG off, measured again, and updated the PR body with those numbers.

We will open a separate issue for the multi-node rendezvous. We have also read through the design @specture724 outlined — thank you for sharing it. Should we end up taking that work, we would be glad to follow that direction.

Thank you both for the review. If any part of this would be better shaped differently, or if there are further measurements you would like to see, please let us know and we will gladly follow up.

[specture724]

We will open a separate issue for the multi-node rendezvous. We have also read through the design @specture724 outlined — thank you for sharing it. Should we end up taking that work, we would be glad to follow that direction.

Thank @swjeong9. We are glad to accept your multi-node work. If you are testing on RDMA NICs, please make sure that GPUs and NICs are binded in the same NUMA node. I remembered that we encountered transfer hang due to GPU Direct RDMA. We just found out how to solve the hang, while didn'd dig into the problem. Just FYI if you encountered the same hang.

[swjeong9]

Thank you, @specture724 — we are glad to take this on. For the multi-node part, we would like to define the reproduction conditions precisely before opening the issue, and will proceed from there.

Thank you also for raising the GPU Direct RDMA hang. Our initial reproduction needed one rank per node, so we scaled out with the smallest single-GPU instances to keep the cost reasonable; EFA is not offered at that size, so the transport was TCP over the standard ENA interface throughout. For the follow-up we would like to extend the matrix to an EFA-enabled configuration, and we will make sure the GPUs and NICs share a NUMA node as you described. We appreciate you raising it ahead of time.

[swjeong9]

Thank you for reviewing this PR and running CI.

I traced the Buildkite unit-test failures to the NPU dependency issue reported in #322. Since this was fixed by #323, I’ve rebased this PR onto the latest main to include that fix.

The previous pre-commit run also reported 13 mypy errors in test_p2p_connector.py. The affected statements are unchanged from the base revision, though the file is modified by this PR.

Would you prefer these existing typing issues to be addressed in this PR, or handled separately to keep this change focused?

[specture724]

IMHO, is this rank division logic the same as ffn_metadata.py: L42-53? If so, make it a helper:

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,
    )

[swjeong9]

Thank you for the review. I agree that consolidating the two expressions into a shared helper makes the code clearer. Applied as suggested.

[specture724]

ditto

[swjeong9]

Applied the same way.

[specture724]

Will ratio == len(subgroup) - 1 always true? If so, remove this attribute and the doc strings in L23-25

[swjeong9]

Yes, that always holds. As suggested, the ratio field on AFDRankMapping and its docstring sentence have been removed. Since the attribute no longer exists, tests/unit/distributed/test_topology_snapshot.py now fills ratio in from subgroup_ranks before comparing against the golden entries, so the captured snapshots stay unchanged.

[specture724]

Will ratio == len(subgroup) - 1 always true? If so, remove this attribute and the doc strings in L23-25

[swjeong9]

Addressed as above.

[specture724]

FYI: I reviewed the tests with the help of AI, and here's the results.

Review of the test changes only. The production change itself looks right to me; I verified the partition separately and will keep that in a separate comment.

First, credit where it is due: I checked that the golden snapshots are genuine. I copied tests/unit/distributed/test_topology_snapshot.py onto upstream/main and ran it there, and all 23 cases pass. So the "byte-for-byte unchanged for existing topologies" claim holds, and the Captured from main @ 603c111 note is accurate. That is exactly the evidence a reviewer needs.

My concern is that the evidence and the permanent test suite are not the same thing. The three new files add roughly 440 lines, and a large part of that is the same guarantee expressed three times.

1. The golden dict duplicates the invariants test in the same file

test_topology_snapshot.py L42-248 is a 176-line literal dict plus the 6 lines that consume it, covering 3 topologies and 10 mappings. test_rank_mapping_invariants at L251-294 in the same file already asserts every one of those fields, across 11 topologies rather than 3. Each golden entry also echoes back role, role_rank, attention_size and ffn_size, which are the inputs.

The one thing the golden adds over the invariants test is that dataclasses.asdict fails if a new field appears on AFDRankMapping. That does not seem worth 180 lines, and a renamed or removed field already breaks the invariants test by name.

2. test_topology_snapshot.py as a whole duplicates test_topology_partition.py

The file's stated purpose is to show that relaxing the divisibility rule leaves the old layouts unchanged. test_topology_partition.py::test_divisible_layouts_keep_the_historical_grouping does that in 8 lines, over every divisible A, F pair up to 8, by asserting the historical formula directly:

assert mapping.subgroup_ranks == (
    ffn_rank,
    *(ffn + ffn_rank * ratio + offset for offset in range(ratio)),
)

_LEGAL_GRID (11 hand-written divisible pairs) is a subset of _GRID. The roster-agreement and partition-coverage assertions in test_rank_mapping_invariants are the same checks as test_every_attention_rank_belongs_to_exactly_one_subgroup and test_both_roles_agree_on_the_subgroup_they_share.

What is genuinely unique to the invariants test is four assertions, and all four hold for non-divisible layouts too, so they can move into the partition file and run over the wider grid:

assert mapping.world_rank == (role_rank if role == "ffn" else ffn + role_rank)
assert mapping.p2p_rank == (role_rank if role == "ffn" else role_rank + min_size)
assert mapping.min_size == min_size
assert sorted(dp_destination_union) == list(range(ffn))

Suggestion: keep the golden file out of the merge and carry those four assertions into test_topology_partition.py. That leaves one new topology test file instead of two.

If you would rather keep a snapshot permanently, a flat table keyed by (A, F, role, role_rank) holding only the six derived fields would be about 10 lines of data instead of 176.

3. Two tests in test_topology_partition.py are already covered

• L37-39 test_validation_accepts_counts_that_do_not_divide only asserts that no exception is raised. Every other test in the file goes through build_rank_mapping, which calls validate_p2p_topology first, and _GRID contains the non-divisible pairs. It cannot fail unless a test elsewhere in the file already fails.
• L41-45 test_validation_still_rejects_fewer_attention_than_ffn_ranks repeats test_p2p_connector.py::test_p2p_topology_validation_errors_are_clear, which pins the same message. The trailing comment in test_topology_snapshot.py says as much.

4. test_p2p_connector.py topology rows

L173-177 adds four parametrize rows for 3A2F and 5A3F. Now that a dedicated partition file exhaustively covers every A, F up to 8 and pins the literal rosters, these rows do not add coverage at the connector level.

The two rows added to test_p2p_ffn_metadata_tracks_each_attention_peer_in_xayf are a different matter and should stay. They exercise the actual behaviour change in p2p.py, where the peer is now read from subgroup_ranks instead of computed from a uniform ratio.

One coupling to note: the new ratio parametrize column at L165 and the assertion at L199 exist because ratio is no longer A // F. ratio is now exactly len(subgroup_ranks) - 1, and the connector already carries that as group_size. If you decide to drop the field, this column goes away with it.

5. test_ffn_metadata.py

No notes. aggregate_ffn_token_counts had no direct coverage before, and 43 lines for six cases including the empty-counts and TP-expansion edges is proportionate.

Summary

Roughly 290 lines of the new test code are restatements of guarantees the remaining tests already make. Concretely: drop test_topology_snapshot.py after review, move its four unique assertions into test_topology_partition.py, drop the two validation tests there, and drop the four topology rows from the connector test.

None of this is a correctness objection. The tests pass and they test true things.

[swjeong9]

Thank you for the detailed review, especially for checking the snapshots against upstream. I agree that the tests overlapped more than necessary, and your distinction between review-time evidence and the permanent test suite is helpful. I have revised the tests along the lines you suggested.

The snapshot file has been removed, and its four unique checks have been moved into test_world_and_p2p_ranks_follow_the_role_layout in the partition file, which runs them over the existing grid of divisible and non-divisible layouts. The redundant validation tests and the added connector topology rows have also been removed, while the two connector metadata cases and the FFN metadata tests remain.

I appreciate the time you have taken to review this.

[specture724]

Thank you for the work, LGTM

[specture724]

Please fix the pre-commit issue

[swjeong9]

The remaining pre-commit errors are addressed by the fix in #328. I’ll follow up here once the outcome of that PR is clear, either by rebasing onto main if it is merged or by applying the fix separately.

[swjeong9]

Rebased onto the latest upstream main, which resolves the pre-commit failures. All checks pass locally.