Skip to content

[2/5] feat: the async GPU connector on the symmetric window - #333

Open
specture724 wants to merge 1 commit into
afd/async-gpu-windowfrom
afd/async-gpu-connector
Open

[2/5] feat: the async GPU connector on the symmetric window#333
specture724 wants to merge 1 commit into
afd/async-gpu-windowfrom
afd/async-gpu-connector

Conversation

@specture724

@specture724 specture724 commented Sep 10, 2026

Copy link
Copy Markdown
Collaborator

Second of a five-PR stack. Based on #332 — the diff shown here is only this
PR's own change; review #332 first.

  1. [1/5] feat: NVSHMEM symmetric window for GPU AFD transport #332 — NVSHMEM symmetric window (transport)
  2. this PR — the async GPU connector
  3. [3/5] feat: run DeepSeek-V4 on the async GPU connector #334 — both GPU runners + DeepSeek-V4
  4. [4/5] feat: capture the FFN experts at bucketed padded shapes #335 — bucketed padded FFN expert graphs
  5. [5/5][WIP] feat: request-aligned DBO ubatch splitting (flag off by default) #336 — request-aligned DBO ubatch splitting (flag off by default)

What this adds

GpuAsyncAFDConnector drives AFD's Attention-to-FFN traffic over the symmetric
window with no control plane at all. The Attention side routes tokens itself
and one-sided-writes the already-grouped rows into every FFN peer's slot; the
FFN side learns the layer index and the row counts from the arriving payload, so
nothing has to be broadcast ahead of a step and the two roles never rendezvous
per layer.

plan_dispatch produces one row per (token, expert slot), sorted by global
expert, so a destination's partials are one contiguous run already grouped for
its grouped GEMM. Nothing about the routing is ever read back to the host:
the per-destination counts go into the slot header on the device, which is what
takes the last synchronize off the layer path.

The dispatch sequence number lives on the device for the same reason the
window's flag protocol uses a constant marker — a replay runs no Python, so a
host-side counter would freeze at whatever was live at capture.

Shared with the NPU async connector

AFDAsyncTopology and build_async_topology move out of the NPU CAM connector
into connectors/async_topology.py so both async connectors share one
role-rank-to-world-rank mapping (logic unchanged — pure hoist). With more than
one async connector now, AFD_ASYNC_CONNECTOR becomes
AFD_ASYNC_NPU_CONNECTOR.

The shared Attention-gate forward gains an opaque-op dispatch path so Dynamo
splits at the MoE round trip instead of tracing into the connector. That path
is gated on the connector type
: CAM keeps its direct calls, because the ops
carry neither router logits nor FlashComm1 token sharding and the CAM protocol
needs both.

Testing

  • tests/unit/connectors/test_async_gpu_connector.py — 33 tests pinning the
    protocol: routing invariants, the flag protocol, dispatch header assembly,
    ring allocation across stages.
  • tests/e2e/async_gpu_connector_e2e.py (2 GPUs) — full round trip against a
    reference, eagerly and with the Attention half captured and replayed.
  • tests/e2e/async_gpu_moe_equivalence.py (1 GPU) — the dispatch/compute/
    combine chain against a naive per-token MoE.

Full unit suite unchanged against main: same 13 pre-existing failures, none
added. pre-commit clean over the PR range.

🤖 Generated with Claude Code

GpuAsyncAFDConnector drives AFD's Attention-to-FFN traffic over the symmetric
window with no control plane at all. The Attention side routes tokens itself
and one-sided-writes the already-grouped rows into every FFN peer's slot; the
FFN side learns the layer and the token counts from the arriving payload, so
nothing has to be broadcast ahead of a step and the two roles never rendezvous
per layer.

`plan_dispatch` produces one row per (token, expert slot), sorted by global
expert, so a destination's partials are one contiguous run already grouped for
its grouped GEMM. Nothing about the routing is ever read back to the host: the
per-destination counts go into the slot header on the device, which is what
takes the last synchronize off the layer path.

The dispatch sequence number lives on the device for the same reason the
window's flag protocol uses a constant marker -- a replay runs no Python, so a
host-side counter would freeze at whatever was live at capture.

AFDAsyncTopology and build_async_topology move out of the NPU CAM connector
into connectors/async_topology.py so both async connectors share one
role-rank-to-world-rank mapping, and AFD_ASYNC_CONNECTOR becomes
AFD_ASYNC_NPU_CONNECTOR now that there is more than one.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: specture724 <specture724@gmail.com>
@specture724
specture724 force-pushed the afd/async-gpu-connector branch from 63b3f75 to 0f09249 Compare September 10, 2026 11:42
@specture724 specture724 changed the title [2/4] feat: the async GPU connector on the symmetric window [2/5] feat: the async GPU connector on the symmetric window Sep 11, 2026
@specture724
specture724 added this pull request to stack #337 September 11, 2026 02:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant