[Core] Extensible (growable) KV cache

tests/utils_/test_extensible_tensor.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.utils.extensible_tensor import ExtensibleTensor
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+
+
+def test_extensible_tensor_grows_without_moving() -> None:
+    buffer = ExtensibleTensor(4096, device="cuda")
+    try:
+        base_ptr = buffer.base_ptr
+        first_view = buffer.resize_(1024)
+        assert first_view.data_ptr() == base_ptr
+        first_view.fill_(7)
+
+        second_view = buffer.resize_(2048)
+        assert second_view.data_ptr() == base_ptr
+        assert torch.equal(second_view[:1024], torch.full_like(second_view[:1024], 7))
+
+        second_view[1024:].fill_(3)
+        assert torch.equal(buffer.tensor, second_view)
+
+        full_view = buffer.full_view()
+        assert full_view.data_ptr() == base_ptr
+        assert full_view.numel() == 4096
+    finally:
+        buffer.free()
+
+
+def test_extensible_tensor_rejects_shrink_and_overflow() -> None:
+    buffer = ExtensibleTensor(1024, device="cuda")
+    try:
+        buffer.resize_(512)
+        with pytest.raises(ValueError, match="grow-only"):
+            buffer.resize_(256)
+        with pytest.raises(ValueError, match="exceeds the segment capacity"):
+            buffer.resize_(1025)
+    finally:
+        buffer.free()
+
+
+def test_segments_grow_in_lockstep_and_zero_new() -> None:
+    """Each segment's committed prefix grows in lockstep.
+
+    Data written to a segment's committed prefix survives a grow; the newly
+    committed range of each segment is zeroed with `zero_new=True` while old
+    bytes are preserved.
+    """
+    et = ExtensibleTensor(max_num_bytes=8192, device="cuda", num_segments=2)
+    try:
+        assert et.num_segments == 2
+        assert et.segment_capacity_bytes == 4096
+
+        et.resize_per_segment_(256, zero_new=True)
+        assert et.bytes_per_segment == 256
+        assert et.num_bytes == 512
+        fv = et.full_view()
+        assert fv.shape == (8192,)
+        # Committed prefixes start zeroed.
+        assert torch.count_nonzero(fv[:256]) == 0
+        assert torch.count_nonzero(fv[4096 : 4096 + 256]) == 0
+
+        pattern_a = torch.arange(256, device="cuda", dtype=torch.uint8)
+        pattern_b = 255 - pattern_a
+        fv[:256].copy_(pattern_a)
+        fv[4096 : 4096 + 256].copy_(pattern_b)
+
+        et.resize_per_segment_(1024, zero_new=True)
+        fv2 = et.full_view()
+        assert fv2.data_ptr() == fv.data_ptr()
+        # Old bytes of both segments preserved; freshly committed ranges zeroed.
+        assert torch.equal(fv2[:256], pattern_a)
+        assert torch.equal(fv2[4096 : 4096 + 256], pattern_b)
+        assert torch.count_nonzero(fv2[256:1024]) == 0
+        assert torch.count_nonzero(fv2[4096 + 256 : 4096 + 1024]) == 0
+    finally:
+        et.free()
+
+
+def test_segments_at_granularity_scale() -> None:
+    """Segments spanning multiple mapping granules commit correctly.
+
+    Uses a segment capacity that is not a multiple of the allocation
+    granularity, so a granule straddles the segment boundary and is shared by
+    the first commit of one segment and a later commit of the other -- it must
+    be mapped exactly once.
+    """
+    probe = ExtensibleTensor(max_num_bytes=1, device="cuda")
+    granularity = probe.capacity_bytes
+    probe.free()
+    # Two segments of 1.5 granules each; the middle granule straddles the
+    # boundary.
+    max_num_bytes = 3 * granularity
+    et = ExtensibleTensor(max_num_bytes=max_num_bytes, device="cuda", num_segments=2)
+    try:
+        seg = et.segment_capacity_bytes
+        assert seg == max_num_bytes // 2
+
+        step = granularity // 2
+        et.resize_per_segment_(step, zero_new=True)
+        fv = et.full_view()
+        fv[:step].fill_(1)
+        fv[seg : seg + step].fill_(2)
+
+        # Grow to the full segment capacity: previously mapped granules
+        # (including the boundary-straddling one) are reused, new ones are
+        # committed and zeroed.
+        et.resize_per_segment_(seg, zero_new=True)
+        fv2 = et.full_view()
+        assert torch.all(fv2[:step] == 1)
+        assert torch.all(fv2[seg : seg + step] == 2)
+        assert torch.count_nonzero(fv2[step:seg]) == 0
+        assert torch.count_nonzero(fv2[seg + step :]) == 0
+    finally:
+        et.free()
+
+
+def test_multi_segment_invalid_usage_raises() -> None:
+    """Prefix-view APIs and invalid segment configs raise for multi-segment
+    buffers."""
+    with pytest.raises(ValueError):
+        ExtensibleTensor(max_num_bytes=100, device="cuda", num_segments=3)
+
+    et = ExtensibleTensor(max_num_bytes=8192, device="cuda", num_segments=2)
+    try:
+        with pytest.raises(ValueError):
+            _ = et.tensor
+        with pytest.raises(ValueError):
+            et.resize_(256)
+
+        et.resize_per_segment_(256)
+        with pytest.raises(ValueError):
+            et.resize_per_segment_(128)  # shrink
+        with pytest.raises(ValueError):
+            et.resize_per_segment_(et.segment_capacity_bytes + 1)  # over capacity
+    finally:
+        et.free()

tests/v1/engine/test_engine_args.py

@@ -49,6 +49,18 @@ def test_prefix_caching_from_cli():
         args = parser.parse_args(["--prefix-caching-hash-algo", "invalid"])
 
 
+def test_extensible_kv_cache_from_cli():
+    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
+
+    args = parser.parse_args([])
+    engine_args = EngineArgs.from_cli_args(args=args)
+    assert not engine_args.enable_extensible_kv_cache
+
+    args = parser.parse_args(["--enable-extensible-kv-cache"])
+    engine_args = EngineArgs.from_cli_args(args=args)
+    assert engine_args.enable_extensible_kv_cache
+
+
 @pytest.mark.skipif(_xxhash is None, reason="xxhash not installed")
 def test_prefix_caching_xxhash_from_cli():
     parser = EngineArgs.add_cli_args(FlexibleArgumentParser())

tests/v1/worker/test_extensible_kv_cache.py

@@ -0,0 +1,202 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Extensible KV cache over the standardized KV cache layout.
+
+Covers the layout-derived buffer segmentation, the committed-prefix
+narrowing used for KV-connector registration views, and VMM-backed
+allocation through the shared allocate+reshape path.
+"""
+
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+from vllm.utils.extensible_tensor import ExtensibleKVCacheBuffers
+from vllm.utils.vmm_driver import get_vmm_driver, vmm_unavailable_reason
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheLayout,
+    KVCacheTensor,
+    MambaSpec,
+    num_outer_segments,
+)
+from vllm.v1.worker.gpu.attn_utils import (
+    _allocate_and_reshape_kv_cache,
+    narrow_kv_caches_to_num_blocks,
+)
+
+BLOCK_SIZE = 16
+NUM_BLOCKS = 8
+NUM_HEADS = 2
+HEAD_SIZE = 32
+
+requires_vmm = pytest.mark.skipif(
+    vmm_unavailable_reason() is not None,
+    reason=f"VMM unavailable: {vmm_unavailable_reason()}",
+)
+
+
+def _attn_spec(**kwargs) -> FullAttentionSpec:
+    return FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_HEADS,
+        head_size=HEAD_SIZE,
+        dtype=torch.float16,
+        **kwargs,
+    )
+
+
+@pytest.mark.parametrize(
+    "layout,num_layer_slots,expected",
+    [
+        (KVCacheLayout.LBHNC, 1, 1),
+        (KVCacheLayout.LBHNC, 4, 4),
+        (KVCacheLayout.LBNHC, 4, 4),
+        (KVCacheLayout.BLHNC, 4, 1),
+        (KVCacheLayout.BLNHC, 4, 1),
+        (KVCacheLayout.BHLNC, 4, 1),
+    ],
+)
+def test_num_outer_segments_layouts(layout, num_layer_slots, expected):
+    """Segments = product of physical dims outer to the block dim."""
+    assert num_outer_segments(_attn_spec(), num_layer_slots, layout) == expected
+
+
+def test_num_outer_segments_separate_kv_head_groups():
+    """NVFP4-style layouts place head groups outside the block dim."""
+    spec = _attn_spec(separate_kv_head_groups=True)
+    assert num_outer_segments(spec, 3, KVCacheLayout.LBHNC) == 3 * spec.num_heads
+
+
+def test_num_outer_segments_mamba():
+    spec = MambaSpec(
+        block_size=BLOCK_SIZE,
+        shapes=((NUM_HEADS, HEAD_SIZE),),
+        dtypes=(torch.float32,),
+    )
+    assert num_outer_segments(spec, 2, KVCacheLayout.LBHNC) == 2
+    assert num_outer_segments(spec, 2, KVCacheLayout.BLHNC) == 1
+
+
+def _single_group_config(spec, num_layer_slots: int = 1) -> KVCacheConfig:
+    layer_names = [f"layer.{i}" for i in range(num_layer_slots)]
+    return KVCacheConfig(
+        num_blocks=NUM_BLOCKS,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=spec.page_size_bytes * NUM_BLOCKS * num_layer_slots,
+                shared_by=[[name] for name in layer_names],
+            )
+        ],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names, spec)],
+    )
+
+
+@dataclass
+class _FakeGroup:
+    kv_cache_spec: object
+    layer_names: list[str]
+    kv_cache_group_id: int = 0
+
+
+@requires_vmm
+@pytest.mark.parametrize("layout", [KVCacheLayout.LBHNC, KVCacheLayout.BLHNC])
+def test_extensible_allocation_and_growth(layout):
+    """Only one block committed at first; extend keeps base pointers."""
+    spec = _attn_spec()
+    config = _single_group_config(spec, num_layer_slots=2)
+    kv_caches, buffers = _allocate_and_reshape_kv_cache(
+        config,
+        torch.device("cuda:0"),
+        layout=layout,
+        extensible=True,
+    )
+    assert isinstance(buffers, ExtensibleKVCacheBuffers)
+    assert buffers.num_blocks_committed == 1
+    # Physical commit is granule-rounded per segment: one committed block
+    # maps at most one granule in each segment.
+    granule = get_vmm_driver().granularity(0)
+    num_segments = 2 if layout is KVCacheLayout.LBHNC else 1
+    assert 0 < buffers.physical_bytes <= num_segments * granule
+
+    ptrs_before = {n: t.data_ptr() for n, t in kv_caches.items()}
+    for name, view in kv_caches.items():
+        # Views span the full declared capacity.
+        assert view.shape[0] == NUM_BLOCKS
+        # Committed prefix is writable.
+        view[0].fill_(1.0)
+
+    buffers.commit(NUM_BLOCKS)
+    assert buffers.num_blocks_committed == NUM_BLOCKS
+    for name, view in kv_caches.items():
+        assert view.data_ptr() == ptrs_before[name]
+        view[NUM_BLOCKS - 1].fill_(2.0)
+        # Earlier contents survive the grow; new blocks were zeroed.
+        assert view[0].eq(1.0).all()
+    torch.cuda.synchronize()
+    buffers.free()
+
+
+@requires_vmm
+def test_extensible_release_and_recommit():
+    """Sleep-style release keeps VA and views valid; recommit re-zeroes."""
+    spec = _attn_spec()
+    config = _single_group_config(spec)
+    kv_caches, buffers = _allocate_and_reshape_kv_cache(
+        config,
+        torch.device("cuda:0"),
+        layout=KVCacheLayout.LBHNC,
+        extensible=True,
+    )
+    buffers.commit(NUM_BLOCKS)
+    view = kv_caches["layer.0"]
+    view.fill_(3.0)
+    torch.cuda.synchronize()
+
+    buffers.release_physical()
+    assert buffers.num_blocks_committed == 0
+    buffers.recommit()
+    assert buffers.num_blocks_committed == NUM_BLOCKS
+    assert view.data_ptr() == kv_caches["layer.0"].data_ptr()
+    assert view.eq(0).all()
+    torch.cuda.synchronize()
+    buffers.free()
+
+
+def test_narrow_kv_caches_to_num_blocks():
+    """Connector views are trimmed to the committed logical block prefix."""
+    spec = _attn_spec()
+    committed = 3
+    full = torch.zeros(NUM_BLOCKS, NUM_HEADS, BLOCK_SIZE, 2 * HEAD_SIZE)
+    groups = [_FakeGroup(spec, ["layer.0"])]
+    narrowed = narrow_kv_caches_to_num_blocks(
+        {"layer.0": full}, groups, [BLOCK_SIZE], committed
+    )
+    assert narrowed["layer.0"].shape[0] == committed
+    assert narrowed["layer.0"].data_ptr() == full.data_ptr()
+
+
+def test_narrow_applies_virtual_block_split():
+    """Kernel-split caches have block_size/kernel ratio more physical blocks."""
+    spec = _attn_spec()
+    kernel_block_size = BLOCK_SIZE // 2
+    physical_blocks = NUM_BLOCKS * 2
+    full = torch.zeros(physical_blocks, NUM_HEADS, kernel_block_size, 2 * HEAD_SIZE)
+    groups = [_FakeGroup(spec, ["layer.0"])]
+    narrowed = narrow_kv_caches_to_num_blocks(
+        {"layer.0": full}, groups, [kernel_block_size], 3
+    )
+    assert narrowed["layer.0"].shape[0] == 3 * 2
+
+
+def test_narrow_skips_out_of_range_groups():
+    spec = _attn_spec()
+    full = torch.zeros(NUM_BLOCKS, NUM_HEADS, BLOCK_SIZE, 2 * HEAD_SIZE)
+    groups = [_FakeGroup(spec, ["layer.0"], kv_cache_group_id=1)]
+    narrowed = narrow_kv_caches_to_num_blocks(
+        {"layer.0": full}, groups, [BLOCK_SIZE], 3
+    )
+    assert narrowed["layer.0"] is full

vllm/config/cache.py

@@ -188,6 +188,18 @@ class CacheConfig:
     gpu_memory_utilization. Note that kv_cache_memory_bytes
     (when not-None) ignores gpu_memory_utilization"""
 
+    enable_extensible_kv_cache: bool = False
+    """Use driver virtual memory to reserve the KV cache address range up
+    front, run warmup and CUDA graph capture with only a small block prefix
+    physically committed, and commit the final size afterwards.
+
+    This makes automatic KV sizing account for the memory that warmup and
+    CUDA graph capture actually consume (including worst-case activation
+    working sets, e.g. with speculative decoding), and avoids warmup-time
+    OOMs. Requires driver VMM support (CUDA or ROCm; falls back to standard
+    allocation with a warning where unavailable, e.g. WSL2).
+    """
+
     kv_offloading_size: float | None = None
     """Size of the KV cache offloading buffer in GiB. When TP > 1, this is
     the total buffer size summed across all TP ranks. By default, this is set
@@ -233,6 +245,8 @@ def compute_hash(self) -> str:
             "kv_cache_max_concurrency",
             # WIP feature toggle not impacting compiled graph shape
             "kv_sharing_fast_prefill",
+            # Runtime memory allocation strategy, not graph structure.
+            "enable_extensible_kv_cache",
         }
 
         from vllm.config.utils import get_hash_factors, hash_factors