diff --git a/.github/workflows/collectivex-sweep.yml b/.github/workflows/collectivex-sweep.yml index 8e9895a667..e39d27fcdc 100644 --- a/.github/workflows/collectivex-sweep.yml +++ b/.github/workflows/collectivex-sweep.yml @@ -32,8 +32,16 @@ on: type: choice default: '' options: ['', dequant] + suites: + description: "Comma-list of suites to resolve (ep-core, kv-transfer); blank = all suites" + type: string + default: '' + skip_queue_pr: + description: "PR number carrying an authorized skip_queue label; shards then request the ci-skip-queue-pr- label so the priority controller admits them ahead of the queue. Blank = normal queueing" + type: string + default: '' concurrency: - group: cx-${{ github.ref }}-${{ inputs.backend }}-${{ inputs.only_sku }} + group: cx-${{ github.ref }}-${{ inputs.backend }}-${{ inputs.only_sku }}-${{ inputs.suites }} cancel-in-progress: false jobs: @@ -61,6 +69,7 @@ jobs: INPUT_EXCLUDE_SKUS: ${{ inputs.exclude_skus }} INPUT_EP_SIZES: ${{ inputs.ep_sizes }} INPUT_MODES: ${{ inputs.modes }} + INPUT_SUITES: ${{ inputs.suites }} RUN_ID: ${{ github.run_id }} RUN_ATTEMPT: ${{ github.run_attempt }} run: | @@ -70,6 +79,7 @@ jobs: [ -n "$INPUT_EXCLUDE_SKUS" ] && args+=(--exclude-skus "$INPUT_EXCLUDE_SKUS") [ -n "$INPUT_EP_SIZES" ] && args+=(--ep-sizes "$INPUT_EP_SIZES") [ -n "$INPUT_MODES" ] && args+=(--modes "$INPUT_MODES") + [ -n "$INPUT_SUITES" ] && args+=(--suites "$INPUT_SUITES") python3 sweep_matrix.py "${args[@]}" --out matrix_full.json >/dev/null python3 - "$GITHUB_OUTPUT" <<'PY' import hashlib @@ -114,22 +124,61 @@ jobs: runs-on: >- ${{ fromJSON( vars.PRIORITY_SCHEDULER_ENABLED == 'true' && - format( - '["self-hosted",{0},{1},{2}]', - toJSON(matrix.sku), - toJSON(format( - 'ci-job-{0}-{1}', - needs.setup.outputs.priority, - matrix.queue-token - )), - toJSON(format('ci-attempt-{0}', github.run_attempt)) + ( + vars.NODE_SLOT_SCHEDULER_ENABLED == 'true' && + ( + inputs.skip_queue_pr != '' && + format( + '["self-hosted",{0},{1},{2},{3},{4}]', + toJSON(matrix.runner), + toJSON(format('nodes:{0}', matrix.nodes)), + toJSON(format( + 'ci-job-{0}-{1}', + needs.setup.outputs.priority, + matrix.queue-token + )), + toJSON(format('ci-attempt-{0}', github.run_attempt)), + toJSON(format('ci-skip-queue-pr-{0}', inputs.skip_queue_pr)) + ) || + format( + '["self-hosted",{0},{1},{2},{3}]', + toJSON(matrix.runner), + toJSON(format('nodes:{0}', matrix.nodes)), + toJSON(format( + 'ci-job-{0}-{1}', + needs.setup.outputs.priority, + matrix.queue-token + )), + toJSON(format('ci-attempt-{0}', github.run_attempt)) + ) + ) || + format( + '["self-hosted",{0},{1},{2}]', + toJSON(matrix.runner), + toJSON(format( + 'ci-job-{0}-{1}', + needs.setup.outputs.priority, + matrix.queue-token + )), + toJSON(format('ci-attempt-{0}', github.run_attempt)) + ) ) || - format('[{0}]', toJSON(matrix.sku)) + format('[{0}]', toJSON(matrix.runner)) ) }} - name: p${{ needs.setup.outputs.priority }} | ${{ matrix.sku }} ${{ matrix.backend }} shard ${{ matrix.id }} - timeout-minutes: 350 + name: p${{ needs.setup.outputs.priority }} | ${{ matrix.runner }} ${{ matrix.backend }} shard ${{ matrix.id }} + # Must sit above every launcher's largest Slurm allocation plus setup, + # or GitHub cancels a healthy shard before the launcher's own guards can + # act: the kv gb200 mnnvl leg holds a 460 minute allocation with a 420 + # minute per-case guard inside it (run 33150394862 was cancelled here at + # 350 minutes while doing honest work), and the kv gb300 legs hold 690 + # minutes with a 660 minute guard because gb300 paces ~1.8x gb200 at the + # top isls. EP shards finish far earlier and are unaffected by the + # ceiling. + timeout-minutes: 720 env: COLLX_BENCH: ${{ matrix.backend }} + COLLX_MODE: ${{ matrix.mode }} + COLLX_IMAGE_OVERRIDE: ${{ matrix.image_ref || '' }} COLLX_NODES: ${{ matrix.nodes }} COLLX_GPUS_PER_NODE: ${{ matrix.gpus_per_node }} COLLX_SCALE_UP_DOMAIN: ${{ matrix.scale_up_domain }} diff --git a/.github/workflows/test-collectivex.yml b/.github/workflows/test-collectivex.yml index e9b7fbd6c7..baac807ca5 100644 --- a/.github/workflows/test-collectivex.yml +++ b/.github/workflows/test-collectivex.yml @@ -41,6 +41,9 @@ jobs: # ep_flashinfer's combine-model switch compares wheel versions with # packaging.version; the benchmark image ships it, this runner does not. pip install packaging + # The kv-transfer workload/geometry tests are numpy math (the cpu + # torch wheel does not pull numpy in). + pip install numpy # Every skip in this suite is torch-gated, so a missing torch turns the oracle # checks into silent passes. Fail here instead, where the cause is obvious. diff --git a/experimental/CollectiveX/README.md b/experimental/CollectiveX/README.md index a6b8a60db5..5177bf7a17 100644 --- a/experimental/CollectiveX/README.md +++ b/experimental/CollectiveX/README.md @@ -1,8 +1,9 @@ # CollectiveX -CollectiveX is an experimental MoE expert-parallel communication benchmark. It measures dispatch, -combine, and paired roundtrip latency across EP libraries and accelerator systems, then uploads -neutral result artifacts. +CollectiveX is an experimental inference-communication benchmark. Its EP suite measures MoE +dispatch, combine, and paired roundtrip latency across EP libraries and accelerator systems; its +KV-transfer suite measures disaggregated-serving KV-cache handoffs across transfer libraries and +fabrics. Both upload neutral result artifacts. CollectiveX schedules benchmarks, executes them on real allocations, and uploads the neutral artifacts each run emits. It does not validate those artifacts, promote, rank, recommend, select, or @@ -149,6 +150,35 @@ 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 — and move bursts of 1 to 32 concurrent +requests' paged KV as per-request layer-major descriptor lists over seed-keyed +random block tables (the post-fragmentation layout vLLM and SGLang post; a +burst posts every request's prepped transfer, then awaits them all), plus one +contiguous bulk row as the wire-speed ceiling. The workload is transcribed from +what vLLM allocates for the model it serves: `kv-dsv4` = 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); ISL +8k to 512k at page sizes 16 and 64 tokens; `pull` (READ, vLLM NixlConnector) and `push` (WRITE, +SGLang disagg) both timed from the initiator with offset-pattern verification on +the destination pool in both directions. Backends: `nixl` (what Dynamo, vLLM, +and SGLang ship), `mooncake` (NVIDIA-only; the wheel links libcuda at import), +and `mori-io` (AMD's native engine) where the registry's `kv_backends` map +enables them; no entry, no legs, mirroring `ll_backends`. A backend entry may +restrict ops, pin an image, or set a NIC filter (mooncake on mi355x is +push-only from AMD's atom-dev image over the GPU-paired Pollara NIC; upstream +ionic RDMA READ is broken; both b300 backends are pinned to one rail — +mooncake because the image's engine draws peer NICs blindly and cross-rail +draws stall ~1 s, nixl because UCX's own two-rail READ selection is unstable +across runs while one rail holds 49 GB/s at p95/p50 ≤ 1.01 — so b300 kv rows +are one-rail measurements). Fabrics: `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). + ## Workflow And Artifacts `.github/workflows/collectivex-sweep.yml` has two jobs. `setup` generates a public-SKU matrix diff --git a/experimental/CollectiveX/bandwidth.py b/experimental/CollectiveX/bandwidth.py index c0b13e54e0..8563b0eea1 100644 --- a/experimental/CollectiveX/bandwidth.py +++ b/experimental/CollectiveX/bandwidth.py @@ -178,6 +178,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-suite 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/experimental/CollectiveX/bench/kv_backend.py b/experimental/CollectiveX/bench/kv_backend.py new file mode 100644 index 0000000000..fec12acdc1 --- /dev/null +++ b/experimental/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/experimental/CollectiveX/bench/kv_mooncake.py b/experimental/CollectiveX/bench/kv_mooncake.py new file mode 100644 index 0000000000..da375e5d6e --- /dev/null +++ b/experimental/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/experimental/CollectiveX/bench/kv_mori_io.py b/experimental/CollectiveX/bench/kv_mori_io.py new file mode 100644 index 0000000000..91e5e98c45 --- /dev/null +++ b/experimental/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/experimental/CollectiveX/bench/kv_nixl.py b/experimental/CollectiveX/bench/kv_nixl.py new file mode 100644 index 0000000000..5bf36699d3 --- /dev/null +++ b/experimental/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/experimental/CollectiveX/bench/kv_pool.py b/experimental/CollectiveX/bench/kv_pool.py new file mode 100644 index 0000000000..403102e7d7 --- /dev/null +++ b/experimental/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/experimental/CollectiveX/bench/kv_workload.py b/experimental/CollectiveX/bench/kv_workload.py new file mode 100644 index 0000000000..83b089ee59 --- /dev/null +++ b/experimental/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/experimental/CollectiveX/bench/run_kv.py b/experimental/CollectiveX/bench/run_kv.py new file mode 100644 index 0000000000..f5f0a1107a --- /dev/null +++ b/experimental/CollectiveX/bench/run_kv.py @@ -0,0 +1,515 @@ +#!/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 +# 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 launcher may lower it (COLLX_KV_POOL_BUDGET, bytes) 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-launcher, not fleet-wide. +POOL_BUDGET = int(os.environ.get("COLLX_KV_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("--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 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"] <= 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", "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/experimental/CollectiveX/configs/kv_sweep.json b/experimental/CollectiveX/configs/kv_sweep.json new file mode 100644 index 0000000000..c2f1b0d11d --- /dev/null +++ b/experimental/CollectiveX/configs/kv_sweep.json @@ -0,0 +1,39 @@ +{ + "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 +} diff --git a/experimental/CollectiveX/configs/platform_config.json b/experimental/CollectiveX/configs/platform_config.json index 506538f4b0..ec06d54b39 100644 --- a/experimental/CollectiveX/configs/platform_config.json +++ b/experimental/CollectiveX/configs/platform_config.json @@ -1,5 +1,11 @@ { - "backend_maturity": {"deepep-v2": "production", "mori": "production", "uccl-ep": "candidate", "nccl-ep": "candidate", "flashinfer-ep": "production"}, + "backend_maturity": { + "deepep-v2": "production", + "mori": "production", + "uccl-ep": "candidate", + "nccl-ep": "candidate", + "flashinfer-ep": "production" + }, "platforms": { "h100-dgxc": { "arch": "sm90", @@ -10,9 +16,33 @@ "scale_up_domain": 8, "scale_up_transport": "nvlink", "launcher": "single-slurm", - "backends": {"deepep-v2": [8, 16], "uccl-ep": [8], "nccl-ep": [8]}, - "ll_backends": {"deepep-v2": [8], "uccl-ep": [8], "nccl-ep": [8]}, - "fabric": {"nic": "ConnectX-7 2x200GbE", "switch": "Arista 7060DX5-64S (Tomahawk4, 25.6T)"}, + "backends": { + "deepep-v2": [ + 8, + 16 + ], + "uccl-ep": [ + 8 + ], + "nccl-ep": [ + 8 + ] + }, + "ll_backends": { + "deepep-v2": [ + 8 + ], + "uccl-ep": [ + 8 + ], + "nccl-ep": [ + 8 + ] + }, + "fabric": { + "nic": "ConnectX-7 2x200GbE", + "switch": "Arista 7060DX5-64S (Tomahawk4, 25.6T)" + }, "operator": { "partition": "hpc-gpu-1", "account": "customer", @@ -34,9 +64,41 @@ "scale_up_domain": 8, "scale_up_transport": "nvlink", "launcher": "single-slurm", - "backends": {"deepep-v2": [8, 16], "uccl-ep": [8], "nccl-ep": [8]}, - "ll_backends": {"deepep-v2": [8], "uccl-ep": [8], "nccl-ep": [8]}, - "fabric": {"nic": "ConnectX-7 400G", "switch": "NVIDIA Quantum-2 QM9790 (25.6T, InfiniBand)"}, + "backends": { + "deepep-v2": [ + 8, + 16 + ], + "uccl-ep": [ + 8 + ], + "nccl-ep": [ + 8 + ] + }, + "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", "squash_dir": "/home/sa-shared/containers" @@ -54,9 +116,42 @@ "scale_up_domain": 8, "scale_up_transport": "nvlink", "launcher": "single-slurm", - "backends": {"deepep-v2": [8, 16], "uccl-ep": [8], "nccl-ep": [8]}, - "ll_backends": {"deepep-v2": [8, 16], "uccl-ep": [8], "nccl-ep": [8]}, - "fabric": {"nic": "8x ConnectX-7-class 400G InfiniBand rails (bare metal, gdrdrv)", "switch": "InfiniBand (model not surveyed)"}, + "backends": { + "deepep-v2": [ + 8, + 16 + ], + "uccl-ep": [ + 8 + ], + "nccl-ep": [ + 8 + ] + }, + "ll_backends": { + "deepep-v2": [ + 8, + 16 + ], + "uccl-ep": [ + 8 + ], + "nccl-ep": [ + 8 + ] + }, + "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", "account": "benchmark", @@ -76,15 +171,44 @@ "scale_up_domain": 8, "scale_up_transport": "nvlink", "launcher": "single-slurm", - "backends": {"deepep-v2": [8, 16], "nccl-ep": [8]}, - "ll_backends": {"nccl-ep": [8]}, - "fabric": {"nic": "ConnectX-8 2x400GbE", "switch": "NVIDIA Spectrum-X SN5600 (51.2T)"}, + "backends": { + "deepep-v2": [ + 8, + 16 + ], + "nccl-ep": [ + 8 + ] + }, + "ll_backends": { + "nccl-ep": [ + 8 + ] + }, + "kv_backends": { + "nixl": { + "fabrics": [ + "rdma" + ], + "device": "mlx5_0" + }, + "mooncake": { + "fabrics": [ + "rdma" + ], + "device": "mlx5_0" + } + }, + "fabric": { + "nic": "ConnectX-8 2x400GbE", + "switch": "NVIDIA Spectrum-X SN5600 (51.2T)" + }, "operator": { "partition": "batch_1", "account": "benchmark", "qos": "batch_1_qos", "squash_dir": "/data/home/sa-shared/sqsh", - "exclude_nodes": "b300-018" + "exclude_nodes": "b300-005,b300-009,b300-018" }, "network": { "socket_ifname": "bond0", @@ -102,15 +226,52 @@ "scale_up_domain": 72, "scale_up_transport": "mnnvl", "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]}, - "fabric": {"nic": "MNNVL (scale-out not used)", "switch": "NVLink NVL72"}, + "backends": { + "deepep-v2": [ + 8, + 16 + ], + "nccl-ep": [ + 8, + 16 + ], + "flashinfer-ep": [ + 8, + 16 + ] + }, + "ll_backends": { + "deepep-v2": [ + 8, + 16 + ], + "nccl-ep": [ + 8 + ] + }, + "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": { @@ -122,9 +283,42 @@ "scale_up_domain": 72, "scale_up_transport": "mnnvl", "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]}, - "fabric": {"nic": "MNNVL (scale-out not used)", "switch": "NVLink NVL72"}, + "backends": { + "deepep-v2": [ + 8, + 16 + ], + "nccl-ep": [ + 8, + 16 + ], + "flashinfer-ep": [ + 8, + 16 + ] + }, + "ll_backends": { + "deepep-v2": [ + 8, + 16 + ], + "nccl-ep": [ + 8 + ] + }, + "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", @@ -132,6 +326,10 @@ "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" } }, "mi325x-tw": { @@ -143,9 +341,23 @@ "scale_up_domain": 8, "scale_up_transport": "xgmi", "launcher": "mi-tw", - "backends": {"mori": [8], "uccl-ep": [8]}, - "ll_backends": {"mori": [8]}, - "fabric": {"nic": "n/a (single-node scale-up)", "switch": "n/a (single-node scale-up)"} + "backends": { + "mori": [ + 8 + ], + "uccl-ep": [ + 8 + ] + }, + "ll_backends": { + "mori": [ + 8 + ] + }, + "fabric": { + "nic": "n/a (single-node scale-up)", + "switch": "n/a (single-node scale-up)" + } }, "mi300x-tw": { "arch": "gfx942", @@ -156,9 +368,23 @@ "scale_up_domain": 8, "scale_up_transport": "xgmi", "launcher": "mi-tw", - "backends": {"mori": [8], "uccl-ep": [8]}, - "ll_backends": {"mori": [8]}, - "fabric": {"nic": "n/a (single-node scale-up)", "switch": "n/a (single-node scale-up)"} + "backends": { + "mori": [ + 8 + ], + "uccl-ep": [ + 8 + ] + }, + "ll_backends": { + "mori": [ + 8 + ] + }, + "fabric": { + "nic": "n/a (single-node scale-up)", + "switch": "n/a (single-node scale-up)" + } }, "mi355x": { "arch": "gfx950", @@ -169,9 +395,37 @@ "scale_up_domain": 8, "scale_up_transport": "xgmi", "launcher": "mi-amds", - "backends": {"mori": [8], "uccl-ep": [8]}, - "ll_backends": {"mori": [8]}, - "fabric": {"nic": "Pollara 400GbE", "switch": "Arista 7060X6-64PE (Tomahawk5, 51.2T)"}, + "runner": "cluster:mi355x-amds", + "backends": { + "mori": [ + 8 + ], + "uccl-ep": [ + 8 + ] + }, + "ll_backends": { + "mori": [ + 8 + ] + }, + "kv_backends": { + "mori-io": [ + "rdma" + ], + "mooncake": { + "fabrics": [ + "rdma" + ], + "ops": "push", + "image_ref": "rocm/atom-dev:nightly_202606101403", + "device": "rdma{gpu}" + } + }, + "fabric": { + "nic": "Pollara 400GbE", + "switch": "Arista 7060X6-64PE (Tomahawk5, 51.2T)" + }, "network": { "socket_ifname": "eno0", "rdma_devices": "rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7", diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index 50acdddd95..3856621587 100644 --- a/experimental/CollectiveX/docs/methodology.md +++ b/experimental/CollectiveX/docs/methodology.md @@ -451,6 +451,107 @@ 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_ref` 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). On b300, mooncake pins `mlx5_0`: the pods +expose two rails whose cross pairs do not route, and the image-provided pre-0.3.12 engine +draws the peer NIC per request with no view of endpoint health, so every cross-rail draw +pays the worker pool's ~1 s inactive-endpoint hold before the per-slice retry walks to the +routable rail (the hold at `worker_pool.cpp` is the stall quantum — the handshake socket +timeout is 60 s and never fires). b300 nixl carries the same pin: for a UCX-backed engine +`device` lands in `UCX_NET_DEVICES` (overriding the blanket 16-device value b300 forwards +from `/etc/environment`, which otherwise swallows every selector the harness sets), because +UCX's own two-rail selection served READs at anywhere from 18 to 83 GB/s across repeated +runs and node pairs while WRITEs held line rate; on one rail both directions sit at +47.7–49.0 GB/s with p95/p50 ≤ 1.01. Published b300 kv rows are therefore one-rail +measurements — ~49 GB/s nixl, ~48 GB/s mooncake ceilings, directly comparable — and the +row's `implementation.nic_filter` records the pin. 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 An implementation-independent oracle uses an expert-specific deterministic transform so wrong expert diff --git a/experimental/CollectiveX/launchers/launch_gb-nv.sh b/experimental/CollectiveX/launchers/launch_gb-nv.sh index 855c16fd60..f6cb092110 100644 --- a/experimental/CollectiveX/launchers/launch_gb-nv.sh +++ b/experimental/CollectiveX/launchers/launch_gb-nv.sh @@ -30,20 +30,58 @@ SCALE_UP_DOMAIN="${COLLX_SCALE_UP_DOMAIN:-72}" NGPUS="${COLLX_NGPUS:-$((NODES * GPN))}" if [ "$PRODUCT" = gb200 ]; then default_time=30; else default_time=90; fi TIME_MIN="${COLLX_TIME:-$default_time}" +case "$COLLX_BENCH" in + nixl | mooncake) + # The five-rung-floor grid measured ~285 minutes on the mnnvl descriptor + # floor (run 33097162900), and the power-of-two batch ladder in + # kv_sweep.json is another ~1.33x of descriptor work grid-wide, so ~380 + # minutes. 460 clears the raised guard below with setup margin; the ask + # stays 2 nodes x 1 GPU, so it still backfills on a contended pool. + # gb300 paces ~1.8x gb200 at isl >= 131072 over mnnvl (run 33244478580 + # and the 2026-08-29 hand retest agree across two node pairs, so it is + # the platform, not a sick pair), projecting ~600 minutes grid-wide. + if [ "$PRODUCT" = gb300 ]; then TIME_MIN=690; else TIME_MIN=460; fi + ;; +esac IMAGE="$COLLX_IMAGE" TS="$(date -u +%Y-%m-%dT%H-%M-%SZ)" -export COLLX_TRANSPORT=mnnvl +# EP on a GB rack always stays inside the NVL72 domain, but a kv-transfer +# shard names its fabric: an rdma leg is real cross-node InfiniBand and must +# get the same fail-closed network profile + validation every other scale-out +# fabric gets, so its transport label must not read mnnvl. +case "$COLLX_BENCH:${COLLX_MODE:-}" in + nixl:rdma | mooncake:rdma) export COLLX_TRANSPORT=mnnvl-rdma ;; + *) export COLLX_TRANSPORT=mnnvl ;; +esac export COLLX_NODES="$NODES" COLLX_GPUS_PER_NODE="$GPN" COLLX_SCALE_UP_DOMAIN="$SCALE_UP_DOMAIN" export COLLX_NGPUS="$NGPUS" case "$COLLX_BENCH" in deepep-v2 | nccl-ep | flashinfer-ep) ;; - *) collx_die "unsupported $PRODUCT EP backend: $COLLX_BENCH" ;; + nixl | mooncake) + # The five-rung-floor grid measured ~285 minutes end to end on the + # mnnvl descriptor floor (run 33097162900), nearly all of it timed + # bursts, and the power-of-two batch ladder in kv_sweep.json is another + # ~1.33x of descriptor work grid-wide, so ~380 minutes projected. The + # guard must clear that with real margin yet still fire before the + # allocation above dies, so the failure stays a clean per-case kill + # instead of a lost allocation. gb300's ~1.8x pacing at the top isls + # projects ~600 minutes, so its guard sits at 660 inside the 690 + # allocation. + if [ "$PRODUCT" = gb300 ]; then default_guard=39600; else default_guard=25200; fi + export COLLX_RUN_TIMEOUT="${COLLX_RUN_TIMEOUT:-$default_guard}" + ;; # kv-transfer suite + *) 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 +# MC_FORCE_MNNVL is mooncake's only reader here: 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, +# kv CI run 3). The mooncake kv row declares the rdma lane, so it opts out. +[ "$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" @@ -75,6 +113,14 @@ 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 rdma kv legs are the only gb-nv shards that leave the NVL domain; prove +# their pinned socket interface and HCAs on the allocation like every other +# scale-out launcher does (mnnvl shards skip, as elsewhere). +if [ "$COLLX_TRANSPORT" != mnnvl ] \ + && ! collx_validate_network_profile_on_job "$JOB_ID" "$NODES" "$COLLX_TRANSPORT"; then + collx_cleanup_allocation + collx_die "network profile validation failed on the allocation" +fi # ---- container-import: squash file resolved on the allocation --------------- SQUASH_FILE="$(collx_ensure_squash_on_job "$JOB_ID" "$SQUASH_DIR" "$IMAGE")" diff --git a/experimental/CollectiveX/launchers/launch_mi-amds.sh b/experimental/CollectiveX/launchers/launch_mi-amds.sh index 1f891295f9..bf9bf36ccb 100644 --- a/experimental/CollectiveX/launchers/launch_mi-amds.sh +++ b/experimental/CollectiveX/launchers/launch_mi-amds.sh @@ -35,7 +35,26 @@ MOUNT_DIR=/ix TS="$(date -u +%Y-%m-%dT%H-%M-%SZ)" case "$COLLX_BENCH" in mori | uccl-ep) ;; - *) collx_die "unsupported AMD EP backend: $COLLX_BENCH" ;; + mori-io | mooncake) + # kv-transfer suite. The five-rung-floor grid's mi355x mori-io leg + # measured ~60 minutes (run 33097162900), and the power-of-two batch + # ladder in kv_sweep.json is another ~1.33x of descriptor work + # grid-wide, so ~80 minutes; the guard clears that with real margin + # yet fires before the raised allocation does, keeping a slow case a + # clean per-case kill instead of a lost allocation. + TIME_MIN=210 + export COLLX_RUN_TIMEOUT="${COLLX_RUN_TIMEOUT:-11400}" + # Mooncake's transfer engine cannot register the power-of-two ladder's + # 53 GiB pool on the ionic NICs (ibv_reg_mr ENOMEM, deterministic on two + # independent allocations: runs 33137809635 and 33150394862), while + # mori-io registers the same pool fine. Cap mooncake's pool at the size + # the mixed ladder proved green; run_kv sheds the top 512k-ISL rungs to + # fit, and every other grid point keeps its full ladder. + if [ "$COLLX_BENCH" = mooncake ]; then + export COLLX_KV_POOL_BUDGET="${COLLX_KV_POOL_BUDGET:-$((20 << 30))}" + fi + ;; + *) collx_die "unsupported AMD backend: $COLLX_BENCH" ;; esac export MORI_DISABLE_AUTO_XGMI="${MORI_DISABLE_AUTO_XGMI:-0}" diff --git a/experimental/CollectiveX/launchers/launch_single-slurm.sh b/experimental/CollectiveX/launchers/launch_single-slurm.sh index 5be8ceda97..33c0d96d88 100644 --- a/experimental/CollectiveX/launchers/launch_single-slurm.sh +++ b/experimental/CollectiveX/launchers/launch_single-slurm.sh @@ -46,11 +46,39 @@ NODES="${COLLX_NODES:-1}"; GPN="${COLLX_GPUS_PER_NODE:-8}" SCALE_UP_DOMAIN="${COLLX_SCALE_UP_DOMAIN:-8}" NGPUS="${COLLX_NGPUS:-$((NODES * GPN))}" TIME_MIN="${COLLX_TIME:-$DEFAULT_TIME}" +case "$COLLX_BENCH" in + nixl | mooncake | mori-io) + # The five-rung-floor grid's slowest lane here (b200 mooncake) measured + # ~96 minutes (run 33097162900), and the power-of-two batch ladder in + # kv_sweep.json is another ~1.33x of descriptor work grid-wide, so + # budget ~130 minutes. A still-short exclusive ask backfills into gaps + # that a multi-hour one waits out (h200 pool contention). The guard + # clears the projected case with real margin yet fires before the + # allocation dies, keeping a slow case a clean per-case kill instead of + # a lost allocation. + TIME_MIN=210 + export COLLX_RUN_TIMEOUT="${COLLX_RUN_TIMEOUT:-11400}" + # b300's virtualized pods stop honoring cuda memory registrations past a + # per-rank total somewhere between 9552 MiB (green) and 14843 MiB (red): + # registration reports no error and the initiator later segfaults + # resolving the region's rkey on the first transfer + # (ucp_worker_add_rkey_config), whether the pool is registered whole or + # in 4 GiB pieces (probes 2026-08-31; BAR1 is 512 GiB, so it is DMA + # mapping capacity, not aperture). Cap the pool well under the wall; + # run_kv sheds the largest batches to fit, and the piece-wise + # registration keeps every piece under the separate ~8 GiB single-MR + # wall the same pods carry. + if [ "$RUNNER" = b300 ]; then + export COLLX_KV_POOL_BUDGET="${COLLX_KV_POOL_BUDGET:-$((8 << 30))}" + fi + ;; +esac IMAGE="$COLLX_IMAGE" TS="$(date -u +%Y-%m-%dT%H-%M-%SZ)" case "$COLLX_BENCH" in deepep-v2 | uccl-ep | nccl-ep) ;; - *) collx_die "unsupported $RUNNER EP backend: $COLLX_BENCH" ;; + nixl | mooncake) ;; # kv-transfer suite + *) collx_die "unsupported $RUNNER backend: $COLLX_BENCH" ;; esac export COLLX_NGPUS="$NGPUS" COLLX_NODES="$NODES" @@ -123,6 +151,8 @@ for allocation_attempt in 1 2 3; do fi retryable=0 [ "$RUNNER:$validation_failure" != h100-dgxc:network ] || retryable=1 + # b300 carries a rotating set of nodes with inactive RDMA ports; retry elsewhere. + [ "$RUNNER:$validation_failure" != b300:network ] || retryable=1 [ "$RUNNER:$validation_failure" != b300:cuda-context ] || retryable=1 # A throttled GPU paces every rank, so retrying on another node is right on every SKU. [ "$validation_failure" != gpu-health ] || retryable=1 diff --git a/experimental/CollectiveX/runtime/common.sh b/experimental/CollectiveX/runtime/common.sh index d786799c73..c6eb6a2985 100644 --- a/experimental/CollectiveX/runtime/common.sh +++ b/experimental/CollectiveX/runtime/common.sh @@ -420,7 +420,10 @@ if [ "${COLLX_NODES:-1}" -gt 1 ] && [ "${COLLX_TRANSPORT:-}" != mnnvl ]; then fi export RANK="$SLURM_PROCID" WORLD_SIZE="$SLURM_NTASKS" export LOCAL_RANK="$SLURM_LOCALID" LOCAL_WORLD_SIZE="$COLLX_GPUS_PER_NODE" -exec python3 bench/run_ep.py "$@" +entry=run_ep +if [ "${1:-}" = "--entrypoint" ]; then entry="${2:-}"; shift 2; fi +case "$entry" in run_ep|run_kv) ;; *) exit 67 ;; esac +exec python3 "bench/${entry}.py" "$@" BASH } @@ -438,7 +441,12 @@ collx_salloc_jobid() { return 1 fi collx_log "scheduler-request=submit" - if ! (salloc "$@" --no-shell) > "$log" 2>&1; then + # Fleet convention (runners/launch_*.sh): a CI allocation is named after its + # GHA runner so operators can squeue/scancel by --name. RUNNER_NAME is set by + # the Actions runner on every step; hand launches without it keep the default. + local name_args=() + [ -z "${RUNNER_NAME:-}" ] || name_args=(--job-name="$RUNNER_NAME") + if ! (salloc ${name_args[@]+"${name_args[@]}"} "$@" --no-shell) > "$log" 2>&1; then collx_log "ERROR: scheduler allocation failed" collx_log_tail "$log" return 1 @@ -482,6 +490,13 @@ collx_cleanup_allocation() { # digest-qualified Docker Hub reference non-interactively. collx_select_image() { local image="$1" + # A registry kv backend may ship only inside a specific image (e.g. AMD's + # atom-dev mooncake build); the matrix carries that ref onto the shard and + # the workflow exports it here. + if [ -n "${COLLX_IMAGE_OVERRIDE:-}" ]; then + collx_log "image override for bench=$COLLX_BENCH: $COLLX_IMAGE_OVERRIDE" + image="$COLLX_IMAGE_OVERRIDE" + fi [[ "$image" =~ ^[A-Za-z0-9._/-]+:[A-Za-z0-9._-]+$ ]] \ || collx_die "configured image reference is malformed" export COLLECTIVEX_IMAGE="$image" diff --git a/experimental/CollectiveX/runtime/config.py b/experimental/CollectiveX/runtime/config.py index ae928550fb..1fbf754e8a 100644 --- a/experimental/CollectiveX/runtime/config.py +++ b/experimental/CollectiveX/runtime/config.py @@ -134,8 +134,50 @@ def case_count(path: str) -> None: print(len(load(path)["cases"]), end="") +def _emit_kv_argv(case: dict, version: object, runner: str, ts: str, index: int) -> None: + """Null-delimited run_kv.py argv for a kv-transfer case. The leading + --entrypoint pair is consumed by the rank wrapper (runtime/common.sh), + which otherwise execs run_ep.py.""" + timing = str(case["timing"]).split(":") + if len(timing) != 3: + print(f"unrecognised kv timing profile {case['timing']!r}", file=sys.stderr) + raise SystemExit(1) + out = ( + f"results/{runner}_kv-{case['backend']}_{case['precision']}_{case['mode']}" + f"_{case['workload']}_{ts}-c{index:03d}.json" + ) + argv = [ + "--entrypoint", "run_kv", + "--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", timing[0], "--reps", timing[1], "--trials", timing[2], + "--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), + "--out", out, + ] + sys.stdout.buffer.write(b"\0".join(part.encode() for part in argv) + b"\0") + + def _emit_argv(case: dict, version: object, runner: str, ts: str, index: int) -> None: - """Emit one null-delimited run_ep.py argv — the only case-to-invocation codec.""" + """Emit one null-delimited benchmark argv — the only case-to-invocation codec.""" + if case.get("suite") == "kv-transfer": + _emit_kv_argv(case, version, runner, ts, index) + return get = lambda key, default="": str(case.get(key) or default) argv = [ "--backend", str(case["backend"]), diff --git a/experimental/CollectiveX/runtime/prepare_backend.sh b/experimental/CollectiveX/runtime/prepare_backend.sh index fed9b6ec03..53fbc10c66 100644 --- a/experimental/CollectiveX/runtime/prepare_backend.sh +++ b/experimental/CollectiveX/runtime/prepare_backend.sh @@ -594,6 +594,29 @@ FICHECK [ "$rc" -eq 0 ] || { collx_log "ERROR: FlashInfer EP one-sided A2A unavailable in this image"; return 1; } } +# NIXL installs from the pinned wheel (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). The named container persists for the +# job, so one install here serves every case srun. On ROCm images nothing to +# do: sglang-rocm bundles nixl-cu12 with a ROCm-built UCX. +nixl_prepare() { + python3 - <<'NIXLCHECK' && return 0 +import sys +try: + import nixl # noqa: F401 +except Exception: + raise SystemExit(1) +NIXLCHECK + # Noble-based images mark python externally managed (PEP 668); the retry flag + # is the uccl-prep pattern (older pips never refuse, so they never reach it). + { python3 -m pip install -q --disable-pip-version-check --no-input 'nixl-cu13==1.3.2' \ + || python3 -m pip install -q --disable-pip-version-check --no-input \ + --break-system-packages '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; } +} + main() { collx_apply_network_profile "${COLLX_NODES:-1}" "${COLLX_TRANSPORT:-}" || return 1 validate_container_network || return 1 @@ -606,6 +629,28 @@ main() { uccl-ep) uccl_prepare || return 1 ;; nccl-ep) nccl_ep_prepare || return 1 ;; flashinfer-ep) flashinfer_ep_prepare || return 1 ;; + nixl) nixl_prepare || return 1 ;; + mooncake) + # 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. + if python3 -c "import mooncake.engine" 2>/dev/null; then + collx_log "mooncake provided by the image" + else + { python3 -m pip install -q --disable-pip-version-check --no-input \ + 'mooncake-transfer-engine==0.3.12.post1' nvidia-cuda-runtime-cu12 \ + || python3 -m pip install -q --disable-pip-version-check --no-input \ + --break-system-packages \ + 'mooncake-transfer-engine==0.3.12.post1' nvidia-cuda-runtime-cu12; } \ + || { collx_log "ERROR: mooncake wheel install failed"; return 1; } + fi + ;; + 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/experimental/CollectiveX/summarize.py b/experimental/CollectiveX/summarize.py index 47e58826b3..62abd17df3 100644 --- a/experimental/CollectiveX/summarize.py +++ b/experimental/CollectiveX/summarize.py @@ -109,15 +109,78 @@ def percentile(block: str, name: str) -> float | 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).""" + 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]) -> list[str]: + """kv-transfer table: paged bandwidth at the bandwidth-bound ISL plus the + contiguous baseline, and the paged latency. Verify failures flip the + outcome column (and the leg already failed in CI).""" + lines = ["", "## CollectiveX KV-transfer results", "", + "| 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 documents: + factors = document["identity"]["case_factors"] + case = factors["case"] + 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"| {document['version']} | {factors['sku']} | `{case['backend']}` | " + f"{case['mode']} | {case['workload']} | {case['precision']} | " + f"{document['outcome']['status']} | {op} | {paged_gbps} | {paged_bmax} | " + f"{bulk_gbps} | {paged_ms} |" + ) + lines.append("") + lines.append("> 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 lines + + def render(documents: list[dict]) -> str: - documents = sorted(documents, key=_identity) - invalid = [d for d in documents if d["outcome"]["status"] != "success"] + kv_documents = sorted( + (d for d in documents + if d["identity"]["case_factors"]["case"].get("suite") == "kv-transfer"), + key=_identity) + documents = sorted( + (d for d in documents + if d["identity"]["case_factors"]["case"].get("suite") != "kv-transfer"), + key=_identity) + invalid = [d for d in documents + kv_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. + # The leg is already red (the benchmark entrypoints return nonzero on a + # non-success outcome); call the count out loudly so it is not lost in the tables. lines.append( - f"> **{len(invalid)} of {len(documents)} outcome(s) INVALID** — " + f"> **{len(invalid)} of {len(documents) + len(kv_documents)} outcome(s) INVALID** — " "the leg fails; see the outcome column below." ) lines.append("") @@ -138,8 +201,10 @@ def render(documents: list[dict]) -> str: f"{phase} | {routing} | {ep} | {topo} | {wire} | " f"{document['outcome']['status']} | {token} | {p50} | {p99} | {min50} | {skew} |" ) - if not documents: + if not documents and not kv_documents: lines.append("\n> No valid native outcome documents found.") + if kv_documents: + lines += render_kv(kv_documents) # The starred columns can hold two different quantities, so the table always says which — and # says so loudly when it holds both, since a mixed column silently compares a steady-state # period against an idle-pipeline latency. diff --git a/experimental/CollectiveX/sweep_matrix.py b/experimental/CollectiveX/sweep_matrix.py index a99613d689..ea0fbdf6fa 100644 --- a/experimental/CollectiveX/sweep_matrix.py +++ b/experimental/CollectiveX/sweep_matrix.py @@ -25,6 +25,7 @@ def _load_config(name: str) -> dict[str, Any]: SWEEP = _load_config("sweep.json") +KV_SWEEP = _load_config("kv_sweep.json") PLATFORMS = _load_config("platform_config.json")["platforms"] # Per-backend production/candidate map for the matrix and docs; see EPBackend.maturity. BACKEND_MATURITY = _load_config("platform_config.json")["backend_maturity"] @@ -97,6 +98,75 @@ def _topology(platform: dict[str, Any], ep: int) -> dict[str, Any]: } +def _kv_backend_spec(value: Any) -> dict[str, Any]: + """kv_backends values are a fabric list (library runs the full sweep) or an + object carrying restrictions: `ops` (a backend that cannot serve one + direction on a fabric, e.g. mooncake on Pollara where ionic RDMA READ is + broken upstream), `image_ref` (a backend only shipped inside a specific + image, e.g. AMD's atom-dev mooncake build), and `device` (an engine NIC + filter template; `{gpu}` expands to the physical GPU index at runtime).""" + if isinstance(value, list): + return {"fabrics": value, "ops": None, "image_ref": None, "device": None} + return { + "fabrics": value["fabrics"], + "ops": value.get("ops"), + "image_ref": value.get("image_ref"), + "device": value.get("device"), + } + + +def _kv_cases(sku: str, platform: dict[str, Any], fabric: str, backend: str, + selected_precisions: set[str], spec: dict[str, Any]) -> list[dict[str, Any]]: + """The kv-transfer cases one (sku, backend, fabric) shard runs. + + A KV leg is 2 nodes x 1 GPU: the per-worker transfer pair an engine + actually forms, not an allocation-wide collective. Capability comes from + the registry's ``kv_backends`` map (backend -> fabrics), so a SKU with no + entry emits nothing — absence is not-yet-enabled, mirroring ll_backends. + """ + timing = KV_SWEEP["timing"] + timing_profile = ":".join(str(timing[key]) for key in ( + "warmup_per_trial", "reps_per_trial", "trials_per_point")) + 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 + # the workload model's own PRESETS (this file stays stdlib-importable for + # the runner-side matrix/extract steps, so it cannot import kv_workload). + for workload, preset_precisions in KV_SWEEP["workloads"].items(): + for precision in preset_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"] or " ".join(KV_SWEEP["ops"]), + "kv_device": spec["device"] or "", + "pool_slack": KV_SWEEP["pool_slack"], + "seed": KV_SWEEP["seed"], + "timing": timing_profile, + "nodes": 2, + "gpus_per_node": 1, + "scale_up_domain": platform["scale_up_domain"], + "scope": "scale-out", + "scale_up_transport": platform["scale_up_transport"], + "scale_out_transport": fabric, + "transport": fabric, + "topology_class": f"{platform['product']}-kv-{fabric}", + } + case["case_id"] = ep_harness.case_id(sku, case) + cases.append(case) + return cases + + def _selected_backends(backend: str) -> list[str]: if backend == "all": return list(SWEEP_BACKENDS) @@ -112,8 +182,14 @@ def resolve_matrix( ep_sizes: str = "", precisions: str = "", modes: str = "", + suites: str = "ep-core,kv-transfer", ) -> dict[str, Any]: """Resolve the fixed sweep into allocation-sized workflow shards.""" + known_suites = {SWEEP["suite"], KV_SWEEP["suite"]} + selected_suites = {value.strip() for value in suites.split(",") if value.strip()} + unknown_suites = sorted(selected_suites - known_suites) + if unknown_suites: + raise SystemExit(f"unknown --suites {unknown_suites}; have {sorted(known_suites)}") selected_eps: set[int] = set() for value in filter(None, (part.strip() for part in ep_sizes.split(","))): if not value.isdigit() or int(value) <= 0: @@ -162,7 +238,7 @@ def resolve_matrix( requested_cases: list[dict[str, Any]] = [] shards: dict[tuple[str, str, str, int, str], list[dict[str, Any]]] = {} - for sku in sorted(PLATFORMS): + for sku in sorted(PLATFORMS) if SWEEP["suite"] in selected_suites else []: if (only_sku and sku != only_sku) or sku in excluded: continue platform = PLATFORMS[sku] @@ -243,7 +319,46 @@ def resolve_matrix( (sku, target, mode, topology["nodes"], precision), [] ).append(case) + # KV-transfer shards: additive, small (4 cases), gated per SKU by the + # registry's kv_backends map and by the same sku/precision selectors. The + # EP --backend filter does not apply — its vocabulary is EP backends. + # A backend-scoped dispatch (--backend deepep-v2 etc.) is an EP-focused run; + # kv shards ride only the full-matrix resolution or an explicit + # --suites kv-transfer with the default backend selector. + kv_shards: dict[tuple[str, str, str], list[dict[str, Any]]] = {} + if KV_SWEEP["suite"] in selected_suites and backend == "all": + for sku in sorted(PLATFORMS): + if (only_sku and sku != only_sku) or sku in excluded: + continue + platform = PLATFORMS[sku] + for kv_backend, raw_spec in sorted(platform.get("kv_backends", {}).items()): + spec = _kv_backend_spec(raw_spec) + for fabric in spec["fabrics"]: + cases = _kv_cases(sku, platform, fabric, kv_backend, + selected_precisions, spec) + for case in cases: + requested_cases.append({ + "sku": sku, "case": case, "disposition": "runnable", + "reason": None, "detail": None, + }) + if cases: + kv_shards[(sku, kv_backend, fabric)] = (cases, spec) + shards_by_sku: dict[str, list[dict[str, Any]]] = {} + for (sku, kv_backend, fabric), (cases, spec) in sorted(kv_shards.items()): + shards_by_sku.setdefault(sku, []).append({ + "id": f"{sku}-kv-{kv_backend}-{fabric}", + "sku": sku, + "runner": PLATFORMS[sku].get("runner", sku), + "backend": kv_backend, + "mode": fabric, + "image_ref": spec["image_ref"] or "", + "launcher": PLATFORMS[sku]["launcher"], + "nodes": 2, + "gpus_per_node": 1, + "scale_up_domain": cases[0]["scale_up_domain"], + "cases": cases, + }) for (sku, target, mode, nodes, precision), cases in sorted(shards.items()): first = cases[0] # Normal-mode shard IDs are unchanged (no mode segment) so existing references @@ -252,6 +367,7 @@ def resolve_matrix( shards_by_sku.setdefault(sku, []).append({ "id": f"{sku}-{target}{mode_segment}-{precision}-n{nodes}", "sku": sku, + "runner": PLATFORMS[sku].get("runner", sku), "backend": target, "mode": mode, "launcher": PLATFORMS[sku]["launcher"], @@ -299,6 +415,8 @@ def main() -> int: parser.add_argument("--modes", default="", help="comma-separated subset of configs/sweep.json modes " "(normal, low-latency); blank = all") + parser.add_argument("--suites", default="ep-core,kv-transfer", + help="comma-separated suites to resolve (ep-core, kv-transfer)") parser.add_argument("--extract-from", default="", metavar="MATRIX") parser.add_argument("--shard-id", default="") parser.add_argument("--out", default="") @@ -319,6 +437,7 @@ def main() -> int: ep_sizes=args.ep_sizes, precisions=args.precisions, modes=args.modes, + suites=args.suites, ) if args.out: Path(args.out).write_text( diff --git a/experimental/CollectiveX/tests/test_kv_suite.py b/experimental/CollectiveX/tests/test_kv_suite.py new file mode 100644 index 0000000000..683ababe56 --- /dev/null +++ b/experimental/CollectiveX/tests/test_kv_suite.py @@ -0,0 +1,515 @@ +#!/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): + matrix = sweep_matrix.resolve_matrix(suites="kv-transfer") + shards = matrix["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) + for shard in shards: + self.assertEqual(shard["nodes"], 2) + self.assertEqual(shard["gpus_per_node"], 1) + # 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. + 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_ref"].startswith("rocm/atom-dev:")) + elif shard["sku"] == "b300": + # b300 pods expose two RDMA rails that are not cross-routable. + # mooncake: the image's engine draws the peer NIC per request + # blind to endpoint health, so every cross-rail draw stalls ~1s. + # nixl: unpinned UCX pull degrades reproducibly across node + # pairs while push stays at line rate; pinning one rail takes + # multi-rail selection out of the measurement. + self.assertEqual({c["ops"] for c in shard["cases"]}, {"pull push"}) + self.assertEqual({c["kv_device"] for c in shard["cases"]}, {"mlx5_0"}) + self.assertEqual(shard["image_ref"], "") + else: + self.assertEqual({c["ops"] for c in shard["cases"]}, {"pull push"}) + self.assertEqual(shard["image_ref"], "") + self.assertEqual({c["kv_device"] for c in shard["cases"]}, {""}) + 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(suites="ep-core") + both = sweep_matrix.resolve_matrix() + ep_ids = [s["id"] for s in ep_only["include"]] + both_ep_ids = [s["id"] for s in both["include"] if "-kv-" not in s["id"]] + self.assertEqual(sorted(ep_ids), sorted(both_ep_ids)) + self.assertEqual( + [c for c in 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_a_backend_scoped_dispatch_is_ep_only(self): + matrix = sweep_matrix.resolve_matrix(backend="deepep-v2") + self.assertFalse([s for s in matrix["include"] if "-kv-" in s["id"]]) + + 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 _case(self): + matrix = sweep_matrix.resolve_matrix(suites="kv-transfer") + return matrix["include"][0] + + @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 test_kv_case_encodes_behind_the_entrypoint_marker(self): + shard = self._case() + case = shard["cases"][0] + argv = self._captured_argv(case, shard["sku"]) + self.assertEqual(argv[:2], ["--entrypoint", "run_kv"]) + pairs = dict(zip(argv[2::2], argv[3::2])) + self.assertEqual(pairs["--backend"], case["backend"]) + self.assertEqual(pairs["--workload-name"], case["workload"]) + self.assertEqual(pairs["--fabric"], case["mode"]) + self.assertEqual(pairs["--case-id"], case["case_id"]) + self.assertEqual( + (pairs["--warmup"], pairs["--reps"], pairs["--trials"]), + tuple(case["timing"].split(":"))) + self.assertEqual(pairs["--batch-sizes"], case["batch_sizes"]) + self.assertEqual(pairs["--kv-device"], case["kv_device"]) + self.assertEqual(pairs["--ops"], case["ops"]) + self.assertIn(shard["sku"], pairs["--out"]) + self.assertTrue(pairs["--out"].startswith("results/")) + + def test_a_malformed_kv_timing_profile_fails_closed(self): + case = dict(self._case()["cases"][0], timing="2:8") + with self.assertRaises(SystemExit): + self._captured_argv(case, "sku") + + +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) + 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"] + saved, run_kv.POOL_BUDGET = run_kv.POOL_BUDGET, budget + try: + points, _isls, _batches = run_kv._grid(args) + finally: + run_kv.POOL_BUDGET = saved + 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"] + saved, run_kv.POOL_BUDGET = run_kv.POOL_BUDGET, budget + try: + points, _isls, _batches = run_kv._grid(args) + finally: + run_kv.POOL_BUDGET = saved + 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/experimental/CollectiveX/tests/test_kv_workload.py b/experimental/CollectiveX/tests/test_kv_workload.py new file mode 100644 index 0000000000..2c99dd0256 --- /dev/null +++ b/experimental/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/experimental/CollectiveX/tests/test_matrix.py b/experimental/CollectiveX/tests/test_matrix.py index 730556dab7..975ffa4da8 100644 --- a/experimental/CollectiveX/tests/test_matrix.py +++ b/experimental/CollectiveX/tests/test_matrix.py @@ -19,7 +19,10 @@ def matrix(**options): - return sweep_matrix.resolve_matrix(**options) + # These are the EP suite's contracts; kv-transfer has its own tests, and + # tests/test_kv_suite.py pins that resolving both suites never perturbs the + # EP half. Scoping here keeps every exhaustive assertion below exact. + return sweep_matrix.resolve_matrix(suites="ep-core", **options) def cells(document, project=("sku", "ep"), **filters):