Skip to content

[GPU] Add MRV2 async GPU connector with NVSHMEM cross-node transport - #580

Draft
specture724 wants to merge 2 commits into
vllm-project:mainfrom
specture724:codex/async-gpu-mrv2
Draft

specture724 wants to merge 2 commits into
vllm-project:mainfrom
specture724:codex/async-gpu-mrv2

Conversation

@specture724

Copy link
Copy Markdown
Collaborator

Purpose

Add an experimental GpuAsyncAFDConnector for the pinned GPU ModelRunnerV2 runtime, including communication between Attention and FFN workers on different machines. Attention replicas can advance independently while native FFN experts serve incoming work.

Summary

Attention: gate/top-k + native MRV2 microbatches
    -> mapped peer writes OR symmetric outbox + NVSHMEM put/signal
FFN: local expert execution + independent compute graphs
    -> separate routed/shared replies -> Attention FP32 combine

Issue

No linked issue. Reference implementation: spike/AsyncGPUConnector at 788d17d.

Scope

  • In scope: Attention-side routing, static FFN expert ownership, MRV2 graph/DBO integration, asynchronous dispatch/combine and drain, mapped-peer and cross-node transport, optional CUDA extension, tests and Slurm recipes.
  • Qualification is limited to H20 / DeepSeek-V2-Lite-Chat BF16, Attention DP2/TP1 and FFN DP2/TP1/EP2. Cross-node NVSHMEM requires equal PE counts on each node.
  • Out of scope: IBGDA/performance qualification, quantization, Attention TP2, V3/V4 hardware qualification, and recovery of individual failed PEs. The V4 recipe remains a candidate.

Implementation Notes

  • Use plugin-owned connector/model-runner components and native local expert kernels. The MRV2 ubatching shim carries the live connector and stage through child forward contexts; no per-layer DP control-plane rendezvous is added.
  • Give each stage/ring its own mutable state. A CUDA/NVSHMEM bridge reads dispatch counters on the device on every graph replay, publishes a complete slot with a uint64 signal, and completes RMA before reusing its symmetric outbox. Receive wait/test provides payload visibility.
  • Keep A-first logical routing ranks; bootstrap node-contiguous NVSHMEM PEs and translate both peer-pointer queries and remote destinations. Reject unequal node counts before native initialization.
  • Keep route weights and accumulation in FP32, and routed/shared BF16 wire replies separate to preserve cancellation precision.
  • Build GPU operations only with AFD_BUILD_GPU_ASYNC_OPS=1; preserve a pure Python installation and same-node operation without the extension. Add the missing NPU patch package marker so fresh wheels include its existing four modules.
  • Document build prerequisites, heap sizing, topology, coordinated shutdown and the measured qualification scope.

Test Plan

  • CPU seam coverage: python -m pytest -q tests/unit -m 'not gpu and not vllm_runtime'.
  • Same-node transport: python -m pytest -sv tests/gpu/test_async_gpu.py.
  • Two-node transport: sbatch tests/gpu/run_async_gpu_multinode.slurm (affine expert oracle; model kernels are checked separately).
  • Model checks: native fixed-prefix probability comparison and tests.e2e.multi_pod.runner --scenario afd-v2-graph-dbo-dp2 --pod-layout 2A0F,0A2F.
  • Build both wheel variants, compare packaged sources, test installation without Torch/vLLM, and exercise installed wheels on real GPUs.

Test Result

Qualified source base: AFD e58fdec5876782b9c6a03683416611deea42af92. Target/runtime: vLLM 0.32.0, target SHA 5df63a7203cdb9328d824a27b682e77c7b9928be; Python 3.12.15, Torch 2.13.0+cu130, NVSHMEM 3.4.5. Native GPU operations built on H20 with CUDA 13 / GCC 10.2.1; no vLLM checkout edits.

Check Result
CPU unit suite 1,584 JUnit events, 45 environment/device skips, zero failures/errors
Same-node transport Five cases passed, zero skips
Two-node IBRC/RoCE transport Eager, 64 graph replays with changed inputs/weights and int32 overflow, precision, independent Attention producers and collective close passed
Mixed local/remote topology 2A2F with 1A1F per node passed 64 graph replays after PE remapping; also passed from the final GPU wheel
Real model, eager / graph+DBO 819 fixed-prefix positions; all 753 high-confidence native argmax predictions match; maximum logprob deltas 0.484016 / 0.429779 (native batch variation: 0.482730)
Live graph+DBO evidence Attention graph replays 80/84 and two-stage steps 83/83; FFN graph replays 4/5, with both stages and both Attention sources observed
Final two-node GSM8K 85/128 = 66.40625%, strict and flexible extraction; 8-shot, concurrency 12; both runners passed and normal teardown left no GPU processes
Packaging Both wheels contain all 102 Python sources byte-identically; pure wheel imports without Torch/vLLM and passed three same-node GPU checks; GPU wheel passed the mixed-node oracle
Static checks Ruff, mypy 3.11/3.12, Markdown, shell syntax, shellcheck and spelling checks passed

The native GSM8K reference (80/128) is the previously qualified identical runtime/checkpoint, not a new baseline run. Scores vary between BF16 generation runs; these results do not establish a performance or accuracy improvement. The multi-pod runner does not collect cross-pod replay counters; the live graph/DBO evidence above comes from the separate model probability request window.

GPU validation used Slurm-allocated GPUs on vllm-sgp-h20-01 and vllm-sgp-h20-03; the allocation was released. NPU hardware was not exercised. Detailed runtime identity, reproduction commands, numerical bounds and limits are in the validation record. Raw logs/wheels remain local under .venv/validation-cross-node/ and are not committed.

Evidence

  • Before: the reference implementation targets the old runner, and mapped-peer-only transfers cannot reach remote PEs. A mixed-role A-first PE layout also reproduced a native NVSHMEM initialization timeout.
    After: native MRV2 graph/DBO and two-node model checks pass; node-contiguous transport IDs unblock the same mixed-role topology without changing logical routing ranks.

Docs Impact

  • Updated configuration/support matrices, connector and runtime module contracts, GPU test instructions and the docs index.
  • Added the CUDA extension build guide, single-/two-node validation record and a two-node 2A2F V2-Lite Slurm recipe.

Merge Danger

Door: two-way. Revert and restart the whole worker group; protocol v5 uses 64-bit signals and cannot mix with the earlier v4 window.

Blast Radius: runtime. The connector is opt-in, but the change touches model proxies, MRV2 graph/ubatching and shared FFN/configuration seams. Architectural review of the compatibility shim and stream/signal ordering is still required before merge. Hardware results qualify the recorded source base; latest-main integration is covered by subsequent CI, not by the archived GPU run.


Essential PR Checklist
  • Purpose is clear and linked to public context when possible.
  • Scope is bounded.
  • Compatibility with vLLM v0.32.0 and the scope of NPU changes are considered.
  • No changes are made to the vLLM source checkout.
  • Plugin-owned classes or explicit dotted class paths are preferred over monkey patches.
  • Complete architectural review of the isolated, documented MRV2 compatibility shim.
  • Imports remain CPU-safe; CUDA-heavy work is delayed or GPU-gated.
  • Validation evidence and skips are included.
  • Documentation impact is stated.

Add Attention-side routing, connector-driven FFN execution, independent
stage buffers and graph-safe NVSHMEM communication for the pinned MRV2
runtime. Keep mapped peer writes locally and use symmetric outboxes with
64-bit signals for remote peers.

Include the optional CUDA bridge, balanced PE bootstrap mapping, CPU/GPU
coverage, Slurm recipes and single-/two-node validation records.

Signed-off-by: Anjie Hou <149605198+specture724@users.noreply.github.com>
Signed-off-by: Anjie Hou <149605198+specture724@users.noreply.github.com>
@github-actions github-actions Bot added documentation Improvements or additions to documentation Ascend Ascend NPU platform and related changes NVIDIA NVIDIA GPU, CUDA, and related changes labels Nov 14, 2026
@jiangkuaixue123

Copy link
Copy Markdown
Collaborator

Is there any performance results?

@specture724

Copy link
Copy Markdown
Collaborator Author

Is there any performance results?

Not yet, I'll attach them while they are done.

@jiangkuaixue123 jiangkuaixue123 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One device-selection issue in the new async connector.

group_name=AFD_ASYNC_GPU_GROUP_NAME,
timeout=timedelta(seconds=self.afd_config.afd_process_group_timeout_s),
)
device = torch.device("cuda", self.local_rank)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Resolve the worker's mapped CUDA device

group.local_rank is a logical rank: pinned vLLM maps it through current_platform.logical_device_id_to_visible_device_id before selecting the worker device, but preserves the logical value passed to this connector. With all four GPUs visible and --device-ids 0 1 for Attention / --device-ids 2 3 for FFN, this line instead puts both roles' window tensors on CUDA 0/1. The first window-geometry NCCL collective then uses duplicate devices, and the window device also disagrees with the FFN model/NVSHMEM allocation. Resolve the mapped visible device here (or pass the actual worker device), and cover a nonidentity --device-ids assignment.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Ascend Ascend NPU platform and related changes documentation Improvements or additions to documentation NVIDIA NVIDIA GPU, CUDA, and related changes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants