Resolve the C++ benchmark conflict by combining NCCL-EP-style random top-k routing and masked selections with the current BF16/FP8 MoERuntime API and format-aware byte accounting.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: efbacae6-f679-430b-bc16-b45ae162fc76
Refresh the Python and C++ benchmark paths for BF16 and FP8 dispatch, current MoERuntime signatures, active kernel sources, portable CUPTI discovery, realistic routing, and safe unified reporting. Remove the merged change to the inactive legacy implementation.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: efbacae6-f679-430b-bc16-b45ae162fc76
## Summary
Replaces the deterministic `(t*K+j) % E` expert assignment in
`mscclpp_ep_bench.cu`
with the same randomized top-k routing scheme used by the NCCL-EP
benchmark, so the
C++ LL benchmark exercises a realistic, non-uniform expert distribution.
## Changes
- Per-token scores drawn from `|N(0,1)| + 1` with a rank-seeded
`std::mt19937(1+rank)`.
- `std::partial_sort` selects the top-K experts per token (descending).
- Randomly masks 10 (token, slot) positions to `-1` to exercise the LL
kernels'
invalid-index handling (guarded at `low_latency.cu:278` and `:726`).
- Byte accounting now counts only valid (`>= 0`) selections.
- Added `<random>` / `<utility>` includes; clang-format-14 clean.
---------
Co-authored-by: Binyang Li <binyli@microsoft.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
## Summary
Adds a unified low-latency (LL) expert-parallel benchmark for comparing
mscclpp EP against NVIDIA NCCL-EP on equal footing, and fixes an LL
combine
performance regression in the feature/ep kernels.
## Commits
- **ep/ll: scale combine grid to numCombinedTokens (fix combine
regression).**
The LL combine host launched a fixed grid of
`ceil(numExperts/kNumWarpGroups)`
(= 43 SMs for 128 experts), but `combineRecv`'s per-token weighted
reduction
strides `tokenIdx` by the grid size, so it only used 43 blocks for 128
tokens.
Scale the grid to `numCombinedTokens` (capped by the device SM count).
Extra
blocks are recv-only (send is guarded by `responsibleExpertIdx <
numExperts`)
and every block still hits `cg::this_grid().sync()`. Measured (e128 t128
d7168
k8, 1-node): combine avg 56 → 43 us (-24%), min 50 → 38 us. Bit-exact
(`test_low_latency_multirank.py`).
- **ep(bench): unified driver + three backends** — mscclpp (Python
MoECommunicator), mscclpp-cpp (pure C++ MoERuntime), nccl-ep
(`ep_bench`).
- **ep(bench): pure-C++ LL benchmark** (`mscclpp_ep_bench.cu`) calling
`MoERuntime::dispatch/combine` directly, with CUPTI kernel timing, built
via CMake.
## Build
No impact on the core library build: the benchmark's
`test/python/ep/CMakeLists.txt` is standalone (no `add_subdirectory`
from any
parent CMake) and this PR does not touch the top-level `CMakeLists.txt`
/
`pyproject.toml` / `setup.py`. The only library change is the combine
grid size.
The Python driver (`run_ep_bench.py`) and the mscclpp Python backend
(`ep_bench_ll.py`) need **no build**. Only the **mscclpp-cpp** backend
needs a
one-time standalone build (it recompiles the two LL translation units +
links the installed `libmscclpp.so`):
```bash
cmake -S test/python/ep -B build \
-DMSCCLPP_EP_NUM_MAX_NVL_PEERS=4 -DCMAKE_CUDA_ARCHITECTURES=100 # GB200 sm_100
cmake --build build -j
# -> build/mscclpp_ep_bench
```
Requires nvcc/CUDA, MPI, and CUPTI (all `find_package REQUIRED`).
## Usage
Common workload: `-e 128 -t 128 -d 7168 -k 8 -w 10 -i 100`. The driver
prints a
unified summary with **host-observed** and **kernel-only** (CUPTI) rows.
```bash
# mscclpp (Python MoECommunicator) — no build needed
python test/python/ep/run_ep_bench.py --ep-lib mscclpp -e 128 --cupti-inproc
# mscclpp-cpp (pure C++ MoERuntime) — after the Build step above
python test/python/ep/run_ep_bench.py --ep-lib mscclpp-cpp -e 128 -t 128 -d 7168 -k 8 -w 10 -i 100
# 2 / 4-node (one NVLink domain): list peer IPs; the driver builds the hostfile + mpirun
python test/python/ep/run_ep_bench.py --ep-lib mscclpp-cpp -e 128 -t 128 -d 7168 -k 8 -w 10 -i 100 \
--nodes "10.0.0.1 10.0.0.2 10.0.0.3 10.0.0.4"
# nccl-ep (NVIDIA reference ep_bench)
python test/python/ep/run_ep_bench.py --ep-lib nccl-ep -e 128 -t 128 -d 7168 -k 8 -w 10 -i 100 \
--nccl-lib-path /path/to/nccl/build/lib
# all backends side-by-side
python test/python/ep/run_ep_bench.py --ep-lib all -e 128 -t 128 -d 7168 -k 8 -w 10 -i 100 \
--nccl-lib-path /path/to/nccl/build/lib
# add --kernel-only for just the CUPTI rows, --dry-run to print the commands
```
Useful flags: `--nodes "<ip ...>"`, `--nproc-per-node 4`,
`--kernel-only`,
`--dry-run`, `--cupti-inproc`, `--mscclpp-cpp-bench <path>`,
`--nccl-ep-bench <path>`,
`--nccl-lib-path <dir>`, `--hpcx <dir>`, `--iface <nic>`.
## Validation
Combine fix bit-exact and benchmarked 1/2/4-node on GB200 NVL72
(sm_100);
dispatch/combine host + kernel times on par with NCCL-EP at all scales.
---------
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
This pull request makes significant improvements to the MoE (Mixture of
Experts) Python API and documentation, focusing on clarifying and
expanding the Expert Parallel (EP) interface, especially around
quantization, dispatch/combine handles, and overlap configuration. The
changes introduce new data structures, update function signatures, and
improve documentation to better reflect the current and planned
capabilities of the system. Additionally, the base development container
is updated to CUDA 13.0, and minor corrections are made to extension
naming.
## Summary
Enables the **high-throughput (HT)** MoE dispatch/combine backend as a
first-class
`MoEMode.HIGH_THROUGHPUT` path on the canonical nanobind
`MoECommunicator` introduced in #818.
After #818 the HT (DeepEP-style) backend was archived under
`src/ext/ep/ht/` and not compiled
(`mode=HIGH_THROUGHPUT` → `NotImplementedError`). This branch
un-archives it, re-binds it under the
nanobind module via a DLPack tensor bridge, adds the GB200 TMA kernel
optimizations, and wires it
through the high-level `MoECommunicator`. The low-latency `MoERuntime`
path is unchanged.
Already merged up-to-date with `feature/ep` (conflict-free): **23
commits, 14 files, +2682 / −690**.
## What's new
- **HT compiles + binds under nanobind.** CMake links Torch and builds
`ht/buffer.cc` +
`ht/kernels/*.cu`; `bindings.cpp` binds `Config`, `EventHandle`, and the
HT runtime (C++
`mscclpp::ep::Buffer`) exposed to Python as **`ExpertParallelRuntime`**
(with a `Buffer` alias).
- **nanobind ⇄ libtorch DLPack bridge.** HT methods return
data-dependent-size `torch::Tensor`
tuples, so tensors cross the boundary as DLPack capsules (`at::toDLPack`
/ `at::fromDLPack`;
`torch.utils.dlpack` on the Python side) — zero-copy, no torch
type-caster needed in nanobind.
- **`MoECommunicator(mode=MoEMode.HIGH_THROUGHPUT)`** end-to-end:
`_init_ht` builds the HT runtime
on a torch `ProcessGroup`; `dispatch`/`combine` route by mode; FLAT
layout; `previous_handle`
cached-dispatch; intranode vs internode auto-selected from the RDMA size
hint. LL continues to use
`MoERuntime` (`comm=CommGroup`).
- **GB200 TMA kernels** (intranode + internode): TMA direct-gather
combine + all-sender dispatch,
per-phase SM knobs, channel-adaptive warp counts.
- **Tests** (`test_intranode_multirank.py`,
`test_internode_multirank.py`): correctness drives the
low-level `ExpertParallelRuntime`; the benchmark drives the public
`MoECommunicator`.
## Validation
1 node × 4 GB200 (e256 / topk8 / hidden=7168 / tokens=4096):
- Correctness: **PASS**, `max|got − expected| = 0.0`.
- `MoECommunicator` HT bench: **dispatch 454.2 µs / combine 476.0 µs**
(matches the raw-runtime
best of 453.8 / 474). The high-level wrapper adds no measurable
overhead.
## Notes / limitations
- HT dispatch is **BF16-only** for now; HT layout is **FLAT** only.
- The internode path compiles and the test is rewritten to
`MoECommunicator`, but it has **only
been validated single-node** so far (one NVL72 rack reachable); a 2-node
run is pending.
- `dispatch_async` / `combine_async` and block-level overlap are not
implemented yet.
## Implementation notes (nanobind)
Re-binding a pybind11 tensor API into nanobind required: object args as
`nb::object` (not
`nb::handle`); `.none()` on every `nb::arg` that may receive `None`
(nanobind rejects `None` for
object args by default); and renaming the reserved Python keyword
`async` → `async_finish`.
internode_ncclep.cuh: add EP_NCCLEP_DRAIN_NOOP compile gate (default 0, inert) - when 1 the NVL receiver keeps all control flow but skips the data copies, to measure the dispatch-time upper bound of eliminating the cross-GPU receiver drain. Probe result (4-node {38,41,59,75}): dispatch inc2 1124us -> DRAIN_NOOP 1048us (-6.8%), agg_bw 836->896 GB/s => confirms real headroom for the cross-GPU peer-map direct-write rework (ceiling ~-16.6% cumulative vs baseline).
test_internode_multirank.py: gate the dispatch range-assert and combine assert behind MSCCLPP_EP_SKIP_VERIFY env so dispatch timing can be measured when recv_x is intentionally incomplete (perf probing).
- src/core/atomicadd_kernel.cu: restore the legacy 3-arg
cuCtxCreate(&proxyAtomicCtx_, 0, cuDevice) in the
'#else' branch of the CUDA_VERSION >= 12050 guard. A prior
edit had corrupted it to 'cuCtxCreate(&proxyAtomicCtx_vice)',
which broke the CUDA 11.8 build (CodeQL CUDA cuda11.8 and
MSCCLPPLang cuda11.8 jobs).
- Apply clang-format to src/ext/ep/* (no logic changes,
fixes the cpplint CI job).
- Apply black to test/python/ext/ep/test_internode_multirank.py
and test_intranode_multirank.py (no logic changes, fixes
the pylint CI job).
The NVLS HT B2 path introduced in 3ab2e43b activated whenever
isNvlsSupported() && num_rdma_ranks > 1. On H100 NDv5 / Azure CX-7 RoCE
that is true (H100 has intra-node NVLink multicast), but there is no
cross-host NVSwitch fabric. mscclpp's GpuIpcMem::create then falls back
to CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR whose handle exchange routes
through /tmp/mscclpp_bootstrap_<pid>.sock -- a master-rank-0 unix-domain
socket worker ranks cannot reach. Symptom on every commit since 3ab2e43b:
RuntimeError: connect() failed for unix socket to
/tmp/mscclpp_bootstrap_<pid>.sock
MSCCLPP_EP_FABRIC_IPC=0 was being silently ignored.
src/ext/ep/buffer.cc: add resolve_fabric_ipc_supported() helper.
Resolution:
1. MSCCLPP_EP_FABRIC_IPC env var (0/off/false/no => off,
1/on/true/yes/force => on, otherwise auto).
2. Auto-detect: requires both
- CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED == 1
- device compute capability >= sm_100 (Blackwell+).
Gate both use_fabric_ipc_alloc (RDMA buffer allocator) and nvls_ht_enabled
(HT B2 multicast region) on fabric_ipc_supported. On H100 both fall back
to cudaMalloc + legacy PortChannel; on GB200 NVL72 both remain enabled.
Diagnostic prints now show fabric_ipc=.
test/python/ext/ep/test_internode_multirank.py: replace hardcoded
NUM_MAX_NVL_PEERS=4 with a runtime _detect_local_world_size() helper
that reads MSCCLPP_EP_LOCAL_WORLD_SIZE / LOCAL_WORLD_SIZE /
OMPI_COMM_WORLD_LOCAL_SIZE, falling back to torch.cuda.device_count().
Makes the test correct on both H100 (8 GPUs/node) and GB200 (4 GPUs/node)
without code changes.
src/core/atomicadd_kernel.cu: use cuCtxCreate_v4 for CUDA >= 12.5 (the
underlying symbol was renamed); preserve legacy 3-arg cuCtxCreate for
older toolkits.
Verified on 2x H100 NDv5 at HEAD:
LL intranode (8 GPUs) PASS
LL internode (16 GPUs, 2 nodes) PASS
HT intranode (8 GPUs) PASS
HT internode (16 GPUs, 2 nodes) PASS
Diagnostic on H100:
[mscclpp_ep] rdma_buffer allocator: cudaMalloc (low_latency=0, nvls=1, fabric_ipc=0)
[mscclpp_ep] NVLS HT multicast: disabled (low_latency=0, num_rdma_ranks=2, nvls_supported=1, fabric_ipc=0)
Run one uncached dispatch to capture (rank_prefix_matrix,
channel_prefix_matrix, num_recv_tokens), then time iters in cached mode.
This replaces notify_dispatch + host busy-wait on mapped pinned counters
with the cheap cached_notify_dispatch (one barrier + memcpy + memset),
matching NCCL-EP ep_bench convention.
Cached mode forces num_experts=0 (buffer.cc:807), so topk_idx must be
None in iters; combine still works because recv_topk_weights is optional.
Per-iter dispatch latency drops ~21% (4247→3373µs). Confirms host-side
notify_dispatch overhead is only ~20% of total dispatch time; the
remaining 14.4× send-total asymmetry vs combine is intrinsic (3× recv/
send byte fan-out × 3.8× dispatch-kernel-vs-combine-kernel work).
At 16 nodes (64 ranks) with topk=8, expected combine values reach
rank*8 = 504, while intermediate partial sums (rank*7 etc.) cross the
bf16 ulp=2 boundary at 256. With the test pattern x = rank*ones and
weights = 1, this produces deterministic +/-1 round-off on certain
ranks (odd local_rank on nodes >= 9), tripping the previous 1e-2
absolute tolerance even though the kernel is correct.
Use tol = max(1e-2, max_exp / 64) which matches the bf16 mantissa
precision and scales with the magnitude of the expected combined
output. The previous tight bound is preserved for small-scale runs
where max_exp < 0.64.
- Internode HT test: accept MSCCLPP_EP_HT_{TOKENS,HIDDEN,TOPK,EXPERTS}
env vars to override the functional-check problem size (was hardcoded
to num_tokens=128, hidden=1024, num_topk=min(4,num_ranks),
num_experts=num_ranks*4).
- Both intranode + internode HT tests: replace dist.all_to_all_single
bookkeeping (per-(src,dst) recv-count matrix used for the
six-metric NVL/RDMA BW breakdown) with dist.all_gather_into_tensor
+ transpose. Functionally identical (gathered[:, rank] gives the
same recv-from-src column) but works on socket-NCCL with
NCCL_IB_DISABLE=1, which is required on rigs where NCCL IB cannot
coexist with mscclpp RDMA. Sends num_ranks^2 int64 instead of
num_ranks per rank — negligible (64 ints at 8 ranks).
Allow tuning the internode HT test cfg from the environment without
editing the source. Supported variables (all optional):
MSCCLPP_EP_NSM (default 152) num channels / SMs
MSCCLPP_EP_NVL_SEND (default 8)
MSCCLPP_EP_NVL_RECV (default 256)
MSCCLPP_EP_RDMA_SEND (default 16)
MSCCLPP_EP_RDMA_RECV (default 128)
The defaults match what we use for 16-node GB200 bench runs (e.g.
NVL_RECV=512 to satisfy the HT combine assert at 16 nodes).
Run `tools/lint.sh cpp` (clang-format 14) and `tools/lint.sh py`
(black) over the EP extension files added by this PR. No functional
changes; pure reformatting to satisfy the cpplint and pylint CI jobs.
Add optional out_packed_recv_x / out_src_info / out_layout_range /
out_count parameters to Buffer::low_latency_dispatch so callers can
hoist the four recv-side allocations out of a hot loop, mirroring the
existing out= path on low_latency_combine.
The bench in test_low_latency_multirank.py preallocates these tensors
once and passes them on every iter so the timed loop reflects kernel
cost, not torch.empty + caching-allocator overhead.
Previously total_send_tokens was Sigma over dst_rank of num_tokens_per_rank
which over-counts intra-node fan-out. NCCL-EP's ep_bench collapses
multiple destinations on the same node into one count; on a single-node
run that means total_send_tokens = number of tokens with at least one
valid expert. Switching to is_token_in_rank.any(dim=1).sum() makes the
send-side BW comparable to NCCL-EP's send: total_bw / nvl_bw line.
Set TORCH_NCCL_ENABLE_MONITORING=0 before importing torch.distributed.
The barrier+destroy_process_group finally block (afbdcd6a) suffices
under torchrun, but under mpirun rank 0 (the TCPStore server) can exit
before non-zero ranks finish teardown, and the background heartbeat
thread polls the store and logs 'recvValue failed / Connection was
likely closed'. Disabling the monitor outright is safe for short-lived
bench runs.
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.
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.
Same change as the intra-node bench (commit 4ed6f229), 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.
- 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.
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.
- 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).
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.
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.
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.
- 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.
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.
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.
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.
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
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.
