[6/N][KV-Cache Layout Refactor] Standardize KV cache layout

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,
+    publish_kv_cache_layout_to_current_process,
+    resolve_kv_cache_layout,
+)
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    compute_layer_kv_cache_shape_bytes,
+    compute_layout_strides,
+    reshape_kv_cache,
 )
-from vllm.v1.kv_cache_interface import FullAttentionSpec
 
 # ============================================================================
 # Backend Configuration
@@ -341,48 +350,38 @@ def _create_kv_cache(
     device: torch.device,
     dtype: torch.dtype,
 ) -> 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)
+    layout = get_kv_cache_layout()
+    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
+    )
+    return reshape_kv_cache(
+        buf,
+        spec,
+        max_num_blocks,
+        config.num_layers,
+        layout,
+        offset=0,
+        layer_stride=layer_stride,
+        block_stride=block_stride,
     )
-
-    # 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 +498,13 @@ 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([[m.name for m in supported]])
+            publish_kv_cache_layout_to_current_process(
+                layout.name, vllm_config.cache_config
+            )
 
             common_metadata = _build_common_attn_metadata(
                 q_lens, kv_lens, config.block_size, device

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/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,15 @@
 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,
+    publish_kv_cache_layout_to_current_process,
+    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 +120,30 @@ 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([[m.name for m in supported]])
+        publish_kv_cache_layout_to_current_process(
+            layout.name, self.vllm_config.cache_config
         )
-        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

tests/compile/passes/test_mla_rope_kvcache_cat_fusion.py

@@ -165,29 +165,14 @@ def build_attn_metadata(self, batch_size: int) -> CommonAttentionMetadata:
         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
-        kv_cache_shape = self.attn_backend.get_kv_cache_shape(
-            num_blocks, self.block_size, self.num_kv_heads, self.head_size
-        )
-        try:
-            kv_cache_stride_order = self.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))
-        ]
-
-        raw_tensor = torch.zeros(
-            num_blocks * self.block_size * self.num_kv_heads * self.head_size,
+        # 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,
         )
-        raw_tensor = raw_tensor.view(kv_cache_shape)
-        kv_cache = raw_tensor.permute(*inv_order)
 
-        self.mla_attn.kv_cache = kv_cache
+        self.mla_attn.bind_kv_cache(kv_cache)
 
         # Build attn metadata
         attn_metadata = self.builder.build(