Repository navigation
EP: GPUNetIO dispatch/combine with pipelining (stacked on #910) - #935
Open
Qinghua Zhou (seagater) wants to merge 162 commits into
Open
Qinghua Zhou (seagater) wants to merge 162 commits into
Qinghua Zhou (seagater) wants to merge 162 commits into
Conversation
Port DeepEP's high-throughput MoE dispatch/combine kernels onto MSCCL++
as an optional build target `mscclpp_ep_cpp`, gated by -DMSCCLPP_BUILD_EXT_EP
(OFF by default). Sources are lifted from DeepEP branch
`chhwang/dev-atomic-add-cleanup` and rebased onto upstream MSCCL++ APIs;
the NVSHMEM / IBGDA dependencies are replaced with `PortChannel` +
`MemoryChannel` + the new `Connection::atomicAdd` primitive.
Scope
-----
Intranode (NVLink-only):
* `Buffer` ctor/dtor: cudaMalloc nvl workspace, export IPC handle,
allocate FIFO + peer-pointer tables, start `ProxyService`.
* `sync()`: import peer IPC handles, upload peer pointer table,
build `MemoryDevice2DeviceSemaphore` + `MemoryChannel` per peer.
* `get_dispatch_layout`, `intranode_dispatch`, `intranode_combine`
ported verbatim (torch::Tensor ABI preserved).
Internode HT (NVLink + RDMA):
* `sync()` RDMA branch: cudaMalloc RDMA buffer + `bootstrap->barrier()`
(replacing NVSHMEM symmetric-heap allocation); register with
`all_transport`, exchange via `sendMemory`/`recvMemory`, build 12 IB
QPs/peer + 16 semaphores/peer + 16 port channels/peer.
* Full `internode.cu` port (notify_dispatch / dispatch / cached_notify
/ combine / get_dispatch_layout). The 4 raw `ChannelTrigger` atomic
sites are rewritten to call the new
`PortChannelDeviceHandle::atomicAdd(offset, value)` API; the single
`nvshmem_fence()` is replaced with `__threadfence_system()` (remote
visibility guaranteed by the subsequent port-channel barrier).
* `internode_dispatch` / `internode_combine` host code ported, with
the torch tensor marshalling and CPU spin-wait on mapped counters.
Low-latency (pure RDMA):
* Not ported. `low_latency_dispatch`, `low_latency_combine`,
`clean_low_latency_buffer`, `get_next_low_latency_combine_buffer`
throw `std::runtime_error`; the Python frontend refuses to
construct a Buffer with `low_latency_mode=True`.
Python layer
------------
* New pybind11 + libtorch Python extension `mscclpp_ep_cpp` (separate
from the nanobind `_mscclpp` because the EP ABI carries
`torch::Tensor` / `at::cuda::CUDAStream`).
* `mscclpp.ext.ep.Buffer` mirrors `deep_ep.Buffer`; exchanges device
IDs, IPC handles and the bootstrap UniqueId over the user's
`torch.distributed` process group before calling `sync()`.
* `mscclpp.ext` auto-imports `ep` if the extension is built.
Build
-----
* `src/ext/ep/CMakeLists.txt`: finds Python + Torch; warns and skips if
`CMAKE_PREFIX_PATH` doesn't point at `torch.utils.cmake_prefix_path`.
Falls back to Torch's bundled pybind11 if a standalone pybind11 is not
installed. Links `libtorch_python` explicitly (without it, `import
mscclpp_ep_cpp` fails with `undefined symbol: THPDtypeType`).
* Top-level `CMakeLists.txt` exposes the `MSCCLPP_BUILD_EXT_EP` option
(default OFF).
Tests
-----
* `test/python/ext/ep/test_ep_smoke.py`: skipped if the extension isn't
built. Covers Config round-trip, low-latency size hint, and the LL
construction guard. Multi-rank functional tests still to do on H100.
Notes
-----
* Builds against the preceding "atomic add" commit which adds
`Connection::atomicAdd` and `PortChannelDeviceHandle::atomicAdd` to
upstream MSCCL++.
* Intranode path verified end-to-end (build + import + smoke tests).
* Internode HT is code-complete but requires real IB hardware to
validate; see `src/ext/ep/README.md` for the detailed port plan and
remaining LL migration.
Port DeepEP's pure-RDMA low-latency (LL) MoE kernels from
csrc/kernels/internode_ll.cu (branch chhwang/dev-atomic-add-cleanup)
into the MSCCL++ EP extension. NVSHMEM / IBGDA device primitives are
replaced with MSCCL++ PortChannelDeviceHandle operations:
nvshmemx_barrier_all_block() -> port-channel signal+wait ring
nvshmemi_ibgda_put_nbi_warp(...) -> lane-0 PortChannel.put(...)
nvshmemi_ibgda_amo_nonfetch_add(...) -> lane-0 PortChannel.atomicAdd(...)
The atomicAdd path relies on the MSCCL++ Connection::atomicAdd /
PortChannelDeviceHandle::atomicAdd API cherry-picked from branch
chhwang/new-atomic-add; the LL dispatch path uses a signed delta
(-num_tokens_sent - 1) which the new int64_t signature supports.
Changes:
* New file src/ext/ep/kernels/internode_ll.cu (~530 lines) with the
three kernels clean_low_latency_buffer, dispatch<kUseFP8,...>,
combine<...> plus their launchers. rdma_buffer_ptr is threaded
through the launchers so the kernel can translate virtual addresses
into registered-memory offsets expected by MSCCL++.
* kernels/api.cuh: replace the single stub signature with full LL
launcher prototypes.
* buffer.cc: replace the four LL throw-stubs
(clean_low_latency_buffer, low_latency_dispatch,
low_latency_combine, get_next_low_latency_combine_buffer) with
torch-Tensor implementations ported from DeepEP/csrc/deep_ep.cpp.
* Drop src/ext/ep/internode_stub.cc and its CMake entry.
* python/mscclpp/ext/ep/buffer.py: remove the low_latency_mode=True
NotImplementedError guard; update docstring.
* test/python/ext/ep/test_ep_smoke.py: rename
test_low_latency_rejected -> test_low_latency_buffer_construct
to reflect that LL construction is now accepted.
* src/ext/ep/README.md: update status matrix, document the
NVSHMEM -> MSCCL++ translation table, and list the known
limitations.
This is a structural port: the kernels compile, link, and pass the
single-rank smoke tests, but end-to-end behaviour on multi-node H100
is not yet validated. Two known caveats:
1. Performance will NOT match IBGDA because MSCCL++ port channels
use a CPU proxy; this port is for functional parity, not latency.
2. Buffer::sync() in LL mode only connects peers that share the
same local GPU id (DeepEP convention), so the LL kernels assume
a one-GPU-per-node topology (num_ranks == num_rdma_ranks).
Multi-GPU-per-node LL layouts will need a follow-up in sync().
Tested:
cmake --build build -j --target mscclpp_ep_cpp # builds clean
pytest test/python/ext/ep/test_ep_smoke.py # 3 passed
Three issues blocked end-to-end intranode validation across multiple ranks. This commit fixes them and adds a 2/4/8-rank functional test. 1. Combine receiver: OOB __shared__ read In the combine receiver warp, the wait loop evaluated `channel_tail_idx[recv_lane_id] <= expected_head` before the `expected_head >= 0` guard. `channel_tail_idx` is a shared array of size `kNumRanks`, but the loop runs on all 32 lanes of a warp, so lanes with `recv_lane_id >= kNumRanks` indexed out of bounds. compute-sanitizer reported "Invalid __shared__ read of size 4 bytes" at combine<bf16,2,768>+0xdd0, surfaced asynchronously as cudaErrorIllegalAddress at the kernel launch site. Swap the operands so the rank-bounds check short-circuits the shared read. 2. Python bindings: UniqueId ABI `mscclpp::UniqueId` is a `std::array<uint8_t, N>` which pybind11 auto-converts to a Python `list`, silently overriding any `py::class_<UniqueId>` wrapper. Expose `create_unique_id` / `connect` as lambdas that produce/consume `py::bytes` and memcpy into a local `UniqueId`. Also coerce `bytes`->`bytearray` at the Python call site for `sync()` whose signature expects `pybind11::bytearray`. 3. Python frontend: communicator required for NVL-only sync `Buffer::sync()` uses `communicator->connect(ipc_config, ...)` on the pure-NVLink path, so the communicator must be initialized even when `num_rdma_ranks == 1` and `low_latency_mode == False`. Always broadcast the unique id and call `runtime.connect()` before `sync()`. Validation on a single H100x8 node via torchrun: - 2 ranks: dispatch 195 tokens, combine diff=0 - 4 ranks: dispatch 371 tokens, combine diff=0 - 8 ranks: dispatch 456 tokens, combine diff=0 Test harness added at test/python/ext/ep/test_intranode_multirank.py.
The `internode` kernels index device-side port channel handles as
`port_channel_handles[channel_id * num_ranks + peer_rank]`, where
`peer_rank` is a global rank in [0, num_ranks). `Buffer::sync` was
building that table by iterating `std::unordered_map<int, MemoryId>`
(and similarly for connections/semaphores), which yields hash order
rather than ascending rank order. Once the cross-node fan-out grew
beyond a single peer, a local rank's trigger for peer `r` landed on
the semaphore/memory pair of a different peer, so RDMA puts and
atomic tail updates went to the wrong destination and the forwarder
spun on a tail counter that never advanced.
Changes:
- Build `sema_ids` and `port_channel_handles` by iterating
`for (int r = 0; r < num_ranks; ++r)` and looking up the
connection / memory id for rank `r`, skipping ranks excluded by
low-latency mode (inserting a placeholder handle so the stride
stays `num_ranks`).
- Tag the RDMA-phase `sendMemory`/`recvMemory`/`connect` calls with
`kRdmaTag = 1` so they do not collide with NVL-phase tag-0
traffic between the same pair of ranks.
- Drop an unused `r` local in the NVL setup loop.
With this fix and a matched `libmscclpp.so` on both nodes, the
2-node x 8-GPU internode HT dispatch path completes successfully
(`[dispatch] OK`). Combine is still under investigation.
Also adds `test/python/ext/ep/test_internode_multirank.py`, a
torchrun-based 2-node functional test that exercises
`get_dispatch_layout` -> `internode_dispatch` -> `internode_combine`
and validates per-source-rank token values end-to-end.
Two issues prevented internode HT combine from completing on 2x8 H100: 1. Wrong prefix matrices passed to internode_combine. Combine runs in the reverse direction of dispatch, so it must consume the receiver-side matrices returned by dispatch (recv_rdma_channel_prefix_matrix, recv_rdma_rank_prefix_sum, recv_gbl_channel_prefix_matrix), not the sender-side rdma_channel_prefix_matrix / gbl_channel_prefix_matrix. This matches DeepEP's deep_ep/buffer.py::internode_combine handle unpacking. Without the fix the NVL forwarder's 'NVL check' timed out because token_start_idx/token_end_idx were computed against the wrong per-channel layout. 2. Cross-rank race between dispatch and combine. Even with the correct matrices, launching combine immediately after dispatch deadlocked the forwarder NVL check (tail stuck one short of expected_head) because peers still had in-flight dispatch proxy traffic while fast ranks had already started combine. A torch.cuda.synchronize() + dist.barrier() between the two calls makes the test pass deterministically on 16 ranks (combine diff == 0, max|expected| up to 60.0). The barrier in the test is a workaround; the real fix belongs in Buffer::internode_dispatch / Buffer::internode_combine so the dispatch->combine handoff fully fences outstanding proxy work across ranks. Marked with an XXX comment in the test.
Refresh status docs and comments now that internode HT dispatch and combine have been validated end-to-end on 2 nodes x 8 H100 GPUs via test/python/ext/ep/test_internode_multirank.py (all 16 ranks recover their per-rank token payloads with zero diff). - src/ext/ep/README.md: consolidate the previously duplicated README into a single document; mark intranode and internode HT dispatch and combine as validated in the status table; add a 'Running the tests' section with torchrun examples for both the intranode and the 2x8 internode setups; record the dispatch->combine torch.cuda.synchronize() + dist.barrier() requirement under Known limitations; mark Phase 2 DONE and keep Phase 3 (LL) as structural port, untested. - python/mscclpp/ext/ep/buffer.py: update the module docstring and the Buffer constructor docstring to say internode HT is validated and clarify that LL mode is untested on multi-node hardware. - src/ext/ep/buffer.cc: drop the stale 'NVSHMEM support not yet ported' and 'low-latency paths still stubbed' comments. mscclpp_ep does not use NVSHMEM at all (PortChannel/MemoryChannel replace it), and the LL paths are a structural port that is present but untested, not stubbed. Note validation on 2x H100x8 in the internode section header.
- Buffer::sync no longer drops non-same-GPU-id peers in low_latency_mode. DeepEP's original filter was safe because its LL path used NVSHMEM; this port drives LL via PortChannel so the kernel indexes port_channel_handles[local_expert*num_ranks + dst_rank] for every dst_rank. All peers now get a real memory/connection/semaphore/port channel entry. - Add test/python/ext/ep/test_low_latency_multirank.py (LL dispatch+combine functional round-trip, BF16 only). Works cross-node in DeepEP's 1-GPU-per-node topology. - Known limitation documented in src/ext/ep/README.md and the test docstring: intra-node 8-GPU LL currently hangs because every peer transfer routes through the CPU proxy over IB loopback between distinct HCAs on the same host, and (separately) CudaIpcConnection::atomicAdd is a 64-bit op which mis-aligns the 32-bit rdma_recv_count slots when used for same-node peers. Proper fix needs a mixed-transport LL variant (MemoryChannel for same-node, PortChannel for cross-node) or 64-bit counters.
Gated behind MSCCLPP_EP_BENCH=1 to keep correctness runs fast. Reports per-iter latency (max across ranks, CUDA-event timed) and aggregate effective bandwidth (sum across ranks, dispatch+combine payload bytes). Tunable via MSCCLPP_EP_BENCH_WARMUP / _ITERS / _TOKENS / _HIDDEN. Bench reuses the Buffer allocated for the correctness phase and self-skips if the requested hidden exceeds the per-peer NVL/RDMA budget.
Previously the optional benchmark measured full round-trip latency. Split it to time dispatch alone (N iters) and combine alone (N iters reusing one dispatch output), reporting per-phase latency (max across ranks) and aggregate effective bandwidth (sum across ranks). Applies to intranode HT, internode HT, and the (currently unreachable on intra-node 8-GPU) LL test. Internode HT keeps the sync+barrier guard between dispatch and combine but excludes it from either phase's timing.
…o int64 The low-latency dispatch/combine kernels signal recv counts via MSCCL++ PortChannel.atomicAdd, which lowers to IB IBV_WR_ATOMIC_FETCH_AND_ADD. That opcode requires the remote address to be 8-byte aligned, but LowLatencyLayout packed the per-expert signaling slots as int32. Odd slots landed at offset %8 == 4; the NIC silently dropped those atomics and the target rank spun forever in recv_hook (observed: even->odd direction works, odd->even does not, across all tested topologies including 2-rank intra-node, 8-rank intra-node, and 2-node 1-GPU-each). Widen dispatch_rdma_recv_count_buffer / combine_rdma_recv_flag_buffer to int64_t, update clean kernel + kernel signatures + next_clean pointers accordingly, and add int64_t overloads for st_na_release / ld_acquire_sys_global in utils.cuh. Also drop the bogus self CUDA-IPC connection in Buffer::sync() that was previously skewing the cross-rank buildAndAddSemaphore handshake order; the kernel's same-rank branch uses a direct warp copy and never touches the self port-channel slot (filled with a zero-initialized placeholder so the [local_expert*num_ranks + dst_rank] indexing still holds).
Dropping the self ipc_cfg connection caused cudaErrorInvalidResourceHandle on multi-node launches. Keep the self connection (needed by other code paths that assume every rank is in the connections map) but continue to skip the self slot in the semaphore + port-channel construction loops so the kernel's [local_expert*num_ranks + dst_rank] indexing hits only peer handles; the self slot is a zero-initialized placeholder since the kernel's same-rank branch uses a direct warp copy.
The prior commit skipped r==rank in the semaphore and port-channel build loops on the theory that the self-slot handshake skew was the cause of LL direction asymmetry. That was wrong (the real bug was int32 atomic alignment), and skipping self breaks other code paths that assume every rank slot is represented -- cross-node HT and LL failed with cudaErrorInvalidResourceHandle at the first barrier after Buffer init. Restore the self-inclusive loop.
When all ranks live on the same host (num_rdma_ranks == 1), the LL
kernels now bypass PortChannel/IB-loopback entirely. In Buffer::sync()
we additionally:
- allGather IPC handles for each rank's rdma_buffer_ptr and
cudaIpcOpenMemHandle them into peer_rdma_bases[]
- build per-peer MemoryChannels over CUDA IPC connections (tag=2)
used only for the LL barrier ring
The three LL kernels (clean / dispatch / combine) gain a kIpcPath
template parameter and two extra args (peer_rdma_bases,
memory_channel_handles). At each peer op:
- put -> peer-mapped warp copy over NVLink
- atomicAdd-like flag store -> single-writer st_na_release on peer ptr
- signal/wait barrier -> MemoryChannel signal/wait
Cross-node LL (num_rdma_ranks > 1) is untouched; the IPC setup block is
a no-op. The host launch wrappers select the variant via use_ipc_path.
Each local expert sends one copy per dispatched token back to its owner, so the bytes actually on the wire during combine match dispatch. The previous num_tokens×hidden under-counted by ~num_topk×, making combine BW look artificially low next to dispatch.
- Report both per-rank and aggregate BW to align with NCCL-EP's ep_bench (which reports per-rank GB/s). - Accept MSCCLPP_EP_LL_TOKENS/HIDDEN/TOPK/EXPERTS_PER_RANK env overrides so we can match external benchmark problem sizes (NCCL-EP LL defaults are num_tokens=128, hidden=7168, top_k=8).
Same alignment with NCCL-EP ep_bench as the LL test: report both per-rank (agg/num_ranks) and aggregate throughput.
LL dispatch/combine are latency-bound at typical problem sizes: for num_experts=32 the previous grid was cell_div(32,3)=11 blocks, i.e. 8% of a 132-SM H100. The recv-side bodies already stride tokens by sm_id, so extra blocks parallelize token work linearly. Extra blocks past num_experts are gated out of the send/count phases by the existing 'responsible_expert_idx < num_experts' check. Cap at the device's SM count (cooperative launch + launch_bounds(960,1) allow one block per SM).
On the PortChannel (cross-node) path the extra blocks don't help: the dispatch recv loop strides tokens per-warp-group (not per-SM), and the additional blocks instead add cooperative-grid sync overhead and increase concurrent host-proxy FIFO traffic. Measured cross-node dispatch regressed from 1013us to 3063us when the unconditional grid bump was active. Keep the scaled grid for the IPC path (intra-node), where combine-recv and dispatch token striding scale with sm_id and the 1.2-1.3x speedup reproduces.
The LL combine benchmark was cloning the ~58 MB dispatch recv buffer
('recv_x.clone()') on every timed iteration, adding ~20 us of D2D
memcpy per sample and masking kernel-level changes. It also called
torch.empty() for the output inside the loop. Both now live outside
the timed region; the kernel is invoked against a persistent bench_out
and the recv_x produced by the most recent dispatch.
NCCL-EP's LL dispatch/combine kernel uses (numWarpGroups=1,
numWarpsPerGroup=32) when num_experts <= device_num_sms, giving each
SM ownership of a single expert and 32 warps to cooperate on its
recv-side per-(expert, src_rank) work. We were using (3, 10) — 3
experts per SM, 10 warps per (expert, rank) pair — which left a
significant amount of recv-side parallelism on the table because each
warp had to walk ~3x more tokens sequentially.
Switching to (1, 32) for both dispatch and combine matches NCCL-EP's
structure for typical EP sizes (num_experts in {32, 64, 256}) where
num_experts <= 132 SMs.
The static_assert kNumMaxTopK + 1 <= kNumWarpGroups * kNumWarpsPerGroup
still holds (9 <= 32) and the wider block also lets the staging loop
process the hidden-dim with one int4 per thread (hidden_bf16_int4=896
fits easily in 992 working threads).
Cross-node LL regressed when (1, 32) was applied uniformly: dispatch 1031us -> 1570us, combine 2553us -> 3484us. Larger grid means more concurrent putWithSignal calls onto the host-proxy FIFO and a costlier cg::this_grid().sync() between phases, both of which dominate the IB path even though more SMs help the recv-side compute. Make (kNumWarpGroups, kNumWarpsPerGroup) path-dependent: (1, 32) when use_ipc_path, (3, 10) otherwise. Restores cross-node performance and keeps the intra-node win.
- Add MSCCLPP_EP_BENCH_EXPERTS / _TOPK env knobs so the bench phase can match NCCL-EP's `ep_bench -a ht` defaults (256 experts, top-8). The functional check above continues to use the smaller (num_ranks*4 experts, topk=4) configuration. - Switch BW accounting from recv_tokens*hidden to bench_tokens*hidden, matching NCCL-EP's `RDMA_send` per-rank byte count. The previous formula counted DeepEP's expanded recv layout (one row per (token,src_rank) pair), inflating reported GB/s ~5x and making cross-stack comparisons misleading.
Same change as the intra-node bench (commit 4ed6f22), applied to the cross-node test: - Add MSCCLPP_EP_BENCH_EXPERTS / _TOPK env knobs so the bench phase can match NCCL-EP's `ep_bench -a ht` defaults (256 experts, top-8). - Switch BW accounting from recv_tokens*hidden to bench_tokens*hidden, matching NCCL-EP's `RDMA_send` per-rank byte count.
Each mscclpp::ProxyService spawns one host-side proxy thread that drains its FIFO and posts IB work requests. With LL combine pushing ~1k put + 60 atomicAdd FIFO entries per iter, that single thread is the wall-clock bottleneck on cross-node runs. Split the channel set across kNumProxyServices=4 separate services so the host-side dispatch parallelism scales linearly. SemaphoreIds and MemoryIds are scoped to a ProxyService, so: - addMemory() is broadcast to every service in the same global order so a single MemoryId still identifies the memory everywhere. - Each (peer_rank, channel_idx) is assigned to one proxy_idx via round-robin; the resulting PortChannel is built on that proxy and inherits its FIFO. The kernel is unchanged: the flat handle array routes the right way automatically. No kernel-level changes, no tuning of QP count, no new env knobs.
… 1 on Blackwell) Override at runtime with MSCCLPP_EP_NUM_PROXIES. N=8 is the knee on H100+IB; N>=12 collapses from CPU oversubscription. Intra-node LL is unchanged.
Add dist.barrier() + dist.destroy_process_group() in a finally block so non-zero ranks don't poll the TCPStore after rank 0 (the store server) exits, which produced noisy 'recvValue failed / Connection was likely closed' stack traces from ProcessGroupNCCL's HeartbeatMonitor. Also pass device_id to init_process_group in the internode test to silence 'Guessing device ID based on global rank' warnings.
Aligns with NCCL-EP's ep_bench convention (BW computed from average time across ranks). Previously we reported only the max time and computed BW per-rank, which made our numbers more pessimistic than NCCL-EP's.
Enable latency-mode TOKEN_MAJOR dispatch/combine support and wire the unified Python EP benchmark to exercise the new layout. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 03ca3fba-cdc2-4c2a-88ec-1ee8aa561b4a
Expose TOKEN_MAJOR top-k metadata as flat row-aligned tensors and make token-major dispatch publish deterministic invalid metadata for non-local lanes. This lets Doug KI token-major compute consume dispatch_out.topk_ids and dispatch_out.weights directly by token row. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 03ca3fba-cdc2-4c2a-88ec-1ee8aa561b4a
## Summary Integrate GPU-initiated networking into the newer feature/ep runtime and add the GPUNetIO-capable rank-major top-k-expanded layout, including topology-aware multi-HCA support and default-enabled expanded fastpaths. - **Base:** `feature/ep-experimental` (`63fea77` at creation). - **Head:** `qinghuazhou/gpunetio_port_channel_merge_feature_ep` (`824823a`). - Preserve the newer `DeviceContext`, `LatencyStorageLayout`, generic dispatch handles, and `MoEMode.LATENCY` APIs rather than restoring the legacy runtime API. ## GPUNetIO - Add the DOCA GPUNetIO/GDAKI device-initiated transport and host service for QP setup, symmetric-memory registration, and remote-key exchange. IPC/NVLink-mapped peers retain mapped-memory communication; GPUNetIO handles peers outside that domain. - Support multiple HCAs and QPs per peer, per-HCA registrations and keys, ordered completion markers, batched writes, and cooperative/deferred queue drains. Keep WQEs CQ-signaled for queue recycling and retain system-scope publication fences. - Auto-select all active HCAs tied for a GPU's best PCI/NUMA affinity, sorted by device name. Nearby GPUs may share the full HCA set instead of splitting it by local rank. - `MSCCLPP_EP_GPUNETIO_HCAS` overrides `MSCCLPP_EP_GPUNETIO_HCA`; otherwise topology selection is automatic. `MSCCLPP_EP_GPUNETIO_QPS_PER_PEER` defaults to the selected HCA count. Ranks must agree on HCA/QP counts; QPs must be in [1, 64] and a positive multiple of the HCA count. - Networking remains opt-in through the build option `MSCCLPP_USE_GPUNETIO` and runtime `MSCCLPP_EP_ENABLE_GPUNETIO=1`. Compact cross-domain rank-major uses BF16 rank-local-reduce; expert-major and compact direct-send remain IPC-only. ## Rank-Major Top-K-Expanded - Expose `DispatchLayout.RANK_MAJOR_TOPK_EXPANDED` with logical rows `[source_rank, original_token, topk_slot, hidden]`. Python tokens are `[world_size * capacity * topk, hidden]`, with flat row-aligned int32 expert IDs and FP32 weights. - Preserve duplicate top-k slots and fixed positions without compaction. Nonlocal, invalid, and padding metadata use the configured invalid ID and zero weight; inactive payload rows are unspecified. Per-rank counts include valid local selections, including duplicates and zero-weight selections. - Alias dispatched tokens and the registered combine-input buffer. The consumer writes **unweighted per-slot expert results** in place; combine applies the original source routing weights exactly once, in top-k order with FP32 accumulation, before writing BF16 output. Invalid IDs and zero weights are skipped before payload reads. - Restrict expanded mode to BF16, fixed configured capacity, top-k 1-9, serialized calls, and the `RANK_LOCAL_REDUCE` mode setting. Preserve existing compact rank-major behavior and buffer contracts. - Use TMA scatter/gather for mapped peers, registered staging and per-HCA pushes for network peers, device-generation completion, and collective lifetime acknowledgments to protect aliased-buffer reuse across CUDA-graph replays. - Add Python validation and benchmark selection via `--ep-layout rank_major_topk_expanded` for MSCCL++ latency mode, plus correctness coverage for padding, duplicate routes, oversized invalid IDs, zero weights/NaN inactive rows, and repeated graph pairs/replays. ## Expanded Fastpaths - Default-enable the IPC specialization only when every peer is mapped and no GPUNetIO context is active, or the network specialization when GPUNetIO is already active outside the IPC domain. Select the protocol collectively before graph capture. - Retain baseline kernels as opt-out paths: `MSCCLPP_EP_EXPANDED_IPC_FASTPATH=0` and `MSCCLPP_EP_EXPANDED_GPUNETIO_FASTPATH=0`. An opt-out disables the corresponding fastpath collectively; these flags do not enable networking itself. - Optimize worker retirement, cached receive readiness, metadata submission, and sparse warp-batched network returns without removing the final peer acknowledgment or outgoing queue drains. ## Validation Validation recorded during the port series: - Latest clean fastpath commit: **55 CPU regressions passed**, covering API/allocation contracts, collective selection/opt-out, retirement/ACK ordering, sparse WQE ticket wrap, per-HCA keys, and metadata transfer-size fallback. - **152 topology-selection checks passed**, including CPU-only CTest with GPUNetIO enabled and disabled. - Full SM100a EP builds passed with GPUNetIO ON/OFF; clean committed snapshots were also compiled on the compile-only host. Expanded kernels compiled for SM90 with GPUNetIO ON/OFF. - Changed-file formatting and diagnostics passed. Whole-repository C++ lint still reports pre-existing formatting debt. - **Distributed GPU correctness and performance remain unverified.** CPU models/stubs and compilation do not establish GPU/NIC ordering or performance. SM90 network-combine fastpaths report register spills (TMA: 36-byte stores/68-byte loads; vector: 40-byte stores/92-byte loads), requiring runtime evaluation. SM100a expanded compilation reported no spills. ## Merge Notes - GitHub currently reports **cannot automatically merge** with the requested base; conflict resolution is still required. - This comparison spans **206 source commits and 87 changed files**, including merged feature/ep architecture history and vendored GPUNetIO dependencies, not only the most recent fastpath commit. The base has 12 commits absent from the head. - Earlier local initialization/IPC diagnostic edits were deliberately excluded from the published commits and are not part of this PR. No runtime deployment or GPU jobs were performed as part of creating it. --------- Co-authored-by: Changho Hwang <changhohwang@microsoft.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: seagater <7475084+seagater@users.noreply.github.com> Co-authored-by: Binyang Li <binyli@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Caio Rocha <caiorocha@microsoft.com> Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> Co-authored-by: Caio Rocha <164253795+caiocbr@users.noreply.github.com> Copilot-Session: 15f71a84-4219-4ae9-a87e-e5fab4205de6 Copilot-Session: 59760d2a-e65f-44b6-b2e6-9c271b834d7e
## Summary - fix rank-major CUDA graph replay synchronization - fuse route weighting into rank-major combine - deduplicate expanded rank-major payloads --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…ntent GPUNetIO API) Port of 3cc3b6e (qinghuazhou/ep-experimental-gpunetio-with-pipeline) onto the EpGpuNetIoDeviceContext / PortChannelDeviceHandle GPUNetIO API. No vendored DOCA changes. - Batched rank-major dispatch with per-destination chunk pipelining (MSCCLPP_EP_GPUNETIO_DISPATCH_CHUNK_ROWS; default 4 for 2-7 nodes, 16 for >=8 nodes) - Parallel staging (MSCCLPP_EP_GPUNETIO_PARALLEL_STAGE, default 1) and staged-row fence (MSCCLPP_EP_GPUNETIO_ROW_FENCE=sys|gpu|gpu+poster, default gpu) - Deferred dispatch send flush via per-(peer, QP) tickets (MSCCLPP_EP_GPUNETIO_DEFER_FLUSH, default on); adds EpGpuNetIoDeviceContext::reservedTicket/waitTicket - Pipelined combine chunks with per-chunk signals (MSCCLPP_EP_GPUNETIO_COMBINE_CHUNK_ROWS) - Optional ASYNC_DISPATCH, EARLY_TAIL, COMBINE_NO_BARRIER, TMA_STAGE=1|2 - Python bench updates Dropped experiments that required DOCA changes (zero-copy gather, dedicated poster, coalesced submit, BlueFlame, NO_DBR, host CQ/SQ) and the buffered timing infrastructure. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 162bd212-6628-4e3a-93b9-190933b019fe
This reverts commit 53dd05b. Conflicts in src/ext/ep/combine/common.cuh resolved by keeping the GPUNetIO combine chunk pipeline (per-chunk landing gate, chunk-base advance) on top of the pre-#913 combine: relaxed ready-epoch load, no publisher/worker grid barrier (MSCCLPP_EP_GPUNETIO_COMBINE_NO_BARRIER is now a no-op), epoch-based synchronizeRankMajorCombine, and one row per token for chunked sends. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 162bd212-6628-4e3a-93b9-190933b019fe
…ain structure Main rewrote the EP extension (MoECommunicator API, latency/throughput runtimes, IPC rank-major protocol). Resolve the conflicts by taking main's EP structure and re-expressing the GPUNetIO features on top of it: - Keep GPUNetIO core: FetchContent dependency, src/gpunetio/host API, GpuNetIoService, PortChannel GPUNetIO device API and unit tests. - Add src/ext/ep/hybrid (namespace mscclpp::ep::hybrid): NVLink+GPUNetIO latency runtime for RANK_MAJOR + RANK_LOCAL_REDUCE with the dispatch chunk pipeline, parallel stage, row fence, deferred flush and combine chunking. Main's LatencyRuntimeContext routes to it only when GPUNetIO is requested (MSCCLPP_EP_ENABLE_GPUNETIO) for cross-IPC-domain runs or when the new layout is used; main's IPC kernels/storage are unchanged otherwise. - Add DispatchLayout::RANK_MAJOR_TOPK_EXPANDED (IPC + GPUNetIO fast paths), exposed through bindings and the Python MoECommunicator/request builder. - Honor MSCCLPP_EP_IPC_DOMAIN_SIZE override. - Drop the old Python EP API and test/python/ep scripts; add python/test/test_ep_hybrid.py. Validated on GB300: main test_ep.py + test_ep_hybrid.py 42 passed at 8 ranks; hybrid tests pass over GPUNetIO at 2 nodes; rank-major GPUNetIO perf at 2/4/8 nodes matches the pre-migration branch. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 162bd212-6628-4e3a-93b9-190933b019fe
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Add validation for MSCCLPP_EP_IPC_DOMAIN_SIZE to ensure it's a positive divisor of the detected IPC domain size. Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Sync with the PortChannel GPUNetIO base branch (PR #910), picking up its mp-unit test fixes, so the EP work can be reviewed as a diff on top of it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 162bd212-6628-4e3a-93b9-190933b019fe
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Qinghua Zhou (seagater)
requested review from
Binyang Li (Binyang2014),
Caio Rocha (caiocbr) and
Mahdieh Ghazi (mahdiehghazim)
and
a balanced review from Copilot
October 22, 2026 21:01
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Unresolved correctness and resource-lifetime issues can cause replay hangs, incorrect combine results, and retained GPU resources.
5 open issues
What changed in this PR
This PR builds on #910’s GPUNetIO PortChannel backend to add cross-IPC-domain expert-parallel dispatch and combine, while preserving the existing IPC path when hybrid transport is not selected.
Changes:
- Adds hybrid latency kernels, pipelined transfers, and automatic multi-HCA selection.
- Exposes BF16
RANK_MAJOR_TOPK_EXPANDEDthrough C++ and Python APIs. - Adds topology and GPU correctness tests, plus build, portability, and communicator-lifetime changes.
| File | Description |
|---|---|
| test/unit/gpunetio_topology_test.cc | Tests HCA affinity selection. |
| test/unit/CMakeLists.txt | Registers topology tests. |
| test/mp_unit/CMakeLists.txt | Configures GPUNetIO test compilation. |
| src/gpunetio/PROVENANCE.md | Updates integration and validation documentation. |
| src/gpunetio/host/gpu_net_io_topology.hpp | Adds HCA affinity scoring. |
| src/ext/ep/moe_runtime.cc | Exposes hybrid runtime buffers. |
| src/ext/ep/latency.cc | Routes latency operations through hybrid runtime. |
| src/ext/ep/include/moe_runtime_context.hpp | Adds hybrid runtime ownership. |
| src/ext/ep/hybrid/topk_expanded.cu | Implements expanded dispatch and combine. |
| src/ext/ep/hybrid/topk_expanded_ipc.cuh | Adds expanded IPC fast paths. |
| src/ext/ep/hybrid/topk_expanded_gpunetio.cuh | Adds expanded network fast paths. |
| src/ext/ep/hybrid/latency.cc | Configures and launches hybrid operations. |
| src/ext/ep/hybrid/include/runtime.hpp | Declares hybrid runtime state. |
| src/ext/ep/hybrid/include/kernels.hpp | Declares hybrid kernel entry points. |
| src/ext/ep/hybrid/include/gpu_net_io.hpp | Adapts GPUNetIO channels for EP. |
| src/ext/ep/hybrid/include/device_context.hpp | Defines transport and pipeline device state. |
| src/ext/ep/hybrid/include/config.hpp | Defines hybrid buffer layouts. |
| src/ext/ep/hybrid/gpu_net_io.cc | Initializes multi-HCA transport resources. |
| src/ext/ep/hybrid/dispatch/rank_major_dispatch.cu | Launches rank-major hybrid dispatch. |
| src/ext/ep/hybrid/common/quantization.cuh | Adds hybrid quantization helpers. |
| src/ext/ep/hybrid/common/latency.cuh | Defines workspace and transport helpers. |
| src/ext/ep/hybrid/common/device_helpers.cuh | Adds device synchronization utilities. |
| src/ext/ep/hybrid/combine/rank_local_reduce_combine.cu | Launches hybrid rank-local combine. |
| src/ext/ep/CMakeLists.txt | Builds hybrid sources and diagnostic options. |
| src/ext/ep/bindings.cpp | Exports the expanded layout enum. |
| src/core/include/communicator.hpp | Removes receive-completion cleanup declaration. |
| src/core/communicator.cc | Changes ordered receive lifetime handling. |
| python/test/test_ep_hybrid.py | Tests hybrid results and graph replay. |
| python/mscclpp/ep/utils.py | Adds transport detection and expanded shapes. |
| python/mscclpp/ep/types.py | Extends configuration validation and layout documentation. |
| python/mscclpp/ep/request_builder.py | Builds expanded requests and routing views. |
| python/mscclpp/ep/communicator.py | Documents expanded combine weighting. |
| python/mscclpp/_core/comm.py | Selects the bootstrap object-broadcast device. |
| pyproject.toml | Enables EP in Python package builds. |
| include/mscclpp/gpu.hpp | Adds HIP driver API aliases. |
| include/mscclpp/gpu_data_types.hpp | Updates vector construction and packed conversions. |
| include/mscclpp/ext/ep/types.hpp | Declares the expanded dispatch layout. |
| CMakeLists.txt | Uses detected GPU architecture lists. |
| cmake/CheckNvidiaGpu.cmake | Processes multiple detected architectures. |
| cmake/check_nvidia_gpu.cu | Reports architectures for all visible GPUs. |
Review effort: Balanced
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
- communicator: restore main's ordered recv futures (#804) and guard lastRecvItems_/connectionInfos_ with a mutex, since deferred futures touch them from whichever thread calls get(). - hybrid combine: add the hidden=2048 dispatch case. - hybrid rank-major: derive the round epoch from a device-side generation advanced by the last combine block, so a CUDA graph that captures a single dispatch/combine round stays correct across replays. - topk-expanded GPUNetIO combine: also push zero-weight valid routes; receivers gather every valid route when weights are pre-applied. - tests: hidden=2048, single-round graph replay with changing routes, and pre-weighted zero-weight expanded routes (fast paths on/off). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 162bd212-6628-4e3a-93b9-190933b019fe
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
…ule_for_ep Bring in the GPUNetIO QP-export failure backport, the ROCm FP8 compile definition, and the GpuNetIoService synchronization/address-type fixes. Resolve PROVENANCE.md by keeping the EP-branch document and adding the backport section and wording from the base branch. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 162bd212-6628-4e3a-93b9-190933b019fe
Replace the end-of-combine block barrier and grid-wide atomic counter with two epoch words: dispatch derives the round epoch from the last retired epoch and publishes it, and combine reads it and retires it. Each kernel only reads one word while block 0 writes the other, so stream order is the only synchronization needed. Single-round CUDA graph replay stays correct, and the ~2 us combine overhead of the previous scheme drops to <1 us. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 162bd212-6628-4e3a-93b9-190933b019fe
- Static FP8 DIRECT_SEND combine: when dequantScale * routerWeight overflows FP32, scale the payload before applying the weight so a finite result stays finite and zero payloads no longer produce NaN. Add a regression with scale 2^-126, weight 4 and 2^-9 / zero payloads. - test_ep_hybrid: validate dispatch results collectively before any rank launches combine, so a single-rank failure fails every rank instead of hanging. - test_ep_hybrid: add cross-domain RANK_MAJOR coverage for batched GPUNetIO dispatch with 2-row chunks plus a partial tail, deferred flush on and off, and single-round graph replays that alternate populated and empty destinations. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 162bd212-6628-4e3a-93b9-190933b019fe
Port the EP GPUNetIO path from the removed GpuNetIoService API to the
IB GpuNetIo endpoint mode introduced on the base branch:
- EpGpuNetIoService::setup now builds per-(peer, queue) PortChannels via
Communicator::connect(EndpointConfig{ib.mode = GpuNetIo}), registered
memory exchange, and ProxyService::portChannel. HCA names are resolved
through getIBDeviceName so MSCCLPP_HCA_DEVICES is honored.
- EpGpuNetIoDeviceContext uses detail::GpuNetIoDeviceHandle (qp, memory
table, atomic scratch) and keeps the raw, unfenced DOCA puts; EP orders
source writes through its own row-fence modes.
- Move gpu_net_io_topology.hpp into src/ext/ep/hybrid/include (the base
branch dropped src/gpunetio); gpunetio_topology_test follows it.
- Drop src/gpunetio/PROVENANCE.md and the redundant mp_unit GPUNetIO
definition (the base propagates MSCCLPP_USE_GPUNETIO publicly).
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 162bd212-6628-4e3a-93b9-190933b019fe
Picks up the internal PortChannel GPUNetIO binding, the atomicAdd fence and the latest main. GpuNetIoDeviceHandle moved out of the detail namespace, so EpGpuNetIoDeviceContext::require now calls GpuNetIoDeviceHandle::require. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 162bd212-6628-4e3a-93b9-190933b019fe
Bring in the GPUNetIO v5.0.0-rc1 pin and the CQ cleanup double-free fix. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ae410543-9b4d-4c07-8fe6-acaaba49d01a
This branch has not been deployed
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.
Summary
Stacked on #910 (
qinghuazhou/gpunetio_submodule). This PR contains only the EP dispatch/combine work that uses the GPUNetIO PortChannel backend added in #910, re-expressed on main's latest EP structure (MoECommunicator/LatencyRuntimeContext).EP GPUNetIO hybrid latency path (
src/ext/ep/hybrid/, namespacemscclpp::ep::hybrid)RANK_MAJOR+RANK_LOCAL_REDUCE, BF16.LatencyRuntimeContextroutes to the hybrid runtime only when GPUNetIO is requested (MSCCLPP_EP_ENABLE_GPUNETIO=1) for a cross-IPC-domain run, or when the newRANK_MAJOR_TOPK_EXPANDEDlayout is used. Main's IPC kernels, storage layout and public API are unchanged otherwise.MSCCLPP_EP_GPUNETIO_DISPATCH_CHUNK_ROWS, default 4 rows for 2-7 nodes, 16 for >= 8 nodes), batched/all-QP dispatch, parallel stage (MSCCLPP_EP_GPUNETIO_PARALLEL_STAGE), per-row fence selection (gpunetio-row-fence: gpu/system), deferred flush (default on).MSCCLPP_EP_GPUNETIO_COMBINE_CHUNK_ROWS) with GPUNetIO puts overlapped with local reduction.src/gpunetio/host/gpu_net_io_topology.hpp,gpunetio_topology_test),MSCCLPP_EP_IPC_DOMAIN_SIZEoverride.New layout
DispatchLayout::RANK_MAJOR_TOPK_EXPANDED(requiresRANK_LOCAL_REDUCE, BF16) with IPC and GPUNetIO fast paths; exported via bindings and the PythonMoECommunicator/ request builder.Tests
python/test/test_ep_hybrid.py: config validation, exact-value dispatch+combine for RANK_MAJOR and TOPK_EXPANDED, CUDA-graph replay.test/unit/gpunetio_topology_test.cc.Validation (GB300, 4 GPUs/node)
python/test/test_ep.py+test_ep_hybrid.py, 8 ranks / one IPC domain: 42 passed.test_ep_hybrid.py2 nodes over GPUNetIO (batched and non-batched): 5 passed.unit_tests54 passed; GPUNetIO CPU tests pass;mp_unit_tests --filter=PortChannelnp=4: 46 passed.