[2/5] feat: the async GPU connector on the symmetric window - #333
Open
specture724 wants to merge 1 commit into
Open
[2/5] feat: the async GPU connector on the symmetric window#333specture724 wants to merge 1 commit into
specture724 wants to merge 1 commit into
Conversation
specture724
requested review from
hsliuustc0106 and
jiangkuaixue123
as code owners
September 10, 2026 09:20
This was referenced Sep 10, 2026
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
force-pushed
the
afd/async-gpu-connector
branch
from
September 10, 2026 11:42
63b3f75 to
0f09249
Compare
specture724
added this pull request to stack #337
September 11, 2026 02:17
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Second of a five-PR stack. Based on #332 — the diff shown here is only this
PR's own change; review #332 first.
What this adds
GpuAsyncAFDConnectordrives AFD's Attention-to-FFN traffic over the symmetricwindow 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_dispatchproduces one row per(token, expert slot), sorted by globalexpert, 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
AFDAsyncTopologyandbuild_async_topologymove out of the NPU CAM connectorinto
connectors/async_topology.pyso both async connectors share onerole-rank-to-world-rank mapping (logic unchanged — pure hoist). With more than
one async connector now,
AFD_ASYNC_CONNECTORbecomesAFD_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 theprotocol: 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 areference, 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, noneadded.
pre-commitclean over the PR range.🤖 Generated with Claude Code