CollectiveX: kv-transfer suite — NIXL + MoRI-IO KV-cache handoff benchmark (stacked on #2489)

.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-<N> 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 }}

.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.

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

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)

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