CollectiveX: kv-transfer suite — NIXL, Mooncake, MoRI-IO KV-cache handoff

.github/workflows/collectivex-sweep.yml

@@ -7,7 +7,7 @@ on:
   workflow_dispatch:
     inputs:
       suites:
-        description: "Comma-list of suites: ep (expert-parallel dispatch/combine), swap-blocks (vLLM block copies, one GPU per pool)"
+        description: "Comma-list of suites: ep (expert-parallel dispatch/combine), swap-blocks (vLLM block copies, one GPU per pool), kv-transfer (disaggregated KV handoff, 2 nodes x 1 GPU; hours per leg)"
         type: string
         default: ep
       backend:
@@ -149,12 +149,18 @@ jobs:
         format('[{0}]', toJSON(matrix.runner))
       ) }}
     name: p${{ needs.setup.outputs.priority }} | ${{ matrix.sku }} ${{ matrix.backend }} shard ${{ matrix.id }}
-    timeout-minutes: 350
+    # A shard with a long allocation (kv-transfer) carries its own ceiling above it, so GitHub
+    # never cancels a healthy leg before the launcher's guards act.
+    timeout-minutes: ${{ matrix.job_timeout_minutes || 350 }}
     env:
       COLLX_BENCH: ${{ matrix.backend }}
-      # Set only on shards whose suite ships its own image (swap-blocks); blank keeps the pool's.
+      # Set only on shards that ship their own image (swap-blocks; mi355x mooncake); blank keeps the pool's.
       COLLX_IMAGE_OVERRIDE: ${{ matrix.image || '' }}
       COLLX_STAGED_IMAGE_DIR: ${{ matrix.staged_image_dir || '' }}
+      # The fabric a shard claims (kv-transfer rdma legs on an MNNVL rack leave the NVL domain)
+      # and its per-case hang guard; blank keeps the launcher defaults.
+      COLLX_FABRIC: ${{ matrix.fabric || '' }}
+      COLLX_RUN_TIMEOUT: ${{ matrix.run_timeout || '' }}
       COLLX_IMPORT_TMPDIR: ${{ matrix.sku == 'h100-dgxc' && '/var/tmp' || '/tmp' }}
       COLLX_NODES: ${{ matrix.nodes }}
       COLLX_GPUS_PER_NODE: ${{ matrix.gpus_per_node }}

collectivex/README.md

@@ -4,7 +4,9 @@ CollectiveX is an experimental MoE expert-parallel communication benchmark. It m
 combine, and paired roundtrip latency across EP libraries and accelerator systems, then uploads
 neutral result artifacts.
 
-A second suite, the [vLLM `swap_blocks` benchmark](docs/swap-blocks.md)
+A second suite, `kv-transfer`, measures disaggregated-serving KV-cache handoffs across transfer
+libraries and fabrics ([below](#kv-cache-transfer-suite)). A third, the
+[vLLM `swap_blocks` benchmark](docs/swap-blocks.md)
 ([中文](docs/swap-blocks_zh.md)), measures pinned CPU↔GPU and same-GPU block copies,
 with its own correctness checks and latency/bandwidth JSON output. It runs through the same
 sweep matrix and pool launchers (`suites: swap-blocks`).
@@ -126,7 +128,8 @@ The matrix covers H100, H200, B200, B300, GB200, GB300, MI300X, MI325X, and MI35
 the requested SKUs, backends, EP sizes, and token ladders, then extracts strict per-shard controls
 and rejects missing, stale, malformed, or altered shard controls. `--only-sku`, `--exclude-skus`,
 `--ep-sizes`, and `--precisions` select a subset. `--suites` picks the suites (`ep` by default;
-`swap-blocks` adds one single-GPU shard per pool from `configs/swap_sweep.json`). Every suite's
+`swap-blocks` adds one single-GPU shard per pool from `configs/swap_sweep.json`; `kv-transfer` adds
+one 2-node shard per registry-enabled backend and fabric from `configs/kv_sweep.json`). Every suite's
 shards take the same path: the pool launcher allocates, `runtime/config.py case-args` encodes each
 case as its entrypoint's argv, and the rank wrapper execs that entrypoint. The matrix is generated
 per dispatch, with no frozen digest or locked case count.
@@ -164,6 +167,27 @@ scale-up ranks per domain. GB EP16 remains MNNVL scale-up and therefore uses LSA
 SKU/backend/EP cell is attempted is a capability fact. Whether it succeeded is decided by the
 benchmark's return code.
 
+## KV-Cache Transfer Suite
+
+`kv-transfer` legs run 2 nodes x 1 GPU, the per-worker prefill/decode pair a disaggregated
+deployment actually forms. Each moves bursts of 1 to 32 concurrent requests' paged KV (vLLM's
+packed block-major descriptor lists over seed-keyed random block tables) plus one contiguous bulk
+row as the wire-speed baseline. The workload is `kv-dsv4`, DeepSeek-V4-Pro's mixed fp8 cache as
+vLLM allocates it, at ISL 2k to 512k and vLLM's 256-token block. `pull` (READ, vLLM
+NixlConnector) and `push` (WRITE, SGLang disagg) are both timed from the initiator, with every
+request verified on the destination pool. The [methodology](docs/methodology.md#kv-cache-transfer-suite)
+has the geometry and verification model.
+
+Backends are `nixl` (what Dynamo, vLLM, and SGLang ship), `mooncake` (the CUDA wheel links libcuda
+at import; mi355x runs AMD's atom-dev build push-only), and `mori-io` (AMD's native engine), where
+the registry's `kv_backends` map enables them. No entry, no legs, mirroring `ll_backends`. An entry
+may restrict ops, pin an image, set a NIC filter, or lower the pool budget. Fabrics are `rdma`
+(torch pools) and, on GB racks, `nixl` also runs `mnnvl` (cuMem FABRIC pools; see the methodology
+for the bulk-vs-paged lane inversion that row exists to publish). On b300's AWS EFA pool `nixl`
+rides its LIBFABRIC plugin instead of UCX, and there is no `mooncake` leg (its wheel's transport is
+verbs RC, which EFA does not provide). The grid and per-pool scheduling
+live in `configs/kv_sweep.json`; dispatch with `suites: kv-transfer`.
+
 ## Workflow And Artifacts
 
 `.github/workflows/collectivex-sweep.yml` has two jobs. `setup` generates a public-SKU matrix

collectivex/bandwidth.py

@@ -32,7 +32,7 @@
 from typing import NamedTuple
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-from summarize import load_results  # noqa: E402  (sibling consumer, stdlib-only)
+from summarize import KV_SUITE, load_results  # noqa: E402  (sibling consumer, stdlib-only)
 
 # stage moves nothing (bytes always 0); isolated_sum is a derived percentile sum, not a real
 # chained rate. Only these three carry a measured latency + bytes.
@@ -184,6 +184,12 @@ def _provenance(document: dict) -> str:
 
 
 def render(documents: list[dict]) -> str:
+    # kv-transfer documents have their own row model (per-transfer, no tokens_per_rank or
+    # routing); this renderer reads only EP rows, and prints nothing for a kv-only leg.
+    documents = [d for d in documents
+                 if d["identity"]["case_factors"]["case"].get("suite") != KV_SUITE]
+    if not documents:
+        return ""
     lines = [
         "## CollectiveX EP bandwidth (per-GPU, wire-basis payload)",
         "",

collectivex/bench/ep_harness.py

@@ -178,6 +178,43 @@ def _component(percentiles, count, *, derived=False, origin=None):
 )
 
 
+def git_run() -> dict | None:
+    """The workflow run that produced an attempt, from the environment the launcher forwards."""
+    run = {
+        "run_id": os.environ.get("GITHUB_RUN_ID"),
+        "run_attempt": os.environ.get("GITHUB_RUN_ATTEMPT"),
+        "source_sha": os.environ.get("COLLECTIVEX_SOURCE_SHA") or os.environ.get("GITHUB_SHA"),
+    }
+    return run if any(run.values()) else None
+
+
+def case_attempt(args, case: dict, run: dict | None, image: str, ok: bool, failure: str,
+                 **sections) -> dict:
+    """The case-attempt document every suite writes: identity, provenance and outcome keyed and
+    validated the same way around the suite's own sections (workload, measurement, ...)."""
+    try:
+        attempt_ordinal = int(os.environ.get("COLLX_ATTEMPT_ID", "1"))
+    except ValueError:
+        attempt_ordinal = 0
+    if attempt_ordinal <= 0:
+        raise ValueError("COLLX_ATTEMPT_ID must be a positive integer")
+    run = run or {}
+    return {
+        "version": args.version,
+        "record_type": "case-attempt",
+        "generated_at": _dt.datetime.now().astimezone().isoformat(),
+        "identity": {
+            "allocation_factors": {key: run.get(key) for key in ("run_attempt", "run_id", "source_sha")},
+            "attempt_ordinal": attempt_ordinal,
+            "case_factors": {"case": case, "sku": args.runner},
+            "case_id": args.case_id,
+        },
+        **sections,
+        "provenance": {"image": image or None, "source_sha": run.get("source_sha")},
+        "outcome": {"reasons": [] if ok else [failure], "status": "success" if ok else "invalid"},
+    }
+
+
 def _write_json_atomic(path: str, value) -> None:
     payload = json.dumps(value, allow_nan=False, ensure_ascii=False, separators=(",", ":")).encode() + b"\n"
     os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
@@ -696,30 +733,16 @@ def _document(self, rows, all_ok):
         computed = case_id(args.runner, scheduled_case)
         if args.case_id != computed:
             raise ValueError(f"scheduled case ID does not match realized factors: {args.case_id} != {computed}")
-        git_run = getattr(args, "git_run", None) or {}
-        try:
-            attempt_ordinal = int(os.environ.get("COLLX_ATTEMPT_ID", "1"))
-        except ValueError:
-            attempt_ordinal = 0
-        if attempt_ordinal <= 0:
-            raise ValueError("COLLX_ATTEMPT_ID must be a positive integer")
-        return {
-            "version": args.version,
-            "record_type": "case-attempt",
-            "generated_at": _dt.datetime.now().astimezone().isoformat(),
-            "identity": {
-                "allocation_factors": {key: git_run.get(key) for key in ("run_attempt", "run_id", "source_sha")},
-                "attempt_ordinal": attempt_ordinal,
-                "case_factors": {"case": scheduled_case, "sku": args.runner},
-                "case_id": args.case_id,
-            },
-            "workload": {
+        return case_attempt(
+            args, scheduled_case, getattr(args, "git_run", None), getattr(args, "image", ""),
+            all_ok, "semantic correctness or routing identity failed",
+            workload={
                 "cross_rank_consistent": self.routing_consistent,
                 "ladder_measured": list(self.ladder),
                 "ladder_dropped": list(self.dropped),
                 "ladder_cap": self.cap,
             },
-            "measurement": {
+            measurement={
                 "combine_dtype": backend.combine_dtype,
                 "combine_semantics": "activation-only",
                 "dispatch_dtype": backend.dispatch_dtype,
@@ -736,7 +759,7 @@ def _document(self, rows, all_ok):
                     "chain_trials": args.chain_trials,
                 },
             },
-            "implementation": {
+            implementation={
                 "fp8_consume": backend.fp8_consume,
                 "kernel_generation": kernel_generation(backend),
                 # The reduction the oracle held the kernel to, and the library version it was
@@ -751,7 +774,7 @@ def _document(self, rows, all_ok):
                 "maturity": backend.maturity or "unknown",
                 "name": backend.name,
             },
-            "topology": {
+            topology={
                 "device_product": getattr(args, "runtime_device_product", None),
                 "gpus_per_node": args.gpus_per_node,
                 "nodes": nodes,
@@ -761,13 +784,8 @@ def _document(self, rows, all_ok):
                 "transport": args.transport,
                 "world_size": self.world_size,
             },
-            "runtime": getattr(args, "runtime", {}),
-            "provenance": {"image": getattr(args, "image", "") or None, "source_sha": git_run.get("source_sha")},
-            "outcome": {
-                "reasons": [] if all_ok else ["semantic correctness or routing identity failed"],
-                "status": "success" if all_ok else "invalid",
-            },
-        }
+            runtime=getattr(args, "runtime", {}),
+        )
 
     def _print_summary(self, rows, doc):
         # Ladder ends plus two interior points: one mid-ladder headline hides low-token behavior.

collectivex/bench/kv_backend.py

@@ -0,0 +1,129 @@
+#!/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
+
+import kv_workload
+
+
+class KVBackend:
+    """One transfer library on one rank, constructed as ``Backend(args, role,
+    device)``. Subclasses implement the five hooks."""
+
+    name = "abstract"
+    #: maturity mirrors EPBackend.maturity ("production" | "candidate").
+    maturity = "candidate"
+    library_version: str | None = None
+    #: the engine NIC filter this case ran under; None = library/UCX choice.
+    nic_filter: str | None = None
+    #: the library's own transport plugin when it has several (NIXL: UCX or
+    #: LIBFABRIC); None = the library has one.
+    transport: str | None = None
+
+    # -- lifecycle ------------------------------------------------------------
+    def register(self, pool, bulk, reg_layout=None) -> None:
+        """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 release(self) -> None:
+        """Drop the transfers made for the grid point just measured (initiator)."""
+
+    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 library_version(dists, module=None) -> str | None:
+    """The first installed distribution's version, else ``module.__version__``."""
+    import importlib.metadata as md
+
+    for name in dists:
+        try:
+            return md.version(name)
+        except md.PackageNotFoundError:
+            pass
+    return getattr(module, "__version__", None)
+
+
+def spans(n: int, cap: int) -> list[tuple[int, int]]:
+    """[start, end) pieces of at most ``cap`` covering ``range(n)``."""
+    return [(i, min(i + cap, n)) for i in range(0, n, cap)]
+
+
+def offset_lists(cfg: dict, local_tables, remote_tables, local_base: int = 0,
+                 remote_base: int = 0) -> tuple[list, list, list]:
+    """(local addrs, remote addrs, sizes) as Python lists: pool offsets plus the
+    given bases (0 for offset-addressed engines)."""
+    return ((local_base + kv_workload.page_offsets(cfg, local_tables)).tolist(),
+            (remote_base + kv_workload.page_offsets(cfg, remote_tables)).tolist(),
+            kv_workload.desc_sizes(cfg).tolist())
+
+
+def time_bursts(transfers, warmup: int, reps: int) -> tuple[list[float], list[float]]:
+    """(burst_ms, request_ms), warmups dropped. ``transfers`` is a list of
+    (post, wait) pairs — one per request. A burst posts every request, then
+    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