Repository navigation
[GPU] Add MRV2 async GPU connector with NVSHMEM cross-node transport - #580
specture724 wants to merge 2 commits into
Conversation
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>
|
Is there any performance results? |
Not yet, I'll attach them while they are done. |
jiangkuaixue123
left a comment
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
[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.
Purpose
Add an experimental
GpuAsyncAFDConnectorfor 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
Issue
No linked issue. Reference implementation:
spike/AsyncGPUConnectorat 788d17d.Scope
Implementation Notes
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.Test Plan
python -m pytest -q tests/unit -m 'not gpu and not vllm_runtime'.python -m pytest -sv tests/gpu/test_async_gpu.py.sbatch tests/gpu/run_async_gpu_multinode.slurm(affine expert oracle; model kernels are checked separately).tests.e2e.multi_pod.runner --scenario afd-v2-graph-dbo-dp2 --pod-layout 2A0F,0A2F.Test Result
Qualified source base: AFD
e58fdec5876782b9c6a03683416611deea42af92. Target/runtime: vLLM 0.32.0, target SHA5df63a7203cdb9328d824a27b682e77c7b9928be; 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.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-01andvllm-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
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
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