[RFC]: Support arbitrary attention-to-FFN rank ratios (including A < F) in P2pNcclAFDConnector

[swjeong9]

Motivation

validate_p2p_topology requires num_attention_ranks >= num_ffn_ranks and num_attention_ranks % num_ffn_ranks == 0. In the wording of the user guide, we read this as requiring ratio = A // F to be an integer of at least 1.

We ran into each of the two conditions:

Condition	What we hit	Layout we wanted
A >= F	The expert weights did not fit on a single GPU, so the FFN role needed several ranks, while we wanted to use as few GPUs as possible for the experiment	DeepSeek-V2-Lite experts are about 25 GB and an L4 has 24 GB, so 1A2F
A % F == 0	We wanted to sweep the ratio, but the step after 2A2F is 4A2F with nothing in between	3A2F, 5A2F

For the first one, 1A2F was rejected, so we moved up to 2A2F and used four GPUs where three would have been enough. We expect the second to show up in the same shape whenever the node count is not a multiple, or when one rank is missing. Deployments that give the two roles different GPUs may also end up with rank counts that do not divide, but we have not verified that ourselves, so we are not offering it as evidence.

Since the benefit of AFD comes from batching tokens from several attention ranks to fill the FFN GEMMs, we think it is natural that A >> F is the recommended regime.

That said, we do not think it is a reason to reject other ratios outright. If there is a reason we have missed, we would be grateful to hear it.

Proposed change

We would like to propose keeping the subgroup structure and the per-group communicators as they are, and generalizing only the partition formula. With 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 empty, and when A >= F && A % F == 0 the result is the same grouping as today.

4A2F → {F0,A0,A1} {F1,A2,A3}   (same as today)
3A2F → {F0,A0,A1} {F1,A2}
1A2F → {F0,F1,A0}

Because the existing layouts are absorbed as a special case of the formula, there is no mode switch, no new config, and no new class. The role with fewer members has exactly one member per group, so a group is always one hub plus several spokes; A >= F behaves as it does today, and under A < F the attention rank splits its tokens across the FFN members of its group and concatenates the results back.

Location	Change
distributed/topology.py	Drop the constraint, generalize the partition arithmetic, add a split-size helper, 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. When a group has a single FFN member, the existing path is taken unchanged
v1/worker/ffn_metadata.py	Aggregate per-FFN token counts with the same partition formula (EP dispatch sizes)

The one new situation under A < F is an FFN rank whose share is zero, which 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 cannot skip the layer; it runs a one-token dummy batch and discards the result. We understand this to be the same approach vLLM takes when it gives an idle DP rank a dummy run.

The connector interface, the model code, the model runner and worker, the config schema, and the expert all-to-all inside the FFN role are all untouched. Under the currently supported condition A >= F && A % F == 0, the execution path also stays exactly as it is.

The implementation and the measurements are done, and we have opened #309 against this RFC so the code is easy to look at. If you would rather it took a different direction, we will rework it accordingly.

Plugin boundary

Plugin-owned: all of this change — topology arithmetic, the GPU connector's boundary transfer, FFN token-count aggregation
Compat helper: none
Compat patch: none — the vLLM sources are not modified
Explicit class path: afd_plugin.connectors.gpu.p2p.P2pNcclAFDConnector (unchanged; no new class)

The NPU connectors (CAMP2pAFDConnector, CAMAsyncAFDConnector) are not modified either. They carry their own validation and their own rank mapping, so this change does not reach them, and A < F on NPU is still rejected exactly as it is today. We do not have the hardware to verify it, so we judged it better to leave it alone.

We think the design documents would need updating as well. docs/design/module/connector_contracts.md states num_attention_ranks >= num_ffn_ranks and an integer A/F ratio as a contract at line 113, and plugin_boundary.md mentions the same values. If you approve the direction, we will prepare updates to both alongside the code.

Risks and alternatives

Existing users (compatibility, correctness, performance)

We placed one rank per pod on four L4 nodes to form 2A2F, and alternated the current code (main) and this proposal on the same nodes. The model is DeepSeek-V2-Lite, with 32 prompts at temperature 0. Below, graph+DBO means CUDA graph capture (cudagraph_mode=FULL_DECODE_ONLY) together with --enable-dbo.

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 proposal
eager	Per-position probability difference (4096)	baseline	identical at every position
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	Peak throughput	95.9 tok/s	96.9 tok/s

Throughput is 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; we also measured single-request latency and it fell in the same range. We are not claiming a performance improvement. The value of this proposal is expressiveness, and on performance our claim stops at the absence of a regression.

Correctness on the new topologies

We brought up 1A2F on three nodes and 1A4F on five, each in eager and in graph+DBO, and collected responses. 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.

The table below is relative to the current code at 2A2F eager, averaged over how much the probability assigned to each token changed.

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; they are a reference for what these configurations look like. Input is 1024 and output 32 tokens; 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

Testing

The partition arithmetic and the transfer wiring can be pinned by CPU unit tests, so CI can catch regressions there. Real communication and performance need GPUs, and A < F needs at least three. The tables above are what we could measure on the hardware we have; we will attach the commands and the full numbers to the PR.

Maintenance

This generalizes the arithmetic of the existing path rather than adding a second path beside it, so there are no two paths to keep in step. The currently supported topologies are handled by the same code.

Alternatives

Alternative	Assessment
Keep the constraint	The situations in the table above remain
A 1-hop expert-routed boundary	More fundamental, but the transfer sizes depend on the routing result, which conflicts with CUDA graph and requires kernel work
Send a fake single token to a peer whose share is zero	Removes the dummy handling, but sends a fake row every layer and puts a branch on the token count into the compiled path

Multi-node deployment

All the measurements above are multi-node. Since docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md line 178 records cross-node use as "not established by the current recipes", we tried it, and initialization hung because every subgroup rendezvous used the single configured host. It reproduces on existing A >= F layouts as well. It is a different kind of problem from this proposal, but the results above are not reproducible without it, so we plan to include the fix in the same PR.

Feedback period

We are not setting a deadline. We will incorporate your comments whenever it suits you.

CC list

@hsliuustc0106 @jiangkuaixue123