From 6a65edd7eb309ab9db7e36f6d5ecad12587679eb Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Mon, 20 Jul 2026 17:45:12 +0800 Subject: [PATCH 01/10] feat(collectivex): add UCCL-EP as a third pluggable EP backend 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. --- experimental/CollectiveX/.gitignore | 2 + experimental/CollectiveX/README.md | 10 +- experimental/CollectiveX/bench/ep_uccl.py | 418 ++++++++++++++++++ experimental/CollectiveX/bench/run_ep.py | 11 +- .../CollectiveX/configs/platform_config.json | 18 +- experimental/CollectiveX/docs/methodology.md | 16 +- .../CollectiveX/launchers/launch_mi-amds.sh | 7 +- .../CollectiveX/launchers/launch_mi-tw.sh | 83 +++- .../launchers/launch_single-slurm.sh | 15 +- experimental/CollectiveX/runtime/common.sh | 54 +++ .../CollectiveX/runtime/prepare_backend.sh | 67 ++- experimental/CollectiveX/sweep_matrix.py | 1 + experimental/CollectiveX/tests/test_matrix.py | 42 +- .../CollectiveX/tests/test_runtime.py | 17 +- 14 files changed, 719 insertions(+), 42 deletions(-) create mode 100644 experimental/CollectiveX/bench/ep_uccl.py diff --git a/experimental/CollectiveX/.gitignore b/experimental/CollectiveX/.gitignore index f34ac59bad..f68f23df47 100644 --- a/experimental/CollectiveX/.gitignore +++ b/experimental/CollectiveX/.gitignore @@ -5,6 +5,8 @@ unsupported/ .shards/ .collx_workloads/ .collx_backend/ +.collx_sources/ +.collx_uccl_pfx/ /matrix_full.json gpucore.* diff --git a/experimental/CollectiveX/README.md b/experimental/CollectiveX/README.md index 27afb8fe38..ff3ea1ee8d 100644 --- a/experimental/CollectiveX/README.md +++ b/experimental/CollectiveX/README.md @@ -20,11 +20,14 @@ uniform routing only. Cases run in one of two modes: unweighted rank-sum combine. It runs the full decode and prefill ladders. - `low-latency` uses each backend's decode-optimized kernel family: on DeepEP the legacy `deep_ep.Buffer` IBGDA `low_latency_dispatch`/`low_latency_combine` (a per-expert padded receive - and a source-side gate-weighted combine); on MoRI the `IntraNodeLL` kernel (single-call, + and a source-side gate-weighted combine); on UCCL-EP the same legacy `Buffer` low-latency kernels + over its CPU-proxy transport; on MoRI the `IntraNodeLL` kernel (single-call, pure-intranode, same compact layout and unweighted rank-sum combine as `IntraNode`). It is a decode-phase-only, per-SKU-capability-gated addition whose runnable set differs from `normal`'s, so - it is enabled from each SKU's `ll_backends` registry entry (currently DeepEP V2 EP8 on H100/H200/B200 - and MoRI EP8 on MI300X/MI325X/MI355X). Scoped single-node EP8 runs over the intra-node NVLink/XGMI + it is enabled from each SKU's `ll_backends` registry entry (currently DeepEP V2 EP8 on H100/H200/B200, + MoRI EP8 on MI300X/MI325X/MI355X, and UCCL-EP EP8 on H100/H200/B200 only — UCCL's low-latency kernel + trips a warp-group assertion on AMD's CU count, so the AMD SKUs keep UCCL-EP normal mode without LL). + Scoped single-node EP8 runs over the intra-node NVLink/XGMI low-latency path (no `/dev/gdrdrv` needed — validated on H200 with it absent); NVSHMEM/IBGDA on the wire is only a multi-node scale-out (EP16) concern. @@ -63,6 +66,7 @@ scale-up domain. |---|---| | DeepEP V2 | `normal` mode is PR #605 `ElasticBuffer` plus exact upstream #630 and #640 fixes: LSA for scale-up and GIN for x86 EP16 scale-out. FP8 dispatch via `use_fp8_dispatch` (blockwise e4m3fn) alongside BF16. `low-latency` mode is the legacy `deep_ep.Buffer` IBGDA decode kernels (per-expert padded layout, weighted combine, `use_fp8` e4m3fn), decode/EP8 only | | MoRI | `normal` mode uses the direct `IntraNode` kernel for scale-up EP8 on every CDNA SKU and pins `InterNodeV1` for EP16 over 2x8 XGMI + RDMA. `low-latency` mode selects the `IntraNodeLL` decode kernel (single-call, pure-intranode, same compact layout and unweighted combine as `IntraNode`), decode/EP8 only. FP8 dispatch is caller-prequantized (per-SKU e4m3fnuz on gfx942, e4m3fn on gfx950); combine stays BF16 (`quant_type=none`) alongside BF16 dispatch | +| UCCL-EP | [UCCL](https://github.com/uccl-project/uccl) EP: a drop-in, API-identical DeepEP replacement whose CPU proxies issue GPUDirect RDMA over plain `libibverbs` (no NVSHMEM/IBGDA), with software message ordering, atomics, and flow control; scale-up is single-node `cudaIpc` over NVLink/XGMI (never MNNVL). `normal` mode is the legacy `Buffer` `dispatch`/`combine` (unweighted rank-sum); `low-latency` reuses the legacy `low_latency_dispatch`/`low_latency_combine` decode kernels (weighted combine), decode/EP8 only. FP8 dispatch is caller-prequantized (blockwise e4m3fn, per-SKU e4m3fnuz on gfx942) with BF16 combine. Runs on NVIDIA and AMD (H100/H200/B200 + MI300X/MI325X/MI355X); EP16 on h100/h200/b200/mi355x is a wall probe against the deepep-v2/mori scale-out limits | DeepEP V2 means the `ElasticBuffer` implementation introduced by [DeepEP PR #605](https://github.com/deepseek-ai/DeepEP/pull/605), not a newer legacy `Buffer` build. diff --git a/experimental/CollectiveX/bench/ep_uccl.py b/experimental/CollectiveX/bench/ep_uccl.py new file mode 100644 index 0000000000..170158a5d9 --- /dev/null +++ b/experimental/CollectiveX/bench/ep_uccl.py @@ -0,0 +1,418 @@ +#!/usr/bin/env python3 +"""UCCL-EP adapter: the drop-in DeepEP-legacy `Buffer` API over UCCL's CPU-proxy transport. + +UCCL-EP (https://github.com/uccl-project/uccl) is an API-identical DeepEP replacement whose +CPU proxies issue GPUDirect RDMA over plain libibverbs (no NVSHMEM/IBGDA); scale-up is +single-node cudaIpc over NVLink/XGMI (never MNNVL). Its Python `Buffer` — installed via UCCL's +`deep_ep_wrapper`, so `import deep_ep` in the isolated UCCL venv resolves to UCCL, NOT DeepSeek's +DeepEP — manages the CPU proxy threads internally (spun up in `__init__`'s `initialize_uccl`, +torn down in `destroy()`'s `destroy_uccl`), so this adapter never calls those functions directly; +it just constructs the Buffer and calls `.destroy()`. + +Both modes use the legacy `Buffer` surface (mirroring bench/ep_deepep_v2.py's legacy-Buffer LL +path, which is ~1:1 reusable here): + normal -> get_dispatch_layout + dispatch + combine; per-token multi-expert recv layout + (expanded per (token, local-expert) for the oracle, reduced back before combine); + activation-only unweighted rank-sum combine. + low-latency -> low_latency_dispatch/low_latency_combine; per-expert padded recv, source-side + weighted-kernel-sum combine. + +FP8 dispatch is caller-prequantized (blockwise e4m3fn, e4m3fnuz on gfx942), combine is BF16 — the +oracle applies the identical per-token cast round-trip via semantic_payload/oracle_x, so the tight +combine gate (COMBINE_REL_TOL = 8*2^-8) is preserved, not loosened. +""" +from __future__ import annotations + +import sys +import types + +import torch +import torch.distributed as dist + +from ep_backend import EPBackend + +try: + # In the isolated UCCL venv `deep_ep` is UCCL's deep_ep_wrapper (a drop-in DeepEP API backed + # by uccl.ep's CPU-proxy runtime), not DeepSeek's DeepEP. Buffer/Config come from it. + import deep_ep # noqa: F401 (UCCL deep_ep_wrapper) + from deep_ep import Buffer, Config # type: ignore +except Exception as exc: # pragma: no cover - requires the benchmark image + print(f"ERROR: UCCL-EP import failed: {exc!r}", file=sys.stderr) + raise + + +# ---- Vendored UCCL FP8 helpers (ep/bench/utils.py) -------------------------------------- +# These live in UCCL's bench dir, not the installed package, so they are vendored VERBATIM so +# the quantization the oracle models is byte-identical to what dispatch transports. Keep in +# lockstep with upstream ep/bench/utils.py if the pinned UCCL commit moves. + +def _fp8_e4m3_dtype() -> "torch.dtype": + """UCCL's arch-keyed FP8 E4M3 dtype: e4m3fnuz on gfx942 (MI300X/MI325X), e4m3fn elsewhere.""" + if hasattr(torch.version, "hip") and torch.version.hip is not None: + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + arch = getattr(props, "gcnArchName", "") + if arch.startswith("gfx942"): + return torch.float8_e4m3fnuz + return torch.float8_e4m3fn + + +def per_token_cast_to_fp8(x: "torch.Tensor"): + """Blockwise (per-128-channel) FP8 quantization: returns (e4m3 [m, n], scales [m, n//128]).""" + assert x.dim() == 2 and x.size(1) % 128 == 0 + m, n = x.shape + fp8_dtype = _fp8_e4m3_dtype() + fp8_max = 240.0 if fp8_dtype == torch.float8_e4m3fnuz else 448.0 + x_view = x.view(m, -1, 128) + x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4) + return (x_view * (fp8_max / x_amax.unsqueeze(2))).to(fp8_dtype).view(m, n), ( + x_amax / fp8_max + ).view(m, -1) + + +def per_token_cast_back(x_fp8: "torch.Tensor", x_scales: "torch.Tensor"): + """Blockwise FP8 -> BF16 dequant mirroring per_token_cast_to_fp8.""" + if x_scales.dtype == torch.int: + x_scales = x_scales.view(dtype=torch.uint8).to(torch.int) << 23 + x_scales = x_scales.view(dtype=torch.float) + x_fp32 = x_fp8.to(torch.float32).view(x_fp8.size(0), -1, 128) + x_scales = x_scales.view(x_fp8.size(0), -1, 1) + return (x_fp32 * x_scales).view(x_fp8.shape).to(torch.bfloat16) + + +@torch.compile(dynamic=False) +def _ll_dequant_static(fp8, scales): + """Static-shape FP32-accumulate dequant of the padded low-latency FP8 recv to BF16. + + Mirror of ep_deepep_v2._ll_dequant_static: the low-latency padded recv shape + ``[num_local_experts, cap*num_ranks, hidden]`` is constant per dispatch, so a static + (dynamic=False) compile fuses to one FP32 pass and stays inside the wall-clock budget that + deep_ep's dynamic-shape per_token_cast_back would overrun. Padding slots decode to NaN + (FP8 padding bytes) — harmless because combine is handle-indexed and never reads padding. + """ + e, s, h = fp8.shape + values = fp8.to(torch.float32).view(e, s, h // 128, 128) + block_scales = scales.to(torch.float32).view(e, s, h // 128, 1) + return (values * block_scales).to(torch.bfloat16).view(e, s, h) + + +# Normal-mode legacy Config launch parameters (DeepEP-legacy Config(num_sms, chunk, nvl_buffer)). +# Bring-up-tunable; these mirror UCCL's own intranode bench (nvl_buffer_size=256) and the DeepEP +# normal default SM budget. num_nvl_bytes is a generous fixed reservation as in UCCL's bench. +_NORMAL_NUM_SMS = 24 +_NORMAL_NVL_BUFFER_SIZE = 256 +_NORMAL_NVL_BYTES = int(2e9) + + +class UCCLEPBackend(EPBackend): + name = "uccl-ep" + # One legacy Buffer under two modes, selected by args.mode: + # normal -> get_dispatch_layout/dispatch/combine; unweighted rank-sum combine. + # low-latency -> low_latency_dispatch/combine decode kernels; source-side weighted combine. + kernel_generation = "uccl-legacy-buffer" + SUPPORTED_MODES = ("normal", "low-latency") + SUPPORTED_PRECISIONS = ("bf16", "fp8") + stage_device_work = False + combine_needs_redispatch = False + dispatch_needs_combine_cleanup = False + combine_weight_semantics = "unweighted-rank-sum" + + def __init__(self, args, rank, world_size, local_rank, device): + super().__init__(args, rank, world_size, local_rank, device) + self.group = dist.group.WORLD + self.experts_per_rank = args.experts // world_size + self._internode = world_size > int(args.scale_up_domain) + self._fp8 = self.precision == "fp8" + # FP8 dispatch dequantizes the received (e4m3, per-128-block scale) payload back to the + # BF16 combine sends — real device work, hence a separately-timed stage component. + self.stage_device_work = self._fp8 + self._fp8_dtype = None + if self._fp8: + self._fp8_dtype = _fp8_e4m3_dtype() + self.dispatch_dtype = ( + "fp8-e4m3fnuz" if self._fp8_dtype == torch.float8_e4m3fnuz else "fp8-e4m3fn" + ) + self.dispatch_value_bytes = 1 + self.dispatch_scale_bytes_per_copy = ((args.hidden + 127) // 128) * 4 + if self.mode == "low-latency": + # Legacy low-latency decode path: a distinct kernel family whose combine multiplies + # by the gate at the source (weighted), not an unweighted rank sum. LL result tensors + # are double-buffered and single-use per dispatch, so every timed combine needs a + # fresh dispatch and every timed dispatch must be drained by its combine. + self.kernel_generation = "uccl-legacy-buffer-ll" + self.combine_weight_semantics = "weighted-kernel-sum" + self.combine_needs_redispatch = True + self.dispatch_needs_combine_cleanup = True + + def buffer_cap(self, args): + if self.mode == "low-latency": + # LL pre-allocates a fixed [num_local_experts, cap*num_ranks, hidden] receive buffer, + # so cap is a hard per-rank dispatch-slot bound (same as ep_deepep_v2's legacy LL). + return 256 + return None + + # ---- buffer construction --------------------------------------------------------------- + + def create_buffer(self, spec): + self.max_tokens = spec.max_tokens_per_rank + if self.mode == "low-latency": + self._create_ll_buffer(spec) + return + args, world_size = self.args, self.world_size + # Normal mode: legacy Buffer with an NVLink scale-up buffer plus (internode) an RDMA + # buffer. UCCL's own intranode bench sizes num_nvl_bytes at a fixed ~2 GB; internode adds + # an RDMA buffer sized from the Config hint. Reservations are generous and backed on + # demand; num_sms/nvl_buffer are bring-up-tunable (see _NORMAL_* above). + self.config = Config(_NORMAL_NUM_SMS, 8, _NORMAL_NVL_BUFFER_SIZE) + num_rdma_bytes = 0 + if self._internode: + # BF16 combine wire is the widest per-value direction; size the RDMA staging for it. + hidden_bytes = args.hidden * 2 + num_rdma_bytes = int( + self.config.get_rdma_buffer_size_hint(hidden_bytes, world_size) + ) + self.buffer = Buffer( + self.group, + _NORMAL_NVL_BYTES, + num_rdma_bytes, + low_latency_mode=False, + num_qps_per_rank=1, + allow_nvlink_for_low_latency_mode=True, + allow_mnnvl=False, + explicitly_destroy=True, + is_intranode=not self._internode, + ) + + def _create_ll_buffer(self, spec): + """Construct the legacy low-latency Buffer (decode kernels over the CPU-proxy transport). + + Distinct from normal mode: LL always allocates the RDMA staging buffer and forces the + proxy path even for single-node EP8. Mirrors ep_deepep_v2._create_ll_buffer. + """ + args, world_size = self.args, self.world_size + assert args.experts % world_size == 0, ( + "low-latency EP requires num_experts divisible by the EP size" + ) + self.num_local_experts = args.experts // world_size + # LL requires the QP-per-rank count to equal the number of local experts. + num_qps_per_rank = self.num_local_experts + if not hasattr(Buffer, "low_latency_dispatch") and not hasattr( + Buffer, "get_low_latency_rdma_size_hint" + ): + raise RuntimeError( + "invalid UCCL-EP LL runtime: Buffer.low_latency_dispatch is absent" + ) + num_rdma_bytes = Buffer.get_low_latency_rdma_size_hint( + self.max_tokens, args.hidden, world_size, args.experts + ) + self.buffer = Buffer( + self.group, + 0, + num_rdma_bytes, + low_latency_mode=True, + num_qps_per_rank=num_qps_per_rank, + allow_nvlink_for_low_latency_mode=True, + explicitly_destroy=True, + is_intranode=not self._internode, + ) + + # ---- FP8 encode/dequant hooks ---------------------------------------------------------- + + def _topk_idx_dtype(self): + return torch.int64 + + def semantic_payload(self, x): + if not self._fp8: + return x + return per_token_cast_back(*per_token_cast_to_fp8(x)) + + def _encode_dispatch(self, x): + if not self._fp8: + return x, None + if self.mode == "low-latency": + # low_latency_dispatch takes BF16 x and casts to e4m3 inside the kernel, so send x + # unquantized; expose the host round-trip as the oracle semantic. + return x, per_token_cast_back(*per_token_cast_to_fp8(x)) + fp8, scales = per_token_cast_to_fp8(x) + # Column-major (TMA-compatible) scale layout the dispatch kernel expects, matching UCCL's + # own bench (`scales.T.contiguous().T`) and the LL scale-contiguity note below. + quantized = (fp8, scales.T.contiguous().T) + return quantized, per_token_cast_back(fp8, scales) + + def _ll_recv_bf16(self, recv_x): + """The padded per-expert receive as BF16 [num_local_experts, cap*num_ranks, hidden]. + + BF16 dispatch already returns that tensor; FP8 returns an (e4m3, per-128-block scale) + tuple, dequantized here with the static-shape compile (the LL fp8 scales come back + column-major / non-contiguous for TMA, so they are made contiguous before the per-block + view). Mirror of ep_deepep_v2._ll_recv_bf16. + """ + if not self._fp8: + return recv_x + fp8, scales = recv_x + return _ll_dequant_static(fp8, scales.contiguous()) + + # ---- transport contract ---------------------------------------------------------------- + + def _ll_dispatch(self, p): + recv_x, recv_count, ll_handle, _event, _hook = self.buffer.low_latency_dispatch( + p.dispatch_x, + p.topk_idx, + self.max_tokens, + self.args.experts, + use_fp8=self._fp8, + ) + return types.SimpleNamespace( + recv_x=recv_x, + recv_count=recv_count, + ll_handle=ll_handle, + ) + + def dispatch(self, p): + if self.mode == "low-latency": + return self._ll_dispatch(p) + # Legacy normal dispatch: compute the layout, then scatter tokens to their experts. + # num_tokens_per_rdma_rank is None intranode (EP8) and populated internode (EP16); pass + # it through so the same call serves both scopes. + (num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, + is_token_in_rank, _) = self.buffer.get_dispatch_layout(p.topk_idx, self.args.experts) + recv_x, recv_topk_idx, recv_topk_weights, _counts, handle, _event = self.buffer.dispatch( + x=p.dispatch_x, + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + is_token_in_rank=is_token_in_rank, + num_tokens_per_expert=num_tokens_per_expert, + topk_idx=p.topk_idx, + topk_weights=p.topk_weights, + config=self.config, + async_finish=False, + ) + return types.SimpleNamespace( + recv_x=recv_x, + recv_topk_idx=recv_topk_idx, + recv_topk_weights=recv_topk_weights, + handle=handle, + ) + + def stage(self, p, h): + if self.mode == "low-latency": + # The timed combine sends the padded per-expert receive back as BF16 (dequant under + # FP8). Value correctness is exercised by the oracle's combine_transformed path. + h.combine_input = self._ll_recv_bf16(h.recv_x) + return + if self._fp8: + h.combine_input = per_token_cast_back(h.recv_x[0], h.recv_x[1]) + else: + h.combine_input = h.recv_x + + def combine(self, p, h): + if self.mode == "low-latency": + combined_x, _event, _hook = self.buffer.low_latency_combine( + h.combine_input, p.topk_idx, p.topk_weights, h.ll_handle + ) + return combined_x[: p.T] + # Normal combine is the activation-only unweighted rank-sum: topk_weights are intentionally + # NOT passed so the kernel sums the per-token expert aggregates across ranks without + # applying the gate (matches combine_weight_semantics and the two-level oracle). + combined_x, _weights, _event = self.buffer.combine( + x=h.combine_input, + handle=h.handle, + config=self.config, + async_finish=False, + ) + return combined_x + + # ---- correctness-oracle views ---------------------------------------------------------- + + def _ll_inspect_dispatch(self, p, h): + """Flat per-slot view over the padded per-expert LL receive (mirror of + ep_deepep_v2._ll_inspect_dispatch).""" + recv_bf16 = self._ll_recv_bf16(h.recv_x) # [E, S, hidden] BF16 + num_slots = recv_bf16.shape[1] + counts = h.recv_count.to(torch.int64) # [E] + slot_valid = ( + torch.arange(num_slots, device=recv_bf16.device).unsqueeze(0) < counts.unsqueeze(1) + ) + slot_expert, slot_j = slot_valid.nonzero(as_tuple=True) + h.slot_expert = slot_expert + h.slot_j = slot_j + local_lo = self.rank * self.num_local_experts + return types.SimpleNamespace( + payload=recv_bf16[slot_expert, slot_j], + expert_ids=local_lo + slot_expert.to(torch.int64), + local_expert_counts=counts, + ) + + def _normal_recv_payload(self, h): + """The received tokens as BF16 [num_recv_tokens, hidden] (dequant under FP8).""" + if self._fp8: + return per_token_cast_back(h.recv_x[0], h.recv_x[1]) + return h.recv_x + + def inspect_dispatch(self, p, h): + if self.mode == "low-latency": + return self._ll_inspect_dispatch(p, h) + # Legacy normal recv: recv_x is [num_recv, hidden] (each received token once) and + # recv_topk_idx/recv_topk_weights are [num_recv, topk] — the oracle's per-received-token + # 2-D contract (it sorts each row over the topk axis and sums the per-expert transforms, + # so token order is free and no per-(token,expert) expansion is needed). recv_topk_idx + # holds LOCAL expert indices [0, experts_per_rank) with non-local masked to -1 (verified + # against UCCL's own test_intranode: every entry is -1 or < epr), so rebase the valid + # locals to the GLOBAL ids the oracle compares by rank*experts_per_rank. + local_idx = h.recv_topk_idx.to(torch.int64) # [num_recv, topk] local ids, -1 non-local + valid = local_idx >= 0 + expert_ids = torch.where( + valid, local_idx + self.rank * self.experts_per_rank, local_idx + ) + return types.SimpleNamespace( + payload=self._normal_recv_payload(h), # [num_recv, hidden] BF16 + expert_ids=expert_ids, + weights=h.recv_topk_weights.to(torch.float32), + local_expert_counts=torch.bincount( + local_idx[valid], minlength=self.experts_per_rank + ), + ) + + def _ll_combine_transformed(self, p, h, transformed): + """Scatter the oracle-transformed rows back into a zeroed padded combine buffer at the + exact (expert, slot) coordinates inspect read them from, then run the weighted LL combine + (mirror of ep_deepep_v2._ll_combine_transformed).""" + if self._fp8: + fp8 = h.recv_x[0] + combine_buf = torch.zeros(fp8.shape, dtype=torch.bfloat16, device=fp8.device) + else: + combine_buf = torch.zeros_like(h.recv_x) + combine_buf[h.slot_expert, h.slot_j] = transformed.to(combine_buf.dtype) + combined_x, _event, _hook = self.buffer.low_latency_combine( + combine_buf, p.topk_idx, p.topk_weights, h.ll_handle + ) + return combined_x[: p.T] + + def combine_transformed(self, p, h, transformed): + if self.mode == "low-latency": + return self._ll_combine_transformed(p, h, transformed) + # `transformed` is the oracle's per-received-token combine input [num_recv, hidden] + # (already summed over the top-k axis) — exactly the per-token buffer legacy combine + # consumes; combine then sums those per-token aggregates across ranks (unweighted). + combined, _weights, _event = self.buffer.combine( + x=transformed.to(torch.bfloat16), + handle=h.handle, + config=self.config, + async_finish=False, + ) + return combined + + def recv_tokens(self, h): + if self.mode == "low-latency": + return int(h.recv_count.sum().item()) + recv = h.recv_x[0] if self._fp8 else h.recv_x + return int(recv.shape[0]) + + def finalize(self, rc): + try: + dist.barrier() + self.buffer.destroy() # tears the CPU proxies down via destroy_uccl internally + dist.barrier() + dist.destroy_process_group() + except Exception: + return 1 + return rc diff --git a/experimental/CollectiveX/bench/run_ep.py b/experimental/CollectiveX/bench/run_ep.py index 084a3462e4..caaf9475d4 100644 --- a/experimental/CollectiveX/bench/run_ep.py +++ b/experimental/CollectiveX/bench/run_ep.py @@ -50,7 +50,7 @@ def _runtime_info(torch, *, vendor: str) -> dict: def main() -> int: ap = argparse.ArgumentParser(description="CollectiveX EP dispatch/combine sweep") - ap.add_argument("--backend", required=True, choices=["deepep-v2", "mori"]) + ap.add_argument("--backend", required=True, choices=["deepep-v2", "mori", "uccl-ep"]) ep_harness.add_common_args(ap) args = ap.parse_args() @@ -90,13 +90,20 @@ def main() -> int: # explicit case dimension; adapters do not infer it from the token ladder. if args.backend == "mori": from ep_mori import MoRIBackend as Backend + elif args.backend == "uccl-ep": + from ep_uccl import UCCLEPBackend as Backend else: from ep_deepep_v2 import DeepEPV2Backend as Backend # MoRI registers the default GPU process group with its SHMEM runtime. Keep that # group device-only so scale-out does not also depend on a host Gloo fabric. if not dist.is_initialized(): - if args.backend == "mori": + if args.backend in ("mori", "uccl-ep"): + # MoRI registers this group with its SHMEM runtime; UCCL-EP is portable across + # NVIDIA (NCCL) and AMD (RCCL) and bootstraps its Buffer + CPU-proxy ranks from + # it. Both take the explicit rank/world_size form and keep the group device-only + # so scale-out does not also depend on a host Gloo fabric. (uccl-ep's own default- + # PG registration, if any, happens in the adapter's create_buffer, mirroring MoRI.) dist.init_process_group( backend="nccl", rank=rank, diff --git a/experimental/CollectiveX/configs/platform_config.json b/experimental/CollectiveX/configs/platform_config.json index 9072bb7522..40a5f383bd 100644 --- a/experimental/CollectiveX/configs/platform_config.json +++ b/experimental/CollectiveX/configs/platform_config.json @@ -9,8 +9,8 @@ "scale_up_domain": 8, "scale_up_transport": "nvlink", "launcher": "single-slurm", - "backends": {"deepep-v2": [8, 16]}, - "ll_backends": {"deepep-v2": [8]}, + "backends": {"deepep-v2": [8, 16], "uccl-ep": [8, 16]}, + "ll_backends": {"deepep-v2": [8], "uccl-ep": [8]}, "fabric": {"nic": "ConnectX-7 2x200GbE", "switch": "Arista 7060DX5-64S (Tomahawk4, 25.6T)"}, "operator": { "partition": "hpc-gpu-1", @@ -33,8 +33,8 @@ "scale_up_domain": 8, "scale_up_transport": "nvlink", "launcher": "single-slurm", - "backends": {"deepep-v2": [8, 16]}, - "ll_backends": {"deepep-v2": [8]}, + "backends": {"deepep-v2": [8, 16], "uccl-ep": [8, 16]}, + "ll_backends": {"deepep-v2": [8], "uccl-ep": [8]}, "fabric": {"nic": "ConnectX-7 400G", "switch": "NVIDIA Quantum-2 QM9790 (25.6T, InfiniBand)"}, "operator": { "partition": "main", @@ -53,8 +53,8 @@ "scale_up_domain": 8, "scale_up_transport": "nvlink", "launcher": "single-slurm", - "backends": {"deepep-v2": [8, 16]}, - "ll_backends": {"deepep-v2": [8]}, + "backends": {"deepep-v2": [8, 16], "uccl-ep": [8, 16]}, + "ll_backends": {"deepep-v2": [8], "uccl-ep": [8]}, "fabric": {"nic": "ConnectX-7 400GbE", "switch": "Whitebox Tomahawk3 leaf + Tomahawk4 (RoCE)"}, "operator": { "partition": "gpu-2", @@ -138,7 +138,7 @@ "scale_up_domain": 8, "scale_up_transport": "xgmi", "launcher": "mi-tw", - "backends": {"mori": [8]}, + "backends": {"mori": [8], "uccl-ep": [8]}, "ll_backends": {"mori": [8]}, "fabric": {"nic": "n/a (single-node scale-up)", "switch": "n/a (single-node scale-up)"} }, @@ -151,7 +151,7 @@ "scale_up_domain": 8, "scale_up_transport": "xgmi", "launcher": "mi-tw", - "backends": {"mori": [8]}, + "backends": {"mori": [8], "uccl-ep": [8]}, "ll_backends": {"mori": [8]}, "fabric": {"nic": "n/a (single-node scale-up)", "switch": "n/a (single-node scale-up)"} }, @@ -164,7 +164,7 @@ "scale_up_domain": 8, "scale_up_transport": "xgmi", "launcher": "mi-amds", - "backends": {"mori": [8]}, + "backends": {"mori": [8], "uccl-ep": [8, 16]}, "ll_backends": {"mori": [8]}, "fabric": {"nic": "Pollara 400GbE", "switch": "Arista 7060X6-64PE (Tomahawk5, 51.2T)"}, "network": { diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index 2177430b17..d398ff02dc 100644 --- a/experimental/CollectiveX/docs/methodology.md +++ b/experimental/CollectiveX/docs/methodology.md @@ -51,7 +51,11 @@ x86 EP16 scale-out uses the hybrid path with GIN and requires two logical scale- represented by two physical RDMA ranks, with eight scale-up ranks per domain. GB EP16 remains MNNVL scale-up and uses LSA. MoRI EP8 uses the direct IntraNode kernel on every CDNA SKU; EP16 uses pinned InterNodeV1 over 2x8 XGMI + RDMA with 96 blocks, 64 RDMA blocks, 8 warps, one QP per PE, and external -input. Those throughput kernels run across the full token ladder in the `normal` mode. +input. UCCL-EP is a drop-in, API-identical DeepEP replacement that keeps the legacy `Buffer` +`dispatch`/`combine` (unweighted rank-sum) but routes it over CPU-proxy GPUDirect RDMA on plain +`libibverbs` — no NVSHMEM/IBGDA — with software message ordering, atomics, and flow control; its +scale-up is single-node `cudaIpc` over NVLink/XGMI (so the scale-up domain is one physical node, +never MNNVL) and its EP16 scale-out uses the same per-SKU RDMA rails as the other backends. Those throughput kernels run across the full token ladder in the `normal` mode. A second `low-latency` mode adds each backend's decode-optimized kernel family. On DeepEP it drives the legacy `deep_ep.Buffer` low-latency decode kernels (`low_latency_dispatch`/`low_latency_combine`), @@ -66,8 +70,10 @@ rank-sum combine as the throughput `IntraNode` kernel, so it differs only by ker does not fit the single-call dispatch/combine contract). Low latency is a decode-phase-only addition whose runnable set is narrower than and distinct from the throughput kernels', so it is enabled cell-by-cell from the registry's `ll_backends` map rather than assumed wherever `normal` runs; it is -currently enabled for DeepEP V2 EP8 on H100/H200/B200 and MoRI -EP8 on MI300X/MI325X/MI355X. Whether a given SKU/backend/EP/mode cell is attempted is a capability +currently enabled for DeepEP V2 EP8 on H100/H200/B200, MoRI +EP8 on MI300X/MI325X/MI355X, and UCCL-EP EP8 on H100/H200/B200 only (the legacy `Buffer` low-latency +kernels over UCCL's CPU-proxy transport; the AMD SKUs keep UCCL-EP normal mode but drop LL, whose +kernel trips a warp-group assertion on AMD's CU count). Whether a given SKU/backend/EP/mode cell is attempted is a capability fact; whether it succeeded is decided only by the emitted artifact. ## Workload Identity @@ -117,8 +123,8 @@ Logical payload bandwidth is: Payload bytes use rank-deduplicated token-rank activations and exclude expert metadata, padding, and backend buffer capacity. BF16 moves 2 bytes per value with no scale payload; an FP8 -dispatch moves 1 byte per value, plus per-128-block FP32 scales for DeepEP's blockwise codec (none -for MoRI's plain e4m3 cast), while combine stays BF16 — so the dispatch and combine directions can carry +dispatch moves 1 byte per value, plus per-128-block FP32 scales for DeepEP's and UCCL-EP's blockwise +codec (none for MoRI's plain e4m3 cast), while combine stays BF16 — so the dispatch and combine directions can carry different byte counts and the roundtrip is their per-field sum. The rank-deduplicated count is exact for the normal-mode layout; the low-latency layout sends one copy per (token, expert) assignment rather than per (token, rank), so for a token whose experts share a destination rank this logical diff --git a/experimental/CollectiveX/launchers/launch_mi-amds.sh b/experimental/CollectiveX/launchers/launch_mi-amds.sh index 70b3263db2..ea4bc6327b 100644 --- a/experimental/CollectiveX/launchers/launch_mi-amds.sh +++ b/experimental/CollectiveX/launchers/launch_mi-amds.sh @@ -34,7 +34,7 @@ NODELIST="${COLLX_NODELIST:-}" MOUNT_DIR=/ix TS="$(date -u +%Y-%m-%dT%H-%M-%SZ)" case "$COLLX_BENCH" in - mori) ;; + mori | uccl-ep) ;; *) collx_die "unsupported AMD EP backend: $COLLX_BENCH" ;; esac @@ -61,6 +61,11 @@ collx_log "runner=$RUNNER nodes=$NODES x ${GPN}gpu world=$NGPUS bench=$COLLX_BEN # ---- repository-stage: compute-visible copy of the checkout ----------------- MOUNT_SRC="$(collx_stage_path "$REPO_ROOT" "$COLLX_STAGE_DIR")" collx_stage_repo "$REPO_ROOT" "$MOUNT_SRC" +# UCCL builds from source (mori ships in the image); stage the pinned tree pre-allocation. +if [ "$COLLX_BENCH" = uccl-ep ]; then + collx_prepare_uccl_source "$MOUNT_SRC" || collx_die "cannot stage the pinned UCCL source" + export COLLX_BACKEND_SOURCE_ROOT=/ix/experimental/CollectiveX/.collx_sources +fi collx_select_image "$IMAGE" # ---- scheduler-allocation + container-import: retry until nodes validate ---- diff --git a/experimental/CollectiveX/launchers/launch_mi-tw.sh b/experimental/CollectiveX/launchers/launch_mi-tw.sh index 15e6dca28d..cfac84d3f6 100644 --- a/experimental/CollectiveX/launchers/launch_mi-tw.sh +++ b/experimental/CollectiveX/launchers/launch_mi-tw.sh @@ -21,7 +21,10 @@ case "$RUNNER" in *) collx_die "launch_mi-tw expects a Slurm-less -tw AMD SKU (mi325x-tw|mi300x-tw), got '${RUNNER}'" ;; esac export COLLX_RUNNER="$RUNNER" COLLX_BENCH="${COLLX_BENCH:-mori}" COLLX_VENDOR=amd -[ "$COLLX_BENCH" = mori ] || collx_die "mi325x-tw supports only the mori backend, got '$COLLX_BENCH'" +case "$COLLX_BENCH" in + mori | uccl-ep) ;; + *) collx_die "the -tw AMD clusters support only the mori and uccl-ep backends, got '$COLLX_BENCH'" ;; +esac # ---- setup: trimmed prologue (no Slurm stage-dir / enroot squash) ----------- # collx_launcher_prologue's collx_prepare_stage_dir requires COLLX_SQUASH_DIR (the @@ -61,6 +64,47 @@ fi collx_log "runner=$RUNNER nodes=1 x ${GPN}gpu world=$NGPUS bench=$COLLX_BENCH image=$IMAGE (${DOCKER[*]}/torchrun)" +# ---- uccl-ep: prepare source + one-time persisted build --------------------- +# UCCL is not in the image, so build it from source. Cases each run in a throwaway +# `docker run --rm`, so build ONCE here into a host-persisted prefix +# ($COLLX_DIR/.collx_uccl_pfx) that every case container puts on PYTHONPATH — the same +# build-once/reuse the enroot paths get from prepare_backend.sh, adapted to Docker. The +# AMD build applies the CDNA managed->pinned-host-memory patch (see prepare_backend.sh). +UCCL_PFX_HOST="$COLLX_DIR/.collx_uccl_pfx" +if [ "$COLLX_BENCH" = uccl-ep ]; then + REPO_ROOT="$(cd "$COLLX_DIR/../.." && pwd)" + collx_prepare_uccl_source "$REPO_ROOT" || collx_die "UCCL source preparation failed" + UCCL_ARCH="$(python3 - "$COLLX_DIR/configs/platform_config.json" "$RUNNER" <<'PY' +import json, sys +print(json.load(open(sys.argv[1]))["platforms"][sys.argv[2]]["arch"]) +PY +)" + if [ ! -d "$UCCL_PFX_HOST/deep_ep" ]; then + collx_log "uccl-ep: one-time from-source build (arch=$UCCL_ARCH, USE_DMABUF, host-atomic path)" + "${DOCKER[@]}" run --rm \ + --device /dev/kfd --device /dev/dri --group-add video --group-add render \ + --ipc host --shm-size 32g --cap-add SYS_PTRACE --security-opt seccomp=unconfined \ + -v "$COLLX_DIR:/cx" -w /cx "$IMAGE" \ + bash -c ' + set -e + { pip install -q nanobind || pip install -q --break-system-packages nanobind; } >&2 + rm -rf /tmp/ub && cp -R "/cx/.collx_sources/uccl-'"$COLLX_UCCL_COMMIT"'" /tmp/ub + # gfx942/gfx950 lack usable managed memory; swap UCCL'"'"'s cudaMallocManaged CPU-proxy + # handles to pinned host memory (coherent + device-accessible on CDNA). + sed -i "s/cudaMallocManaged/cudaMallocHost/g" /tmp/ub/ep/src/uccl_ep.cc /tmp/ub/ep/src/uccl_proxy.cpp + cd /tmp/ub/ep && env USE_DMABUF=1 PER_EXPERT_BATCHING=1 PYTORCH_ROCM_ARCH="'"$UCCL_ARCH"'" python3 setup.py install >&2 + cd /tmp/ub/ep/deep_ep_wrapper && python3 setup.py install >&2 + SP="$(python3 -c "import site;print(site.getsitepackages()[0])")" + mkdir -p /cx/.collx_uccl_pfx && cp -R "$SP"/deep_ep* "$SP"/uccl* /cx/.collx_uccl_pfx/ + python3 -c "import torch,sys; sys.path.insert(0,\"/cx/.collx_uccl_pfx\"); import deep_ep; from deep_ep import Buffer; assert hasattr(Buffer,\"get_dispatch_layout\")" >&2 + ' >&2 \ + || collx_die "uccl-ep from-source build failed" + collx_log "uccl-ep: build persisted to $UCCL_PFX_HOST" + else + collx_log "uccl-ep: reusing persisted build at $UCCL_PFX_HOST" + fi +fi + # ---- execute: one Docker+torchrun invocation per case ----------------------- # The shard control and results dir live under the CX source tree the workflow # checked out; mount that tree so run_ep.py's `results/*.json` land where the @@ -74,19 +118,30 @@ ncases="$(python3 "$COLLX_RUNTIME_DIR/config.py" case-count "$COLLX_SHARD_FILE") || collx_die "cannot count cases in $COLLX_SHARD_FILE" [ "$ncases" -gt 0 ] || collx_die "shard $COLLX_SHARD_FILE declares no cases" -# MoRI's SDMA "anvil" transport (hsaKmtCreateQueueExt with HSA_QUEUE_SDMA_BY_ENG_ID) -# fails at init on the mi300x-tw nodes' kernel thunk (anvil.cpp:193, both nodes), so -# disable it there and let MoRI fall back to the hipIpc/P2P intra-node path (correct -# results, normal latency). mi325x-tw's thunk accepts the SDMA queue, so keep it on. -mori_sdma_default=1 -[ "$RUNNER" = mi300x-tw ] && mori_sdma_default=0 -docker_env=( - -e MORI_DISABLE_AUTO_XGMI="${MORI_DISABLE_AUTO_XGMI:-0}" - -e MORI_ENABLE_SDMA="${MORI_ENABLE_SDMA:-$mori_sdma_default}" - -e MORI_APP_LOG_LEVEL="${MORI_APP_LOG_LEVEL:-info}" - -e HSA_NO_SCRATCH_RECLAIM=1 - -e COLLECTIVEX_SOURCE_SHA="${COLLECTIVEX_SOURCE_SHA:-}" -) +if [ "$COLLX_BENCH" = uccl-ep ]; then + # uccl-ep imports the host-persisted build via PYTHONPATH; CDNA needs the aggressive + # host-atomic EP path (matches prepare_backend.sh's uccl_prepare AMD branch). + docker_env=( + -e PYTHONPATH=/cx/.collx_uccl_pfx + -e UCCL_EP_ENABLE_AGGRESSIVE_ATOMIC="${UCCL_EP_ENABLE_AGGRESSIVE_ATOMIC:-1}" + -e HSA_NO_SCRATCH_RECLAIM=1 + -e COLLECTIVEX_SOURCE_SHA="${COLLECTIVEX_SOURCE_SHA:-}" + ) +else + # MoRI's SDMA "anvil" transport (hsaKmtCreateQueueExt with HSA_QUEUE_SDMA_BY_ENG_ID) + # fails at init on the mi300x-tw nodes' kernel thunk (anvil.cpp:193, both nodes), so + # disable it there and let MoRI fall back to the hipIpc/P2P intra-node path (correct + # results, normal latency). mi325x-tw's thunk accepts the SDMA queue, so keep it on. + mori_sdma_default=1 + [ "$RUNNER" = mi300x-tw ] && mori_sdma_default=0 + docker_env=( + -e MORI_DISABLE_AUTO_XGMI="${MORI_DISABLE_AUTO_XGMI:-0}" + -e MORI_ENABLE_SDMA="${MORI_ENABLE_SDMA:-$mori_sdma_default}" + -e MORI_APP_LOG_LEVEL="${MORI_APP_LOG_LEVEL:-info}" + -e HSA_NO_SCRATCH_RECLAIM=1 + -e COLLECTIVEX_SOURCE_SHA="${COLLECTIVEX_SOURCE_SHA:-}" + ) +fi final_rc=0 for ((ci = 0; ci < ncases; ci++)); do diff --git a/experimental/CollectiveX/launchers/launch_single-slurm.sh b/experimental/CollectiveX/launchers/launch_single-slurm.sh index c776eef250..f1f81ab986 100644 --- a/experimental/CollectiveX/launchers/launch_single-slurm.sh +++ b/experimental/CollectiveX/launchers/launch_single-slurm.sh @@ -47,7 +47,7 @@ TIME_MIN="${COLLX_TIME:-$DEFAULT_TIME}" IMAGE="$COLLX_IMAGE" TS="$(date -u +%Y-%m-%dT%H-%M-%SZ)" case "$COLLX_BENCH" in - deepep-v2) ;; + deepep-v2 | uccl-ep) ;; *) collx_die "unsupported $RUNNER EP backend: $COLLX_BENCH" ;; esac @@ -76,10 +76,15 @@ collx_select_image "$IMAGE" MOUNT_SRC="$(collx_stage_path "$REPO_ROOT" "${COLLX_STAGE_DIR:-}")" collx_stage_repo "$REPO_ROOT" "$MOUNT_SRC" CONTAINER_MOUNTS="$MOUNT_SRC:/ix" -# ---- backend-setup: pinned DeepEP source + isolated build cache ------------- -# The backend case above admits only deepep-v2, so its staging is unconditional. -collx_prepare_deepep_source "$MOUNT_SRC" \ - || collx_die "cannot stage the pinned backend source" +# ---- backend-setup: pinned backend source + isolated build cache ------------- +# Stage the pinned source for the selected from-source backend before allocation (the +# submit host has network; compute nodes may not). +case "$COLLX_BENCH" in + deepep-v2) collx_prepare_deepep_source "$MOUNT_SRC" \ + || collx_die "cannot stage the pinned DeepEP source" ;; + uccl-ep) collx_prepare_uccl_source "$MOUNT_SRC" \ + || collx_die "cannot stage the pinned UCCL source" ;; +esac export COLLX_BACKEND_SOURCE_ROOT=/ix/experimental/CollectiveX/.collx_sources collx_prepare_backend_cache "$COLLX_SQUASH_DIR" \ || collx_die "cannot prepare the isolated backend cache" diff --git a/experimental/CollectiveX/runtime/common.sh b/experimental/CollectiveX/runtime/common.sh index b9adb97320..706ffaeb65 100644 --- a/experimental/CollectiveX/runtime/common.sh +++ b/experimental/CollectiveX/runtime/common.sh @@ -14,6 +14,9 @@ collx_die() { printf '[collectivex] FATAL: %s\n' "$*" >&2; exit 1; } COLLX_DEEPEP_V2_REPO="https://github.com/deepseek-ai/DeepEP" COLLX_DEEPEP_V2_COMMIT="fa8a9b16898204afd347c663b89e65ef87dc6ce6" +COLLX_UCCL_REPO="https://github.com/uccl-project/uccl" +COLLX_UCCL_COMMIT="fc1b582031221645ea9fce58aeb57187713145e3" + # Print bounded command output without maintaining a parallel failure taxonomy. collx_log_tail() { local log_path="$1" @@ -159,6 +162,8 @@ collx_apply_network_profile() { unset EP_NIC_NAME EP_OVERRIDE_RDMA_SL unset MORI_RDMA_DEVICES unset MORI_RDMA_TC MORI_IO_TC MORI_RDMA_SL MORI_IO_SL + unset UCCL_SOCKET_IFNAME UCCL_IB_HCA UCCL_IB_GID_INDEX UCCL_IB_SL UCCL_IB_TC + unset UCCL_IB_MAX_INFLIGHT_BYTES UCCL_IB_MAX_INFLIGHT_NORMAL UCCL_EP_ENABLE_AGGRESSIVE_ATOMIC # Single-node and MNNVL runs need only the scrub above; everything past this # point is the scale-out path, so no per-branch scale-out guards remain. Single-node # low-latency also takes this early return: the decode kernels run over the intra-node @@ -191,6 +196,15 @@ collx_apply_network_profile() { fi export NCCL_IB_HCA="=$COLLX_RDMA_DEVICES" export MORI_RDMA_DEVICES="$rdma_names" EP_NIC_NAME="$ep_nic" + # UCCL-EP's CPU proxies read UCCL_* selectors but fall back to NCCL_IB_HCA / NCCL_IB_GID_INDEX + # (set above / by the link-layer helper) for the device list and GID, so only the socket iface + # and, on AMD, the strict Pollara/Broadcom flow control + CDNA host-atomic path need setting. + export UCCL_SOCKET_IFNAME="${COLLX_SOCKET_IFNAME:-}" + if [ "${COLLX_VENDOR:-nvidia}" = amd ]; then + export UCCL_IB_MAX_INFLIGHT_BYTES="${UCCL_IB_MAX_INFLIGHT_BYTES:-2097152}" + export UCCL_IB_MAX_INFLIGHT_NORMAL="${UCCL_IB_MAX_INFLIGHT_NORMAL:-1}" + export UCCL_EP_ENABLE_AGGRESSIVE_ATOMIC="${UCCL_EP_ENABLE_AGGRESSIVE_ATOMIC:-1}" + fi # The selector enumerates individual ports. NCCL's default dual-port fusion # would collapse each card into one "fused" device, and any fused device # disables NCCL GIN (init.cc nicFused gate) — the deep_ep EP16 hybrid path @@ -216,11 +230,13 @@ collx_apply_network_profile() { export NCCL_IB_SL="$COLLX_RDMA_SERVICE_LEVEL" export EP_OVERRIDE_RDMA_SL="$COLLX_RDMA_SERVICE_LEVEL" export MORI_RDMA_SL="$COLLX_RDMA_SERVICE_LEVEL" MORI_IO_SL="$COLLX_RDMA_SERVICE_LEVEL" + export UCCL_IB_SL="$COLLX_RDMA_SERVICE_LEVEL" fi if [ -n "${COLLX_RDMA_TRAFFIC_CLASS:-}" ]; then [[ "$COLLX_RDMA_TRAFFIC_CLASS" =~ ^[0-9]+$ ]] && [ "$COLLX_RDMA_TRAFFIC_CLASS" -le 255 ] \ || collx_die "invalid private RDMA traffic class" export MORI_RDMA_TC="$COLLX_RDMA_TRAFFIC_CLASS" MORI_IO_TC="$COLLX_RDMA_TRAFFIC_CLASS" + export UCCL_IB_TC="$COLLX_RDMA_TRAFFIC_CLASS" fi local nic_handler=gpu export NVSHMEM_IB_ENABLE_IBGDA=1 NVSHMEM_IBGDA_NIC_HANDLER="$nic_handler" @@ -500,6 +516,44 @@ collx_materialize_deepep_source() { rm -rf -- "$destination" && cp -R -- "$source" "$destination" } +# Fetch the pinned UCCL tree before allocating GPUs. Like the DeepEP fetch, this runs on the +# submit host (which has network) because compute nodes may not reach GitHub. The EP extension +# needs the main tree (ep/ + top-level util/ + include/) but NOT the thirdparty submodules +# (rccl/mscclpp, for other targets), so this skips them — faster and sufficient. NB: build the +# whole tree, not ep/ alone: the ROCm path (common_hip.hpp) includes top-level util/gpu_rt.h. +collx_prepare_uccl_source() { + local mount_src="$1" root source temporary log + root="$mount_src/experimental/CollectiveX/.collx_sources" + source="$root/uccl-$COLLX_UCCL_COMMIT" + [ ! -d "$source" ] || return 0 + mkdir -p -- "$root" && chmod 700 "$root" || return 1 + temporary="$(mktemp -d "$root/.uccl.XXXXXX")" || return 1 + log="$(collx_private_log_path backend-source-uccl)" || return 1 + git config --global --add safe.directory '*' >> "$log" 2>&1 || true + if GIT_TERMINAL_PROMPT=0 git init -q "$temporary" > "$log" 2>&1 \ + && git -C "$temporary" remote add origin "$COLLX_UCCL_REPO" >> "$log" 2>&1 \ + && GIT_TERMINAL_PROMPT=0 git -C "$temporary" fetch -q --no-tags --depth 1 \ + origin "$COLLX_UCCL_COMMIT" >> "$log" 2>&1 \ + && git -C "$temporary" -c advice.detachedHead=false checkout -q --detach FETCH_HEAD \ + >> "$log" 2>&1 \ + && [ "$(git -C "$temporary" rev-parse HEAD)" = "$COLLX_UCCL_COMMIT" ] \ + && mv -- "$temporary" "$source" >> "$log" 2>&1; then + return 0 + fi + rm -rf -- "$temporary" + collx_log "ERROR: UCCL source preparation failed" + collx_log_tail "$log" + return 1 +} + +collx_materialize_uccl_source() { + local destination="$1" source + [ -n "${COLLX_BACKEND_SOURCE_ROOT:-}" ] || return 1 + source="$COLLX_BACKEND_SOURCE_ROOT/uccl-$COLLX_UCCL_COMMIT" + [ -d "$source" ] || return 1 + rm -rf -- "$destination" && cp -R -- "$source" "$destination" +} + collx_prepare_implicit_stage_base() { python3 "$COLLX_RUNTIME_DIR/stage.py" implicit-stage-base "${1:-}" "${2:-}" } diff --git a/experimental/CollectiveX/runtime/prepare_backend.sh b/experimental/CollectiveX/runtime/prepare_backend.sh index 132203fec1..2647b44113 100644 --- a/experimental/CollectiveX/runtime/prepare_backend.sh +++ b/experimental/CollectiveX/runtime/prepare_backend.sh @@ -16,7 +16,7 @@ collx_log "backend preparation: runner=$COLLX_RUNNER bench=$COLLX_BENCH nodes=${ readonly -a RANK_ENV_VARS=( PATH VIRTUAL_ENV LD_LIBRARY_PATH PYTHONPATH CUDA_HOME CPATH NVCC_PREPEND_FLAGS NVSHMEM_DIR EP_NCCL_ROOT_DIR EP_NVSHMEM_ROOT_DIR EP_JIT_CACHE_DIR - EP_REUSE_NCCL_COMM NCCL_CUMEM_ENABLE + EP_REUSE_NCCL_COMM NCCL_CUMEM_ENABLE UCCL_EP_ENABLE_AGGRESSIVE_ATOMIC ) readonly -a DEEPEP_RANK_UNSETS=(EP_SUPPRESS_NCCL_CHECK) @@ -256,6 +256,70 @@ deepep_prepare() { collx_log "DeepEP V2 ready ($COLLX_DEEPEP_V2_COMMIT, ElasticBuffer, NCCL Device API; LSA/Gin selected by adapter)" } +# ---- UCCL-EP lifecycle ------------------------------------------------------ + +# Registry arch string for the runner (gfx942/gfx950 on AMD) for PYTORCH_ROCM_ARCH. +uccl_rocm_arch() { + python3 - "$COLLX_RUNNER" <<'PY' +import json, sys +print(json.load(open("configs/platform_config.json"))["platforms"][sys.argv[1]]["arch"]) +PY +} + +uccl_probe() { + # import torch FIRST so libc10 is resident before the uccl.ep extension dlopens (it links + # libc10/libtorch); importing deep_ep before torch fails with "libc10.so: cannot open". + python3 - <<'PY' +import torch # noqa: F401 +import deep_ep +from deep_ep import Buffer +assert hasattr(Buffer, "low_latency_dispatch") and hasattr(Buffer, "get_dispatch_layout") +PY +} + +# Direct in-container source build against the image's torch — validated on h200 (sglang +# cu130). NOT `build.sh` (that spins up its own Docker image to make a wheel and cannot run +# inside enroot/pyxis). Installs into the image's system python; single-slurm and mi-amds run +# the writable container as remapped root, so this needs no venv. verbs/nl/numa dev headers +# ship in the sglang/rocm images; only nanobind must be added. +uccl_prepare() { + local source_dir="/tmp/collectivex-uccl-$COLLX_UCCL_COMMIT" arch_env + command -v python3 >/dev/null || { collx_log "ERROR: python3 unavailable for UCCL build"; return 1; } + collx_log "UCCL-EP: building $COLLX_UCCL_COMMIT from source (USE_DMABUF, PER_EXPERT_BATCHING)" + # Plain install first; some sglang/rocm image variants mark the system env externally-managed + # (PEP 668), so fall back to --break-system-packages (a no-op on older pip that lacks the flag). + { python3 -m pip install -q --disable-pip-version-check --no-input nanobind \ + || python3 -m pip install -q --disable-pip-version-check --no-input \ + --break-system-packages nanobind; } >&2 2>&1 \ + || { collx_log "ERROR: UCCL nanobind install failed"; return 1; } + collx_materialize_uccl_source "$source_dir" \ + || { collx_log "ERROR: UCCL staged source is invalid"; return 1; } + if [ "${COLLX_VENDOR:-nvidia}" = amd ]; then + arch_env="PYTORCH_ROCM_ARCH=$(uccl_rocm_arch)" + # CDNA requires the aggressive host-atomic EP path; persist it for the rank tasks too. + export UCCL_EP_ENABLE_AGGRESSIVE_ATOMIC=1 + # Managed/unified memory (cudaMallocManaged) is unavailable on our CDNA nodes (hipMallocManaged + # fails even for 4 KiB, regardless of XNACK / --privileged / memlock). UCCL's HIP CPU-proxy path + # uses it for the d2h channel handles + proxy atomic buffer; pinned host memory (cudaMallocHost) + # is coherent + device-accessible on gfx942/gfx950 and is already used elsewhere in UCCL (e.g. + # the RDMA scratch), so swap the two on the runtime path before building. Validated on mi300x-tw + # (bf16/fp8 normal green). NB: build the WHOLE tree (materialize copies it) — the ROCm path + # includes top-level util/gpu_rt.h. + sed -i 's/cudaMallocManaged/cudaMallocHost/g' \ + "$source_dir/ep/src/uccl_ep.cc" "$source_dir/ep/src/uccl_proxy.cpp" \ + || { collx_log "ERROR: UCCL AMD managed-memory patch failed"; return 1; } + else + arch_env="TORCH_CUDA_ARCH_LIST=$(cuda_arch)" + fi + ( cd "$source_dir/ep" \ + && env USE_DMABUF=1 PER_EXPERT_BATCHING=1 "$arch_env" python3 setup.py install ) >&2 2>&1 \ + || { collx_log "ERROR: UCCL ep extension build failed"; return 1; } + ( cd "$source_dir/ep/deep_ep_wrapper" && python3 setup.py install ) >&2 2>&1 \ + || { collx_log "ERROR: UCCL deep_ep_wrapper build failed"; return 1; } + uccl_probe || { collx_log "ERROR: UCCL import probe failed"; return 1; } + collx_log "UCCL-EP ready ($COLLX_UCCL_COMMIT, deep_ep wrapper over uccl.ep CPU-proxy runtime)" +} + # ---- container boundary ---------------------------------------------------- write_rank_env() { @@ -313,6 +377,7 @@ main() { python3 -c "import mori" \ || { collx_log "ERROR: MoRI backend import failed"; return 1; } ;; + uccl-ep) uccl_prepare || return 1 ;; *) collx_log "ERROR: unknown backend preparation request" return 1 diff --git a/experimental/CollectiveX/sweep_matrix.py b/experimental/CollectiveX/sweep_matrix.py index 5d3737cf9c..2dff3f5497 100644 --- a/experimental/CollectiveX/sweep_matrix.py +++ b/experimental/CollectiveX/sweep_matrix.py @@ -35,6 +35,7 @@ def _load_config(name: str) -> dict[str, Any]: BACKEND_PRECISIONS = { "deepep-v2": ("bf16", "fp8"), "mori": ("bf16", "fp8"), + "uccl-ep": ("bf16", "fp8"), } # Short shard-ID slug per non-normal mode. Normal-mode shard IDs carry no mode # segment so existing references stay valid; a low-latency shard adds "-ll". diff --git a/experimental/CollectiveX/tests/test_matrix.py b/experimental/CollectiveX/tests/test_matrix.py index 96a0e44d2f..f07f09f87e 100644 --- a/experimental/CollectiveX/tests/test_matrix.py +++ b/experimental/CollectiveX/tests/test_matrix.py @@ -100,7 +100,7 @@ def test_runnable_cases_fan_out_over_backend_precisions(self): if precision in sweep_matrix.BACKEND_PRECISIONS[cell[1]] } self.assertEqual(precisions, expected, cell) - # Both current backends realize BF16 and FP8. + # Every backend (deepep-v2, mori, uccl-ep) realizes BF16 and FP8. self.assertEqual( {precision for precisions in by_cell.values() for precision in precisions}, {"bf16", "fp8"}, @@ -161,6 +161,46 @@ def test_ll_backends_is_a_well_formed_subset_of_backends(self): self.assertTrue(degrees) self.assertLessEqual(set(degrees), set(platform["backends"][backend])) + def test_uccl_ep_rollout_shape(self): + # UCCL-EP's intended rollout, locked here: runnable on exactly the six supported + # SKUs; EP16 (cross-node) on h100/h200/b200/mi355x; EP8-only on the -tw pair (no + # cross-node fabric, so their EP16 cells are unsupported placeholders); and no rows + # at all on b300/gb200/gb300, where the backend is not offered (a missing backends + # key means "not provided", exactly like NVIDIA SKUs carrying no mori rows). LL + # (decode) on every supported SKU at EP8. + document = matrix(backend="all") + runnable = { + (item["sku"], item["case"]["ep"]) + for item in document["requested_cases"] + if item["case"]["backend"] == "uccl-ep" and item["disposition"] == "runnable" + } + unsupported = { + (item["sku"], item["case"]["ep"]) + for item in document["requested_cases"] + if item["case"]["backend"] == "uccl-ep" and item["disposition"] == "unsupported" + } + supported_skus = { + "h100-dgxc", "h200-dgxc", "b200-dgxc", "mi355x", "mi325x-tw", "mi300x-tw", + } + self.assertEqual({sku for sku, _ in runnable}, supported_skus) + self.assertEqual( + {sku for sku, ep in runnable if ep == 16}, + {"h100-dgxc", "h200-dgxc", "b200-dgxc", "mi355x"}, + ) + self.assertEqual(unsupported, {("mi325x-tw", 16), ("mi300x-tw", 16)}) + offered = {sku for sku, _ in runnable | unsupported} + for absent in ("b300", "gb200", "gb300"): + self.assertNotIn(absent, offered) + # uccl-ep low-latency is enabled only on NVIDIA; the AMD SKUs keep normal mode but drop + # LL (UCCL's low-latency kernel trips a warp-group assertion on AMD's CU count). + ll_skus = { + item["sku"] + for item in document["requested_cases"] + if item["case"]["backend"] == "uccl-ep" + and item["case"]["mode"] == "low-latency" + } + self.assertEqual(ll_skus, {"h100-dgxc", "h200-dgxc", "b200-dgxc"}) + def test_invalid_filters_fail_closed(self): for options in ( {"exclude_skus": "unknown"}, diff --git a/experimental/CollectiveX/tests/test_runtime.py b/experimental/CollectiveX/tests/test_runtime.py index 71acce484a..a3795ecae1 100644 --- a/experimental/CollectiveX/tests/test_runtime.py +++ b/experimental/CollectiveX/tests/test_runtime.py @@ -340,7 +340,7 @@ def _run_ep_parser() -> argparse.ArgumentParser: # Mirror of the parser bench/run_ep.py builds in main(). parser = argparse.ArgumentParser() parser.add_argument( - "--backend", required=True, choices=["deepep-v2", "mori"] + "--backend", required=True, choices=["deepep-v2", "mori", "uccl-ep"] ) ep_harness.add_common_args(parser) return parser @@ -400,6 +400,21 @@ def test_low_latency_case_round_trips_through_the_run_ep_parser(self) -> None: self.assertEqual((args.mode, args.phase, args.scope), ("low-latency", "decode", "scale-up")) self.assertEqual(args.case_id, ll_case["case_id"]) + def test_uccl_ep_case_round_trips_through_the_run_ep_parser(self) -> None: + # A uccl-ep case flows through the same generic codec; run_ep's --backend choices + # must accept "uccl-ep" and the result filename must carry the backend token so a + # uccl-ep leg never collides with the deepep-v2/mori legs of the same cell. + uccl_case = { + **self.CASE, + "backend": "uccl-ep", + "case_id": "h200-dgxc-uccl-ep-deepseek-v3-normal-decode-ep16-uniform-bf16", + } + argv = self._case_argv(["16", "2", "8", "8"], case=uccl_case) + args = self._run_ep_parser().parse_args(argv) + self.assertEqual(args.backend, "uccl-ep") + self.assertEqual(args.case_id, uccl_case["case_id"]) + self.assertEqual(args.out, "results/h200-dgxc_uccl-ep_bf16_decode_TS-c000.json") + # logical_byte_provenance is where FP8 changes MEASUREMENT semantics (asymmetric # per-direction byte counts), so its arithmetic and guards are pinned here on CPU. From faa0f92e84a85f3782fc8c0611adcb72f9c943fd Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Mon, 20 Jul 2026 17:49:31 +0800 Subject: [PATCH 02/10] ci(collectivex): make uccl-ep a dispatchable sweep backend choice --- .github/workflows/collectivex-sweep.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/collectivex-sweep.yml b/.github/workflows/collectivex-sweep.yml index 5a149bd966..1038c9b5c7 100644 --- a/.github/workflows/collectivex-sweep.yml +++ b/.github/workflows/collectivex-sweep.yml @@ -10,7 +10,7 @@ on: description: "EP library to sweep — 'all' runs every EP backend in one matrix" type: choice default: all - options: [all, deepep-v2, mori] + options: [all, deepep-v2, mori, uccl-ep] only_sku: description: Restrict to one GHA runner pool; blank = all type: string From 9221172e281b628d118d52ce33ba89c379d020fe Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:14:20 +0800 Subject: [PATCH 03/10] fix(collectivex): uccl-ep EP16 internode config (rdma-chunked-send constraint) 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. --- experimental/CollectiveX/bench/ep_uccl.py | 61 +++++++++++++++++------ 1 file changed, 47 insertions(+), 14 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_uccl.py b/experimental/CollectiveX/bench/ep_uccl.py index 170158a5d9..4dc352d549 100644 --- a/experimental/CollectiveX/bench/ep_uccl.py +++ b/experimental/CollectiveX/bench/ep_uccl.py @@ -101,6 +101,15 @@ def _ll_dequant_static(fp8, scales): _NORMAL_NUM_SMS = 24 _NORMAL_NVL_BUFFER_SIZE = 256 _NORMAL_NVL_BYTES = int(2e9) +# Internode (EP16) buffer-sizing Config, straight from UCCL's own test_internode bench +# compute_buffer_sizes (nvl_chunk=8/512, rdma_chunk=16/512). Used only to size the NVLink+RDMA +# staging; dispatch/combine themselves run on the per-world-size recommended configs. +_INTERNODE_SIZE_CFG = (8, 512, 16, 512) + + +def _align_buffer_bytes(size, margin=1.2, alignment=128): + """Safety margin + alignment for a buffer-size hint (mirrors the UCCL bench helper).""" + return ((int(size * margin) + alignment - 1) // alignment) * alignment class UCCLEPBackend(EPBackend): @@ -158,28 +167,52 @@ def create_buffer(self, spec): self._create_ll_buffer(spec) return args, world_size = self.args, self.world_size - # Normal mode: legacy Buffer with an NVLink scale-up buffer plus (internode) an RDMA - # buffer. UCCL's own intranode bench sizes num_nvl_bytes at a fixed ~2 GB; internode adds - # an RDMA buffer sized from the Config hint. Reservations are generous and backed on - # demand; num_sms/nvl_buffer are bring-up-tunable (see _NORMAL_* above). - self.config = Config(_NORMAL_NUM_SMS, 8, _NORMAL_NVL_BUFFER_SIZE) - num_rdma_bytes = 0 if self._internode: - # BF16 combine wire is the widest per-value direction; size the RDMA staging for it. + # Internode (EP16) scale-out: the RDMA combine kernel asserts + # num_max_rdma_chunked_send_tokens >= num_warps_per_forwarder, which a hand-rolled + # Config does NOT satisfy (its rdma-chunked-send default is 6). Mirror UCCL's own + # internode bench: size NVLink+RDMA from a generous sizing Config, give each rank + # num_sms QPs, and drive dispatch/combine with the per-world-size RECOMMENDED configs + # (these set rdma-chunked-send to 20/12 for EP16 and satisfy the kernel constraints). + num_sms = Buffer.num_sms + self.dispatch_config = Buffer.get_dispatch_config(world_size) + self.combine_config = Buffer.get_combine_config(world_size) + self.config = self.combine_config hidden_bytes = args.hidden * 2 - num_rdma_bytes = int( - self.config.get_rdma_buffer_size_hint(hidden_bytes, world_size) + size_config = Config(num_sms, *_INTERNODE_SIZE_CFG) + num_nvl_bytes = _align_buffer_bytes( + size_config.get_nvl_buffer_size_hint(hidden_bytes, world_size) + ) + num_rdma_bytes = _align_buffer_bytes( + size_config.get_rdma_buffer_size_hint(hidden_bytes, world_size) + ) + self.buffer = Buffer( + self.group, + num_nvl_bytes, + num_rdma_bytes, + low_latency_mode=False, + num_qps_per_rank=num_sms, + allow_nvlink_for_low_latency_mode=True, + allow_mnnvl=False, + explicitly_destroy=True, + is_intranode=False, ) + return + # Intranode (EP8) scale-up: validated recipe — one fixed ~2 GB NVLink buffer, no RDMA, a + # single QP, and the legacy 3-arg Config (rdma-chunked params unused with no RDMA path). + self.config = Config(_NORMAL_NUM_SMS, 8, _NORMAL_NVL_BUFFER_SIZE) + self.dispatch_config = self.config + self.combine_config = self.config self.buffer = Buffer( self.group, _NORMAL_NVL_BYTES, - num_rdma_bytes, + 0, low_latency_mode=False, num_qps_per_rank=1, allow_nvlink_for_low_latency_mode=True, allow_mnnvl=False, explicitly_destroy=True, - is_intranode=not self._internode, + is_intranode=True, ) def _create_ll_buffer(self, spec): @@ -283,7 +316,7 @@ def dispatch(self, p): num_tokens_per_expert=num_tokens_per_expert, topk_idx=p.topk_idx, topk_weights=p.topk_weights, - config=self.config, + config=self.dispatch_config, async_finish=False, ) return types.SimpleNamespace( @@ -316,7 +349,7 @@ def combine(self, p, h): combined_x, _weights, _event = self.buffer.combine( x=h.combine_input, handle=h.handle, - config=self.config, + config=self.combine_config, async_finish=False, ) return combined_x @@ -396,7 +429,7 @@ def combine_transformed(self, p, h, transformed): combined, _weights, _event = self.buffer.combine( x=transformed.to(torch.bfloat16), handle=h.handle, - config=self.config, + config=self.combine_config, async_finish=False, ) return combined From b3c5e993e1d61f83e7a4f45566ee7c5e92e386a9 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:48:40 +0800 Subject: [PATCH 04/10] fix(collectivex): scope uccl-ep to EP8 (cross-node EP16 is wall-clock-limited) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- experimental/CollectiveX/README.md | 2 +- .../CollectiveX/configs/platform_config.json | 8 +++---- experimental/CollectiveX/tests/test_matrix.py | 23 ++++++++++--------- 3 files changed, 17 insertions(+), 16 deletions(-) diff --git a/experimental/CollectiveX/README.md b/experimental/CollectiveX/README.md index ff3ea1ee8d..77504b5a8c 100644 --- a/experimental/CollectiveX/README.md +++ b/experimental/CollectiveX/README.md @@ -66,7 +66,7 @@ scale-up domain. |---|---| | DeepEP V2 | `normal` mode is PR #605 `ElasticBuffer` plus exact upstream #630 and #640 fixes: LSA for scale-up and GIN for x86 EP16 scale-out. FP8 dispatch via `use_fp8_dispatch` (blockwise e4m3fn) alongside BF16. `low-latency` mode is the legacy `deep_ep.Buffer` IBGDA decode kernels (per-expert padded layout, weighted combine, `use_fp8` e4m3fn), decode/EP8 only | | MoRI | `normal` mode uses the direct `IntraNode` kernel for scale-up EP8 on every CDNA SKU and pins `InterNodeV1` for EP16 over 2x8 XGMI + RDMA. `low-latency` mode selects the `IntraNodeLL` decode kernel (single-call, pure-intranode, same compact layout and unweighted combine as `IntraNode`), decode/EP8 only. FP8 dispatch is caller-prequantized (per-SKU e4m3fnuz on gfx942, e4m3fn on gfx950); combine stays BF16 (`quant_type=none`) alongside BF16 dispatch | -| UCCL-EP | [UCCL](https://github.com/uccl-project/uccl) EP: a drop-in, API-identical DeepEP replacement whose CPU proxies issue GPUDirect RDMA over plain `libibverbs` (no NVSHMEM/IBGDA), with software message ordering, atomics, and flow control; scale-up is single-node `cudaIpc` over NVLink/XGMI (never MNNVL). `normal` mode is the legacy `Buffer` `dispatch`/`combine` (unweighted rank-sum); `low-latency` reuses the legacy `low_latency_dispatch`/`low_latency_combine` decode kernels (weighted combine), decode/EP8 only. FP8 dispatch is caller-prequantized (blockwise e4m3fn, per-SKU e4m3fnuz on gfx942) with BF16 combine. Runs on NVIDIA and AMD (H100/H200/B200 + MI300X/MI325X/MI355X); EP16 on h100/h200/b200/mi355x is a wall probe against the deepep-v2/mori scale-out limits | +| UCCL-EP | [UCCL](https://github.com/uccl-project/uccl) EP: a drop-in, API-identical DeepEP replacement whose CPU proxies issue GPUDirect RDMA over plain `libibverbs` (no NVSHMEM/IBGDA), with software message ordering, atomics, and flow control; scale-up is single-node `cudaIpc` over NVLink/XGMI (never MNNVL). `normal` mode is the legacy `Buffer` `dispatch`/`combine` (unweighted rank-sum); `low-latency` reuses the legacy `low_latency_dispatch`/`low_latency_combine` decode kernels (weighted combine), decode/EP8 only. FP8 dispatch is caller-prequantized (blockwise e4m3fn, per-SKU e4m3fnuz on gfx942) with BF16 combine. Runs on NVIDIA and AMD (H100/H200/B200 + MI300X/MI325X/MI355X), EP8 scale-up. Cross-node EP16 is functional (the internode RDMA path connects and the light case passes correctness) but its CPU-proxy throughput overruns the standardized per-case wall-clock budget on heavy token counts, so EP16 is an unsupported coverage row for now | DeepEP V2 means the `ElasticBuffer` implementation introduced by [DeepEP PR #605](https://github.com/deepseek-ai/DeepEP/pull/605), not a newer legacy `Buffer` build. diff --git a/experimental/CollectiveX/configs/platform_config.json b/experimental/CollectiveX/configs/platform_config.json index 40a5f383bd..45925f9991 100644 --- a/experimental/CollectiveX/configs/platform_config.json +++ b/experimental/CollectiveX/configs/platform_config.json @@ -9,7 +9,7 @@ "scale_up_domain": 8, "scale_up_transport": "nvlink", "launcher": "single-slurm", - "backends": {"deepep-v2": [8, 16], "uccl-ep": [8, 16]}, + "backends": {"deepep-v2": [8, 16], "uccl-ep": [8]}, "ll_backends": {"deepep-v2": [8], "uccl-ep": [8]}, "fabric": {"nic": "ConnectX-7 2x200GbE", "switch": "Arista 7060DX5-64S (Tomahawk4, 25.6T)"}, "operator": { @@ -33,7 +33,7 @@ "scale_up_domain": 8, "scale_up_transport": "nvlink", "launcher": "single-slurm", - "backends": {"deepep-v2": [8, 16], "uccl-ep": [8, 16]}, + "backends": {"deepep-v2": [8, 16], "uccl-ep": [8]}, "ll_backends": {"deepep-v2": [8], "uccl-ep": [8]}, "fabric": {"nic": "ConnectX-7 400G", "switch": "NVIDIA Quantum-2 QM9790 (25.6T, InfiniBand)"}, "operator": { @@ -53,7 +53,7 @@ "scale_up_domain": 8, "scale_up_transport": "nvlink", "launcher": "single-slurm", - "backends": {"deepep-v2": [8, 16], "uccl-ep": [8, 16]}, + "backends": {"deepep-v2": [8, 16], "uccl-ep": [8]}, "ll_backends": {"deepep-v2": [8], "uccl-ep": [8]}, "fabric": {"nic": "ConnectX-7 400GbE", "switch": "Whitebox Tomahawk3 leaf + Tomahawk4 (RoCE)"}, "operator": { @@ -164,7 +164,7 @@ "scale_up_domain": 8, "scale_up_transport": "xgmi", "launcher": "mi-amds", - "backends": {"mori": [8], "uccl-ep": [8, 16]}, + "backends": {"mori": [8], "uccl-ep": [8]}, "ll_backends": {"mori": [8]}, "fabric": {"nic": "Pollara 400GbE", "switch": "Arista 7060X6-64PE (Tomahawk5, 51.2T)"}, "network": { diff --git a/experimental/CollectiveX/tests/test_matrix.py b/experimental/CollectiveX/tests/test_matrix.py index f07f09f87e..f73c771327 100644 --- a/experimental/CollectiveX/tests/test_matrix.py +++ b/experimental/CollectiveX/tests/test_matrix.py @@ -162,12 +162,13 @@ def test_ll_backends_is_a_well_formed_subset_of_backends(self): self.assertLessEqual(set(degrees), set(platform["backends"][backend])) def test_uccl_ep_rollout_shape(self): - # UCCL-EP's intended rollout, locked here: runnable on exactly the six supported - # SKUs; EP16 (cross-node) on h100/h200/b200/mi355x; EP8-only on the -tw pair (no - # cross-node fabric, so their EP16 cells are unsupported placeholders); and no rows - # at all on b300/gb200/gb300, where the backend is not offered (a missing backends - # key means "not provided", exactly like NVIDIA SKUs carrying no mori rows). LL - # (decode) on every supported SKU at EP8. + # UCCL-EP's rollout, locked here: EP8 runnable on exactly the six supported SKUs, and + # EP16 an unsupported coverage row on every one of them. uccl-ep is EP8-only: the + # -tw pair has no cross-node fabric, and on the fabric SKUs cross-node EP16 is + # functional but its CPU-proxy throughput overruns the standardized per-case + # wall-clock budget (the internode Config fix landed; EP16 stays scoped out of the + # sweep, mirroring the mori EP16 re-wall). No rows at all on b300/gb200/gb300, where + # the backend is not offered. LL (decode) on every NVIDIA supported SKU at EP8. document = matrix(backend="all") runnable = { (item["sku"], item["case"]["ep"]) @@ -182,12 +183,12 @@ def test_uccl_ep_rollout_shape(self): supported_skus = { "h100-dgxc", "h200-dgxc", "b200-dgxc", "mi355x", "mi325x-tw", "mi300x-tw", } + # EP8 runnable on all six; nothing runnable at EP16. self.assertEqual({sku for sku, _ in runnable}, supported_skus) - self.assertEqual( - {sku for sku, ep in runnable if ep == 16}, - {"h100-dgxc", "h200-dgxc", "b200-dgxc", "mi355x"}, - ) - self.assertEqual(unsupported, {("mi325x-tw", 16), ("mi300x-tw", 16)}) + self.assertEqual({sku for sku, ep in runnable if ep == 8}, supported_skus) + self.assertEqual({sku for sku, ep in runnable if ep == 16}, set()) + # EP16 is an honest unsupported coverage row on every supported SKU. + self.assertEqual(unsupported, {(sku, 16) for sku in supported_skus}) offered = {sku for sku, _ in runnable | unsupported} for absent in ("b300", "gb200", "gb300"): self.assertNotIn(absent, offered) From dead72665f9f9f9cac12239b6ed56e416a6f71c0 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Mon, 20 Jul 2026 21:39:59 +0800 Subject: [PATCH 05/10] ci(collectivex): raise sweep max-parallel to 30 (avoid free-pool starvation) --- .github/workflows/collectivex-sweep.yml | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/.github/workflows/collectivex-sweep.yml b/.github/workflows/collectivex-sweep.yml index 1038c9b5c7..9a9d9b33fb 100644 --- a/.github/workflows/collectivex-sweep.yml +++ b/.github/workflows/collectivex-sweep.yml @@ -84,10 +84,12 @@ jobs: if: ${{ fromJSON(needs.setup.outputs.n) > 0 }} strategy: fail-fast: false - # Fixed global cap; real throttling is each SKU pool's runner count and its - # cluster's Slurm partition. setup interleaves shards across SKUs, so the - # first jobs under this cap spread over pools instead of queuing on one. - max-parallel: 10 + # Global cap; real throttling is each SKU pool's runner count and its cluster's + # Slurm partition. setup interleaves shards across SKUs so the cap spreads over + # pools instead of queuing on one. Kept above the total per-pool runner count so a + # saturated pool's dispatched-but-waiting legs cannot starve pools that have free + # runners (10 was too low when the shared fleet is busy with concurrent sweeps). + max-parallel: 30 matrix: ${{ fromJSON(needs.setup.outputs.matrix) }} runs-on: ${{ matrix.sku }} timeout-minutes: 350 From 564067542d76b75b1cae366c66b29dfecefb04c7 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Tue, 21 Jul 2026 01:20:10 +0800 Subject: [PATCH 06/10] fix(collectivex): install uccl-ep deep_ep wrapper with --no-deps (ROCm 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. --- experimental/CollectiveX/launchers/launch_mi-tw.sh | 4 +++- experimental/CollectiveX/runtime/prepare_backend.sh | 9 ++++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/experimental/CollectiveX/launchers/launch_mi-tw.sh b/experimental/CollectiveX/launchers/launch_mi-tw.sh index cfac84d3f6..8d8d7e8366 100644 --- a/experimental/CollectiveX/launchers/launch_mi-tw.sh +++ b/experimental/CollectiveX/launchers/launch_mi-tw.sh @@ -93,7 +93,9 @@ PY # handles to pinned host memory (coherent + device-accessible on CDNA). sed -i "s/cudaMallocManaged/cudaMallocHost/g" /tmp/ub/ep/src/uccl_ep.cc /tmp/ub/ep/src/uccl_proxy.cpp cd /tmp/ub/ep && env USE_DMABUF=1 PER_EXPERT_BATCHING=1 PYTORCH_ROCM_ARCH="'"$UCCL_ARCH"'" python3 setup.py install >&2 - cd /tmp/ub/ep/deep_ep_wrapper && python3 setup.py install >&2 + # --no-deps: the wrapper install_requires=["uccl"] pulls the PyPI uccl->uccl-cu12 wheel + # (absent on ROCm); our from-source ep build already provides uccl.ep in site-packages. + cd /tmp/ub/ep/deep_ep_wrapper && { pip install -q --no-deps . || pip install -q --no-deps --break-system-packages . ; } >&2 SP="$(python3 -c "import site;print(site.getsitepackages()[0])")" mkdir -p /cx/.collx_uccl_pfx && cp -R "$SP"/deep_ep* "$SP"/uccl* /cx/.collx_uccl_pfx/ python3 -c "import torch,sys; sys.path.insert(0,\"/cx/.collx_uccl_pfx\"); import deep_ep; from deep_ep import Buffer; assert hasattr(Buffer,\"get_dispatch_layout\")" >&2 diff --git a/experimental/CollectiveX/runtime/prepare_backend.sh b/experimental/CollectiveX/runtime/prepare_backend.sh index 2647b44113..9975dddd5b 100644 --- a/experimental/CollectiveX/runtime/prepare_backend.sh +++ b/experimental/CollectiveX/runtime/prepare_backend.sh @@ -314,7 +314,14 @@ uccl_prepare() { ( cd "$source_dir/ep" \ && env USE_DMABUF=1 PER_EXPERT_BATCHING=1 "$arch_env" python3 setup.py install ) >&2 2>&1 \ || { collx_log "ERROR: UCCL ep extension build failed"; return 1; } - ( cd "$source_dir/ep/deep_ep_wrapper" && python3 setup.py install ) >&2 2>&1 \ + # Install the wrapper WITHOUT its deps: install_requires=["uccl"] resolves to the PyPI + # uccl metapackage, which depends on the prebuilt uccl-cu12 wheel — absent on ROCm (hard + # fail) and wrong even on CUDA, since our from-source ep build already provides uccl.ep in + # site-packages/uccl. --no-deps makes the source build authoritative on both vendors. + ( cd "$source_dir/ep/deep_ep_wrapper" \ + && { python3 -m pip install -q --disable-pip-version-check --no-input --no-deps . \ + || python3 -m pip install -q --disable-pip-version-check --no-input \ + --no-deps --break-system-packages . ; } ) >&2 2>&1 \ || { collx_log "ERROR: UCCL deep_ep_wrapper build failed"; return 1; } uccl_probe || { collx_log "ERROR: UCCL import probe failed"; return 1; } collx_log "UCCL-EP ready ($COLLX_UCCL_COMMIT, deep_ep wrapper over uccl.ep CPU-proxy runtime)" From e4094b7fb6052fe30f0850aa4e9a2912521fa818 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Tue, 21 Jul 2026 01:58:58 +0800 Subject: [PATCH 07/10] fix(collectivex): -tw uccl-ep persist prefix to node-local /tmp (cleanup 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. --- .../CollectiveX/launchers/launch_mi-tw.sh | 29 +++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/experimental/CollectiveX/launchers/launch_mi-tw.sh b/experimental/CollectiveX/launchers/launch_mi-tw.sh index 8d8d7e8366..3f9af87ba2 100644 --- a/experimental/CollectiveX/launchers/launch_mi-tw.sh +++ b/experimental/CollectiveX/launchers/launch_mi-tw.sh @@ -64,13 +64,16 @@ fi collx_log "runner=$RUNNER nodes=1 x ${GPN}gpu world=$NGPUS bench=$COLLX_BENCH image=$IMAGE (${DOCKER[*]}/torchrun)" -# ---- uccl-ep: prepare source + one-time persisted build --------------------- -# UCCL is not in the image, so build it from source. Cases each run in a throwaway -# `docker run --rm`, so build ONCE here into a host-persisted prefix -# ($COLLX_DIR/.collx_uccl_pfx) that every case container puts on PYTHONPATH — the same -# build-once/reuse the enroot paths get from prepare_backend.sh, adapted to Docker. The -# AMD build applies the CDNA managed->pinned-host-memory patch (see prepare_backend.sh). -UCCL_PFX_HOST="$COLLX_DIR/.collx_uccl_pfx" +# ---- uccl-ep: prepare source + build (persisted node-local, outside the job root) -------- +# UCCL is not in the image, so build from source. Cases each run in a throwaway `docker run +# --rm`, so build ONCE into a prefix every case container puts on PYTHONPATH. The prefix lives +# in a NODE-LOCAL /tmp path keyed on the pinned commit, NOT under the isolated job root: the +# build container writes it as root, and the workflow's cleanup step (non-root runner) cannot +# rm root-owned files under the job root (that reds an otherwise-green leg). /tmp is outside +# cleanup's scope and lets a second leg on the same node reuse the build. The AMD build applies +# the CDNA managed->pinned-host-memory patch (see prepare_backend.sh). +UCCL_PFX_HOST="/tmp/collx-uccl-pfx-$COLLX_UCCL_COMMIT" +UCCL_PFX_MOUNT=() if [ "$COLLX_BENCH" = uccl-ep ]; then REPO_ROOT="$(cd "$COLLX_DIR/../.." && pwd)" collx_prepare_uccl_source "$REPO_ROOT" || collx_die "UCCL source preparation failed" @@ -79,12 +82,14 @@ import json, sys print(json.load(open(sys.argv[1]))["platforms"][sys.argv[2]]["arch"]) PY )" + mkdir -p "$UCCL_PFX_HOST" + UCCL_PFX_MOUNT=(-v "$UCCL_PFX_HOST:/uccl_pfx") if [ ! -d "$UCCL_PFX_HOST/deep_ep" ]; then collx_log "uccl-ep: one-time from-source build (arch=$UCCL_ARCH, USE_DMABUF, host-atomic path)" "${DOCKER[@]}" run --rm \ --device /dev/kfd --device /dev/dri --group-add video --group-add render \ --ipc host --shm-size 32g --cap-add SYS_PTRACE --security-opt seccomp=unconfined \ - -v "$COLLX_DIR:/cx" -w /cx "$IMAGE" \ + -v "$COLLX_DIR:/cx" -v "$UCCL_PFX_HOST:/uccl_pfx" -w /cx "$IMAGE" \ bash -c ' set -e { pip install -q nanobind || pip install -q --break-system-packages nanobind; } >&2 @@ -97,8 +102,8 @@ PY # (absent on ROCm); our from-source ep build already provides uccl.ep in site-packages. cd /tmp/ub/ep/deep_ep_wrapper && { pip install -q --no-deps . || pip install -q --no-deps --break-system-packages . ; } >&2 SP="$(python3 -c "import site;print(site.getsitepackages()[0])")" - mkdir -p /cx/.collx_uccl_pfx && cp -R "$SP"/deep_ep* "$SP"/uccl* /cx/.collx_uccl_pfx/ - python3 -c "import torch,sys; sys.path.insert(0,\"/cx/.collx_uccl_pfx\"); import deep_ep; from deep_ep import Buffer; assert hasattr(Buffer,\"get_dispatch_layout\")" >&2 + rm -rf /uccl_pfx/* && cp -R "$SP"/deep_ep* "$SP"/uccl* /uccl_pfx/ + python3 -c "import torch,sys; sys.path.insert(0,\"/uccl_pfx\"); import deep_ep; from deep_ep import Buffer; assert hasattr(Buffer,\"get_dispatch_layout\")" >&2 ' >&2 \ || collx_die "uccl-ep from-source build failed" collx_log "uccl-ep: build persisted to $UCCL_PFX_HOST" @@ -124,7 +129,7 @@ if [ "$COLLX_BENCH" = uccl-ep ]; then # uccl-ep imports the host-persisted build via PYTHONPATH; CDNA needs the aggressive # host-atomic EP path (matches prepare_backend.sh's uccl_prepare AMD branch). docker_env=( - -e PYTHONPATH=/cx/.collx_uccl_pfx + -e PYTHONPATH=/uccl_pfx -e UCCL_EP_ENABLE_AGGRESSIVE_ATOMIC="${UCCL_EP_ENABLE_AGGRESSIVE_ATOMIC:-1}" -e HSA_NO_SCRATCH_RECLAIM=1 -e COLLECTIVEX_SOURCE_SHA="${COLLECTIVEX_SOURCE_SHA:-}" @@ -169,7 +174,7 @@ for ((ci = 0; ci < ncases; ci++)); do --cap-add SYS_PTRACE --security-opt seccomp=unconfined \ --network host \ "${docker_env[@]}" \ - -v "$COLLX_DIR:/cx" -v "$argv_file:/cx-argv:ro" -w /cx \ + -v "$COLLX_DIR:/cx" -v "$argv_file:/cx-argv:ro" ${UCCL_PFX_MOUNT[@]+"${UCCL_PFX_MOUNT[@]}"} -w /cx \ "$IMAGE" \ bash -c 'xargs -0 torchrun --standalone --nproc-per-node='"$NGPUS"' bench/run_ep.py < /cx-argv'; then case_ok=1; break From b24698834d84c36871eeeb0f4f4b483fa4788a7d Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Tue, 21 Jul 2026 12:59:24 +0800 Subject: [PATCH 08/10] =?UTF-8?q?fix(collectivex):=20address=20uccl-ep=20r?= =?UTF-8?q?eview=20=E2=80=94=20AMD=20SM=20profile,=20-tw=20cache=20keying,?= =?UTF-8?q?=20UCCL=20GID=20env,=20FP8=20LL=20docs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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. --- experimental/CollectiveX/README.md | 2 +- experimental/CollectiveX/bench/ep_uccl.py | 24 +++++-- .../CollectiveX/launchers/launch_mi-tw.sh | 67 ++++++++++++------- experimental/CollectiveX/runtime/common.sh | 14 ++-- 4 files changed, 71 insertions(+), 36 deletions(-) diff --git a/experimental/CollectiveX/README.md b/experimental/CollectiveX/README.md index 77504b5a8c..0048c1fe0c 100644 --- a/experimental/CollectiveX/README.md +++ b/experimental/CollectiveX/README.md @@ -66,7 +66,7 @@ scale-up domain. |---|---| | DeepEP V2 | `normal` mode is PR #605 `ElasticBuffer` plus exact upstream #630 and #640 fixes: LSA for scale-up and GIN for x86 EP16 scale-out. FP8 dispatch via `use_fp8_dispatch` (blockwise e4m3fn) alongside BF16. `low-latency` mode is the legacy `deep_ep.Buffer` IBGDA decode kernels (per-expert padded layout, weighted combine, `use_fp8` e4m3fn), decode/EP8 only | | MoRI | `normal` mode uses the direct `IntraNode` kernel for scale-up EP8 on every CDNA SKU and pins `InterNodeV1` for EP16 over 2x8 XGMI + RDMA. `low-latency` mode selects the `IntraNodeLL` decode kernel (single-call, pure-intranode, same compact layout and unweighted combine as `IntraNode`), decode/EP8 only. FP8 dispatch is caller-prequantized (per-SKU e4m3fnuz on gfx942, e4m3fn on gfx950); combine stays BF16 (`quant_type=none`) alongside BF16 dispatch | -| UCCL-EP | [UCCL](https://github.com/uccl-project/uccl) EP: a drop-in, API-identical DeepEP replacement whose CPU proxies issue GPUDirect RDMA over plain `libibverbs` (no NVSHMEM/IBGDA), with software message ordering, atomics, and flow control; scale-up is single-node `cudaIpc` over NVLink/XGMI (never MNNVL). `normal` mode is the legacy `Buffer` `dispatch`/`combine` (unweighted rank-sum); `low-latency` reuses the legacy `low_latency_dispatch`/`low_latency_combine` decode kernels (weighted combine), decode/EP8 only. FP8 dispatch is caller-prequantized (blockwise e4m3fn, per-SKU e4m3fnuz on gfx942) with BF16 combine. Runs on NVIDIA and AMD (H100/H200/B200 + MI300X/MI325X/MI355X), EP8 scale-up. Cross-node EP16 is functional (the internode RDMA path connects and the light case passes correctness) but its CPU-proxy throughput overruns the standardized per-case wall-clock budget on heavy token counts, so EP16 is an unsupported coverage row for now | +| UCCL-EP | [UCCL](https://github.com/uccl-project/uccl) EP: a drop-in, API-identical DeepEP replacement whose CPU proxies issue GPUDirect RDMA over plain `libibverbs` (no NVSHMEM/IBGDA), with software message ordering, atomics, and flow control; scale-up is single-node `cudaIpc` over NVLink/XGMI (never MNNVL). `normal` mode is the legacy `Buffer` `dispatch`/`combine` (unweighted rank-sum); `low-latency` reuses the legacy `low_latency_dispatch`/`low_latency_combine` decode kernels (weighted combine), decode/EP8 only. FP8 dispatch is caller-prequantized in `normal` mode (blockwise e4m3fn, per-SKU e4m3fnuz on gfx942); in `low-latency` mode the caller sends BF16 and the decode kernel quantizes to e4m3 internally (`use_fp8`). Combine is BF16. Runs on NVIDIA and AMD (H100/H200/B200 + MI300X/MI325X/MI355X), EP8 scale-up. Cross-node EP16 is functional (the internode RDMA path connects and the light case passes correctness) but its CPU-proxy throughput overruns the standardized per-case wall-clock budget on heavy token counts, so EP16 is an unsupported coverage row for now | DeepEP V2 means the `ElasticBuffer` implementation introduced by [DeepEP PR #605](https://github.com/deepseek-ai/DeepEP/pull/605), not a newer legacy `Buffer` build. diff --git a/experimental/CollectiveX/bench/ep_uccl.py b/experimental/CollectiveX/bench/ep_uccl.py index 4dc352d549..0acb3335ce 100644 --- a/experimental/CollectiveX/bench/ep_uccl.py +++ b/experimental/CollectiveX/bench/ep_uccl.py @@ -17,9 +17,11 @@ low-latency -> low_latency_dispatch/low_latency_combine; per-expert padded recv, source-side weighted-kernel-sum combine. -FP8 dispatch is caller-prequantized (blockwise e4m3fn, e4m3fnuz on gfx942), combine is BF16 — the -oracle applies the identical per-token cast round-trip via semantic_payload/oracle_x, so the tight -combine gate (COMBINE_REL_TOL = 8*2^-8) is preserved, not loosened. +FP8 dispatch is caller-prequantized in normal mode (blockwise e4m3fn, e4m3fnuz on gfx942); in +low-latency mode the caller sends BF16 and the decode kernel quantizes to e4m3 internally +(``use_fp8``). Combine is always BF16 — the oracle applies the identical per-token cast round-trip +via semantic_payload/oracle_x in both modes, so the tight combine gate (COMBINE_REL_TOL = 8*2^-8) +is preserved, not loosened. """ from __future__ import annotations @@ -96,9 +98,8 @@ def _ll_dequant_static(fp8, scales): # Normal-mode legacy Config launch parameters (DeepEP-legacy Config(num_sms, chunk, nvl_buffer)). -# Bring-up-tunable; these mirror UCCL's own intranode bench (nvl_buffer_size=256) and the DeepEP -# normal default SM budget. num_nvl_bytes is a generous fixed reservation as in UCCL's bench. -_NORMAL_NUM_SMS = 24 +# These mirror UCCL's own intranode bench (nvl_buffer_size=256); num_nvl_bytes is a generous fixed +# reservation as in that bench. The SM budget is vendor-keyed (see _normal_num_sms). _NORMAL_NVL_BUFFER_SIZE = 256 _NORMAL_NVL_BYTES = int(2e9) # Internode (EP16) buffer-sizing Config, straight from UCCL's own test_internode bench @@ -112,6 +113,14 @@ def _align_buffer_bytes(size, margin=1.2, alignment=128): return ((int(size * margin) + alignment - 1) // alignment) * alignment +def _normal_num_sms() -> int: + """Intranode normal-mode SM budget, keyed by vendor exactly as UCCL's own bench does + (ep/bench/test_intranode.py: ``num_sms = 24 if torch.version.cuda else 64``): 24 on CUDA, + 64 on HIP/ROCm. A flat 24 would materially understate AMD, whose wider CU count wants the + larger grid — the same reason upstream branches on the vendor.""" + return 24 if torch.version.cuda else 64 + + class UCCLEPBackend(EPBackend): name = "uccl-ep" # One legacy Buffer under two modes, selected by args.mode: @@ -200,7 +209,8 @@ def create_buffer(self, spec): return # Intranode (EP8) scale-up: validated recipe — one fixed ~2 GB NVLink buffer, no RDMA, a # single QP, and the legacy 3-arg Config (rdma-chunked params unused with no RDMA path). - self.config = Config(_NORMAL_NUM_SMS, 8, _NORMAL_NVL_BUFFER_SIZE) + # SM budget is vendor-keyed (24 CUDA / 64 HIP), matching UCCL's intranode bench. + self.config = Config(_normal_num_sms(), 8, _NORMAL_NVL_BUFFER_SIZE) self.dispatch_config = self.config self.combine_config = self.config self.buffer = Buffer( diff --git a/experimental/CollectiveX/launchers/launch_mi-tw.sh b/experimental/CollectiveX/launchers/launch_mi-tw.sh index 3f9af87ba2..089ebb850e 100644 --- a/experimental/CollectiveX/launchers/launch_mi-tw.sh +++ b/experimental/CollectiveX/launchers/launch_mi-tw.sh @@ -72,7 +72,6 @@ collx_log "runner=$RUNNER nodes=1 x ${GPN}gpu world=$NGPUS bench=$COLLX_BENCH im # rm root-owned files under the job root (that reds an otherwise-green leg). /tmp is outside # cleanup's scope and lets a second leg on the same node reuse the build. The AMD build applies # the CDNA managed->pinned-host-memory patch (see prepare_backend.sh). -UCCL_PFX_HOST="/tmp/collx-uccl-pfx-$COLLX_UCCL_COMMIT" UCCL_PFX_MOUNT=() if [ "$COLLX_BENCH" = uccl-ep ]; then REPO_ROOT="$(cd "$COLLX_DIR/../.." && pwd)" @@ -82,30 +81,50 @@ import json, sys print(json.load(open(sys.argv[1]))["platforms"][sys.argv[2]]["arch"]) PY )" - mkdir -p "$UCCL_PFX_HOST" + # Cache key = pinned UCCL commit + the image's CONTENT id + GPU arch. Torch/ROCm are baked into + # the image, so its content id (not the mutable tag) captures an ABI change even under a + # re-pushed tag, and the arch captures a cross-SKU reuse hazard. Keying on the commit alone (as + # before) would let a bumped image or a different arch silently reuse an ABI-stale build. + UCCL_IMAGE_ID="$("${DOCKER[@]}" image inspect --format '{{.Id}}' "$IMAGE" 2>/dev/null \ + || printf '%s' "$IMAGE")" + UCCL_CACHE_KEY="$(printf '%s\0%s\0%s' "$COLLX_UCCL_COMMIT" "$UCCL_IMAGE_ID" "$UCCL_ARCH" \ + | sha1sum | cut -c1-16)" + UCCL_PFX_HOST="/tmp/collx-uccl-pfx-$UCCL_CACHE_KEY" UCCL_PFX_MOUNT=(-v "$UCCL_PFX_HOST:/uccl_pfx") - if [ ! -d "$UCCL_PFX_HOST/deep_ep" ]; then - collx_log "uccl-ep: one-time from-source build (arch=$UCCL_ARCH, USE_DMABUF, host-atomic path)" - "${DOCKER[@]}" run --rm \ - --device /dev/kfd --device /dev/dri --group-add video --group-add render \ - --ipc host --shm-size 32g --cap-add SYS_PTRACE --security-opt seccomp=unconfined \ - -v "$COLLX_DIR:/cx" -v "$UCCL_PFX_HOST:/uccl_pfx" -w /cx "$IMAGE" \ - bash -c ' - set -e - { pip install -q nanobind || pip install -q --break-system-packages nanobind; } >&2 - rm -rf /tmp/ub && cp -R "/cx/.collx_sources/uccl-'"$COLLX_UCCL_COMMIT"'" /tmp/ub - # gfx942/gfx950 lack usable managed memory; swap UCCL'"'"'s cudaMallocManaged CPU-proxy - # handles to pinned host memory (coherent + device-accessible on CDNA). - sed -i "s/cudaMallocManaged/cudaMallocHost/g" /tmp/ub/ep/src/uccl_ep.cc /tmp/ub/ep/src/uccl_proxy.cpp - cd /tmp/ub/ep && env USE_DMABUF=1 PER_EXPERT_BATCHING=1 PYTORCH_ROCM_ARCH="'"$UCCL_ARCH"'" python3 setup.py install >&2 - # --no-deps: the wrapper install_requires=["uccl"] pulls the PyPI uccl->uccl-cu12 wheel - # (absent on ROCm); our from-source ep build already provides uccl.ep in site-packages. - cd /tmp/ub/ep/deep_ep_wrapper && { pip install -q --no-deps . || pip install -q --no-deps --break-system-packages . ; } >&2 - SP="$(python3 -c "import site;print(site.getsitepackages()[0])")" - rm -rf /uccl_pfx/* && cp -R "$SP"/deep_ep* "$SP"/uccl* /uccl_pfx/ - python3 -c "import torch,sys; sys.path.insert(0,\"/uccl_pfx\"); import deep_ep; from deep_ep import Buffer; assert hasattr(Buffer,\"get_dispatch_layout\")" >&2 - ' >&2 \ - || collx_die "uccl-ep from-source build failed" + # Readiness is a `.ready` marker written LAST (only after the in-container import verification + # passes), never the mere existence of deep_ep/: an interrupted copy or a failed import must not + # leave a half-populated cache a later job reuses blind. Build into a private temp dir, then + # publish atomically with `mv -T` (a concurrent leg that loses the rename just drops its temp). + if [ ! -f "$UCCL_PFX_HOST/.ready" ]; then + collx_log "uccl-ep: one-time from-source build (arch=$UCCL_ARCH, key=$UCCL_CACHE_KEY, USE_DMABUF, host-atomic path)" + rm -rf "$UCCL_PFX_HOST" # clear any partial/aborted prior attempt (no .ready) + uccl_build_tmp="$(mktemp -d /tmp/collx-uccl-pfx.XXXXXX)" || collx_die "uccl-ep: mktemp failed" + if "${DOCKER[@]}" run --rm \ + --device /dev/kfd --device /dev/dri --group-add video --group-add render \ + --ipc host --shm-size 32g --cap-add SYS_PTRACE --security-opt seccomp=unconfined \ + -v "$COLLX_DIR:/cx" -v "$uccl_build_tmp:/uccl_pfx" -w /cx "$IMAGE" \ + bash -c ' + set -e + { pip install -q nanobind || pip install -q --break-system-packages nanobind; } >&2 + rm -rf /tmp/ub && cp -R "/cx/.collx_sources/uccl-'"$COLLX_UCCL_COMMIT"'" /tmp/ub + # gfx942/gfx950 lack usable managed memory; swap UCCL'"'"'s cudaMallocManaged CPU-proxy + # handles to pinned host memory (coherent + device-accessible on CDNA). + sed -i "s/cudaMallocManaged/cudaMallocHost/g" /tmp/ub/ep/src/uccl_ep.cc /tmp/ub/ep/src/uccl_proxy.cpp + cd /tmp/ub/ep && env USE_DMABUF=1 PER_EXPERT_BATCHING=1 PYTORCH_ROCM_ARCH="'"$UCCL_ARCH"'" python3 setup.py install >&2 + # --no-deps: the wrapper install_requires=["uccl"] pulls the PyPI uccl->uccl-cu12 wheel + # (absent on ROCm); our from-source ep build already provides uccl.ep in site-packages. + cd /tmp/ub/ep/deep_ep_wrapper && { pip install -q --no-deps . || pip install -q --no-deps --break-system-packages . ; } >&2 + SP="$(python3 -c "import site;print(site.getsitepackages()[0])")" + rm -rf /uccl_pfx/* && cp -R "$SP"/deep_ep* "$SP"/uccl* /uccl_pfx/ + python3 -c "import torch,sys; sys.path.insert(0,\"/uccl_pfx\"); import deep_ep; from deep_ep import Buffer; assert hasattr(Buffer,\"get_dispatch_layout\")" >&2 + touch /uccl_pfx/.ready # publish gate: written only after the import check succeeds + ' >&2; then + mv -T "$uccl_build_tmp" "$UCCL_PFX_HOST" 2>/dev/null || rm -rf "$uccl_build_tmp" + else + rm -rf "$uccl_build_tmp" + collx_die "uccl-ep from-source build failed" + fi + [ -f "$UCCL_PFX_HOST/.ready" ] || collx_die "uccl-ep: build did not publish a ready cache" collx_log "uccl-ep: build persisted to $UCCL_PFX_HOST" else collx_log "uccl-ep: reusing persisted build at $UCCL_PFX_HOST" diff --git a/experimental/CollectiveX/runtime/common.sh b/experimental/CollectiveX/runtime/common.sh index 706ffaeb65..26ed30bd89 100644 --- a/experimental/CollectiveX/runtime/common.sh +++ b/experimental/CollectiveX/runtime/common.sh @@ -134,12 +134,15 @@ collx_require_vars() { collx_export_gid_index_for_link_layer() { local link_layer="$1" - unset NVSHMEM_IB_GID_INDEX NCCL_IB_GID_INDEX + unset NVSHMEM_IB_GID_INDEX NCCL_IB_GID_INDEX UCCL_IB_GID_INDEX [ -n "${COLLX_IB_GID_INDEX:-}" ] || return 0 case "$link_layer" in roce) export NVSHMEM_IB_GID_INDEX="$COLLX_IB_GID_INDEX" export NCCL_IB_GID_INDEX="$COLLX_IB_GID_INDEX" + # UCCL-EP reads only its own UCCL_IB_GID_INDEX (it does NOT consult NCCL_IB_GID_INDEX), so + # RoCE runs must set it here or the CPU proxies fall back to GID 0 and mis-address the fabric. + export UCCL_IB_GID_INDEX="$COLLX_IB_GID_INDEX" ;; infiniband) ;; *) collx_die "unsupported RDMA link layer" ;; @@ -196,9 +199,12 @@ collx_apply_network_profile() { fi export NCCL_IB_HCA="=$COLLX_RDMA_DEVICES" export MORI_RDMA_DEVICES="$rdma_names" EP_NIC_NAME="$ep_nic" - # UCCL-EP's CPU proxies read UCCL_* selectors but fall back to NCCL_IB_HCA / NCCL_IB_GID_INDEX - # (set above / by the link-layer helper) for the device list and GID, so only the socket iface - # and, on AMD, the strict Pollara/Broadcom flow control + CDNA host-atomic path need setting. + # UCCL-EP's CPU proxies read ONLY the UCCL_* selectors; they do NOT consult NCCL_IB_HCA / + # NCCL_IB_GID_INDEX. So the device list and GID must come through UCCL's own vars: the HCA list + # here, and the GID index via collx_export_gid_index_for_link_layer (RoCE, called at the end of + # this function). The socket iface and, on AMD, the strict Pollara/Broadcom flow control + CDNA + # host-atomic path are likewise UCCL_* vars. + export UCCL_IB_HCA="$rdma_names" export UCCL_SOCKET_IFNAME="${COLLX_SOCKET_IFNAME:-}" if [ "${COLLX_VENDOR:-nvidia}" = amd ]; then export UCCL_IB_MAX_INFLIGHT_BYTES="${UCCL_IB_MAX_INFLIGHT_BYTES:-2097152}" From 51b3a6f5e79e3ad2d50933f0e7311e76bed2c0c3 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Tue, 21 Jul 2026 14:24:57 +0800 Subject: [PATCH 09/10] fix(collectivex): correct uccl-ep HCA selector + FP8-mode docs; drop stale gitignore MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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. --- experimental/CollectiveX/.gitignore | 1 - experimental/CollectiveX/README.md | 5 +++-- experimental/CollectiveX/docs/methodology.md | 7 +++++-- experimental/CollectiveX/runtime/common.sh | 13 +++++++------ 4 files changed, 15 insertions(+), 11 deletions(-) diff --git a/experimental/CollectiveX/.gitignore b/experimental/CollectiveX/.gitignore index f68f23df47..795c9f1a3f 100644 --- a/experimental/CollectiveX/.gitignore +++ b/experimental/CollectiveX/.gitignore @@ -6,7 +6,6 @@ unsupported/ .collx_workloads/ .collx_backend/ .collx_sources/ -.collx_uccl_pfx/ /matrix_full.json gpucore.* diff --git a/experimental/CollectiveX/README.md b/experimental/CollectiveX/README.md index 0048c1fe0c..ea87b89e89 100644 --- a/experimental/CollectiveX/README.md +++ b/experimental/CollectiveX/README.md @@ -13,8 +13,9 @@ responsibility. The full measurement methodology is in [docs/methodology.md](doc The workload uses packed placement and one pinned `fixed-profile` resource configuration per backend/topology; there is no tuning sweep. Combine is always BF16; dispatch precision is a swept -dimension — a BF16 control plus a caller-prequantized FP8 dispatch on every backend. Coverage is -uniform routing only. Cases run in one of two modes: +dimension — a BF16 control plus an FP8 dispatch on every backend, caller-prequantized in `normal` +mode (in `low-latency` the DeepEP and UCCL-EP kernels quantize internally from BF16; MoRI stays +caller-prequantized). Coverage is uniform routing only. Cases run in one of two modes: - `normal` uses `layout-and-dispatch-v1`, rank-deduplicated token payloads, and activation-only, unweighted rank-sum combine. It runs the full decode and prefill ladders. diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index d398ff02dc..332b5cef03 100644 --- a/experimental/CollectiveX/docs/methodology.md +++ b/experimental/CollectiveX/docs/methodology.md @@ -21,8 +21,11 @@ It does not predict serving throughput without a separate correlation study. The implemented workload is `deepseek-v3`: hidden 7168, top-k 8, 256 routed experts, packed placement, and one pinned fixed resource profile per backend/topology. Combine is always BF16; -dispatch precision is a swept dimension — a BF16 control and a caller-prequantized FP8 dispatch -(`bf16`, `fp8`). Every case uses the normal `layout-and-dispatch-v1` semantics. +dispatch precision is a swept dimension — a BF16 control and an FP8 dispatch (`bf16`, `fp8`), +caller-prequantized in `normal` mode (the `low-latency` kernels quantize FP8 internally from BF16 on +DeepEP and UCCL-EP, and stay caller-prequantized on MoRI). `normal`-mode cases use the +`layout-and-dispatch-v1` semantics; `low-latency` cases use each backend's decode-kernel semantics +(detailed below). - `ep-core`: uniform routing over the workload's token ladders — for `deepseek-v3`, decode T=1..512 powers of two and prefill T=1024..8192 powers of two. Ladders are model-specific and diff --git a/experimental/CollectiveX/runtime/common.sh b/experimental/CollectiveX/runtime/common.sh index 26ed30bd89..c65bee18fe 100644 --- a/experimental/CollectiveX/runtime/common.sh +++ b/experimental/CollectiveX/runtime/common.sh @@ -199,12 +199,13 @@ collx_apply_network_profile() { fi export NCCL_IB_HCA="=$COLLX_RDMA_DEVICES" export MORI_RDMA_DEVICES="$rdma_names" EP_NIC_NAME="$ep_nic" - # UCCL-EP's CPU proxies read ONLY the UCCL_* selectors; they do NOT consult NCCL_IB_HCA / - # NCCL_IB_GID_INDEX. So the device list and GID must come through UCCL's own vars: the HCA list - # here, and the GID index via collx_export_gid_index_for_link_layer (RoCE, called at the end of - # this function). The socket iface and, on AMD, the strict Pollara/Broadcom flow control + CDNA - # host-atomic path are likewise UCCL_* vars. - export UCCL_IB_HCA="$rdma_names" + # UCCL-EP's EP transport reads UCCL_IB_HCA and falls back to NCCL_IB_HCA (ep/src/rdma.cpp), and + # its filter honors the same leading '=' exact-match and ':port' syntax as NCCL. So mirror the + # exact-match selector already set on NCCL_IB_HCA above — a bare name list would prefix-match + # (mlx5_1 -> mlx5_1,mlx5_10..19) and drop the port. The GID index, by contrast, has NO NCCL + # fallback in UCCL's EP path (it reads only UCCL_IB_GID_INDEX, ep/include/rdma_util.hpp), so + # collx_export_gid_index_for_link_layer must set that UCCL_* var explicitly for RoCE. + export UCCL_IB_HCA="=$COLLX_RDMA_DEVICES" export UCCL_SOCKET_IFNAME="${COLLX_SOCKET_IFNAME:-}" if [ "${COLLX_VENDOR:-nvidia}" = amd ]; then export UCCL_IB_MAX_INFLIGHT_BYTES="${UCCL_IB_MAX_INFLIGHT_BYTES:-2097152}" From c8df0dbafd838950a9aa412e44f5f12789c86430 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Tue, 21 Jul 2026 15:03:45 +0800 Subject: [PATCH 10/10] =?UTF-8?q?fix(collectivex):=20resolve=20uccl-ep=20r?= =?UTF-8?q?eview=20threads=20=E2=80=94=20LL=20guard=20+=20cross-allocation?= =?UTF-8?q?=20build=20cache?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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. --- experimental/CollectiveX/bench/ep_uccl.py | 5 +- .../CollectiveX/launchers/launch_mi-amds.sh | 9 ++ .../CollectiveX/runtime/prepare_backend.sh | 97 +++++++++++++++++-- 3 files changed, 99 insertions(+), 12 deletions(-) diff --git a/experimental/CollectiveX/bench/ep_uccl.py b/experimental/CollectiveX/bench/ep_uccl.py index 0acb3335ce..821b23c25d 100644 --- a/experimental/CollectiveX/bench/ep_uccl.py +++ b/experimental/CollectiveX/bench/ep_uccl.py @@ -238,11 +238,12 @@ def _create_ll_buffer(self, spec): self.num_local_experts = args.experts // world_size # LL requires the QP-per-rank count to equal the number of local experts. num_qps_per_rank = self.num_local_experts - if not hasattr(Buffer, "low_latency_dispatch") and not hasattr( + if not hasattr(Buffer, "low_latency_dispatch") or not hasattr( Buffer, "get_low_latency_rdma_size_hint" ): raise RuntimeError( - "invalid UCCL-EP LL runtime: Buffer.low_latency_dispatch is absent" + "invalid UCCL-EP LL runtime: Buffer.low_latency_dispatch / " + "get_low_latency_rdma_size_hint absent" ) num_rdma_bytes = Buffer.get_low_latency_rdma_size_hint( self.max_tokens, args.hidden, world_size, args.experts diff --git a/experimental/CollectiveX/launchers/launch_mi-amds.sh b/experimental/CollectiveX/launchers/launch_mi-amds.sh index ea4bc6327b..9f0f2a137b 100644 --- a/experimental/CollectiveX/launchers/launch_mi-amds.sh +++ b/experimental/CollectiveX/launchers/launch_mi-amds.sh @@ -123,6 +123,15 @@ for allocation_attempt in 1 2 3; do done unset COLLX_SALLOC_ATTEMPT COLLX_NETWORK_VALIDATION_ATTEMPT CONTAINER_MOUNTS="$MOUNT_SRC:$MOUNT_DIR$DEVICE_MOUNTS" +# uccl-ep builds from source, so give it the same cross-allocation backend cache the single-slurm +# launcher provides (built once per arch/image/commit under /cx-cache, reused each allocation). +# mori ships in the image and needs no cache, so its mounts are left untouched. +if [ "$COLLX_BENCH" = uccl-ep ]; then + collx_prepare_backend_cache "$SQUASH_DIR" \ + || collx_die "cannot prepare the isolated backend cache" + CONTAINER_MOUNTS="$CONTAINER_MOUNTS,$COLLX_PREPARED_BACKEND_CACHE:/cx-cache" + export COLLX_BACKEND_CACHE_ROOT=/cx-cache +fi # ---- container-launch -> artifact-collection (shared tail) ------------------ COLLX_DISTRIBUTED_CONTAINER_ARGS=(--container-writable --container-remap-root) diff --git a/experimental/CollectiveX/runtime/prepare_backend.sh b/experimental/CollectiveX/runtime/prepare_backend.sh index 9975dddd5b..5ebf6b22a9 100644 --- a/experimental/CollectiveX/runtime/prepare_backend.sh +++ b/experimental/CollectiveX/runtime/prepare_backend.sh @@ -279,12 +279,46 @@ PY # Direct in-container source build against the image's torch — validated on h200 (sglang # cu130). NOT `build.sh` (that spins up its own Docker image to make a wheel and cannot run -# inside enroot/pyxis). Installs into the image's system python; single-slurm and mi-amds run -# the writable container as remapped root, so this needs no venv. verbs/nl/numa dev headers -# ship in the sglang/rocm images; only nanobind must be added. -uccl_prepare() { - local source_dir="/tmp/collectivex-uccl-$COLLX_UCCL_COMMIT" arch_env - command -v python3 >/dev/null || { collx_log "ERROR: python3 unavailable for UCCL build"; return 1; } +# inside enroot/pyxis). single-slurm and mi-amds run the writable container as remapped root, +# so the build needs no venv. verbs/nl/numa dev headers ship in the sglang/rocm images; only +# nanobind must be added. The built deep_ep/uccl packages are persisted under a cache root and +# put on PYTHONPATH (which write_rank_env carries to the ranks), so later allocations reuse them +# without recompiling — the same copy+PYTHONPATH scheme the mi-tw Docker launcher already uses. + +# Cache root keyed by cpu + build arch + image + pinned commit, under the shared /cx-cache mount +# ($COLLX_BACKEND_CACHE_ROOT). Returns non-zero when no shared cache is mounted (manual runs), so +# the caller falls back to a node-local build. Mirrors deepep_cache_root. +uccl_cache_root() { + local arch="$1" cpu base image + cpu="$(uname -m)" + [[ "$cpu" =~ ^[A-Za-z0-9._-]+$ ]] || return 1 + base="${COLLX_BACKEND_CACHE_ROOT:-}" + [[ "$base" = /* ]] || return 1 + image="$(printf '%s' "${COLLECTIVEX_IMAGE:-manual}" | tr -cs 'A-Za-z0-9_.-' '-')" + arch="$(printf '%s' "$arch" | tr -cs 'A-Za-z0-9_.-' '-')" + printf '%s/uccl-ep-%s-%s-%s-%s' \ + "$base" "$cpu" "${arch#-}" "${image#-}" "${COLLX_UCCL_COMMIT:0:12}" +} + +# Put the persisted build ($root/site) on PYTHONPATH for the probe and the rank tasks; mirror the +# minimal runtime bits of deepep_activate. CDNA additionally needs the aggressive host-atomic path. +uccl_activate() { + local site="$1/site" + [ -d "$site" ] || { collx_log "ERROR: UCCL cache site is unavailable"; return 1; } + export PYTHONPATH="$site${PYTHONPATH:+:$PYTHONPATH}" + [ "${COLLX_VENDOR:-nvidia}" != amd ] || export UCCL_EP_ENABLE_AGGRESSIVE_ATOMIC=1 +} + +# Build UCCL from source into $root/site (fresh root, with a .ready marker written LAST). The +# build installs into the image's system python as a sandbox, then copies the built deep_ep/uccl +# packages into the cache; the runtime imports them via PYTHONPATH (uccl_activate), so cache-hit +# and cache-miss paths import identically. Only nanobind is added to the image. +uccl_install() { + local root="$1" arch="$2" source_dir="/tmp/collectivex-uccl-$COLLX_UCCL_COMMIT" arch_env sp + if [ -e "$root" ] || [ -L "$root" ]; then + rm -rf "$root" || { collx_log "ERROR: incomplete UCCL cache-reset failed"; return 1; } + fi + mkdir -m 700 "$root" || { collx_log "ERROR: UCCL cache-create failed"; return 1; } collx_log "UCCL-EP: building $COLLX_UCCL_COMMIT from source (USE_DMABUF, PER_EXPERT_BATCHING)" # Plain install first; some sglang/rocm image variants mark the system env externally-managed # (PEP 668), so fall back to --break-system-packages (a no-op on older pip that lacks the flag). @@ -295,9 +329,7 @@ uccl_prepare() { collx_materialize_uccl_source "$source_dir" \ || { collx_log "ERROR: UCCL staged source is invalid"; return 1; } if [ "${COLLX_VENDOR:-nvidia}" = amd ]; then - arch_env="PYTORCH_ROCM_ARCH=$(uccl_rocm_arch)" - # CDNA requires the aggressive host-atomic EP path; persist it for the rank tasks too. - export UCCL_EP_ENABLE_AGGRESSIVE_ATOMIC=1 + arch_env="PYTORCH_ROCM_ARCH=$arch" # Managed/unified memory (cudaMallocManaged) is unavailable on our CDNA nodes (hipMallocManaged # fails even for 4 KiB, regardless of XNACK / --privileged / memlock). UCCL's HIP CPU-proxy path # uses it for the d2h channel handles + proxy atomic buffer; pinned host memory (cudaMallocHost) @@ -309,7 +341,7 @@ uccl_prepare() { "$source_dir/ep/src/uccl_ep.cc" "$source_dir/ep/src/uccl_proxy.cpp" \ || { collx_log "ERROR: UCCL AMD managed-memory patch failed"; return 1; } else - arch_env="TORCH_CUDA_ARCH_LIST=$(cuda_arch)" + arch_env="TORCH_CUDA_ARCH_LIST=$arch" fi ( cd "$source_dir/ep" \ && env USE_DMABUF=1 PER_EXPERT_BATCHING=1 "$arch_env" python3 setup.py install ) >&2 2>&1 \ @@ -323,6 +355,51 @@ uccl_prepare() { || python3 -m pip install -q --disable-pip-version-check --no-input \ --no-deps --break-system-packages . ; } ) >&2 2>&1 \ || { collx_log "ERROR: UCCL deep_ep_wrapper build failed"; return 1; } + sp="$(python3 -c 'import site; print(site.getsitepackages()[0])')" \ + || { collx_log "ERROR: UCCL site-packages resolution failed"; return 1; } + mkdir -p "$root/site" \ + && cp -R "$sp"/deep_ep* "$sp"/uccl* "$root/site/" \ + || { collx_log "ERROR: UCCL cache population failed"; return 1; } + : > "$root/.ready" +} + +# UCCL-EP lifecycle: build once per (arch, image, commit) into the shared /cx-cache behind an +# flock + .ready marker, reused on every later allocation (mirrors deepep_prepare); fall back to +# a node-local build when no shared cache is mounted (e.g. a manual run). +uccl_prepare() { + local arch root ready lock_path + command -v python3 >/dev/null || { collx_log "ERROR: python3 unavailable for UCCL build"; return 1; } + if [ "${COLLX_VENDOR:-nvidia}" = amd ]; then + arch="$(uccl_rocm_arch)" || return 1 + else + arch="$(cuda_arch)" || return 1 + fi + if root="$(uccl_cache_root "$arch")"; then + ready="$root/.ready"; lock_path="${root}.lock" + command -v flock >/dev/null \ + || { collx_log "ERROR: flock is required for UCCL-EP caching"; return 1; } + mkdir -p "${root%/*}" || return 1 + collx_log "UCCL-EP: preparing $COLLX_UCCL_COMMIT (shared cache $root)" + if ! ( + [ ! -L "$lock_path" ] || { collx_log "ERROR: UCCL cache lock is unsafe"; exit 1; } + (umask 077; : >> "$lock_path") && chmod 600 "$lock_path" \ + || { collx_log "ERROR: UCCL cache-lock-create failed"; exit 1; } + exec 9<>"$lock_path" || { collx_log "ERROR: UCCL cache-lock-open failed"; exit 1; } + flock 9 || { collx_log "ERROR: UCCL cache-lock-acquire failed"; exit 1; } + if [ ! -f "$ready" ] || [ ! -d "$root/site" ]; then + uccl_install "$root" "$arch" || exit 1 + fi + ); then + collx_log "ERROR: shared UCCL-EP environment is incomplete"; return 1 + fi + else + root="/tmp/collectivex-uccl-cache-$COLLX_UCCL_COMMIT" + collx_log "UCCL-EP: preparing $COLLX_UCCL_COMMIT (node-local $root; no shared cache mounted)" + if [ ! -f "$root/.ready" ] || [ ! -d "$root/site" ]; then + uccl_install "$root" "$arch" || return 1 + fi + fi + uccl_activate "$root" || return 1 uccl_probe || { collx_log "ERROR: UCCL import probe failed"; return 1; } collx_log "UCCL-EP ready ($COLLX_UCCL_COMMIT, deep_ep wrapper over uccl.ep CPU-proxy runtime)" }