CollectiveX: add UCCL-EP as a third pluggable EP backend#2288
Merged
Conversation
UCCL-EP (github.com/uccl-project/uccl) is a drop-in, DeepEP-API-identical expert-parallel backend over a GPU->CPU-proxy GPUDirect RDMA transport (libibverbs; no NVSHMEM/IBGDA) with cudaIpc scale-up. It runs on both NVIDIA and AMD, slotting in beside deepep-v2 (NVIDIA) and mori (AMD). - adapter bench/ep_uccl.py over the legacy Buffer surface: normal (unweighted-rank-sum combine) and low-latency (weighted-kernel-sum); caller-prequantized blockwise FP8 (e4m3fn / e4m3fnuz per arch) with BF16 combine, matching the existing fp8-dispatch/bf16-combine gate unchanged - registered on h100/h200/b200 (EP8+EP16), mi355x (EP8+EP16), mi300x-tw/mi325x-tw (EP8). LL is NVIDIA-only (AMD LL trips a warp-group assert); AMD SKUs keep mori-only LL - from-source in-image build seam across all three launchers: prepare_backend.sh uccl_prepare (Slurm/enroot + mi-amds) and the launch_mi-tw.sh Docker path (build-once-to-host + PYTHONPATH reuse); common.sh pins the source + emits UCCL_* network env. The AMD build swaps cudaMallocManaged->cudaMallocHost (CDNA lacks usable managed memory) - tests + docs updated; .collx_sources/ and .collx_uccl_pfx/ gitignored Adding a backend extends the matrix under the existing schema (new series / case_ids), so no sweep-version bump.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
…nstraint) The internode combine kernel asserts num_max_rdma_chunked_send_tokens >= num_warps_per_forwarder (ep/src/internode.cu:3091). The adapter's hand-rolled Config(24, 8, 256) left num_max_rdma_chunked_send_tokens at its default (6), which fails for EP16 (EP8 intranode has no RDMA path, so it was unaffected). Mirror UCCL's own test_internode bench for the internode path: size the NVLink+RDMA staging from a generous sizing Config, give each rank num_sms QPs, and drive dispatch/combine with the per-world-size recommended configs (Buffer.get_dispatch_config/get_combine_config), which set rdma-chunked-send to 20/12 for EP16 and satisfy the kernel constraints. Intranode EP8 keeps its validated single-NVLink-buffer recipe unchanged.
…-limited) The internode Config fix lands uccl-ep cross-node RDMA (the CPU proxy connects to all peers and the light EP16 case passes dispatch/combine/correctness), but its CPU-proxy throughput overruns the standardized 900s per-case wall-clock budget (COLLX_RUN_TIMEOUT) on heavy token counts. Rather than weaken the guard for one backend, scope uccl-ep to EP8 in the registry and mark EP16 an honest unsupported coverage row on every supported SKU — mirroring the mori EP16 re-wall. The internode Config fix stays in ep_uccl.py so EP16 is one budget change away if pursued later. deepep-v2 and mori EP16 are untouched; uccl-ep LL stays NVIDIA-only EP8.
…m build)
The wrapper's setup.py has install_requires=["uccl"], which resolves to the PyPI
uccl metapackage -> the prebuilt uccl-cu12 wheel. That wheel exists only for
CUDA, so `python setup.py install` hard-fails on ROCm ("Could not find suitable
distribution for uccl-cu12"); NVIDIA masked the bug by resolving the cu12 wheel.
Our from-source ep build already installs uccl.ep into site-packages/uccl, so the
wrapper needs no PyPI dep. Install it with `pip install --no-deps` in both AMD
build paths (prepare_backend.sh uccl_prepare + launch_mi-tw.sh), making the
source build authoritative on both vendors. Validated on-metal mi325x-tw (gfx942):
ep build + wrapper --no-deps + import deep_ep all clean.
…nup perms)
The -tw docker build wrote the persisted deep_ep/uccl prefix into
$COLLX_DIR/.collx_uccl_pfx, i.e. under the workflow's isolated job root. The
build container runs as root, so those files are root-owned; the workflow's
"Cleanup isolated workspace" step (non-root runner) then fails to `rm -rf` them
("Permission denied"), reddening a leg whose benchmark actually succeeded
(status=success, rc=0). Move the prefix to a node-local /tmp path keyed on the
pinned commit (outside cleanup's scope) and bind-mount it read-write into the
build container and read-only into each case container (PYTHONPATH=/uccl_pfx).
Bonus: a second -tw leg on the same node reuses the build. Empty-array mount
expansion is set -u safe so the mori path is unaffected.
…keying, UCCL GID env, FP8 LL docs - ep_uccl.py: vendor-key the normal-mode SM budget (24 CUDA / 64 HIP) to match UCCL's intranode bench (num_sms = 24 if torch.version.cuda else 64); a flat 24 understated AMD. - launch_mi-tw.sh: key the uccl-ep build cache on commit + image content-id + arch (not commit alone) and publish atomically via a .ready marker written only after the import check, so an interrupted copy or failed import can no longer poison future jobs. - common.sh: UCCL-EP reads only UCCL_* selectors — export UCCL_IB_GID_INDEX (RoCE) and UCCL_IB_HCA explicitly; drop the false NCCL_IB_HCA/NCCL_IB_GID_INDEX fallback claim. - README + ep_uccl.py docstring: FP8 dispatch is caller-prequantized in normal mode; in low-latency the caller sends BF16 and the decode kernel quantizes to e4m3 internally.
…stale gitignore - common.sh: UCCL-EP reads UCCL_IB_HCA and falls back to NCCL_IB_HCA, and its filter honors the leading '=' exact-match and ':port' syntax (ep/src/rdma.cpp). Export UCCL_IB_HCA="=$COLLX_RDMA_DEVICES" to mirror NCCL_IB_HCA rather than a bare name list that prefix-matches (mlx5_1 -> mlx5_1,10..19) and drops the port; fix the comment. Only the GID index has no NCCL fallback in UCCL's EP path (it reads only UCCL_IB_GID_INDEX, ep/include/rdma_util.hpp), so that stays an explicit UCCL_* export. - README.md / docs/methodology.md: FP8 dispatch is caller-prequantized only in normal mode; the low-latency DeepEP and UCCL-EP kernels quantize internally from BF16 (MoRI LL stays caller- prequantized). Drop the blanket "every case uses layout-and-dispatch-v1 semantics" — low-latency cases use each backend's decode-kernel semantics. - .gitignore: drop the stale .collx_uccl_pfx/ entry; the -tw build persists to node-local /tmp.
…llocation build cache - ep_uccl.py: the low-latency capability guard used `and`, so it raised only when BOTH LL entry points were absent; De Morgan makes the intended check `or` — fail closed at buffer construction whenever EITHER low_latency_dispatch or get_low_latency_rdma_size_hint is missing, instead of a raw AttributeError deeper in the timed path. Name both in the message. - prepare_backend.sh: cache the from-source UCCL build the way deepep_prepare does — build once per (cpu, arch, image, commit) under $COLLX_BACKEND_CACHE_ROOT behind an flock + .ready marker, persist the built deep_ep/uccl packages and put them on PYTHONPATH (write_rank_env carries it to the ranks), reused on every later allocation. Node-local fallback when no shared cache is mounted. Ends the full recompile on every uccl-ep shard on single-slurm/mi-amds. - launch_mi-amds.sh: mount the /cx-cache backend cache for uccl-ep (mirrors single-slurm); the mori path is left untouched.
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
1 participant
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.
What
Adds UCCL-EP (github.com/uccl-project/uccl) as a third pluggable expert-parallel backend in CollectiveX, alongside
deepep-v2(NVIDIA) andmori(AMD).UCCL-EP is a drop-in, DeepEP-API-identical backend over a GPU→CPU-proxy GPUDirect RDMA transport (plain
libibverbs; no NVSHMEM/IBGDA), withcudaIpcscale-up. It runs on both NVIDIA and AMD, exercising a completely different transport than the two existing backends on the same matrix.Changes
bench/ep_uccl.pyover the legacyBuffersurface — normal (unweighted-rank-sum combine) and low-latency (weighted-kernel-sum). Caller-prequantized blockwise FP8 (e4m3fn/e4m3fnuzper arch) with BF16 combine, so the existing fp8-dispatch/bf16-combine correctness gate (COMBINE_REL_TOL) is preserved unchanged.configs/platform_config.json): uccl-ep EP8 on h100/h200/b200 + mi355x + mi300x-tw/mi325x-tw. LL is NVIDIA-only EP8 (AMD LL trips a warp-group assertion). b300/gb200/gb300 and EP16 emit honestunsupportedrows.prepare_backend.shuccl_prepare(Slurm/enroot + mi-amds) and thelaunch_mi-tw.shDocker path (build-once-to-host +PYTHONPATHreuse).common.shpins the source and emitsUCCL_*network env. The AMD build swapscudaMallocManaged→cudaMallocHost(CDNA lacks usable managed memory).uccl-epadded as a dispatchableCollectiveX Sweepbackend choice..collx_sources/and.collx_uccl_pfx/gitignored.Adding a backend extends the matrix under the existing schema (new series / case_ids) — no sweep-version bump.
Validation (on-metal + CI)
EP16 (cross-node): functional but wall-clock-limited → scoped to EP8
The internode combine kernel asserts
num_max_rdma_chunked_send_tokens >= num_warps_per_forwarder. A follow-up commit fixes the adapter's Config (mirror UCCL's owntest_internodebench: size buffers from a generous config,num_qps_per_rank = num_sms, drive dispatch/combine with the recommended per-world-size configs). With that fix the CPU proxy connects cross-node to all 16 peers and the light EP16 case passes correctness — but its CPU-proxy throughput overruns the standardized 900s per-case wall-clock budget on heavy token counts. Rather than weaken the guard for one backend, uccl-ep is scoped to EP8 and EP16 is an honestunsupportedcoverage row (mirroring the mori EP16 re-wall); the internode Config fix stays in the adapter so EP16 is one budget change away if pursued.