[RFC]: Async GPU Connector

[specture724]

[RFC]: Async GPU Connector

Motivation

CAMAsyncAFDConnector already removed the global synchronization between Attention DP replicas on NPU (DP wave, coordinator, DPMetadata, the AFD control plane). On CUDA the only option so far has been the synchronous connector P2pNcclAFDConnector, which requires every replica to agree on token count and layer index; attention cost is O(∑sᵢ²), so under an uneven request mix the replicas can never line up and the fast one waits for the slow one.

The two vLLM engine patches that make async possible (compat/patches/async_dp_engine.py, async_dp_forward_context.py) are already written and platform-neutral. Their only switch is is_afd_async_dp(), which used to hard-code connector == "CAMAsyncAFDConnector". Reusing the mechanism on GPU only requires relaxing that predicate to a set — not a single line of the engine patches changes. What was actually missing is the connector itself, the self-driven FFN loop, and an Attention-side gate for CUDA.

Proposed change

Add GpuAsyncAFDConnector, semantically aligned with CAM: control_plane = None, token-level dispatch/combine, in-band metadata, anonymous arrival, and a self-driven FFN pull loop. The baseline limitations are aligned as well: eager only, prefill only, single node, incompatible with native DBO.

Dispatch routing, the weighted combine reduction, and the expert GEMM are all assembled from existing torch and vLLM ops; the transport is NVSHMEM peer-pointer mapping plus a plain copy_.

Transport: bootstrapping NVSHMEM ourselves

The NVSHMEM backend of _symmetric_memory bootstraps on the default process group and carves teams out of NVSHMEM_TEAM_WORLD with nvshmem_team_split_strided. The two AFD roles are two independent vllm serve processes whose default groups cover only their own role's ranks, so the cross-role AFD group is never a strided subset of it and team creation always fails.

nvshmem_rt.py therefore binds libnvshmem_host.so.3 directly via ctypes: rank 0 mints a uniqueid, broadcasts it over the AFD process group's store, and every rank initializes with the same id — the AFD world simply is NVSHMEM_TEAM_WORLD, no team split needed, and the PE numbering equals the AFD world rank (asserted with nvshmem_my_pe/n_pes after init). The host library's ABI is pinned by static asserts in its own headers (uniqueid 128 B, init_attr 144 B, version = (1<<16)+sizeof), so ctypes is enough and no compiled extension is required.

The data plane uses only three primitives: nvshmem_malloc (symmetric allocation), nvshmem_ptr (map a peer's copy into this process), and a __cuda_array_interface__ view that wraps a raw pointer as a tensor. Writing to a remote is a copy_ into the mapped view, one-sided by construction. The price is that a NULL from nvshmem_ptr is a hard error — only a single node with NVLink/P2P reachability is supported.

Window and slot layout (symm_window.py)

window = [ flag[num_regions * ring_depth] | slot(0,0) | slot(0,1) | ... ]
slot    = [ header | route_table | routed_x | shared_idx | shared_x ]   # every field 256B-aligned

• num_regions = max(attn_size, ffn_size); both sides take the larger so the symmetric allocation matches. A sender only ever writes region = its own role_rank, and the in-slot offsets are computable locally, so no remote atomics are needed.
• Dispatch (A→F) and combine (F→A) share one slot layout, so a single spec sizes both directions.
• Write ordering: the payload and the flag copy_ are issued on the same stream, and same-stream D2D copies complete in issue order, so "seeing the flag" implies "the payload is complete". This holds only for NVLink-mapped memory; a cross-node transport would need an explicit fence, as noted in the module docstring.
• The flag's value is the sender's monotonically increasing seq; the receiver keeps a host-side _seen mirror, and poll() fetches all flags in one D2H and returns the first slot whose value changed — ABA-safe by construction.

The header is a fixed 12 int32 words plus expert_counts[expert_per_rank], with a self-checking magic AFDG and version; flags bit1 is the shutdown sentinel. layer_idx / num_tokens / routed_tokens / shared_tokens / group list all travel in band — the FFN side knows nothing beforehand and relies entirely on the header.

Capacity and ring depth

routed_cap = ceil(max_num_batched_tokens * topk / ffn_size) * routed_cap_multiplier  # default 2.0
token_cap  = max_num_batched_tokens

routed_cap_multiplier is necessary: real gates are not balanced (DeepSeek-V2-Lite on 2 FFN ranks was measured at roughly 1.33× the even split), and the worst case is ffn_size (every partial landing on one rank). On overflow write_slot raises immediately — never a silent truncation.

ring_depth is derived from an invariant, not a performance knob: Attention is strictly send-then-recv per layer, so each (peer, stage) has at most one request in flight and the default equals the number of stages (1 without ubatching). Exhaustion can only mean the topology config does not match the actual peer count, which is an error.

Data flow

Attention-side send_attn_output (async_gpu.py:447):

1. plan_dispatch flattens topk_ids and runs argsort(stable=True) — one sort yields both groupings at once: segmented by destination FFN rank, and grouped by local expert within each segment, exactly the order the receiver needs. bincount gives the group_list, cumsum gives the segment offsets.
2. One D2H per send via counts.cpu() (offsets are a prefix sum, cheaper to redo on the host than to fetch a second tensor).
3. For each FFN rank: index_select the routed rows of that segment plus that rank's shared rows (shared tokens are split round-robin as arange(ffn_rank, num_tokens, ffn_size)), then write them one-sided into the peer's window together with the header and route table.
4. topk_weights never goes on the wire; it stays in the Attention-side pending FIFO and the weighting happens during combine (matching NPU — the FFN side does not multiply).

FFN-side recv_attn_output → recv_ffn_work_item: a poll() hit returns zero-copy views into the window, with no movement at all; layer_idx and group_list come back with the header in one shot. The CAM Python side calls .item() twice on TokenNums_Rankid_Layeridx (two D2H syncs on the critical path per layer per receive); this path does not replicate that structure.

Attention-side recv_ffn_output: it waits only on the FFN ranks that actually received data at dispatch time (expected_ffn), so an empty message never creates a wait. Each arrival is accumulated into an fp32 accumulator via combine_scatter (a weighted index_add_ scatter); the shared branch is index_add_ed directly by shared_idx.

Attention gate and FFN expert compute

connector.select_experts calls vLLM's grouped_topk directly, letting CAM and CUDA share the same deepseek_v2_attention_gate path (deepseek_v2.py:512).

What the FFN side receives is already routed and already grouped by local expert, so it must not route again. compute_attention_gate_moe_ffn (models/gpu/deepseek_v2_attention_gate.py) treats it as a topk == 1 problem: give every row an expert id expanded from group_list and a unit weight, then call vLLM's fused_experts — what comes out is exactly this rank's local expert output. FusedMoE therefore needs no patch at all, which dissolves the largest unknown risk of the earlier draft.

The shared expert goes through runner._shared_experts._layer(...) rather than SharedExperts.forward: the latter is a stateful scheduler for the runner's own multi-stream pipeline and returns None when the expected ordering does not match, while AFD feeds shared tokens as their own batch, so that machinery does not apply. routed_scaling_factor is applied the same way as the NPU version (scale the routed branch unless fp16, where the shared branch is scaled down instead).

FFN main loop

The control_plane is None branch in ffn_worker.py changes from NotImplementedError to calling execute_connector_driven_step(): the inner loop drains at most num_layers work items (that is only a drain granularity — adjacent work items may belong to different layers of different replicas), and returns on a poll timeout so the outer while gets to check the shutdown event. The unconditional per-step torch.cuda.synchronize() is gone — it would erase all overlap; ordering is carried by the connector's own streams. Shutdown uses the header's shutdown bit (announce_shutdown broadcasts to every opposite-role peer) plus the recv_poll_timeout_ms timeout as a second line of defense, and ConnectorShutdown is recognized as a normal exit in the worker loop.

Plugin boundary

Plugin-owned

• afd_plugin/connectors/gpu/async_gpu.py — GpuAsyncAFDConnector, GpuAsyncExtraInfo, GpuAsyncTransferState, GpuAsyncFFNWorkItem, plus the two pure functions plan_dispatch / combine_scatter (unit-testable without a device).
• afd_plugin/connectors/gpu/symm_window.py — slot layout, header encode/decode, one-sided write, poll.
• afd_plugin/connectors/gpu/nvshmem_rt.py — ctypes NVSHMEM binding (init / malloc / peer_ptr / tensor views). No C++/CUDA sources, no new torch ops.
• afd_plugin/connectors/async_topology.py — AFDAsyncTopology and build_async_topology extracted from npu/async_cam.py, now shared by both async connectors (CAM imports from here instead; behavior unchanged).
• afd_plugin/model_executor/models/gpu/deepseek_v2_attention_gate.py — compute_attention_gate_moe_ffn, same name and signature as the NPU version.
• v1/worker/ffn_worker.py, v1/worker/ffn_model_runner.py (new execute_connector_driven_step / _ffn_forward_connector_driven), v1/worker/attention_model_runner.py (drop the two control_plane is not None asserts; _send_dp_metadata returns early when there is no control plane).
• config.py / connectors/factory.py: the single-valued AFD_ASYNC_CONNECTOR constant becomes the AFD_ASYNC_CONNECTORS set (two references: is_afd_async_dp and config validation); register the connector.
• models/deepseek_v2.py: the async forward schedule and the gate helper are shared by both platforms (the module still lives under models/npu/ but is now platform-neutral); the GPU branch of compute_gate_on_attention forwards to the GPU gate once it has a group_list, and keeps the original error otherwise (that caller is the control-plane path).
• Recipes recipe/gpu/GpuAsyncAFDConnector/deepseek_v2_lite/{1a1f,2a2f}_eager_async.sh.

Compat helper

None. NVSHMEM availability is checked by the explicit errors in nvshmem_rt._find_host_library() and peer_ptr().

Compat patch

• async_dp_engine.py / async_dp_forward_context.py: unchanged. Relaxing the is_afd_async_dp() predicate makes them apply to GPU automatically. This is the largest leverage point of the design.
• One new runtime monkey patch that needs careful review: _dp_batch_coordination_disabled in attention_model_runner.py temporarily replaces the module attribute gpu_model_runner.coordinate_batch_across_dp for the duration of the super()._determine_batch_execution_and_padding call. Upstream all-reduces the batch shape across the DP group whenever data_parallel_size > 1; async AFD deliberately lets each replica advance on its own, so an idle replica never joins that collective and a busy one blocks in it forever — which is exactly where 2A2F hangs before reaching the first MoE layer. The replacement returns the single-rank answer (num_tokens_across_dp=None), so upstream takes precisely the same branch as dp_size == 1. The scope is limited to that one call and only applies when control_plane is None, but per the AGENTS.md patching rules it needs # ### PATCH START/END markers and a long-term plan (an upstream skip hook, or a subclass override instead).
• in_profile_run in deepseek_v2_async_cam_forward.py becomes getattr(..., False): the field exists only on the Ascend forward context, not on the CUDA one. This deviates from the AGENTS.md rule against reaching into upstream structures with getattr, and is the price of sharing one schedule across platforms — please review whether that is acceptable, or whether it should dispatch per platform instead.

Explicit class path

• Config name "GpuAsyncAFDConnector" → afd_plugin.connectors.gpu.async_gpu.GpuAsyncAFDConnector
• NVSHMEM world group name afd_async_gpu; connector_extra_config fields: attn_ranks_per_dp, ring_depth, routed_cap_multiplier, recv_poll_timeout_ms, async_moe_ubatching, async_moe_num_ubatches, async_moe_split (unknown fields are always rejected).

Risks and alternatives

Alternatives already ruled out

• DeepEP: its receive buffer is statically partitioned as [num_local_experts, max_tokens_per_rank * num_ranks, hidden], every rank in the group is assigned experts and reserves a send budget, whereas an AFD Attention rank holds zero experts; dispatch is also round-synchronous with no layer tag on the wire, which cannot carry the asynchronous premise of "two replicas parked on different layers".
• ZMQ sideband + bidirectional NCCL (early draft): needs a host-side header channel, slab slot management, and one comm per pair, and still misses the ×topk/ffn_size bandwidth gain of token-level dispatch.
• A custom wait_any CUDA kernel plus four afd_gpu ops (the previous version of this file): the implementation showed a host-side poll is already enough to run, so no compiled artifact is introduced for now. See the performance risk below.
• A CUDA IPC fallback: maintaining two window protocols and two flag semantics costs more than it returns.

Risk	Status / mitigation
The gain is only ~5% (see "Measured performance"), far below the 194% ASAP reports	What is missing is the overlap and scheduling work of the paper's §3.3/§3.4, none of which this RFC implements. This is the main decision to make before merging: land the overlap work first, or land correctness first and iterate
Host-side poll: SymmWindow.poll costs one blocking D2H per call, on the critical path of both sides	Marked in the code with ponytail: and an upgrade path — replace it with a device-side wait_any kernel spinning on the flag array. It was not the main cause of the measured knee (both connectors saturate at the same throughput), but it is a prerequisite for the overlap work
Every transfer runs on the default stream; the recv/compute/send multi-stream and event chain from the early RFC version is not implemented	This is the primary source of the gap to the paper, and should be the next step
poll scans linearly and returns the first hit, favoring low region indices	The region count is currently ≤ max(A,F), so the scale is tiny; if starvation ever shows up, rotating the scan start is a one-liner
NVSHMEM becomes a hard dependency, with the host library ABI bound via ctypes	The ABI is pinned by NVSHMEM's own version macros and static asserts; PE numbering is asserted with my_pe/n_pes after init; a missing library produces explicit installation guidance
Single node only: nvshmem_ptr requires direct P2P	An explicit error rather than a silent fallback; cross-node would need IB/RDMA plus an explicit fence (the same-stream write-ordering assumption no longer holds), which is out of scope for the first version
Window memory comes out of the KV cache budget	ring_depth is already tightened to 1–2 by the invariant; routed_cap_multiplier is tunable; init logs the slot/total size at logger.info (note: the logger must sit under the vllm.* tree, or this — the only report of that allocation — is dropped)
routed_cap is an expected capacity, not an upper bound	write_slot raises on overflow; the worst case is multiplier = ffn_size
close() does not free the symmetric memory	nvshmem_free is collective and would require both roles to shut down in lockstep; marked ponytail: and reclaimed at process exit. It must be added if windows are ever recreated in-process
The coordinate_batch_across_dp monkey patch (see Plugin boundary)	Scope is minimized, but it needs patch markers and an upstream plan
No startup feature validation: GPU async does not currently enforce eager / prefill-only / compute_gate_on_attention=true / no DBO	The recipe scripts set them by hand. A GPU counterpart of compat/npu/feature_validation.py is recommended — this is an open item before merging
Reusing a staging buffer for async H2D once caused silent data corruption: with one shared pinned header buffer, the next peer in the send loop overwrote it on the host while the previous copy was still in flight, and the receiver read someone else's token counts and ran out of bounds	Now staged one row per (peer, ring); the lesson is that the source buffer of any async copy must be allocated at in-flight granularity. The header's echo_seq check (one of the two seq levels in the design) has also been added, turning this class of mismatch into an explicit error
Numerical agreement with the synchronous connector	tests/e2e/async_gpu_moe_equivalence.py already compares the full dispatch→compute→combine chain against a naive per-token MoE reference; a 2A2F end-to-end, per-token comparison against P2pNcclAFDConnector is still needed
Shared experts are assigned round-robin by token_idx % ffn_size, which may be uneven	Accepted for the first version; load balancing comes later
Maintenance surface: a second GPU connector	The work-item protocol has the same names and signatures as CAM (recv_ffn_work_item / send_ffn_work_item_output, still duck-typed for now), and the topology code is merged into async_topology.py, so both can be folded in together when #107 promotes the work-item protocol to a formal interface

Existing verification

• Unit tests tests/unit/connectors/test_async_gpu_connector.py: slot fields are disjoint and in bounds, every offset is viewable under its dtype, header round-trip with magic/length validation, the shutdown bit, every partial routed exactly once, each destination segment grouped by local expert, identity experts recombining to the weighted sum, an expert count not divisible by ffn_size, and an empty destination.
• Two-GPU tests/e2e/async_gpu_connector_e2e.py (two processes over the connector's public API end to end, with a real grouped GEMM, reproducing the deployment topology: a private size-1 default group per process).
• Recipes 1a1f / 2a2f, eager prefill.

Feedback period

Currently I'll focus on this RFC

CC list

@jiangkuaixue123