[Core] Extensible (growable) KV cache

.buildkite/test_areas/disaggregated.yaml

@@ -243,6 +243,9 @@ steps:
     DECODE_BLOCK_SIZE: "256"
     MODEL_NAMES: "deepseek-ai/DeepSeek-V4-Flash"
     VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--kv-cache-dtype,fp8"
+    # The V4 indexer only supports block-outermost layouts; let it pick.
+    PREFILLER_KV_LAYOUT: ""
+    DECODER_KV_LAYOUT: ""
   commands:
     - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
     - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh

benchmarks/attention_benchmarks/runner.py

@@ -12,6 +12,7 @@
 import statistics
 import types
 from contextlib import contextmanager
+from math import prod
 
 import torch
 from batch_spec import parse_batch_spec, reorder_for_flashinfer
@@ -35,12 +36,20 @@
     VllmConfig,
     set_current_vllm_config,
 )
+from vllm.platforms import current_platform
 from vllm.v1.attention.backends.utils import (
     CommonAttentionMetadata,
-    get_kv_cache_layout,
-    set_kv_cache_layout,
+    get_supported_kv_cache_layouts,
+    resolve_kv_cache_layout,
+)
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheLayout,
+    KVCacheTensor,
+    compute_layer_kv_cache_shape_bytes,
+    compute_layout_strides,
+    create_kv_cache_views,
 )
-from vllm.v1.kv_cache_interface import FullAttentionSpec
 
 # ============================================================================
 # Backend Configuration
@@ -340,49 +349,42 @@ def _create_kv_cache(
     backend_class,
     device: torch.device,
     dtype: torch.dtype,
+    layout: KVCacheLayout,
 ) -> list:
-    """Create KV cache tensors for all layers using the backend's methods.
-
-    Uses the backend's get_kv_cache_shape() and get_kv_cache_stride_order()
-    to create the cache with the correct shape and memory layout.
-    """
-    # Get the logical shape from the backend
-    cache_shape = backend_class.get_kv_cache_shape(
-        num_blocks=max_num_blocks,
+    """Create KV cache tensors for all layers using the standard allocator."""
+    if config.kv_cache_dtype.startswith("fp8"):
+        cache_dtype = current_platform.fp8_dtype()
+    else:
+        cache_dtype = dtype
+    spec = FullAttentionSpec(
         block_size=config.block_size,
         num_kv_heads=config.num_kv_heads,
         head_size=config.head_dim,
+        dtype=cache_dtype,
+    )
+    # Apply the backend's page customization, as the worker does for the real spec.
+    spec = backend_class.customize_spec(spec)
+    total_bytes = (
+        prod(compute_layer_kv_cache_shape_bytes(spec, max_num_blocks))
+        * config.num_layers
+    )
+    buf = torch.zeros(total_bytes, device=device, dtype=torch.int8)
+    layer_stride, block_stride, _, _, _ = compute_layout_strides(
+        spec, max_num_blocks, config.num_layers, layout
+    )
+    tensor = KVCacheTensor(
+        size=total_bytes,
+        layers=[str(i) for i in range(config.num_layers)],
+        layer_stride=layer_stride,
+        block_stride=block_stride,
+    )
+    return create_kv_cache_views(
+        buf,
+        spec,
+        max_num_blocks,
+        layout,
+        tensor,
     )
-
-    # Get the stride order for custom memory layout
-    try:
-        stride_order = backend_class.get_kv_cache_stride_order()
-        assert len(stride_order) == len(cache_shape)
-    except (AttributeError, NotImplementedError):
-        stride_order = tuple(range(len(cache_shape)))
-
-    # Permute shape to physical layout order
-    physical_shape = tuple(cache_shape[i] for i in stride_order)
-
-    # Compute inverse permutation to get back to logical view
-    inv_order = [stride_order.index(i) for i in range(len(stride_order))]
-
-    # Use fp8 dtype for cache when requested.
-    cache_dtype = dtype
-    if config.kv_cache_dtype == "fp8":
-        from vllm.platforms import current_platform
-
-        cache_dtype = current_platform.fp8_dtype()
-
-    cache_list = []
-    for _ in range(config.num_layers):
-        # Allocate in physical layout order (contiguous in memory)
-        cache = torch.zeros(*physical_shape, device=device, dtype=cache_dtype)
-        # Permute to logical view
-        cache = cache.permute(*inv_order)
-        cache_list.append(cache)
-
-    return cache_list
 
 
 # ============================================================================
@@ -499,13 +501,12 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
             backend_class, impl, layer = _create_backend_impl(
                 backend_cfg, config, device, dtype
             )
-
             # Set KV cache layout if the backend requires a specific one
-            # (e.g., FlashInfer requires HND on SM100/Blackwell for TRTLLM attention)
-            required_layout = backend_class.get_required_kv_cache_layout()
-            if required_layout is not None:
-                set_kv_cache_layout(required_layout)
-                get_kv_cache_layout.cache_clear()
+            # (e.g., FlashInfer requires LBHNC on SM100/Blackwell for TRTLLM attention)
+            supported = get_supported_kv_cache_layouts([backend_class])
+            layout = resolve_kv_cache_layout(
+                vllm_config.cache_config, [[m.name for m in supported]]
+            )
 
             common_metadata = _build_common_attn_metadata(
                 q_lens, kv_lens, config.block_size, device
@@ -542,7 +543,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
             )
 
             cache_list = _create_kv_cache(
-                config, max_num_blocks, backend_class, device, dtype
+                config, max_num_blocks, backend_class, device, dtype, layout
             )
 
             timing_stats, mem_stats = _run_single_benchmark(

docs/features/mooncake_store_connector_usage.md

@@ -227,7 +227,6 @@ Strict isolation requires a Mooncake master started with `--enable_multi_tenants
 
 - `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`.
 - `lookup_async` (bool): Run the external prefix-cache lookup on a background thread so it never blocks the scheduler step. The request is held until the in-flight lookup completes, then resumed on a later step. Default: `false`.
-- `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`.
 - `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.
 - `cache_prefix` (str): Namespace prepended to every store key. Lets separate deployments share one Mooncake master without polluting each other — instances configured with different prefixes never see each other's cached blocks, even for identical prompts. All instances that should share a prefix cache must use the same value. Default: `""` (no prefix; keys are byte-identical to the unprefixed format).
 

docs/features/nixl_connector_compatibility.md

@@ -61,7 +61,7 @@ th:not(:first-child) {
 See [Configuration Notes](#configuration-notes) below for what must match and
 what may differ.
 
-<sup>2</sup> Requires `FLASH_ATTN` or `FLASHINFER` backend **and** `HND` KV cache layout. Enable via `--kv-transfer-config '{"kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}'`.
+<sup>2</sup> Cross-layer contiguity is achieved by using a `BLHNC` layout (set via `VLLM_KV_CACHE_LAYOUT=BLHNC`).
 
 <sup>3</sup> Supported only when HMA is **not** required (i.e., non-hybrid models). Block IDs are remapped automatically. Only P block size < D block size is supported.
 
@@ -100,9 +100,9 @@ By default, a **compatibility hash** is checked during handshake. P and D instan
 
 ### KV cache layout
 
-- NixlConnector defaults to **`HND`** layout for optimal transfer performance (non-MLA models).
-- `NHD` layout is supported but does **not** allow heterogeneous TP head splitting.
-- Experimental `HND` ↔ `NHD` permute: enable via `--kv-transfer-config '{"enable_permute_local_kv": true}'`. Not supported with HMA.
+- NixlConnector defaults to **`LBHNC`** (head-major, formerly `HND`) layout for optimal transfer performance (non-MLA models).
+- `LBNHC` (token-major, formerly `NHD`) layout is supported but does **not** allow heterogeneous TP head splitting.
+- Experimental `LBHNC` ↔ `LBNHC` permute: enable via `--kv-transfer-config '{"enable_permute_local_kv": true}'`. Not supported with HMA.
 
 ### Quantized KV cache
 

docs/features/nixl_connector_usage.md

@@ -404,21 +404,12 @@ GB-series GPUs support multi-node NVLink. NIXL supports this capability, but KVC
 
 ### Heterogeneous KV Layout support
 
-Support use case: Prefill with 'HND' and decode with 'NHD' with experimental configuration
+Support use case: Prefill with `LBHNC` and decode with `LBNHC` with experimental configuration
 
 ```bash
 --kv-transfer-config '{..., "enable_permute_local_kv":"True"}'
 ```
 
-### Cross layers blocks
-
-By default, this feature is disabled. On attention backends that support this feature, each logical block is contiguous in physical memory. This reduces the number of buffers that need to be transferred.
-To enable this feature:
-
-```bash
---kv-transfer-config '{..., "kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}'
-```
-
 ## Metrics Reference
 
 vLLM periodically logs a `KV Transfer metrics` line summarising NIXL transfer

tests/basic_correctness/test_basic_correctness.py

@@ -164,7 +164,6 @@ def test_models(
         max_model_len=8192,
         enforce_eager=enforce_eager,
         enable_prompt_embeds=enable_prompt_embeds,
-        gpu_memory_utilization=0.7,
         async_scheduling=async_scheduling,
         distributed_executor_backend=model_executor,
         attention_config={"backend": backend},
@@ -254,8 +253,9 @@ def test_models_distributed(
             tensor_parallel_size=2,
             distributed_executor_backend=distributed_executor_backend,
             enable_prompt_embeds=enable_prompt_embeds,
-            gpu_memory_utilization=0.7,
             attention_config=attention_config,
+            # The HF reference model loads below while this engine is alive.
+            gpu_memory_utilization=0.7,
         ) as vllm_model:
             if enable_prompt_embeds:
                 with hf_runner(model, dtype=dtype) as hf_model:

tests/compile/fullgraph/test_basic_correctness.py

@@ -46,8 +46,6 @@ class TestSetting:
                 "bfloat16",
                 "--max-model-len",
                 "2048",
-                "--gpu-memory-utilization",
-                "0.98",
             ],
             pp_size=1,
             tp_size=1,

tests/compile/fullgraph/test_full_cudagraph.py

@@ -83,20 +83,20 @@ def llm_pair(request):
     with temporary_environ(env_vars):
         full = LLM(
             model=model,
-            gpu_memory_utilization=0.43,
             trust_remote_code=True,
             max_model_len=1024,
             max_num_seqs=128,
+            gpu_memory_utilization=0.43,  # both engines alive concurrently
             compilation_config=CompilationConfig(**backend_config.comp_config),
             generation_config="vllm",
             seed=42,
         )
         piecewise = LLM(
             model=model,
-            gpu_memory_utilization=0.43,
             trust_remote_code=True,
             max_model_len=1024,
             max_num_seqs=128,
+            gpu_memory_utilization=0.43,  # both engines alive concurrently
             compilation_config=CompilationConfig(
                 cudagraph_mode=CUDAGraphMode.PIECEWISE
             ),

tests/compile/fullgraph/test_multimodal_compile.py

@@ -38,7 +38,6 @@ def test_qwen2_5_vl_compilation(vllm_runner, monkeypatch):
         vllm_runner(
             "Qwen/Qwen2.5-VL-3B-Instruct",
             max_model_len=2048,
-            gpu_memory_utilization=0.8,
             compilation_config={
                 "mode": CompilationMode.VLLM_COMPILE,
                 "compile_mm_encoder": True,
@@ -63,7 +62,6 @@ def test_qwen2_5_vl_no_vit_compilation(vllm_runner, monkeypatch):
         vllm_runner(
             "Qwen/Qwen2.5-VL-3B-Instruct",
             max_model_len=2048,
-            gpu_memory_utilization=0.8,
             compilation_config={
                 "mode": CompilationMode.VLLM_COMPILE,
                 "compile_mm_encoder": False,
@@ -99,7 +97,6 @@ def test_mllama4_vit_compilation(vllm_runner, monkeypatch):
         vllm_runner(
             "meta-llama/Llama-4-Scout-17B-16E-Instruct",
             max_model_len=512,
-            gpu_memory_utilization=0.8,
             tensor_parallel_size=8,
             compilation_config={
                 "mode": CompilationMode.VLLM_COMPILE,

tests/compile/passes/test_fusion_attn.py

@@ -7,7 +7,11 @@
 
 from tests.compile.backend import LazyInitPass, TestBackend
 from tests.utils import TestFP8Layer, flat_product
-from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
+from tests.v1.attention.utils import (
+    BatchSpec,
+    create_common_attn_metadata,
+    dense_kv_cache_views,
+)
 from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
 from vllm.compilation.passes.fusion.attn_quant_fusion import (
     ATTN_OP,
@@ -39,7 +43,14 @@
 from vllm.utils.flashinfer import has_flashinfer
 from vllm.v1.attention.backend import AttentionMetadata
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.v1.kv_cache_interface import AttentionSpec, get_kv_quant_mode
+from vllm.v1.attention.backends.utils import (
+    get_supported_kv_cache_layouts,
+    resolve_kv_cache_layout,
+)
+from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
+    get_kv_quant_mode,
+)
 
 DEVICE_TYPE = current_platform.device_type
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -108,32 +119,29 @@ def build_attn_metadata(self, batch_size: int) -> AttentionMetadata:
         max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
         num_blocks = batch_size * max_blocks
 
-        # Fetch the attention backend and kv cache shape and stride order
-        attn_backend = self.attn.attn_backend
-        kv_cache_shape = attn_backend.get_kv_cache_shape(
-            num_blocks,
-            self.block_size,
-            self.num_kv_heads,
-            self.head_size,
-            cache_dtype_str=self.attn.kv_cache_dtype,
+        spec = AttentionSpec(
+            block_size=self.block_size,
+            num_kv_heads=self.num_kv_heads,
+            head_size=self.head_size,
+            dtype=self.attn.kv_cache_torch_dtype,
+            kv_quant_mode=get_kv_quant_mode(self.attn.kv_cache_dtype),
+        )
+        supported = get_supported_kv_cache_layouts([self.attn.attn_backend])
+        layout = resolve_kv_cache_layout(
+            self.vllm_config.cache_config, [[m.name for m in supported]]
         )
-        try:
-            kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
-        except (AttributeError, NotImplementedError):
-            kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
-
-        kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
-        inv_order = [
-            kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
-        ]
-
-        # Create dummy KV cache
         raw_tensor = torch.zeros(
-            kv_cache_shape,
-            dtype=self.attn.kv_cache_torch_dtype,
+            num_blocks * spec.page_size_bytes,
+            dtype=torch.int8,
             device=self.device,
         )
-        kv_cache = raw_tensor.permute(*inv_order)
+        kv_cache = dense_kv_cache_views(
+            raw_tensor,
+            spec,
+            num_blocks,
+            num_layers=1,
+            layout=layout,
+        )[0]
 
         self.attn.kv_cache = kv_cache
 

tests/compile/passes/test_mla_attn_quant_fusion.py

@@ -150,27 +150,14 @@ def build_attn_metadata(self, batch_size: int) -> AttentionMetadata:
         max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
         num_blocks = batch_size * max_blocks
 
-        # MLA KV cache is 3D: (num_blocks, block_size, head_size)
-        attn_backend = self.mla_attn.attn_backend
-        kv_cache_shape = attn_backend.get_kv_cache_shape(
-            num_blocks, self.block_size, 1, self.head_size
-        )
-        try:
-            kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
-        except (AttributeError, NotImplementedError):
-            kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
-
-        ordered_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
-        inv_order = [
-            kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
-        ]
-
-        raw_tensor = torch.zeros(
-            ordered_shape, dtype=self.kv_cache_dtype, device=self.device
+        # MLA KV cache is 4D: (num_blocks, num_heads=1, block_size, head_size)
+        kv_cache = torch.zeros(
+            (num_blocks, 1, self.block_size, self.head_size),
+            dtype=self.kv_cache_dtype,
+            device=self.device,
         )
-        kv_cache = raw_tensor.permute(*inv_order)
 
-        self.mla_attn.kv_cache = kv_cache
+        self.mla_attn.bind_kv_cache(kv_cache)
 
         self.attn_metadata = self.builder.build(
             common_prefix_len=0, common_attn_metadata=common_attn_metadata