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

[swjeong9]

Purpose

This implements what we discussed in RFC #308. It removes the A >= F and A % F == 0 constraints in validate_p2p_topology and generalizes the subgroup partition to a balanced block distribution over G = min(A, F) groups. The currently supported topologies are the special case of that formula and produce the same groups. With no mode switch and no new config, A < F and non-divisible ratios become expressible. We measured 1A2F and 1A4F.

It also carries a fix for subgroup rendezvous failing in multi-node deployments. Since docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md line 178 records cross-node use as "not established by the current recipes", we went and checked. Initialization hangs whenever num_ffn_ranks >= 2 and the FFN ranks sit on different nodes, and it reproduces on the currently supported 4A2F as well. The cause is that every subgroup's StatelessProcessGroup uses the single additional_config.afd.host as its rendezvous address. The socket for a group is opened by that group's first member (an FFN rank); on one host every first member shares the address so nothing goes wrong, but once the ranks are spread out, the first member of the second subgroup cannot bind that address and the other members cannot find it there. We considered sending this separately, since it is a different kind of change, but all of the GPU evidence in this PR is multi-node and none of it is reproducible without the fix. We are happy to split it out if you prefer.

Issue

• Related issue(s): [RFC]: Support arbitrary attention-to-FFN rank ratios (including A < F) in P2pNcclAFDConnector #308
• Closing keyword, only if fully resolved: Closes [RFC]: Support arbitrary attention-to-FFN rank ratios (including A < F) in P2pNcclAFDConnector #308

Scope

• In scope
  • distributed/topology.py: drop the constraint, generalize the partition arithmetic, add split_send_sizes, correct the rank offset for the DP metadata channel.
  • connectors/gpu/p2p.py: drive the boundary transfer from the group roster and the split sizes; add two slice ops used only under A < F; the address used to create each subgroup.
  • v1/worker/ffn_metadata.py: aggregate per-FFN token counts with the same partition formula.
• Out of scope
  • A routing-aware 1-hop boundary. This keeps the 2-hop boundary and does not touch the expert all-to-all inside the FFN role.
  • The NPU connectors (CAMP2pAFDConnector, CAMAsyncAFDConnector). They carry the same constraint, but we have no hardware to verify a change there.
  • The rendezvous of the full AFD group and the p2p control-plane group. They keep using the configured host, which is correct because world rank 0 owns that address.
  • Performance. This PR is about expressiveness; our claim stops at the absence of a regression.

Implementation Notes

• Partition: with G = min(A, F) groups, attention a joins group a*G//A and FFN f joins group f*G//F. Both maps are surjective, so no subgroup is left empty. When A >= F && A % F == 0 the grouping is the same as today.
• Transfer: the role with fewer members has one member per group, so a group is always one hub plus several spokes. A >= F behaves as before; under A < F the attention rank splits its tokens across the FFN members and concatenates the results back. Sizes are computed in one place, split_send_sizes; the sender uses the whole list and the receiver its own entry, so both reach the same answer without extra communication.
• An FFN rank whose share is zero: this happens when the token count is smaller than the number of FFN members in the group. That rank still has to take part in the expert all-to-all inside the FFN role, so it runs a one-token dummy batch and discards the result. We understand this to be the same approach vLLM takes for an idle DP rank's dummy run.
• torch.compile: the attention-side transfer sits inside the compiled region, so we kept to two rules. First, slice boundaries are passed as op arguments and the slicing happens inside the op implementation. When we passed a view with a non-zero storage offset to the existing op, the generated code delivered it at offset 0 and the second FFN member received the first member's tokens; we hit this on GPU. Second, the split sizes use a closed form rather than a branch on the remainder, and we do not skip the transfer to a peer whose share is zero, because a branch on the token count freezes at its trace-time value.
• Existing path preserved: when a group has a single FFN member — every topology supported today — the tensor is not sliced and goes through the existing op unchanged. The slice ops are only reached under A < F.
• Rendezvous address: at initialization each rank exchanges its own address once over the AFD group that already exists, and each subgroup is then created at its first member's address. A rank resolves the address of its own interface toward the configured rendezvous host: connect on a UDP socket makes the kernel pick a route without sending a packet, so the peer need not exist yet and nothing goes on the wire. We judged this safer than a name lookup on a node with several interfaces. On a single host the resolved value equals the configured host (127.0.0.1 stays 127.0.0.1), so existing recipes behave exactly as before. Having each rank take its own address from config would also work; we chose automatic resolution because the alternative grows the config surface and has to be filled in per deployment. We are glad to switch to the explicit form if you prefer it.
• No config fields or environment variables are added. The vLLM sources are not modified and no compat patch is added.
• The custom op count goes from two to four. The new slice ops subsume a whole-tensor transfer, so in principle they could replace the existing two. That would change the code path for every currently supported topology and call for re-validation, so we left them in place for now. If you would rather have them merged, we will do that.

Test Plan

• CPU
  • A snapshot pinning the rank mapping of the existing topologies against change.
  • Partition properties over every combination of A, F in 1..6, plus literal examples.
  • Values, properties, and input validation for split_send_sizes.
  • Who sends how many rows to whom and in what order on the generalized topologies, with all ranks connected through a fake wire so the round trip checks token ordering too.
  • Tracing once with the same compile options vLLM uses, then calling with other token counts to check the split.
  • The custom op registration schema.
  • Per-FFN token-count aggregation.
  • Three behaviors of the address resolution helper: loopback resolves to loopback, a local LAN address comes back unchanged, and it does not fail when no peer is listening.
• GPU
  • Bring-up of 2A2F spread across nodes, before and after the rendezvous fix.
  • Output and performance on a supported topology, alternating two images on the same layout.
  • Bring-up and output on the new topologies.

Test Result

CPU: pytest tests/unit/ — 1075 passed, 51 skipped, 2 failed. The two failures are in tests/unit/test_e2e_process_utils.py, which needs os.pidfd_open; our development machine is macOS, and they fail there on main as well, without this change. We modified one existing test: the one asserting the list of registered custom ops, to include the two new ops.

GPU: L4 (g6.xlarge), one rank per pod, vLLM multi-node DP, DeepSeek-V2-Lite. Node counts differ per configuration (1A2F three, 2A2F four, 1A4F five). We sent 32 prompts sequentially at temperature 0. Below, graph+DBO means CUDA graph capture (cudagraph_mode=FULL_DECODE_ONLY) together with --enable-dbo.

Rendezvous fix

On 2A2F across four nodes, the members of the second subgroup timed out at the rendezvous after 120 s before the fix, and came up and served after it. Everything below was collected with the fix in place.

Existing topology (2A2F, current code vs this PR)

We checked output equality as follows. We took the token sequence the current code had generated (128 tokens per prompt), fed it back to both runs as input, and collected the probability each run assigned to that token at each position. This scores rather than generates, so both runs look at the same positions and all 32 x 128 = 4096 of them line up one to one.

Mode	Metric	Current code	This PR
eager	Per-position probability difference (4096)	baseline	identical at every position
eager	Single-request TPOT	75.1 ms	74.6 ms
eager	Peak throughput	81.3 tok/s	80.1 tok/s
graph+DBO	Per-position probability difference (4096)	baseline	identical at every position
graph+DBO	Single-request TPOT (median of 3)	112.8 ms	115.0 ms
graph+DBO	Peak throughput	95.9 tok/s	96.9 tok/s

Latency and throughput are measured with 1024 input and 32 output tokens. Repeating the same configuration moves the number by roughly 8%, so we read the differences above as noise.

New topologies

A different topology means a different batch composition and therefore different bf16 rounding, so the generated token sequences diverge in places. Every divergence we saw was at a position where the top two candidates had nearly equal probability. To measure the size of the difference we used the same method as above, feeding the token sequence generated by the current code at 2A2F eager into each configuration and comparing per-position probabilities.

Configuration	Probability difference (mean over 4096)
1A2F eager	0.9%
1A2F graph+DBO	0.9%
1A4F eager	0.9%
1A4F graph+DBO	1.0%

To establish a reference point, we kept the topology at 2A2F and changed only the execution mode of the current code to graph+DBO, then measured the same way. That gave 0.9%. In other words, the difference introduced by changing the topology is the same size as the difference that changing only the execution mode already produces.

We are also recording the latency and throughput each configuration produced. The GPU counts and batch compositions differ, so these are not values to compare against each other. Latency is 4 requests sent one at a time, and throughput is 256 requests sent at once.

Configuration	Nodes	TTFT	TPOT	Peak throughput
1A2F graph+DBO	3	418 ms	83.5 ms	96.9 tok/s
1A4F graph+DBO	5	581 ms	162.8 ms	68.3 tok/s

1A2F and 1A4F cannot start on the current code — they raise the ValueError at topology.py:56 — and came up on this PR to produce the results above.

Docs Impact

• Files updated: none.
• docs/design/module/connector_contracts.md states num_attention_ranks >= num_ffn_ranks as a contract at line 113, and docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md also describes the supported topologies, so both look like they need updating. Tell us the shape you would like and we will include it in this PR.

---

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.