From 77c3b6355624ee3e4f2af00d16c432aa7ca8a3e0 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Mon, 28 Sep 2026 22:52:42 +0800 Subject: [PATCH 1/4] CollectiveX: add the kv-transfer suite (NIXL, Mooncake, MoRI-IO KV-cache handoff) A leg is 2 nodes x 1 GPU moving bursts of concurrent requests' paged KV (vLLM's packed block-major DSV4 descriptors over random block tables), pull and push, against a single-descriptor bulk baseline, with every request pattern-verified. It resolves in sweep_matrix.py (--suites kv-transfer) from configs/kv_sweep.json and the registry's kv_backends map, runs on each pool's own launcher, and reaches run_kv through the suite codec. Per-pool allocation and hang-guard budgets live in kv_sweep.json and per-backend pool budgets in the registry, not in launcher branches. No b300 rows: the current b300 pool is the EFA cluster. --- .github/workflows/collectivex-sweep.yml | 10 +- collectivex/README.md | 28 +- collectivex/bandwidth.py | 4 + collectivex/bench/kv_backend.py | 97 ++++ collectivex/bench/kv_mooncake.py | 181 +++++++ collectivex/bench/kv_mori_io.py | 133 +++++ collectivex/bench/kv_nixl.py | 151 ++++++ collectivex/bench/kv_pool.py | 124 +++++ collectivex/bench/kv_workload.py | 222 ++++++++ collectivex/bench/run_kv.py | 518 ++++++++++++++++++ collectivex/configs/kv_sweep.json | 53 ++ collectivex/configs/platform_config.json | 21 +- collectivex/docs/methodology.md | 104 ++++ collectivex/launchers/launch_gb-nv.sh | 15 +- collectivex/launchers/launch_mi-amds.sh | 2 +- collectivex/launchers/launch_single-slurm.sh | 2 +- collectivex/runtime/common.sh | 5 +- collectivex/runtime/config.py | 31 ++ collectivex/runtime/prepare_backend.sh | 38 ++ collectivex/summarize.py | 87 ++- collectivex/sweep_matrix.py | 115 +++- collectivex/tests/test_kv_suite.py | 534 +++++++++++++++++++ collectivex/tests/test_kv_workload.py | 274 ++++++++++ collectivex/tests/test_runtime.py | 2 + 24 files changed, 2729 insertions(+), 22 deletions(-) create mode 100644 collectivex/bench/kv_backend.py create mode 100644 collectivex/bench/kv_mooncake.py create mode 100644 collectivex/bench/kv_mori_io.py create mode 100644 collectivex/bench/kv_nixl.py create mode 100644 collectivex/bench/kv_pool.py create mode 100644 collectivex/bench/kv_workload.py create mode 100644 collectivex/bench/run_kv.py create mode 100644 collectivex/configs/kv_sweep.json create mode 100644 collectivex/tests/test_kv_suite.py create mode 100644 collectivex/tests/test_kv_workload.py diff --git a/.github/workflows/collectivex-sweep.yml b/.github/workflows/collectivex-sweep.yml index 77789fe46f..b098e710fe 100644 --- a/.github/workflows/collectivex-sweep.yml +++ b/.github/workflows/collectivex-sweep.yml @@ -7,7 +7,7 @@ on: workflow_dispatch: inputs: suites: - description: "Comma-list of suites: ep (expert-parallel dispatch/combine), swap-blocks (vLLM block copies, one GPU per pool)" + description: "Comma-list of suites: ep (expert-parallel dispatch/combine), swap-blocks (vLLM block copies, one GPU per pool), kv-transfer (disaggregated KV handoff, 2 nodes x 1 GPU; hours per leg)" type: string default: ep backend: @@ -149,12 +149,18 @@ jobs: format('[{0}]', toJSON(matrix.runner)) ) }} name: p${{ needs.setup.outputs.priority }} | ${{ matrix.sku }} ${{ matrix.backend }} shard ${{ matrix.id }} - timeout-minutes: 350 + # A shard with a long allocation (kv-transfer) carries its own ceiling above it, so GitHub + # never cancels a healthy leg before the launcher's guards act. + timeout-minutes: ${{ matrix.job_timeout_minutes || 350 }} env: COLLX_BENCH: ${{ matrix.backend }} # Set only on shards whose suite ships its own image (swap-blocks); blank keeps the pool's. COLLX_IMAGE_OVERRIDE: ${{ matrix.image || '' }} COLLX_STAGED_IMAGE_DIR: ${{ matrix.staged_image_dir || '' }} + # The fabric a shard claims (kv-transfer rdma legs on an MNNVL rack leave the NVL domain) + # and its per-case hang guard; blank keeps the launcher defaults. + COLLX_FABRIC: ${{ matrix.fabric || '' }} + COLLX_RUN_TIMEOUT: ${{ matrix.run_timeout || '' }} COLLX_IMPORT_TMPDIR: ${{ matrix.sku == 'h100-dgxc' && '/var/tmp' || '/tmp' }} COLLX_NODES: ${{ matrix.nodes }} COLLX_GPUS_PER_NODE: ${{ matrix.gpus_per_node }} diff --git a/collectivex/README.md b/collectivex/README.md index 42ad7b44b2..54e788adf1 100644 --- a/collectivex/README.md +++ b/collectivex/README.md @@ -4,7 +4,9 @@ CollectiveX is an experimental MoE expert-parallel communication benchmark. It m combine, and paired roundtrip latency across EP libraries and accelerator systems, then uploads neutral result artifacts. -A second suite, the [vLLM `swap_blocks` benchmark](docs/swap-blocks.md) +A second suite, `kv-transfer`, measures disaggregated-serving KV-cache handoffs across transfer +libraries and fabrics ([below](#kv-cache-transfer-suite)). A third, the +[vLLM `swap_blocks` benchmark](docs/swap-blocks.md) ([中文](docs/swap-blocks_zh.md)), measures pinned CPU↔GPU and same-GPU block copies, with its own correctness checks and latency/bandwidth JSON output. It runs through the same sweep matrix and pool launchers (`suites: swap-blocks`). @@ -126,7 +128,8 @@ The matrix covers H100, H200, B200, B300, GB200, GB300, MI300X, MI325X, and MI35 the requested SKUs, backends, EP sizes, and token ladders, then extracts strict per-shard controls and rejects missing, stale, malformed, or altered shard controls. `--only-sku`, `--exclude-skus`, `--ep-sizes`, and `--precisions` select a subset. `--suites` picks the suites (`ep` by default; -`swap-blocks` adds one single-GPU shard per pool from `configs/swap_sweep.json`). Every suite's +`swap-blocks` adds one single-GPU shard per pool from `configs/swap_sweep.json`; `kv-transfer` adds +one 2-node shard per registry-enabled backend and fabric from `configs/kv_sweep.json`). Every suite's shards take the same path: the pool launcher allocates, `runtime/config.py case-args` encodes each case as its entrypoint's argv, and the rank wrapper execs that entrypoint. The matrix is generated per dispatch, with no frozen digest or locked case count. @@ -164,6 +167,27 @@ scale-up ranks per domain. GB EP16 remains MNNVL scale-up and therefore uses LSA SKU/backend/EP cell is attempted is a capability fact. Whether it succeeded is decided by the benchmark's return code. +## KV-Cache Transfer Suite + +`kv-transfer` legs run 2 nodes x 1 GPU, the per-worker prefill/decode pair a disaggregated +deployment actually forms. Each moves bursts of 1 to 32 concurrent requests' paged KV as vLLM's +packed block-major descriptor lists over seed-keyed random block tables (a burst posts every +request's prepped transfer, then awaits them all), plus one contiguous bulk row as the wire-speed +baseline. The workload is transcribed from what vLLM allocates for the model it serves: `kv-dsv4` +is DeepSeek-V4-Pro's mixed cache (30 Compressed Sparse Attention layers at 4 tokens per 576 B entry +plus their 132 B indexer entries, 31 Heavily Compressed Attention layers at 128 tokens per entry, +and the 128-token sliding-window cache on all 61 layers; fp8 by architecture), at ISL 2k to 512k +and vLLM's 256-token block. `pull` (READ, vLLM NixlConnector) and `push` (WRITE, SGLang disagg) are +both timed from the initiator, with pattern verification of every request on the destination pool. + +Backends are `nixl` (what Dynamo, vLLM, and SGLang ship), `mooncake` (the CUDA wheel links libcuda +at import; mi355x runs AMD's atom-dev build push-only), and `mori-io` (AMD's native engine), where +the registry's `kv_backends` map enables them. No entry, no legs, mirroring `ll_backends`. An entry +may restrict ops, pin an image, set a NIC filter, or lower the pool budget. Fabrics are `rdma` +(torch pools) and, on GB racks, `mnnvl` (cuMem FABRIC pools; see the methodology for the +bulk-vs-paged lane inversion that row exists to publish). The grid and per-pool scheduling live in +`configs/kv_sweep.json`; dispatch with `suites: kv-transfer`. + ## Workflow And Artifacts `.github/workflows/collectivex-sweep.yml` has two jobs. `setup` generates a public-SKU matrix diff --git a/collectivex/bandwidth.py b/collectivex/bandwidth.py index 93f5d14cec..97cf33d4fc 100644 --- a/collectivex/bandwidth.py +++ b/collectivex/bandwidth.py @@ -193,6 +193,10 @@ def render(documents: list[dict]) -> str: "marks an extrapolated alpha, and rungs failing the correctness gate are excluded.", "", ] + # kv-transfer documents have their own row model (per-transfer, no tokens_per_rank or + # routing); this renderer reads only EP rows. + documents = [d for d in documents + if d["identity"]["case_factors"]["case"].get("suite") != "kv-transfer"] for document in sorted(documents, key=_sort_key): case = document["identity"]["case_factors"]["case"] ep = _ep(document) diff --git a/collectivex/bench/kv_backend.py b/collectivex/bench/kv_backend.py new file mode 100644 index 0000000000..fec12acdc1 --- /dev/null +++ b/collectivex/bench/kv_backend.py @@ -0,0 +1,97 @@ +#!/usr/bin/env python3 +"""Backend contract for the KV-cache transfer suite. + +The harness owns the data (kv_pool pools, pattern fill, verification) and the +protocol (rank 0 = target, rank 1 = initiator, lockstep barriers); an adapter +owns registration, connection, and posting. Transfers are one-sided from the +initiator, so completion is host-visible and timing is wall clock around +post-to-complete; no CUDA events, because no local kernel participates. +`pull` (READ) is the vLLM NixlConnector shape, `push` (WRITE) the SGLang disagg +shape; the measured quantity is the same completion either way. Connection +payloads ride the harness object exchange, never adapter side channels. +""" + +from __future__ import annotations + +import time + + +class KVBackend: + """One transfer library on one rank. Subclasses implement the five hooks.""" + + name = "abstract" + #: maturity mirrors EPBackend.maturity ("production" | "candidate"). + maturity = "candidate" + library_version: str | None = None + + def __init__(self, args, role: str, device) -> None: + self.args = args + self.role = role + self.device = device + + # -- lifecycle ------------------------------------------------------------ + def register(self, pool, bulk, reg_layout=None) -> None: + """Register the pool + bulk tensors with the library. + + ``reg_layout`` is the pool's shared region layout — (base, + packed_bytes, nbytes) triples, contiguous from zero and valid for + every planned config (run_kv._harmonize). Adapters may use it to + split one oversized registration into pieces cut on the descriptor + grid, so no descriptor straddles two pieces; ignoring it is valid. + """ + raise NotImplementedError + + def publish(self) -> dict: + """Payload the peer needs to reach this rank (addresses, packed descs).""" + raise NotImplementedError + + def connect(self, peer: dict) -> None: + """Consume the peer's payload; after this, transfers may be prepared.""" + raise NotImplementedError + + def teardown(self) -> None: # pragma: no cover - adapter-specific + pass + + # -- transfers (initiator only) -------------------------------------------- + def make_paged(self, cfg: dict, op: str, local_tables, remote_tables): + """Return (post, wait, prep_seconds) for one request's paged KV. + + ``post()`` submits the whole descriptor list asynchronously; ``wait()`` + blocks until it completes — split so a batch of requests overlaps like + a decode step admitting several requests at once. Preparation cost + (descriptor build + handle creation) is amortized by engines through + prepped-handle reuse, so it is reported separately, never inside the + timed transfer. + """ + raise NotImplementedError + + def make_bulk(self, nbytes: int, op: str): + """Return (post, wait, prep_seconds) for one contiguous transfer of + ``nbytes`` — the single-descriptor contiguous baseline row (logical + payload over host-observed completion; not a proven physical wire + rate — backends may split large operations internally).""" + raise NotImplementedError + + +def time_bursts(transfers, warmup: int, reps: int) -> tuple[list[float], list[float]]: + """(burst_ms, request_ms), warmups dropped. ``transfers`` is a list of + (post, wait) pairs — one per request. A burst posts every request, then + drains the waits in posting order; burst_ms is post-of-first to + completion-of-last, and request_ms records each individual request's + host-observed completion offset from the burst start. Because the waits + drain in posting order, a request's mark upper-bounds its true completion + (a later request that finished early is observed at its wait's turn).""" + burst_ms: list[float] = [] + request_ms: list[float] = [] + for rep in range(warmup + reps): + start = time.perf_counter() + for post, _ in transfers: + post() + marks = [] + for _, wait in transfers: + wait() + marks.append((time.perf_counter() - start) * 1e3) + if rep >= warmup: + burst_ms.append(marks[-1]) + request_ms.extend(marks) + return burst_ms, request_ms diff --git a/collectivex/bench/kv_mooncake.py b/collectivex/bench/kv_mooncake.py new file mode 100644 index 0000000000..da375e5d6e --- /dev/null +++ b/collectivex/bench/kv_mooncake.py @@ -0,0 +1,181 @@ +#!/usr/bin/env python3 +"""Mooncake TransferEngine adapter for the KV-transfer suite. + +P2PHANDSHAKE metadata (no etcd); the peer session id is ip:rpc_port. The +binding is sync-only in the shape production uses it: SGLang's mooncake +connector posts blocking calls from a transfer thread pool (the binding +releases the GIL), so post() here hands the sync call to a worker thread and +wait() joins it. On CUDA images the wheel links libcudart.so.12, which cu13 +images do not carry; when the plain import fails the adapter dlopens it from +the nvidia-cuda-runtime-cu12 package and retries, so no launcher-side +LD_LIBRARY_PATH seam is needed. ROCm runs the image-provided build (AMD's +atom-dev tree; upstream wheels link libcuda.so.1), where transfers require +the GPU-paired NIC filter the registry passes through --kv-device. +""" + +from __future__ import annotations + +import ctypes +import os +import time +from concurrent.futures import ThreadPoolExecutor + +import kv_workload +from kv_backend import KVBackend + +BATCH_CAP = 8192 + + +def _preload_cudart() -> None: + try: + ctypes.CDLL("libcudart.so.12", mode=ctypes.RTLD_GLOBAL) + return + except OSError: + pass + import importlib.metadata as md + import pathlib + + for entry in md.files("nvidia-cuda-runtime-cu12") or []: + if entry.name == "libcudart.so.12": + ctypes.CDLL(str(pathlib.Path(entry.locate()).resolve()), mode=ctypes.RTLD_GLOBAL) + return + raise RuntimeError("libcudart.so.12 unavailable; install nvidia-cuda-runtime-cu12") + + +def _import_engine(): + """Plain import first (ROCm images ship a self-contained build); dlopen + the CUDA runtime and retry only when the wheel's link fails.""" + try: + from mooncake.engine import TransferEngine + except ImportError: + _preload_cudart() + from mooncake.engine import TransferEngine + return TransferEngine + + +def _engine_version(): + """The engine build actually imported, not the pin prepare_backend.sh + attempted: image-provided builds (b300's pre-0.3.12 lineage, AMD's + atom-dev tree) register under varying dist names or none at all, and + a null here is what let an image wheel masquerade as the pinned one.""" + import importlib.metadata as md + + for dist in ("mooncake-transfer-engine", "mooncake"): + try: + return md.version(dist) + except Exception: + pass + try: + import mooncake + + return getattr(mooncake, "__version__", None) + except Exception: + return None + + +def _physical_gpu_index() -> int: + """The physical GPU index behind this rank's visible device 0: GPU-paired + NIC selection (rdma{gpu}) needs the host-level index, which the Slurm + visibility mask carries.""" + for var in ("ROCR_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"): + first = os.environ.get(var, "").split(",")[0].strip() + if first.isdigit(): + return int(first) + return 0 + + +class MooncakeBackend(KVBackend): + name = "mooncake" + maturity = "production" + + def __init__(self, args, role, device): + super().__init__(args, role, device) + TransferEngine = _import_engine() + + self.library_version = _engine_version() + # Same-fabric GB pairs: the NVLink-IPC transport claims cross-node + # segments inside one NVLink domain and then fails the address import + # (nvlink_transport "Requested address not found", first kv CI run on + # gb200). This row measures the rdma lane, so pin the transport off; + # the ROCm twin (MC_USE_HIP_IPC) misclaims the same way on mi355x. + os.environ.setdefault("MC_USE_NVLINK_IPC", "0") + # The engine fails any single sync call after 30 s. Under measured + # contention (the x86 high-batch collapse this suite publishes) a + # chunked call on a thrashing lane can legitimately exceed that, so + # give the library guard 4x headroom; the runtime's per-case guard + # still bounds a truly wedged transfer. + os.environ.setdefault("MC_TRANSFER_TIMEOUT", "120") + self._engine = TransferEngine() + self._ip = kv_workload.iface_ipv4(args.socket_ifname) + local = f"{self._ip}:{args.kv_mc_port + (0 if role == 'target' else 1)}" + nic_filter = (getattr(args, "kv_device", "") or "").replace( + "{gpu}", str(_physical_gpu_index())) + self.nic_filter = nic_filter or None + rc = self._engine.initialize(local, "P2PHANDSHAKE", "rdma", nic_filter) + if rc != 0: + raise RuntimeError(f"mooncake initialize failed rc={rc} " + f"nic_filter={nic_filter!r}") + self._pool = None + self._bulk = None + self._peer = None + workers = max(int(v) for v in str(getattr(args, "batch_sizes", "1")).split()) + self._exec = ThreadPoolExecutor(max_workers=workers) + + def register(self, pool, bulk, reg_layout=None) -> None: + self._pool, self._bulk = pool, bulk + if self._engine.register_memory(pool.ptr, pool.nbytes) != 0 \ + or self._engine.register_memory(bulk.ptr, bulk.nbytes) != 0: + raise RuntimeError("mooncake memory registration failed") + + def publish(self) -> dict: + return {"session": f"{self._ip}:{self._engine.get_rpc_port()}", + "pool_base": self._pool.ptr, "bulk_base": self._bulk.ptr} + + def connect(self, peer: dict) -> None: + self._peer = peer + + def _split(self, run, prep_s): + """(post, wait, prep_s) around a blocking call via the worker pool.""" + pending: list = [] + + def post(): + pending.append(self._exec.submit(run)) + + def wait(): + pending.pop(0).result() + + return post, wait, prep_s + + def make_paged(self, cfg, op, local_tables, remote_tables): + start = time.perf_counter() + local = (self._pool.ptr + kv_workload.page_offsets(cfg, local_tables)).tolist() + remote = (self._peer["pool_base"] + kv_workload.page_offsets(cfg, remote_tables)).tolist() + sizes = kv_workload.desc_sizes(cfg).tolist() + chunks = [(i, min(i + BATCH_CAP, len(local))) for i in range(0, len(local), BATCH_CAP)] + session = self._peer["session"] + func = self._engine.batch_transfer_sync_read if op == "pull" \ + else self._engine.batch_transfer_sync_write + + def run(): + for i, j in chunks: + rc = func(session, local[i:j], remote[i:j], sizes[i:j]) + if rc != 0: + raise RuntimeError(f"mooncake batch transfer failed rc={rc}") + + return self._split(run, time.perf_counter() - start) + + def make_bulk(self, nbytes, op): + session = self._peer["session"] + func = self._engine.transfer_sync_read if op == "pull" \ + else self._engine.transfer_sync_write + local, remote = self._bulk.ptr, self._peer["bulk_base"] + + def run(): + rc = func(session, local, remote, nbytes) + if rc != 0: + raise RuntimeError(f"mooncake bulk transfer failed rc={rc}") + + return self._split(run, 0.0) + + def teardown(self) -> None: + self._exec.shutdown(wait=False) diff --git a/collectivex/bench/kv_mori_io.py b/collectivex/bench/kv_mori_io.py new file mode 100644 index 0000000000..91e5e98c45 --- /dev/null +++ b/collectivex/bench/kv_mori_io.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +"""MoRI-IO adapter (AMD's native P2P engine). Transfers address (region, +offset, size), so the paged list becomes offset lists over one registration. +Engine/Memory descriptors are packed blobs exchanged through the harness; the +engine's own control plane binds host/port from the SKU's socket interface. +Posts are capped at ``BATCH_CAP`` offsets per batch call to bound SQ/WR usage +and awaited together, the shape the SGLang MoRI-IO connector posts. +""" + +from __future__ import annotations + +import time + +import kv_workload +from kv_backend import KVBackend + +BATCH_CAP = 16384 + + +class MoRIIOBackend(KVBackend): + name = "mori-io" + maturity = "production" + + def __init__(self, args, role, device): + super().__init__(args, role, device) + from mori.io import (BackendType, IOEngine, IOEngineConfig, + MemoryLocationType, PollCqMode, RdmaBackendConfig) + + self._gpu_location = MemoryLocationType.GPU + + try: + import mori + + self.library_version = getattr(mori, "__version__", None) + except Exception: + self.library_version = None + self._mori_io = __import__("mori.io", fromlist=["EngineDesc", "MemoryDesc"]) + host = kv_workload.iface_ipv4(args.socket_ifname) if args.socket_ifname else "" + port = int(args.kv_mori_port) + (0 if role == "target" else 1) + self._engine = IOEngine(key=role, config=IOEngineConfig(host=host, port=port)) + self._engine.create_backend(BackendType.RDMA, RdmaBackendConfig( + qp_per_transfer=int(args.kv_mori_qp), + post_batch_size=-1, + num_worker_threads=1, + poll_cq_mode=PollCqMode.POLLING, + enable_notification=False, + enable_transfer_chunking=bool(args.kv_mori_chunking), + chunk_bytes=65536, + max_chunks_per_transfer=64, + )) + self._pool_mem = None + self._bulk_mem = None + self._sessions = None + + def register(self, pool, bulk, reg_layout=None) -> None: + self._pool_mem = self._engine.register_memory( + pool.ptr, pool.nbytes, pool.device, self._gpu_location) + self._bulk_mem = self._engine.register_memory( + bulk.ptr, bulk.nbytes, bulk.device, self._gpu_location) + + def publish(self) -> dict: + return { + "engine": bytes(self._engine.get_engine_desc().pack()), + "pool": bytes(self._pool_mem.pack()), + "bulk": bytes(self._bulk_mem.pack()), + } + + def connect(self, peer: dict) -> None: + self._engine.register_remote_engine(self._mori_io.EngineDesc.unpack(peer["engine"])) + remote_pool = self._mori_io.MemoryDesc.unpack(peer["pool"]) + remote_bulk = self._mori_io.MemoryDesc.unpack(peer["bulk"]) + self._sessions = { + "pool": self._engine.create_session(self._pool_mem, remote_pool), + "bulk": self._engine.create_session(self._bulk_mem, remote_bulk), + } + + @staticmethod + def _wait(statuses): + for status in statuses: + status.Wait() + if not status.Succeeded(): + raise RuntimeError(f"mori-io transfer failed: {status.Message()}") + + def make_paged(self, cfg, op, local_tables, remote_tables): + start = time.perf_counter() + local = kv_workload.page_offsets(cfg, local_tables).tolist() + remote = kv_workload.page_offsets(cfg, remote_tables).tolist() + sizes = kv_workload.desc_sizes(cfg).tolist() + chunks = [(i, min(i + BATCH_CAP, len(local))) for i in range(0, len(local), BATCH_CAP)] + session = self._sessions["pool"] + func = session.batch_read if op == "pull" else session.batch_write + engine = self._engine + prep_s = time.perf_counter() - start + statuses: list = [] + + def post(): + statuses.clear() + statuses.extend( + func(local[i:j], remote[i:j], sizes[i:j], engine.allocate_transfer_uid()) + for i, j in chunks + ) + + def wait(): + self._wait(statuses) + + return post, wait, prep_s + + # Verbs providers cap a single WR's message size (1 GiB on the Pollara path: + # a 2.3 GB single-WR bulk read dies ibv_post_send EINVAL). Split client-side; + # MoRI's own enable_transfer_chunking covers the same ground server-config-side, + # but a library-default row must not depend on a tuned engine config. + BULK_WR_CAP = 1 << 30 + + def make_bulk(self, nbytes, op): + session = self._sessions["bulk"] + func = session.read if op == "pull" else session.write + engine = self._engine + spans = [(offset, min(offset + self.BULK_WR_CAP, nbytes)) + for offset in range(0, nbytes, self.BULK_WR_CAP)] + + statuses: list = [] + + def post(): + statuses.clear() + statuses.extend( + func(start, start, end - start, engine.allocate_transfer_uid()) + for start, end in spans + ) + + def wait(): + self._wait(statuses) + + return post, wait, 0.0 diff --git a/collectivex/bench/kv_nixl.py b/collectivex/bench/kv_nixl.py new file mode 100644 index 0000000000..5bf36699d3 --- /dev/null +++ b/collectivex/bench/kv_nixl.py @@ -0,0 +1,151 @@ +#!/usr/bin/env python3 +"""NIXL (UCX) adapter: the library Dynamo, vLLM NixlConnector, and SGLang +disagg ship. Agent metadata rides the harness exchange (`add_remote_agent`), +not NIXL's TCP listener, so the adapter needs no port and no listener race. +Remote descriptors are built locally from the peer's published pool base; both +block tables are seed-keyed, the same information a decode worker gets from the +prefill side's block table message. +""" + +from __future__ import annotations + +import time + +import numpy as np + +import kv_workload +from kv_backend import KVBackend + +# b300's CX NICs refuse cuda registrations somewhere between 7083 and 8847 MiB +# (an ~8 GiB MR wall); UCX surfaces no error and the initiator later segfaults +# in ucp_worker_add_rkey_config resolving the region's rkey. Registering the +# pool in pieces below the wall sidesteps it everywhere; each region is cut on +# its own packed-block grid so no transfer descriptor straddles two pieces. +REG_CHUNK_BYTES = 4 << 30 + + +def reg_spans(nbytes: int, layout, + cap: int = REG_CHUNK_BYTES) -> list[tuple[int, int]]: + """(offset, length) registration pieces covering ``nbytes`` exactly. + + ``layout`` is the pool's shared region layout — (base, packed_bytes, + region_nbytes) triples, contiguous from zero and valid for every planned + config (run_kv._harmonize). Each region is cut into pieces of the largest + multiple of its packed_bytes at most ``cap``; without a layout the pool + is registered whole.""" + if not layout: + return [(0, nbytes)] + spans = [] + for base, packed, region_nbytes in layout: + chunk = max(cap // packed, 1) * packed + spans.extend((base + off, min(chunk, region_nbytes - off)) + for off in range(0, region_nbytes, chunk)) + covered = sum(length for _, length in spans) + if covered < nbytes: # tail the layout does not describe + spans.append((covered, nbytes - covered)) + return spans + + +class NIXLBackend(KVBackend): + name = "nixl" + maturity = "production" + + def __init__(self, args, role, device): + super().__init__(args, role, device) + from nixl._api import nixl_agent, nixl_agent_config + + try: + import importlib.metadata as md + + for dist_name in ("nixl", "nixl-cu13", "nixl-cu12"): + try: + self.library_version = md.version(dist_name) + break + except md.PackageNotFoundError: + continue + except Exception: + self.library_version = None + # The registry pin run_kv hands to UCX_NET_DEVICES for this case; + # None means UCX chose among the operator inventory itself. + self.nic_filter = getattr(args, "kv_device", "") or None + # prog thread on, listener off: metadata goes through the harness exchange. + self._agent = nixl_agent(role, nixl_agent_config(True, False, 0, backends=["UCX"])) + self._handles = [] + self._pool = None + self._bulk = None + self._peer = None + + def register(self, pool, bulk, reg_layout=None) -> None: + self._pool, self._bulk = pool, bulk + entries = [(pool.ptr + off, length, pool.device, f"pool{i}") + for i, (off, length) in + enumerate(reg_spans(pool.nbytes, reg_layout))] + # bulk rides one whole-request descriptor, so it can never be split; + # BULK_CAP bounds it. + entries.append((bulk.ptr, bulk.nbytes, bulk.device, "bulk")) + reg = self._agent.get_reg_descs(entries, mem_type="cuda") + if self._agent.register_memory(reg) is None: + raise RuntimeError("nixl memory registration failed") + + def publish(self) -> dict: + return { + "agent": bytes(self._agent.get_agent_metadata()), + "pool_base": self._pool.ptr, + "bulk_base": self._bulk.ptr, + "dev": self._pool.device, + } + + def connect(self, peer: dict) -> None: + self._peer = peer + remote = self._agent.add_remote_agent(peer["agent"]) + self._remote_name = remote.decode() if isinstance(remote, (bytes, bytearray)) else str(remote) + + def _make(self, local_np: np.ndarray, remote_np: np.ndarray, op: str): + start = time.perf_counter() + local_descs = self._agent.get_xfer_descs(local_np, mem_type="cuda") + remote_descs = self._agent.get_xfer_descs(remote_np, mem_type="cuda") + handle = self._agent.initialize_xfer( + "READ" if op == "pull" else "WRITE", + local_descs, remote_descs, self._remote_name, + ) + prep_s = time.perf_counter() - start + self._handles.append(handle) + agent = self._agent + + def post(): + if agent.transfer(handle) == "ERR": + raise RuntimeError("nixl post failed") + + def wait(): + while True: + state = agent.check_xfer_state(handle) + if state == "DONE": + return + if state == "ERR": + raise RuntimeError("nixl transfer errored") + + return post, wait, prep_s + + def make_paged(self, cfg, op, local_tables, remote_tables): + local_np = kv_workload.desc_array(self._pool.ptr, cfg, local_tables, self._pool.device) + remote_np = kv_workload.desc_array(self._peer["pool_base"], cfg, remote_tables, + self._peer["dev"]) + return self._make(local_np, remote_np, op) + + def make_bulk(self, nbytes, op): + local_np = np.array([[self._bulk.ptr, nbytes, self._bulk.device]], dtype=np.uint64) + remote_np = np.array([[self._peer["bulk_base"], nbytes, self._peer["dev"]]], + dtype=np.uint64) + return self._make(local_np, remote_np, op) + + def teardown(self) -> None: + for handle in self._handles: + try: + self._agent.release_xfer_handle(handle) + except Exception: + pass + if self._peer is not None: + try: + self._agent.remove_remote_agent(self._remote_name) + except Exception: + pass diff --git a/collectivex/bench/kv_pool.py b/collectivex/bench/kv_pool.py new file mode 100644 index 0000000000..403102e7d7 --- /dev/null +++ b/collectivex/bench/kv_pool.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +"""Pool allocators for the KV suite. + +The rdma lanes use plain torch (cudaMalloc) pools. The mnnvl lane needs cuMem +FABRIC allocations: UCX's cross-node cuda_ipc only engages on fabric-mappable +memory (cudaMalloc pools silently ride the IB rails instead), and fabric +handles need a live nvidia-imex domain. Both expose the same surface: raw +``ptr``/``nbytes``/``device``, pattern fill, byte fill, and ``read8`` for +kv_workload.verify_transfer. Adapters register raw pointers, never tensors. +""" + +from __future__ import annotations + +import ctypes +from ctypes import byref, c_int, c_size_t, c_ulonglong, c_void_p + +import numpy as np + +import kv_workload + +CU_MEM_ALLOCATION_TYPE_PINNED = 1 +CU_MEM_HANDLE_TYPE_FABRIC = 0x8 +CU_MEM_LOCATION_TYPE_DEVICE = 1 +CU_MEM_ACCESS_FLAGS_PROT_READWRITE = 3 + + +class TorchPool: + def __init__(self, nbytes: int, device: int): + import torch + + self._t = torch.empty(nbytes, dtype=torch.uint8, device=f"cuda:{device}") + self._torch = torch + self.ptr, self.nbytes, self.device = self._t.data_ptr(), nbytes, device + + def fill_pattern(self) -> None: + kv_workload.fill_pattern(self._t) + self._torch.cuda.synchronize() + + def fill_byte(self, value: int) -> None: + self._t.fill_(value) + self._torch.cuda.synchronize() + + def read8(self, offset: int): + return self._t[offset : offset + 8].cpu().numpy().tobytes() + + +class _AllocProp(ctypes.Structure): + _fields_ = [("type", c_int), ("requestedHandleTypes", c_int), + ("location_type", c_int), ("location_id", c_int), + ("win32HandleMetaData", c_void_p), + ("compressionType", ctypes.c_ubyte), + ("gpuDirectRDMACapable", ctypes.c_ubyte), + ("usage", ctypes.c_ushort), + ("reserved", ctypes.c_ubyte * 4)] + + +class _AccessDesc(ctypes.Structure): + _fields_ = [("location_type", c_int), ("location_id", c_int), ("flags", c_int)] + + +_PATTERNS: dict[int, np.ndarray] = {} + + +def _pattern(nbytes: int) -> np.ndarray: + if nbytes not in _PATTERNS: + chunks = nbytes // 256 + vals = ((np.arange(chunks, dtype=np.int64) * 131 + 7) & 0xFF).astype(np.uint8) + _PATTERNS[nbytes] = np.repeat(vals, 256) + return _PATTERNS[nbytes] + + +class FabricPool: + def __init__(self, nbytes: int, device: int): + cu = self._cu = ctypes.CDLL("libcuda.so.1") + self._check(cu.cuInit(0), "cuInit") + dev = c_int() + self._check(cu.cuDeviceGet(byref(dev), device), "cuDeviceGet") + ctx = c_void_p() + self._check(cu.cuDevicePrimaryCtxRetain(byref(ctx), dev), "cuDevicePrimaryCtxRetain") + self._check(cu.cuCtxSetCurrent(ctx), "cuCtxSetCurrent") + prop = _AllocProp(type=CU_MEM_ALLOCATION_TYPE_PINNED, + requestedHandleTypes=CU_MEM_HANDLE_TYPE_FABRIC, + location_type=CU_MEM_LOCATION_TYPE_DEVICE, + location_id=device, gpuDirectRDMACapable=1) + gran = c_size_t() + self._check(cu.cuMemGetAllocationGranularity(byref(gran), byref(prop), 0), "granularity") + size = (nbytes + gran.value - 1) // gran.value * gran.value + handle = c_ulonglong() + code = cu.cuMemCreate(byref(handle), c_size_t(size), byref(prop), 0) + if code != 0: + raise RuntimeError( + f"cuMemCreate(FABRIC) -> CUresult {code}: no IMEX fabric access on this " + "allocation; the mnnvl lane cannot run here") + ptr = c_ulonglong() + self._check(cu.cuMemAddressReserve(byref(ptr), c_size_t(size), 0, 0, 0), "reserve") + self._check(cu.cuMemMap(ptr, c_size_t(size), 0, handle, 0), "map") + access = _AccessDesc(location_type=CU_MEM_LOCATION_TYPE_DEVICE, location_id=device, + flags=CU_MEM_ACCESS_FLAGS_PROT_READWRITE) + self._check(cu.cuMemSetAccess(ptr, c_size_t(size), byref(access), 1), "setAccess") + self.ptr, self.nbytes, self.device = ptr.value, size, device + + def _check(self, code: int, what: str) -> None: + if code != 0: + raise RuntimeError(f"{what} -> CUresult {code}") + + def _h2d(self, host: np.ndarray) -> None: + self._check(self._cu.cuMemcpyHtoD_v2( + c_ulonglong(self.ptr), host.ctypes.data_as(c_void_p), c_size_t(host.nbytes)), "h2d") + + def fill_pattern(self) -> None: + self._h2d(_pattern(self.nbytes)) + + def fill_byte(self, value: int) -> None: + self._h2d(np.full(self.nbytes, value, dtype=np.uint8)) + + def read8(self, offset: int): + out = np.empty(8, dtype=np.uint8) + self._check(self._cu.cuMemcpyDtoH_v2( + out.ctypes.data_as(c_void_p), c_ulonglong(self.ptr + offset), c_size_t(8)), "d2h") + return out.tobytes() + + +def create(fabric: str, nbytes: int, device: int): + return FabricPool(nbytes, device) if fabric == "mnnvl" else TorchPool(nbytes, device) diff --git a/collectivex/bench/kv_workload.py b/collectivex/bench/kv_workload.py new file mode 100644 index 0000000000..83b089ee59 --- /dev/null +++ b/collectivex/bench/kv_workload.py @@ -0,0 +1,222 @@ +#!/usr/bin/env python3 +"""Workload model for the KV-cache transfer suite. + +A transfer is one request's paged KV in the shape vLLM's packed DSV4 NIXL path +actually registers and posts: per cache group, the physical block is the +transfer unit, and one contiguous descriptor covers ALL of that group's layers +for the block (block-major `[block][layer]` layout, `packed_bytes = layers x +page_bytes` per descriptor). Fragmentation is real but block-granular: +seed-keyed random block tables per side scatter each request's blocks over the +pool, exactly what a fragmented allocator hands a connector. What this model +deliberately does NOT do is explode each (layer, page) into its own descriptor +— vLLM's connector asserts one descriptor per packed physical block, and the +per-(layer, page) shape inflates descriptor counts by ~2 orders of magnitude, +which inverts backend and fabric conclusions on descriptor-bound lanes. + +Geometry for ``dsv4`` is transcribed from vLLM (validated against commit +32ad1400d7): every token-state is 584 B of content (448 B NoPE + 128 B RoPE + +8 B fp8 scale, the ``fp8_ds_mla`` layout), and each block's page is padded to +a 576 B multiple (FlashMLA packing — alignment applies at PAGE granularity, +not per state). The config's ``compress_ratios`` interleave 30 Compressed +Sparse Attention layers (4 tokens per state) with 31 Heavily Compressed +Attention layers (128 tokens per state); CSA layers add a lightning-indexer +cache (132 B per state: 128 fp8 + 4 scale bytes); and all 61 layers keep a +128-token sliding window whose block size is FIXED at 64 tokens because the +window shares its physical tensor with the CSA cache (a 256-token CSA block is +64 states, so the shared tensor's block covers 64 window tokens — the window +page equals the CSA page byte for byte). HCA's 128-token states force the +model block size to a multiple of 128; vLLM serves DSV4 at 256. The dtype mix +is architectural, so the preset pins precision to "fp8". + +Pattern correctness: byte at offset o of a pool is derived from o (constant per +256-byte chunk), so any block's expected contents follow from its offset alone, +at any alignment. +""" + +from __future__ import annotations + +import fcntl +import math +import socket +import struct + +import numpy as np + +PRESETS = { + "dsv4": dict( + model_class="deepseek-v4-pro", + precisions=("fp8",), # vLLM's fp8_ds_mla states + fp8 indexer, baked in + model_layers=61, + alignment=576, # vLLM pads each block's page to this (FlashMLA packing) + groups=( + dict(name="c4a", layers=30, tokens_per_state=4, state_bytes=584), + dict(name="c4a-idx", layers=30, tokens_per_state=4, state_bytes=132), + dict(name="c128a", layers=31, tokens_per_state=128, state_bytes=584), + dict(name="swa", layers=61, tokens_per_state=1, state_bytes=584, + block_tokens=64, window_tokens=128), + ), + ), +} + + +def _round_up(value: int, align: int) -> int: + return -(-value // align) * align + + +def plan_config(preset: str, precision: str, isl: int, block_tokens: int, + pool_slack: float = 2.0, batch_max: int = 1) -> dict: + """Resolve one (preset, precision, isl, block size) point into regions. + + A region is one vLLM cache group. Every region gets: layers, page_bytes + (one layer's padded page for one block), packed_bytes (the transfer unit — + one descriptor covering all the group's layers for one physical block), + blocks_req (descriptors for one request), pool_blocks (sized so + ``batch_max`` concurrent requests hold disjoint blocks, plus fragmentation + head-room), and its base offset in the one contiguous pool allocation. + """ + shape = PRESETS[preset] + if precision not in shape["precisions"]: + raise ValueError(f"{preset} runs {shape['precisions']}, not {precision}") + pool_slack = max(pool_slack, batch_max * 1.25) + regions = [] + offset = 0 + for group in shape["groups"]: + group_block = group.get("block_tokens", block_tokens) + if group_block < group["tokens_per_state"] \ + or group_block % group["tokens_per_state"]: + raise ValueError( + f"{preset} block size {block_tokens} does not hold whole " + f"{group['name']} states ({group['tokens_per_state']} tokens each)") + states = group_block // group["tokens_per_state"] + page_bytes = _round_up(states * group["state_bytes"], shape["alignment"]) + packed_bytes = group["layers"] * page_bytes + tokens = min(isl, group["window_tokens"]) if "window_tokens" in group else isl + blocks_req = math.ceil(tokens / group_block) + pool_blocks = int(blocks_req * pool_slack) + 8 + regions.append(dict(name=group["name"], layers=group["layers"], + block_tokens=group_block, page_bytes=page_bytes, + packed_bytes=packed_bytes, blocks_req=blocks_req, + pool_blocks=pool_blocks, base=offset)) + offset += pool_blocks * packed_bytes + + return dict( + preset=preset, + precision=precision, + isl=isl, + page_tokens=block_tokens, # row label: the model block size in tokens + layers=shape["model_layers"], + page_bytes=regions[0]["packed_bytes"], # one primary-region descriptor + regions=regions, + pool_bytes=offset, + req_bytes=sum(r["blocks_req"] * r["packed_bytes"] for r in regions), + descs=sum(r["blocks_req"] for r in regions), + ) + + +def block_table(cfg: dict, seed: int, request: int = 0) -> dict: + """Per-region block tables (deterministic, seed-keyed): region name -> the + random block permutation a fragmented allocator would hand the request. + Requests in one batch slice disjoint ranges of a single permutation, as a + real allocator's live requests never alias blocks.""" + rng = np.random.default_rng(seed) + tables = {} + for region in cfg["regions"]: + low = request * region["blocks_req"] + tables[region["name"]] = ( + rng.permutation(region["pool_blocks"])[low : low + region["blocks_req"]] + ) + if len(tables[region["name"]]) < region["blocks_req"]: + raise ValueError(f"pool too small for batch request {request} " + f"in region {region['name']}") + return tables + + +def table_seed(cfg: dict, side: str) -> int: + """Both ranks derive both sides' tables from the config alone — no exchange.""" + base = cfg["isl"] * 31 + cfg["page_tokens"] + len(cfg["preset"]) * 7 + return base + (1000 if side == "local" else 0) + + +def page_offsets(cfg: dict, tables: dict) -> np.ndarray: + """Block-major byte offsets (relative to the pool base) across all regions: + one offset per packed physical block, the descriptor vLLM posts.""" + parts = [] + for region in cfg["regions"]: + offsets = (tables[region["name"]].astype(np.uint64) + * np.uint64(region["packed_bytes"]) + np.uint64(region["base"])) + parts.append(offsets) + return np.concatenate(parts) + + +def desc_sizes(cfg: dict) -> np.ndarray: + """Per-descriptor byte sizes aligned with page_offsets' ordering.""" + return np.concatenate([ + np.full(region["blocks_req"], region["packed_bytes"], dtype=np.uint64) + for region in cfg["regions"] + ]) + + +def desc_array(base: int, cfg: dict, tables: dict, dev: int) -> np.ndarray: + """(addr, len, devId) uint64 rows for descriptor-list APIs (NIXL's numpy form).""" + out = np.empty((cfg["descs"], 3), dtype=np.uint64) + out[:, 0] = np.uint64(base) + page_offsets(cfg, tables) + out[:, 1] = desc_sizes(cfg) + out[:, 2] = dev + return out + + +def _chunk_byte(offset: int) -> int: + return ((offset >> 8) * 131 + 7) & 0xFF + + +def fill_pattern(pool_u8) -> None: + """Paint the offset-derived pattern over the whole pool (torch uint8 tensor).""" + import torch + + chunks = pool_u8.numel() // 256 + view = pool_u8[: chunks * 256].view(chunks, 256) + vals = ((torch.arange(chunks, device=pool_u8.device, dtype=torch.int64) * 131 + 7) & 0xFF) + view.copy_(vals.to(torch.uint8)[:, None].expand(chunks, 256)) + + +def verify_transfer(read8, cfg: dict, dst_tables: dict, src_tables: dict, + samples: int = 16, seed: int = 7) -> tuple[bool, str]: + """On the destination pool: packed block (region, dst[i]) must hold the + source pool's pattern at (region, src[i])'s offset. Each sample probes one + layer's page inside the packed block, so the checks range over the whole + descriptor. ``read8(offset)`` returns 8 destination-pool bytes (see + kv_pool). Compared per byte, so any page alignment verifies exactly.""" + rng = np.random.default_rng(seed) + for _ in range(samples): + region = cfg["regions"][int(rng.integers(len(cfg["regions"])))] + dst, src = dst_tables[region["name"]], src_tables[region["name"]] + layer = int(rng.integers(region["layers"])) + i = int(rng.integers(len(dst))) + delta = layer * region["page_bytes"] + src_off = int(src[i]) * region["packed_bytes"] + region["base"] + delta + dst_off = int(dst[i]) * region["packed_bytes"] + region["base"] + delta + expected = bytes(_chunk_byte(src_off + j) for j in range(8)) + got = bytes(read8(dst_off)) + if got != expected: + return False, (f"region={region['name']} layer={layer} i={i} " + f"expected={list(expected)} got={list(got)}") + return True, "" + + +def pcts(samples_ms: list[float]) -> dict: + ordered = sorted(samples_ms) + n = len(ordered) + return { + "p50": ordered[n // 2], + "p95": ordered[min(n - 1, int(n * 0.95))], + "min": ordered[0], + "max": ordered[-1], + "n": n, + } + + +def iface_ipv4(iface: str) -> str: + """IPv4 of a named interface (SIOCGIFADDR); the TCP bootstrap address.""" + sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + packed = struct.pack("256s", iface.encode()[:15]) + return socket.inet_ntoa(fcntl.ioctl(sock.fileno(), 0x8915, packed)[20:24]) diff --git a/collectivex/bench/run_kv.py b/collectivex/bench/run_kv.py new file mode 100644 index 0000000000..19311e7445 --- /dev/null +++ b/collectivex/bench/run_kv.py @@ -0,0 +1,518 @@ +#!/usr/bin/env python3 +"""CollectiveX KV-cache transfer benchmark entrypoint (2 ranks, 1 per node). + +Rank 0 is the target (owns the pool the initiator pulls from / pushes into), +rank 1 the initiator (posts every one-sided transfer and is the timed side). +The control plane is a gloo process group: payload exchange by object gather, +lockstep by barrier — no shared-FS or side-channel protocols. Data never rides +gloo. + +Per (isl, page_tokens, batch) point the initiator preps one transfer per +request in the burst (disjoint block-table slices), posts them all, then awaits +them all — a decode step admitting B requests at once. Verification covers both +directions (pull on the initiator's pool, push on the target's, exchanged as +verdict objects) and both pools are repainted between points so every verify +reads a clean pattern. Points whose pool would not fit POOL_BUDGET shed their +largest batches, so one grid covers dense GQA-bf16 and DSv4's ~2% cache alike. +""" + +from __future__ import annotations + +import argparse +import datetime as _dt +import json +import os +import socket +import sys + +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path[:0] = [HERE, os.path.dirname(HERE)] + +import ep_harness # noqa: E402 (case_id/is_case_id + atomic write; stdlib-only) +import kv_workload # noqa: E402 +from kv_backend import time_bursts # noqa: E402 + +BULK_CAP = 8 << 30 +# Default pool ceiling per rank: fits the fleet's smallest HBM (h200, 141 GB) +# next to the bulk buffer; grid points shed their largest batches to stay under +# it. A registry kv_backends entry lowers it (`pool_budget`, reaching here as +# --pool-budget) where an engine/NIC pairing cannot register a pool this +# large: mooncake on the mi355x ionic NICs fails ibv_reg_mr with ENOMEM between +# the 20 GiB pool the mixed batch ladder planned (green) and the 53 GiB the +# power-of-two ladder plans (red on two independent allocations), while +# mori-io registers the same pool fine, so the cap is per backend, not +# fleet-wide. +POOL_BUDGET = 64 << 30 +# Burst posting ceiling: a burst posts batch x descs descriptors, and the +# per-descriptor floor makes time linear in that product. On the packed +# block-major geometry a request is only ceil(isl/block) descriptors per +# group (a 512k-ISL block-256 request is ~6.1k), so the production grid sits +# far under this; the budget stays as the fail-closed guard for future +# presets or small block sizes. +DESC_BUDGET = 2_250_000 +# The LADDER_FLOOR smallest requested batches ride over DESC_BUDGET anyway: +# the frontier chart draws its line through the batch ladder at the largest +# measured ISL, and budget shedding alone leaves that ladder 2-3 points — +# not interpretable. Five rungs keep every point chartable while bounding +# the overrun (on the power-of-two ladder, 512k page-16 tops out at batch +# 16, ~15x budget for that one burst; grid-wide the floor costs ~1.33x the +# mixed twelve-rung grid and is priced into all three kv launcher guards). +LADDER_FLOOR = 5 + + +def add_kv_args(ap: argparse.ArgumentParser) -> None: + ap.add_argument("--workload-name", required=True, help="kv-, e.g. kv-dsv4") + ap.add_argument("--precision", required=True, choices=["bf16", "fp8"]) + ap.add_argument("--fabric", default="rdma", choices=["rdma", "mnnvl"], + help="which lane the SKU row claims; mnnvl additionally sets " + "UCX_CUDA_IPC_ENABLE_MNNVL=y for the UCX-backed libraries") + ap.add_argument("--isl-ladder", default="512 4096 32768") + ap.add_argument("--page-tokens", default="256", + help="vLLM block size in tokens; dsv4 needs a multiple of " + "128 (HCA states) and vLLM serves it at 256") + ap.add_argument("--ops", default="pull push") + ap.add_argument("--batch-sizes", default="1", + help="requests per burst; each is a separate prepped transfer, " + "posted together then awaited together") + ap.add_argument("--warmup", type=int, default=2) + ap.add_argument("--reps", type=int, default=8) + ap.add_argument("--trials", type=int, default=3) + ap.add_argument("--pool-slack", type=float, default=2.0) + ap.add_argument("--pool-budget", type=int, default=POOL_BUDGET, + help="per-rank pool ceiling in bytes; points shed batches to fit") + ap.add_argument("--seed", type=int, default=67) + ap.add_argument("--runner", required=True) + ap.add_argument("--case-id", default="", help="scheduled case ID; computed when omitted") + ap.add_argument("--suite", default="kv-transfer") + ap.add_argument("--version", type=int, default=1) + ap.add_argument("--out", default="") + ap.add_argument("--gpus-per-node", type=int, default=8) + ap.add_argument("--scale-up-domain", type=int, default=8) + ap.add_argument("--scale-up-transport", default="") + ap.add_argument("--topology-class", default="") + ap.add_argument("--socket-ifname", default=os.environ.get("COLLX_SOCKET_IFNAME", "")) + ap.add_argument("--kv-mori-qp", type=int, default=1) + ap.add_argument("--kv-mori-chunking", action="store_true") + ap.add_argument("--kv-device", default="", + help="engine NIC filter template; {gpu} expands to the " + "physical GPU index (GPU-paired NICs, e.g. Pollara). " + "For nixl it is a literal netdev comma-list pinning " + "UCX_NET_DEVICES below the operator inventory") + ap.add_argument("--kv-mori-port", type=int, default=48810) + ap.add_argument("--kv-mc-port", type=int, default=48830) + + +def export_ucx_selectors(environ=os.environ, device: str = "") -> None: + """Pin the UCX fabric to the operator's validated RDMA selectors. + + UCX auto-selection is a wrong-fabric trap on several SKUs (b200-nscale's + quad-port aux card, b300's storage IB), and the launcher's network profile + only exports the COLLX_* names. Explicit UCX_* values always win. + + ``device`` is the case's registry NIC pin (kv_device) for a UCX-backed + engine: a literal netdev comma-list that narrows UCX below the operator + inventory, for rail-isolated pods where multi-rail selection is the + variance source under measurement. Unlike the inventory it overrides a + host-inherited UCX_NET_DEVICES: b300 ships a blanket 16-device value in + /etc/environment (forwarded by srun --export=ALL) that would otherwise + silently swallow the pin, the same way its UCX_TLS=rc is dropped below. + """ + devices = device or environ.get("COLLX_RDMA_DEVICES", "") + if devices and (device or "UCX_NET_DEVICES" not in environ): + environ["UCX_NET_DEVICES"] = ",".join( + dev if ":" in dev else f"{dev}:1" + for dev in devices.split(",") if dev) + gid = environ.get("COLLX_IB_GID_INDEX", "") + if gid and "UCX_IB_GID_INDEX" not in environ: + environ["UCX_IB_GID_INDEX"] = str(gid) + # A host-inherited positive UCX_TLS list without the cuda transports (b300 + # ships UCX_TLS=rc cluster-wide in /etc/environment, forwarded by srun + # --export=ALL) makes ucp close the cuda mds; UCX then classifies VRAM as + # host memory and NIXL registration fails with NIXL_ERR_BACKEND. Extending + # the list with cuda_copy,cuda_ipc is not enough: the initiator then + # segfaults in ucp_worker_add_rkey_config resolving the cuda rkey on the + # first ucp_get_nbx. Drop the list and let UCX auto-select; the wire stays + # pinned through UCX_NET_DEVICES above. + tls = environ.get("UCX_TLS", "") + if tls and tls != "all" and not tls.startswith("^") and "cuda" not in tls: + del environ["UCX_TLS"] + + +def exchange_verdict(dist, role, verify_side, verify): + """One rank verifies its destination pool; every rank returns that verdict. + + Bulk rows have no verifying side (verify_side "none"): every rank gathers + None and the row passes by construction, without a gather-of-nothing crash. + """ + verdict = None + if role == verify_side: + passed, detail = verify() + verdict = {"passed": passed, "detail": detail} + gathered = [None, None] + dist.all_gather_object(gathered, verdict) + return next((v for v in gathered if v is not None), {"passed": True, "detail": ""}) + + +def kv_case(args) -> dict: + return { + "backend": args.backend, + "workload": args.workload_name, + "mode": args.fabric, + "phase": "xfer", + "ep": 2, + "routing": "paged", + "precision": args.precision, + } + + +def _grid(args) -> tuple[list[tuple[dict, list[int]]], list[int], list[int]]: + """(cfg, allowed_batches) per (isl, page) point. Batches whose burst would + exceed DESC_BUDGET are shed first (the LADDER_FLOOR smallest requested + batches are always kept, so every point carries a chartable batch ladder + even where a single request nearly fills the budget), then the point is + planned for the largest surviving batch whose pool fits the pool budget. + Smaller batches share that cfg (and pool), so batch is the only variable + across a point's rows.""" + preset = args.workload_name.removeprefix("kv-") + isls = [int(v) for v in args.isl_ladder.split()] + pages = [int(v) for v in args.page_tokens.split()] + batches = sorted({int(v) for v in args.batch_sizes.split()}) + points = [] + for isl in isls: + for page in pages: + # Per-request descriptor count is independent of batch_max. + probe = kv_workload.plan_config(preset, args.precision, isl, page, + args.pool_slack) + allowed = [batch for batch in batches + if batch in batches[:LADDER_FLOOR] + or batch * probe["descs"] <= DESC_BUDGET] + while allowed: + cfg = kv_workload.plan_config(preset, args.precision, isl, page, + args.pool_slack, batch_max=allowed[-1]) + if cfg["pool_bytes"] <= args.pool_budget: + break + allowed.pop() + if allowed: + points.append((cfg, allowed)) + return points, isls, batches + + +def _harmonize(points) -> list[tuple[int, int, int]]: + """Rewrite every cfg's regions onto one shared pool layout and return it + as (base, packed_bytes, nbytes) triples, contiguous from zero. + + The shard's configs share one physical pool, but plan_config lays each + config's regions out independently, so region bases shift with ISL and no + registration cut point is on every config's descriptor grid at once. + Giving every region the largest pool_blocks any config plans for it makes + the bases config-invariant: a cut on a region's own packed grid is then + between descriptors for every config, which is what lets a backend split + an oversized registration (b300 NICs refuse cuda registrations past ~8 + GiB) without a descriptor ever straddling two pieces. Configs planning a + different page size carry a different packed grid, so each page family + gets its own slab after the previous one. The union can run past the + largest single config's per-config pool budget check by the smaller + configs' head-room; the budget's slack absorbs that.""" + layout: list[tuple[int, int, int]] = [] + offset = 0 + families: dict[int, list[dict]] = {} + for cfg, _ in points: + families.setdefault(cfg["page_tokens"], []).append(cfg) + for cfgs in families.values(): + shared = [] + for i, region in enumerate(cfgs[0]["regions"]): + packed = region["packed_bytes"] + blocks = max(cfg["regions"][i]["pool_blocks"] for cfg in cfgs) + shared.append((offset, packed, blocks * packed)) + offset += blocks * packed + for cfg in cfgs: + for region, (base, packed, nbytes) in zip(cfg["regions"], shared): + region["base"] = base + region["pool_blocks"] = nbytes // packed + layout.extend(shared) + for cfg, _ in points: + cfg["pool_bytes"] = offset + return layout + + +def main() -> int: + ap = argparse.ArgumentParser(description="CollectiveX KV-cache transfer sweep") + ap.add_argument("--backend", required=True, choices=["nixl", "mori-io", "mooncake"]) + add_kv_args(ap) + args = ap.parse_args() + + case = kv_case(args) + computed_case_id = ep_harness.case_id(args.runner, case) + if args.case_id and args.case_id != computed_case_id: + print(f"ERROR: scheduled case ID does not match factors: " + f"{args.case_id} != {computed_case_id}", file=sys.stderr) + return 2 + args.case_id = args.case_id or computed_case_id + + if args.fabric == "mnnvl": + os.environ.setdefault("UCX_CUDA_IPC_ENABLE_MNNVL", "y") + if args.socket_ifname: + os.environ.setdefault("GLOO_SOCKET_IFNAME", args.socket_ifname) + export_ucx_selectors( + device=args.kv_device if args.backend == "nixl" else "") + + import torch + import torch.distributed as dist + + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "2")) + if world_size != 2: + print(f"ERROR: kv-transfer runs exactly 2 ranks, got {world_size}", file=sys.stderr) + return 2 + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + role = "target" if rank == 0 else "initiator" + # A single grid point's timed stretch can run past gloo's 30-minute + # default recv timeout (a slow lane's large-ISL bursts, while the target + # rank waits silently at the next gather). Size the control-plane timeout + # to the per-case hang guard so the guard, not gloo, decides when a run + # died. + grace_s = int(os.environ.get("COLLX_RUN_TIMEOUT") or "21600") + dist.init_process_group("gloo", rank=rank, world_size=world_size, + timeout=_dt.timedelta(seconds=grace_s)) + + if args.backend == "mori-io": + from kv_mori_io import MoRIIOBackend as Backend + elif args.backend == "mooncake": + from kv_mooncake import MooncakeBackend as Backend + else: + from kv_nixl import NIXLBackend as Backend + + points, isls, batches = _grid(args) + reg_layout = _harmonize(points) + ops = args.ops.split() + pool_bytes = max(cfg["pool_bytes"] for cfg, _ in points) + bulk_bytes = min(max(cfg["req_bytes"] for cfg, _ in points), BULK_CAP) + + # RDMA registration pins the whole pool; a small inherited soft memlock + # limit fails it with an unhelpful ENOMEM/EIO deep inside the library + # (Slurm propagates the SUBMITTER's limits into steps). Raise soft to hard + # when possible; otherwise fail here with the actual numbers. + import resource + + soft, hard = resource.getrlimit(resource.RLIMIT_MEMLOCK) + need = pool_bytes + bulk_bytes + if soft != resource.RLIM_INFINITY and (hard == resource.RLIM_INFINITY or soft < hard): + resource.setrlimit(resource.RLIMIT_MEMLOCK, (hard, hard)) + soft = hard + if soft != resource.RLIM_INFINITY and soft < need: + print(f"ERROR: RLIMIT_MEMLOCK {soft} < {need} needed to register the KV pools; " + "submit with --propagate=NONE or raise the limit", file=sys.stderr) + return 2 + + import kv_pool + + pool = kv_pool.create(args.fabric, pool_bytes, local_rank) + bulk = kv_pool.create(args.fabric, bulk_bytes, local_rank) + + def repaint(): + pool.fill_pattern() + bulk.fill_byte(0xAB if role == "target" else 0xCD) + + repaint() + backend = Backend(args, role, device) + backend.register(pool, bulk, reg_layout=reg_layout) + payloads = [None, None] + dist.all_gather_object(payloads, backend.publish()) + backend.connect(payloads[1 - rank]) + dist.barrier() + if rank == 1: + print(f"[run_kv] backend={args.backend} workload={args.workload_name} " + f"precision={args.precision} fabric={args.fabric} isls={isls} " + f"batches={batches} pool={pool_bytes >> 20}MiB case={args.case_id}", + flush=True) + for cfg, allowed in points: + if allowed != batches: + print(f"[run_kv] budgets cap isl={cfg['isl']} " + f"page={cfg['page_tokens']} at batch<={allowed[-1]}", flush=True) + + rows: list[dict] = [] + + def verify_burst(cfg, table_pairs): + """Every request in the burst must land: a passing request 0 says + nothing about the others, and concurrent same-session requests are + exactly where corruption would hide.""" + for r, (dst, src) in enumerate(table_pairs): + passed, detail = kv_workload.verify_transfer(pool.read8, cfg, dst, src) + if not passed: + return False, f"request={r} {detail}" + return True, "" + + def measure(make, cfg_row: dict, op: str, verify_side: str, table_pairs=None): + """One grid point: initiator times bursts, then the verifying side checks.""" + if role == "initiator": + made = make() # one (post, wait, prep_seconds) per request in the burst + prep_s = sum(m[2] for m in made) + pairs = [m[:2] for m in made] + samples: list[float] = [] + request_samples: list[float] = [] + for _ in range(args.trials): + burst_ms, request_ms = time_bursts(pairs, args.warmup, args.reps) + samples.extend(burst_ms) + request_samples.extend(request_ms) + dist.barrier() # transfers complete before anyone inspects pools + verdict = exchange_verdict( + dist, role, verify_side, + lambda: verify_burst(cfg_row["_cfg"], table_pairs)) + repaint() + dist.barrier() + if role != "initiator": + return None + stats = kv_workload.pcts(samples) + request_stats = kv_workload.pcts(request_samples) + prep_ms = prep_s * 1e3 + gbps = cfg_row["req_bytes"] * cfg_row["batch"] / stats["p50"] / 1e6 + # The cold-path rate: a burst whose descriptors and handles are built + # fresh (unique block tables, no prepped-handle reuse) pays prep once. + gbps_incl_prep = (cfg_row["req_bytes"] * cfg_row["batch"] + / (stats["p50"] + prep_ms) / 1e6) + return { + **{k: v for k, v in cfg_row.items() if not k.startswith("_")}, + "op": op, + "prep_ms": round(prep_ms, 3), + "latency_ms": {k: round(v, 3) for k, v in stats.items()}, + # Host-observed completion of each individual request within its + # burst (waits drain in posting order, so each is an upper bound). + "request_ms": {k: round(v, 3) for k, v in request_stats.items()}, + "gbps_p50": round(gbps, 2), + "gbps_p50_incl_prep": round(gbps_incl_prep, 2), + "verify": verdict, + } + + for cfg, allowed in points: + seed_t = kv_workload.table_seed(cfg, "remote") + seed_i = kv_workload.table_seed(cfg, "local") + target_tables = [kv_workload.block_table(cfg, seed_t, r) for r in range(allowed[-1])] + initiator_tables = [kv_workload.block_table(cfg, seed_i, r) for r in range(allowed[-1])] + base = { + "kind": "paged", "preset": cfg["preset"], "isl": cfg["isl"], + "page_tokens": cfg["page_tokens"], "layers": cfg["layers"], + "page_bytes": cfg["page_bytes"], "descs": cfg["descs"], + "req_bytes": cfg["req_bytes"], "_cfg": cfg, + } + for batch in allowed: + for op in ops: + make = None + if role == "initiator": + make = lambda op=op, batch=batch: [ + backend.make_paged(cfg, op, initiator_tables[r], target_tables[r]) + for r in range(batch)] + # pull lands on the initiator's pool; push on the target's. + # Every request in the burst is checked against its own tables. + verify_side = "initiator" if op == "pull" else "target" + table_pairs = [ + (initiator_tables[r], target_tables[r]) if op == "pull" + else (target_tables[r], initiator_tables[r]) + for r in range(batch)] + row = measure(make, {**base, "batch": batch}, op, verify_side, + table_pairs) + if row is not None: + rows.append(row) + print(f"[run_kv] {json.dumps(row)}", flush=True) + + for isl in isls: + preset = args.workload_name.removeprefix("kv-") + block_tokens = int(args.page_tokens.split()[0]) + cfg = kv_workload.plan_config(preset, args.precision, isl, block_tokens, + args.pool_slack) + nbytes = min(cfg["req_bytes"], bulk_bytes) + base = {"kind": "bulk", "preset": preset, "isl": isl, "page_tokens": None, + "layers": cfg["layers"], "page_bytes": None, "descs": 1, "batch": 1, + "req_bytes": nbytes, "_cfg": cfg} + for op in ops: + make = (lambda op=op, n=nbytes: [backend.make_bulk(n, op)]) if role == "initiator" else None + row = measure(make, base, op, verify_side="none", table_pairs=None) + if row is not None: + rows.append(row) + print(f"[run_kv] {json.dumps(row)}", flush=True) + + backend.teardown() + + gathered: list = [None, None] + dist.all_gather_object(gathered, rows if rank == 1 else None) + rows = gathered[1] or [] + hosts: list = [None, None] + dist.all_gather_object(hosts, socket.gethostname()) + all_ok = bool(rows) and all(r["verify"]["passed"] for r in rows) + + if rank == 0: + doc = { + "version": args.version, + "record_type": "case-attempt", + "generated_at": _dt.datetime.now().astimezone().isoformat(), + "identity": { + "allocation_factors": { + "run_attempt": os.environ.get("GITHUB_RUN_ATTEMPT"), + "run_id": os.environ.get("GITHUB_RUN_ID"), + "source_sha": os.environ.get("COLLECTIVEX_SOURCE_SHA") or os.environ.get("GITHUB_SHA"), + }, + "attempt_ordinal": int(os.environ.get("COLLX_ATTEMPT_ID", "1")), + "case_factors": {"case": {**case, "suite": args.suite}, "sku": args.runner}, + "case_id": args.case_id, + }, + "workload": { + "isl_ladder": isls, + "page_tokens": [int(v) for v in args.page_tokens.split()], + "batch_sizes": batches, + "ops": ops, + "preset": kv_workload.PRESETS[args.workload_name.removeprefix("kv-")], + }, + "measurement": { + "payload_unit": "request-kv-bytes", + "rows": rows, + "sampling": { + "reps_per_trial": args.reps, + "trials": args.trials, + "warmup_per_trial": args.warmup, + }, + }, + "implementation": { + "name": args.backend, + "fabric": args.fabric, + "library_version": getattr(backend, "library_version", None), + "maturity": getattr(backend, "maturity", "candidate"), + "nic_filter": getattr(backend, "nic_filter", None), + }, + "topology": { + "device_product": torch.cuda.get_device_name(device), + "gpus_per_node": args.gpus_per_node, + "hosts": hosts, + "nodes": 2, + "ranks_per_node": 1, + "scale_up_domain": args.scale_up_domain, + "scale_up_transport": args.scale_up_transport or None, + "topology_class": args.topology_class or None, + "world_size": world_size, + }, + "runtime": { + "framework": str(torch.__version__), + "vendor": "amd" if torch.version.hip else "nvidia", + }, + "provenance": { + "image": os.environ.get("COLLECTIVEX_IMAGE") or None, + "source_sha": os.environ.get("COLLECTIVEX_SOURCE_SHA") or os.environ.get("GITHUB_SHA"), + }, + "outcome": { + "reasons": [] if all_ok else ["transfer verification failed"], + "status": "success" if all_ok else "invalid", + }, + } + if args.out: + ep_harness._write_json_atomic(args.out, doc) + print(f"[run_kv] status={doc['outcome']['status']} rows={len(rows)}" + + (f" -> {args.out}" if args.out else ""), flush=True) + + flag = torch.tensor([int(all_ok)]) + dist.all_reduce(flag, op=dist.ReduceOp.MIN) + dist.barrier() + return 0 if int(flag.item()) else 3 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/collectivex/configs/kv_sweep.json b/collectivex/configs/kv_sweep.json new file mode 100644 index 0000000000..cab200b161 --- /dev/null +++ b/collectivex/configs/kv_sweep.json @@ -0,0 +1,53 @@ +{ + "version": 2, + "suite": "kv-transfer", + "workloads": { + "kv-dsv4": [ + "fp8" + ] + }, + "isl_ladder": [ + 2048, + 8192, + 32768, + 65536, + 131072, + 524288 + ], + "page_tokens": [ + 256 + ], + "batch_sizes": [ + 1, + 2, + 4, + 8, + 16, + 32 + ], + "ops": [ + "pull", + "push" + ], + "timing": { + "warmup_per_trial": 2, + "reps_per_trial": 16, + "trials_per_point": 3 + }, + "pool_slack": 2.0, + "seed": 67, + "scheduling": { + "default": { + "allocation_minutes": 210, + "run_timeout": 11400 + }, + "gb200": { + "allocation_minutes": 460, + "run_timeout": 25200 + }, + "gb300": { + "allocation_minutes": 690, + "run_timeout": 39600 + } + } +} diff --git a/collectivex/configs/platform_config.json b/collectivex/configs/platform_config.json index c1e288c384..b3fae34691 100644 --- a/collectivex/configs/platform_config.json +++ b/collectivex/configs/platform_config.json @@ -36,6 +36,7 @@ "launcher": "single-slurm", "backends": {"deepep-v2": [8, 16], "uccl-ep": [8], "nccl-ep": [8, 16]}, "ll_backends": {"deepep-v2": [8], "uccl-ep": [8], "nccl-ep": [8]}, + "kv_backends": {"nixl": ["rdma"], "mooncake": ["rdma"]}, "fabric": {"nic": "ConnectX-7 400G", "switch": "NVIDIA Quantum-2 QM9790 (25.6T, InfiniBand)"}, "operator": { "partition": "main", @@ -56,6 +57,7 @@ "launcher": "single-slurm", "backends": {"deepep-v2": [8, 16], "uccl-ep": [8, 16], "nccl-ep": [8, 16]}, "ll_backends": {"deepep-v2": [8, 16], "uccl-ep": [8], "nccl-ep": [8, 16]}, + "kv_backends": {"nixl": ["rdma"], "mooncake": ["rdma"]}, "fabric": {"nic": "8x ConnectX-7-class 400G InfiniBand rails (bare metal, gdrdrv)", "switch": "InfiniBand (model not surveyed)"}, "operator": { "partition": "batch_1", @@ -100,13 +102,18 @@ "launcher": "gb-nv", "backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16], "flashinfer-ep": [8, 16]}, "ll_backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16]}, - "fabric": {"nic": "MNNVL (scale-out not used)", "switch": "NVLink NVL72"}, + "kv_backends": {"nixl": ["rdma", "mnnvl"], "mooncake": ["rdma"]}, + "fabric": {"nic": "4x ConnectX-7 NDR400 InfiniBand (KV scale-out; EP stays MNNVL)", "switch": "NVLink NVL72 + Quantum-2"}, "operator": { "partition": "batch", "account": "benchmark", "storage_roots": [ "/mnt/lustre01/users-public/sa-shared" ] + }, + "network": { + "socket_ifname": "enP6p9s0np0", + "rdma_devices": "mlx5_0,mlx5_1,mlx5_3,mlx5_4" } }, "gb300": { @@ -120,13 +127,18 @@ "launcher": "gb-nv", "backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16], "flashinfer-ep": [8, 16]}, "ll_backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16]}, - "fabric": {"nic": "MNNVL (scale-out not used)", "switch": "NVLink NVL72"}, + "kv_backends": {"nixl": ["rdma", "mnnvl"], "mooncake": ["rdma"]}, + "fabric": {"nic": "4x ConnectX-8 XDR800 InfiniBand (KV scale-out; EP stays MNNVL)", "switch": "NVLink NVL72 + Quantum-X800"}, "operator": { "partition": "batch_1", "account": "benchmark", "squash_dir": "/data/home/sa-shared/collectivex/containers", "enroot_cache_path": "/data/home/sa-shared/collectivex/enroot-cache", "exclude_nodes": "im-gb300-r01-c003,im-gb300-r01-c005" + }, + "network": { + "socket_ifname": "enP22s22f0np0", + "rdma_devices": "mlx5_0,mlx5_1,mlx5_2,mlx5_3" } }, "mi300x": { @@ -181,6 +193,11 @@ "launcher": "mi-amds", "backends": {"mori": [8, 16], "uccl-ep": [8]}, "ll_backends": {"mori": [8]}, + "kv_backends": { + "mori-io": ["rdma"], + "mooncake": {"fabrics": ["rdma"], "ops": "push", "image": "rocm/atom-dev:nightly_202606101403", + "device": "rdma{gpu}", "pool_budget": 21474836480} + }, "fabric": {"nic": "Pollara 400GbE", "switch": "Arista 7060X6-64PE (Tomahawk5, 51.2T)"}, "network": { "socket_ifname": "eno0", diff --git a/collectivex/docs/methodology.md b/collectivex/docs/methodology.md index 3433e2ae27..fd69c8d15b 100644 --- a/collectivex/docs/methodology.md +++ b/collectivex/docs/methodology.md @@ -532,6 +532,110 @@ rates are named `rate_at_latency_percentile`: bytes or tokens divided by the mat percentile. They are lower-tail service rates at p99 latency, not p99 percentiles of an inverted rate distribution. +## KV-Cache Transfer Suite + +`suite: kv-transfer` measures the prefill→decode KV handoff of disaggregated serving as the +libraries engines ship perform it: one-sided RDMA over registered GPU pools, initiated by one side +(`pull` = READ, the vLLM NixlConnector shape; `push` = WRITE, the SGLang disagg shape). A leg is +2 nodes x 1 GPU — the per-worker pair — with rank 0 owning the target pool and rank 1 posting and +timing every transfer. Control is a gloo group (payload exchange + lockstep barriers); data never +rides it. + +The transferred object is a burst of `batch` concurrent requests' paged KV in the shape vLLM's +packed DSV4 NIXL path actually registers and posts: per cache group, the physical block is the +transfer unit, and ONE contiguous descriptor covers all of that group's layers for the block +(block-major `[block][layer]` layout, `packed_bytes = layers x page_bytes` per descriptor — +vLLM's connector asserts exactly this shape). Per request, `isl` tokens (2k through 512k) at the +production block size of 256 tokens, addressed through seed-keyed random block tables on BOTH +sides (batched requests slice disjoint ranges of one permutation, as live requests never alias +blocks) — fragmentation is real but block-granular, the post-fragmentation layout a fragmented +allocator hands a connector. The suite deliberately does NOT explode each (layer, page) into its +own descriptor: that shape inflates descriptor counts by ~2 orders of magnitude (~2.1M vs ~6.1k +per 512k-ISL request) and inverts backend and fabric conclusions on descriptor-bound lanes. Each +request is its own prepped transfer; a burst posts all of them, then awaits all, the way a decode +step admits several requests at once. + +Workload presets are transcribed from what vLLM actually allocates for the model class, region by +region (geometry validated against vLLM commit `32ad1400d7`). `kv-dsv4` is DeepSeek-V4-Pro as +vLLM serves it (MXFP4 checkpoints included — quantization covers weights, the cache layout is +architectural): every token-state is 584 B of content (448 B NoPE + 128 B RoPE + 8 B fp8 scale, +the `fp8_ds_mla` layout), and each block's page is padded to a 576 B multiple at PAGE granularity +(FlashMLA packing), not per state. The config's `compress_ratios` interleave 30 Compressed Sparse +Attention layers (4 tokens per state) with 31 Heavily Compressed Attention layers (128 tokens per +state); CSA layers add a 132 B/state lightning-indexer cache (128 fp8 + 4 scale bytes); and all +61 layers keep a 128-token sliding window whose block size is FIXED at 64 tokens because the +window shares its physical tensor with the CSA cache — its page equals the CSA page byte for +byte. HCA's 128-token states force the model block size to a multiple of 128 (the model fails +closed on anything else); vLLM serves DSV4 at 256. Precision is pinned fp8 because the dtype mix +is architectural, and the whole thing computes to a few percent of an equivalent dense GQA-bf16 +cache. Each lane also reports a `bulk` row — one single-descriptor transfer of the request's +total bytes per ISL — as the contiguous baseline the paged rows are read against: logical payload +over host-observed completion of one contiguous post, NOT a proven physical wire rate (backends +may split large operations internally). Two budgets shed a point's largest batches rather than +dropping the point, and the smallest batches always survive, so a single request stays measurable +everywhere and every point keeps a chartable batch ladder: a per-rank pool budget (64 GiB, sized +to the fleet's smallest HBM, a hard memory limit the batch floor never overrides — on the packed +geometry this is the budget that actually bites, at the 512k point's largest batches), and a +per-burst descriptor budget (posting time is linear in batch x descriptors on the per-descriptor +floor; the packed grid sits far under it, and it stays as the fail-closed guard for future +presets or smaller block sizes). + +Timing is host wall clock around post→completion — completion of a one-sided transfer is +host-visible and no local kernel participates, so CUDA events have nothing to bracket. Descriptor +build + handle creation are reported separately as `prep_ms` (engines amortize them through +prepped-handle reuse), never inside the timed transfer; because that amortization does NOT hold +for admissions with unique block tables and handle churn, each row also reports +`gbps_p50_incl_prep`, the cold-path rate with prep paid once per burst, so a lane whose prep +rivals its transfer time cannot hide it. Every point reports pooled trials x reps burst +percentiles, GB/s at p50 (burst-aggregate), and per-request completion marks: `request_ms` is +each individual request's host-observed completion offset from its burst's start (waits drain in +posting order, so each mark upper-bounds that request's true completion) — the per-request +latency distribution, distinct from the burst quantities, which are capacity numbers. A +verification verdict closes each point: the destination pool is pattern-checked after `pull` on +the initiator and after `push` on the target (an offset-derived byte pattern makes any block's +expected contents computable from its offset alone), covering every request in the burst against +its own block tables — concurrent same-session requests are exactly where corruption would hide, +so a passing request 0 is never taken as evidence for the others. Both pools are repainted +between points. A failed verify flips the document `invalid` and the leg red. + +A registry backend can carry restrictions: `ops` when a fabric serves one direction only +(mooncake on mi355x runs `push` — AMD's atom-dev build moves WRITE at healthy rates over the +GPU-paired Pollara NIC, while upstream ionic RDMA READ completes with retry-exceeded and one +failed READ poisons the engine, which is also why ATOM's production connector is write-only), +`image` when the build ships only inside a specific image, and `device` for engine NIC +filters (`{gpu}` expands to the physical GPU index; registering GPU memory on a non-paired +NIC fails and cross-rail pairs are unroutable), and `pool_budget` when the engine cannot +register the default pool on that pool's NICs (mooncake on the mi355x ionic NICs fails +`ibv_reg_mr` with ENOMEM past ~20 GiB while mori-io registers the same pool; points shed their +largest batches to fit). No kv backend is enabled on b300: the current b300 pool is the AWS +EFA cluster, and the one-rail rows measured on the earlier RoCE b300 pool do not carry over. The summary's `op` column names the +measured direction. + +Fabrics are a case dimension. `rdma` runs on torch (cudaMalloc) pools. `mnnvl` allocates the +pools with cuMem FABRIC handles (kv_pool.FabricPool; needs a live nvidia-imex domain), because +UCX's cross-node cuda_ipc only engages on fabric-mappable memory: on cudaMalloc pools the flag +is silently inert and the transfer rides the IB rails with byte-identical numbers. On GB200 the +mnnvl lane pays ~3.9 µs per descriptor copy where the IB lane does not, so the two lanes invert +with descriptor count: mnnvl leads on coalesced transfers and falls behind the rails as the +descriptor list grows. Per-lane bandwidth figures for the packed geometry are what the published +rows carry; figures measured under the retired per-(layer, page) geometry are not comparable and +are not restated here. + +Other lane facts, measured on the metal: single-WR bulk transfers above the provider's max +message size must be split (the MoRI adapter caps WRs at 1 GiB); and Mooncake is NVIDIA-only at +the binary level (the wheel links libcuda.so.1 at import; measured failing on mi355x). + + ## Correctness + +Scheduling is data. `configs/kv_sweep.json` holds the grid and, per pool, the allocation and the +per-case hang guard (`scheduling`): a leg runs for hours, gb200's mnnvl descriptor floor alone +measured ~285 minutes on the five-rung grid, and gb300 paces ~1.8x gb200 at ISL >= 131072 over +mnnvl, so gb200 asks 460 minutes with a 420-minute guard and gb300 690 with 660; every other pool +asks 210 with 190. The guard fires before the allocation dies, so a slow case is a clean per-case +kill, and each shard carries a GitHub job ceiling above its allocation. The asks stay 2 nodes x 1 +GPU, short enough to backfill on a contended pool. KV legs run only when a dispatch names +`kv-transfer` in `suites`. + ## Correctness An implementation-independent oracle uses an expert-specific deterministic transform so wrong expert diff --git a/collectivex/launchers/launch_gb-nv.sh b/collectivex/launchers/launch_gb-nv.sh index 3ff131e226..5bd04686aa 100644 --- a/collectivex/launchers/launch_gb-nv.sh +++ b/collectivex/launchers/launch_gb-nv.sh @@ -24,14 +24,18 @@ if [ "$PRODUCT" = gb200 ]; then default_time=30; else default_time=90; fi TIME_MIN="${COLLX_TIME:-$default_time}" IMAGE="$COLLX_IMAGE" case "$COLLX_BENCH" in - deepep-v2 | nccl-ep | flashinfer-ep | swap-blocks) ;; + deepep-v2 | nccl-ep | flashinfer-ep | swap-blocks | nixl | mooncake) ;; *) collx_die "unsupported $PRODUCT backend: $COLLX_BENCH" ;; esac collx_require_vars COLLX_IMAGE COLLX_IMAGE_PLATFORM COLLX_PARTITION COLLX_ACCOUNT COLLX_SQUASH_DIR COLLX_STAGE_DIR [ "$PRODUCT" != gb300 ] || collx_require_vars COLLX_ENROOT_CACHE_PATH PARTITION="$COLLX_PARTITION"; ACCOUNT="$COLLX_ACCOUNT"; SQUASH_DIR="$COLLX_SQUASH_DIR" [ -z "${COLLX_ENROOT_CACHE_PATH:-}" ] || export ENROOT_CACHE_PATH="$COLLX_ENROOT_CACHE_PATH" -export NCCL_CUMEM_ENABLE=1 NCCL_MNNVL_ENABLE=1 MC_FORCE_MNNVL=1 +export NCCL_CUMEM_ENABLE=1 NCCL_MNNVL_ENABLE=1 +# Mooncake is MC_FORCE_MNNVL's only reader, and it makes the engine install ONLY its cross-node +# NVLink transport, which cannot open another host's segments in the pinned wheel +# (cudaIpcOpenMemHandle: invalid resource handle). The mooncake kv row declares the rdma lane. +[ "$COLLX_BENCH" = mooncake ] || export MC_FORCE_MNNVL=1 collx_apply_network_profile "$NODES" "$COLLX_TRANSPORT" collx_log "$PRODUCT nodes=$NODES x ${GPN}gpu world=$NGPUS bench=$COLLX_BENCH" @@ -48,6 +52,13 @@ allocation=(--partition="$PARTITION" --account="$ACCOUNT" --nodes="$NODES" [ -z "${COLLX_EXCLUDE_NODES:-}" ] || allocation+=(--exclude="$COLLX_EXCLUDE_NODES") collx_salloc_jobid "${allocation[@]}" [ -n "$JOB_ID" ] || collx_die "no JOB_ID from salloc" +# The kv rdma legs are the only gb-nv shards that leave the NVL domain; prove their pinned socket +# interface and HCAs on the allocation as every other scale-out launcher does. +if [ "$COLLX_TRANSPORT" != mnnvl ] \ + && ! collx_validate_network_profile_on_job "$JOB_ID" "$NODES" "$COLLX_TRANSPORT"; then + collx_log_tail "${COLLX_NETWORK_PROFILE_LOG:-}" + collx_die "allocated nodes failed the network profile" +fi SQUASH_FILE="$(collx_ensure_squash_on_job "$JOB_ID" "$SQUASH_DIR" "$IMAGE")" diff --git a/collectivex/launchers/launch_mi-amds.sh b/collectivex/launchers/launch_mi-amds.sh index b88d5f6230..4b56a51ad2 100644 --- a/collectivex/launchers/launch_mi-amds.sh +++ b/collectivex/launchers/launch_mi-amds.sh @@ -28,7 +28,7 @@ EXCLUDE_NODES="${COLLX_EXCLUDE_NODES:-}" NODELIST="${COLLX_NODELIST:-}" MOUNT_DIR=/ix case "$COLLX_BENCH" in - mori | uccl-ep | swap-blocks) ;; + mori | uccl-ep | swap-blocks | mori-io | mooncake) ;; *) collx_die "unsupported AMD backend: $COLLX_BENCH" ;; esac diff --git a/collectivex/launchers/launch_single-slurm.sh b/collectivex/launchers/launch_single-slurm.sh index 6cae774fc7..6c848dbfda 100644 --- a/collectivex/launchers/launch_single-slurm.sh +++ b/collectivex/launchers/launch_single-slurm.sh @@ -40,7 +40,7 @@ collx_set_placement 1 8 8 nvlink TIME_MIN="${COLLX_TIME:-$DEFAULT_TIME}" IMAGE="$COLLX_IMAGE" case "$COLLX_BENCH" in - deepep-v2 | uccl-ep | nccl-ep | swap-blocks) ;; + deepep-v2 | uccl-ep | nccl-ep | swap-blocks | nixl | mooncake) ;; *) collx_die "unsupported $RUNNER backend: $COLLX_BENCH" ;; esac diff --git a/collectivex/runtime/common.sh b/collectivex/runtime/common.sh index 6a7f59e3bf..1f1603f571 100644 --- a/collectivex/runtime/common.sh +++ b/collectivex/runtime/common.sh @@ -72,6 +72,9 @@ collx_set_placement() { export COLLX_GPUS_PER_NODE="$GPN" COLLX_SCALE_UP_DOMAIN="$SCALE_UP_DOMAIN" COLLX_TRANSPORT="$4" [ "$4" = mnnvl ] || [ "$NODES" -le 1 ] || COLLX_TRANSPORT="$4-rdma" + # A shard that names the rdma fabric on an MNNVL rack (the kv-transfer rdma legs) leaves the + # NVL domain, so it takes the same network profile and validation as any other scale-out. + [ "$4:${COLLX_FABRIC:-}" != mnnvl:rdma ] || [ "$NODES" -le 1 ] || COLLX_TRANSPORT=mnnvl-rdma export COLLX_TRANSPORT } @@ -481,7 +484,7 @@ export RANK="$SLURM_PROCID" WORLD_SIZE="$SLURM_NTASKS" export LOCAL_RANK="$SLURM_LOCALID" LOCAL_WORLD_SIZE="$COLLX_GPUS_PER_NODE" # config.py case-args leads every argv with the suite's entrypoint. [ "${1:-}" = --entrypoint ] || exit 67 -case "${2:-}" in run_ep|run_swap_blocks) ;; *) exit 67 ;; esac +case "${2:-}" in run_ep|run_swap_blocks|run_kv) ;; *) exit 67 ;; esac entry="$2"; shift 2 exec python3 "bench/$entry.py" "$@" BASH diff --git a/collectivex/runtime/config.py b/collectivex/runtime/config.py index e3b0157586..e70fab14b7 100644 --- a/collectivex/runtime/config.py +++ b/collectivex/runtime/config.py @@ -170,11 +170,42 @@ def _swap_argv(case: dict, version: object, runner: str) -> list[str]: return argv + _flag_pairs(case, "layout seed device max_payload_bytes warmup iterations") +def _kv_argv(case: dict, version: object, runner: str) -> list[str]: + argv = [ + "--backend", str(case["backend"]), + "--workload-name", str(case["workload"]), + "--precision", str(case["precision"]), + "--fabric", str(case["mode"]), + "--isl-ladder", str(case["isl_ladder"]), + "--page-tokens", str(case["page_tokens"]), + "--batch-sizes", str(case["batch_sizes"]), + "--kv-device", str(case.get("kv_device", "")), + "--ops", str(case["ops"]), + "--warmup", str(case["warmup"]), + "--reps", str(case["reps"]), + "--trials", str(case["trials"]), + "--pool-slack", str(case["pool_slack"]), + "--seed", str(case["seed"]), + "--runner", runner, + "--case-id", str(case["case_id"]), + "--suite", str(case["suite"]), + "--gpus-per-node", str(case["gpus_per_node"]), + "--scale-up-domain", str(case["scale_up_domain"]), + "--scale-up-transport", str(case["scale_up_transport"]), + "--topology-class", str(case["topology_class"]), + "--version", str(version), + ] + if "pool_budget" in case: + argv += ["--pool-budget", str(case["pool_budget"])] + return argv + + # suite -> (bench/.py, argv codec, output flag). The rank wrapper in # runtime/common.sh execs the entrypoint the leading --entrypoint pair names. _SUITES = { "ep-core": ("run_ep", _ep_argv, "--out"), "swap-blocks": ("run_swap_blocks", _swap_argv, "--output"), + "kv-transfer": ("run_kv", _kv_argv, "--out"), } diff --git a/collectivex/runtime/prepare_backend.sh b/collectivex/runtime/prepare_backend.sh index 95427b9c20..33a159f153 100644 --- a/collectivex/runtime/prepare_backend.sh +++ b/collectivex/runtime/prepare_backend.sh @@ -558,6 +558,38 @@ FICHECK [ "$rc" -eq 0 ] || { collx_log "ERROR: FlashInfer EP one-sided A2A unavailable in this image"; return 1; } } +# The kv-transfer wheels install into the named container, which persists for the job, so one +# install here serves every case srun. Noble-based images mark python externally managed +# (PEP 668); the retry flag is the uccl-prep pattern (older pips never refuse, so never reach it). +pip_install() { + python3 -m pip install -q --disable-pip-version-check --no-input "$@" \ + || python3 -m pip install -q --disable-pip-version-check --no-input --break-system-packages "$@" +} + +# nixl-cuXX directly: the `nixl` meta package depends on BOTH cu12 and cu13 variants, and an +# unpinned install under the image's stale pip resolved 1.0.1. ROCm images need nothing: +# sglang-rocm bundles nixl-cu12 with a ROCm-built UCX. +nixl_prepare() { + python3 -c "import nixl" 2>/dev/null && return 0 + pip_install 'nixl-cu13==1.3.2' \ + || { collx_log "ERROR: nixl wheel install failed"; return 1; } + python3 -c "import nixl" \ + || { collx_log "ERROR: nixl import failed after install"; return 1; } +} + +# ROCm builds ship inside the image (upstream wheels link libcuda.so.1; AMD's atom-dev image +# carries a working build), so an importable mooncake.engine wins. Otherwise install the pinned +# CUDA wheel; it links libcudart.so.12, which the adapter dlopens from the runtime package at +# import, so no LD_LIBRARY_PATH seam is needed. +mooncake_prepare() { + if python3 -c "import mooncake.engine" 2>/dev/null; then + collx_log "mooncake provided by the image" + return 0 + fi + pip_install 'mooncake-transfer-engine==0.3.12.post1' nvidia-cuda-runtime-cu12 \ + || { collx_log "ERROR: mooncake wheel install failed"; return 1; } +} + main() { collx_apply_network_profile "${COLLX_NODES:-1}" "${COLLX_TRANSPORT:-}" || return 1 validate_container_network || return 1 @@ -575,6 +607,12 @@ main() { python3 -c "from vllm._custom_ops import swap_blocks" \ || { collx_log "ERROR: vLLM swap_blocks import failed"; return 1; } ;; + nixl) nixl_prepare || return 1 ;; + mooncake) mooncake_prepare || return 1 ;; + mori-io) + python3 -c "import mori.io" \ + || { collx_log "ERROR: MoRI-IO import failed"; return 1; } + ;; *) collx_log "ERROR: unknown backend preparation request" return 1 diff --git a/collectivex/summarize.py b/collectivex/summarize.py index 2d8f9db36a..de4409264a 100644 --- a/collectivex/summarize.py +++ b/collectivex/summarize.py @@ -109,18 +109,87 @@ def percentile(block: str, name: str) -> float | str: ) -def render(documents: list[dict]) -> str: - documents = sorted(documents, key=_identity) +KV_SUITE = "kv-transfer" + + +def _invalid_banner(documents: list[dict]) -> list[str]: + # The leg is already red (each benchmark entrypoint returns nonzero on a non-success + # outcome); call the count out loudly so it is not lost in the per-row table. invalid = [d for d in documents if d["outcome"]["status"] != "success"] - lines = ["## CollectiveX EP results", ""] - if invalid: - # The leg is already red (ep_harness.run_sweep returns nonzero on a non-success - # outcome); call the count out loudly so it is not lost in the per-row table. + if not invalid: + return [] + return [f"> **{len(invalid)} of {len(documents)} outcome(s) INVALID** — " + "the leg fails; see the outcome column below.", ""] + + +def _kv_cell(rows: list[dict], kind: str, op: str, batch: str = "min"): + """The largest-ISL row of a (kind, op) family -- the bandwidth-bound point -- at its smallest + or largest measured batch. Paged cells read the largest measured block size (the production + one when several ran).""" + matching = [r for r in rows if r.get("kind") == kind and r.get("op") == op] + if kind == "paged" and matching: + block = max(r["page_tokens"] for r in matching) + matching = [r for r in matching if r["page_tokens"] == block] + if not matching: + return "-", "-" + isl = max(r["isl"] for r in matching) + pick = min if batch == "min" else max + row = pick((r for r in matching if r["isl"] == isl), key=lambda r: r.get("batch", 1)) + return row["gbps_p50"], row["latency_ms"]["p50"] + + +def render_kv(documents: list[dict]) -> str: + """kv-transfer table: paged bandwidth at the bandwidth-bound ISL plus the contiguous + baseline, and the paged latency.""" + def key(document): + factors = document["identity"]["case_factors"] + case = factors["case"] + return factors["sku"], case["backend"], case["mode"], case["workload"], case["precision"] + + lines = ["## CollectiveX KV-transfer results", "", *_invalid_banner(documents), + "| ver | sku | backend | fabric | workload | precision | outcome | op " + "| paged GB/s b1 | paged GB/s bmax | contig GB/s | paged ms b1 |", + "|--:|---|---|---|---|---|---|---|--:|--:|--:|--:|"] + for document in sorted(documents, key=key): + sku, backend, fabric, workload, precision = key(document) + rows = document["measurement"]["rows"] + # Cells read the pull lane when measured, else the push lane (a backend may serve one + # direction only, e.g. mooncake on Pollara, where upstream ionic RDMA READ is broken); + # the op column names which lane the row's numbers come from. + op = next((candidate for candidate in ("pull", "push") + if _kv_cell(rows, "paged", candidate)[0] != "-"), "pull") + paged_gbps, paged_ms = _kv_cell(rows, "paged", op) + paged_bmax, _ = _kv_cell(rows, "paged", op, batch="max") + bulk_gbps, _ = _kv_cell(rows, "bulk", op) lines.append( - f"> **{len(invalid)} of {len(documents)} outcome(s) INVALID** — " - "the leg fails; see the outcome column below." + f"| {document['version']} | {sku} | `{backend}` | {fabric} | {workload} | " + f"{precision} | {document['outcome']['status']} | {op} | {paged_gbps} | " + f"{paged_bmax} | {bulk_gbps} | {paged_ms} |" ) - lines.append("") + lines.append( + "\n> Paged rows move requests' KV as vLLM's packed block-major descriptor lists (one " + "contiguous descriptor per physical block per cache group) over randomized block " + "tables; b1/bmax = requests posted per burst (GB/s is burst-aggregate); contig is the " + "single-descriptor contiguous baseline (host-observed goodput, not proven wire " + "utilization); op names the measured direction. GB/s at the largest ISL " + "(bandwidth-bound)." + ) + return "\n".join(lines) + + +def render(documents: list[dict]) -> str: + """One table per suite present; the EP table also renders when nothing was found.""" + kv = [d for d in documents if d["identity"]["case_factors"]["case"].get("suite") == KV_SUITE] + ep = [d for d in documents if d not in kv] + parts = [render_ep(ep)] if ep or not kv else [] + if kv: + parts.append(render_kv(kv)) + return "\n\n".join(parts) + + +def render_ep(documents: list[dict]) -> str: + documents = sorted(documents, key=_identity) + lines = ["## CollectiveX EP results", "", *_invalid_banner(documents)] lines += [ "| ver | sku | backend | mode | precision | suite | phase | routing | ep | topo " "| wire | outcome | T* | p50* us | p99* us | min50 us | skew us |", diff --git a/collectivex/sweep_matrix.py b/collectivex/sweep_matrix.py index af65520911..9d7526db4f 100644 --- a/collectivex/sweep_matrix.py +++ b/collectivex/sweep_matrix.py @@ -30,8 +30,9 @@ def _load_config(name: str) -> dict[str, Any]: SWEEP = _load_config("sweep.json") SWAP_SWEEP = _load_config("swap_sweep.json") +KV_SWEEP = _load_config("kv_sweep.json") PLATFORMS = _load_config("platform_config.json")["platforms"] -SUITES = ("ep", "swap-blocks") +SUITES = ("ep", "swap-blocks", "kv-transfer") SWEEP_BACKENDS = tuple(dict.fromkeys( @@ -179,6 +180,108 @@ def _swap_shard(sku: str, profile_name: str) -> tuple[list[dict[str, Any]], dict return _runnable(sku, cases), shard +def _kv_backend_spec(value: Any) -> dict[str, Any]: + """A registry kv_backends value is a fabric list (the library runs the full sweep) or an + object that restricts it: `ops` (a backend that cannot serve one direction on a fabric, e.g. + mooncake on Pollara, where ionic RDMA READ is broken upstream), `image` (a backend shipped + only inside a specific image, e.g. AMD's atom-dev mooncake build), `device` (an engine NIC + filter template; `{gpu}` expands to the physical GPU index at runtime), and `pool_budget` + (bytes, where the engine cannot register the default pool on that pool's NICs).""" + if isinstance(value, list): + value = {"fabrics": value} + return { + "fabrics": value["fabrics"], + "ops": value.get("ops") or " ".join(KV_SWEEP["ops"]), + "image": value.get("image"), + "device": value.get("device") or "", + "pool_budget": value.get("pool_budget"), + } + + +def _kv_shards(sku: str, selected_precisions: set[str]) -> tuple[list, list[dict[str, Any]]]: + """The kv-transfer shards one pool runs: one per (backend, fabric) its registry enables. + + A KV leg is 2 nodes x 1 GPU: the per-worker transfer pair an engine actually forms, not an + allocation-wide collective. A pool with no ``kv_backends`` entry emits nothing -- absence is + not-yet-enabled, mirroring ll_backends. The allocation and the per-case hang guard come from + kv_sweep.json's scheduling block: KV grids run for hours, and gb300 paces ~1.8x gb200 at the + top ISLs over mnnvl, so both are sized per pool rather than fleet-wide. + """ + platform = PLATFORMS[sku] + timing = KV_SWEEP["timing"] + scheduling = KV_SWEEP["scheduling"].get(sku, KV_SWEEP["scheduling"]["default"]) + requested, shards = [], [] + for backend, raw in sorted(platform.get("kv_backends", {}).items()): + spec = _kv_backend_spec(raw) + for fabric in spec["fabrics"]: + cases = [] + # A workload's dtype mix can be architectural (dsv4's fp8 slots), so the sweep config + # maps each workload to its precisions; a test pins the map to kv_workload's PRESETS. + for workload, workload_precisions in KV_SWEEP["workloads"].items(): + for precision in workload_precisions: + if selected_precisions and precision not in selected_precisions: + continue + case = { + "suite": KV_SWEEP["suite"], + "workload": workload, + "backend": backend, + "routing": "paged", + "precision": precision, + "phase": "xfer", + "ep": 2, + "mode": fabric, + "isl_ladder": " ".join(map(str, KV_SWEEP["isl_ladder"])), + "page_tokens": " ".join(map(str, KV_SWEEP["page_tokens"])), + "batch_sizes": " ".join(map(str, KV_SWEEP["batch_sizes"])), + "ops": spec["ops"], + "kv_device": spec["device"], + "pool_slack": KV_SWEEP["pool_slack"], + "seed": KV_SWEEP["seed"], + "warmup": timing["warmup_per_trial"], + "reps": timing["reps_per_trial"], + "trials": timing["trials_per_point"], + "nodes": 2, + "gpus_per_node": 1, + "scale_up_domain": platform["scale_up_domain"], + "scale_up_transport": platform["scale_up_transport"], + "topology_class": f"{platform['product']}-kv-{fabric}", + } + if spec["pool_budget"]: + case["pool_budget"] = spec["pool_budget"] + case["case_id"] = ep_harness.case_id(sku, case) + cases.append(case) + if not cases: + continue + requested += [ + {"sku": sku, "case": case, "disposition": "runnable", "reason": None, + "detail": None} + for case in cases + ] + shard = { + "id": f"{sku}-kv-{backend}-{fabric}", + "sku": sku, + "runner": platform.get("runner_label", sku), + "backend": backend, + "suite": KV_SWEEP["suite"], + "mode": fabric, + "fabric": fabric, + "launcher": platform["launcher"], + "nodes": 2, + "gpus_per_node": 1, + "scale_up_domain": platform["scale_up_domain"], + "allocation_minutes": scheduling["allocation_minutes"], + "run_timeout": scheduling["run_timeout"], + # The GitHub job must outlive the allocation, or it cancels a healthy shard + # before the launcher's own guards act. + "job_timeout_minutes": scheduling["allocation_minutes"] + 30, + "cases": cases, + } + if spec["image"]: + shard["image"] = spec["image"] + shards.append(shard) + return requested, shards + + def resolve_matrix( backend: str = "all", only_sku: str = "", @@ -193,8 +296,10 @@ def resolve_matrix( selected_suites = _comma_subset("suites", suites, SUITES) if not selected_suites: raise SystemExit("--suites selects no suite") - if "ep" not in selected_suites and (backend != "all" or ep_sizes or precisions or modes): + if "ep" not in selected_suites and (backend != "all" or ep_sizes or modes): raise SystemExit("EP filters need the ep suite in --suites") + if not selected_suites & {"ep", "kv-transfer"} and precisions: + raise SystemExit("--precisions needs the ep or kv-transfer suite in --suites") if swap_profile not in SWAP_SWEEP["profiles"]: raise SystemExit( f"unknown --swap-profile {swap_profile!r}; have {sorted(SWAP_SWEEP['profiles'])}" @@ -330,6 +435,12 @@ def resolve_matrix( requested, shard = _swap_shard(sku, swap_profile) requested_cases += requested shards_by_sku.setdefault(sku, []).append(shard) + if "kv-transfer" in selected_suites: + for sku in selected_skus: + requested, kv_shards = _kv_shards(sku, selected_precisions) + requested_cases += requested + if kv_shards: + shards_by_sku.setdefault(sku, []).extend(kv_shards) include = [ shards_by_sku[sku][index] for index in range(max(map(len, shards_by_sku.values()), default=0)) diff --git a/collectivex/tests/test_kv_suite.py b/collectivex/tests/test_kv_suite.py new file mode 100644 index 0000000000..cdd90fa410 --- /dev/null +++ b/collectivex/tests/test_kv_suite.py @@ -0,0 +1,534 @@ +#!/usr/bin/env python3 +"""The kv-transfer suite's scheduling, argv codec, and summary contracts. + +Three seams keep KV legs honest end to end: sweep_matrix must emit kv shards +only for SKUs whose registry carries `kv_backends` (and must not perturb the EP +matrix at all); config.py must encode a kv case into run_kv argv behind the +`--entrypoint` marker the rank wrapper dispatches on; and summarize must render +kv documents in their own table instead of crashing the EP renderer. +""" +from __future__ import annotations + +import io +import sys +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +sys.path[:0] = [str(ROOT), str(ROOT / "bench"), str(ROOT / "runtime")] + +import config as runtime_config # noqa: E402 +import ep_harness # noqa: E402 +import summarize # noqa: E402 +import sweep_matrix # noqa: E402 + + +class KVMatrix(unittest.TestCase): + def test_kv_shards_only_where_the_registry_enables_them(self): + shards = sweep_matrix.resolve_matrix(suites="kv-transfer")["include"] + self.assertTrue(shards, "registry carries kv_backends but no shard resolved") + enabled = { + sku for sku, platform in sweep_matrix.PLATFORMS.items() + if platform.get("kv_backends") + } + self.assertEqual({shard["sku"] for shard in shards}, enabled) + default = sweep_matrix.KV_SWEEP["scheduling"]["default"] + for shard in shards: + self.assertEqual((shard["nodes"], shard["gpus_per_node"]), (2, 1)) + self.assertEqual(shard["launcher"], sweep_matrix.PLATFORMS[shard["sku"]]["launcher"]) + scheduling = sweep_matrix.KV_SWEEP["scheduling"].get(shard["sku"], default) + self.assertEqual(shard["allocation_minutes"], scheduling["allocation_minutes"]) + self.assertEqual(shard["run_timeout"], scheduling["run_timeout"]) + # The guard fires inside the allocation, and the GitHub job outlives it. + self.assertLess(shard["run_timeout"], shard["allocation_minutes"] * 60) + self.assertGreater(shard["job_timeout_minutes"], shard["allocation_minutes"]) + # the suite sweeps DeepSeek-V4-Pro's shape; its dtype mix is + # architectural, so one workload x one precision + self.assertEqual( + {(c["workload"], c["precision"]) for c in shard["cases"]}, + {("kv-dsv4", "fp8")}) + if shard["sku"] == "mi355x" and shard["backend"] == "mooncake": + # AMD's atom-dev build: push-only (upstream ionic RDMA READ is + # broken), GPU-paired NIC filter, shipped inside a pinned image, + # and a pool the ionic NICs can register. + self.assertEqual({c["ops"] for c in shard["cases"]}, {"push"}) + self.assertEqual({c["kv_device"] for c in shard["cases"]}, {"rdma{gpu}"}) + self.assertTrue(shard["image"].startswith("rocm/atom-dev:")) + self.assertEqual({c["pool_budget"] for c in shard["cases"]}, {20 << 30}) + else: + self.assertEqual({c["ops"] for c in shard["cases"]}, {"pull push"}) + self.assertNotIn("image", shard) + self.assertEqual({c["kv_device"] for c in shard["cases"]}, {""}) + self.assertFalse([c for c in shard["cases"] if "pool_budget" in c]) + for case in shard["cases"]: + self.assertEqual(case["suite"], "kv-transfer") + self.assertEqual(case["ep"], 2) + self.assertEqual( + case["case_id"], ep_harness.case_id(shard["sku"], case)) + + def test_kv_never_perturbs_the_ep_matrix(self): + ep_only = sweep_matrix.resolve_matrix() + both = sweep_matrix.resolve_matrix(suites="ep,kv-transfer") + self.assertEqual( + ep_only["include"], + [s for s in both["include"] if s.get("suite") != "kv-transfer"]) + self.assertEqual( + ep_only["requested_cases"], + [c for c in both["requested_cases"] if c["case"].get("suite") != "kv-transfer"]) + + def test_unknown_suite_fails_closed(self): + with self.assertRaises(SystemExit): + sweep_matrix.resolve_matrix(suites="kv-transfr") + + def test_precision_filter_applies_to_kv(self): + # dsv4 is fp8-only, so a bf16-scoped dispatch has no kv legs at all. + matrix = sweep_matrix.resolve_matrix(suites="kv-transfer", precisions="bf16") + self.assertEqual(matrix["include"], []) + matrix = sweep_matrix.resolve_matrix(suites="kv-transfer", precisions="fp8") + self.assertTrue(matrix["include"]) + + def test_ep_only_filters_need_the_ep_suite(self): + for options in ({"backend": "deepep-v2"}, {"modes": "normal"}, {"ep_sizes": "8"}): + with self.subTest(options=options), self.assertRaises(SystemExit): + sweep_matrix.resolve_matrix(suites="kv-transfer", **options) + + def test_the_workload_map_matches_the_presets(self): + # sweep_matrix stays stdlib-only for the bare-runner matrix step, so it cannot import + # kv_workload; pin its workload -> precision map to the workload model instead. + import kv_workload + + for workload, precisions in sweep_matrix.KV_SWEEP["workloads"].items(): + preset = workload.removeprefix("kv-") + with self.subTest(workload=workload): + self.assertIn(preset, kv_workload.PRESETS) + for precision in precisions: + kv_workload.plan_config(preset, precision, 2048, 256) + + def test_every_kv_backend_passes_its_launcher_identity_gate(self): + # The launchers collx_die on unknown COLLX_BENCH values before anything + # runs; a registry kv backend its launcher rejects is a dead shard + # (this exact gap shipped once — every kv leg died at the gate). + launchers = Path(sweep_matrix.__file__).parent / "launchers" + for sku, platform in sweep_matrix.PLATFORMS.items(): + source = (launchers / f"launch_{platform['launcher']}.sh").read_text() + for backend in platform.get("kv_backends", {}): + with self.subTest(sku=sku, backend=backend): + self.assertRegex(source, rf"(^|[ |]){backend}( |\)|\s*\|)") + + +class KVArgvCodec(unittest.TestCase): + def _shard(self, backend="nixl"): + shards = sweep_matrix.resolve_matrix(suites="kv-transfer")["include"] + return next(shard for shard in shards if shard["backend"] == backend) + + @staticmethod + def _captured_argv(case, sku): + class _Stdout: + buffer = io.BytesIO() + + saved, sys.stdout = sys.stdout, _Stdout() + try: + runtime_config._emit_argv(case, 1, sku, "20260807", 0) + return sys.stdout.buffer.getvalue().decode().split("\0")[:-1] + finally: + sys.stdout = saved + + def _parsed(self, shard): + import argparse + + import run_kv + + argv = self._captured_argv(shard["cases"][0], shard["sku"]) + self.assertEqual(argv[:2], ["--entrypoint", "run_kv"]) + parser = argparse.ArgumentParser() + parser.add_argument("--backend", required=True, choices=["nixl", "mori-io", "mooncake"]) + run_kv.add_kv_args(parser) + return parser.parse_args(argv[2:]) + + def test_kv_case_round_trips_through_the_run_kv_parser(self): + import run_kv + + shard = self._shard() + case = shard["cases"][0] + args = self._parsed(shard) + self.assertEqual((args.backend, args.workload_name, args.precision, args.fabric), + (case["backend"], case["workload"], case["precision"], case["mode"])) + self.assertEqual((args.warmup, args.reps, args.trials), + (case["warmup"], case["reps"], case["trials"])) + self.assertEqual((args.batch_sizes, args.kv_device, args.ops), + (case["batch_sizes"], case["kv_device"], case["ops"])) + self.assertEqual(args.case_id, case["case_id"]) + self.assertEqual(args.pool_budget, run_kv.POOL_BUDGET) + self.assertEqual(args.out, f"results/{case['case_id']}_20260807-c000.json") + # run_kv recomputes the identity from the same factors and refuses a mismatch. + self.assertEqual(ep_harness.case_id(shard["sku"], run_kv.kv_case(args)), case["case_id"]) + + def test_a_pool_budget_reaches_run_kv(self): + shard = next(s for s in sweep_matrix.resolve_matrix(suites="kv-transfer")["include"] + if s["cases"][0].get("pool_budget")) + self.assertEqual(self._parsed(shard).pool_budget, shard["cases"][0]["pool_budget"]) + + +class _StubDist: + """all_gather_object across a simulated 2-rank pair.""" + + def __init__(self, other_value): + self.other = other_value + + def all_gather_object(self, out, mine): + out[0], out[1] = mine, self.other + + +class VerdictExchange(unittest.TestCase): + """Bulk rows have no verifying side; that path crashed on the metal (gb200 + smoke 22840: StopIteration on both ranks) before this contract existed.""" + + def test_the_verifying_rank_supplies_the_verdict(self): + import run_kv + + verdict = run_kv.exchange_verdict( + _StubDist(None), "initiator", "initiator", lambda: (False, "bad page")) + self.assertEqual(verdict, {"passed": False, "detail": "bad page"}) + + def test_the_other_rank_receives_it(self): + import run_kv + + verdict = run_kv.exchange_verdict( + _StubDist({"passed": False, "detail": "bad page"}), "target", "initiator", + lambda: (True, "")) + self.assertEqual(verdict["passed"], False) + + def test_a_row_with_no_verifying_side_passes_without_a_gather_crash(self): + import run_kv + + verdict = run_kv.exchange_verdict( + _StubDist(None), "initiator", "none", + lambda: (_ for _ in ()).throw(AssertionError("must not verify"))) + self.assertEqual(verdict, {"passed": True, "detail": ""}) + + +class UCXSelectors(unittest.TestCase): + """run_kv pins UCX to the operator's validated RDMA selectors — UCX + auto-selection is a wrong-fabric trap (b200-nscale's aux quad-port card, + b300's storage IB) — while explicit UCX_* values always win.""" + + def test_registry_selectors_map_to_ucx(self): + import run_kv + + env = {"COLLX_RDMA_DEVICES": "mlx5_0,mlx5_10", "COLLX_IB_GID_INDEX": "3"} + run_kv.export_ucx_selectors(env) + self.assertEqual(env["UCX_NET_DEVICES"], "mlx5_0:1,mlx5_10:1") + self.assertEqual(env["UCX_IB_GID_INDEX"], "3") + + def test_explicit_ucx_env_wins(self): + import run_kv + + env = {"COLLX_RDMA_DEVICES": "mlx5_0", "UCX_NET_DEVICES": "rdma0:1", + "COLLX_IB_GID_INDEX": "3", "UCX_IB_GID_INDEX": "1"} + run_kv.export_ucx_selectors(env) + self.assertEqual(env["UCX_NET_DEVICES"], "rdma0:1") + self.assertEqual(env["UCX_IB_GID_INDEX"], "1") + + def test_registry_device_pin_narrows_the_inventory(self): + # A kv_device pin on a UCX-backed case wins over the operator's full + # RDMA inventory: rail-isolated pods (b300) publish one-rail rows. + import run_kv + + env = {"COLLX_RDMA_DEVICES": "mlx5_0,mlx5_1"} + run_kv.export_ucx_selectors(env, device="mlx5_0") + self.assertEqual(env["UCX_NET_DEVICES"], "mlx5_0:1") + + def test_device_pin_overrides_host_inherited_ucx_env(self): + # b300 ships a blanket 16-device UCX_NET_DEVICES in /etc/environment + # (forwarded by srun --export=ALL); left standing it silently swallows + # the registry pin, so the pin wins — same treatment as UCX_TLS=rc. + import run_kv + + env = {"UCX_NET_DEVICES": "rdma0:1"} + run_kv.export_ucx_selectors(env, device="mlx5_0") + self.assertEqual(env["UCX_NET_DEVICES"], "mlx5_0:1") + + def test_empty_device_pin_keeps_the_inventory_path(self): + import run_kv + + env = {"COLLX_RDMA_DEVICES": "mlx5_0,mlx5_1"} + run_kv.export_ucx_selectors(env, device="") + self.assertEqual(env["UCX_NET_DEVICES"], "mlx5_0:1,mlx5_1:1") + + def test_ports_in_selectors_pass_through(self): + import run_kv + + env = {"COLLX_RDMA_DEVICES": "mlx5_18:1,mlx5_19"} + run_kv.export_ucx_selectors(env) + self.assertEqual(env["UCX_NET_DEVICES"], "mlx5_18:1,mlx5_19:1") + + def test_positive_tls_list_without_cuda_is_dropped(self): + # b300 exports UCX_TLS=rc cluster-wide; an RC-only context closes the + # cuda mds and NIXL registration of VRAM fails with NIXL_ERR_BACKEND. + # Extending the list with cuda transports segfaults UCX rkey-config + # resolution on the first GET, so the list is dropped entirely. + import run_kv + + env = {"UCX_TLS": "rc"} + run_kv.export_ucx_selectors(env) + self.assertNotIn("UCX_TLS", env) + + def test_tls_lists_already_covering_cuda_stay_untouched(self): + import run_kv + + for tls in ("^tcp", "rc,cuda_copy", "all"): + env = {"UCX_TLS": tls} + run_kv.export_ucx_selectors(env) + self.assertEqual(env["UCX_TLS"], tls) + env = {} + run_kv.export_ucx_selectors(env) + self.assertNotIn("UCX_TLS", env) + + +def _kv_document(status="success", sku="b200-nscale"): + def row(kind, page, op, gbps, p50, batch=1): + return {"kind": kind, "preset": "dsv4", "isl": 32768, "page_tokens": page, + "op": op, "descs": 1, "req_bytes": 1, "batch": batch, "prep_ms": 0.1, + "latency_ms": {"p50": p50, "p95": p50, "min": p50, "max": p50, "n": 48}, + "request_ms": {"p50": p50, "p95": p50, "min": p50, "max": p50, + "n": 48 * batch}, + "gbps_p50": gbps, "gbps_p50_incl_prep": gbps, + "verify": {"passed": status == "success", "detail": ""}} + + return { + "version": 1, + "record_type": "case-attempt", + "identity": {"case_factors": {"sku": sku, "case": { + "suite": "kv-transfer", "backend": "nixl", "workload": "kv-dsv4", + "mode": "rdma", "phase": "xfer", "ep": 2, "routing": "paged", + "precision": "fp8"}}}, + "measurement": {"rows": [ + row("paged", 256, "pull", 43.4, 53.1), + row("paged", 256, "pull", 96.2, 21.4, batch=16), + # a smaller measured block must lose to the production block size + row("paged", 128, "pull", 12.4, 185.2), + row("bulk", None, "pull", 48.3, 47.7), + row("paged", 256, "push", 48.4, 47.6), + ]}, + "topology": {"gpus_per_node": 1, "scale_up_domain": 8, "nodes": 2}, + "outcome": {"status": status, "reasons": []}, + } + + +class BurstTiming(unittest.TestCase): + def test_a_burst_posts_every_request_before_waiting_on_any(self): + from kv_backend import time_bursts + + order = [] + pairs = [(lambda i=i: order.append(("post", i)), + lambda i=i: order.append(("wait", i))) for i in range(3)] + burst_ms, request_ms = time_bursts(pairs, warmup=1, reps=2) + self.assertEqual(len(burst_ms), 2) + # one completion mark per request per kept rep, in posting order + self.assertEqual(len(request_ms), 2 * 3) + self.assertEqual(order[:6], [("post", 0), ("post", 1), ("post", 2), + ("wait", 0), ("wait", 1), ("wait", 2)]) + + def test_the_burst_sample_is_the_last_request_mark(self): + from kv_backend import time_bursts + + pairs = [(lambda: None, lambda: None)] * 2 + burst_ms, request_ms = time_bursts(pairs, warmup=0, reps=1) + self.assertEqual(burst_ms[0], request_ms[-1]) + # marks are offsets from the burst start, so they never decrease + self.assertEqual(request_ms, sorted(request_ms)) + + +class KVGrid(unittest.TestCase): + @staticmethod + def _args(**overrides): + import argparse + + base = dict(workload_name="kv-dsv4", precision="fp8", + isl_ladder="8192 32768 131072 524288", page_tokens="256", + batch_sizes="1 2 4 8 16 32 64", pool_slack=2.0) + import run_kv + + base["pool_budget"] = run_kv.POOL_BUDGET + base.update(overrides) + return argparse.Namespace(**base) + + def test_the_packed_grid_sheds_only_where_the_pool_budget_bites(self): + # Packed block-major geometry: a 512k-ISL block-256 request is 6,146 + # descriptors, so no batch on this ladder nears DESC_BUDGET. Only the + # 512k point sheds, and via the pool budget: its batch-32 pool plans + # ~118 GB against the 64 GiB budget, batch 16 fits at ~59 GB. + import run_kv + + points, isls, batches = run_kv._grid(self._args()) + self.assertEqual((isls, batches), + ([8192, 32768, 131072, 524288], [1, 2, 4, 8, 16, 32, 64])) + allowed = {cfg["isl"]: allowed for cfg, allowed in points} + self.assertEqual(allowed[8192], [1, 2, 4, 8, 16, 32, 64]) + self.assertEqual(allowed[32768], [1, 2, 4, 8, 16, 32, 64]) + self.assertEqual(allowed[131072], [1, 2, 4, 8, 16, 32, 64]) + self.assertEqual(allowed[524288], [1, 2, 4, 8, 16]) + for cfg, batch_list in points: + self.assertEqual(cfg["descs"], 3 * -(-cfg["isl"] // 256) + 2) + self.assertLessEqual(cfg["pool_bytes"], run_kv.POOL_BUDGET) + for batch in batch_list[run_kv.LADDER_FLOOR:]: + self.assertLessEqual(batch * cfg["descs"], run_kv.DESC_BUDGET) + + def test_descriptor_budget_sheds_batches_but_keeps_a_chartable_ladder(self): + # DESC_BUDGET stays as the fail-closed guard for future presets whose + # bursts are descriptor-bound. Pin it to 4 requests' descriptors at + # the largest ISL: batches above the per-point allowance shed, but the + # LADDER_FLOOR smallest batches always survive so every point keeps a + # chartable batch ladder (the frontier draws its line through the + # ladder at the largest measured ISL). + import kv_workload + import run_kv + + probe = kv_workload.plan_config("dsv4", "fp8", 524288, 256) + saved, run_kv.DESC_BUDGET = run_kv.DESC_BUDGET, 4 * probe["descs"] + try: + points, _isls, _batches = run_kv._grid(self._args()) + finally: + run_kv.DESC_BUDGET = saved + allowed = {cfg["isl"]: allowed for cfg, allowed in points} + self.assertEqual(allowed[8192], [1, 2, 4, 8, 16, 32, 64]) # 98 descs/req + self.assertEqual(allowed[32768], [1, 2, 4, 8, 16, 32]) # 386 + self.assertEqual(allowed[131072], [1, 2, 4, 8, 16]) # 1538, floor + self.assertEqual(allowed[524288], [1, 2, 4, 8, 16]) # 6146, floor + + def test_pool_budget_sheds_largest_batches_not_the_point(self): + # A point whose largest batch cannot fit the pool budget must survive + # with the batches that do: pin the budget between the batch-4 and + # batch-16 pool sizes. + import kv_workload + import run_kv + + args = self._args(isl_ladder="32768", batch_sizes="1 4 16") + budget = kv_workload.plan_config("dsv4", "fp8", 32768, 256, + 2.0, batch_max=4)["pool_bytes"] + args.pool_budget = budget + points, _isls, _batches = run_kv._grid(args) + self.assertEqual(points[0][1], [1, 4]) + self.assertLessEqual(points[0][0]["pool_bytes"], budget) + + def test_pool_budget_overrides_the_ladder_floor(self): + # The descriptor floor keeps the LADDER_FLOOR smallest batches, but + # the pool budget is a hard memory limit and must still shed a + # floor-kept batch. Pin the budget to the 512k point's batch-1 pool + # size: every larger batch survives the descriptor floor, then the + # pool loop must drop them all, leaving [1]. + import kv_workload + import run_kv + + args = self._args(isl_ladder="524288") + budget = kv_workload.plan_config("dsv4", "fp8", 524288, 256, + 2.0, batch_max=1)["pool_bytes"] + args.pool_budget = budget + points, _isls, _batches = run_kv._grid(args) + self.assertEqual(points[0][1], [1]) + self.assertLessEqual(points[0][0]["pool_bytes"], budget) + + +class RegistrationChunking(unittest.TestCase): + # b300's NICs refuse cuda registrations past ~8 GiB, so the NIXL adapter + # registers the pool in pieces. The pieces must never cut through a + # descriptor of ANY planned config, which _harmonize guarantees by giving + # every config one shared region layout. + + def test_harmonize_makes_region_bases_config_invariant(self): + import run_kv + + points, _isls, _batches = run_kv._grid(KVGrid._args()) + layout = run_kv._harmonize(points) + total = sum(nbytes for _, _, nbytes in layout) + running = 0 + for base, _packed, nbytes in layout: + self.assertEqual(base, running) + running += nbytes + for cfg, _ in points: + self.assertEqual(cfg["pool_bytes"], total) + for region, (base, packed, nbytes) in zip(cfg["regions"], layout): + self.assertEqual(region["base"], base) + self.assertEqual(region["packed_bytes"], packed) + self.assertEqual(region["pool_blocks"], nbytes // packed) + self.assertLessEqual(region["blocks_req"], region["pool_blocks"]) + + def test_reg_spans_cut_each_region_on_its_own_packed_grid(self): + import kv_nixl + import run_kv + + points, _isls, _batches = run_kv._grid(KVGrid._args()) + layout = run_kv._harmonize(points) + total = sum(nbytes for _, _, nbytes in layout) + spans = kv_nixl.reg_spans(total, layout) + # The full test grid plans a pool far past one chunk. + self.assertGreater(len(spans), 1) + # Exact in-order coverage, no gap, no overlap. + self.assertEqual(spans[0][0], 0) + for (a_off, a_len), (b_off, _) in zip(spans, spans[1:]): + self.assertEqual(a_off + a_len, b_off) + self.assertEqual(sum(length for _, length in spans), total) + for off, length in spans: + base, packed, _ = next(entry for entry in reversed(layout) + if entry[0] <= off) + self.assertEqual((off - base) % packed, 0) + self.assertLessEqual(length, max(kv_nixl.REG_CHUNK_BYTES, packed)) + + def test_no_descriptor_straddles_a_registration_cut(self): + # Every block any config can ever address must land whole inside one + # registered piece; a tiny cap on a small grid forces many cuts. + import bisect + + import kv_nixl + import run_kv + + args = KVGrid._args(isl_ladder="2048 8192", batch_sizes="1 4") + points, _isls, _batches = run_kv._grid(args) + layout = run_kv._harmonize(points) + total = sum(nbytes for _, _, nbytes in layout) + spans = kv_nixl.reg_spans(total, layout, cap=1 << 24) + self.assertGreater(len(spans), len(layout)) + starts = [off for off, _ in spans] + straddles = [] + for cfg, _ in points: + for region in cfg["regions"]: + packed = region["packed_bytes"] + for block in range(region["pool_blocks"]): + off = region["base"] + block * packed + s_off, s_len = spans[bisect.bisect_right(starts, off) - 1] + if off + packed > s_off + s_len: + straddles.append((region["name"], block)) + self.assertEqual(straddles, []) + + def test_without_a_layout_the_pool_registers_whole(self): + import kv_nixl + + self.assertEqual(kv_nixl.reg_spans(123456, None), [(0, 123456)]) + self.assertEqual(kv_nixl.reg_spans(123456, []), [(0, 123456)]) + + +class KVSummary(unittest.TestCase): + def test_kv_documents_render_their_own_table(self): + text = summarize.render([_kv_document()]) + self.assertIn("KV-transfer results", text) + self.assertIn("| pull | 43.4 | 96.2 | 48.3 | 53.1 |", text) + + def test_a_push_only_document_reads_its_push_lane(self): + doc = _kv_document() + doc["measurement"]["rows"] = [ + row for row in doc["measurement"]["rows"] if row["op"] == "push" + ] + text = summarize.render([doc]) + self.assertIn("| push | 48.4 |", text) + self.assertNotIn("INVALID", text) + + def test_kv_invalid_counts_in_the_banner(self): + text = summarize.render([_kv_document(status="invalid")]) + self.assertIn("INVALID", text) + + def test_ep_documents_do_not_grow_a_kv_table(self): + self.assertNotIn("KV-transfer results", summarize.render([])) + + +if __name__ == "__main__": + unittest.main() diff --git a/collectivex/tests/test_kv_workload.py b/collectivex/tests/test_kv_workload.py new file mode 100644 index 0000000000..2c99dd0256 --- /dev/null +++ b/collectivex/tests/test_kv_workload.py @@ -0,0 +1,274 @@ +#!/usr/bin/env python3 +"""Geometry and correctness math of the KV-transfer workload model. + +The packed block-major layout is the contract: per cache-group region, one +contiguous descriptor covers all the group's layers for one physical block +(vLLM's packed DSV4 NIXL shape), block tables are seed-keyed permutations both +ranks derive independently (batched requests slicing disjoint ranges of one +permutation), and an offset-derived pattern makes any byte's expected value +computable from its offset alone. These tests pin that math with hand-computed +cases validated against vLLM commit 32ad1400d7 (state content 584 B, page +padded to a 576 B multiple at block granularity, one descriptor per packed +block); the torch fill path is exercised on metal by the suite itself (a wrong +fill fails every verify row loudly). +""" +from __future__ import annotations + +import sys +import unittest +from pathlib import Path + +import numpy as np + +ROOT = Path(__file__).resolve().parents[1] +sys.path[:0] = [str(ROOT), str(ROOT / "bench")] + +import kv_workload # noqa: E402 + + +def _read8(pool: np.ndarray): + return lambda offset: pool[offset : offset + 8].tobytes() + + +class Geometry(unittest.TestCase): + def test_dsv4_regions_by_hand(self): + # isl=512, block=256. Every token-state is 584 B (448 NoPE + 128 RoPE + # + 8 fp8 scale); pages pad to a 576 B multiple at BLOCK granularity. + # C4A: 64 states -> round_up(64*584, 576) = 37,440; its indexer keeps + # 132 B states -> round_up(64*132, 576) = 8,640; C128A: 2 states -> + # round_up(2*584, 576) = 1,728; the sliding window's block is fixed at + # 64 tokens (it shares C4A's physical tensor) -> 37,440 on all 61 + # layers, capped at 128 window tokens. One descriptor per block spans + # the group's layers. + cfg = kv_workload.plan_config("dsv4", "fp8", 512, 256) + regions = {r["name"]: r for r in cfg["regions"]} + self.assertEqual([r["name"] for r in cfg["regions"]], + ["c4a", "c4a-idx", "c128a", "swa"]) + self.assertEqual( + (regions["c4a"]["layers"], regions["c4a"]["page_bytes"], + regions["c4a"]["packed_bytes"], regions["c4a"]["blocks_req"]), + (30, 37_440, 30 * 37_440, 2)) + self.assertEqual( + (regions["c4a-idx"]["layers"], regions["c4a-idx"]["page_bytes"], + regions["c4a-idx"]["blocks_req"]), (30, 8_640, 2)) + self.assertEqual( + (regions["c128a"]["layers"], regions["c128a"]["page_bytes"], + regions["c128a"]["blocks_req"]), (31, 1_728, 2)) + self.assertEqual( + (regions["swa"]["layers"], regions["swa"]["block_tokens"], + regions["swa"]["blocks_req"]), (61, 64, 2)) + self.assertEqual(cfg["descs"], 2 + 2 + 2 + 2) + self.assertEqual(cfg["req_bytes"], + 2 * (30 * 37_440 + 30 * 8_640 + 31 * 1_728 + 61 * 37_440)) + # regions tile one contiguous pool + self.assertEqual(cfg["pool_bytes"], + sum(r["pool_blocks"] * r["packed_bytes"] + for r in cfg["regions"])) + + def test_alignment_pads_the_page_not_each_state(self): + # 64 states * 584 B = 37,376 -> padded once per page to 37,440. The + # old per-entry 576 B model would give 64 * 576 = 36,864 — vLLM pads + # at page granularity, not per state. + cfg = kv_workload.plan_config("dsv4", "fp8", 512, 256) + c4a = {r["name"]: r for r in cfg["regions"]}["c4a"] + self.assertEqual(c4a["page_bytes"], 37_440) + self.assertNotEqual(c4a["page_bytes"], 64 * 576) + + def test_swa_shares_the_c4a_page_size(self): + # Both block types live in one physical tensor: a 64-token window + # block (1 token/state) and a 256-token C4A block (4 tokens/state) + # are the same 64 states -> byte-identical pages. + cfg = kv_workload.plan_config("dsv4", "fp8", 512, 256) + regions = {r["name"]: r for r in cfg["regions"]} + self.assertEqual(regions["swa"]["page_bytes"], regions["c4a"]["page_bytes"]) + + def test_one_descriptor_per_block_at_the_big_isl(self): + # 512k tokens at block 256: 2048 blocks per non-window group + 2 + # window blocks = 6,146 descriptors per request — the packed shape + # vLLM's connector asserts, not a per-(layer, page) explosion. + cfg = kv_workload.plan_config("dsv4", "fp8", 524_288, 256) + self.assertEqual(cfg["descs"], 2048 * 3 + 2) + + def test_dsv4_window_caps_at_128_tokens(self): + small = kv_workload.plan_config("dsv4", "fp8", 64, 256) + large = kv_workload.plan_config("dsv4", "fp8", 32_768, 256) + window = {r["name"]: r for r in large["regions"]}["swa"] + self.assertEqual(window["blocks_req"], 2) # 128 tokens / 64 per block + self.assertEqual({r["name"]: r for r in small["regions"]}["swa"]["blocks_req"], + 1) # min(isl, 128) = 64 tokens + + def test_block_sizes_that_split_a_state_fail_closed(self): + # C128A's 128-token states force the model block size to a multiple + # of 128; vLLM serves DSV4 at 256. The old 16/64-token sweep values + # cannot hold a whole HCA state and must be rejected. + for block in (16, 64, 192): + with self.assertRaises(ValueError): + kv_workload.plan_config("dsv4", "fp8", 512, block) + self.assertEqual( + {r["name"]: r for r in + kv_workload.plan_config("dsv4", "fp8", 512, 128)["regions"] + }["c128a"]["page_bytes"], 1_152) # 1 state, 584 -> padded + + def test_dsv4_precision_is_architectural(self): + with self.assertRaises(ValueError): + kv_workload.plan_config("dsv4", "bf16", 512, 256) + + def test_partial_last_block_rounds_up(self): + # 300 tokens at 256/block -> 2 blocks for every non-window group. + cfg = kv_workload.plan_config("dsv4", "fp8", 300, 256) + self.assertEqual(cfg["regions"][0]["blocks_req"], 2) + + def test_batch_max_grows_the_pool_for_disjoint_requests(self): + cfg = kv_workload.plan_config("dsv4", "fp8", 512, 256, batch_max=16) + for region in cfg["regions"]: + self.assertGreaterEqual(region["pool_blocks"], 16 * region["blocks_req"]) + + +class Tables(unittest.TestCase): + def test_deterministic_and_distinct_per_side(self): + cfg = kv_workload.plan_config("dsv4", "fp8", 4096, 256) + local = kv_workload.block_table(cfg, kv_workload.table_seed(cfg, "local")) + remote = kv_workload.block_table(cfg, kv_workload.table_seed(cfg, "remote")) + again = kv_workload.block_table(cfg, kv_workload.table_seed(cfg, "local")) + for region in cfg["regions"]: + name, blocks_req = region["name"], region["blocks_req"] + self.assertTrue((local[name] == again[name]).all()) + self.assertFalse((local[name] == remote[name]).all()) + # distinct in-range blocks (fragmented, never aliased) + self.assertEqual(len(set(local[name].tolist())), blocks_req) + self.assertTrue((local[name] < region["pool_blocks"]).all()) + + def test_batched_requests_slice_disjoint_blocks(self): + cfg = kv_workload.plan_config("dsv4", "fp8", 512, 256, batch_max=4) + seed = kv_workload.table_seed(cfg, "local") + tables = [kv_workload.block_table(cfg, seed, request=r) for r in range(4)] + for region in cfg["regions"]: + blocks = [t[region["name"]].tolist() for t in tables] + union = set().union(*map(set, blocks)) + self.assertEqual(len(union), 4 * region["blocks_req"]) + + def test_a_request_beyond_the_pool_fails_closed(self): + cfg = kv_workload.plan_config("dsv4", "fp8", 512, 256) # slack for ~2 requests + with self.assertRaises(ValueError): + kv_workload.block_table(cfg, 1, request=8) + + def test_block_major_offsets(self): + # One offset per packed block: block b sits at b * packed_bytes. + cfg = dict(regions=[dict(name="kv", packed_bytes=512, blocks_req=2, + pool_blocks=3, base=0)], descs=2) + offsets = kv_workload.page_offsets(cfg, {"kv": np.array([2, 0])}) + self.assertEqual(offsets.tolist(), [2 * 512, 0]) + + def test_second_region_offsets_start_at_its_base(self): + cfg = dict(regions=[ + dict(name="a", packed_bytes=256, blocks_req=1, pool_blocks=2, base=0), + dict(name="b", packed_bytes=128, blocks_req=1, pool_blocks=2, base=512), + ], descs=2) + offsets = kv_workload.page_offsets(cfg, {"a": np.array([1]), "b": np.array([1])}) + self.assertEqual(offsets.tolist(), [256, 512 + 128]) + + def test_desc_array_carries_per_region_packed_sizes(self): + cfg = dict(regions=[ + dict(name="a", packed_bytes=256, blocks_req=2, pool_blocks=4, base=0), + dict(name="b", packed_bytes=132, blocks_req=1, pool_blocks=4, base=1024), + ], descs=3) + tables = {"a": np.array([1, 3]), "b": np.array([2])} + descs = kv_workload.desc_array(10_000, cfg, tables, dev=5) + self.assertEqual(descs[:, 0].tolist(), + [10_000 + 256, 10_000 + 768, 10_000 + 1024 + 264]) + self.assertEqual(descs[:, 1].tolist(), [256, 256, 132]) + self.assertEqual(descs[:, 2].tolist(), [5, 5, 5]) + + +class Verify(unittest.TestCase): + def _painted_destination(self, cfg, dst_tables, src_tables): + """A destination pool where every dst block holds its src block's pattern.""" + pool = np.zeros(cfg["pool_bytes"], dtype=np.uint8) + for region in cfg["regions"]: + size = region["packed_bytes"] + for dst, src in zip(dst_tables[region["name"]], src_tables[region["name"]]): + dst_off = int(dst) * size + region["base"] + src_off = int(src) * size + region["base"] + src_bytes = src_off + np.arange(size, dtype=np.int64) + pool[dst_off : dst_off + size] = ((src_bytes >> 8) * 131 + 7) & 0xFF + return pool + + def _tables(self, cfg): + dst = kv_workload.block_table(cfg, kv_workload.table_seed(cfg, "local")) + src = kv_workload.block_table(cfg, kv_workload.table_seed(cfg, "remote")) + return dst, src + + def test_a_faithful_transfer_verifies_across_unaligned_pages(self): + # dsv4's page sizes are 576 B multiples, never 256 B multiples, so + # per-layer probes land at any byte alignment and exercise the + # per-byte expectation model. + cfg = kv_workload.plan_config("dsv4", "fp8", 512, 256) + dst, src = self._tables(cfg) + pool = self._painted_destination(cfg, dst, src) + ok, detail = kv_workload.verify_transfer(_read8(pool), cfg, dst, src) + self.assertTrue(ok, detail) + + def test_one_missing_transfer_fails_with_its_coordinates(self): + cfg = kv_workload.plan_config("dsv4", "fp8", 512, 256) + dst, src = self._tables(cfg) + pool = self._painted_destination(cfg, dst, src) + pool[:] = 0 # a transfer that never happened + ok, detail = kv_workload.verify_transfer(_read8(pool), cfg, dst, src) + self.assertFalse(ok) + self.assertIn("expected", detail) + + def test_direction_matters(self): + # Verifying with the tables swapped must fail: dst blocks hold src + # pattern, not their own. + cfg = kv_workload.plan_config("dsv4", "fp8", 512, 256) + dst, src = self._tables(cfg) + pool = self._painted_destination(cfg, dst, src) + ok, _ = kv_workload.verify_transfer(_read8(pool), cfg, src, dst) + self.assertFalse(ok) + + def test_fabric_pool_pattern_matches_the_verify_model(self): + # kv_pool's host-built pattern (the mnnvl fill path) and the verify + # model must agree byte for byte, or every mnnvl row fails verify. + import kv_pool + + pattern = kv_pool._pattern(1024) + for offset in (0, 8, 256, 512, 1016): + expected = kv_workload._chunk_byte(offset) + self.assertTrue((pattern[offset : offset + 8] == expected).all(), offset) + + +class SweepConfigConsistency(unittest.TestCase): + def test_kv_sweep_precisions_match_the_workload_model(self): + # sweep_matrix schedules from the JSON map (it must stay stdlib-only); + # the workload model owns the truth and plan_config fail-closes on a + # mismatch at runtime. This pins the two together at PR time. + import json + + sweep = json.loads((ROOT / "configs" / "kv_sweep.json").read_text()) + for workload, precisions in sweep["workloads"].items(): + preset = kv_workload.PRESETS[workload.removeprefix("kv-")] + self.assertEqual(tuple(precisions), preset["precisions"], workload) + + def test_kv_sweep_block_sizes_are_plannable(self): + # A sweep block size the model rejects (splitting an HCA state) would + # kill every kv leg at the first grid point. + import json + + sweep = json.loads((ROOT / "configs" / "kv_sweep.json").read_text()) + for workload, precisions in sweep["workloads"].items(): + for block in sweep["page_tokens"]: + kv_workload.plan_config(workload.removeprefix("kv-"), + precisions[0], 512, block) + + +class Percentiles(unittest.TestCase): + def test_pcts(self): + stats = kv_workload.pcts([5.0, 1.0, 3.0, 2.0, 4.0]) + self.assertEqual(stats["p50"], 3.0) + self.assertEqual(stats["min"], 1.0) + self.assertEqual(stats["max"], 5.0) + self.assertEqual(stats["n"], 5) + + +if __name__ == "__main__": + unittest.main() diff --git a/collectivex/tests/test_runtime.py b/collectivex/tests/test_runtime.py index 5cebe5a248..ba80ecbb22 100644 --- a/collectivex/tests/test_runtime.py +++ b/collectivex/tests/test_runtime.py @@ -600,6 +600,8 @@ def test_the_rank_wrapper_execs_only_a_named_entrypoint(self) -> None: "ARGV bench/run_swap_blocks.py --layout random"), (["--entrypoint", "run_ep", "--backend", "mori"], 0, "ARGV bench/run_ep.py --backend mori"), + (["--entrypoint", "run_kv", "--backend", "nixl"], 0, + "ARGV bench/run_kv.py --backend nixl"), (["--entrypoint", "../../bin/sh"], 67, ""), (["--backend", "mori"], 67, ""), ): From 3b42529845d2eeea28770d8502704623401bb45c Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Mon, 28 Sep 2026 23:20:50 +0800 Subject: [PATCH 2/4] CollectiveX: share the case-attempt envelope and trim the kv plumbing run_ep and run_kv build their identity, provenance and outcome through one ep_harness.case_attempt (kv now validates COLLX_ATTEMPT_ID as EP does). The kv codec and shard builder use the shared helpers, run_kv folds its row tail and drops the never-set MoRI QP/chunking flags (library defaults were already the only value), and the UCX selector and pool-budget tests are table-driven. EP documents, emitted argv and resolved matrices are unchanged. --- collectivex/bench/ep_harness.py | 74 +++++++++------ collectivex/bench/kv_mori_io.py | 5 +- collectivex/bench/run_ep.py | 8 +- collectivex/bench/run_kv.py | 55 ++++------- collectivex/docs/methodology.md | 13 +-- collectivex/runtime/config.py | 35 ++----- collectivex/sweep_matrix.py | 34 ++----- collectivex/tests/test_kv_suite.py | 147 +++++++++-------------------- 8 files changed, 135 insertions(+), 236 deletions(-) diff --git a/collectivex/bench/ep_harness.py b/collectivex/bench/ep_harness.py index a526288bf6..39cad1a4f9 100644 --- a/collectivex/bench/ep_harness.py +++ b/collectivex/bench/ep_harness.py @@ -178,6 +178,43 @@ def _component(percentiles, count, *, derived=False, origin=None): ) +def git_run() -> dict | None: + """The workflow run that produced an attempt, from the environment the launcher forwards.""" + run = { + "run_id": os.environ.get("GITHUB_RUN_ID"), + "run_attempt": os.environ.get("GITHUB_RUN_ATTEMPT"), + "source_sha": os.environ.get("COLLECTIVEX_SOURCE_SHA") or os.environ.get("GITHUB_SHA"), + } + return run if any(run.values()) else None + + +def case_attempt(args, case: dict, run: dict | None, image: str, ok: bool, failure: str, + **sections) -> dict: + """The case-attempt document every suite writes: identity, provenance and outcome keyed and + validated the same way around the suite's own sections (workload, measurement, ...).""" + try: + attempt_ordinal = int(os.environ.get("COLLX_ATTEMPT_ID", "1")) + except ValueError: + attempt_ordinal = 0 + if attempt_ordinal <= 0: + raise ValueError("COLLX_ATTEMPT_ID must be a positive integer") + run = run or {} + return { + "version": args.version, + "record_type": "case-attempt", + "generated_at": _dt.datetime.now().astimezone().isoformat(), + "identity": { + "allocation_factors": {key: run.get(key) for key in ("run_attempt", "run_id", "source_sha")}, + "attempt_ordinal": attempt_ordinal, + "case_factors": {"case": case, "sku": args.runner}, + "case_id": args.case_id, + }, + **sections, + "provenance": {"image": image or None, "source_sha": run.get("source_sha")}, + "outcome": {"reasons": [] if ok else [failure], "status": "success" if ok else "invalid"}, + } + + def _write_json_atomic(path: str, value) -> None: payload = json.dumps(value, allow_nan=False, ensure_ascii=False, separators=(",", ":")).encode() + b"\n" os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) @@ -696,30 +733,16 @@ def _document(self, rows, all_ok): computed = case_id(args.runner, scheduled_case) if args.case_id != computed: raise ValueError(f"scheduled case ID does not match realized factors: {args.case_id} != {computed}") - git_run = getattr(args, "git_run", None) or {} - try: - attempt_ordinal = int(os.environ.get("COLLX_ATTEMPT_ID", "1")) - except ValueError: - attempt_ordinal = 0 - if attempt_ordinal <= 0: - raise ValueError("COLLX_ATTEMPT_ID must be a positive integer") - return { - "version": args.version, - "record_type": "case-attempt", - "generated_at": _dt.datetime.now().astimezone().isoformat(), - "identity": { - "allocation_factors": {key: git_run.get(key) for key in ("run_attempt", "run_id", "source_sha")}, - "attempt_ordinal": attempt_ordinal, - "case_factors": {"case": scheduled_case, "sku": args.runner}, - "case_id": args.case_id, - }, - "workload": { + return case_attempt( + args, scheduled_case, getattr(args, "git_run", None), getattr(args, "image", ""), + all_ok, "semantic correctness or routing identity failed", + workload={ "cross_rank_consistent": self.routing_consistent, "ladder_measured": list(self.ladder), "ladder_dropped": list(self.dropped), "ladder_cap": self.cap, }, - "measurement": { + measurement={ "combine_dtype": backend.combine_dtype, "combine_semantics": "activation-only", "dispatch_dtype": backend.dispatch_dtype, @@ -736,7 +759,7 @@ def _document(self, rows, all_ok): "chain_trials": args.chain_trials, }, }, - "implementation": { + implementation={ "fp8_consume": backend.fp8_consume, "kernel_generation": kernel_generation(backend), # The reduction the oracle held the kernel to, and the library version it was @@ -751,7 +774,7 @@ def _document(self, rows, all_ok): "maturity": backend.maturity or "unknown", "name": backend.name, }, - "topology": { + topology={ "device_product": getattr(args, "runtime_device_product", None), "gpus_per_node": args.gpus_per_node, "nodes": nodes, @@ -761,13 +784,8 @@ def _document(self, rows, all_ok): "transport": args.transport, "world_size": self.world_size, }, - "runtime": getattr(args, "runtime", {}), - "provenance": {"image": getattr(args, "image", "") or None, "source_sha": git_run.get("source_sha")}, - "outcome": { - "reasons": [] if all_ok else ["semantic correctness or routing identity failed"], - "status": "success" if all_ok else "invalid", - }, - } + runtime=getattr(args, "runtime", {}), + ) def _print_summary(self, rows, doc): # Ladder ends plus two interior points: one mid-ladder headline hides low-token behavior. diff --git a/collectivex/bench/kv_mori_io.py b/collectivex/bench/kv_mori_io.py index 91e5e98c45..b1a2b3102b 100644 --- a/collectivex/bench/kv_mori_io.py +++ b/collectivex/bench/kv_mori_io.py @@ -38,13 +38,14 @@ def __init__(self, args, role, device): host = kv_workload.iface_ipv4(args.socket_ifname) if args.socket_ifname else "" port = int(args.kv_mori_port) + (0 if role == "target" else 1) self._engine = IOEngine(key=role, config=IOEngineConfig(host=host, port=port)) + # Library defaults only: four QPs plus transfer chunking wedged on the metal. self._engine.create_backend(BackendType.RDMA, RdmaBackendConfig( - qp_per_transfer=int(args.kv_mori_qp), + qp_per_transfer=1, post_batch_size=-1, num_worker_threads=1, poll_cq_mode=PollCqMode.POLLING, enable_notification=False, - enable_transfer_chunking=bool(args.kv_mori_chunking), + enable_transfer_chunking=False, chunk_bytes=65536, max_chunks_per_transfer=64, )) diff --git a/collectivex/bench/run_ep.py b/collectivex/bench/run_ep.py index fd438de49a..e501fac8d0 100644 --- a/collectivex/bench/run_ep.py +++ b/collectivex/bench/run_ep.py @@ -86,13 +86,7 @@ def main() -> int: device_name = torch.cuda.get_device_name(device) args.runtime_device_product = device_name args.image = os.environ.get("COLLECTIVEX_IMAGE", "") - _run = { - "run_id": os.environ.get("GITHUB_RUN_ID"), - "run_attempt": os.environ.get("GITHUB_RUN_ATTEMPT"), - "source_sha": os.environ.get("COLLECTIVEX_SOURCE_SHA") - or os.environ.get("GITHUB_SHA"), - } - args.git_run = _run if any(_run.values()) else None + args.git_run = ep_harness.git_run() module, class_name = BACKENDS[args.backend] Backend = getattr(importlib.import_module(module), class_name) diff --git a/collectivex/bench/run_kv.py b/collectivex/bench/run_kv.py index 19311e7445..9ebe4a80a6 100644 --- a/collectivex/bench/run_kv.py +++ b/collectivex/bench/run_kv.py @@ -91,8 +91,6 @@ def add_kv_args(ap: argparse.ArgumentParser) -> None: ap.add_argument("--scale-up-transport", default="") ap.add_argument("--topology-class", default="") ap.add_argument("--socket-ifname", default=os.environ.get("COLLX_SOCKET_IFNAME", "")) - ap.add_argument("--kv-mori-qp", type=int, default=1) - ap.add_argument("--kv-mori-chunking", action="store_true") ap.add_argument("--kv-device", default="", help="engine NIC filter template; {gpu} expands to the " "physical GPU index (GPU-paired NICs, e.g. Pollara). " @@ -334,6 +332,11 @@ def repaint(): rows: list[dict] = [] + def record(row): + if row is not None: + rows.append(row) + print(f"[run_kv] {json.dumps(row)}", flush=True) + def verify_burst(cfg, table_pairs): """Every request in the burst must land: a passing request 0 says nothing about the others, and concurrent same-session requests are @@ -410,11 +413,7 @@ def measure(make, cfg_row: dict, op: str, verify_side: str, table_pairs=None): (initiator_tables[r], target_tables[r]) if op == "pull" else (target_tables[r], initiator_tables[r]) for r in range(batch)] - row = measure(make, {**base, "batch": batch}, op, verify_side, - table_pairs) - if row is not None: - rows.append(row) - print(f"[run_kv] {json.dumps(row)}", flush=True) + record(measure(make, {**base, "batch": batch}, op, verify_side, table_pairs)) for isl in isls: preset = args.workload_name.removeprefix("kv-") @@ -427,10 +426,7 @@ def measure(make, cfg_row: dict, op: str, verify_side: str, table_pairs=None): "req_bytes": nbytes, "_cfg": cfg} for op in ops: make = (lambda op=op, n=nbytes: [backend.make_bulk(n, op)]) if role == "initiator" else None - row = measure(make, base, op, verify_side="none", table_pairs=None) - if row is not None: - rows.append(row) - print(f"[run_kv] {json.dumps(row)}", flush=True) + record(measure(make, base, op, verify_side="none", table_pairs=None)) backend.teardown() @@ -442,28 +438,17 @@ def measure(make, cfg_row: dict, op: str, verify_side: str, table_pairs=None): all_ok = bool(rows) and all(r["verify"]["passed"] for r in rows) if rank == 0: - doc = { - "version": args.version, - "record_type": "case-attempt", - "generated_at": _dt.datetime.now().astimezone().isoformat(), - "identity": { - "allocation_factors": { - "run_attempt": os.environ.get("GITHUB_RUN_ATTEMPT"), - "run_id": os.environ.get("GITHUB_RUN_ID"), - "source_sha": os.environ.get("COLLECTIVEX_SOURCE_SHA") or os.environ.get("GITHUB_SHA"), - }, - "attempt_ordinal": int(os.environ.get("COLLX_ATTEMPT_ID", "1")), - "case_factors": {"case": {**case, "suite": args.suite}, "sku": args.runner}, - "case_id": args.case_id, - }, - "workload": { + doc = ep_harness.case_attempt( + args, {**case, "suite": args.suite}, ep_harness.git_run(), + os.environ.get("COLLECTIVEX_IMAGE", ""), all_ok, "transfer verification failed", + workload={ "isl_ladder": isls, "page_tokens": [int(v) for v in args.page_tokens.split()], "batch_sizes": batches, "ops": ops, "preset": kv_workload.PRESETS[args.workload_name.removeprefix("kv-")], }, - "measurement": { + measurement={ "payload_unit": "request-kv-bytes", "rows": rows, "sampling": { @@ -472,14 +457,14 @@ def measure(make, cfg_row: dict, op: str, verify_side: str, table_pairs=None): "warmup_per_trial": args.warmup, }, }, - "implementation": { + implementation={ "name": args.backend, "fabric": args.fabric, "library_version": getattr(backend, "library_version", None), "maturity": getattr(backend, "maturity", "candidate"), "nic_filter": getattr(backend, "nic_filter", None), }, - "topology": { + topology={ "device_product": torch.cuda.get_device_name(device), "gpus_per_node": args.gpus_per_node, "hosts": hosts, @@ -490,19 +475,11 @@ def measure(make, cfg_row: dict, op: str, verify_side: str, table_pairs=None): "topology_class": args.topology_class or None, "world_size": world_size, }, - "runtime": { + runtime={ "framework": str(torch.__version__), "vendor": "amd" if torch.version.hip else "nvidia", }, - "provenance": { - "image": os.environ.get("COLLECTIVEX_IMAGE") or None, - "source_sha": os.environ.get("COLLECTIVEX_SOURCE_SHA") or os.environ.get("GITHUB_SHA"), - }, - "outcome": { - "reasons": [] if all_ok else ["transfer verification failed"], - "status": "success" if all_ok else "invalid", - }, - } + ) if args.out: ep_harness._write_json_atomic(args.out, doc) print(f"[run_kv] status={doc['outcome']['status']} rows={len(rows)}" diff --git a/collectivex/docs/methodology.md b/collectivex/docs/methodology.md index fd69c8d15b..d2185a8482 100644 --- a/collectivex/docs/methodology.md +++ b/collectivex/docs/methodology.md @@ -627,14 +627,11 @@ the binary level (the wheel links libcuda.so.1 at import; measured failing on mi ## Correctness -Scheduling is data. `configs/kv_sweep.json` holds the grid and, per pool, the allocation and the -per-case hang guard (`scheduling`): a leg runs for hours, gb200's mnnvl descriptor floor alone -measured ~285 minutes on the five-rung grid, and gb300 paces ~1.8x gb200 at ISL >= 131072 over -mnnvl, so gb200 asks 460 minutes with a 420-minute guard and gb300 690 with 660; every other pool -asks 210 with 190. The guard fires before the allocation dies, so a slow case is a clean per-case -kill, and each shard carries a GitHub job ceiling above its allocation. The asks stay 2 nodes x 1 -GPU, short enough to backfill on a contended pool. KV legs run only when a dispatch names -`kv-transfer` in `suites`. +Scheduling is data: `configs/kv_sweep.json` holds the grid and, per pool, the allocation and the +per-case hang guard. GB pools get the long budgets, because the mnnvl descriptor floor makes a leg +run for hours and gb300 paces ~1.8x gb200 at ISL >= 131072. The guard fires inside the allocation, +so a slow case is a clean per-case kill, and each shard carries a GitHub job ceiling above its +allocation. KV legs run only when a dispatch names `kv-transfer` in `suites`. ## Correctness diff --git a/collectivex/runtime/config.py b/collectivex/runtime/config.py index e70fab14b7..61f8edb905 100644 --- a/collectivex/runtime/config.py +++ b/collectivex/runtime/config.py @@ -171,33 +171,14 @@ def _swap_argv(case: dict, version: object, runner: str) -> list[str]: def _kv_argv(case: dict, version: object, runner: str) -> list[str]: - argv = [ - "--backend", str(case["backend"]), - "--workload-name", str(case["workload"]), - "--precision", str(case["precision"]), - "--fabric", str(case["mode"]), - "--isl-ladder", str(case["isl_ladder"]), - "--page-tokens", str(case["page_tokens"]), - "--batch-sizes", str(case["batch_sizes"]), - "--kv-device", str(case.get("kv_device", "")), - "--ops", str(case["ops"]), - "--warmup", str(case["warmup"]), - "--reps", str(case["reps"]), - "--trials", str(case["trials"]), - "--pool-slack", str(case["pool_slack"]), - "--seed", str(case["seed"]), - "--runner", runner, - "--case-id", str(case["case_id"]), - "--suite", str(case["suite"]), - "--gpus-per-node", str(case["gpus_per_node"]), - "--scale-up-domain", str(case["scale_up_domain"]), - "--scale-up-transport", str(case["scale_up_transport"]), - "--topology-class", str(case["topology_class"]), - "--version", str(version), - ] - if "pool_budget" in case: - argv += ["--pool-budget", str(case["pool_budget"])] - return argv + argv = _flag_pairs( + case, + "backend workload_name precision fabric isl_ladder page_tokens batch_sizes kv_device ops " + "warmup reps trials pool_slack seed case_id suite gpus_per_node scale_up_domain " + "scale_up_transport topology_class", + workload_name="workload", fabric="mode", + ) + ["--runner", runner, "--version", str(version)] + return argv + (_flag_pairs(case, "pool_budget") if "pool_budget" in case else []) # suite -> (bench/.py, argv codec, output flag). The rank wrapper in diff --git a/collectivex/sweep_matrix.py b/collectivex/sweep_matrix.py index 9d7526db4f..cfb89a59af 100644 --- a/collectivex/sweep_matrix.py +++ b/collectivex/sweep_matrix.py @@ -252,33 +252,17 @@ def _kv_shards(sku: str, selected_precisions: set[str]) -> tuple[list, list[dict cases.append(case) if not cases: continue - requested += [ - {"sku": sku, "case": case, "disposition": "runnable", "reason": None, - "detail": None} - for case in cases - ] - shard = { - "id": f"{sku}-kv-{backend}-{fabric}", - "sku": sku, - "runner": platform.get("runner_label", sku), - "backend": backend, - "suite": KV_SWEEP["suite"], - "mode": fabric, - "fabric": fabric, - "launcher": platform["launcher"], - "nodes": 2, - "gpus_per_node": 1, - "scale_up_domain": platform["scale_up_domain"], - "allocation_minutes": scheduling["allocation_minutes"], - "run_timeout": scheduling["run_timeout"], + requested += _runnable(sku, cases) + shards.append(_shard( + sku, f"{sku}-kv-{backend}-{fabric}", backend, cases, + suite=KV_SWEEP["suite"], mode=fabric, fabric=fabric, + allocation_minutes=scheduling["allocation_minutes"], + run_timeout=scheduling["run_timeout"], # The GitHub job must outlive the allocation, or it cancels a healthy shard # before the launcher's own guards act. - "job_timeout_minutes": scheduling["allocation_minutes"] + 30, - "cases": cases, - } - if spec["image"]: - shard["image"] = spec["image"] - shards.append(shard) + job_timeout_minutes=scheduling["allocation_minutes"] + 30, + **({"image": spec["image"]} if spec["image"] else {}), + )) return requested, shards diff --git a/collectivex/tests/test_kv_suite.py b/collectivex/tests/test_kv_suite.py index cdd90fa410..9805cc8afd 100644 --- a/collectivex/tests/test_kv_suite.py +++ b/collectivex/tests/test_kv_suite.py @@ -209,80 +209,42 @@ def test_a_row_with_no_verifying_side_passes_without_a_gather_crash(self): class UCXSelectors(unittest.TestCase): """run_kv pins UCX to the operator's validated RDMA selectors — UCX - auto-selection is a wrong-fabric trap (b200-nscale's aux quad-port card, - b300's storage IB) — while explicit UCX_* values always win.""" - - def test_registry_selectors_map_to_ucx(self): - import run_kv - - env = {"COLLX_RDMA_DEVICES": "mlx5_0,mlx5_10", "COLLX_IB_GID_INDEX": "3"} - run_kv.export_ucx_selectors(env) - self.assertEqual(env["UCX_NET_DEVICES"], "mlx5_0:1,mlx5_10:1") - self.assertEqual(env["UCX_IB_GID_INDEX"], "3") - - def test_explicit_ucx_env_wins(self): - import run_kv - - env = {"COLLX_RDMA_DEVICES": "mlx5_0", "UCX_NET_DEVICES": "rdma0:1", - "COLLX_IB_GID_INDEX": "3", "UCX_IB_GID_INDEX": "1"} - run_kv.export_ucx_selectors(env) - self.assertEqual(env["UCX_NET_DEVICES"], "rdma0:1") - self.assertEqual(env["UCX_IB_GID_INDEX"], "1") - - def test_registry_device_pin_narrows_the_inventory(self): - # A kv_device pin on a UCX-backed case wins over the operator's full - # RDMA inventory: rail-isolated pods (b300) publish one-rail rows. + auto-selection is a wrong-fabric trap (b200-nscale's aux quad-port card) — + while explicit UCX_* values always win, except over a case's own NIC pin.""" + + # (environment, kv_device pin) -> the UCX variables export_ucx_selectors leaves. + CASES = ( + # registry selectors map to UCX, ports default to 1 and pass through + ({"COLLX_RDMA_DEVICES": "mlx5_0,mlx5_10", "COLLX_IB_GID_INDEX": "3"}, "", + {"UCX_NET_DEVICES": "mlx5_0:1,mlx5_10:1", "UCX_IB_GID_INDEX": "3"}), + ({"COLLX_RDMA_DEVICES": "mlx5_18:1,mlx5_19"}, "", + {"UCX_NET_DEVICES": "mlx5_18:1,mlx5_19:1"}), + # explicit UCX env wins over the inventory + ({"COLLX_RDMA_DEVICES": "mlx5_0", "UCX_NET_DEVICES": "rdma0:1", + "COLLX_IB_GID_INDEX": "3", "UCX_IB_GID_INDEX": "1"}, "", + {"UCX_NET_DEVICES": "rdma0:1", "UCX_IB_GID_INDEX": "1"}), + # a case pin narrows the inventory and overrides a host-inherited blanket value + # (forwarded by srun --export=ALL from /etc/environment), which would swallow it + ({"COLLX_RDMA_DEVICES": "mlx5_0,mlx5_1"}, "mlx5_0", {"UCX_NET_DEVICES": "mlx5_0:1"}), + ({"UCX_NET_DEVICES": "rdma0:1"}, "mlx5_0", {"UCX_NET_DEVICES": "mlx5_0:1"}), + # A positive UCX_TLS list without cuda (a cluster-wide UCX_TLS=rc) closes the cuda + # mds and NIXL VRAM registration fails; extending it segfaults rkey resolution on the + # first GET, so it is dropped. Lists already covering cuda stay untouched. + ({"UCX_TLS": "rc"}, "", {}), + ({"UCX_TLS": "^tcp"}, "", {"UCX_TLS": "^tcp"}), + ({"UCX_TLS": "rc,cuda_copy"}, "", {"UCX_TLS": "rc,cuda_copy"}), + ({"UCX_TLS": "all"}, "", {"UCX_TLS": "all"}), + ({}, "", {}), + ) + + def test_selectors(self): import run_kv - env = {"COLLX_RDMA_DEVICES": "mlx5_0,mlx5_1"} - run_kv.export_ucx_selectors(env, device="mlx5_0") - self.assertEqual(env["UCX_NET_DEVICES"], "mlx5_0:1") - - def test_device_pin_overrides_host_inherited_ucx_env(self): - # b300 ships a blanket 16-device UCX_NET_DEVICES in /etc/environment - # (forwarded by srun --export=ALL); left standing it silently swallows - # the registry pin, so the pin wins — same treatment as UCX_TLS=rc. - import run_kv - - env = {"UCX_NET_DEVICES": "rdma0:1"} - run_kv.export_ucx_selectors(env, device="mlx5_0") - self.assertEqual(env["UCX_NET_DEVICES"], "mlx5_0:1") - - def test_empty_device_pin_keeps_the_inventory_path(self): - import run_kv - - env = {"COLLX_RDMA_DEVICES": "mlx5_0,mlx5_1"} - run_kv.export_ucx_selectors(env, device="") - self.assertEqual(env["UCX_NET_DEVICES"], "mlx5_0:1,mlx5_1:1") - - def test_ports_in_selectors_pass_through(self): - import run_kv - - env = {"COLLX_RDMA_DEVICES": "mlx5_18:1,mlx5_19"} - run_kv.export_ucx_selectors(env) - self.assertEqual(env["UCX_NET_DEVICES"], "mlx5_18:1,mlx5_19:1") - - def test_positive_tls_list_without_cuda_is_dropped(self): - # b300 exports UCX_TLS=rc cluster-wide; an RC-only context closes the - # cuda mds and NIXL registration of VRAM fails with NIXL_ERR_BACKEND. - # Extending the list with cuda transports segfaults UCX rkey-config - # resolution on the first GET, so the list is dropped entirely. - import run_kv - - env = {"UCX_TLS": "rc"} - run_kv.export_ucx_selectors(env) - self.assertNotIn("UCX_TLS", env) - - def test_tls_lists_already_covering_cuda_stay_untouched(self): - import run_kv - - for tls in ("^tcp", "rc,cuda_copy", "all"): - env = {"UCX_TLS": tls} - run_kv.export_ucx_selectors(env) - self.assertEqual(env["UCX_TLS"], tls) - env = {} - run_kv.export_ucx_selectors(env) - self.assertNotIn("UCX_TLS", env) + for env, device, expected in self.CASES: + with self.subTest(env=env, device=device): + env = dict(env) + run_kv.export_ucx_selectors(env, device=device) + self.assertEqual({k: v for k, v in env.items() if k.startswith("UCX_")}, expected) def _kv_document(status="success", sku="b200-nscale"): @@ -396,37 +358,22 @@ def test_descriptor_budget_sheds_batches_but_keeps_a_chartable_ladder(self): self.assertEqual(allowed[131072], [1, 2, 4, 8, 16]) # 1538, floor self.assertEqual(allowed[524288], [1, 2, 4, 8, 16]) # 6146, floor - def test_pool_budget_sheds_largest_batches_not_the_point(self): - # A point whose largest batch cannot fit the pool budget must survive - # with the batches that do: pin the budget between the batch-4 and - # batch-16 pool sizes. - import kv_workload - import run_kv - - args = self._args(isl_ladder="32768", batch_sizes="1 4 16") - budget = kv_workload.plan_config("dsv4", "fp8", 32768, 256, - 2.0, batch_max=4)["pool_bytes"] - args.pool_budget = budget - points, _isls, _batches = run_kv._grid(args) - self.assertEqual(points[0][1], [1, 4]) - self.assertLessEqual(points[0][0]["pool_bytes"], budget) - - def test_pool_budget_overrides_the_ladder_floor(self): - # The descriptor floor keeps the LADDER_FLOOR smallest batches, but - # the pool budget is a hard memory limit and must still shed a - # floor-kept batch. Pin the budget to the 512k point's batch-1 pool - # size: every larger batch survives the descriptor floor, then the - # pool loop must drop them all, leaving [1]. + def test_pool_budget_sheds_largest_batches_even_below_the_ladder_floor(self): + # A point whose largest batch cannot fit the pool budget survives with the batches that + # do. The budget is a hard memory limit, so it sheds batches the descriptor floor keeps: + # pinned to the 512k point's batch-1 pool, only [1] remains. import kv_workload import run_kv - args = self._args(isl_ladder="524288") - budget = kv_workload.plan_config("dsv4", "fp8", 524288, 256, - 2.0, batch_max=1)["pool_bytes"] - args.pool_budget = budget - points, _isls, _batches = run_kv._grid(args) - self.assertEqual(points[0][1], [1]) - self.assertLessEqual(points[0][0]["pool_bytes"], budget) + for isl, batches, fit_batch, expected in ((32768, "1 4 16", 4, [1, 4]), + (524288, "1 2 4 8 16 32 64", 1, [1])): + with self.subTest(isl=isl): + args = self._args(isl_ladder=str(isl), batch_sizes=batches) + args.pool_budget = kv_workload.plan_config( + "dsv4", "fp8", isl, 256, 2.0, batch_max=fit_batch)["pool_bytes"] + points, _isls, _batches = run_kv._grid(args) + self.assertEqual(points[0][1], expected) + self.assertLessEqual(points[0][0]["pool_bytes"], args.pool_budget) class RegistrationChunking(unittest.TestCase): From 061e3ebbad182c8a87aa100597d86578191cedf6 Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Tue, 29 Sep 2026 00:24:45 +0800 Subject: [PATCH 3/4] CollectiveX: harden the kv suite and trim its plumbing and tests Correctness: - Wire the sweep seed into the block tables and record it. - Salt each rank's pool pattern so a loopback transfer fails verification. - Paint the fabric pool on-device instead of from a pool-sized host array. - Fail closed on an empty grid and on a multi-page-family pool over budget. - Release NIXL transfer handles after each row. - Keep hostname rendezvous on MNNVL: the GB network block is kv-rdma only. - Never set a kv job timeout below the fleet-wide 350 minutes. Trims: shared library_version/spans/offset_lists helpers, the one-call registry-spec helper inlined, the kv precision filter and dead MC_FORCE_MNNVL export dropped, pip_install reused, one-pass summarize split, table-driven and registry-independent kv tests (the launcher-grep test removed, the argv codec moved onto the real case-args seam), and condensed kv docs with the stale b300 and 576 B claims fixed. --- .github/workflows/collectivex-sweep.yml | 2 +- collectivex/README.md | 22 +- collectivex/bandwidth.py | 12 +- collectivex/bench/kv_backend.py | 41 ++- collectivex/bench/kv_mooncake.py | 45 +-- collectivex/bench/kv_mori_io.py | 38 +- collectivex/bench/kv_nixl.py | 23 +- collectivex/bench/kv_pool.py | 43 ++- collectivex/bench/kv_workload.py | 39 +- collectivex/bench/run_kv.py | 84 +++-- collectivex/configs/kv_sweep.json | 1 - collectivex/configs/platform_config.json | 2 +- collectivex/docs/methodology.md | 100 +++-- collectivex/launchers/launch_gb-nv.sh | 10 +- collectivex/runtime/common.sh | 9 +- collectivex/runtime/prepare_backend.sh | 32 +- collectivex/summarize.py | 14 +- collectivex/sweep_matrix.py | 56 +-- collectivex/tests/test_kv_suite.py | 448 ++++++++--------------- collectivex/tests/test_kv_workload.py | 195 +++++----- collectivex/tests/test_runtime.py | 62 +++- 21 files changed, 579 insertions(+), 699 deletions(-) diff --git a/.github/workflows/collectivex-sweep.yml b/.github/workflows/collectivex-sweep.yml index b098e710fe..abf69b016f 100644 --- a/.github/workflows/collectivex-sweep.yml +++ b/.github/workflows/collectivex-sweep.yml @@ -154,7 +154,7 @@ jobs: timeout-minutes: ${{ matrix.job_timeout_minutes || 350 }} env: COLLX_BENCH: ${{ matrix.backend }} - # Set only on shards whose suite ships its own image (swap-blocks); blank keeps the pool's. + # Set only on shards that ship their own image (swap-blocks; mi355x mooncake); blank keeps the pool's. COLLX_IMAGE_OVERRIDE: ${{ matrix.image || '' }} COLLX_STAGED_IMAGE_DIR: ${{ matrix.staged_image_dir || '' }} # The fabric a shard claims (kv-transfer rdma legs on an MNNVL rack leave the NVL domain) diff --git a/collectivex/README.md b/collectivex/README.md index 54e788adf1..f0bf460ba8 100644 --- a/collectivex/README.md +++ b/collectivex/README.md @@ -170,23 +170,21 @@ benchmark's return code. ## KV-Cache Transfer Suite `kv-transfer` legs run 2 nodes x 1 GPU, the per-worker prefill/decode pair a disaggregated -deployment actually forms. Each moves bursts of 1 to 32 concurrent requests' paged KV as vLLM's -packed block-major descriptor lists over seed-keyed random block tables (a burst posts every -request's prepped transfer, then awaits them all), plus one contiguous bulk row as the wire-speed -baseline. The workload is transcribed from what vLLM allocates for the model it serves: `kv-dsv4` -is DeepSeek-V4-Pro's mixed cache (30 Compressed Sparse Attention layers at 4 tokens per 576 B entry -plus their 132 B indexer entries, 31 Heavily Compressed Attention layers at 128 tokens per entry, -and the 128-token sliding-window cache on all 61 layers; fp8 by architecture), at ISL 2k to 512k -and vLLM's 256-token block. `pull` (READ, vLLM NixlConnector) and `push` (WRITE, SGLang disagg) are -both timed from the initiator, with pattern verification of every request on the destination pool. +deployment actually forms. Each moves bursts of 1 to 32 concurrent requests' paged KV (vLLM's +packed block-major descriptor lists over seed-keyed random block tables) plus one contiguous bulk +row as the wire-speed baseline. The workload is `kv-dsv4`, DeepSeek-V4-Pro's mixed fp8 cache as +vLLM allocates it, at ISL 2k to 512k and vLLM's 256-token block. `pull` (READ, vLLM +NixlConnector) and `push` (WRITE, SGLang disagg) are both timed from the initiator, with every +request verified on the destination pool. The [methodology](docs/methodology.md#kv-cache-transfer-suite) +has the geometry and verification model. Backends are `nixl` (what Dynamo, vLLM, and SGLang ship), `mooncake` (the CUDA wheel links libcuda at import; mi355x runs AMD's atom-dev build push-only), and `mori-io` (AMD's native engine), where the registry's `kv_backends` map enables them. No entry, no legs, mirroring `ll_backends`. An entry may restrict ops, pin an image, set a NIC filter, or lower the pool budget. Fabrics are `rdma` -(torch pools) and, on GB racks, `mnnvl` (cuMem FABRIC pools; see the methodology for the -bulk-vs-paged lane inversion that row exists to publish). The grid and per-pool scheduling live in -`configs/kv_sweep.json`; dispatch with `suites: kv-transfer`. +(torch pools) and, on GB racks, `nixl` also runs `mnnvl` (cuMem FABRIC pools; see the methodology +for the bulk-vs-paged lane inversion that row exists to publish). The grid and per-pool scheduling +live in `configs/kv_sweep.json`; dispatch with `suites: kv-transfer`. ## Workflow And Artifacts diff --git a/collectivex/bandwidth.py b/collectivex/bandwidth.py index 97cf33d4fc..c7e57825ca 100644 --- a/collectivex/bandwidth.py +++ b/collectivex/bandwidth.py @@ -32,7 +32,7 @@ from typing import NamedTuple sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) -from summarize import load_results # noqa: E402 (sibling consumer, stdlib-only) +from summarize import KV_SUITE, load_results # noqa: E402 (sibling consumer, stdlib-only) # stage moves nothing (bytes always 0); isolated_sum is a derived percentile sum, not a real # chained rate. Only these three carry a measured latency + bytes. @@ -184,6 +184,12 @@ def _provenance(document: dict) -> str: def render(documents: list[dict]) -> str: + # kv-transfer documents have their own row model (per-transfer, no tokens_per_rank or + # routing); this renderer reads only EP rows, and prints nothing for a kv-only leg. + documents = [d for d in documents + if d["identity"]["case_factors"]["case"].get("suite") != KV_SUITE] + if not documents: + return "" lines = [ "## CollectiveX EP bandwidth (per-GPU, wire-basis payload)", "", @@ -193,10 +199,6 @@ def render(documents: list[dict]) -> str: "marks an extrapolated alpha, and rungs failing the correctness gate are excluded.", "", ] - # kv-transfer documents have their own row model (per-transfer, no tokens_per_rank or - # routing); this renderer reads only EP rows. - documents = [d for d in documents - if d["identity"]["case_factors"]["case"].get("suite") != "kv-transfer"] for document in sorted(documents, key=_sort_key): case = document["identity"]["case_factors"]["case"] ep = _ep(document) diff --git a/collectivex/bench/kv_backend.py b/collectivex/bench/kv_backend.py index fec12acdc1..b4d07ee1bd 100644 --- a/collectivex/bench/kv_backend.py +++ b/collectivex/bench/kv_backend.py @@ -15,19 +15,19 @@ import time +import kv_workload + class KVBackend: - """One transfer library on one rank. Subclasses implement the five hooks.""" + """One transfer library on one rank, constructed as ``Backend(args, role, + device)``. Subclasses implement the five hooks.""" name = "abstract" #: maturity mirrors EPBackend.maturity ("production" | "candidate"). maturity = "candidate" library_version: str | None = None - - def __init__(self, args, role: str, device) -> None: - self.args = args - self.role = role - self.device = device + #: the engine NIC filter this case ran under; None = library/UCX choice. + nic_filter: str | None = None # -- lifecycle ------------------------------------------------------------ def register(self, pool, bulk, reg_layout=None) -> None: @@ -49,6 +49,9 @@ def connect(self, peer: dict) -> None: """Consume the peer's payload; after this, transfers may be prepared.""" raise NotImplementedError + def release(self) -> None: + """Drop the transfers made for the grid point just measured (initiator).""" + def teardown(self) -> None: # pragma: no cover - adapter-specific pass @@ -73,6 +76,32 @@ def make_bulk(self, nbytes: int, op: str): raise NotImplementedError +def library_version(dists, module=None) -> str | None: + """The first installed distribution's version, else ``module.__version__``.""" + import importlib.metadata as md + + for name in dists: + try: + return md.version(name) + except md.PackageNotFoundError: + pass + return getattr(module, "__version__", None) + + +def spans(n: int, cap: int) -> list[tuple[int, int]]: + """[start, end) pieces of at most ``cap`` covering ``range(n)``.""" + return [(i, min(i + cap, n)) for i in range(0, n, cap)] + + +def offset_lists(cfg: dict, local_tables, remote_tables, local_base: int = 0, + remote_base: int = 0) -> tuple[list, list, list]: + """(local addrs, remote addrs, sizes) as Python lists: pool offsets plus the + given bases (0 for offset-addressed engines).""" + return ((local_base + kv_workload.page_offsets(cfg, local_tables)).tolist(), + (remote_base + kv_workload.page_offsets(cfg, remote_tables)).tolist(), + kv_workload.desc_sizes(cfg).tolist()) + + def time_bursts(transfers, warmup: int, reps: int) -> tuple[list[float], list[float]]: """(burst_ms, request_ms), warmups dropped. ``transfers`` is a list of (post, wait) pairs — one per request. A burst posts every request, then diff --git a/collectivex/bench/kv_mooncake.py b/collectivex/bench/kv_mooncake.py index da375e5d6e..62069d6d2d 100644 --- a/collectivex/bench/kv_mooncake.py +++ b/collectivex/bench/kv_mooncake.py @@ -21,7 +21,7 @@ from concurrent.futures import ThreadPoolExecutor import kv_workload -from kv_backend import KVBackend +from kv_backend import KVBackend, library_version, offset_lists, spans BATCH_CAP = 8192 @@ -53,26 +53,6 @@ def _import_engine(): return TransferEngine -def _engine_version(): - """The engine build actually imported, not the pin prepare_backend.sh - attempted: image-provided builds (b300's pre-0.3.12 lineage, AMD's - atom-dev tree) register under varying dist names or none at all, and - a null here is what let an image wheel masquerade as the pinned one.""" - import importlib.metadata as md - - for dist in ("mooncake-transfer-engine", "mooncake"): - try: - return md.version(dist) - except Exception: - pass - try: - import mooncake - - return getattr(mooncake, "__version__", None) - except Exception: - return None - - def _physical_gpu_index() -> int: """The physical GPU index behind this rank's visible device 0: GPU-paired NIC selection (rdma{gpu}) needs the host-level index, which the Slurm @@ -89,10 +69,15 @@ class MooncakeBackend(KVBackend): maturity = "production" def __init__(self, args, role, device): - super().__init__(args, role, device) TransferEngine = _import_engine() + import mooncake - self.library_version = _engine_version() + # The engine build actually imported, not the pin prepare_backend.sh + # attempted: image-provided builds (b300's pre-0.3.12 lineage, AMD's + # atom-dev tree) register under varying dist names or none at all, and + # a null here is what let an image wheel masquerade as the pinned one. + self.library_version = library_version( + ("mooncake-transfer-engine", "mooncake"), mooncake) # Same-fabric GB pairs: the NVLink-IPC transport claims cross-node # segments inside one NVLink domain and then fails the address import # (nvlink_transport "Requested address not found", first kv CI run on @@ -105,11 +90,12 @@ def __init__(self, args, role, device): # give the library guard 4x headroom; the runtime's per-case guard # still bounds a truly wedged transfer. os.environ.setdefault("MC_TRANSFER_TIMEOUT", "120") + if not args.socket_ifname: + raise RuntimeError("mooncake needs --socket-ifname for its P2P handshake address") self._engine = TransferEngine() self._ip = kv_workload.iface_ipv4(args.socket_ifname) local = f"{self._ip}:{args.kv_mc_port + (0 if role == 'target' else 1)}" - nic_filter = (getattr(args, "kv_device", "") or "").replace( - "{gpu}", str(_physical_gpu_index())) + nic_filter = args.kv_device.replace("{gpu}", str(_physical_gpu_index())) self.nic_filter = nic_filter or None rc = self._engine.initialize(local, "P2PHANDSHAKE", "rdma", nic_filter) if rc != 0: @@ -118,7 +104,7 @@ def __init__(self, args, role, device): self._pool = None self._bulk = None self._peer = None - workers = max(int(v) for v in str(getattr(args, "batch_sizes", "1")).split()) + workers = max(int(v) for v in args.batch_sizes.split()) self._exec = ThreadPoolExecutor(max_workers=workers) def register(self, pool, bulk, reg_layout=None) -> None: @@ -148,10 +134,9 @@ def wait(): def make_paged(self, cfg, op, local_tables, remote_tables): start = time.perf_counter() - local = (self._pool.ptr + kv_workload.page_offsets(cfg, local_tables)).tolist() - remote = (self._peer["pool_base"] + kv_workload.page_offsets(cfg, remote_tables)).tolist() - sizes = kv_workload.desc_sizes(cfg).tolist() - chunks = [(i, min(i + BATCH_CAP, len(local))) for i in range(0, len(local), BATCH_CAP)] + local, remote, sizes = offset_lists(cfg, local_tables, remote_tables, + self._pool.ptr, self._peer["pool_base"]) + chunks = spans(len(local), BATCH_CAP) session = self._peer["session"] func = self._engine.batch_transfer_sync_read if op == "pull" \ else self._engine.batch_transfer_sync_write diff --git a/collectivex/bench/kv_mori_io.py b/collectivex/bench/kv_mori_io.py index b1a2b3102b..b569eb9467 100644 --- a/collectivex/bench/kv_mori_io.py +++ b/collectivex/bench/kv_mori_io.py @@ -12,7 +12,7 @@ import time import kv_workload -from kv_backend import KVBackend +from kv_backend import KVBackend, library_version, offset_lists, spans BATCH_CAP = 16384 @@ -22,21 +22,15 @@ class MoRIIOBackend(KVBackend): maturity = "production" def __init__(self, args, role, device): - super().__init__(args, role, device) - from mori.io import (BackendType, IOEngine, IOEngineConfig, - MemoryLocationType, PollCqMode, RdmaBackendConfig) + import mori + from mori.io import (BackendType, EngineDesc, IOEngine, IOEngineConfig, + MemoryDesc, MemoryLocationType, PollCqMode, RdmaBackendConfig) self._gpu_location = MemoryLocationType.GPU - - try: - import mori - - self.library_version = getattr(mori, "__version__", None) - except Exception: - self.library_version = None - self._mori_io = __import__("mori.io", fromlist=["EngineDesc", "MemoryDesc"]) + self._engine_desc, self._memory_desc = EngineDesc, MemoryDesc + self.library_version = library_version((), mori) host = kv_workload.iface_ipv4(args.socket_ifname) if args.socket_ifname else "" - port = int(args.kv_mori_port) + (0 if role == "target" else 1) + port = args.kv_mori_port + (0 if role == "target" else 1) self._engine = IOEngine(key=role, config=IOEngineConfig(host=host, port=port)) # Library defaults only: four QPs plus transfer chunking wedged on the metal. self._engine.create_backend(BackendType.RDMA, RdmaBackendConfig( @@ -67,9 +61,9 @@ def publish(self) -> dict: } def connect(self, peer: dict) -> None: - self._engine.register_remote_engine(self._mori_io.EngineDesc.unpack(peer["engine"])) - remote_pool = self._mori_io.MemoryDesc.unpack(peer["pool"]) - remote_bulk = self._mori_io.MemoryDesc.unpack(peer["bulk"]) + self._engine.register_remote_engine(self._engine_desc.unpack(peer["engine"])) + remote_pool = self._memory_desc.unpack(peer["pool"]) + remote_bulk = self._memory_desc.unpack(peer["bulk"]) self._sessions = { "pool": self._engine.create_session(self._pool_mem, remote_pool), "bulk": self._engine.create_session(self._bulk_mem, remote_bulk), @@ -84,10 +78,8 @@ def _wait(statuses): def make_paged(self, cfg, op, local_tables, remote_tables): start = time.perf_counter() - local = kv_workload.page_offsets(cfg, local_tables).tolist() - remote = kv_workload.page_offsets(cfg, remote_tables).tolist() - sizes = kv_workload.desc_sizes(cfg).tolist() - chunks = [(i, min(i + BATCH_CAP, len(local))) for i in range(0, len(local), BATCH_CAP)] + local, remote, sizes = offset_lists(cfg, local_tables, remote_tables) + chunks = spans(len(local), BATCH_CAP) session = self._sessions["pool"] func = session.batch_read if op == "pull" else session.batch_write engine = self._engine @@ -116,16 +108,14 @@ def make_bulk(self, nbytes, op): session = self._sessions["bulk"] func = session.read if op == "pull" else session.write engine = self._engine - spans = [(offset, min(offset + self.BULK_WR_CAP, nbytes)) - for offset in range(0, nbytes, self.BULK_WR_CAP)] - + pieces = spans(nbytes, self.BULK_WR_CAP) statuses: list = [] def post(): statuses.clear() statuses.extend( func(start, start, end - start, engine.allocate_transfer_uid()) - for start, end in spans + for start, end in pieces ) def wait(): diff --git a/collectivex/bench/kv_nixl.py b/collectivex/bench/kv_nixl.py index 5bf36699d3..ceb23afcfc 100644 --- a/collectivex/bench/kv_nixl.py +++ b/collectivex/bench/kv_nixl.py @@ -14,7 +14,7 @@ import numpy as np import kv_workload -from kv_backend import KVBackend +from kv_backend import KVBackend, library_version # b300's CX NICs refuse cuda registrations somewhere between 7083 and 8847 MiB # (an ~8 GiB MR wall); UCX surfaces no error and the initiator later segfaults @@ -51,23 +51,12 @@ class NIXLBackend(KVBackend): maturity = "production" def __init__(self, args, role, device): - super().__init__(args, role, device) from nixl._api import nixl_agent, nixl_agent_config - try: - import importlib.metadata as md - - for dist_name in ("nixl", "nixl-cu13", "nixl-cu12"): - try: - self.library_version = md.version(dist_name) - break - except md.PackageNotFoundError: - continue - except Exception: - self.library_version = None + self.library_version = library_version(("nixl", "nixl-cu13", "nixl-cu12")) # The registry pin run_kv hands to UCX_NET_DEVICES for this case; # None means UCX chose among the operator inventory itself. - self.nic_filter = getattr(args, "kv_device", "") or None + self.nic_filter = args.kv_device or None # prog thread on, listener off: metadata goes through the harness exchange. self._agent = nixl_agent(role, nixl_agent_config(True, False, 0, backends=["UCX"])) self._handles = [] @@ -138,12 +127,16 @@ def make_bulk(self, nbytes, op): dtype=np.uint64) return self._make(local_np, remote_np, op) - def teardown(self) -> None: + def release(self) -> None: for handle in self._handles: try: self._agent.release_xfer_handle(handle) except Exception: pass + self._handles.clear() + + def teardown(self) -> None: + self.release() if self._peer is not None: try: self._agent.remove_remote_agent(self._remote_name) diff --git a/collectivex/bench/kv_pool.py b/collectivex/bench/kv_pool.py index 403102e7d7..94c6181e05 100644 --- a/collectivex/bench/kv_pool.py +++ b/collectivex/bench/kv_pool.py @@ -32,8 +32,8 @@ def __init__(self, nbytes: int, device: int): self._torch = torch self.ptr, self.nbytes, self.device = self._t.data_ptr(), nbytes, device - def fill_pattern(self) -> None: - kv_workload.fill_pattern(self._t) + def fill_pattern(self, salt: int = 0) -> None: + kv_workload.fill_pattern(self._t, salt) self._torch.cuda.synchronize() def fill_byte(self, value: int) -> None: @@ -58,17 +58,6 @@ class _AccessDesc(ctypes.Structure): _fields_ = [("location_type", c_int), ("location_id", c_int), ("flags", c_int)] -_PATTERNS: dict[int, np.ndarray] = {} - - -def _pattern(nbytes: int) -> np.ndarray: - if nbytes not in _PATTERNS: - chunks = nbytes // 256 - vals = ((np.arange(chunks, dtype=np.int64) * 131 + 7) & 0xFF).astype(np.uint8) - _PATTERNS[nbytes] = np.repeat(vals, 256) - return _PATTERNS[nbytes] - - class FabricPool: def __init__(self, nbytes: int, device: int): cu = self._cu = ctypes.CDLL("libcuda.so.1") @@ -103,15 +92,31 @@ def _check(self, code: int, what: str) -> None: if code != 0: raise RuntimeError(f"{what} -> CUresult {code}") - def _h2d(self, host: np.ndarray) -> None: + def _sync(self) -> None: + # Device-side memset/copy are asynchronous to the host; the pool must be + # painted before the barrier that lets the peer transfer. + self._check(self._cu.cuCtxSynchronize(), "sync") + + def fill_pattern(self, salt: int = 0) -> None: + """Upload one pattern period, then double it in place: the pattern is + PATTERN_PERIOD-periodic and every copy's destination sits at a multiple + of the period, so a copy of [0, n) lands the right bytes. No host + buffer of pool size is ever built.""" + tile = kv_workload.pattern_tile(salt) + filled = min(tile.nbytes, self.nbytes) self._check(self._cu.cuMemcpyHtoD_v2( - c_ulonglong(self.ptr), host.ctypes.data_as(c_void_p), c_size_t(host.nbytes)), "h2d") - - def fill_pattern(self) -> None: - self._h2d(_pattern(self.nbytes)) + c_ulonglong(self.ptr), tile.ctypes.data_as(c_void_p), c_size_t(filled)), "h2d") + while filled < self.nbytes: + n = min(filled, self.nbytes - filled) + self._check(self._cu.cuMemcpyDtoD_v2( + c_ulonglong(self.ptr + filled), c_ulonglong(self.ptr), c_size_t(n)), "d2d") + filled += n + self._sync() def fill_byte(self, value: int) -> None: - self._h2d(np.full(self.nbytes, value, dtype=np.uint8)) + self._check(self._cu.cuMemsetD8_v2( + c_ulonglong(self.ptr), ctypes.c_ubyte(value), c_size_t(self.nbytes)), "memset") + self._sync() def read8(self, offset: int): out = np.empty(8, dtype=np.uint8) diff --git a/collectivex/bench/kv_workload.py b/collectivex/bench/kv_workload.py index 83b089ee59..755a63b24c 100644 --- a/collectivex/bench/kv_workload.py +++ b/collectivex/bench/kv_workload.py @@ -29,8 +29,9 @@ is architectural, so the preset pins precision to "fp8". Pattern correctness: byte at offset o of a pool is derived from o (constant per -256-byte chunk), so any block's expected contents follow from its offset alone, -at any alignment. +256-byte chunk) and the owning rank's salt, so any block's expected contents +follow from its source offset and source rank alone, at any alignment; the +salt makes a transfer that reads its own pool (a loopback) fail verify. """ from __future__ import annotations @@ -131,9 +132,9 @@ def block_table(cfg: dict, seed: int, request: int = 0) -> dict: return tables -def table_seed(cfg: dict, side: str) -> int: - """Both ranks derive both sides' tables from the config alone — no exchange.""" - base = cfg["isl"] * 31 + cfg["page_tokens"] + len(cfg["preset"]) * 7 +def table_seed(cfg: dict, side: str, seed: int = 0) -> int: + """Both ranks derive both sides' tables from the config and the sweep seed — no exchange.""" + base = cfg["isl"] * 31 + cfg["page_tokens"] + len(cfg["preset"]) * 7 + seed * 7919 return base + (1000 if side == "local" else 0) @@ -165,24 +166,34 @@ def desc_array(base: int, cfg: dict, tables: dict, dev: int) -> np.ndarray: return out -def _chunk_byte(offset: int) -> int: - return ((offset >> 8) * 131 + 7) & 0xFF +# The pattern depends on (offset >> 8) mod 256 only, so it repeats every 64 KiB. +PATTERN_PERIOD = 1 << 16 -def fill_pattern(pool_u8) -> None: +def _chunk_byte(offset: int, salt: int = 0) -> int: + return ((offset >> 8) * 131 + 7 + 101 * salt) & 0xFF + + +def pattern_tile(salt: int = 0) -> np.ndarray: + """One PATTERN_PERIOD of the pattern from offset 0 (host uint8).""" + chunks = np.arange(PATTERN_PERIOD // 256, dtype=np.int64) + return np.repeat(((chunks * 131 + 7 + 101 * salt) & 0xFF).astype(np.uint8), 256) + + +def fill_pattern(pool_u8, salt: int = 0) -> None: """Paint the offset-derived pattern over the whole pool (torch uint8 tensor).""" import torch chunks = pool_u8.numel() // 256 view = pool_u8[: chunks * 256].view(chunks, 256) - vals = ((torch.arange(chunks, device=pool_u8.device, dtype=torch.int64) * 131 + 7) & 0xFF) + vals = (torch.arange(chunks, device=pool_u8.device, dtype=torch.int64) * 131 + 7 + 101 * salt) & 0xFF view.copy_(vals.to(torch.uint8)[:, None].expand(chunks, 256)) def verify_transfer(read8, cfg: dict, dst_tables: dict, src_tables: dict, - samples: int = 16, seed: int = 7) -> tuple[bool, str]: + src_salt: int = 0, samples: int = 16, seed: int = 7) -> tuple[bool, str]: """On the destination pool: packed block (region, dst[i]) must hold the - source pool's pattern at (region, src[i])'s offset. Each sample probes one + source pool's pattern (salted with ``src_salt``) at (region, src[i])'s offset. Each sample probes one layer's page inside the packed block, so the checks range over the whole descriptor. ``read8(offset)`` returns 8 destination-pool bytes (see kv_pool). Compared per byte, so any page alignment verifies exactly.""" @@ -195,7 +206,7 @@ def verify_transfer(read8, cfg: dict, dst_tables: dict, src_tables: dict, delta = layer * region["page_bytes"] src_off = int(src[i]) * region["packed_bytes"] + region["base"] + delta dst_off = int(dst[i]) * region["packed_bytes"] + region["base"] + delta - expected = bytes(_chunk_byte(src_off + j) for j in range(8)) + expected = bytes(_chunk_byte(src_off + j, src_salt) for j in range(8)) got = bytes(read8(dst_off)) if got != expected: return False, (f"region={region['name']} layer={layer} i={i} " @@ -217,6 +228,6 @@ def pcts(samples_ms: list[float]) -> dict: def iface_ipv4(iface: str) -> str: """IPv4 of a named interface (SIOCGIFADDR); the TCP bootstrap address.""" - sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) packed = struct.pack("256s", iface.encode()[:15]) - return socket.inet_ntoa(fcntl.ioctl(sock.fileno(), 0x8915, packed)[20:24]) + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock: + return socket.inet_ntoa(fcntl.ioctl(sock.fileno(), 0x8915, packed)[20:24]) diff --git a/collectivex/bench/run_kv.py b/collectivex/bench/run_kv.py index 9ebe4a80a6..78ea0d30d8 100644 --- a/collectivex/bench/run_kv.py +++ b/collectivex/bench/run_kv.py @@ -28,7 +28,7 @@ HERE = os.path.dirname(os.path.abspath(__file__)) sys.path[:0] = [HERE, os.path.dirname(HERE)] -import ep_harness # noqa: E402 (case_id/is_case_id + atomic write; stdlib-only) +import ep_harness # noqa: E402 (stdlib-only; safe before torch) import kv_workload # noqa: E402 from kv_backend import time_bursts # noqa: E402 @@ -283,9 +283,20 @@ def main() -> int: from kv_nixl import NIXLBackend as Backend points, isls, batches = _grid(args) + if not points: + print("ERROR: no grid point fits DESC_BUDGET and the pool budget", file=sys.stderr) + return 2 reg_layout = _harmonize(points) ops = args.ops.split() - pool_bytes = max(cfg["pool_bytes"] for cfg, _ in points) + pool_bytes = points[0][0]["pool_bytes"] # _harmonize gives every cfg the union pool + preset = args.workload_name.removeprefix("kv-") + pages = [int(v) for v in args.page_tokens.split()] + # One page family overruns the budget only by head-room (see _harmonize); + # several families each get their own slab and can overrun it outright. + if len({cfg["page_tokens"] for cfg, _ in points}) > 1 and pool_bytes > args.pool_budget: + print(f"ERROR: harmonized pool {pool_bytes} B across page sizes exceeds " + f"--pool-budget {args.pool_budget} B", file=sys.stderr) + return 2 bulk_bytes = min(max(cfg["req_bytes"] for cfg, _ in points), BULK_CAP) # RDMA registration pins the whole pool; a small inherited soft memlock @@ -296,7 +307,7 @@ def main() -> int: soft, hard = resource.getrlimit(resource.RLIMIT_MEMLOCK) need = pool_bytes + bulk_bytes - if soft != resource.RLIM_INFINITY and (hard == resource.RLIM_INFINITY or soft < hard): + if soft != hard: # soft <= hard always, so this raises soft to hard resource.setrlimit(resource.RLIMIT_MEMLOCK, (hard, hard)) soft = hard if soft != resource.RLIM_INFINITY and soft < need: @@ -310,7 +321,7 @@ def main() -> int: bulk = kv_pool.create(args.fabric, bulk_bytes, local_rank) def repaint(): - pool.fill_pattern() + pool.fill_pattern(salt=rank) bulk.fill_byte(0xAB if role == "target" else 0xCD) repaint() @@ -340,14 +351,16 @@ def record(row): def verify_burst(cfg, table_pairs): """Every request in the burst must land: a passing request 0 says nothing about the others, and concurrent same-session requests are - exactly where corruption would hide.""" + exactly where corruption would hide. The source is always the peer, + so its salt is the peer's rank.""" for r, (dst, src) in enumerate(table_pairs): - passed, detail = kv_workload.verify_transfer(pool.read8, cfg, dst, src) + passed, detail = kv_workload.verify_transfer(pool.read8, cfg, dst, src, + src_salt=1 - rank) if not passed: return False, f"request={r} {detail}" return True, "" - def measure(make, cfg_row: dict, op: str, verify_side: str, table_pairs=None): + def measure(make, cfg, cfg_row: dict, op: str, verify_side: str, table_pairs=None): """One grid point: initiator times bursts, then the verifying side checks.""" if role == "initiator": made = make() # one (post, wait, prep_seconds) per request in the burst @@ -359,10 +372,11 @@ def measure(make, cfg_row: dict, op: str, verify_side: str, table_pairs=None): burst_ms, request_ms = time_bursts(pairs, args.warmup, args.reps) samples.extend(burst_ms) request_samples.extend(request_ms) + backend.release() dist.barrier() # transfers complete before anyone inspects pools verdict = exchange_verdict( dist, role, verify_side, - lambda: verify_burst(cfg_row["_cfg"], table_pairs)) + lambda: verify_burst(cfg, table_pairs)) repaint() dist.barrier() if role != "initiator": @@ -376,7 +390,7 @@ def measure(make, cfg_row: dict, op: str, verify_side: str, table_pairs=None): gbps_incl_prep = (cfg_row["req_bytes"] * cfg_row["batch"] / (stats["p50"] + prep_ms) / 1e6) return { - **{k: v for k, v in cfg_row.items() if not k.startswith("_")}, + **cfg_row, "op": op, "prep_ms": round(prep_ms, 3), "latency_ms": {k: round(v, 3) for k, v in stats.items()}, @@ -389,23 +403,21 @@ def measure(make, cfg_row: dict, op: str, verify_side: str, table_pairs=None): } for cfg, allowed in points: - seed_t = kv_workload.table_seed(cfg, "remote") - seed_i = kv_workload.table_seed(cfg, "local") + seed_t = kv_workload.table_seed(cfg, "remote", args.seed) + seed_i = kv_workload.table_seed(cfg, "local", args.seed) target_tables = [kv_workload.block_table(cfg, seed_t, r) for r in range(allowed[-1])] initiator_tables = [kv_workload.block_table(cfg, seed_i, r) for r in range(allowed[-1])] base = { "kind": "paged", "preset": cfg["preset"], "isl": cfg["isl"], "page_tokens": cfg["page_tokens"], "layers": cfg["layers"], "page_bytes": cfg["page_bytes"], "descs": cfg["descs"], - "req_bytes": cfg["req_bytes"], "_cfg": cfg, + "req_bytes": cfg["req_bytes"], } for batch in allowed: for op in ops: - make = None - if role == "initiator": - make = lambda op=op, batch=batch: [ - backend.make_paged(cfg, op, initiator_tables[r], target_tables[r]) - for r in range(batch)] + # Called by measure on the initiator only, within this iteration. + make = lambda: [backend.make_paged(cfg, op, initiator_tables[r], target_tables[r]) + for r in range(batch)] # pull lands on the initiator's pool; push on the target's. # Every request in the burst is checked against its own tables. verify_side = "initiator" if op == "pull" else "target" @@ -413,28 +425,25 @@ def measure(make, cfg_row: dict, op: str, verify_side: str, table_pairs=None): (initiator_tables[r], target_tables[r]) if op == "pull" else (target_tables[r], initiator_tables[r]) for r in range(batch)] - record(measure(make, {**base, "batch": batch}, op, verify_side, table_pairs)) + record(measure(make, cfg, {**base, "batch": batch}, op, verify_side, + table_pairs)) for isl in isls: - preset = args.workload_name.removeprefix("kv-") - block_tokens = int(args.page_tokens.split()[0]) - cfg = kv_workload.plan_config(preset, args.precision, isl, block_tokens, - args.pool_slack) + cfg = kv_workload.plan_config(preset, args.precision, isl, pages[0], args.pool_slack) nbytes = min(cfg["req_bytes"], bulk_bytes) base = {"kind": "bulk", "preset": preset, "isl": isl, "page_tokens": None, "layers": cfg["layers"], "page_bytes": None, "descs": 1, "batch": 1, - "req_bytes": nbytes, "_cfg": cfg} + "req_bytes": nbytes} for op in ops: - make = (lambda op=op, n=nbytes: [backend.make_bulk(n, op)]) if role == "initiator" else None - record(measure(make, base, op, verify_side="none", table_pairs=None)) + record(measure(lambda: [backend.make_bulk(nbytes, op)], cfg, base, op, + verify_side="none")) backend.teardown() gathered: list = [None, None] - dist.all_gather_object(gathered, rows if rank == 1 else None) - rows = gathered[1] or [] - hosts: list = [None, None] - dist.all_gather_object(hosts, socket.gethostname()) + dist.all_gather_object(gathered, (socket.gethostname(), rows if rank == 1 else None)) + hosts = [host for host, _ in gathered] + rows = gathered[1][1] or [] all_ok = bool(rows) and all(r["verify"]["passed"] for r in rows) if rank == 0: @@ -443,10 +452,11 @@ def measure(make, cfg_row: dict, op: str, verify_side: str, table_pairs=None): os.environ.get("COLLECTIVEX_IMAGE", ""), all_ok, "transfer verification failed", workload={ "isl_ladder": isls, - "page_tokens": [int(v) for v in args.page_tokens.split()], + "page_tokens": pages, "batch_sizes": batches, "ops": ops, - "preset": kv_workload.PRESETS[args.workload_name.removeprefix("kv-")], + "seed": args.seed, + "preset": kv_workload.PRESETS[preset], }, measurement={ "payload_unit": "request-kv-bytes", @@ -460,9 +470,9 @@ def measure(make, cfg_row: dict, op: str, verify_side: str, table_pairs=None): implementation={ "name": args.backend, "fabric": args.fabric, - "library_version": getattr(backend, "library_version", None), - "maturity": getattr(backend, "maturity", "candidate"), - "nic_filter": getattr(backend, "nic_filter", None), + "library_version": backend.library_version, + "maturity": backend.maturity, + "nic_filter": backend.nic_filter, }, topology={ "device_product": torch.cuda.get_device_name(device), @@ -485,10 +495,10 @@ def measure(make, cfg_row: dict, op: str, verify_side: str, table_pairs=None): print(f"[run_kv] status={doc['outcome']['status']} rows={len(rows)}" + (f" -> {args.out}" if args.out else ""), flush=True) - flag = torch.tensor([int(all_ok)]) - dist.all_reduce(flag, op=dist.ReduceOp.MIN) + # all_ok is rank-invariant (both ranks hold rank 1's gathered rows); the + # barrier keeps a failing rank 1 from exiting before rank 0 writes --out. dist.barrier() - return 0 if int(flag.item()) else 3 + return 0 if all_ok else 3 if __name__ == "__main__": diff --git a/collectivex/configs/kv_sweep.json b/collectivex/configs/kv_sweep.json index cab200b161..4b01212ba7 100644 --- a/collectivex/configs/kv_sweep.json +++ b/collectivex/configs/kv_sweep.json @@ -1,5 +1,4 @@ { - "version": 2, "suite": "kv-transfer", "workloads": { "kv-dsv4": [ diff --git a/collectivex/configs/platform_config.json b/collectivex/configs/platform_config.json index b3fae34691..a0fe91444e 100644 --- a/collectivex/configs/platform_config.json +++ b/collectivex/configs/platform_config.json @@ -195,7 +195,7 @@ "ll_backends": {"mori": [8]}, "kv_backends": { "mori-io": ["rdma"], - "mooncake": {"fabrics": ["rdma"], "ops": "push", "image": "rocm/atom-dev:nightly_202606101403", + "mooncake": {"fabrics": ["rdma"], "ops": ["push"], "image": "rocm/atom-dev:nightly_202606101403", "device": "rdma{gpu}", "pool_budget": 21474836480} }, "fabric": {"nic": "Pollara 400GbE", "switch": "Arista 7060X6-64PE (Tomahawk5, 51.2T)"}, diff --git a/collectivex/docs/methodology.md b/collectivex/docs/methodology.md index d2185a8482..e873c7d811 100644 --- a/collectivex/docs/methodology.md +++ b/collectivex/docs/methodology.md @@ -546,39 +546,37 @@ packed DSV4 NIXL path actually registers and posts: per cache group, the physica transfer unit, and ONE contiguous descriptor covers all of that group's layers for the block (block-major `[block][layer]` layout, `packed_bytes = layers x page_bytes` per descriptor — vLLM's connector asserts exactly this shape). Per request, `isl` tokens (2k through 512k) at the -production block size of 256 tokens, addressed through seed-keyed random block tables on BOTH -sides (batched requests slice disjoint ranges of one permutation, as live requests never alias -blocks) — fragmentation is real but block-granular, the post-fragmentation layout a fragmented -allocator hands a connector. The suite deliberately does NOT explode each (layer, page) into its -own descriptor: that shape inflates descriptor counts by ~2 orders of magnitude (~2.1M vs ~6.1k -per 512k-ISL request) and inverts backend and fabric conclusions on descriptor-bound lanes. Each -request is its own prepped transfer; a burst posts all of them, then awaits all, the way a decode -step admits several requests at once. +production block size of 256 tokens, addressed through random block tables on BOTH sides, keyed +by the sweep's `seed` and the point (batched requests slice disjoint ranges of one permutation, +as live requests never alias blocks) — fragmentation is real but block-granular, the layout a +fragmented allocator hands a connector. The suite deliberately does NOT explode each (layer, +page) into its own descriptor: at block 256 that shape inflates descriptor counts ~30x (~186k vs +~6.1k per 512k-ISL request) and inverts backend and fabric conclusions on descriptor-bound lanes. +Each request is its own prepped transfer; a burst posts all of them, then awaits all, the way a +decode step admits several requests at once. Workload presets are transcribed from what vLLM actually allocates for the model class, region by region (geometry validated against vLLM commit `32ad1400d7`). `kv-dsv4` is DeepSeek-V4-Pro as -vLLM serves it (MXFP4 checkpoints included — quantization covers weights, the cache layout is -architectural): every token-state is 584 B of content (448 B NoPE + 128 B RoPE + 8 B fp8 scale, -the `fp8_ds_mla` layout), and each block's page is padded to a 576 B multiple at PAGE granularity -(FlashMLA packing), not per state. The config's `compress_ratios` interleave 30 Compressed Sparse -Attention layers (4 tokens per state) with 31 Heavily Compressed Attention layers (128 tokens per -state); CSA layers add a 132 B/state lightning-indexer cache (128 fp8 + 4 scale bytes); and all -61 layers keep a 128-token sliding window whose block size is FIXED at 64 tokens because the -window shares its physical tensor with the CSA cache — its page equals the CSA page byte for -byte. HCA's 128-token states force the model block size to a multiple of 128 (the model fails -closed on anything else); vLLM serves DSV4 at 256. Precision is pinned fp8 because the dtype mix -is architectural, and the whole thing computes to a few percent of an equivalent dense GQA-bf16 -cache. Each lane also reports a `bulk` row — one single-descriptor transfer of the request's -total bytes per ISL — as the contiguous baseline the paged rows are read against: logical payload -over host-observed completion of one contiguous post, NOT a proven physical wire rate (backends -may split large operations internally). Two budgets shed a point's largest batches rather than -dropping the point, and the smallest batches always survive, so a single request stays measurable -everywhere and every point keeps a chartable batch ladder: a per-rank pool budget (64 GiB, sized -to the fleet's smallest HBM, a hard memory limit the batch floor never overrides — on the packed -geometry this is the budget that actually bites, at the 512k point's largest batches), and a -per-burst descriptor budget (posting time is linear in batch x descriptors on the per-descriptor -floor; the packed grid sits far under it, and it stays as the fail-closed guard for future -presets or smaller block sizes). +vLLM serves it, with precision pinned fp8 because the cache layout is architectural (MXFP4 +checkpoints quantize weights only): every token-state is 584 B of content (448 B NoPE + 128 B +RoPE + 8 B fp8 scale, the `fp8_ds_mla` layout), and each block's page is padded to a 576 B +multiple at PAGE granularity (FlashMLA packing), not per state. The config's `compress_ratios` +interleave 30 Compressed Sparse Attention layers (4 tokens per state) with 31 Heavily Compressed +Attention layers (128 tokens per state); CSA layers add a 132 B/state lightning-indexer cache +(128 fp8 + 4 scale bytes); and all 61 layers keep a 128-token sliding window whose block size is +FIXED at 64 tokens because the window shares its physical tensor with the CSA cache — its page +equals the CSA page byte for byte. HCA's 128-token states force the model block size to a +multiple of 128; vLLM serves DSV4 at 256. + +Each lane also reports a `bulk` row — one single-descriptor transfer of the request's total bytes +per ISL — as the contiguous baseline the paged rows are read against: logical payload over +host-observed completion of one contiguous post, NOT a proven physical wire rate (backends may +split large operations internally). Two budgets shed a point's largest batches. The per-burst +descriptor budget never sheds the five smallest; the packed grid sits far under it, and it stays +as the fail-closed guard for future presets or smaller block sizes. The per-rank pool budget +(64 GiB, sized to the fleet's smallest HBM) is a hard memory limit: it can shed below that floor, +drops a point whose single request does not fit, and on the packed grid is the budget that +bites, at the 512k point's largest batches. Timing is host wall clock around post→completion — completion of a one-sided transfer is host-visible and no local kernel participates, so CUDA events have nothing to bracket. Descriptor @@ -591,25 +589,26 @@ percentiles, GB/s at p50 (burst-aggregate), and per-request completion marks: `r each individual request's host-observed completion offset from its burst's start (waits drain in posting order, so each mark upper-bounds that request's true completion) — the per-request latency distribution, distinct from the burst quantities, which are capacity numbers. A -verification verdict closes each point: the destination pool is pattern-checked after `pull` on -the initiator and after `push` on the target (an offset-derived byte pattern makes any block's -expected contents computable from its offset alone), covering every request in the burst against -its own block tables — concurrent same-session requests are exactly where corruption would hide, -so a passing request 0 is never taken as evidence for the others. Both pools are repainted -between points. A failed verify flips the document `invalid` and the leg red. +verification verdict closes each row: the destination pool is spot-checked after `pull` on the +initiator and after `push` on the target, 16 random 8-byte probes per request, each on a random +layer's page. Each rank paints its pool with an offset-derived byte pattern salted by rank, so +any block's expected contents follow from its source offset and the source rank's salt, and a +transfer that never happened fails even where both sides' tables pick the same block. Every +request in the burst is checked against its own block tables — concurrent same-session requests +are exactly where corruption would hide, so a passing request 0 is never taken as evidence for +the others. Both pools are repainted after every row. A failed verify flips the document +`invalid` and the leg red. A registry backend can carry restrictions: `ops` when a fabric serves one direction only (mooncake on mi355x runs `push` — AMD's atom-dev build moves WRITE at healthy rates over the GPU-paired Pollara NIC, while upstream ionic RDMA READ completes with retry-exceeded and one -failed READ poisons the engine, which is also why ATOM's production connector is write-only), -`image` when the build ships only inside a specific image, and `device` for engine NIC -filters (`{gpu}` expands to the physical GPU index; registering GPU memory on a non-paired -NIC fails and cross-rail pairs are unroutable), and `pool_budget` when the engine cannot -register the default pool on that pool's NICs (mooncake on the mi355x ionic NICs fails -`ibv_reg_mr` with ENOMEM past ~20 GiB while mori-io registers the same pool; points shed their -largest batches to fit). No kv backend is enabled on b300: the current b300 pool is the AWS -EFA cluster, and the one-rail rows measured on the earlier RoCE b300 pool do not carry over. The summary's `op` column names the -measured direction. +failed READ poisons the engine, which is also why ATOM's production connector is write-only); +`image` when the build ships only inside a specific image; `device` for engine NIC filters +(`{gpu}` expands to the physical GPU index; registering GPU memory on a non-paired NIC fails and +cross-rail pairs are unroutable); and `pool_budget` when the engine cannot register the default +pool on that pool's NICs (mooncake on the mi355x ionic NICs fails `ibv_reg_mr` with ENOMEM past +~20 GiB while mori-io registers the same pool; points shed their largest batches to fit). The +summary's `op` column names the measured direction. Fabrics are a case dimension. `rdma` runs on torch (cudaMalloc) pools. `mnnvl` allocates the pools with cuMem FABRIC handles (kv_pool.FabricPool; needs a live nvidia-imex domain), because @@ -622,16 +621,15 @@ rows carry; figures measured under the retired per-(layer, page) geometry are no are not restated here. Other lane facts, measured on the metal: single-WR bulk transfers above the provider's max -message size must be split (the MoRI adapter caps WRs at 1 GiB); and Mooncake is NVIDIA-only at -the binary level (the wheel links libcuda.so.1 at import; measured failing on mi355x). +message size must be split (the MoRI adapter caps WRs at 1 GiB); and the upstream PyPI Mooncake +wheel is NVIDIA-only (it links libcuda.so.1 at import), so mi355x runs AMD's atom-dev build. - ## Correctness - Scheduling is data: `configs/kv_sweep.json` holds the grid and, per pool, the allocation and the per-case hang guard. GB pools get the long budgets, because the mnnvl descriptor floor makes a leg run for hours and gb300 paces ~1.8x gb200 at ISL >= 131072. The guard fires inside the allocation, -so a slow case is a clean per-case kill, and each shard carries a GitHub job ceiling above its -allocation. KV legs run only when a dispatch names `kv-transfer` in `suites`. +so a slow case is a clean per-case kill, and each shard's GitHub job ceiling is +max(350, allocation + 30) minutes, so the job outlives its allocation. KV legs run only when a +dispatch names `kv-transfer` in `suites`. ## Correctness diff --git a/collectivex/launchers/launch_gb-nv.sh b/collectivex/launchers/launch_gb-nv.sh index 5bd04686aa..7ecbc0d12b 100644 --- a/collectivex/launchers/launch_gb-nv.sh +++ b/collectivex/launchers/launch_gb-nv.sh @@ -32,10 +32,6 @@ collx_require_vars COLLX_IMAGE COLLX_IMAGE_PLATFORM COLLX_PARTITION COLLX_ACCOUN PARTITION="$COLLX_PARTITION"; ACCOUNT="$COLLX_ACCOUNT"; SQUASH_DIR="$COLLX_SQUASH_DIR" [ -z "${COLLX_ENROOT_CACHE_PATH:-}" ] || export ENROOT_CACHE_PATH="$COLLX_ENROOT_CACHE_PATH" export NCCL_CUMEM_ENABLE=1 NCCL_MNNVL_ENABLE=1 -# Mooncake is MC_FORCE_MNNVL's only reader, and it makes the engine install ONLY its cross-node -# NVLink transport, which cannot open another host's segments in the pinned wheel -# (cudaIpcOpenMemHandle: invalid resource handle). The mooncake kv row declares the rdma lane. -[ "$COLLX_BENCH" = mooncake ] || export MC_FORCE_MNNVL=1 collx_apply_network_profile "$NODES" "$COLLX_TRANSPORT" collx_log "$PRODUCT nodes=$NODES x ${GPN}gpu world=$NGPUS bench=$COLLX_BENCH" @@ -52,10 +48,8 @@ allocation=(--partition="$PARTITION" --account="$ACCOUNT" --nodes="$NODES" [ -z "${COLLX_EXCLUDE_NODES:-}" ] || allocation+=(--exclude="$COLLX_EXCLUDE_NODES") collx_salloc_jobid "${allocation[@]}" [ -n "$JOB_ID" ] || collx_die "no JOB_ID from salloc" -# The kv rdma legs are the only gb-nv shards that leave the NVL domain; prove their pinned socket -# interface and HCAs on the allocation as every other scale-out launcher does. -if [ "$COLLX_TRANSPORT" != mnnvl ] \ - && ! collx_validate_network_profile_on_job "$JOB_ID" "$NODES" "$COLLX_TRANSPORT"; then +# No-op inside the NVL domain; the kv rdma legs (mnnvl-rdma) prove their socket iface and HCAs. +if ! collx_validate_network_profile_on_job "$JOB_ID" "$NODES" "$COLLX_TRANSPORT"; then collx_log_tail "${COLLX_NETWORK_PROFILE_LOG:-}" collx_die "allocated nodes failed the network profile" fi diff --git a/collectivex/runtime/common.sh b/collectivex/runtime/common.sh index 1f1603f571..664293a0dc 100644 --- a/collectivex/runtime/common.sh +++ b/collectivex/runtime/common.sh @@ -70,11 +70,11 @@ collx_set_placement() { TS="$(date -u +%Y-%m-%dT%H-%M-%SZ)" export COLLX_NGPUS="$NGPUS" COLLX_NODES="$NODES" export COLLX_GPUS_PER_NODE="$GPN" COLLX_SCALE_UP_DOMAIN="$SCALE_UP_DOMAIN" - COLLX_TRANSPORT="$4" - [ "$4" = mnnvl ] || [ "$NODES" -le 1 ] || COLLX_TRANSPORT="$4-rdma" # A shard that names the rdma fabric on an MNNVL rack (the kv-transfer rdma legs) leaves the # NVL domain, so it takes the same network profile and validation as any other scale-out. - [ "$4:${COLLX_FABRIC:-}" != mnnvl:rdma ] || [ "$NODES" -le 1 ] || COLLX_TRANSPORT=mnnvl-rdma + COLLX_TRANSPORT="$4" + { [ "$4" = mnnvl ] && [ "${COLLX_FABRIC:-}" != rdma ]; } || [ "$NODES" -le 1 ] \ + || COLLX_TRANSPORT="$4-rdma" export COLLX_TRANSPORT } @@ -428,6 +428,9 @@ collx_allocation_nodes_csv() { collx_resolve_slurm_rendezvous() { local job_id="$1" master_addr master_port socket_ifname="${COLLX_SOCKET_IFNAME:-}" [[ "$job_id" =~ ^[1-9][0-9]*$ ]] || collx_die "invalid rendezvous allocation" + # Inside an MNNVL domain the socket interface is never validated (the registry pins it for the + # kv rdma legs only), so rank zero keeps resolving by hostname there. + [ "${COLLX_TRANSPORT:-}" != mnnvl ] || socket_ifname="" # Relative node zero hosts global rank 0. Prefer the address on the validated socket interface: # a short hostname may resolve onto a management network that ranks cannot use. if [[ "$socket_ifname" =~ ^[A-Za-z][A-Za-z0-9_.-]{0,31}$ ]]; then diff --git a/collectivex/runtime/prepare_backend.sh b/collectivex/runtime/prepare_backend.sh index 33a159f153..994e8ed77f 100644 --- a/collectivex/runtime/prepare_backend.sh +++ b/collectivex/runtime/prepare_backend.sh @@ -340,6 +340,13 @@ uccl_activate() { [ "${COLLX_VENDOR:-nvidia}" != amd ] || export UCCL_EP_ENABLE_AGGRESSIVE_ATOMIC=1 } +# Noble-based images mark python externally managed (PEP 668); older pips never refuse, so they +# never reach the retry flag. +pip_install() { + python3 -m pip install -q --disable-pip-version-check --no-input "$@" \ + || python3 -m pip install -q --disable-pip-version-check --no-input --break-system-packages "$@" +} + uccl_install() { local root="$1" arch="$2" source_dir="/tmp/collectivex-uccl-$COLLX_UCCL_COMMIT" arch_env sp if [ -e "$root" ] || [ -L "$root" ]; then @@ -347,10 +354,7 @@ uccl_install() { 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)" - # Some sglang/rocm images mark the system env externally-managed (PEP 668). - { 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 \ + pip_install nanobind >&2 2>&1 \ || { collx_log "ERROR: UCCL nanobind install failed"; return 1; } collx_materialize_source "uccl-$COLLX_UCCL_COMMIT" "$source_dir" \ || { collx_log "ERROR: UCCL staged source is invalid"; return 1; } @@ -371,9 +375,7 @@ uccl_install() { # --no-deps: the wrapper's install_requires=["uccl"] resolves to the PyPI uccl-cu12 wheel, absent # on ROCm and wrong on CUDA too, since the from-source ep build already provides uccl.ep. ( 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 \ + && pip_install --no-deps . ) >&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; } @@ -559,16 +561,8 @@ FICHECK } # The kv-transfer wheels install into the named container, which persists for the job, so one -# install here serves every case srun. Noble-based images mark python externally managed -# (PEP 668); the retry flag is the uccl-prep pattern (older pips never refuse, so never reach it). -pip_install() { - python3 -m pip install -q --disable-pip-version-check --no-input "$@" \ - || python3 -m pip install -q --disable-pip-version-check --no-input --break-system-packages "$@" -} - -# nixl-cuXX directly: the `nixl` meta package depends on BOTH cu12 and cu13 variants, and an -# unpinned install under the image's stale pip resolved 1.0.1. ROCm images need nothing: -# sglang-rocm bundles nixl-cu12 with a ROCm-built UCX. +# install here serves every case srun. nixl-cuXX directly: the `nixl` meta package depends on +# BOTH cu12 and cu13 variants, and an unpinned install under the image's stale pip resolved 1.0.1. nixl_prepare() { python3 -c "import nixl" 2>/dev/null && return 0 pip_install 'nixl-cu13==1.3.2' \ @@ -586,8 +580,10 @@ mooncake_prepare() { collx_log "mooncake provided by the image" return 0 fi - pip_install 'mooncake-transfer-engine==0.3.12.post1' nvidia-cuda-runtime-cu12 \ + pip_install 'mooncake-transfer-engine==0.3.12.post1' 'nvidia-cuda-runtime-cu12==12.9.79' \ || { collx_log "ERROR: mooncake wheel install failed"; return 1; } + python3 -c "import mooncake.engine" \ + || { collx_log "ERROR: mooncake import failed after install"; return 1; } } main() { diff --git a/collectivex/summarize.py b/collectivex/summarize.py index de4409264a..4047f48c24 100644 --- a/collectivex/summarize.py +++ b/collectivex/summarize.py @@ -124,17 +124,13 @@ def _invalid_banner(documents: list[dict]) -> list[str]: def _kv_cell(rows: list[dict], kind: str, op: str, batch: str = "min"): """The largest-ISL row of a (kind, op) family -- the bandwidth-bound point -- at its smallest - or largest measured batch. Paged cells read the largest measured block size (the production - one when several ran).""" + or largest measured batch.""" matching = [r for r in rows if r.get("kind") == kind and r.get("op") == op] - if kind == "paged" and matching: - block = max(r["page_tokens"] for r in matching) - matching = [r for r in matching if r["page_tokens"] == block] if not matching: return "-", "-" isl = max(r["isl"] for r in matching) pick = min if batch == "min" else max - row = pick((r for r in matching if r["isl"] == isl), key=lambda r: r.get("batch", 1)) + row = pick((r for r in matching if r["isl"] == isl), key=lambda r: r["batch"]) return row["gbps_p50"], row["latency_ms"]["p50"] @@ -179,8 +175,10 @@ def key(document): def render(documents: list[dict]) -> str: """One table per suite present; the EP table also renders when nothing was found.""" - kv = [d for d in documents if d["identity"]["case_factors"]["case"].get("suite") == KV_SUITE] - ep = [d for d in documents if d not in kv] + kv, ep = [], [] + for document in documents: + is_kv = document["identity"]["case_factors"]["case"].get("suite") == KV_SUITE + (kv if is_kv else ep).append(document) parts = [render_ep(ep)] if ep or not kv else [] if kv: parts.append(render_kv(kv)) diff --git a/collectivex/sweep_matrix.py b/collectivex/sweep_matrix.py index cfb89a59af..d80b95f402 100644 --- a/collectivex/sweep_matrix.py +++ b/collectivex/sweep_matrix.py @@ -180,25 +180,7 @@ def _swap_shard(sku: str, profile_name: str) -> tuple[list[dict[str, Any]], dict return _runnable(sku, cases), shard -def _kv_backend_spec(value: Any) -> dict[str, Any]: - """A registry kv_backends value is a fabric list (the library runs the full sweep) or an - object that restricts it: `ops` (a backend that cannot serve one direction on a fabric, e.g. - mooncake on Pollara, where ionic RDMA READ is broken upstream), `image` (a backend shipped - only inside a specific image, e.g. AMD's atom-dev mooncake build), `device` (an engine NIC - filter template; `{gpu}` expands to the physical GPU index at runtime), and `pool_budget` - (bytes, where the engine cannot register the default pool on that pool's NICs).""" - if isinstance(value, list): - value = {"fabrics": value} - return { - "fabrics": value["fabrics"], - "ops": value.get("ops") or " ".join(KV_SWEEP["ops"]), - "image": value.get("image"), - "device": value.get("device") or "", - "pool_budget": value.get("pool_budget"), - } - - -def _kv_shards(sku: str, selected_precisions: set[str]) -> tuple[list, list[dict[str, Any]]]: +def _kv_shards(sku: str) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: """The kv-transfer shards one pool runs: one per (backend, fabric) its registry enables. A KV leg is 2 nodes x 1 GPU: the per-worker transfer pair an engine actually forms, not an @@ -212,15 +194,18 @@ def _kv_shards(sku: str, selected_precisions: set[str]) -> tuple[list, list[dict scheduling = KV_SWEEP["scheduling"].get(sku, KV_SWEEP["scheduling"]["default"]) requested, shards = [], [] for backend, raw in sorted(platform.get("kv_backends", {}).items()): - spec = _kv_backend_spec(raw) + # A registry value is a fabric list (the full sweep) or an object that restricts it: `ops` + # (a direction the fabric cannot serve, e.g. ionic RDMA READ on Pollara), `image` (a build + # shipped only in one image), `device` (an engine NIC filter; `{gpu}` expands to the + # physical GPU index at runtime) and `pool_budget` (bytes, where the NICs cannot register + # the default pool). + spec = raw if isinstance(raw, dict) else {"fabrics": raw} for fabric in spec["fabrics"]: cases = [] # A workload's dtype mix can be architectural (dsv4's fp8 slots), so the sweep config - # maps each workload to its precisions; a test pins the map to kv_workload's PRESETS. + # maps each workload to its precisions; a test proves every mapped cell plans. for workload, workload_precisions in KV_SWEEP["workloads"].items(): for precision in workload_precisions: - if selected_precisions and precision not in selected_precisions: - continue case = { "suite": KV_SWEEP["suite"], "workload": workload, @@ -233,8 +218,8 @@ def _kv_shards(sku: str, selected_precisions: set[str]) -> tuple[list, list[dict "isl_ladder": " ".join(map(str, KV_SWEEP["isl_ladder"])), "page_tokens": " ".join(map(str, KV_SWEEP["page_tokens"])), "batch_sizes": " ".join(map(str, KV_SWEEP["batch_sizes"])), - "ops": spec["ops"], - "kv_device": spec["device"], + "ops": " ".join(spec.get("ops", KV_SWEEP["ops"])), + "kv_device": spec.get("device", ""), "pool_slack": KV_SWEEP["pool_slack"], "seed": KV_SWEEP["seed"], "warmup": timing["warmup_per_trial"], @@ -246,22 +231,21 @@ def _kv_shards(sku: str, selected_precisions: set[str]) -> tuple[list, list[dict "scale_up_transport": platform["scale_up_transport"], "topology_class": f"{platform['product']}-kv-{fabric}", } - if spec["pool_budget"]: + if "pool_budget" in spec: case["pool_budget"] = spec["pool_budget"] case["case_id"] = ep_harness.case_id(sku, case) cases.append(case) - if not cases: - continue requested += _runnable(sku, cases) shards.append(_shard( sku, f"{sku}-kv-{backend}-{fabric}", backend, cases, - suite=KV_SWEEP["suite"], mode=fabric, fabric=fabric, + suite=KV_SWEEP["suite"], fabric=fabric, allocation_minutes=scheduling["allocation_minutes"], run_timeout=scheduling["run_timeout"], - # The GitHub job must outlive the allocation, or it cancels a healthy shard - # before the launcher's own guards act. - job_timeout_minutes=scheduling["allocation_minutes"] + 30, - **({"image": spec["image"]} if spec["image"] else {}), + # The GitHub job must outlive the allocation plus its salloc queue wait, or it + # cancels a healthy shard before the launcher's own guards act; never below the + # fleet-wide 350. + job_timeout_minutes=max(350, scheduling["allocation_minutes"] + 30), + **({"image": spec["image"]} if "image" in spec else {}), )) return requested, shards @@ -280,10 +264,8 @@ def resolve_matrix( selected_suites = _comma_subset("suites", suites, SUITES) if not selected_suites: raise SystemExit("--suites selects no suite") - if "ep" not in selected_suites and (backend != "all" or ep_sizes or modes): + if "ep" not in selected_suites and (backend != "all" or ep_sizes or precisions or modes): raise SystemExit("EP filters need the ep suite in --suites") - if not selected_suites & {"ep", "kv-transfer"} and precisions: - raise SystemExit("--precisions needs the ep or kv-transfer suite in --suites") if swap_profile not in SWAP_SWEEP["profiles"]: raise SystemExit( f"unknown --swap-profile {swap_profile!r}; have {sorted(SWAP_SWEEP['profiles'])}" @@ -421,7 +403,7 @@ def resolve_matrix( shards_by_sku.setdefault(sku, []).append(shard) if "kv-transfer" in selected_suites: for sku in selected_skus: - requested, kv_shards = _kv_shards(sku, selected_precisions) + requested, kv_shards = _kv_shards(sku) requested_cases += requested if kv_shards: shards_by_sku.setdefault(sku, []).extend(kv_shards) diff --git a/collectivex/tests/test_kv_suite.py b/collectivex/tests/test_kv_suite.py index 9805cc8afd..f6123bfd34 100644 --- a/collectivex/tests/test_kv_suite.py +++ b/collectivex/tests/test_kv_suite.py @@ -1,70 +1,95 @@ #!/usr/bin/env python3 -"""The kv-transfer suite's scheduling, argv codec, and summary contracts. +"""The kv-transfer suite's scheduling, control-plane, grid, and summary contracts. -Three seams keep KV legs honest end to end: sweep_matrix must emit kv shards -only for SKUs whose registry carries `kv_backends` (and must not perturb the EP -matrix at all); config.py must encode a kv case into run_kv argv behind the -`--entrypoint` marker the rank wrapper dispatches on; and summarize must render -kv documents in their own table instead of crashing the EP renderer. +sweep_matrix emits kv shards only where the registry carries `kv_backends` (and +never perturbs the EP matrix); run_kv's pure helpers (verdict exchange, UCX +selectors, burst timing, grid budgets, registration layout) are exercised on +CPU; summarize renders kv documents in their own table. The argv codec's +round trip lives with the other suites' in test_runtime.CaseArgvContract. """ from __future__ import annotations -import io +import argparse +import bisect import sys import unittest from pathlib import Path +from unittest import mock ROOT = Path(__file__).resolve().parents[1] -sys.path[:0] = [str(ROOT), str(ROOT / "bench"), str(ROOT / "runtime")] +sys.path[:0] = [str(ROOT), str(ROOT / "bench"), str(ROOT / "tests")] -import config as runtime_config # noqa: E402 import ep_harness # noqa: E402 +import kv_nixl # noqa: E402 +import kv_workload # noqa: E402 +import run_kv # noqa: E402 import summarize # noqa: E402 import sweep_matrix # noqa: E402 +from kv_backend import time_bursts # noqa: E402 +from test_chain import document as ep_document # noqa: E402 class KVMatrix(unittest.TestCase): - def test_kv_shards_only_where_the_registry_enables_them(self): + _BASE = {"scale_up_domain": 8, "scale_up_transport": "nvlink"} + PLATFORMS = { + # a fabric list runs the full sweep + "full": {**_BASE, "launcher": "single-slurm", "product": "full", + "kv_backends": {"nixl": ["rdma", "mnnvl"]}}, + # an object restricts it + "restricted": {**_BASE, "launcher": "mi-amds", "product": "restricted", + "kv_backends": {"mooncake": { + "fabrics": ["rdma"], "ops": ["push"], "image": "pinned:tag", + "device": "rdma{gpu}", "pool_budget": 123}}}, + # no entry, no legs + "absent": {**_BASE, "launcher": "single-slurm", "product": "absent"}, + } + + def _kv(self): + sweep = {**sweep_matrix.KV_SWEEP, "scheduling": { + "default": {"allocation_minutes": 100, "run_timeout": 5000}, + "restricted": {"allocation_minutes": 690, "run_timeout": 39600}, + }} + with mock.patch.object(sweep_matrix, "PLATFORMS", self.PLATFORMS), \ + mock.patch.object(sweep_matrix, "KV_SWEEP", sweep): + return sweep_matrix.resolve_matrix(suites="kv-transfer") + + def test_registry_entries_drive_the_shards(self): + matrix = self._kv() + shards = {(s["sku"], s["backend"], s["cases"][0]["mode"]): s for s in matrix["include"]} + self.assertEqual(set(shards), {("full", "nixl", "rdma"), ("full", "nixl", "mnnvl"), + ("restricted", "mooncake", "rdma")}) + self.assertEqual(len(matrix["requested_cases"]), + sum(len(s["cases"]) for s in shards.values())) + for (sku, _, fabric), shard in shards.items(): + with self.subTest(sku=sku, fabric=fabric): + self.assertEqual((shard["nodes"], shard["gpus_per_node"]), (2, 1)) + self.assertEqual(shard["launcher"], self.PLATFORMS[sku]["launcher"]) + for case in shard["cases"]: + self.assertEqual((case["suite"], case["ep"], case["mode"]), + ("kv-transfer", 2, fabric)) + self.assertEqual(case["case_id"], ep_harness.case_id(sku, case)) + full = shards[("full", "nixl", "rdma")] + self.assertNotIn("image", full) + self.assertEqual({(c["ops"], c["kv_device"], "pool_budget" in c) for c in full["cases"]}, + {("pull push", "", False)}) + restricted = shards[("restricted", "mooncake", "rdma")] + self.assertEqual(restricted["image"], "pinned:tag") + self.assertEqual({(c["ops"], c["kv_device"], c["pool_budget"]) + for c in restricted["cases"]}, {("push", "rdma{gpu}", 123)}) + # scheduling: the default vs a per-pool override; the GitHub job ceiling is + # max(350, allocation + 30) so it always outlives the allocation + self.assertEqual((full["allocation_minutes"], full["run_timeout"], + full["job_timeout_minutes"]), (100, 5000, 350)) + self.assertEqual((restricted["allocation_minutes"], restricted["run_timeout"], + restricted["job_timeout_minutes"]), (690, 39600, 720)) + + def test_the_real_registry_keeps_each_guard_inside_its_allocation(self): shards = sweep_matrix.resolve_matrix(suites="kv-transfer")["include"] self.assertTrue(shards, "registry carries kv_backends but no shard resolved") - enabled = { - sku for sku, platform in sweep_matrix.PLATFORMS.items() - if platform.get("kv_backends") - } - self.assertEqual({shard["sku"] for shard in shards}, enabled) - default = sweep_matrix.KV_SWEEP["scheduling"]["default"] for shard in shards: - self.assertEqual((shard["nodes"], shard["gpus_per_node"]), (2, 1)) - self.assertEqual(shard["launcher"], sweep_matrix.PLATFORMS[shard["sku"]]["launcher"]) - scheduling = sweep_matrix.KV_SWEEP["scheduling"].get(shard["sku"], default) - self.assertEqual(shard["allocation_minutes"], scheduling["allocation_minutes"]) - self.assertEqual(shard["run_timeout"], scheduling["run_timeout"]) - # The guard fires inside the allocation, and the GitHub job outlives it. - self.assertLess(shard["run_timeout"], shard["allocation_minutes"] * 60) - self.assertGreater(shard["job_timeout_minutes"], shard["allocation_minutes"]) - # the suite sweeps DeepSeek-V4-Pro's shape; its dtype mix is - # architectural, so one workload x one precision - self.assertEqual( - {(c["workload"], c["precision"]) for c in shard["cases"]}, - {("kv-dsv4", "fp8")}) - if shard["sku"] == "mi355x" and shard["backend"] == "mooncake": - # AMD's atom-dev build: push-only (upstream ionic RDMA READ is - # broken), GPU-paired NIC filter, shipped inside a pinned image, - # and a pool the ionic NICs can register. - self.assertEqual({c["ops"] for c in shard["cases"]}, {"push"}) - self.assertEqual({c["kv_device"] for c in shard["cases"]}, {"rdma{gpu}"}) - self.assertTrue(shard["image"].startswith("rocm/atom-dev:")) - self.assertEqual({c["pool_budget"] for c in shard["cases"]}, {20 << 30}) - else: - self.assertEqual({c["ops"] for c in shard["cases"]}, {"pull push"}) - self.assertNotIn("image", shard) - self.assertEqual({c["kv_device"] for c in shard["cases"]}, {""}) - self.assertFalse([c for c in shard["cases"] if "pool_budget" in c]) - for case in shard["cases"]: - self.assertEqual(case["suite"], "kv-transfer") - self.assertEqual(case["ep"], 2) - self.assertEqual( - case["case_id"], ep_harness.case_id(shard["sku"], case)) + with self.subTest(shard=shard["id"]): + self.assertLess(shard["run_timeout"], shard["allocation_minutes"] * 60) + self.assertGreater(shard["job_timeout_minutes"], shard["allocation_minutes"]) def test_kv_never_perturbs_the_ep_matrix(self): ep_only = sweep_matrix.resolve_matrix() @@ -76,98 +101,6 @@ def test_kv_never_perturbs_the_ep_matrix(self): ep_only["requested_cases"], [c for c in both["requested_cases"] if c["case"].get("suite") != "kv-transfer"]) - def test_unknown_suite_fails_closed(self): - with self.assertRaises(SystemExit): - sweep_matrix.resolve_matrix(suites="kv-transfr") - - def test_precision_filter_applies_to_kv(self): - # dsv4 is fp8-only, so a bf16-scoped dispatch has no kv legs at all. - matrix = sweep_matrix.resolve_matrix(suites="kv-transfer", precisions="bf16") - self.assertEqual(matrix["include"], []) - matrix = sweep_matrix.resolve_matrix(suites="kv-transfer", precisions="fp8") - self.assertTrue(matrix["include"]) - - def test_ep_only_filters_need_the_ep_suite(self): - for options in ({"backend": "deepep-v2"}, {"modes": "normal"}, {"ep_sizes": "8"}): - with self.subTest(options=options), self.assertRaises(SystemExit): - sweep_matrix.resolve_matrix(suites="kv-transfer", **options) - - def test_the_workload_map_matches_the_presets(self): - # sweep_matrix stays stdlib-only for the bare-runner matrix step, so it cannot import - # kv_workload; pin its workload -> precision map to the workload model instead. - import kv_workload - - for workload, precisions in sweep_matrix.KV_SWEEP["workloads"].items(): - preset = workload.removeprefix("kv-") - with self.subTest(workload=workload): - self.assertIn(preset, kv_workload.PRESETS) - for precision in precisions: - kv_workload.plan_config(preset, precision, 2048, 256) - - def test_every_kv_backend_passes_its_launcher_identity_gate(self): - # The launchers collx_die on unknown COLLX_BENCH values before anything - # runs; a registry kv backend its launcher rejects is a dead shard - # (this exact gap shipped once — every kv leg died at the gate). - launchers = Path(sweep_matrix.__file__).parent / "launchers" - for sku, platform in sweep_matrix.PLATFORMS.items(): - source = (launchers / f"launch_{platform['launcher']}.sh").read_text() - for backend in platform.get("kv_backends", {}): - with self.subTest(sku=sku, backend=backend): - self.assertRegex(source, rf"(^|[ |]){backend}( |\)|\s*\|)") - - -class KVArgvCodec(unittest.TestCase): - def _shard(self, backend="nixl"): - shards = sweep_matrix.resolve_matrix(suites="kv-transfer")["include"] - return next(shard for shard in shards if shard["backend"] == backend) - - @staticmethod - def _captured_argv(case, sku): - class _Stdout: - buffer = io.BytesIO() - - saved, sys.stdout = sys.stdout, _Stdout() - try: - runtime_config._emit_argv(case, 1, sku, "20260807", 0) - return sys.stdout.buffer.getvalue().decode().split("\0")[:-1] - finally: - sys.stdout = saved - - def _parsed(self, shard): - import argparse - - import run_kv - - argv = self._captured_argv(shard["cases"][0], shard["sku"]) - self.assertEqual(argv[:2], ["--entrypoint", "run_kv"]) - parser = argparse.ArgumentParser() - parser.add_argument("--backend", required=True, choices=["nixl", "mori-io", "mooncake"]) - run_kv.add_kv_args(parser) - return parser.parse_args(argv[2:]) - - def test_kv_case_round_trips_through_the_run_kv_parser(self): - import run_kv - - shard = self._shard() - case = shard["cases"][0] - args = self._parsed(shard) - self.assertEqual((args.backend, args.workload_name, args.precision, args.fabric), - (case["backend"], case["workload"], case["precision"], case["mode"])) - self.assertEqual((args.warmup, args.reps, args.trials), - (case["warmup"], case["reps"], case["trials"])) - self.assertEqual((args.batch_sizes, args.kv_device, args.ops), - (case["batch_sizes"], case["kv_device"], case["ops"])) - self.assertEqual(args.case_id, case["case_id"]) - self.assertEqual(args.pool_budget, run_kv.POOL_BUDGET) - self.assertEqual(args.out, f"results/{case['case_id']}_20260807-c000.json") - # run_kv recomputes the identity from the same factors and refuses a mismatch. - self.assertEqual(ep_harness.case_id(shard["sku"], run_kv.kv_case(args)), case["case_id"]) - - def test_a_pool_budget_reaches_run_kv(self): - shard = next(s for s in sweep_matrix.resolve_matrix(suites="kv-transfer")["include"] - if s["cases"][0].get("pool_budget")) - self.assertEqual(self._parsed(shard).pool_budget, shard["cases"][0]["pool_budget"]) - class _StubDist: """all_gather_object across a simulated 2-rank pair.""" @@ -183,28 +116,21 @@ class VerdictExchange(unittest.TestCase): """Bulk rows have no verifying side; that path crashed on the metal (gb200 smoke 22840: StopIteration on both ranks) before this contract existed.""" - def test_the_verifying_rank_supplies_the_verdict(self): - import run_kv - - verdict = run_kv.exchange_verdict( - _StubDist(None), "initiator", "initiator", lambda: (False, "bad page")) - self.assertEqual(verdict, {"passed": False, "detail": "bad page"}) - - def test_the_other_rank_receives_it(self): - import run_kv - - verdict = run_kv.exchange_verdict( - _StubDist({"passed": False, "detail": "bad page"}), "target", "initiator", - lambda: (True, "")) - self.assertEqual(verdict["passed"], False) + def test_every_rank_returns_the_verifying_ranks_verdict(self): + bad = {"passed": False, "detail": "bad page"} - def test_a_row_with_no_verifying_side_passes_without_a_gather_crash(self): - import run_kv + def never(): + raise AssertionError("must not verify") - verdict = run_kv.exchange_verdict( - _StubDist(None), "initiator", "none", - lambda: (_ for _ in ()).throw(AssertionError("must not verify"))) - self.assertEqual(verdict, {"passed": True, "detail": ""}) + for name, other, role, side, verify, expected in ( + ("verifying rank", None, "initiator", "initiator", lambda: (False, "bad page"), bad), + ("other rank", bad, "target", "initiator", never, bad), + ("no verifying side", None, "initiator", "none", never, + {"passed": True, "detail": ""}), + ): + with self.subTest(name): + self.assertEqual( + run_kv.exchange_verdict(_StubDist(other), role, side, verify), expected) class UCXSelectors(unittest.TestCase): @@ -238,8 +164,6 @@ class UCXSelectors(unittest.TestCase): ) def test_selectors(self): - import run_kv - for env, device, expected in self.CASES: with self.subTest(env=env, device=device): env = dict(env) @@ -247,111 +171,58 @@ def test_selectors(self): self.assertEqual({k: v for k, v in env.items() if k.startswith("UCX_")}, expected) -def _kv_document(status="success", sku="b200-nscale"): - def row(kind, page, op, gbps, p50, batch=1): - return {"kind": kind, "preset": "dsv4", "isl": 32768, "page_tokens": page, - "op": op, "descs": 1, "req_bytes": 1, "batch": batch, "prep_ms": 0.1, - "latency_ms": {"p50": p50, "p95": p50, "min": p50, "max": p50, "n": 48}, - "request_ms": {"p50": p50, "p95": p50, "min": p50, "max": p50, - "n": 48 * batch}, - "gbps_p50": gbps, "gbps_p50_incl_prep": gbps, - "verify": {"passed": status == "success", "detail": ""}} - - return { - "version": 1, - "record_type": "case-attempt", - "identity": {"case_factors": {"sku": sku, "case": { - "suite": "kv-transfer", "backend": "nixl", "workload": "kv-dsv4", - "mode": "rdma", "phase": "xfer", "ep": 2, "routing": "paged", - "precision": "fp8"}}}, - "measurement": {"rows": [ - row("paged", 256, "pull", 43.4, 53.1), - row("paged", 256, "pull", 96.2, 21.4, batch=16), - # a smaller measured block must lose to the production block size - row("paged", 128, "pull", 12.4, 185.2), - row("bulk", None, "pull", 48.3, 47.7), - row("paged", 256, "push", 48.4, 47.6), - ]}, - "topology": {"gpus_per_node": 1, "scale_up_domain": 8, "nodes": 2}, - "outcome": {"status": status, "reasons": []}, - } - - class BurstTiming(unittest.TestCase): def test_a_burst_posts_every_request_before_waiting_on_any(self): - from kv_backend import time_bursts - order = [] pairs = [(lambda i=i: order.append(("post", i)), lambda i=i: order.append(("wait", i))) for i in range(3)] burst_ms, request_ms = time_bursts(pairs, warmup=1, reps=2) - self.assertEqual(len(burst_ms), 2) - # one completion mark per request per kept rep, in posting order - self.assertEqual(len(request_ms), 2 * 3) self.assertEqual(order[:6], [("post", 0), ("post", 1), ("post", 2), ("wait", 0), ("wait", 1), ("wait", 2)]) + # warmups dropped; one completion mark per request per kept rep, offsets + # from the burst start (never decreasing), and the burst is its last mark + self.assertEqual((len(burst_ms), len(request_ms)), (2, 2 * 3)) + for rep, burst in enumerate(burst_ms): + marks = request_ms[3 * rep : 3 * rep + 3] + self.assertEqual(marks, sorted(marks)) + self.assertEqual(burst, marks[-1]) - def test_the_burst_sample_is_the_last_request_mark(self): - from kv_backend import time_bursts - pairs = [(lambda: None, lambda: None)] * 2 - burst_ms, request_ms = time_bursts(pairs, warmup=0, reps=1) - self.assertEqual(burst_ms[0], request_ms[-1]) - # marks are offsets from the burst start, so they never decrease - self.assertEqual(request_ms, sorted(request_ms)) +def _grid_args(**overrides): + base = dict(workload_name="kv-dsv4", precision="fp8", + isl_ladder="8192 32768 131072 524288", page_tokens="256", + batch_sizes="1 2 4 8 16 32 64", pool_slack=2.0, + pool_budget=run_kv.POOL_BUDGET, seed=67) + base.update(overrides) + return argparse.Namespace(**base) class KVGrid(unittest.TestCase): - @staticmethod - def _args(**overrides): - import argparse - - base = dict(workload_name="kv-dsv4", precision="fp8", - isl_ladder="8192 32768 131072 524288", page_tokens="256", - batch_sizes="1 2 4 8 16 32 64", pool_slack=2.0) - import run_kv - - base["pool_budget"] = run_kv.POOL_BUDGET - base.update(overrides) - return argparse.Namespace(**base) - def test_the_packed_grid_sheds_only_where_the_pool_budget_bites(self): # Packed block-major geometry: a 512k-ISL block-256 request is 6,146 # descriptors, so no batch on this ladder nears DESC_BUDGET. Only the # 512k point sheds, and via the pool budget: its batch-32 pool plans # ~118 GB against the 64 GiB budget, batch 16 fits at ~59 GB. - import run_kv - - points, isls, batches = run_kv._grid(self._args()) + points, isls, batches = run_kv._grid(_grid_args()) self.assertEqual((isls, batches), ([8192, 32768, 131072, 524288], [1, 2, 4, 8, 16, 32, 64])) allowed = {cfg["isl"]: allowed for cfg, allowed in points} - self.assertEqual(allowed[8192], [1, 2, 4, 8, 16, 32, 64]) - self.assertEqual(allowed[32768], [1, 2, 4, 8, 16, 32, 64]) - self.assertEqual(allowed[131072], [1, 2, 4, 8, 16, 32, 64]) - self.assertEqual(allowed[524288], [1, 2, 4, 8, 16]) + self.assertEqual(allowed, {8192: batches, 32768: batches, 131072: batches, + 524288: [1, 2, 4, 8, 16]}) for cfg, batch_list in points: + # one descriptor per packed block: 3 full-ISL groups + 2 window blocks self.assertEqual(cfg["descs"], 3 * -(-cfg["isl"] // 256) + 2) self.assertLessEqual(cfg["pool_bytes"], run_kv.POOL_BUDGET) - for batch in batch_list[run_kv.LADDER_FLOOR:]: - self.assertLessEqual(batch * cfg["descs"], run_kv.DESC_BUDGET) def test_descriptor_budget_sheds_batches_but_keeps_a_chartable_ladder(self): # DESC_BUDGET stays as the fail-closed guard for future presets whose # bursts are descriptor-bound. Pin it to 4 requests' descriptors at # the largest ISL: batches above the per-point allowance shed, but the # LADDER_FLOOR smallest batches always survive so every point keeps a - # chartable batch ladder (the frontier draws its line through the - # ladder at the largest measured ISL). - import kv_workload - import run_kv - + # chartable batch ladder. probe = kv_workload.plan_config("dsv4", "fp8", 524288, 256) - saved, run_kv.DESC_BUDGET = run_kv.DESC_BUDGET, 4 * probe["descs"] - try: - points, _isls, _batches = run_kv._grid(self._args()) - finally: - run_kv.DESC_BUDGET = saved + with mock.patch.object(run_kv, "DESC_BUDGET", 4 * probe["descs"]): + points, _isls, _batches = run_kv._grid(_grid_args()) allowed = {cfg["isl"]: allowed for cfg, allowed in points} self.assertEqual(allowed[8192], [1, 2, 4, 8, 16, 32, 64]) # 98 descs/req self.assertEqual(allowed[32768], [1, 2, 4, 8, 16, 32]) # 386 @@ -359,16 +230,12 @@ def test_descriptor_budget_sheds_batches_but_keeps_a_chartable_ladder(self): self.assertEqual(allowed[524288], [1, 2, 4, 8, 16]) # 6146, floor def test_pool_budget_sheds_largest_batches_even_below_the_ladder_floor(self): - # A point whose largest batch cannot fit the pool budget survives with the batches that - # do. The budget is a hard memory limit, so it sheds batches the descriptor floor keeps: - # pinned to the 512k point's batch-1 pool, only [1] remains. - import kv_workload - import run_kv - + # The budget is a hard memory limit, so it sheds batches the descriptor + # floor keeps: pinned to the 512k point's batch-1 pool, only [1] remains. for isl, batches, fit_batch, expected in ((32768, "1 4 16", 4, [1, 4]), (524288, "1 2 4 8 16 32 64", 1, [1])): with self.subTest(isl=isl): - args = self._args(isl_ladder=str(isl), batch_sizes=batches) + args = _grid_args(isl_ladder=str(isl), batch_sizes=batches) args.pool_budget = kv_workload.plan_config( "dsv4", "fp8", isl, 256, 2.0, batch_max=fit_batch)["pool_bytes"] points, _isls, _batches = run_kv._grid(args) @@ -377,15 +244,12 @@ def test_pool_budget_sheds_largest_batches_even_below_the_ladder_floor(self): class RegistrationChunking(unittest.TestCase): - # b300's NICs refuse cuda registrations past ~8 GiB, so the NIXL adapter - # registers the pool in pieces. The pieces must never cut through a - # descriptor of ANY planned config, which _harmonize guarantees by giving - # every config one shared region layout. + # The NIXL adapter registers an oversized pool in pieces. The pieces must + # never cut through a descriptor of ANY planned config, which _harmonize + # guarantees by giving every config one shared region layout. def test_harmonize_makes_region_bases_config_invariant(self): - import run_kv - - points, _isls, _batches = run_kv._grid(KVGrid._args()) + points, _isls, _batches = run_kv._grid(_grid_args()) layout = run_kv._harmonize(points) total = sum(nbytes for _, _, nbytes in layout) running = 0 @@ -395,46 +259,28 @@ def test_harmonize_makes_region_bases_config_invariant(self): for cfg, _ in points: self.assertEqual(cfg["pool_bytes"], total) for region, (base, packed, nbytes) in zip(cfg["regions"], layout): - self.assertEqual(region["base"], base) - self.assertEqual(region["packed_bytes"], packed) - self.assertEqual(region["pool_blocks"], nbytes // packed) + self.assertEqual((region["base"], region["packed_bytes"], region["pool_blocks"]), + (base, packed, nbytes // packed)) self.assertLessEqual(region["blocks_req"], region["pool_blocks"]) - def test_reg_spans_cut_each_region_on_its_own_packed_grid(self): - import kv_nixl - import run_kv - - points, _isls, _batches = run_kv._grid(KVGrid._args()) + def test_spans_tile_the_pool_and_no_descriptor_straddles_a_cut(self): + # A tiny cap on a small grid forces many cuts; every block any config + # can address must land whole inside one registered piece. + cap = 1 << 24 + points, _isls, _batches = run_kv._grid( + _grid_args(isl_ladder="2048 8192", batch_sizes="1 4")) layout = run_kv._harmonize(points) total = sum(nbytes for _, _, nbytes in layout) - spans = kv_nixl.reg_spans(total, layout) - # The full test grid plans a pool far past one chunk. - self.assertGreater(len(spans), 1) - # Exact in-order coverage, no gap, no overlap. + spans = kv_nixl.reg_spans(total, layout, cap=cap) + self.assertGreater(len(spans), len(layout)) + # exact in-order coverage, no gap, no overlap, each piece within the cap self.assertEqual(spans[0][0], 0) for (a_off, a_len), (b_off, _) in zip(spans, spans[1:]): self.assertEqual(a_off + a_len, b_off) self.assertEqual(sum(length for _, length in spans), total) for off, length in spans: - base, packed, _ = next(entry for entry in reversed(layout) - if entry[0] <= off) - self.assertEqual((off - base) % packed, 0) - self.assertLessEqual(length, max(kv_nixl.REG_CHUNK_BYTES, packed)) - - def test_no_descriptor_straddles_a_registration_cut(self): - # Every block any config can ever address must land whole inside one - # registered piece; a tiny cap on a small grid forces many cuts. - import bisect - - import kv_nixl - import run_kv - - args = KVGrid._args(isl_ladder="2048 8192", batch_sizes="1 4") - points, _isls, _batches = run_kv._grid(args) - layout = run_kv._harmonize(points) - total = sum(nbytes for _, _, nbytes in layout) - spans = kv_nixl.reg_spans(total, layout, cap=1 << 24) - self.assertGreater(len(spans), len(layout)) + packed = next(entry for entry in reversed(layout) if entry[0] <= off)[1] + self.assertLessEqual(length, max(cap, packed)) starts = [off for off, _ in spans] straddles = [] for cfg, _ in points: @@ -448,16 +294,37 @@ def test_no_descriptor_straddles_a_registration_cut(self): self.assertEqual(straddles, []) def test_without_a_layout_the_pool_registers_whole(self): - import kv_nixl + for layout in (None, []): + self.assertEqual(kv_nixl.reg_spans(123456, layout), [(0, 123456)]) + + +def _kv_document(status="success"): + def row(kind, page, op, gbps, p50, batch=1): + return {"kind": kind, "isl": 32768, "page_tokens": page, "op": op, "batch": batch, + "latency_ms": {"p50": p50}, "gbps_p50": gbps} - self.assertEqual(kv_nixl.reg_spans(123456, None), [(0, 123456)]) - self.assertEqual(kv_nixl.reg_spans(123456, []), [(0, 123456)]) + return { + "version": 1, + "identity": {"case_factors": {"sku": "b200-nscale", "case": { + "suite": "kv-transfer", "backend": "nixl", "workload": "kv-dsv4", + "mode": "rdma", "precision": "fp8"}}}, + "measurement": {"rows": [ + row("paged", 256, "pull", 43.4, 53.1), + row("paged", 256, "pull", 96.2, 21.4, batch=16), + # a smaller measured block must lose to the production block size + row("paged", 128, "pull", 12.4, 185.2), + row("bulk", None, "pull", 48.3, 47.7), + row("paged", 256, "push", 48.4, 47.6), + ]}, + "outcome": {"status": status, "reasons": []}, + } class KVSummary(unittest.TestCase): def test_kv_documents_render_their_own_table(self): text = summarize.render([_kv_document()]) self.assertIn("KV-transfer results", text) + self.assertNotIn("EP results", text) self.assertIn("| pull | 43.4 | 96.2 | 48.3 | 53.1 |", text) def test_a_push_only_document_reads_its_push_lane(self): @@ -470,11 +337,18 @@ def test_a_push_only_document_reads_its_push_lane(self): self.assertNotIn("INVALID", text) def test_kv_invalid_counts_in_the_banner(self): - text = summarize.render([_kv_document(status="invalid")]) - self.assertIn("INVALID", text) + self.assertIn("INVALID", summarize.render([_kv_document(status="invalid")])) + + def test_mixed_documents_render_one_table_per_suite(self): + text = summarize.render([ep_document(with_period=True), _kv_document()]) + self.assertIn("EP results", text) + self.assertIn("KV-transfer results", text) + self.assertLess(text.index("EP results"), text.index("KV-transfer results")) - def test_ep_documents_do_not_grow_a_kv_table(self): - self.assertNotIn("KV-transfer results", summarize.render([])) + def test_an_empty_render_has_only_the_ep_table(self): + text = summarize.render([]) + self.assertIn("EP results", text) + self.assertNotIn("KV-transfer results", text) if __name__ == "__main__": diff --git a/collectivex/tests/test_kv_workload.py b/collectivex/tests/test_kv_workload.py index 2c99dd0256..17a39d7a85 100644 --- a/collectivex/tests/test_kv_workload.py +++ b/collectivex/tests/test_kv_workload.py @@ -5,9 +5,9 @@ contiguous descriptor covers all the group's layers for one physical block (vLLM's packed DSV4 NIXL shape), block tables are seed-keyed permutations both ranks derive independently (batched requests slicing disjoint ranges of one -permutation), and an offset-derived pattern makes any byte's expected value -computable from its offset alone. These tests pin that math with hand-computed -cases validated against vLLM commit 32ad1400d7 (state content 584 B, page +permutation), and an offset-derived, per-rank-salted pattern makes any byte's +expected value computable from its source offset and source rank alone. +These tests pin that math with hand-computed cases validated against vLLM commit 32ad1400d7 (state content 584 B, page padded to a 576 B multiple at block granularity, one descriptor per packed block); the torch fill path is exercised on metal by the suite itself (a wrong fill fails every verify row loudly). @@ -54,9 +54,11 @@ def test_dsv4_regions_by_hand(self): self.assertEqual( (regions["c128a"]["layers"], regions["c128a"]["page_bytes"], regions["c128a"]["blocks_req"]), (31, 1_728, 2)) + # the window shares C4A's physical tensor, so its page equals C4A's self.assertEqual( (regions["swa"]["layers"], regions["swa"]["block_tokens"], - regions["swa"]["blocks_req"]), (61, 64, 2)) + regions["swa"]["page_bytes"], regions["swa"]["blocks_req"]), + (61, 64, 37_440, 2)) self.assertEqual(cfg["descs"], 2 + 2 + 2 + 2) self.assertEqual(cfg["req_bytes"], 2 * (30 * 37_440 + 30 * 8_640 + 31 * 1_728 + 61 * 37_440)) @@ -65,37 +67,10 @@ def test_dsv4_regions_by_hand(self): sum(r["pool_blocks"] * r["packed_bytes"] for r in cfg["regions"])) - def test_alignment_pads_the_page_not_each_state(self): - # 64 states * 584 B = 37,376 -> padded once per page to 37,440. The - # old per-entry 576 B model would give 64 * 576 = 36,864 — vLLM pads - # at page granularity, not per state. - cfg = kv_workload.plan_config("dsv4", "fp8", 512, 256) - c4a = {r["name"]: r for r in cfg["regions"]}["c4a"] - self.assertEqual(c4a["page_bytes"], 37_440) - self.assertNotEqual(c4a["page_bytes"], 64 * 576) - - def test_swa_shares_the_c4a_page_size(self): - # Both block types live in one physical tensor: a 64-token window - # block (1 token/state) and a 256-token C4A block (4 tokens/state) - # are the same 64 states -> byte-identical pages. - cfg = kv_workload.plan_config("dsv4", "fp8", 512, 256) - regions = {r["name"]: r for r in cfg["regions"]} - self.assertEqual(regions["swa"]["page_bytes"], regions["c4a"]["page_bytes"]) - - def test_one_descriptor_per_block_at_the_big_isl(self): - # 512k tokens at block 256: 2048 blocks per non-window group + 2 - # window blocks = 6,146 descriptors per request — the packed shape - # vLLM's connector asserts, not a per-(layer, page) explosion. - cfg = kv_workload.plan_config("dsv4", "fp8", 524_288, 256) - self.assertEqual(cfg["descs"], 2048 * 3 + 2) - - def test_dsv4_window_caps_at_128_tokens(self): + def test_a_short_request_uses_only_its_own_window_tokens(self): + # min(isl, 128) = 64 tokens -> one 64-token window block small = kv_workload.plan_config("dsv4", "fp8", 64, 256) - large = kv_workload.plan_config("dsv4", "fp8", 32_768, 256) - window = {r["name"]: r for r in large["regions"]}["swa"] - self.assertEqual(window["blocks_req"], 2) # 128 tokens / 64 per block - self.assertEqual({r["name"]: r for r in small["regions"]}["swa"]["blocks_req"], - 1) # min(isl, 128) = 64 tokens + self.assertEqual({r["name"]: r for r in small["regions"]}["swa"]["blocks_req"], 1) def test_block_sizes_that_split_a_state_fail_closed(self): # C128A's 128-token states force the model block size to a multiple @@ -125,22 +100,26 @@ def test_batch_max_grows_the_pool_for_disjoint_requests(self): class Tables(unittest.TestCase): - def test_deterministic_and_distinct_per_side(self): + def test_deterministic_and_distinct_per_side_and_seed(self): cfg = kv_workload.plan_config("dsv4", "fp8", 4096, 256) - local = kv_workload.block_table(cfg, kv_workload.table_seed(cfg, "local")) - remote = kv_workload.block_table(cfg, kv_workload.table_seed(cfg, "remote")) - again = kv_workload.block_table(cfg, kv_workload.table_seed(cfg, "local")) + + def table(side, seed=67): + return kv_workload.block_table(cfg, kv_workload.table_seed(cfg, side, seed)) + + local, remote, again, reseeded = table("local"), table("remote"), table("local"), \ + table("local", seed=68) for region in cfg["regions"]: name, blocks_req = region["name"], region["blocks_req"] self.assertTrue((local[name] == again[name]).all()) self.assertFalse((local[name] == remote[name]).all()) + self.assertFalse((local[name] == reseeded[name]).all()) # distinct in-range blocks (fragmented, never aliased) self.assertEqual(len(set(local[name].tolist())), blocks_req) self.assertTrue((local[name] < region["pool_blocks"]).all()) def test_batched_requests_slice_disjoint_blocks(self): cfg = kv_workload.plan_config("dsv4", "fp8", 512, 256, batch_max=4) - seed = kv_workload.table_seed(cfg, "local") + seed = kv_workload.table_seed(cfg, "local", 67) tables = [kv_workload.block_table(cfg, seed, request=r) for r in range(4)] for region in cfg["regions"]: blocks = [t[region["name"]].tolist() for t in tables] @@ -152,21 +131,6 @@ def test_a_request_beyond_the_pool_fails_closed(self): with self.assertRaises(ValueError): kv_workload.block_table(cfg, 1, request=8) - def test_block_major_offsets(self): - # One offset per packed block: block b sits at b * packed_bytes. - cfg = dict(regions=[dict(name="kv", packed_bytes=512, blocks_req=2, - pool_blocks=3, base=0)], descs=2) - offsets = kv_workload.page_offsets(cfg, {"kv": np.array([2, 0])}) - self.assertEqual(offsets.tolist(), [2 * 512, 0]) - - def test_second_region_offsets_start_at_its_base(self): - cfg = dict(regions=[ - dict(name="a", packed_bytes=256, blocks_req=1, pool_blocks=2, base=0), - dict(name="b", packed_bytes=128, blocks_req=1, pool_blocks=2, base=512), - ], descs=2) - offsets = kv_workload.page_offsets(cfg, {"a": np.array([1]), "b": np.array([1])}) - self.assertEqual(offsets.tolist(), [256, 512 + 128]) - def test_desc_array_carries_per_region_packed_sizes(self): cfg = dict(regions=[ dict(name="a", packed_bytes=256, blocks_req=2, pool_blocks=4, base=0), @@ -181,84 +145,89 @@ def test_desc_array_carries_per_region_packed_sizes(self): class Verify(unittest.TestCase): + SRC_SALT, DST_SALT = 1, 0 + + @staticmethod + def _pattern(nbytes, salt): + # The expected byte model, written independently of kv_workload. + chunks = np.arange(nbytes, dtype=np.int64) >> 8 + return ((chunks * 131 + 7 + 101 * salt) & 0xFF).astype(np.uint8) + def _painted_destination(self, cfg, dst_tables, src_tables): """A destination pool where every dst block holds its src block's pattern.""" - pool = np.zeros(cfg["pool_bytes"], dtype=np.uint8) + pool = self._pattern(cfg["pool_bytes"], self.DST_SALT) + src_pool = self._pattern(cfg["pool_bytes"], self.SRC_SALT) for region in cfg["regions"]: size = region["packed_bytes"] for dst, src in zip(dst_tables[region["name"]], src_tables[region["name"]]): dst_off = int(dst) * size + region["base"] src_off = int(src) * size + region["base"] - src_bytes = src_off + np.arange(size, dtype=np.int64) - pool[dst_off : dst_off + size] = ((src_bytes >> 8) * 131 + 7) & 0xFF + pool[dst_off : dst_off + size] = src_pool[src_off : src_off + size] return pool - def _tables(self, cfg): - dst = kv_workload.block_table(cfg, kv_workload.table_seed(cfg, "local")) - src = kv_workload.block_table(cfg, kv_workload.table_seed(cfg, "remote")) - return dst, src + def _setup(self): + cfg = kv_workload.plan_config("dsv4", "fp8", 512, 256) + dst = kv_workload.block_table(cfg, kv_workload.table_seed(cfg, "local", 67)) + src = kv_workload.block_table(cfg, kv_workload.table_seed(cfg, "remote", 67)) + return cfg, dst, src, self._painted_destination(cfg, dst, src) + + def _verify(self, pool, cfg, dst, src, salt=SRC_SALT): + return kv_workload.verify_transfer(_read8(pool), cfg, dst, src, src_salt=salt) def test_a_faithful_transfer_verifies_across_unaligned_pages(self): # dsv4's page sizes are 576 B multiples, never 256 B multiples, so # per-layer probes land at any byte alignment and exercise the # per-byte expectation model. - cfg = kv_workload.plan_config("dsv4", "fp8", 512, 256) - dst, src = self._tables(cfg) - pool = self._painted_destination(cfg, dst, src) - ok, detail = kv_workload.verify_transfer(_read8(pool), cfg, dst, src) + cfg, dst, src, pool = self._setup() + ok, detail = self._verify(pool, cfg, dst, src) self.assertTrue(ok, detail) - def test_one_missing_transfer_fails_with_its_coordinates(self): - cfg = kv_workload.plan_config("dsv4", "fp8", 512, 256) - dst, src = self._tables(cfg) - pool = self._painted_destination(cfg, dst, src) - pool[:] = 0 # a transfer that never happened - ok, detail = kv_workload.verify_transfer(_read8(pool), cfg, dst, src) - self.assertFalse(ok) - self.assertIn("expected", detail) - - def test_direction_matters(self): - # Verifying with the tables swapped must fail: dst blocks hold src - # pattern, not their own. - cfg = kv_workload.plan_config("dsv4", "fp8", 512, 256) - dst, src = self._tables(cfg) - pool = self._painted_destination(cfg, dst, src) - ok, _ = kv_workload.verify_transfer(_read8(pool), cfg, src, dst) - self.assertFalse(ok) - - def test_fabric_pool_pattern_matches_the_verify_model(self): - # kv_pool's host-built pattern (the mnnvl fill path) and the verify - # model must agree byte for byte, or every mnnvl row fails verify. - import kv_pool - - pattern = kv_pool._pattern(1024) - for offset in (0, 8, 256, 512, 1016): - expected = kv_workload._chunk_byte(offset) - self.assertTrue((pattern[offset : offset + 8] == expected).all(), offset) + def test_wrong_transfers_fail(self): + cfg, dst, src, pool = self._setup() + untouched = self._pattern(cfg["pool_bytes"], self.DST_SALT) + for name, args in ( + # a transfer that never happened leaves the destination's own + # repainted pattern, which the source salt tells apart + ("never happened", (untouched, cfg, dst, src)), + # dst blocks hold the src pattern, not their own + ("tables swapped", (pool, cfg, src, dst)), + # a loopback read of the destination's own pool + ("wrong source rank", (pool, cfg, dst, src, self.DST_SALT)), + ): + with self.subTest(name): + ok, detail = self._verify(*args) + self.assertFalse(ok) + self.assertIn("expected", detail) + + def test_the_fabric_pool_tile_matches_the_verify_model(self): + # FabricPool (the mnnvl fill path) doubles pattern_tile across the pool; + # it must agree byte for byte with the verify model, or every mnnvl row + # fails verify. + for salt in (0, 1): + tile = kv_workload.pattern_tile(salt) + self.assertEqual(tile.nbytes, kv_workload.PATTERN_PERIOD) + for offset in (0, 8, 256, 1016, kv_workload.PATTERN_PERIOD - 8): + expected = bytes(kv_workload._chunk_byte(offset + j, salt) for j in range(8)) + self.assertEqual(tile[offset : offset + 8].tobytes(), expected, (salt, offset)) + # periodic: the byte one period on is the same byte + self.assertEqual(kv_workload._chunk_byte(kv_workload.PATTERN_PERIOD + 300, salt), + kv_workload._chunk_byte(300, salt)) class SweepConfigConsistency(unittest.TestCase): - def test_kv_sweep_precisions_match_the_workload_model(self): - # sweep_matrix schedules from the JSON map (it must stay stdlib-only); - # the workload model owns the truth and plan_config fail-closes on a - # mismatch at runtime. This pins the two together at PR time. - import json - - sweep = json.loads((ROOT / "configs" / "kv_sweep.json").read_text()) - for workload, precisions in sweep["workloads"].items(): - preset = kv_workload.PRESETS[workload.removeprefix("kv-")] - self.assertEqual(tuple(precisions), preset["precisions"], workload) - - def test_kv_sweep_block_sizes_are_plannable(self): - # A sweep block size the model rejects (splitting an HCA state) would - # kill every kv leg at the first grid point. - import json - - sweep = json.loads((ROOT / "configs" / "kv_sweep.json").read_text()) - for workload, precisions in sweep["workloads"].items(): - for block in sweep["page_tokens"]: - kv_workload.plan_config(workload.removeprefix("kv-"), - precisions[0], 512, block) + def test_every_scheduled_workload_point_is_plannable(self): + # sweep_matrix stays stdlib-only, so it schedules from kv_sweep.json's + # workload -> precision map without importing this model; a precision + # or block size plan_config rejects would kill every kv leg at its + # first grid point. + import sweep_matrix + + for workload, precisions in sweep_matrix.KV_SWEEP["workloads"].items(): + for precision in precisions: + for block in sweep_matrix.KV_SWEEP["page_tokens"]: + with self.subTest(workload=workload, precision=precision, block=block): + kv_workload.plan_config(workload.removeprefix("kv-"), precision, + 512, block) class Percentiles(unittest.TestCase): diff --git a/collectivex/tests/test_runtime.py b/collectivex/tests/test_runtime.py index ba80ecbb22..7d0a29dc71 100644 --- a/collectivex/tests/test_runtime.py +++ b/collectivex/tests/test_runtime.py @@ -447,14 +447,14 @@ class CaseArgvContract(unittest.TestCase): "suite": "ep-core", "workload": "deepseek-v3", } - def _run_ep_parser(self) -> argparse.ArgumentParser: - import run_ep + def _entrypoint_parser(self, entrypoint: str = "run_ep") -> argparse.ArgumentParser: + module = __import__(entrypoint) # Capture the real entrypoint's parser before it initializes any GPU runtime. with mock.patch.object(argparse.ArgumentParser, "parse_args", autospec=True, side_effect=SystemExit) as parse: with self.assertRaises(SystemExit): - run_ep.main() + module.main() return parse.call_args.args[0] def _decode(self, stdout: bytes, entrypoint: str = "run_ep") -> list: @@ -476,9 +476,9 @@ def _case_argv(self, placement: list, case: dict | None = None, ) return self._decode(result.stdout, entrypoint) - def test_case_args_round_trips_through_the_run_ep_parser(self) -> None: + def test_case_args_round_trips_through_the_entrypoint_parser(self) -> None: argv = self._case_argv(["16", "2", "8", "8"]) - args = self._run_ep_parser().parse_args(argv) + args = self._entrypoint_parser().parse_args(argv) self.assertEqual( (args.backend, args.mode, args.phase, args.routing, args.scope), ("deepep-v2", "normal", "decode", "uniform", "scale-out"), @@ -501,7 +501,7 @@ def test_a_legacy_colon_string_profile_still_decodes(self) -> None: # run_ep's own defaults rather than being duplicated in the codec. for profile, chain in (("8:256:32:128:4:16", (128, 4, 16)), ("8:256:32", None)): with self.subTest(timing=profile): - args = self._run_ep_parser().parse_args(self._case_argv( + args = self._entrypoint_parser().parse_args(self._case_argv( ["16", "2", "8", "8"], case={**self.CASE, "timing": profile}, )) self.assertEqual((args.iters, args.trials, args.warmup), (8, 256, 32)) @@ -534,13 +534,13 @@ def test_a_malformed_timing_profile_cannot_reach_a_run(self) -> None: ["16", "2", "8", "8"], case={**self.CASE, "timing": "8:256:32:128:4:x"}, ) with contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit): - self._run_ep_parser().parse_args(argv) + self._entrypoint_parser().parse_args(argv) def test_case_args_fails_closed_on_placement_mismatch(self) -> None: with self.assertRaises(subprocess.CalledProcessError): self._case_argv(["8", "1", "8", "8"]) - def test_each_backend_round_trips_through_the_run_ep_parser(self) -> None: + def test_each_backend_round_trips_through_the_entrypoint_parser(self) -> None: # The codec is backend-agnostic, so one loop replaces three near-identical tests: # run_ep's --backend choices must accept each name, and the filename must carry the # backend token or two legs of one cell collide in results/. @@ -550,7 +550,7 @@ def test_each_backend_round_trips_through_the_run_ep_parser(self) -> None: **self.CASE, "backend": backend, "case_id": f"h200-dgxc-{backend}-deepseek-v3-normal-decode-ep16-uniform-bf16", } - args = self._run_ep_parser().parse_args( + args = self._entrypoint_parser().parse_args( self._case_argv(["16", "2", "8", "8"], case=case) ) self.assertEqual(args.backend, backend) @@ -579,6 +579,50 @@ def test_a_swap_blocks_case_round_trips_through_its_own_parser(self) -> None: with self.assertRaises(subprocess.CalledProcessError): self._case_argv(["8", "1", "8", "8"], case=case, entrypoint="run_swap_blocks") + def _kv_case(self) -> dict: + import sweep_matrix + + return next(shard for shard in sweep_matrix.resolve_matrix( + suites="kv-transfer", only_sku="h200-dgxc")["include"] + if shard["backend"] == "nixl")["cases"][0] + + def test_a_kv_case_round_trips_through_the_run_kv_parser(self) -> None: + import run_kv + + case = self._kv_case() + parser = self._entrypoint_parser("run_kv") + for extra, budget in (({}, run_kv.POOL_BUDGET), ({"pool_budget": 123}, 123)): + with self.subTest(extra=extra): + argv = self._case_argv(["2", "2", "1", "8"], case={**case, **extra}, + entrypoint="run_kv") + args = parser.parse_args(argv) + self.assertEqual( + (args.backend, args.workload_name, args.precision, args.fabric), + (case["backend"], case["workload"], case["precision"], case["mode"])) + self.assertEqual((args.warmup, args.reps, args.trials, args.seed), + (case["warmup"], case["reps"], case["trials"], case["seed"])) + self.assertEqual((args.batch_sizes, args.kv_device, args.ops), + (case["batch_sizes"], case["kv_device"], case["ops"])) + self.assertEqual(args.pool_budget, budget) + self.assertEqual((args.case_id, args.runner), (case["case_id"], "h200-dgxc")) + self.assertEqual(args.out, f"results/{case['case_id']}_TS-c000.json") + # run_kv recomputes the identity from the runner and its own factors + self.assertEqual(ep_harness.case_id(args.runner, run_kv.kv_case(args)), + case["case_id"]) + # a kv case is 2 ranks on 2 nodes x 1 GPU; any other allocation is refused + with self.assertRaises(subprocess.CalledProcessError): + self._case_argv(["8", "1", "8", "8"], case=case, entrypoint="run_kv") + + def test_run_kv_refuses_a_case_id_its_factors_do_not_produce(self) -> None: + import run_kv + + argv = self._case_argv(["2", "2", "1", "8"], case=self._kv_case(), entrypoint="run_kv") + argv[argv.index("--case-id") + 1] = "h200-dgxc-somebody-else" + with mock.patch.object(sys, "argv", ["run_kv.py", *argv]), \ + contextlib.redirect_stderr(io.StringIO()) as err: + self.assertEqual(run_kv.main(), 2) + self.assertIn("does not match factors", err.getvalue()) + def test_a_case_from_an_unknown_suite_cannot_reach_a_run(self) -> None: with self.assertRaises(subprocess.CalledProcessError): self._case_argv(["16", "2", "8", "8"], case={**self.CASE, "suite": "turbo"}) From 8a015008fd48ba003148a41a271105c63e7b600a Mon Sep 17 00:00:00 2001 From: Oseltamivir <58582368+Oseltamivir@users.noreply.github.com> Date: Tue, 29 Sep 2026 06:57:38 +0800 Subject: [PATCH 4/4] CollectiveX: run the kv-transfer nixl leg on b300's EFA pool over LIBFABRIC b300 is now the AWS p6-b300 EFA pool. EFA is not a verbs HCA and UCX has no transport for it, so the NIXL adapter selects the wheel's LIBFABRIC plugin when the network profile marks the pool rdma_fabric=efa, and rows record the plugin in implementation.transport. A two-node hand probe on the pool moved 94 GB/s READ and 97 GB/s WRITE per GPU at 1 GiB, verified. No mooncake leg: the PyPI wheel's transport is verbs RC only. --- collectivex/README.md | 4 +++- collectivex/bench/kv_backend.py | 3 +++ collectivex/bench/kv_nixl.py | 15 +++++++++++---- collectivex/bench/run_kv.py | 1 + collectivex/configs/platform_config.json | 1 + collectivex/docs/methodology.md | 6 ++++++ 6 files changed, 25 insertions(+), 5 deletions(-) diff --git a/collectivex/README.md b/collectivex/README.md index f0bf460ba8..39cd7b2249 100644 --- a/collectivex/README.md +++ b/collectivex/README.md @@ -183,7 +183,9 @@ at import; mi355x runs AMD's atom-dev build push-only), and `mori-io` (AMD's nat the registry's `kv_backends` map enables them. No entry, no legs, mirroring `ll_backends`. An entry may restrict ops, pin an image, set a NIC filter, or lower the pool budget. Fabrics are `rdma` (torch pools) and, on GB racks, `nixl` also runs `mnnvl` (cuMem FABRIC pools; see the methodology -for the bulk-vs-paged lane inversion that row exists to publish). The grid and per-pool scheduling +for the bulk-vs-paged lane inversion that row exists to publish). On b300's AWS EFA pool `nixl` +rides its LIBFABRIC plugin instead of UCX, and there is no `mooncake` leg (its wheel's transport is +verbs RC, which EFA does not provide). The grid and per-pool scheduling live in `configs/kv_sweep.json`; dispatch with `suites: kv-transfer`. ## Workflow And Artifacts diff --git a/collectivex/bench/kv_backend.py b/collectivex/bench/kv_backend.py index b4d07ee1bd..61aef4148d 100644 --- a/collectivex/bench/kv_backend.py +++ b/collectivex/bench/kv_backend.py @@ -28,6 +28,9 @@ class KVBackend: library_version: str | None = None #: the engine NIC filter this case ran under; None = library/UCX choice. nic_filter: str | None = None + #: the library's own transport plugin when it has several (NIXL: UCX or + #: LIBFABRIC); None = the library has one. + transport: str | None = None # -- lifecycle ------------------------------------------------------------ def register(self, pool, bulk, reg_layout=None) -> None: diff --git a/collectivex/bench/kv_nixl.py b/collectivex/bench/kv_nixl.py index ceb23afcfc..07148ca987 100644 --- a/collectivex/bench/kv_nixl.py +++ b/collectivex/bench/kv_nixl.py @@ -1,7 +1,10 @@ #!/usr/bin/env python3 -"""NIXL (UCX) adapter: the library Dynamo, vLLM NixlConnector, and SGLang -disagg ship. Agent metadata rides the harness exchange (`add_remote_agent`), -not NIXL's TCP listener, so the adapter needs no port and no listener race. +"""NIXL adapter: the library Dynamo, vLLM NixlConnector, and SGLang disagg +ship. UCX carries it on verbs fabrics (IB, RoCE); on AWS EFA, which is not a +verbs HCA and has no UCX transport, the wheel's LIBFABRIC plugin carries it +over the host libfabric the cluster's enroot hook mounts. Agent metadata rides +the harness exchange (`add_remote_agent`), not NIXL's TCP listener, so the +adapter needs no port and no listener race. Remote descriptors are built locally from the peer's published pool base; both block tables are seed-keyed, the same information a decode worker gets from the prefill side's block table message. @@ -9,6 +12,7 @@ from __future__ import annotations +import os import time import numpy as np @@ -57,8 +61,11 @@ def __init__(self, args, role, device): # The registry pin run_kv hands to UCX_NET_DEVICES for this case; # None means UCX chose among the operator inventory itself. self.nic_filter = args.kv_device or None + # The network profile marks EFA pools; there FI_PROVIDER=efa is already set. + self.transport = "LIBFABRIC" if os.environ.get("COLLX_RDMA_FABRIC") == "efa" else "UCX" # prog thread on, listener off: metadata goes through the harness exchange. - self._agent = nixl_agent(role, nixl_agent_config(True, False, 0, backends=["UCX"])) + self._agent = nixl_agent(role, nixl_agent_config(True, False, 0, + backends=[self.transport])) self._handles = [] self._pool = None self._bulk = None diff --git a/collectivex/bench/run_kv.py b/collectivex/bench/run_kv.py index 78ea0d30d8..7ee2071bff 100644 --- a/collectivex/bench/run_kv.py +++ b/collectivex/bench/run_kv.py @@ -473,6 +473,7 @@ def measure(make, cfg, cfg_row: dict, op: str, verify_side: str, table_pairs=Non "library_version": backend.library_version, "maturity": backend.maturity, "nic_filter": backend.nic_filter, + "transport": backend.transport, }, topology={ "device_product": torch.cuda.get_device_name(device), diff --git a/collectivex/configs/platform_config.json b/collectivex/configs/platform_config.json index a0fe91444e..17890b6374 100644 --- a/collectivex/configs/platform_config.json +++ b/collectivex/configs/platform_config.json @@ -80,6 +80,7 @@ "launcher": "single-slurm", "backends": {"deepep-v2": [8], "nccl-ep": [8]}, "ll_backends": {"nccl-ep": [8]}, + "kv_backends": {"nixl": ["rdma"]}, "fabric": {"nic": "AWS EFA (p6-b300.48xlarge, 16 x 400G per node)", "switch": "AWS EFA fabric (SRD)"}, "operator": { "partition": "batch_1", diff --git a/collectivex/docs/methodology.md b/collectivex/docs/methodology.md index e873c7d811..874cfb5f79 100644 --- a/collectivex/docs/methodology.md +++ b/collectivex/docs/methodology.md @@ -623,6 +623,12 @@ are not restated here. Other lane facts, measured on the metal: single-WR bulk transfers above the provider's max message size must be split (the MoRI adapter caps WRs at 1 GiB); and the upstream PyPI Mooncake wheel is NVIDIA-only (it links libcuda.so.1 at import), so mi355x runs AMD's atom-dev build. +AWS EFA (b300) is not a verbs HCA and UCX has no transport for it, so on a pool whose network +profile says `rdma_fabric: efa` the NIXL adapter selects the wheel's LIBFABRIC plugin, which +loads the host libfabric the cluster's enroot hook mounts and rails a GPU across its local EFA +devices; the row's `implementation.transport` records which plugin carried it. The same PyPI +Mooncake wheel links only libibverbs/libmlx5 (verbs RC, which EFA does not offer), so b300 carries +no mooncake leg. Scheduling is data: `configs/kv_sweep.json` holds the grid and, per pool, the allocation and the per-case hang guard. GB pools get the long budgets, because the mnnvl descriptor floor makes a leg