Skip to content

[3/5] feat: run DeepSeek-V4 on the async GPU connector - #334

Open
specture724 wants to merge 1 commit into
afd/async-gpu-connectorfrom
afd/async-gpu-v4
Open

[3/5] feat: run DeepSeek-V4 on the async GPU connector#334
specture724 wants to merge 1 commit into
afd/async-gpu-connectorfrom
afd/async-gpu-v4

Conversation

@specture724

@specture724 specture724 commented Sep 10, 2026

Copy link
Copy Markdown
Collaborator

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

  1. [1/5] feat: NVSHMEM symmetric window for GPU AFD transport #332 — NVSHMEM symmetric window (transport)
  2. [2/5] feat: the async GPU connector on the symmetric window #333 — the async GPU connector
  3. this PR — 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

Wires the connector into both AFD roles and gives DeepSeek-V4-Flash — 256
experts at topk 6, where the per-layer control-plane round trip costs more than
the expert compute it guards — a recipe that uses it.

Both GPU runners previously asserted control_plane is not None, and the FFN
worker loop raised NotImplementedError without one, so the connector could not
run at all:

  • FFN side pulls one work item at a time from the connector's own receive
    loop, taking the layer index and row counts from the arriving payload, and
    returns on an idle poll so the worker loop still sees its shutdown event.
  • Attention side skips vLLM's cross-DP batch agreement. Async AFD lets each
    replica advance alone, so an idle replica never joins that all-reduce and a
    busy one blocks in it forever — which is exactly where a 2A2F run hung before
    reaching the first MoE layer.

The V4 adapter learns the expert-routed dispatch protocol (the gate runs on the
Attention side, so the wire carries topk ids and weights instead of token ids)
and compute_ffn_output takes a device-side group_list so the FFN side runs
only the grouped GEMM over its local experts. The V2 adapter gets the same GPU
entry point, which is what the connector's e2e tests use as their reference.

Behaviour change worth a look

layers.N.ffn.gate.* now loads on both roles. The gate's parameters live
under .ffn so the checkpoint names resolve, but with
compute_gate_on_attention the Attention side is what runs it.
test_v4_gate_loads_on_both_roles pins the new contract.

The rest

Found bringing 2A2F up on real weights: the FFN role must force vLLM's NoDP MoE
prepare/finalize under EP with DP>1, shared experts must be skipped on an empty
shared slice, the SWIGLUOAI clamp has to reach the routed experts, and the DP
coordinator's startup wait needs to be long enough for the second role's weights
to load.

Testing

Full unit suite unchanged against main: same 13 pre-existing failures, none
added. pre-commit clean over the PR range. The recipe is
recipe/gpu/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh (4 GPUs).

🤖 Generated with Claude Code

Wires the connector into both AFD roles and gives DeepSeek-V4-Flash -- 256
experts at topk 6, where the per-layer control-plane round trip costs more
than the expert compute it guards -- a recipe that uses it.

Both GPU runners previously asserted `control_plane is not None`, and the FFN
worker loop raised NotImplementedError without one, so the connector could not
run at all. The FFN side now pulls one work item at a time from the
connector's own receive loop, taking the layer index and row counts from the
arriving payload, and returns on an idle poll so the worker loop still sees
its shutdown event. The Attention side skips vLLM's cross-DP batch agreement:
async AFD lets each replica advance alone, so an idle replica never joins that
all-reduce and a busy one would block in it forever -- which is exactly where
a 2A2F run hung before reaching the first MoE layer.

The V4 adapter learns the expert-routed dispatch protocol (the gate runs on
the Attention side, so the wire carries topk ids and weights instead of token
ids) and `compute_ffn_output` takes a device-side group_list so the FFN side
runs only the grouped GEMM over its local experts. The V2 adapter gets the
same GPU entry point, which is what the connector's e2e tests use as their
reference.

The rest was found bringing 2A2F up on real weights: the FFN role must force
vLLM's NoDP MoE prepare/finalize under EP with DP>1, shared experts must be
skipped on an empty shared slice, the SWIGLUOAI clamp has to reach the routed
experts, the V4 Attention-side gate loads under its checkpoint path, and the
DP coordinator's startup wait needs to be long enough for the second role's
weights to load.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: specture724 <specture724@gmail.com>
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