Two connectors (NIXL packed registration, SimpleCPUOffload) derived KV
geometry from untyped_storage().nbytes() and broke under the extensible
KV cache, where storages span reserved capacity. Document the contract
so out-of-tree connectors avoid the same pattern.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
The worker derived per-block sizes from storage.nbytes() // num_blocks
and viewed whole storages as (num_blocks, block_bytes). With the
extensible KV cache, registration-view storages span the reserved
capacity while num_blocks is the committed count, so block strides were
wrong and tail rows pointed into unmapped virtual memory. Derive the
per-block size from the registration views' committed extent instead
(summing a layer's state tensors for Mamba), keep the bounded-storage
size for packed layouts, and slice each segment to its committed block
prefix. Byte-identical behavior when committed == capacity.
No sleep/wake override is needed for this connector: it holds VA-stable
views plus its own pinned CPU pool (default no-op hooks are correct,
like OffloadingConnector).
Validated on GPU: cold-vs-CPU-reload greedy outputs match 5/5 with the
extensible cache (and 5/5 baseline), incl. kv_cache_memory_bytes mode.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
- Assert one packed row per logical block at allocation: the reshape
view construction, NIXL's packed registration math, and the packed
storage bounding all rely on bytes_per_block == block_stride, so make
the constraint explicit at the source instead of implicit in three
places.
- Make kv_cache_config a required argument of
narrow_kv_caches_to_num_blocks so future callers cannot silently skip
the packed storage bounding.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
Root-caused via a minimal 2-process NIXL repro on GB200: UCX transfers
succeed for VMM-backed regions mapped as a single physical allocation
but fail (remote-endpoint invalidation, NIXL_ERR_REMOTE_DISCONNECT) for
regions spanning multiple incrementally-committed cuMemCreate handles -
exactly what the 1-block -> warmup-prefix -> final commit sequence
produces. Before deferred connector registration, extend_kv_cache now
releases the warmup-time chunks and re-commits each segment prefix as
one physical allocation (contents at that point are only warmup garbage;
no requests have been served).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
NIXL 1P1D validation on GB200 showed the decode side invalidating the
prefill agent on the very first KV pull: intra-node UCX uses CUDA IPC,
and cuMemCreate allocations are only exportable to other processes when
created with requestedHandleTypes=POSIX_FILE_DESCRIPTOR, which the alloc
props did not set. When a KV connector is configured, request the POSIX
FD handle type alongside the GPU-direct-RDMA-capable flag (renamed
rdma_capable -> shareable), keeping the fallback-with-warning where such
allocations are unavailable.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
HIP mirrors the CUDA driver's VMM API (hipMemAddressReserve /
hipMemCreate / hipMemMap / hipMemSetAccess / ...) with identical call
signatures, struct layouts, and constants, so the backend only supplies
the library, symbol names, error-string convention, and implicit-context
handling; DLPack views use kDLROCM. The worker-side probe gates actual
use, so unsupported ROCm stacks still fall back gracefully.
Untested on AMD hardware so far.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
Connectors must not register KV cache memory (e.g. RDMA memory regions)
before the final size is physically committed. With the V2 runner:
- Defer ensure_kv_transfer_initialized + connector creation/registration
from initialize_from_config to extend_kv_cache, which now receives the
final (pristine, post-warmup-sizing) per-rank kv_cache_config through
the executor RPC instead of a bare block count.
- Register views narrowed along each layer's block dim to the committed
block count (narrow_kv_caches_to_num_blocks), so connectors only see
physically backed memory. Committed blocks form a prefix of each layout
segment, so a narrow covers exactly the committed bytes (e.g. NIXL's
separate K/V regions land on the two committed prefixes).
- Allocate physical chunks with the gpuDirectRDMACapable flag when a KV
connector is configured, falling back with a warning where GDR-capable
VMM allocations are unavailable.
- Warmup runs against the no-op connector (it is disabled during warmup
anyway); V1 runner + connectors + extensible remains rejected.
Validated e2e on GPU with ExampleConnector (shared-storage): deferred
registration, then a real external-cache save + hit through the narrowed
registered views with identical output.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
The VMM-backed KV cache lives outside the torch/CuMem allocators, so
sleep now discards its physical pages directly (release_physical: unmap
and release handles, keeping the VA reservation so tensor views and
captured graphs stay pointer-valid) and wake_up recommits the same block
count with freshly zeroed pages, matching the CuMem discard semantics.
Validated e2e on GPU: sleep(level=1) frees weights + KV physical memory
(0.17 GiB residual), wake_up restores and generation output is identical.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
The packed (block_stride) backing is block-major by construction: block b
occupies the b-th block_stride-byte row, holding every layer's page. Back
it with one shared single-segment ExtensibleTensor so a prefix of blocks
commits naturally, instead of rejecting the layout.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
- Extract the driver ctypes bindings into vllm/utils/vmm_driver.py behind
a small VmmDriver interface (CUDA implementation; struct layouts and
call signatures are shared with HIP for a future ROCm backend).
- Probe VMM support on the workers (driver loads, VA reservation works)
and fall back to standard KV cache allocation with a warning instead of
failing on platforms without VMM (e.g. WSL2, non-GPU workers).
- Support kv_cache_memory_bytes: the requested size is committed as-is
after warmup (single sizing pass), still avoiding warmup-time OOM.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
Move the grow-only buffer collection from the V2 attn_utils module to
vllm/utils/extensible_tensor.py and use it from the V1 runner as well
(replacing the _extensible_kv_cache_* attribute trio). Both runners now
expose the same `extensible_kv_buffers` attribute, so worker-level
features (memory measurement, sleep, connector deferral) can treat the
runners uniformly.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
Adapt #47363's extensible KV cache to the V2 model runner:
- V2 allocation (gpu/attn_utils.py): reserve each KV cache tensor's full
virtual range with ExtensibleTensor, committing a per-segment block
prefix. Segment counts are derived from each backend's physical layout
(block dim / stride order), with hybrid attention+Mamba forced
block-major to match the re-strided layout.
- V2 warmup writes to real block IDs (a contiguous prefix starting at 1),
unlike V1's all-zero dummy block tables, so warmup_kernels and
run_mixed_prefill_decode_warmup now commit exactly the block prefix
they touch via a new ensure_kv_cache_blocks() hook.
- Post-warmup measurement: instead of only the CUDA graph pool bytes,
the worker measures actual non-KV memory in use after ALL warmup
(retained worst-case activation segments, NCCL buffers, CUDA graphs)
and reports the excess over the profiling estimate
(CompilationTimes.cuda_graph renamed to warmup_memory). The engine's
second sizing pass then commits a KV cache that leaves room for the
real runtime working set - including the worst-case spec-decode
logits all-gather that memory profiling misses today.
- Gate extensible mode against KV connectors and sleep mode; drop it
from the V2-unsupported feature list.
- Extend tests/v1/worker/test_extensible_kv_cache.py with V2 coverage
(segment inference, staged prefix commits, hybrid re-stride layout).
Co-authored-by: Zhuohan Li <zhuohan123@gmail.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0133xqsNmqLHG9Pyhr5wSp1D
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
Reserve the KV cache address range with CUDA virtual memory, commit a
minimal prefix before CUDA graph capture, measure real post-capture
memory usage, then commit the final KV cache size with stable tensor
addresses. Opt-in via --enable-extensible-kv-cache.
Squashed pick of vllm-project/vllm#47363.
Co-authored-by: Zhuohan Li <zhuohan123@gmail.com>
Signed-off-by: Nick Hill <nickhill@us.ibm.com>
Nick Hill