[Core] Extensible (growable) KV cache

[njhill]

Stacked on #44458 (lwilkinson/kv-layout/core-standardize). Only the commits above that PR's head belong to this one; review #44458 first.
This is a draft — it cannot merge until #44458 lands, after which the base will be retargeted to main.

Summary

Opt-in growable KV cache (--enable-extensible-kv-cache) built on the standardized KV cache layout from #44458.

vLLM sizes the KV cache from a profiling run, then allocates it up front. Any memory the engine consumes after profiling — CUDA graph pools, the spec-decode logits all-gather, workspace growth — has to be guessed at via
headroom margins, and a wrong guess is either an OOM or wasted HBM.

This change reserves virtual address space for the KV cache capacity but commits physical pages incrementally (CUDA cuMem* / ROCm hipMem*):

1. Allocate the VA reservation, commit one block per layout segment.
2. Run warmup and CUDA graph capture.
3. Measure actual post-warmup free memory and re-size.
4. extend_kv_cache() maps the remaining pages under the same base pointer.

Because each block keeps a fixed offset within its layout segment and the base pointer never moves, captured graphs and layer views stay valid across the growth — no re-view, no re-capture.

What #44458 buys this

The layout standardization removes most of what made the earlier version of this work invasive. Buffer segmentation is now a pure function of KVCacheLayout.stride_order (new num_outer_segments() helper) instead of per-backend shape/stride/block-dim probing; connector registration views are a plain narrow(0) of the logical [B, H, N, C] view; and the packed / block_stride special cases disappear entirely. The integration is ~30%
smaller than the equivalent against main.

Views are authoritative, not storage.nbytes()

The one contract change worth reviewer attention. Under an extensible cache the untyped storage spans the reserved capacity, of which only each view's per-segment block prefix is physically backed. Any code deriving block geometry or registration extents from untyped_storage().nbytes() will compute a size that includes unbacked pages.

Updated accordingly, with the contract documented on KVConnectorBase_V1.register_kv_caches:

• NixlBaseConnectorWorker.register_kv_caches — collects per-region committed spans and registers those; falls back to the single per-allocation registration when the storage is fully backed. Block-major packing is
  detected from view strides rather than nbytes().
• get_kv_cache_block_regions (v1/worker/utils.py) — same, which also covers the CPU-offload connector.

KV-transfer init is deferred until after extend_kv_cache, so connectors only ever register committed memory. On the NIXL path, extend also defragments (release + single-chunk recommit per segment): UCX cannot transfer a VMM region spanning multiple cuMemCreate handles.

Relationship to existing work

• [Core] Demo implementation of extensible kv cache memory #47363 (@zhuohan123, "Demo implementation of extensible kv cache memory) — this is the productionization of that demo, not a competing implementation. It adds the V1 + V2 runner integration, the measured post-warmup sizing pass, VMM probe/fallback, sleep-mode support, KV connector support, and a ROCm backend. Happy to fold this into [Core] Demo implementation of extensible kv cache memory #47363 instead if that's preferred.
• [4/N][KV-Cache Layout Refactor] Standardize KV cache layout #44458 — strictly a dependency; no overlap in changed behavior.
• Searched open PRs for extensible/growable/VMM KV cache work; nothing else overlaps.

Contributors: @zhuohan123 (original demo), @ZJY0516 (packed-storage bounding, sleep gate).

Testing

Unless noted, run against this branch at its current head.

Unit — pytest tests/v1/worker/test_extensible_kv_cache.py tests/utils_/test_extensible_tensor.py tests/v1/engine/test_engine_args.py -v → 29/29 pass (GB200, aarch64).

e2e, 1×GB200, CUDA graphs on — outputs byte-identical to a non-extensible baseline for: standard generation, sleep(level=1)/wake, and --kv-cache-memory-bytes.

e2e, sizing under pressure — Qwen3.5-4B TP4, MTP num_speculative_tokens=100, chunk 256, gpu-memory-utilization=0.92 (a config that OOMs without this change): passes in both eager and cudagraph modes. Sizing 35,277→34,721 tokens (eager), 35,243→34,592 (cudagraph). Functional check 64/43 answers correct.

e2e, NIXL P/D disaggregation — Qwen3-0.6B, HND, block 128, gsm8k via P/D proxy:

config	gsm8k	expected
1P1D eager	0.422	~0.41
1P1D cudagraph	0.418	~0.41
P(TP1)→D(TP2)	0.415	~0.41
P(TP2)→D(TP2)	0.415	~0.41

Model evaluation — the feature is opt-in and does not change model outputs; the byte-identical-output comparison above and the gsm8k table are the evidence for that.

Known gaps

• The V2 encoder-cache profiling reservation (last commit) is validated on the equivalent branch against main (Qwen2-VL-2B: 16,384-token budget reserved, 238,148 blocks committed) but has not been re-run on this branch — the cluster battery above has no multimodal configuration. Will re-run before marking ready.
• No unit coverage for commit(defragment=True) or the connector sleep gate; both are exercised only by the e2e NIXL runs. Adding these before ready.
• Multi-node NIXL (NIC RDMA path with the GDR + POSIX-FD allocation flags) is unvalidated; the repro harness is localhost-only.
• The ROCm hipMem* backend is written but untested on AMD hardware.
• Encoder-decoder mixed layouts raise rather than falling back.

AI assistance

AI assistance (Claude Code) was used for this change: for the port onto the #44458 layout model, the connector geometry audit, and drafting. I have reviewed every changed line, and ran the tests and cluster validation above
myself.