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