CollectiveX: kv-transfer suite — NIXL + MoRI-IO KV-cache handoff benchmark (stacked on #2489) - #2510
CollectiveX: kv-transfer suite — NIXL + MoRI-IO KV-cache handoff benchmark (stacked on #2489)#2510Oseltamivir wants to merge 25 commits into
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Additional findings (outside current diff — PR may have been updated during review):
-
🟡
experimental/CollectiveX/configs/platform_config.json:105-112— Enabling kv_backends on gb200 routes its KV leg through launch_gb-nv.sh, which unconditionally exports COLLX_TRANSPORT=mnnvl for every shard and never calls collx_validate_network_profile_on_job — unlike launch_single-slurm.sh/launch_mi-amds.sh, which branch to an -rdma transport for scale-out and validate the fabric. Because transport stays mnnvl, collx_apply_network_profile, the rank wrapper's network branch, and prepare_backend.sh's validate_container_network all skip the fail-closed HCA/interface checks for this leg, even though this PR's own fabric note calls the gb200 KV leg real cross-node InfiniBand. The transfer itself likely still works (run_kv pins UCX/gloo selectors from the operator config independently of transport), so the concrete loss is the missing pre-flight fabric validation, not a guaranteed break.Extended reasoning...
What's happening:
configs/platform_config.jsonnow setskv_backends: {nixl: [rdma]}ongb200(line 108) alongside a newnetwork.rdma_devices/socket_ifnameblock and a fabric note that explicitly calls the KV leg "4x ConnectX-7 NDR400 InfiniBand (KV scale-out; EP stays MNNVL)". That's a genuine cross-node RDMA transfer.sweep_matrix._kv_cases()schedules this as a 2-node x 1-GPU shard usingPLATFORMS["gb200"]["launcher"], which isgb-nv.launchers/launch_gb-nv.sh(untouched by this PR) unconditionally doesexport COLLX_TRANSPORT=mnnvlfor every shard it runs and never callscollx_validate_network_profile_on_job. Compare that tolaunch_single-slurm.shandlaunch_mi-amds.sh, which branchCOLLX_TRANSPORTto an-rdmavariant whenNODES>1and then validate the fabric on the job before proceeding.Why this was fine before, and why it isn't now: gb200/gb300 previously only ran the EP suite, whose EP16 always stays inside the 72-GPU MNNVL scale-up domain, so
mnnvlwas always the correct transport label for anything gb-nv launched. This PR is the first thing that schedules a real scale-out RDMA leg (KV transfer) on a gb-nv-launched SKU, and the launcher has no branch to distinguish that case from the EP/MNNVL case.Concrete effect: because
COLLX_TRANSPORTstaysmnnvlfor the KV shard, three fail-closed validation paths all skip:collx_apply_network_profile(runtime/common.sh) early-returns on itsnodes>1 && transport!=mnnvlgate, so it never validates or exportsNCCL_IB_HCA/GLOO_SOCKET_IFNAMEfor this leg.- The rank wrapper's own network branch in
common.shis gated the same way and is skipped. prepare_backend.sh'svalidate_container_networkis likewise gated ontransport != mnnvland returns early.
So the gb200 KV leg is the only scale-out RDMA row in the registry that never gets the "prove the configured socket interface and RDMA HCA actually exist on every allocated node" check every other scale-out fabric (b200-nscale, mi355x, and the EP16 x86 rows) gets.
Step-by-step to see it:
platform_config.json:105-112— gb200 haslauncher: gb-nvandkv_backends: {nixl: [rdma]}.sweep_matrix._kv_cases()builds a case withnodes=2, gpus_per_node=1for this SKU/backend.- That case dispatches through
launch_gb-nv.sh, which doesexport COLLX_TRANSPORT=mnnvlwith noNODES-based branch (contrastlaunch_single-slurm.sh's scale-out branch). collx_apply_network_profile 2 mnnvlis called somewhere downstream; its gate[ "$nodes" -gt 1 ] && [ "$transport" != mnnvl ]evaluates2 -gt 1 && mnnvl != mnnvl→ false, so it returns immediately without validating anything.- Same story for
validate_container_networkand the rank wrapper's branch — both keyed off the sametransport != mnnvltest. - Net effect: the KV shard runs without ever confirming the InfiniBand interfaces/HCAs named in the new
networkblock actually exist on the allocated nodes.
Why it's not a hard break:
run_kv.py'sexport_ucx_selectors()pinsUCX_NET_DEVICES/UCX_IB_GID_INDEXdirectly fromCOLLX_RDMA_DEVICES/COLLX_IB_GID_INDEX, independent of the mnnvl branch, and it derivesGLOO_SOCKET_IFNAMEsimilarly. So the actual UCX/gloo transfer likely still selects the right devices and runs correctly on a healthy node — the loss is specifically the pre-flight, fail-closed proof (that methodology.md documents as required for every non-MNNVL scale-out node) that those devices exist and are up, not a guaranteed crash or silently wrong measurement.How to fix: give
launch_gb-nv.shthe sameNODES>1branch the other two launchers have — select an-rdmatransport variant for the KV shard and callcollx_validate_network_profile_on_jobon it — so a bad or missing IB config on a gb200 node fails the shard early instead of running the transfer unvalidated. Note this is independent of the separately-reportedCOLLX_BENCHallowlist issue on the same launcher family: that gate is a backend-name check that happens after theCOLLX_TRANSPORT=mnnvlassignment, so fixing it alone would still leave this transport hardcoding in place.
…date gb-nv rdma legs Review findings on #2510, both real. The launchers collx_die on COLLX_BENCH values outside their EP enum, so every kv shard died at the identity stage; nixl/mooncake/mori-io are now accepted where the registry schedules them, pinned by a test that greps each SKU's launcher for its kv backends. launch_gb-nv.sh also exported COLLX_TRANSPORT=mnnvl unconditionally, which made a gb200 kv rdma leg the only scale-out fabric that skipped collx_apply_network_profile, the rank wrapper's network branch, and validate_container_network. kv rdma shards now carry mnnvl-rdma (the workflow exports the shard mode) and the launcher proves the pinned socket interface and HCAs on the allocation before running, like every other scale-out launcher; mnnvl shards keep skipping, as elsewhere.
|
On the gb-nv transport finding from the review: fixed in dac114e. kv rdma shards on gb-nv now carry COLLX_TRANSPORT=mnnvl-rdma (the workflow exports the shard mode), which re-enables collx_apply_network_profile, the rank wrapper's network branch, and validate_container_network for those legs, and the launcher gained the same on-allocation collx_validate_network_profile_on_job check the other scale-out launchers run. mnnvl shards keep the mnnvl label and skip, as before. |
…ff benchmark) Disaggregated serving's prefill->decode KV handoff becomes a first-class suite beside ep-core: 2-node x 1-GPU legs move one request's paged KV as layer-major descriptor lists over seed-keyed random block tables (the post-fragmentation layout vLLM/SGLang connectors post), pull and push, against a single-descriptor bulk wire ceiling, with offset-pattern verification on the destination pool in both directions. Workloads name production shapes (kv-mla: 61x576 DeepSeek/Kimi latent; kv-gqa: 94x1024 Qwen3-235B class) at bf16 and fp8; page sizes 16/64 pin the dominant variable measured on the metal (a ~1.5us/descriptor floor puts 16-token MLA pages at ~25% of wire while 64-token pages saturate). Backends: nixl (the library Dynamo/vLLM/SGLang ship; pip wheel on NVIDIA images, bundled with a ROCm UCX in the mi355x image) and mori-io (AMD's native engine; batch posts capped at 16384 offsets, bulk WRs split at 1 GiB below provider max-message limits). Capability is registry-gated via kv_backends, mirroring ll_backends; kv shards resolve only on --backend all dispatches and never perturb the EP matrix (pinned by test).
Slurm propagates the submitter's soft RLIMIT_MEMLOCK into job steps, and a 3.8 GiB soft limit fails a 12.7 GB pool MR deep inside the library (MoRI errno 12, UCX EIO on ucp_mem_map) with no hint at the cause. Raise soft to hard when hard allows -- covering every launcher path without touching their salloc flags -- and fail with the actual numbers when it does not.
Validated on the metal over the Pollara rails: 24-25 GB/s paged at library defaults (qp=1, no chunking), pattern-verified on every row. NIXL stays off this SKU for now -- both the image-bundled and a source-built UCX stack fall back to TCP (~0.8 GB/s) for ROCm memory on the ionic provider; the GDR path there is still being diagnosed and the row would misrepresent the library.
UCX auto-selection is a wrong-fabric trap (b200-nscale's quad-port aux card, b300's storage IB). run_kv maps COLLX_RDMA_DEVICES / COLLX_IB_GID_INDEX into UCX_NET_DEVICES / UCX_IB_GID_INDEX before any backend instantiates UCX; explicit UCX_* values still win.
The probe finally scheduled once the request matched what a KV leg needs (gpu:1 + bounded mem on a mixed partition): 36/36 rows pattern-verified over the IB VFs -- bulk 43.4 GB/s, paged-64 32-36, paged-16 9.2-16.2 (the per-descriptor floor rides a little heavier on SR-IOV).
Mooncake TransferEngine, probe-validated on b200-nscale/h200/gb200 (36/36 verified rows each; its batch path has the lowest per-descriptor floor and beats NIXL at 16-token MLA pages on B200) and measured infeasible on AMD (the wheel links libcuda.so.1 at import). The adapter dlopens libcudart.so.12 from nvidia-cuda-runtime-cu12, so cu13 images need no LD_LIBRARY_PATH seam. Fabrics become real: pools move behind kv_pool (torch for rdma, cuMem FABRIC via ctypes for mnnvl -- UCX's cross-node cuda_ipc only engages on fabric-mappable memory; on cudaMalloc the flag is inert and silently rides the IB rails), adapters register raw pointers, and gb200 gains the mnnvl row. Measured on the rack: bulk 636-707 GB/s (7.9x the IB rails) but ~3.9us per descriptor copy, so paged rows land BELOW the IB lane -- the inversion the two fabric rows exist to publish. run_kv smoke-ran end to end on b200 (nixl + mooncake, rdma) and gb200 (nixl, mnnvl): status=success artifacts on all three. The smoke caught a bulk-row verdict crash (no verifying side -> StopIteration on both ranks), fixed as exchange_verdict with the no-verifier case pinned by test. Also trims prose to its factual core and drops the dead --link-gbps parameter.
…date gb-nv rdma legs Review findings on #2510, both real. The launchers collx_die on COLLX_BENCH values outside their EP enum, so every kv shard died at the identity stage; nixl/mooncake/mori-io are now accepted where the registry schedules them, pinned by a test that greps each SKU's launcher for its kv backends. launch_gb-nv.sh also exported COLLX_TRANSPORT=mnnvl unconditionally, which made a gb200 kv rdma leg the only scale-out fabric that skipped collx_apply_network_profile, the rank wrapper's network branch, and validate_container_network. kv rdma shards now carry mnnvl-rdma (the workflow exports the shard mode) and the launcher proves the pinned socket interface and HCAs on the allocation before running, like every other scale-out launcher; mnnvl shards keep skipping, as elsewhere.
Per the model config: the MLA cache is expressed as MQA (1 kv head x head_dim 512) plus rope 64, and the DSA indexer keeps its own k cache (index_head_dim 128) that vLLM's connector merges into the transfer regions, so the per-token-per-layer transfer unit becomes 704 elements across the same 61 layers. Workload name and case ids stay kv-mla; the probe evidence in the PR was measured at the prior 576-element shape and the first CI dispatch re-baselines at this one.
…cross batch sizes kv-dsv4 replaces kv-mla: V4-Pro has no MLA; the preset transcribes vLLM's cache specs region by region (30 CSA layers at 4 tokens per 576B fp8_ds_mla entry plus their 132B lightning-indexer entries, 31 HCA layers at 128 tokens per entry, and the 128-token sliding-window cache on all 61 layers), pinned fp8 because the dtype mix is architectural. Transfers gain a batch dimension: each request in a burst is its own prepped transfer over a disjoint slice of one block-table permutation; bursts post all requests then await all, timed as one completion (split post/wait backend contract). Grid points shed batches that cannot fit the per-rank pool budget instead of dropping out. Mooncake posts its sync calls from a worker pool, the SGLang connector shape. Verification is per-byte so unaligned regions (132B indexer) check exactly.
The kv precision gate moved sweep_matrix onto kv_workload, which imports numpy at module top; the matrix/extract steps run on bare runners where that import is not a given. The sweep config now maps each workload to its precisions directly (a test pins the map to the workload model's PRESETS, and plan_config still fail-closes on a mismatch at runtime). The CollectiveX test job installs numpy: the cpu torch wheel does not pull it in, so every kv test module had been failing at import since the suite landed.
A kv-transfer validation run had no way in: the default dispatch resolves every suite (the full EP matrix rides along), and only_sku still drags that SKU's EP shards. The input passes straight to sweep_matrix --suites and joins the concurrency key so suite-scoped runs do not queue behind full sweeps.
The noble-based sweep images mark the container python externally managed (PEP 668), so the bare mooncake install refused on b200 and gb200 in the first kv CI run (nixl dodged it only because the image bundles nixl). Same try-then-retry shape the uccl prep already uses; older pips never refuse, so they never reach the flag.
… pairs On a same-rack GB200 pair the engine's NVLink-IPC transport claims the cross-node segments (one NVLink domain) and fails the address import (nvlink_transport 'Requested address not found', first kv CI run). The row declares the rdma lane, so the adapter sets MC_USE_NVLINK_IPC=0; the ROCm twin knob (MC_USE_HIP_IPC) misclaims the same way on mi355x. Transfer failures now carry the library rc.
launch_gb-nv.sh exported MC_FORCE_MNNVL=1 for every bench; mooncake is its only reader and it responds by installing only the cross-node NVLink transport, which cannot open another host's segments in the pinned wheel (cudaIpcOpenMemHandle: invalid resource handle; found 6 HCAs but skipped rdma, kv CI runs 1-3 on gb200). The mooncake kv row declares the rdma lane, so the mooncake bench opts out; EP benches keep the export untouched.
Drops the kv-gqa preset and its cases: the suite measures the deployment shape of interest, and the dense-counterpoint numbers told their story (the descriptor-floor gap is visible against each lane's own bulk ceiling). One workload x one precision per shard; the pool-budget shedding test pins the mechanism against a synthetic budget since dsv4 alone never nears 64 GiB.
run_kv copied its --version argv string straight into the document while ep_harness types the same flag as int; consumers comparing against the numeric matrix version had to coerce. Mirrors the EP entrypoint.
The first production charts show a bandwidth bump at batch 4 on some lanes where scaling should be monotonic; batches 2 and 8 bracket it to separate a real concurrency sweet spot from sampling noise, and a fourth trial tightens the p50. Worst lane (gb200 mnnvl descriptor floor) computes to about 31 minutes per case, well inside the 5400 s case guard.
…ptor budget The ladder becomes 8k / 32k / 128k / 512k. A 512k page-16 request alone is ~2.1M descriptors and burst posting time is linear in batch x descriptors on the per-descriptor floor, so points now shed batches whose burst exceeds DESC_BUDGET (the smallest batch always survives, keeping a single request measurable at every point) before the pool-budget fit; re-planning the pool for the surviving batch also drops the 512k pool from ~57 GB to ~8 GB. The budget is sized to keep every 32k cell of the previous grid; the fourth sampling trial reverts to three now that the dense grid has pinned the mooncake batch-4 peak as stable signal. Slowest lane (gb200 mnnvl floor) computes to ~62 minutes per case against the 5400 s guard.
One octave past the mooncake collapse (does batch 32 keep falling?), past nixl's flat line, and toward mori-io's saturation. The descriptor budget sheds the new rung wherever a burst cannot afford it (32k page-16 stays capped at 16, the 128k/512k caps are unchanged), so only 8k and 32k page-64 points grow; the slowest lane computes to ~73 minutes against the 5400 s guard, and the largest pool stays ~7.4 GB.
The descriptor budget places the new rung where a burst affords it (8k both pages; 32k page-64 keeps 32, larger ISLs unchanged). The gb200 mnnvl descriptor floor makes the dense kv grid ~90 minutes of legitimate work, so kv benches on gb-nv raise the per-case hang guard to 7200 s; every other lane computes to 45 minutes or less under the default guard.
A kv case is ~45 minutes of work (x86) or ~90 (gb mnnvl floor) but the allocation asked for the fleet-wide 300; on a contended pool a 5-hour exclusive 2-node ask cannot backfill into gaps between long benchmark jobs and pends for hours (both h200 legs burned two full GHA budgets waiting behind the pool's post-wedge backlog). kv benches now ask 150 (single-slurm) or 180 (gb-nv) minutes, which Slurm can slot into windows the big jobs leave open.
72c1c3b to
f2f36c3
Compare
…v build
A kv_backends entry can now restrict a backend: ops for one-direction
fabrics, image_ref for builds that ship only inside a specific image, and
device for engine NIC filters ({gpu} expands to the physical GPU index).
mooncake on mi355x uses all three: AMD's atom-dev build moves WRITE at wire
speed over the GPU-paired Pollara NIC (probed 30.7-33.6 GB/s bulk, above
mori-io), while upstream ionic RDMA READ completes with retry-exceeded and
one failed READ poisons the engine, which is why ATOM's production
connector is write-only. The shard's image ref rides the matrix into
COLLX_IMAGE_OVERRIDE and collx_select_image; prepare_backend keeps an
image-provided mooncake instead of pinning the CUDA wheel; the adapter
imports without the cudart preload when the build is self-contained and
passes the resolved NIC filter to initialize. The summary table gains an
op column naming the measured direction.
The engine fails any single sync transfer after 30 s; on the collapse-prone h200 lane a chunked page-16 batch-16 call can legitimately exceed that while thrashing (the exact contention the batch dimension publishes), flaking the shard on run-to-run variance. MC_TRANSFER_TIMEOUT=120 keeps it a hang guard with 4x headroom; the runtime's per-case guard still bounds a wedged case.
… a scaling step at every point The descriptor budget shed 512k page-16 to batch 1 (a single request is ~2.1M descriptors), collapsing the requests-per-burst axis to one point there. Keep the two smallest requested batches over the budget, bounded and deliberate, and size the gb-nv per-case guard (7200 -> 9000 s) for the ~15 extra minutes the mnnvl descriptor floor spends on the new batch-2 burst. Pool budget stays a hard limit the floor never overrides.
… guards Verification previously checked only request 0 of a burst, so a batch>=2 row never proved the other requests moved; every request is now checked against its own block tables. The gb-nv kv hang guard rises to 10500 s: the mnnvl grid measured ~135 minutes of bursts against the old 9000 s guard (11% margin), and the guard must still fire before the 180-minute allocation dies. A new test pins the pool budget between the 512k page-16 batch-1 and batch-2 pool sizes and asserts the floor-kept batch is shed, closing the untested pool-overrides-floor claim; the grid test fixture gains the 131072 rung so tests exercise the shipped ladder.
1 participant
Adds a second suite: the prefill-to-decode KV handoff of disaggregated serving, measured with the transfer libraries engines actually ship (NIXL, Mooncake, MoRI-IO) on real fabrics. A leg is 2 nodes x 1 GPU moving bursts of 1 to 64 concurrent requests' paged KV as per-request layer-major descriptor lists over seed-keyed random block tables (the post-fragmentation layout vLLM/SGLang post; batched requests slice disjoint ranges of one permutation), pull and push, against a one-descriptor bulk wire-ceiling row. Each request is its own prepped transfer; a burst posts all, then awaits all, the way a decode step admits several requests at once.
The workload is transcribed from what vLLM allocates for the model it serves, region by region. kv-dsv4 is DeepSeek-V4-Pro as vLLM serves it (MXFP4 checkpoints included, since quantization covers weights while the cache layout is architectural): 30 Compressed Sparse Attention layers at 4 tokens per 576 B fp8_ds_mla entry plus their 132 B lightning-indexer entries, interleaved with 31 Heavily Compressed Attention layers at 128 tokens per entry, plus the 128-token sliding-window cache on all 61 layers; fp8 pinned because the dtype mix is architectural. Pages 16/64, ISL 8k/32k/128k/512k, batches 1/2/4/8/16/32/64; two budgets shed a point's largest batches instead of dropping the point (a 64 GiB per-rank pool budget, and a per-burst descriptor budget, since a 512k page-16 request alone is ~2.1M descriptors and posting time is linear in batch x descriptors), with the two smallest batches always kept so a single request stays measurable everywhere and the batch axis keeps a one-to-two scaling step at every point; both directions pattern-verified per byte (the 132 B indexer entries land at any alignment), failed verify = invalid artifact = red leg. Details in docs/methodology.md.
Every registry row was probe-validated on the metal, run_kv smoke-ran end to end at the final geometry on b200, gb200, and mi355x, and the suite runs in CI: the sweep workflow gained a suites input (suite-scoped dispatches without the EP matrix). CI sweeps validated the suite end to end across b200-nscale, h200-dgxc, gb200 (rdma and mnnvl), and mi355x; the full nine-shard matrix is all green in a single reference run, 31565324148 (on the expanded grid with the two-batch floor), every row pattern-verified.
Measured rows (kv-dsv4 pull at ISL 32k, aggregate GB/s, CI run 31180525317)
Findings the suite exists to publish (final grid):
Fail-closed exclusions, with evidence: mi355x nixl (image and source-built UCX both ride TCP at 0.8 GB/s on the ionic provider; ib transports never selected even with GDR forced) and mooncake on AMD (the wheel links libcuda.so.1 at import; a ROCm source build hits upstream HIP-IPC bugs on cross-node paths).
Hardening from the probe and CI campaigns
run_kv raises soft RLIMIT_MEMLOCK to hard before registration; UCX is pinned to the registry RDMA selectors (auto-select is a wrong-fabric trap); NIXL metadata rides the harness exchange (no listener race); MoRI bulk WRs split at 1 GiB and batch posts are capped, library defaults only (qp4 + chunking wedged on metal); wheels pinned (nixl-cu13==1.3.2, mooncake-transfer-engine==0.3.12.post1 with nvidia-cuda-runtime-cu12 dlopened by the adapter). Scheduling: kv benches ask backfill-friendly allocation times (150/180 min instead of the fleet 300, since a multi-hour exclusive 2-node ask starves on contended pools), and gb-nv kv cases carry a 9000 s hang guard for the mnnvl descriptor floor's ~105 minutes of legitimate work (the batch-axis floor includes a 512k page-16 batch-2 burst). Mooncake's engine guards each sync call with a 30 s transfer timeout; page-16 high-batch bursts brushed it once on the h200 collapse lane, so the adapter raises MC_TRANSFER_TIMEOUT to 120 s (a real hang still fails the case through the harness's own per-case guard). Two CI-only issues the hand probes could not see: the noble-based b200/gb200 images refuse bare pip (PEP 668), so the kv wheel installs retry with --break-system-packages; and launch_gb-nv.sh exported MC_FORCE_MNNVL=1 for every bench, which makes mooncake (its only reader) install only its cross-node NVLink transport and fail cudaIpcOpenMemHandle on remote segments, so the mooncake bench opts out of that export.
Wiring
kv shards resolve from configs/kv_sweep.json x the registry kv_backends map (absence = off), only on --backend all dispatches, and never perturb the EP matrix (test-pinned byte equality). The sweep config maps each workload to its precisions (one workload, kv-dsv4 at fp8, its dtype mix being architectural); sweep_matrix stays stdlib-only for the bare-runner matrix/extract steps, a test pins the map to the workload model's PRESETS, and plan_config fail-closes on any mismatch at runtime. config.py emits run_kv argv behind an --entrypoint marker the rank wrapper dispatches on; launchers unchanged. summarize renders a second table with b1 and bmax columns. Tests: 119 passed / 7 skipped on the rebased tree, and the CollectiveX test job installs numpy (the cpu torch wheel does not pull it in).
The mi355x mooncake row shipped push-only from AMD's atom-dev image (a kv_backends entry can restrict ops, pin an image_ref, and set a NIC filter; upstream ionic RDMA READ is broken, which is also why ATOM's own connector is write-only), validated all green in run 31364338000 with push scaling monotonic to batch 32 and a 40.7 GB/s bulk ceiling. Follow-ups: h100/b300 rows, mi355x nixl per UCX-ionic bring-up, gb200 cross-rack once associations exist, a multi-descriptor bulk ceiling row (single-descriptor bulk underestimates multi-rail lanes), 8x8 aggregate rows, layer-streamed mode, UCCL p2p. App-side rendering shipped separately (InferenceX-app #688 merged, #689 open).