[6/N][KV-Cache Layout Refactor] Standardize KV cache layout (#51718)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

.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 logging
 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 @@ from vllm.config import (
     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,10 @@ 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, [[m.name for m in supported]])
 
             common_metadata = _build_common_attn_metadata(
                 q_lens, kv_lens, config.block_size, device
@@ -542,7 +541,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

@@ -246,7 +246,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).
 - `save_decode_cache` (bool): Enable offloading decode tokens' KV cache. A `kv_consumer` does not save during prefill; when decode starts, it fills any missing block-aligned prompt prefix before appending completed decode blocks. Default: `false`.

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

rust/Cargo.lock

@@ -2340,15 +2340,6 @@ version = "0.4.29"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897"
 
-[[package]]
-name = "lz4_flex"
-version = "0.14.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ecbdfe44b1bd960b68170b417450a628c43f7cf56bb3c5317e61cb230ee7f226"
-dependencies = [
- "twox-hash",
-]
-
 [[package]]
 name = "logos"
 version = "0.15.1"
@@ -2415,6 +2406,15 @@ dependencies = [
  "logos-codegen 0.16.1",
 ]
 
+[[package]]
+name = "lz4_flex"
+version = "0.14.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ecbdfe44b1bd960b68170b417450a628c43f7cf56bb3c5317e61cb230ee7f226"
+dependencies = [
+ "twox-hash",
+]
+
 [[package]]
 name = "macro_rules_attribute"
 version = "0.2.2"

tests/compile/passes/test_fusion_attn.py

@@ -7,7 +7,11 @@ import torch._dynamo
 
 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.platforms import current_platform
 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 @@ class AttentionQuantPatternModel(torch.nn.Module):
         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,
-        )
-        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,
+        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, [[m.name for m in supported]]
+        )
+        raw_tensor = torch.zeros(
+            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 @@ class MLAAttentionQuantPatternModel(torch.nn.Module):
         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
+        # 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,
         )
-        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
-        )
-        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 @@ class MLARoPEKVCacheCatTestModel(torch.nn.Module):
         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(

tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py

@@ -10,7 +10,11 @@ from packaging.version import Version
 
 import vllm.config
 from tests.compile.backend import TestBackend
-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._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
 from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP
 from vllm.compilation.passes.fusion.qk_norm_rope_kvcache_fusion import (
@@ -40,7 +44,7 @@ from vllm.v1.attention.backend import (
     CommonAttentionMetadata,
 )
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.v1.kv_cache_interface import AttentionSpec
+from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
 
 pytestmark = pytest.mark.skip_global_cleanup
 
@@ -147,20 +151,25 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module):
             self.attn._k_scale = self.attn._k_scale.to(device)
             self.attn._v_scale = self.attn._v_scale.to(device)
 
-        self.builder = self.attn.attn_backend.get_builder_cls()(
-            kv_cache_spec=AttentionSpec(
+        # Mirror the worker spec-collection loop: the backend publishes the
+        # packing it needs.
+        self.kv_cache_spec = self.attn.attn_backend.customize_spec(
+            AttentionSpec(
                 block_size=self.block_size,
                 num_kv_heads=self.num_kv_heads,
                 head_size=head_size,
                 dtype=self.kv_cache_dtype,
-            ),
+            )
+        )
+        self.builder = self.attn.attn_backend.get_builder_cls()(
+            kv_cache_spec=self.kv_cache_spec,
             layer_names=[self.attn.layer_name],
             vllm_config=vllm_config,
             device=device,
         )
 
     def build_attn_metadata(
-        self, batch_size: int, kv_stride_order: tuple[int, ...] | None = None
+        self, batch_size: int, layout: KVCacheLayout
     ) -> CommonAttentionMetadata:
         batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
         common_attn_metadata = create_common_attn_metadata(
@@ -170,27 +179,18 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module):
         max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
         num_blocks = batch_size * max_blocks
 
-        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
-        )
-        # Caller can force a physical layout; else use the backend's.
-        if kv_stride_order is None:
-            try:
-                kv_stride_order = attn_backend.get_kv_cache_stride_order()
-            except (AttributeError, NotImplementedError):
-                kv_stride_order = tuple(range(len(kv_cache_shape)))
-
-        kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_stride_order)
-        inv_order = [kv_stride_order.index(i) for i in range(len(kv_stride_order))]
-
         raw_tensor = torch.zeros(
-            2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
-            dtype=self.kv_cache_dtype,
+            num_blocks * self.kv_cache_spec.page_size_bytes,
+            dtype=torch.int8,
             device=self.device,
         )
-        raw_tensor = raw_tensor.view(kv_cache_shape)
-        kv_cache = raw_tensor.permute(*inv_order)
+        kv_cache = dense_kv_cache_views(
+            raw_tensor,
+            self.kv_cache_spec,
+            num_blocks,
+            num_layers=1,
+            layout=layout,
+        )[0]
 
         # Store as a bare tensor (not wrapped in a list) to match production
         # `bind_kv_cache` behavior.  `get_attention_context` returns this
@@ -272,7 +272,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(
     block_size: int,
     is_neox: bool,
     use_shuffle_kv_layout: str,
-    kv_stride_order: tuple[int, ...],
+    kv_layout: KVCacheLayout,
     dtype: torch.dtype,
     kv_cache_dtype: str,
     rms_norm_eps: float,
@@ -345,7 +345,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(
         # Run unfused (eager) forward
         with set_forward_context(None, vllm_config):
             forward_context = get_forward_context()
-            attn_metadata = model.build_attn_metadata(num_tokens, kv_stride_order)
+            attn_metadata = model.build_attn_metadata(num_tokens, kv_layout)
             forward_context.slot_mapping = {
                 model.layer_name: attn_metadata.slot_mapping
             }
@@ -360,7 +360,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(
         with set_forward_context(None, vllm_config):
             model_fused = torch.compile(model, backend=backend)
             forward_context = get_forward_context()
-            attn_metadata = model_fused.build_attn_metadata(num_tokens, kv_stride_order)
+            attn_metadata = model_fused.build_attn_metadata(num_tokens, kv_layout)
             forward_context.slot_mapping = {
                 model.layer_name: attn_metadata.slot_mapping
             }
@@ -437,10 +437,11 @@ _FUSION_CONFIGS = [
 )
 @pytest.mark.parametrize("num_tokens", [5, 2048])
 @pytest.mark.parametrize(
-    "kv_stride_order",
-    # FA/unified use the 4D packed cache (num_blocks, num_kv_heads, block_size,
-    # 2*head_size); no separate K/V dim to permute.
-    [pytest.param((0, 1, 2, 3), id="packed_4d")],
+    "kv_layout",
+    [
+        pytest.param(KVCacheLayout.LBHNC, id="head_major"),
+        pytest.param(KVCacheLayout.LBNHC, id="token_major"),
+    ],
 )
 @pytest.mark.parametrize("enable_aiter_triton_rope", [True, False])
 @pytest.mark.parametrize("block_size", [16, 32, 64])
@@ -467,7 +468,7 @@ def test_qk_norm_rope_kvcache_fusion(
     is_neox: bool,
     attn_backend: AttentionBackendEnum,
     enable_aiter_triton_rope: bool,
-    kv_stride_order: tuple[int, ...],
+    kv_layout: KVCacheLayout,
     block_size: int,
     dtype: torch.dtype,
     kv_cache_dtype: str,
@@ -486,7 +487,7 @@ def test_qk_norm_rope_kvcache_fusion(
         block_size=block_size,
         is_neox=is_neox,
         use_shuffle_kv_layout="0",
-        kv_stride_order=kv_stride_order,
+        kv_layout=kv_layout,
         dtype=dtype,
         kv_cache_dtype=kv_cache_dtype,
         rms_norm_eps=rms_norm_eps,

tests/compile/passes/test_rope_kvcache_fusion.py

@@ -7,7 +7,11 @@ from torch._higher_order_ops import auto_functionalized
 
 import vllm.config
 from tests.compile.backend import TestBackend
-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._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
 from vllm.compilation.passes.fusion import rope_kvcache_fusion
 from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP
@@ -38,6 +42,10 @@ from vllm.v1.attention.backend import (
     CommonAttentionMetadata,
 )
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheLayout,
+)
 
 INDEX_SELECT_OP = torch.ops.aten.index.Tensor
 VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update
@@ -149,9 +157,20 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):
             FP8_DTYPE if kv_cache_dtype_str.startswith("fp8") else self.dtype
         )
 
+        # Mirror the worker spec-collection loop: the backend publishes its
+        # packing (e.g. ROCM_ATTN's separate K/V head groups).
+        self.kv_cache_spec = self.attn_backend.customize_spec(
+            FullAttentionSpec(
+                block_size=self.block_size,
+                num_kv_heads=self.num_kv_heads,
+                head_size=self.head_size,
+                dtype=self.kv_cache_dtype,
+            )
+        )
+
         # Initialize attn MetadataBuilder
         self.builder = self.attn_backend.get_builder_cls()(
-            kv_cache_spec=self.attn.get_kv_cache_spec(vllm_config),
+            kv_cache_spec=self.kv_cache_spec,
             layer_names=[self.attn.layer_name],
             vllm_config=vllm_config,
             device=device,
@@ -168,30 +187,25 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):
         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)))
+        # A backend's preferred supported layout wins over the default, mirroring
+        # selector-time resolution (e.g. ROCM_ATTN's head groups force LHBNC).
+        layout = KVCacheLayout.LBNHC
+        supported = self.attn_backend.supported_kv_cache_layouts()
+        if supported:
+            layout = supported[0]
 
-        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(
-            2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
-            dtype=self.kv_cache_dtype,
+            num_blocks * self.kv_cache_spec.page_size_bytes,
+            dtype=torch.int8,
             device=self.device,
         )
-        raw_tensor = raw_tensor.view(kv_cache_shape)
-        kv_cache = raw_tensor.permute(*inv_order)
-
-        self.attn.kv_cache = kv_cache
+        self.attn.kv_cache = dense_kv_cache_views(
+            raw_tensor,
+            self.kv_cache_spec,
+            num_blocks,
+            num_layers=1,
+            layout=layout,
+        )[0]
 
         # Build attn metadata
         attn_metadata = self.builder.build(

tests/distributed/test_kvlayout.py

@@ -21,7 +21,7 @@ def test_get_kv_connector_cache_layout_without_kv_connector():
     with set_current_vllm_config(vllm_config):
         # Test with default settings
         layout = get_kv_connector_cache_layout()
-        assert layout == "NHD"
+        assert layout is None
 
 
 def test_get_kv_connector_cache_layout_with_lmcache_connector():
@@ -35,7 +35,7 @@ def test_get_kv_connector_cache_layout_with_lmcache_connector():
     with set_current_vllm_config(vllm_config):
         # Test with default settings
         layout = get_kv_connector_cache_layout()
-        assert layout == "NHD"
+        assert layout is None
 
 
 def test_get_kv_connector_cache_layout_with_nixl_connector():
@@ -52,7 +52,7 @@ def test_get_kv_connector_cache_layout_with_nixl_connector():
     with set_current_vllm_config(vllm_config):
         # Test with default settings
         layout = get_kv_connector_cache_layout()
-        assert layout == "HND"
+        assert layout == "LBHNC"
 
 
 def test_get_kv_connector_cache_layout_with_multi_connector():
@@ -75,4 +75,4 @@ def test_get_kv_connector_cache_layout_with_multi_connector():
     with set_current_vllm_config(vllm_config):
         # Test with default settings
         layout = get_kv_connector_cache_layout()
-        assert layout == "HND"
+        assert layout == "LBHNC"

tests/kernels/attention/test_cache.py

@@ -356,21 +356,21 @@ def test_reshape_and_cache_flash(
             dequant_nvfp4_kv_cache,
         )
 
-        def dequant_nvfp4_cache_nhd(data_cache, scale_cache, global_scale):
-            # data_cache:  [N, T, H, data_dim]  NHD (contiguous inner dims)
-            # scale_cache: [N, T, H, scale_dim] NHD (contiguous inner dims)
-            # Permute to HND layout for the dequant utility.
+        def dequant_nvfp4_cache(data_cache, scale_cache, global_scale):
+            # data_cache:  [B, N, H, data_dim]  logical view (layout in strides)
+            # scale_cache: [B, N, H, scale_dim] logical view (layout in strides)
+            # Permute to [B, H, N, dim] for the dequant utility.
             data_hnd = data_cache.permute(0, 2, 1, 3)
             scale_hnd = scale_cache.permute(0, 2, 1, 3)
             result_hnd = dequant_nvfp4_kv_cache(
                 data_hnd, scale_hnd, global_scale, head_size, block_size
             )
-            return result_hnd.permute(0, 2, 1, 3)  # back to [N, T, H, D]
+            return result_hnd.permute(0, 2, 1, 3)  # back to [B, N, H, dim]
 
-        result_key_cache = dequant_nvfp4_cache_nhd(
+        result_key_cache = dequant_nvfp4_cache(
             nvfp4_key_data, key_scale_cache, k_scale.item()
         )
-        result_value_cache = dequant_nvfp4_cache_nhd(
+        result_value_cache = dequant_nvfp4_cache(
             nvfp4_value_data, value_scale_cache, v_scale.item()
         )
 

tests/kernels/attention/test_minimax_m3.py

@@ -5,10 +5,10 @@
 import pytest
 import torch
 
+from tests.v1.attention.utils import dense_kv_cache_views
 from vllm import _custom_ops as ops
-from vllm.models.minimax_m3.common.indexer import (
-    MiniMaxM3IndexerBackend,
-)
+from vllm.config import CacheConfig
+from vllm.config.cache import _layout_from_name
 from vllm.models.minimax_m3.common.ops.index_topk import (
     minimax_m3_index_decode,
     minimax_m3_index_score,
@@ -19,15 +19,15 @@ from vllm.models.minimax_m3.common.ops.sparse_attn import (
     minimax_m3_sparse_attn,
     minimax_m3_sparse_attn_decode,
 )
-from vllm.models.minimax_m3.common.sparse_attention import (
-    MiniMaxM3SparseBackend,
-    MiniMaxM3SparseTritonImpl,
-)
+from vllm.models.minimax_m3.common.sparse_attention import MiniMaxM3SparseTritonImpl
 from vllm.platforms import current_platform
-from vllm.v1.attention.backends.utils import set_kv_cache_layout
-from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
-from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache
-from vllm.v1.worker.utils import AttentionGroup
+from vllm.v1.attention.backends.utils import record_kv_cache_layout
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheLayout,
+    MLAAttentionSpec,
+    compute_layer_kv_cache_shape_bytes,
+)
 
 if not (current_platform.is_cuda() or current_platform.is_rocm()):
     pytest.skip(
@@ -37,35 +37,46 @@ if not (current_platform.is_cuda() or current_platform.is_rocm()):
 
 
 @pytest.fixture
-def kv_layout(request):
-    """Set the global KV cache layout for one test and restore it after."""
-    set_kv_cache_layout(request.param)
-    try:
-        yield request.param
-    finally:
-        set_kv_cache_layout(None)
+def kv_layout(request) -> KVCacheLayout:
+    """Resolve the requested layout name (legacy NHD/HND aliases included)."""
+    return _layout_from_name(request.param)
 
 
-def _stride_order_for(backend: type[MiniMaxM3SparseBackend], ndim: int) -> tuple:
-    """Mirror the allocator's stride-order resolution (identity fallback)."""
-    try:
-        stride_order = backend.get_kv_cache_stride_order()
-        assert len(stride_order) == ndim
-    except (AttributeError, NotImplementedError):
-        stride_order = tuple(range(ndim))
+def _layer_stride_order(layout: KVCacheLayout, ndim: int) -> tuple[int, ...]:
+    """Per-layer physical stride order for the given layout; the 3-dim
+    indexer side cache (H=1) is contiguous, so identity."""
+    if ndim == 3:
+        return (0, 1, 2)
+    stride_order = layout.layer_view_order
+    assert len(stride_order) == ndim
     return stride_order
 
 
+def _main_spec() -> FullAttentionSpec:
+    return FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_DIM,
+        head_size_v=HEAD_DIM,
+        dtype=DTYPE,
+    )
+
+
+def _main_kv_logical_shape(num_pages: int) -> tuple[int, ...]:
+    """Standardized per-layer logical shape [B, H, N, C] for the main cache,
+    derived the same way the production allocator does."""
+    shape_bytes = compute_layer_kv_cache_shape_bytes(_main_spec(), num_pages)
+    return (*shape_bytes[:-1], shape_bytes[-1] // DTYPE.itemsize)
+
+
 def _allocate_main_kv_via_contract(
-    num_pages: int, device: torch.device | str = "cuda"
+    num_pages: int, layout: KVCacheLayout, device: torch.device | str = "cuda"
 ) -> torch.Tensor:
     """Build the main KV cache exactly as the production allocator does for the
-    currently active layout: allocate the physical (permuted) tensor, then
-    expose the inverse-permuted logical-NHD view the backend sees."""
-    logical_shape = MiniMaxM3SparseBackend.get_kv_cache_shape(
-        num_pages, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
-    )
-    stride_order = _stride_order_for(MiniMaxM3SparseBackend, len(logical_shape))
+    given layout: allocate the physical (permuted) tensor, then expose the
+    inverse-permuted logical [B, H, N, C] view the backend sees."""
+    logical_shape = _main_kv_logical_shape(num_pages)
+    stride_order = _layer_stride_order(layout, len(logical_shape))
     physical_shape = tuple(logical_shape[i] for i in stride_order)
     inv_order = [stride_order.index(i) for i in range(len(stride_order))]
     raw = torch.randn(physical_shape, device=device, dtype=DTYPE)
@@ -817,7 +828,7 @@ def _reference_sparse_attn(
     ],
 )
 def test_prefill_sparse_attention_correctness(
-    kv_layout: str,
+    kv_layout: KVCacheLayout,
     q_lens: tuple[int, ...],
     kv_lens: tuple[int, ...],
 ):
@@ -851,7 +862,7 @@ def test_prefill_sparse_attention_correctness(
     # Allocate the main KV cache through the backend layout contract so the
     # physical storage matches the active layout (contiguous NHD or strided
     # HND), while the kernels and reference see the logical-NHD view.
-    kv_cache = _allocate_main_kv_via_contract(num_pages)
+    kv_cache = _allocate_main_kv_via_contract(num_pages, kv_layout)
 
     # Build sparse block indices with the same contract as the real M3 indexer:
     # one forced local block, then score-selected older causal blocks.
@@ -904,38 +915,31 @@ def test_prefill_sparse_attention_correctness(
     assert error.max().item() < 1.7e-2
 
 
-def test_main_backend_layout_contract():
-    """The main sparse backend exposes the logical-NHD shape and the
-    flash_attn-style stride order for each layout."""
+def test_main_cache_layout_contract():
+    """The standardized per-layer logical shape is [B, H, N, C] with packed
+    K/V content, and the legacy layout aliases resolve to the expected
+    per-layer stride orders."""
     nb, bs, h, d = 7, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
-    logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
+    logical = _main_kv_logical_shape(nb)
     assert logical == (nb, h, bs, 2 * d)
     # The old separate K/V-axis shape is no longer the logical shape.
     assert logical != (nb, 2, bs, h, d)
 
-    try:
-        set_kv_cache_layout("HND")
-        assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 2, 3)
-        set_kv_cache_layout("NHD")
-        assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 2, 1, 3)
-    finally:
-        set_kv_cache_layout(None)
+    assert KVCacheLayout.LBHNC.layer_view_order == (0, 1, 2, 3)
+    assert KVCacheLayout.LBNHC.layer_view_order == (0, 2, 1, 3)
 
-    for layout in ("NHD", "HND"):
-        try:
-            set_kv_cache_layout(layout)
-            order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
-        finally:
-            set_kv_cache_layout(None)
+    cache_config = CacheConfig()
+    record_kv_cache_layout(cache_config, "HND")
+    assert cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.LBHNC
+    cache_config = CacheConfig()
+    record_kv_cache_layout(cache_config, "NHD")
+    assert cache_config.get_resolved_kv_cache_layout() is KVCacheLayout.LBNHC
+
+    for layout in (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC):
+        order = layout.layer_view_order
         # Valid permutation: no duplicates, covers every axis.
         assert set(order) == set(range(len(order)))
 
-    # M3 has no cross-layer KV blocks.
-    with pytest.raises(NotImplementedError):
-        MiniMaxM3SparseBackend.get_kv_cache_stride_order(
-            include_num_layers_dimension=True
-        )
-
 
 def test_aiter_sparse_pa_layout_contract(monkeypatch):
     """The shuffle-only AITER path retains separately contiguous K/V storage."""
@@ -948,75 +952,53 @@ def test_aiter_sparse_pa_layout_contract(monkeypatch):
         lambda: True,
     )
 
-    nb, bs, h, d = 7, BLOCK_SIZE, 1, HEAD_DIM
-    logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
-    order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
-    assert logical == (nb, 2, bs, h, d)
-    assert order == (1, 0, 2, 3, 4)
+    backend = sparse_attn_mod.MiniMaxM3SparseBackend
+    assert KVCacheLayout.LHBNC in backend.supported_kv_cache_layouts()
 
-    physical_shape = tuple(logical[i] for i in order)
-    inv_order = [order.index(i) for i in range(len(order))]
-    raw = torch.empty(physical_shape, device="cuda", dtype=DTYPE)
-    logical_view = raw.permute(*inv_order)
-    key_cache, value_cache = logical_view.unbind(1)
+    nb, h = 7, 1
+    spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=h,
+        head_size=HEAD_DIM,
+        dtype=DTYPE,
+        num_head_slots=2,
+        state_content_bytes=h * HEAD_DIM * DTYPE.itemsize,
+    )
+    raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8)
+    view = dense_kv_cache_views(raw, spec, nb, 1, KVCacheLayout.LHBNC)[0]
+    key_cache, value_cache = view.unbind(1)
     assert key_cache.is_contiguous()
     assert value_cache.is_contiguous()
 
 
-def test_aiter_sparse_pa_rejects_multiple_kv_heads(monkeypatch):
-    """Do not pair AITER's separated cache layout with the Triton fallback."""
-    import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
-
-    monkeypatch.setattr(sparse_attn_mod.rocm_aiter_ops, "is_enabled", lambda: True)
-    monkeypatch.setattr(
-        sparse_attn_mod.rocm_aiter_ops,
-        "is_shuffle_kv_cache_enabled",
-        lambda: True,
+def test_indexer_cache_squeezes_to_contiguous_3d():
+    """The indexer side cache is standardized 4D with H=1: under both layouts
+    the allocator's logical view stays contiguous and squeezes (as
+    `MiniMaxM3IndexerCache.bind_kv_cache` does) to the 3-dim
+    [num_blocks, block_size, head_dim] cache the kernels consume."""
+    nb = 5
+    ispec = MLAAttentionSpec(
+        block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
     )
+    shape_bytes = compute_layer_kv_cache_shape_bytes(ispec, nb)
+    assert shape_bytes == (nb, 1, BLOCK_SIZE, HEAD_DIM * DTYPE.itemsize)
+    assert _layer_stride_order(KVCacheLayout.LBNHC, 3) == (0, 1, 2)
 
-    with pytest.raises(ValueError, match="num_kv_heads == 1"):
-        MiniMaxM3SparseBackend.get_kv_cache_shape(7, BLOCK_SIZE, 2, HEAD_DIM)
-
-
-def test_main_backend_unknown_layout_raises(monkeypatch):
-    """An unrecognized layout (injected past env-var validation) is rejected."""
-    import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
-
-    monkeypatch.setattr(sparse_attn_mod, "get_kv_cache_layout", lambda: "BOGUS")
-    with pytest.raises(ValueError, match="Unknown cache layout format"):
-        MiniMaxM3SparseBackend.get_kv_cache_stride_order()
-
-
-def test_indexer_backend_stride_order_is_identity():
-    """The 3-dim indexer cache must not inherit the parent's 4-element stride
-    order; it overrides to the 3-element identity so the allocator keeps the
-    contiguous layout."""
-    assert MiniMaxM3IndexerBackend.get_kv_cache_stride_order() == (0, 1, 2)
-
-    # Cross-layer (per-layer-stacked) KV blocks are not supported.
-    with pytest.raises(NotImplementedError):
-        MiniMaxM3IndexerBackend.get_kv_cache_stride_order(
-            include_num_layers_dimension=True
-        )
-
-    # The stride order matches the 3-dim indexer shape rank.
-    indexer_shape = MiniMaxM3IndexerBackend.get_kv_cache_shape(
-        5, BLOCK_SIZE, 1, HEAD_DIM
-    )
-    assert len(indexer_shape) == 3
-    assert _stride_order_for(MiniMaxM3IndexerBackend, len(indexer_shape)) == (0, 1, 2)
+    for layout in (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC):
+        iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8)
+        view = dense_kv_cache_views(iraw, ispec, nb, 1, layout, BLOCK_SIZE)[0]
+        assert tuple(view.shape) == (nb, 1, BLOCK_SIZE, HEAD_DIM)
+        indexer_cache = view.squeeze(1)
+        assert tuple(indexer_cache.shape) == (nb, BLOCK_SIZE, HEAD_DIM)
+        assert indexer_cache.is_contiguous()
 
 
 def test_hnd_allocation_is_packed_head_major():
     """Under HND the backend-visible logical view is the packed head-major
     physical allocation."""
     nb, bs, h, d = 4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
-    logical = MiniMaxM3SparseBackend.get_kv_cache_shape(nb, bs, h, d)
-    try:
-        set_kv_cache_layout("HND")
-        stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
-    finally:
-        set_kv_cache_layout(None)
+    logical = _main_kv_logical_shape(nb)
+    stride_order = KVCacheLayout.LBHNC.layer_view_order
 
     physical_shape = tuple(logical[i] for i in stride_order)
     assert physical_shape == (nb, h, bs, 2 * d)
@@ -1042,27 +1024,13 @@ def test_main_cache_is_block_first_and_unpadded():
     """The allocator's contiguous-view branch (not the padded-strided branch)
     is used for the main GQA cache: its spec is unpadded and the physical
     layout keeps num_blocks as the first dimension under both layouts."""
-    from vllm.v1.kv_cache_interface import FullAttentionSpec
-
-    spec = FullAttentionSpec(
-        block_size=BLOCK_SIZE,
-        num_kv_heads=NUM_KV_HEADS,
-        head_size=HEAD_DIM,
-        head_size_v=HEAD_DIM,
-        dtype=DTYPE,
-    )
+    spec = _main_spec()
     # Unpadded -> allocator uses kv_tensor.view(...) rather than as_strided().
     assert spec.page_size_padded is None
 
-    logical = MiniMaxM3SparseBackend.get_kv_cache_shape(
-        4, BLOCK_SIZE, NUM_KV_HEADS, HEAD_DIM
-    )
-    for layout in ("NHD", "HND"):
-        try:
-            set_kv_cache_layout(layout)
-            order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
-        finally:
-            set_kv_cache_layout(None)
+    logical = _main_kv_logical_shape(4)
+    for layout in (KVCacheLayout.LBNHC, KVCacheLayout.LBHNC):
+        order = layout.layer_view_order
         inv_order = [order.index(i) for i in range(len(order))]
         # Physical first dim is num_blocks (block-first); required by the
         # padded-strided branch's block-first assumption if it were ever taken.
@@ -1275,7 +1243,7 @@ def test_aiter_decode_sparse_block_table_supports_spec_decode():
 @pytest.mark.parametrize("decode_query_len", [1, 4])
 @pytest.mark.parametrize("num_padded_reqs", [0, 2])
 def test_decode_sparse_attention_correctness(
-    kv_layout: str,
+    kv_layout: KVCacheLayout,
     seq_lens_list: tuple[int, ...],
     decode_query_len: int,
     num_padded_reqs: int,
@@ -1287,7 +1255,7 @@ def test_decode_sparse_attention_correctness(
     q, block_table, seq_lens, topk_idx, num_pages = _build_decode_inputs(
         seq_lens_list, decode_query_len, num_padded_reqs
     )
-    kv_cache = _allocate_main_kv_via_contract(num_pages)
+    kv_cache = _allocate_main_kv_via_contract(num_pages, kv_layout)
 
     actual = torch.empty_like(q)
     minimax_m3_sparse_attn_decode(
@@ -1364,93 +1332,20 @@ def test_decode_wrong_layout_breaks_parity():
     assert (actual.float() - expected.float()).abs().max().item() > 1.7e-2
 
 
-def _make_attn_group(backend, spec):
-    return AttentionGroup(
-        backend=backend,
-        layer_names=["main"],
-        kv_cache_spec=spec,
-        kv_cache_group_id=0,
-    )
-
-
 def test_main_cache_byte_identical_through_production_allocator():
-    """AC-2: drive the real allocator (`_reshape_kv_cache`) for the M3 main
+    """AC-2: drive the real allocator (`create_kv_cache_views`) for the M3 main
     `FullAttentionSpec` under HND and assert the backend-visible view has the
-    same shape, stride, and storage offset as the packed-HND allocation; the
-    indexer `MLAAttentionSpec` allocates through the same path to its 3-dim
-    shape."""
+    same shape, stride, and storage offset as the packed-HND allocation."""
     nb = 4
-    spec = FullAttentionSpec(
-        block_size=BLOCK_SIZE,
-        num_kv_heads=NUM_KV_HEADS,
-        head_size=HEAD_DIM,
-        head_size_v=HEAD_DIM,
-        dtype=DTYPE,
-    )
+    spec = _main_spec()
     raw = torch.zeros(nb * spec.page_size_bytes, dtype=torch.int8)
-    group = _make_attn_group(MiniMaxM3SparseBackend, spec)
-    try:
-        set_kv_cache_layout("HND")
-        kv_caches = _reshape_kv_cache([group], {"main": raw}, "auto", [BLOCK_SIZE], {})
-    finally:
-        set_kv_cache_layout(None)
-    view = kv_caches["main"]
+    view = dense_kv_cache_views(raw, spec, nb, 1, KVCacheLayout.LBHNC, BLOCK_SIZE)[0]
 
     oracle = raw.view(DTYPE).view((nb, NUM_KV_HEADS, BLOCK_SIZE, 2 * HEAD_DIM))
     assert tuple(view.shape) == tuple(oracle.shape)
     assert view.stride() == oracle.stride()
     assert view.storage_offset() == oracle.storage_offset()
 
-    # Indexer cache allocates through the same path under both layouts.
-    ispec = MLAAttentionSpec(
-        block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
-    )
-    for layout in ("NHD", "HND"):
-        iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8)
-        igroup = AttentionGroup(
-            backend=MiniMaxM3IndexerBackend,
-            layer_names=["idx"],
-            kv_cache_spec=ispec,
-            kv_cache_group_id=0,
-        )
-        try:
-            set_kv_cache_layout(layout)
-            iout = _reshape_kv_cache([igroup], {"idx": iraw}, "auto", [BLOCK_SIZE], {})
-        finally:
-            set_kv_cache_layout(None)
-        assert tuple(iout["idx"].shape) == (nb, BLOCK_SIZE, HEAD_DIM)
-
-
-def test_indexer_inherited_stride_order_trips_allocator_assert():
-    """AC-4 negative: without the indexer override, the inherited 4-element
-    stride order trips the allocator's `len(stride_order) == len(shape)` assert
-    for the 3-dim indexer shape; the `AssertionError` is NOT swallowed by the
-    allocator's `(AttributeError, NotImplementedError)` fallback."""
-
-    class _BrokenIndexerBackend(MiniMaxM3IndexerBackend):
-        # Simulate inheriting the parent's 4-element stride order.
-        get_kv_cache_stride_order = staticmethod(
-            MiniMaxM3SparseBackend.get_kv_cache_stride_order
-        )
-
-    nb = 4
-    ispec = MLAAttentionSpec(
-        block_size=BLOCK_SIZE, num_kv_heads=1, head_size=HEAD_DIM, dtype=DTYPE
-    )
-    iraw = torch.zeros(nb * ispec.page_size_bytes, dtype=torch.int8)
-    igroup = AttentionGroup(
-        backend=_BrokenIndexerBackend,
-        layer_names=["idx"],
-        kv_cache_spec=ispec,
-        kv_cache_group_id=0,
-    )
-    try:
-        set_kv_cache_layout("HND")
-        with pytest.raises(AssertionError):
-            _reshape_kv_cache([igroup], {"idx": iraw}, "auto", [BLOCK_SIZE], {})
-    finally:
-        set_kv_cache_layout(None)
-
 
 def test_padded_main_cache_is_flagged():
     """AC-2.1 negative: the M3 main cache relies on the allocator's
@@ -1465,11 +1360,7 @@ def test_padded_main_cache_is_flagged():
         )
         assert inv_order[0] == 0, "main GQA cache must remain block-first"
 
-    try:
-        set_kv_cache_layout("HND")
-        stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
-    finally:
-        set_kv_cache_layout(None)
+    stride_order = KVCacheLayout.LBHNC.layer_view_order
 
     good = FullAttentionSpec(
         block_size=BLOCK_SIZE,
@@ -1493,13 +1384,13 @@ def test_padded_main_cache_is_flagged():
 
 
 @pytest.mark.parametrize("kv_layout", ["NHD", "HND"], indirect=True)
-def test_reshape_and_cache_flash_write_persists(kv_layout: str):
+def test_reshape_and_cache_flash_write_persists(kv_layout: KVCacheLayout):
     """AC-5 write path: the `reshape_and_cache_flash` write site now consumes
     packed-content K/V split views. Writing through those views must persist
     into the bound storage under both layouts."""
     torch.manual_seed(0)
     num_pages = 4
-    kv_cache = _allocate_main_kv_via_contract(num_pages)
+    kv_cache = _allocate_main_kv_via_contract(num_pages, kv_layout)
     with torch.no_grad():
         kv_cache.zero_()
 

tests/kernels/attention/test_rocm_aiter_fa.py

@@ -310,18 +310,23 @@ def test_aiter_mha_backend_contract():
 
 
 def test_aiter_mha_backend_validates_kv_cache_block_size():
-    """The backend should reject KV cache shapes that cannot be gathered
+    """The backend should reject KV cache specs that cannot be gathered
     correctly by the ROCm kernel."""
     from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
+    from vllm.v1.kv_cache_interface import FullAttentionSpec
 
-    assert AiterFlashAttentionBackend.get_kv_cache_shape(8, 16, 8, 128) == (
-        8,
-        8,
-        16,
-        256,
-    )
     with pytest.raises(ValueError, match="Block size must be a multiple of 16"):
-        AiterFlashAttentionBackend.get_kv_cache_shape(8, 15, 8, 128)
+        AiterFlashAttentionBackend.customize_spec(
+            FullAttentionSpec(
+                block_size=15, num_kv_heads=8, head_size=128, dtype=torch.bfloat16
+            )
+        )
+    # block_size == 1 is the per-token page-size probe and must pass through.
+    AiterFlashAttentionBackend.customize_spec(
+        FullAttentionSpec(
+            block_size=1, num_kv_heads=8, head_size=128, dtype=torch.bfloat16
+        )
+    )
 
 
 def test_aiter_mha_backend_supports_compute_capability_matches_mi3xx_probe():

tests/kernels/attention/test_trtllm_kvfp8_dequant.py

@@ -11,6 +11,7 @@ import pytest
 import torch
 
 from vllm.platforms import current_platform
+from vllm.v1.kv_cache_interface import KVCacheLayout
 
 if current_platform.is_rocm():
     pytest.skip(
@@ -34,51 +35,40 @@ def to_float8(x, dtype=None):
     return x_scl_sat.to(dtype), scale.float().reciprocal()
 
 
-def make_contiguous_kv_cache(num_blocks, num_kv_heads, block_size, head_size):
-    """Create a standard contiguous fp8 KV cache (HND layout)."""
-    raw = torch.randn(
-        num_blocks,
-        2,
-        num_kv_heads,
-        block_size,
-        head_size,
-        dtype=torch.bfloat16,
-        device="cuda",
-    )
-    kv_cache, scale = to_float8(raw)
-    return kv_cache, scale
-
-
-def make_cross_layer_kv_cache(
+def make_random_kv_cache(
     num_blocks,
     num_kv_heads,
     block_size,
     head_size,
-    num_layers=4,
+    layout=KVCacheLayout.LBHNC,
+    num_layers=None,
 ):
     """
-    Create a non-contiguous per-layer view mimicking cross-layer allocation.
+    Create a random fp8 KV cache mimicking the production allocation.
 
-    Physical layout: (num_blocks, 2, num_kv_heads, num_layers, block_size, head_size)
+    Physical layout: [L, B, H, N, 2 * hs] permuted by ``layout.stride_order``
     Returned view:   (num_blocks, 2, num_kv_heads, block_size, head_size)
-    with non-contiguous strides on dims 0, 1, 2 (they skip over num_layers).
+    with non-contiguous strides under cross-layer layouts (they skip over
+    num_layers), matching the actual forward path.
     """
-    raw = torch.randn(
+    logical_4d = (num_blocks, num_kv_heads, block_size, 2 * head_size)
+    if num_layers is None:
+        num_layers = 1 if layout.is_layer_compact else 2
+    logical_5d = (num_layers, *logical_4d)
+    physical_5d = tuple(logical_5d[i] for i in layout.stride_order)
+    inv_order = [layout.stride_order.index(i) for i in range(5)]
+
+    raw_phys = torch.randn(*physical_5d, dtype=torch.bfloat16, device="cuda")
+    fp8_phys, scale = to_float8(raw_phys)
+    fp8_4d = fp8_phys.permute(*inv_order)[0]
+    kv_5d = fp8_4d.view(
         num_blocks,
-        2,
         num_kv_heads,
-        num_layers,
         block_size,
+        2,
         head_size,
-        dtype=torch.bfloat16,
-        device="cuda",
-    )
-    fp8_full, scale = to_float8(raw)
-    layer_view = fp8_full[:, :, :, 0, :, :]
-    assert not layer_view.is_contiguous(), (
-        f"Expected non-contiguous view, got strides {layer_view.stride()}"
-    )
-    return layer_view, scale
+    ).permute(0, 3, 1, 2, 4)
+    return kv_5d, scale
 
 
 def ref_dequant(kv_cache, block_tables, k_scale, v_scale, dequant_dtype):
@@ -114,7 +104,7 @@ def ref_dequant(kv_cache, block_tables, k_scale, v_scale, dequant_dtype):
 @pytest.mark.parametrize("block_size", [16, 32])
 @pytest.mark.parametrize("batch_size", [1, 4])
 @pytest.mark.parametrize("num_pages_per_seq", [3, 8])
-@pytest.mark.parametrize("contiguous", [True, False])
+@pytest.mark.parametrize("layout", list(KVCacheLayout))
 @torch.inference_mode()
 def test_trtllm_kvfp8_dequant(
     num_kv_heads: int,
@@ -122,7 +112,7 @@ def test_trtllm_kvfp8_dequant(
     block_size: int,
     batch_size: int,
     num_pages_per_seq: int,
-    contiguous: bool,
+    layout: KVCacheLayout,
 ):
     from vllm.v1.attention.backends.flashinfer import (
         trtllm_prefill_attn_kvfp8_dequant,
@@ -130,20 +120,13 @@ def test_trtllm_kvfp8_dequant(
 
     torch.set_default_device("cuda")
 
-    if contiguous:
-        kv_cache, scale = make_contiguous_kv_cache(
-            NUM_BLOCKS,
-            num_kv_heads,
-            block_size,
-            head_size,
-        )
-    else:
-        kv_cache, scale = make_cross_layer_kv_cache(
-            NUM_BLOCKS,
-            num_kv_heads,
-            block_size,
-            head_size,
-        )
+    kv_cache, scale = make_random_kv_cache(
+        NUM_BLOCKS,
+        num_kv_heads,
+        block_size,
+        head_size,
+        layout=layout,
+    )
 
     k_scale = scale.clone()
     v_scale = scale.clone()
@@ -187,7 +170,7 @@ def test_block_tables_with_zero_pages():
     torch.set_default_device("cuda")
     num_kv_heads, block_size, head_size = 8, 16, 64
 
-    kv_cache, scale = make_contiguous_kv_cache(
+    kv_cache, scale = make_random_kv_cache(
         NUM_BLOCKS,
         num_kv_heads,
         block_size,
@@ -234,7 +217,7 @@ def test_all_zero_block_tables():
     torch.set_default_device("cuda")
     num_kv_heads, block_size, head_size = 4, 16, 64
 
-    kv_cache, scale = make_contiguous_kv_cache(
+    kv_cache, scale = make_random_kv_cache(
         NUM_BLOCKS,
         num_kv_heads,
         block_size,
@@ -266,7 +249,7 @@ def test_different_k_v_scales():
     torch.set_default_device("cuda")
     num_kv_heads, block_size, head_size = 8, 16, 64
 
-    kv_cache, _ = make_contiguous_kv_cache(
+    kv_cache, _ = make_random_kv_cache(
         NUM_BLOCKS,
         num_kv_heads,
         block_size,
@@ -299,7 +282,7 @@ def test_single_page_per_seq():
     torch.set_default_device("cuda")
     num_kv_heads, block_size, head_size = 8, 16, 128
 
-    kv_cache, scale = make_contiguous_kv_cache(
+    kv_cache, scale = make_random_kv_cache(
         NUM_BLOCKS,
         num_kv_heads,
         block_size,
@@ -332,7 +315,7 @@ def test_large_page_indices():
     num_kv_heads, block_size, head_size = 8, 16, 128
     large_num_blocks = 32768
 
-    kv_cache, scale = make_contiguous_kv_cache(
+    kv_cache, scale = make_random_kv_cache(
         large_num_blocks,
         num_kv_heads,
         block_size,
@@ -369,7 +352,7 @@ def test_large_block_size():
     torch.set_default_device("cuda")
     num_kv_heads, block_size, head_size = 4, 64, 128
 
-    kv_cache, scale = make_contiguous_kv_cache(
+    kv_cache, scale = make_random_kv_cache(
         NUM_BLOCKS,
         num_kv_heads,
         block_size,
@@ -401,7 +384,7 @@ def test_large_block_size():
 def test_cross_layer_many_layers():
     """
     Non-contiguous with 36 layers -- matches real gpt-oss-120b.
-    Strides are far from contiguous (factor of 36 in the gaps).
+    Cross-layer (BLHNC) strides are far from contiguous (factor of 36 in the gaps).
     """
     from vllm.v1.attention.backends.flashinfer import (
         trtllm_prefill_attn_kvfp8_dequant,
@@ -409,14 +392,14 @@ def test_cross_layer_many_layers():
 
     torch.set_default_device("cuda")
     num_kv_heads, block_size, head_size = 8, 16, 64
-    num_layers = 36
 
-    kv_cache, scale = make_cross_layer_kv_cache(
+    kv_cache, scale = make_random_kv_cache(
         NUM_BLOCKS,
         num_kv_heads,
         block_size,
         head_size,
-        num_layers=num_layers,
+        layout=KVCacheLayout.BLHNC,
+        num_layers=36,
     )
     k_scale = v_scale = scale.clone()
 

tests/models/multimodal/generation/test_voxtral_realtime.py

@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from fractions import Fraction
+
 import pytest
 import pytest_asyncio
 from mistral_common.protocol.transcription.request import (
@@ -75,7 +77,8 @@ def assert_encoder_kv_cache_spec(engine: LLM) -> None:
     assert audio_config.block_pool_size == 4
     assert isinstance(spec, SlidingWindowSpec)
     assert spec.block_size == 16
-    assert spec.num_kv_heads == 128
+    assert spec.num_kv_heads == 32
+    assert spec.tokens_per_state == Fraction(1, 4)
     # cdiv(750, 4) == 188 pooled tokens cover the model's window; the extra
     # +1 is an eviction margin (see whisper_causal.py get_kv_cache_spec).
     assert spec.sliding_window == cdiv(750, 4) + 1 == 189

tests/models/test_initialization.py

@@ -8,6 +8,7 @@ import pytest
 
 from vllm import LLM
 from vllm.utils.mem_constants import GiB_bytes
+from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
 from vllm.v1.core.kv_cache_utils import (
     generate_scheduler_kv_cache_config,
     get_kv_cache_configs,
@@ -82,6 +83,12 @@ def can_initialize(
     # Avoid calling model.forward()
     def _initialize_kv_caches_v1(self, vllm_config):
         kv_cache_specs = self.model_executor.get_kv_cache_specs()
+        layout = resolve_kv_cache_layout(
+            vllm_config,
+            self.model_executor.get_supported_kv_cache_layouts(),
+            [spec for worker_specs in kv_cache_specs for spec in worker_specs.values()],
+        )
+        self.model_executor.set_kv_cache_layout(layout.name)
         kv_cache_configs = get_kv_cache_configs(
             vllm_config,
             kv_cache_specs,

tests/quantization/test_turboquant.py

@@ -294,6 +294,7 @@ class TestTurboQuantKVCacheSpec:
             head_size_v=128,
             num_kv_heads=4,
             sliding_window=None,
+            get_attn_backend=lambda: TurboQuantAttentionBackend,
         )
         vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=32))
 

tests/test_attention_backend_registry.py

@@ -38,11 +38,6 @@ class CustomAttentionBackend(AttentionBackend):
         """Mock builder class."""
         return None
 
-    @staticmethod
-    def get_required_kv_cache_layout():
-        """Mock KV cache layout."""
-        return None
-
 
 class CustomMambaAttentionImpl(AttentionImpl):
     """Mock custom mamba attention implementation for testing."""
@@ -71,11 +66,6 @@ class CustomMambaAttentionBackend(AttentionBackend):
         """Mock builder class."""
         return None
 
-    @staticmethod
-    def get_required_kv_cache_layout():
-        """Mock KV cache layout."""
-        return None
-
 
 def test_custom_is_not_alias_of_any_backend():
     # Get all members of AttentionBackendEnum

tests/v1/attention/test_attention_backends.py

@@ -31,10 +31,7 @@ from vllm.v1.attention.backend import (
     CommonAttentionMetadata,
 )
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.v1.attention.backends.utils import (
-    set_kv_cache_layout,
-)
-from vllm.v1.kv_cache_interface import FullAttentionSpec
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout
 
 BACKENDS_TO_TEST = [
     AttentionBackendEnum.FLASH_ATTN,
@@ -44,6 +41,14 @@ BACKENDS_TO_TEST = [
     "FLEX_ATTENTION_SLOW",
 ]
 
+
+def _actual_backend(backend) -> AttentionBackendEnum:
+    """Resolve pseudo-backends (FLEX_ATTENTION_SLOW) to their real enum."""
+    if backend == "FLEX_ATTENTION_SLOW":
+        return AttentionBackendEnum.FLEX_ATTENTION
+    return backend
+
+
 DEVICE_TYPE = current_platform.device_type
 
 # Use the platform's preferred FP8 type so the stored cache matches what the
@@ -118,6 +123,7 @@ def create_and_prepopulate_kv_cache(
     device: torch.device,
     num_blocks: int,
     common_attn_metadata: CommonAttentionMetadata,
+    layout: KVCacheLayout,
     randomize_blocks: bool = True,
     kv_cache_dtype: str = "auto",
 ) -> torch.Tensor:
@@ -126,19 +132,22 @@ def create_and_prepopulate_kv_cache(
     Args:
         k_contexts: List of key context tensors for each sequence
         v_contexts: List of value context tensors for each sequence
-        seq_lens: List of sequence lengths
         block_size: Size of each block
         num_kv_heads: Number of KV heads
         head_size: Size of each head
         dtype: Data type for the cache
         device: Device to create the cache on
         num_blocks: Total number of blocks in the cache
-        block_table: Block table tensor to populate
+        common_attn_metadata: Provides seq lens, block table and slot mapping
+        layout: Physical layout to allocate in; the cache is returned as the
+                logical ``[B, H, N, C]`` view (as ``create_kv_cache_views`` does)
         randomize_blocks: Whether to randomly permute blocks
                           or use sequential order
+        kv_cache_dtype: Cache dtype string; fp8 caches use fp8 storage
 
     Returns:
-        Tuple of (kv_cache, updated_block_table)
+        A 4D tensor in logical ``(num_blocks, num_kv_heads, block_size,
+        2 * head_size)`` order with strides determined by ``layout``.
     """
     batch_size = len(k_contexts)
     seq_lens = common_attn_metadata.seq_lens.cpu()
@@ -155,26 +164,33 @@ def create_and_prepopulate_kv_cache(
     fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
     storage_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype] if fp8_kv_cache else dtype
 
-    kv_cache = torch.zeros(
-        num_blocks,
-        block_size,
-        num_kv_heads,
-        2 * head_size,
-        dtype=storage_dtype,
-        device=device,
-    )
-    kv_cache_flat = kv_cache.view(-1, num_kv_heads, 2 * head_size)
+    # Logical 5D shape is always [L, B, H, N, C]. Cross-layer layouts need
+    # at least two layers to reproduce the inter-layer gaps in a layer view.
+    logical_4d = (num_blocks, num_kv_heads, block_size, 2 * head_size)
+    num_layers = 1 if layout.is_layer_compact else 2
+    logical_5d = (num_layers, *logical_4d)
+    physical_5d = tuple(logical_5d[i] for i in layout.stride_order)
+    inv_order = [layout.stride_order.index(i) for i in range(5)]
 
-    # Populate the cache with the context tokens
+    kv_cache_physical = torch.zeros(physical_5d, dtype=storage_dtype, device=device)
+    # Permute to logical [L, B, H, N, C], then select a layer. This mirrors
+    # create_kv_cache_views and retains cross-layer strides in the 4D view.
+    kv_cache = kv_cache_physical.permute(*inv_order)[0]
+
+    # Write context tokens into the cache via the logical view:
+    # kv_cache[block, :, token_in_block, :] routes correctly regardless
+    # of physical layout.
     # Start from block_id=1 since block_id=0 is considered the null block
     start_block_idx = 1
     for i in range(batch_size):
         k_context, v_context = k_contexts[i], v_contexts[i]
-        start = start_block_idx * block_size
-        end = start + k_context.shape[0]
-        kv_cache_flat[start:end, :, :head_size] = k_context
-        kv_cache_flat[start:end, :, head_size:] = v_context
-
+        t = torch.arange(k_context.shape[0], device=device)
+        blk = start_block_idx + t // block_size
+        off = t % block_size
+        # Advanced indexing on (blk, off) yields [T, H, hs] destinations.
+        # index_put is dtype-strict; cast like the scalar path would.
+        kv_cache[blk, :, off, :head_size] = k_context.to(kv_cache.dtype)
+        kv_cache[blk, :, off, head_size:] = v_context.to(kv_cache.dtype)
         # Stay block aligned and allocate enough blocks for the new tokens
         start_block_idx += cdiv(int(seq_lens[i]), block_size)
 
@@ -214,9 +230,6 @@ def create_and_prepopulate_kv_cache(
             i, block_indices
         ] * block_size + token_inter_block_offsets.to(device)
 
-    # Transpose to logical (num_blocks, num_kv_heads, block_size, 2*hs)
-    kv_cache = kv_cache.transpose(1, 2).contiguous()
-
     if fp8_kv_cache:
         kv_cache = kv_cache.view(torch.uint8)
 
@@ -236,6 +249,22 @@ class MockAttentionLayer:
         self._v_scale_float = 1.0
 
 
+def _clone_kv_cache_in_layout(
+    kv_cache: torch.Tensor, layout: KVCacheLayout
+) -> torch.Tensor:
+    """Copy a logical [B, H, N, C] cache into a fresh allocation in `layout`."""
+    logical_5d = (1, *kv_cache.shape)
+    physical = torch.zeros(
+        tuple(logical_5d[i] for i in layout.stride_order),
+        dtype=kv_cache.dtype,
+        device=kv_cache.device,
+    )
+    inv_order = [layout.stride_order.index(i) for i in range(5)]
+    view = physical.permute(*inv_order)[0]
+    view.copy_(kv_cache)
+    return view
+
+
 def run_attention_backend(
     backend: AttentionBackendEnum,
     kv_cache_spec: FullAttentionSpec,
@@ -254,18 +283,17 @@ def run_attention_backend(
 ) -> torch.Tensor:
     """Run attention computation using the specified backend's AttentionImpl."""
 
-    # Handle special case for FLEX_ATTENTION_SLOW
-    actual_backend = backend
-
-    use_direct_block_mask = is_torch_equal_or_newer("2.9.0.dev0")
+    use_direct_block_mask = not current_platform.is_rocm() and is_torch_equal_or_newer(
+        "2.9.0.dev0"
+    )
     if backend == "FLEX_ATTENTION_SLOW":
-        actual_backend = AttentionBackendEnum.FLEX_ATTENTION
         use_direct_block_mask = False
+    backend = _actual_backend(backend)
 
-    builder_cls, impl_cls = try_get_attention_backend(actual_backend)
+    builder_cls, impl_cls = try_get_attention_backend(backend)
 
     # Mock flashinfer's get_per_layer_parameters if needed
-    if actual_backend == AttentionBackendEnum.FLASHINFER:
+    if backend == AttentionBackendEnum.FLASHINFER:
         import unittest.mock
 
         from vllm.v1.attention.backends.utils import PerLayerParameters
@@ -295,7 +323,7 @@ def run_attention_backend(
     else:
         # Build metadata
         builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
-        if actual_backend == AttentionBackendEnum.FLEX_ATTENTION:
+        if backend == AttentionBackendEnum.FLEX_ATTENTION:
             builder.direct_build = use_direct_block_mask
         attn_metadata = builder.build(
             common_prefix_len=0,
@@ -311,17 +339,19 @@ def run_attention_backend(
     )
     head_size = vllm_config.model_config.get_head_size()
     scale = 1.0 / (head_size**0.5)
-    impl = impl_cls(
-        num_heads=num_heads,
-        head_size=head_size,
-        scale=scale,
-        num_kv_heads=num_kv_heads,
-        alibi_slopes=None,
-        sliding_window=sliding_window,
-        attn_type=attn_type,
-        kv_cache_dtype=kv_cache_dtype,
-        **({"sinks": sinks} if sinks is not None else {}),
-    )
+    # Impls capture the current vllm config at construction, as in model loading.
+    with set_current_vllm_config(vllm_config):
+        impl = impl_cls(
+            num_heads=num_heads,
+            head_size=head_size,
+            scale=scale,
+            num_kv_heads=num_kv_heads,
+            alibi_slopes=None,
+            sliding_window=sliding_window,
+            attn_type=attn_type,
+            kv_cache_dtype=kv_cache_dtype,
+            **({"sinks": sinks} if sinks is not None else {}),
+        )
 
     # Create mock layer and output buffer
     mock_layer = MockAttentionLayer(device)
@@ -333,7 +363,7 @@ def run_attention_backend(
     # Run forward pass
     # NOTE: The query, key, and value are already shaped correctly
     # in the calling test function.
-    if not try_backend_includes_kv_cache_update(actual_backend):
+    if not try_backend_includes_kv_cache_update(backend):
         impl.do_kv_cache_update(
             mock_layer, key, value, kv_cache, attn_metadata.slot_mapping
         )
@@ -347,7 +377,7 @@ def run_attention_backend(
 def _test_backend_correctness(
     batch_spec: BatchSpec,
     model: str,
-    backend_to_test: list[AttentionBackendEnum | str],
+    backend_to_test: list[AttentionBackendEnum],
     mask_mod,
     *,
     causal: bool = True,
@@ -358,6 +388,7 @@ def _test_backend_correctness(
     tensor_parallel_size: int = 1,
     kv_cache_dtype: str = "auto",
     use_sinks: bool = False,
+    layout: KVCacheLayout | None = None,
 ):
     """
     Test that all backends produce similar outputs to a reference implementation
@@ -535,6 +566,22 @@ def _test_backend_correctness(
     common_attn_metadata.causal = causal
 
     # 3. Simulate Paged KV Cache and a realistic slot_mapping
+    # Mirror selector-time resolution locally (caller-requested layout, then
+    # any backend-required layout, then the default chain).
+    declared = [
+        supported
+        for backend in backend_to_test
+        if (
+            supported := _actual_backend(backend)
+            .get_class()
+            .supported_kv_cache_layouts()
+        )
+        is not None
+    ]
+    if layout is None:
+        # Mirror the resolver: the shared cache uses the declared sets' preferred
+        # layout; backends that don't support it get a per-backend copy below.
+        layout = min(declared, key=len)[0] if declared else KVCacheLayout.LBNHC
     kv_cache = create_and_prepopulate_kv_cache(
         k_contexts=k_contexts,
         v_contexts=v_contexts,
@@ -545,6 +592,7 @@ def _test_backend_correctness(
         device=device,
         num_blocks=vllm_config.cache_config.num_gpu_blocks or 1000,
         common_attn_metadata=common_attn_metadata,
+        layout=layout,
         randomize_blocks=True,
         kv_cache_dtype=kv_cache_dtype,
     )
@@ -553,62 +601,44 @@ def _test_backend_correctness(
     # Note: flex_attention has known Triton kernel compatibility issues
     # with test infrastructures
     for backend_name in backend_to_test:
-        reset_kv_cache_layout = False
-
-        # Resolve backend class for both enum and string names.
-        actual_backend = backend_name
-        if backend_name == "FLEX_ATTENTION_SLOW":
-            actual_backend = AttentionBackendEnum.FLEX_ATTENTION
-        if hasattr(actual_backend, "get_class"):
-            backend_cls = actual_backend.get_class()
-        else:
-            backend_cls = None
+        backend_cls = _actual_backend(backend_name).get_class()
 
         if is_quantized_kv_cache(kv_cache_dtype) and (
-            backend_cls is None
-            or not backend_cls.supports_kv_cache_dtype(kv_cache_dtype)
+            not backend_cls.supports_kv_cache_dtype(kv_cache_dtype)
         ):
             continue
 
-        if backend_name == AttentionBackendEnum.FLASHINFER:
-            set_kv_cache_layout("HND")
-            reset_kv_cache_layout = True
-
         kv_cache_for_backend = kv_cache
-        if backend_cls is not None:
-            try:
-                stride_order = backend_cls.get_kv_cache_stride_order()
-            except (AttributeError, NotImplementedError):
-                stride_order = tuple(range(kv_cache.ndim))
-            if stride_order != tuple(range(kv_cache.ndim)):
-                # Apply stride order like runtime does in
-                # _reshape_kv_cache (attn_utils.py:182-210): permute to physical
-                # layout, make contiguous, then permute to logical layout.
-                inv_order = [stride_order.index(i) for i in range(len(stride_order))]
-                kv_cache_for_backend = (
-                    kv_cache.permute(*stride_order).contiguous().permute(*inv_order)
-                )
+        backend_layout = layout
 
-        try:
-            backend_output = run_attention_backend(
-                backend_name,
-                kv_cache_spec,
-                ["placeholder"],
-                vllm_config,
-                device,
-                common_attn_metadata,
-                query_vllm,
-                key_vllm,
-                value_vllm,
-                kv_cache_for_backend,
-                sliding_window=sliding_window,
-                attn_type=attn_type,
-                kv_cache_dtype=kv_cache_dtype,
-                sinks=sinks,
-            )
-        finally:
-            if reset_kv_cache_layout:
-                set_kv_cache_layout(None)
+        backend_supported = backend_cls.supported_kv_cache_layouts()
+        if backend_supported is not None and backend_layout not in backend_supported:
+            # Production resolution never pairs this backend with the shared layout
+            # (e.g. flex is LBNHC-only next to FlashInfer's LBHNC); give it a copy
+            # of the cache in a layout it supports.
+            backend_layout = backend_supported[0]
+            kv_cache_for_backend = _clone_kv_cache_in_layout(kv_cache, backend_layout)
+
+        # FlashInfer reads the layout at plan time; set it to match
+        # the physical order of the test cache.
+        vllm_config.cache_config.kv_cache_layout = backend_layout.name
+
+        backend_output = run_attention_backend(
+            backend_name,
+            kv_cache_spec,
+            ["placeholder"],
+            vllm_config,
+            device,
+            common_attn_metadata,
+            query_vllm,
+            key_vllm,
+            value_vllm,
+            kv_cache_for_backend,
+            sliding_window=sliding_window,
+            attn_type=attn_type,
+            kv_cache_dtype=kv_cache_dtype,
+            sinks=sinks,
+        )
 
         # Check shape and dtype consistency
         assert backend_output.shape == sdpa_output.shape, (
@@ -637,6 +667,38 @@ def _test_backend_correctness(
         )
 
 
+@pytest.mark.parametrize("layout", ["BLHNC", "BHLNC"])
+@pytest.mark.parametrize("batch_spec_name", ["small_decode", "small_prefill"])
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
+def test_flashinfer_cross_layer_layout(
+    default_vllm_config,
+    layout: str,
+    batch_spec_name: str,
+    kv_cache_dtype: str,
+):
+    if AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST:
+        pytest.skip("FlashInfer is not installed")
+
+    def causal_mask_mod(
+        b: torch.Tensor,
+        h: torch.Tensor,
+        q_idx: torch.Tensor,
+        kv_idx: torch.Tensor,
+        *,
+        context_len: int,
+    ):
+        return (q_idx + context_len) >= kv_idx
+
+    _test_backend_correctness(
+        batch_spec=BATCH_SPECS[batch_spec_name],
+        model="meta-llama/Meta-Llama-3-8B",
+        backend_to_test=[AttentionBackendEnum.FLASHINFER],
+        mask_mod=causal_mask_mod,
+        kv_cache_dtype=kv_cache_dtype,
+        layout=KVCacheLayout[layout],
+    )
+
+
 @pytest.mark.parametrize(
     "batch_spec_name",
     [
@@ -984,14 +1046,12 @@ if current_platform.is_rocm():
     SLIDING_WINDOW_BACKENDS_TO_TEST = [
         AttentionBackendEnum.FLEX_ATTENTION,
         AttentionBackendEnum.TRITON_ATTN,
-        "FLEX_ATTENTION_SLOW",
     ]
 else:
     SLIDING_WINDOW_BACKENDS_TO_TEST = [
         AttentionBackendEnum.FLASH_ATTN,
         AttentionBackendEnum.FLEX_ATTENTION,
         AttentionBackendEnum.TRITON_ATTN,
-        "FLEX_ATTENTION_SLOW",
     ]
 
 
@@ -1009,7 +1069,10 @@ else:
 @pytest.mark.parametrize("model", ["microsoft/Phi-tiny-MoE-instruct"])
 @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
 def test_sliding_window_backend_correctness(
-    default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int
+    default_vllm_config,
+    batch_spec_name: str,
+    model: str,
+    tensor_parallel_size: int,
 ):
     """Test backend's correctness with sliding window attention."""
 
@@ -1071,7 +1134,10 @@ def test_sliding_window_backend_correctness(
 @pytest.mark.parametrize("model", ["google/embeddinggemma-300m"])
 @pytest.mark.parametrize("tensor_parallel_size", [1, 2])
 def test_sliding_window_encoder_backend_correctness(
-    default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int
+    default_vllm_config,
+    batch_spec_name: str,
+    model: str,
+    tensor_parallel_size: int,
 ):
     """Test backend's correctness with sliding window attention."""
 
@@ -1107,7 +1173,6 @@ def test_sliding_window_encoder_backend_correctness(
 NON_CAUSAL_BACKENDS_TO_TEST = [
     AttentionBackendEnum.FLASH_ATTN,
     AttentionBackendEnum.FLEX_ATTENTION,
-    "FLEX_ATTENTION_SLOW",
 ]
 
 if current_platform.is_rocm():
@@ -1128,7 +1193,9 @@ if current_platform.is_rocm():
 )
 @pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])
 def test_non_causal_backend_correctness(
-    default_vllm_config, batch_spec_name: str, model: str
+    default_vllm_config,
+    batch_spec_name: str,
+    model: str,
 ):
     """Test backend's correctness with non-causal (bidirectional) decoder
     attention, as used by DFlash speculative decoding."""

tests/v1/attention/test_group_head_counts.py

@@ -1,6 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Attention metadata geometry must come from the builder's own group."""
+"""Scheduler metadata sizes a scratchpad from the query head count, so it must
+come from the builder's own group: the model-wide ``get_num_attention_heads()``
+is wrong for models that vary it per layer (e.g. Laguna), and too small a
+scratchpad is indexed past its end.
+"""
 
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
@@ -13,9 +17,8 @@ from vllm.v1.attention.backends.cpu_attn import (
     CPUAttentionBackendImpl,
     CPUAttentionMetadataBuilder,
 )
-from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadataBuilder
 
-requires_cpu = pytest.mark.skipif(
+pytestmark = pytest.mark.skipif(
     not current_platform.is_cpu(), reason="CPU attention backend"
 )
 
@@ -66,7 +69,6 @@ def _build(layer_num_heads: list[int]) -> CPUAttentionMetadataBuilder:
         )
 
 
-@requires_cpu
 @pytest.mark.parametrize("group_num_heads", [MODEL_WIDE_NUM_HEADS, 64, 16])
 def test_num_heads_comes_from_the_group(group_num_heads):
     """The group's own count wins, even when it is not the model-wide one."""
@@ -74,56 +76,7 @@ def test_num_heads_comes_from_the_group(group_num_heads):
     assert builder.num_heads == group_num_heads
 
 
-@requires_cpu
 def test_mixed_head_counts_in_one_group_are_rejected():
     """Grouping guarantees uniformity; a mixed group means that broke."""
     with pytest.raises(AssertionError, match="share num_heads"):
         _build([MODEL_WIDE_NUM_HEADS, 64])
-
-
-def test_flash_attention_geometry_comes_from_the_group():
-    """FA3 must use group geometry for layers without an ``impl`` wrapper."""
-    layers = {f"layer_{i}": SimpleNamespace(num_heads=16) for i in range(2)}
-    vllm_config = MagicMock()
-    vllm_config.model_config.get_num_attention_heads.return_value = MODEL_WIDE_NUM_HEADS
-    vllm_config.model_config.get_num_kv_heads.return_value = NUM_KV_HEADS
-    vllm_config.model_config.get_head_size.return_value = 128
-    vllm_config.model_config.rswa_window = None
-    vllm_config.model_config.is_mm_prefix_lm = False
-    vllm_config.parallel_config.cp_kv_cache_interleave_size = 1
-    vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs.return_value = (
-        False
-    )
-    vllm_config.compilation_config.max_cudagraph_capture_size = None
-    kv_cache_spec = SimpleNamespace(
-        block_size=16,
-        num_kv_heads=2,
-        head_size=64,
-        dtype=torch.bfloat16,
-    )
-
-    with (
-        patch(
-            "vllm.v1.attention.backends.utils.get_layers_from_vllm_config",
-            return_value=layers,
-        ),
-        patch(
-            "vllm.distributed.parallel_state.get_dcp_group",
-            side_effect=AssertionError,
-        ),
-        patch(
-            "vllm.v1.attention.backends.flash_attn.get_flash_attn_version",
-            return_value=3,
-        ),
-    ):
-        builder = FlashAttentionMetadataBuilder(
-            kv_cache_spec=kv_cache_spec,
-            layer_names=list(layers),
-            vllm_config=vllm_config,
-            device=torch.device("cpu"),
-        )
-
-    assert builder.aot_schedule
-    assert builder.num_heads_q == 16
-    assert builder.num_heads_kv == 2
-    assert builder.headdim == 64

tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py

@@ -34,19 +34,19 @@ def test_indexer_warmup_normalizes_zero_compress_ratios():
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
-def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_storage_block_size():
+def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_num_states():
     """Regression test: DeepseekV4 compression path must compute slot_mapping from
     compressed positions, not reuse the uncompressed common metadata mapping.
     """
     device = torch.device("cuda")
 
-    # storage_block_size = block_size // compress_ratio = 256 // 4 = 64
+    # num_states = block_size // tokens_per_state = 256 // 4 = 64
     kv_cache_spec = MLAAttentionSpec(
         block_size=256,
         num_kv_heads=1,
         head_size=128,
         dtype=torch.bfloat16,
-        compress_ratio=4,
+        tokens_per_state=4,
     )
     vllm_config = create_vllm_config(max_model_len=1024)
     max_num_blocks = kv_cache_spec.max_num_blocks_per_req(vllm_config, 1024)
@@ -95,7 +95,7 @@ def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_storage_block_
     valid_slots = md.slot_mapping[md.slot_mapping >= 0]
     assert valid_slots.numel() == 10  # 40 tokens / compress_ratio 4
 
-    storage_bs = kv_cache_spec.storage_block_size  # 64
+    storage_bs = kv_cache_spec.num_states  # 64
     # Compressed positions 60..63 land in block 5, positions 64..69 in block 7.
     expected = torch.tensor(
         [

tests/v1/attention/test_mla_backends.py

@@ -250,6 +250,15 @@ def _prefill_backend_dimension_params():
                         )
                     )
                 )
+            elif (
+                current_platform.is_rocm()
+                and prefill_backend == MLAPrefillBackendEnum.FLASH_ATTN
+            ):
+                # Pre-existing on main: flash-attn MLA prefill returns wrong
+                # outputs and faults on ROCm, so it fails on the base commit too.
+                marks.append(
+                    pytest.mark.skip(reason="FLASH_ATTN MLA prefill broken on ROCm")
+                )
             params.append(
                 pytest.param(
                     prefill_backend,
@@ -392,7 +401,7 @@ def create_and_prepopulate_kv_cache(
             kv_entry_size = head_size
 
         kv_cache = torch.zeros(
-            num_blocks, block_size, kv_entry_size, dtype=torch.uint8, device=device
+            num_blocks, 1, block_size, kv_entry_size, dtype=torch.uint8, device=device
         )
         scale_tensor = (
             scale
@@ -401,9 +410,9 @@ def create_and_prepopulate_kv_cache(
         )
         scale_tensor = scale_tensor.to(device=device, dtype=torch.float32)
     else:
-        # Create MLA KV cache: (num_blocks, block_size, head_size)
+        # Create MLA KV cache: (num_blocks, num_heads=1, block_size, head_size)
         kv_cache = torch.zeros(
-            num_blocks, block_size, head_size, dtype=dtype, device=device
+            num_blocks, 1, block_size, head_size, dtype=dtype, device=device
         )
         kv_cache_flat = kv_cache.view(-1, head_size)
 
@@ -424,7 +433,7 @@ def create_and_prepopulate_kv_cache(
             ops.concat_and_cache_mla(
                 kv_c_context,
                 k_pe_context.squeeze(1),
-                kv_cache,
+                kv_cache.squeeze(1),
                 slots,
                 kv_cache_dtype=kv_cache_dtype,
                 scale=scale_tensor,
@@ -543,6 +552,10 @@ class MockSparseMLAAttentionLayer:
         kv_cache_dtype = getattr(self.impl, "kv_cache_dtype", "auto")
         fp8_attention = kv_cache_dtype.startswith("fp8")
 
+        # Impls see the bind-time-squeezed [B, N, C] cache; mirror bind_kv_cache.
+        if kv_cache.ndim == 4:
+            kv_cache = kv_cache.squeeze(1)
+
         # Write to KV cache
         if kv_cache.numel() > 0:
             ops.concat_and_cache_mla(
@@ -678,6 +691,10 @@ class MockMLAAttentionLayer(MLAAttention):
         output: torch.Tensor,
     ) -> torch.Tensor:
         """Replicates MLAAttention.forward_impl logic for testing."""
+        # Impls see the bind-time-squeezed [B, N, C] cache; mirror bind_kv_cache.
+        if kv_cache.ndim == 4:
+            kv_cache = kv_cache.squeeze(1)
+
         # Write to KV cache
         kv_cache_dtype = getattr(self.impl, "kv_cache_dtype", "auto")
         fp8_attention = kv_cache_dtype.startswith("fp8")

tests/v1/attention/test_mm_prefix.py

@@ -20,7 +20,10 @@ import pytest
 import torch
 
 from vllm.platforms import current_platform
-from vllm.v1.attention.backends.utils import fill_mm_prefix_query_ranges
+from vllm.v1.attention.backends.utils import (
+    fill_mm_prefix_query_ranges,
+)
+from vllm.v1.kv_cache_interface import KVCacheLayout
 
 
 def _fa4_available() -> bool:
@@ -585,6 +588,7 @@ def test_mm_prefix_kv_cache_path(head_size: int):
         device=DEVICE,
         num_blocks=2048,
         common_attn_metadata=common,
+        layout=KVCacheLayout.LBNHC,
         randomize_blocks=True,
     )
 

tests/v1/attention/test_sparse_mla_backends.py

@@ -813,7 +813,7 @@ def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens():
     assert attn_metadata.req_id_per_token.shape[0] == num_batch_tokens
 
     q = torch.zeros(num_mqa_tokens, 4, 576, dtype=torch.bfloat16, device=device)
-    kv_cache = torch.zeros(40 * block_size, 576, dtype=torch.bfloat16, device=device)
+    kv_cache = torch.zeros(40, block_size, 576, dtype=torch.bfloat16, device=device)
     topk_indices = torch.randint(
         0,
         block_size * 10,
@@ -1157,7 +1157,8 @@ def test_sparse_backend_prefill_correctness(
             q=query_cat,
             kv_c_normed=kv_c_cat,
             k_pe=k_pe_cat,
-            kv_c_and_k_pe_cache=kv_cache,
+            # Impls see the bind-time-squeezed [B, N, C] cache; mirror bind_kv_cache.
+            kv_c_and_k_pe_cache=kv_cache.squeeze(1),
             attn_metadata=metadata,
             k_scale=torch.tensor(1.0, device=device),
             output=out_buffer,

tests/v1/attention/test_trtllm_attention_integration.py

@@ -9,6 +9,7 @@ import pytest
 import torch
 from torch.nn.attention.flex_attention import create_block_mask, flex_attention
 
+from tests.v1.attention.test_attention_backends import create_and_prepopulate_kv_cache
 from tests.v1.attention.utils import (
     BatchSpec,
     create_common_attn_metadata,
@@ -16,14 +17,9 @@ from tests.v1.attention.utils import (
 )
 from vllm.config import set_current_vllm_config
 from vllm.platforms import current_platform
-from vllm.utils.math_utils import cdiv
 from vllm.utils.torch_utils import nvfp4_kv_cache_full_dim, set_random_seed
-from vllm.v1.attention.backends.utils import (
-    PerLayerParameters,
-    get_kv_cache_layout,
-    set_kv_cache_layout,
-)
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode
+from vllm.v1.attention.backends.utils import PerLayerParameters
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout, KVQuantMode
 
 if not current_platform.is_device_capability_family(100):
     pytest.skip(
@@ -87,96 +83,6 @@ def _mock_get_per_layer_parameters(vllm_config, layer_names, impl_cls):
     }
 
 
-def _create_hnd_kv_cache(
-    k_contexts,
-    v_contexts,
-    block_size,
-    num_kv_heads,
-    head_size,
-    dtype,
-    device,
-    num_blocks,
-    common_attn_metadata,
-    kv_in_head_dim=False,
-):
-    """Create and populate a packed KV cache with HND-compatible strides.
-
-    When kv_in_head_dim=False (default), returns (B, H, N, 2*hs) with K/V
-    packed in the content dim. When kv_in_head_dim=True, returns
-    (B, 2*H, N, hs) with K/V as separate head groups.
-    """
-    seq_lens = common_attn_metadata.seq_lens.cpu()
-    query_lens = (
-        common_attn_metadata.query_start_loc_cpu[1:]
-        - common_attn_metadata.query_start_loc_cpu[:-1]
-    )
-    block_table = common_attn_metadata.block_table_tensor
-    slot_mapping = common_attn_metadata.slot_mapping
-    batch_size = len(k_contexts)
-
-    # kv_in_head_dim: (B, N, 2*H, hs) — K/V as separate head groups
-    # else:           (B, N, H, 2*hs) — K/V packed in content dim
-    n_heads, content = (
-        (2 * num_kv_heads, head_size)
-        if kv_in_head_dim
-        else (num_kv_heads, 2 * head_size)
-    )
-    kv_cache = torch.zeros(
-        num_blocks,
-        block_size,
-        n_heads,
-        content,
-        dtype=dtype,
-        device=device,
-    )
-    kv_cache_flat = kv_cache.view(-1, n_heads, content)
-
-    start_block_idx = 1
-    for i in range(batch_size):
-        k_ctx, v_ctx = k_contexts[i], v_contexts[i]
-        start = start_block_idx * block_size
-        end = start + k_ctx.shape[0]
-        if kv_in_head_dim:
-            kv_cache_flat[start:end, :num_kv_heads] = k_ctx
-            kv_cache_flat[start:end, num_kv_heads:] = v_ctx
-        else:
-            kv_cache_flat[start:end, :, :head_size] = k_ctx
-            kv_cache_flat[start:end, :, head_size:] = v_ctx
-        start_block_idx += cdiv(int(seq_lens[i]), block_size)
-
-    blocks_end = start_block_idx
-
-    # Randomly permute blocks (starting from block 1; block 0 is null).
-    perm = torch.randperm(blocks_end - 1) + 1
-    inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
-    inv_perm[1:] = torch.argsort(perm) + 1
-    kv_cache[1:blocks_end] = kv_cache[perm]
-
-    # Build block table.
-    start_block_idx = 1
-    for i in range(batch_size):
-        n_blocks = cdiv(int(seq_lens[i]), block_size)
-        block_table[i, :n_blocks] = inv_perm[
-            start_block_idx : start_block_idx + n_blocks
-        ]
-        start_block_idx += n_blocks
-
-    # Build slot mapping that is consistent with the block table.
-    for i in range(batch_size):
-        ctx_len = int(seq_lens[i]) - int(query_lens[i])
-        token_offsets = torch.arange(int(query_lens[i])) + ctx_len
-        block_indices = token_offsets // block_size
-        intra_block_offsets = token_offsets % block_size
-        start = common_attn_metadata.query_start_loc_cpu[i]
-        end = common_attn_metadata.query_start_loc_cpu[i + 1]
-        slot_mapping[start:end] = block_table[
-            i, block_indices
-        ] * block_size + intra_block_offsets.to(device)
-
-    # Transpose to canonical: (B, H, N, 2*hs) or (B, 2*H, N, hs)
-    return kv_cache.transpose(1, 2).contiguous()
-
-
 def _create_nvfp4_hnd_kv_cache(
     k_contexts,
     v_contexts,
@@ -191,7 +97,7 @@ def _create_nvfp4_hnd_kv_cache(
 ):
     """Create an nvfp4 KV cache by quantizing bf16 context via
     reshape_and_cache_flash, using the same block-table layout as
-    _create_hnd_kv_cache.
+    create_and_prepopulate_kv_cache.
 
     The returned tensor is dtype ``uint8`` with head-group layout
       ``(num_blocks, 2 * num_kv_heads, block_size, full_dim)``
@@ -217,11 +123,10 @@ def _create_nvfp4_hnd_kv_cache(
             during quantization.
 
     Returns:
-        ``torch.Tensor``: The nvfp4 kv_cache tensor (uint8, HND-strided).
+        ``torch.Tensor``: The nvfp4 kv_cache tensor (uint8, LBHNC-strided).
     """
-    # First create a bf16 HND cache so block tables are populated.
-    # Use kv_in_head_dim=True so K/V are separate head groups (B, 2*H, N, hs).
-    bf16_cache = _create_hnd_kv_cache(
+    # First create a bf16 cache so block tables are populated.
+    bf16_cache = create_and_prepopulate_kv_cache(
         k_contexts,
         v_contexts,
         block_size,
@@ -231,7 +136,7 @@ def _create_nvfp4_hnd_kv_cache(
         device,
         num_blocks,
         common_attn_metadata,
-        kv_in_head_dim=True,
+        layout=KVCacheLayout.LBHNC,
     )
 
     # (num_blocks, 2 * num_kv_heads, block_size, full_dim) — K heads first, then V heads
@@ -244,7 +149,7 @@ def _create_nvfp4_hnd_kv_cache(
     k_cache, v_cache = nvfp4_cache.split(num_kv_heads, dim=1)
 
     # Flatten bf16 context into tokens and quantize via reshape_and_cache_flash.
-    # bf16_cache is (B, 2*H, N, hs); split K/V on head dim.
+    # bf16_cache is (B, H, N, 2*hs); split K/V on the content dim.
     block_table = common_attn_metadata.block_table_tensor
     seq_lens = common_attn_metadata.seq_lens.cpu()
     query_lens = (
@@ -260,10 +165,16 @@ def _create_nvfp4_hnd_kv_cache(
         # Gather context tokens from the bf16 cache using block table.
         n_ctx_blocks = (ctx_len + block_size - 1) // block_size
         blocks = block_table[i, :n_ctx_blocks]
-        # bf16_cache is (B, 2*H, N, hs); split K and V head groups.
-        k_bf16, v_bf16 = bf16_cache[blocks].split(num_kv_heads, dim=1)
-        k_ctx = k_bf16.transpose(1, 2).reshape(-1, num_kv_heads, head_size)[:ctx_len]
-        v_ctx = v_bf16.transpose(1, 2).reshape(-1, num_kv_heads, head_size)[:ctx_len]
+        k_ctx = (
+            bf16_cache[blocks, :, :, :head_size]
+            .transpose(1, 2)
+            .reshape(-1, num_kv_heads, head_size)[:ctx_len]
+        )
+        v_ctx = (
+            bf16_cache[blocks, :, :, head_size:]
+            .transpose(1, 2)
+            .reshape(-1, num_kv_heads, head_size)[:ctx_len]
+        )
         # Build slot mapping for these context tokens.
         token_offsets = torch.arange(ctx_len, device=device)
         block_indices = token_offsets // block_size
@@ -383,7 +294,7 @@ def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL)
         )
     else:
         kv_scale_val = 1.0
-        kv_cache = _create_hnd_kv_cache(
+        kv_cache = create_and_prepopulate_kv_cache(
             k_contexts,
             v_contexts,
             BLOCK_SIZE,
@@ -393,127 +304,117 @@ def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL)
             device,
             NUM_GPU_BLOCKS,
             common_attn_metadata,
+            layout=KVCacheLayout.LBHNC,
         )
 
     # 3. Run through FlashInfer with TRTLLM enabled
-    set_kv_cache_layout("HND")
-    get_kv_cache_layout.cache_clear()
+    vllm_config.cache_config.kv_cache_layout = "LBHNC"
 
-    try:
-        is_nvfp4 = kv_cache_dtype == "nvfp4"
-        kv_quant_mode = KVQuantMode.NVFP4 if is_nvfp4 else KVQuantMode.NONE
-        spec_dtype = torch.uint8 if is_nvfp4 else dtype
-        kv_cache_spec = FullAttentionSpec(
-            block_size=BLOCK_SIZE,
-            num_kv_heads=num_kv_heads,
-            head_size=head_size,
-            dtype=spec_dtype,
-            kv_quant_mode=kv_quant_mode,
+    is_nvfp4 = kv_cache_dtype == "nvfp4"
+    kv_quant_mode = KVQuantMode.NVFP4 if is_nvfp4 else KVQuantMode.NONE
+    spec_dtype = torch.uint8 if is_nvfp4 else dtype
+    kv_cache_spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=num_kv_heads,
+        head_size=head_size,
+        dtype=spec_dtype,
+        kv_quant_mode=kv_quant_mode,
+    )
+    layer_names = ["test_layer_0"]
+
+    with (
+        set_current_vllm_config(vllm_config),
+        unittest.mock.patch(
+            "vllm.utils.flashinfer.supports_trtllm_attention",
+            return_value=True,
+        ),
+        unittest.mock.patch(
+            "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
+            _mock_get_per_layer_parameters,
+        ),
+    ):
+        builder = FlashInferMetadataBuilder(
+            kv_cache_spec, layer_names, vllm_config, device
+        )
+        attn_metadata = builder.build(
+            common_prefix_len=0,
+            common_attn_metadata=common_attn_metadata,
         )
-        layer_names = ["test_layer_0"]
 
-        with (
-            set_current_vllm_config(vllm_config),
-            unittest.mock.patch(
-                "vllm.utils.flashinfer.supports_trtllm_attention",
-                return_value=True,
-            ),
-            unittest.mock.patch(
-                "vllm.v1.attention.backends.flashinfer.get_per_layer_parameters",
-                _mock_get_per_layer_parameters,
-            ),
-        ):
-            builder = FlashInferMetadataBuilder(
-                kv_cache_spec, layer_names, vllm_config, device
+        # Verify the correct TRTLLM metadata types were produced.
+        has_prefills = any(ql > 1 for ql in batch_spec.query_lens)
+        has_decodes = any(ql == 1 for ql in batch_spec.query_lens)
+
+        if has_prefills:
+            assert isinstance(attn_metadata.prefill, TRTLLMPrefill), (
+                f"Expected TRTLLMPrefill, got {type(attn_metadata.prefill)}"
             )
-            attn_metadata = builder.build(
-                common_prefix_len=0,
-                common_attn_metadata=common_attn_metadata,
+        if has_decodes:
+            assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode), (
+                f"Expected FlashInferTrtllmAPIDecode, got {type(attn_metadata.decode)}"
             )
+            assert attn_metadata.decode.kernel == FlashInferDecodeKernel.TRTLLM_GEN
 
-            # Verify the correct TRTLLM metadata types were produced.
-            has_prefills = any(ql > 1 for ql in batch_spec.query_lens)
-            has_decodes = any(ql == 1 for ql in batch_spec.query_lens)
+        impl = FlashInferImpl(
+            num_heads=num_q_heads,
+            head_size=head_size,
+            scale=scale,
+            num_kv_heads=num_kv_heads,
+            alibi_slopes=None,
+            sliding_window=None,
+            kv_cache_dtype=kv_cache_dtype,
+        )
 
-            if has_prefills:
-                assert isinstance(attn_metadata.prefill, TRTLLMPrefill), (
-                    f"Expected TRTLLMPrefill, got {type(attn_metadata.prefill)}"
-                )
-            if has_decodes:
-                assert isinstance(attn_metadata.decode, FlashInferTrtllmAPIDecode), (
-                    "Expected FlashInferTrtllmAPIDecode, got "
-                    f"{type(attn_metadata.decode)}"
-                )
-                assert attn_metadata.decode.kernel == FlashInferDecodeKernel.TRTLLM_GEN
-
-            impl = FlashInferImpl(
-                num_heads=num_q_heads,
-                head_size=head_size,
-                scale=scale,
-                num_kv_heads=num_kv_heads,
-                alibi_slopes=None,
-                sliding_window=None,
-                kv_cache_dtype=kv_cache_dtype,
-            )
-
-            mock_layer = MockAttentionLayer(device)
-            if is_nvfp4:
-                # For nvfp4, k_scale/v_scale are the global quantization
-                # scales (amax/448) used by reshape_and_cache_flash.
-                kv_scale_t = torch.tensor(
-                    kv_scale_val, dtype=torch.float32, device=device
-                )
-                mock_layer._k_scale = kv_scale_t
-                mock_layer._v_scale = kv_scale_t
-                mock_layer._k_scale_float = kv_scale_val
-                mock_layer._v_scale_float = kv_scale_val
-            output = torch.empty_like(query_vllm)
-
-            impl.do_kv_cache_update(
-                mock_layer,
-                key_vllm,
-                value_vllm,
-                kv_cache,
-                attn_metadata.slot_mapping,
-            )
-
-            # nvfp4 trtllm kernel requires FP8 queries. In the real
-            # pipeline the attention layer handles this; here we
-            # quantize manually.
-            if is_nvfp4:
-                finfo = torch.finfo(torch.float8_e4m3fn)
-                q_amax = query_vllm.abs().amax().clamp(min=1e-12)
-                q_s = (finfo.max / q_amax * 0.1).item()
-                query_vllm = (
-                    (query_vllm * q_s)
-                    .clamp(finfo.min, finfo.max)
-                    .to(torch.float8_e4m3fn)
-                )
-                mock_layer._q_scale = torch.tensor(
-                    1.0 / q_s, dtype=torch.float32, device=device
-                )
-                mock_layer._q_scale_float = 1.0 / q_s
-
-            output = impl.forward(
-                mock_layer,
-                query_vllm,
-                key_vllm,
-                value_vllm,
-                kv_cache,
-                attn_metadata,
-                output=output,
-            )
-
-        # 4. Compare against SDPA reference
+        mock_layer = MockAttentionLayer(device)
         if is_nvfp4:
-            atol, rtol = 1.0, 1.0  # nvfp4 has higher quantization error
-        else:
-            atol, rtol = 1e-2, 1e-2
-        torch.testing.assert_close(output, sdpa_output, atol=atol, rtol=rtol)
+            # For nvfp4, k_scale/v_scale are the global quantization
+            # scales (amax/448) used by reshape_and_cache_flash.
+            kv_scale_t = torch.tensor(kv_scale_val, dtype=torch.float32, device=device)
+            mock_layer._k_scale = kv_scale_t
+            mock_layer._v_scale = kv_scale_t
+            mock_layer._k_scale_float = kv_scale_val
+            mock_layer._v_scale_float = kv_scale_val
+        output = torch.empty_like(query_vllm)
 
-    finally:
-        set_kv_cache_layout(None)
-        get_kv_cache_layout.cache_clear()
+        impl.do_kv_cache_update(
+            mock_layer,
+            key_vllm,
+            value_vllm,
+            kv_cache,
+            attn_metadata.slot_mapping,
+        )
+
+        # nvfp4 trtllm kernel requires FP8 queries. In the real
+        # pipeline the attention layer handles this; here we
+        # quantize manually.
+        if is_nvfp4:
+            finfo = torch.finfo(torch.float8_e4m3fn)
+            q_amax = query_vllm.abs().amax().clamp(min=1e-12)
+            q_s = (finfo.max / q_amax * 0.1).item()
+            query_vllm = (
+                (query_vllm * q_s).clamp(finfo.min, finfo.max).to(torch.float8_e4m3fn)
+            )
+            mock_layer._q_scale = torch.tensor(
+                1.0 / q_s, dtype=torch.float32, device=device
+            )
+            mock_layer._q_scale_float = 1.0 / q_s
+
+        output = impl.forward(
+            mock_layer,
+            query_vllm,
+            key_vllm,
+            value_vllm,
+            kv_cache,
+            attn_metadata,
+            output=output,
+        )
+
+    # 4. Compare against SDPA reference
+    if is_nvfp4:
+        atol, rtol = 1.0, 1.0  # nvfp4 has higher quantization error
+    else:
+        atol, rtol = 1e-2, 1e-2
+    torch.testing.assert_close(output, sdpa_output, atol=atol, rtol=rtol)
 
 
 @pytest.mark.parametrize(

tests/v1/attention/utils.py

@@ -32,7 +32,12 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.kv_cache_interface import (
     EncoderOnlyAttentionSpec,
     FullAttentionSpec,
+    KVCacheLayout,
+    KVCacheSpec,
+    KVCacheTensor,
     MambaSpec,
+    compute_layout_strides,
+    create_kv_cache_views,
     get_kv_quant_mode,
 )
 
@@ -215,6 +220,9 @@ def create_vllm_config(
     #   (these may be set during initialization normally)
     cache_config.num_gpu_blocks = num_gpu_blocks
     cache_config.num_cpu_blocks = 0
+    # Builders read the resolved layout from the config; tests that exercise a
+    # different layout overwrite this field.
+    cache_config.kv_cache_layout = "LBNHC"
 
     parallel_config = ParallelConfig(
         tensor_parallel_size=tensor_parallel_size,
@@ -417,3 +425,31 @@ class MockMambaBuilder(BaseMambaAttentionMetadataBuilder[BaseMambaAttentionMetad
             is_prefilling=torch.tensor(is_prefilling, dtype=torch.bool)
         )
         return builder.build(0, common_metadata)
+
+
+def dense_kv_cache_views(
+    raw: torch.Tensor,
+    spec: KVCacheSpec,
+    num_blocks: int,
+    num_layers: int,
+    layout: KVCacheLayout,
+    kernel_block_size: int | None = None,
+) -> list[torch.Tensor]:
+    """``create_kv_cache_views`` for a dense allocation of ``num_layers`` layers."""
+    layer_stride, block_stride, _, _, _ = compute_layout_strides(
+        spec, num_blocks, num_layers, layout
+    )
+    tensor = KVCacheTensor(
+        size=raw.numel(),
+        layers=[str(i) for i in range(num_layers)],
+        layer_stride=layer_stride,
+        block_stride=block_stride,
+    )
+    return create_kv_cache_views(
+        raw,
+        spec,
+        num_blocks,
+        layout,
+        tensor,
+        kernel_block_size=kernel_block_size,
+    )

tests/v1/core/test_contiguous_kv_packing.py

@@ -1,380 +1,237 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for contiguous KV cache packing."""
+"""Tests for contiguous KV cache packing.
+
+Every cache group packs its layers densely into one block; groups overlay each other
+(a block ID is owned by one group at a time), so the packed block stride is the
+largest group's packing. The layout decides whether the layer dim sits outside the
+block dim (a contiguous region per layer) or inside it (all layers' pages within each
+block); the allocation is the same either way.
+"""
 
 from unittest.mock import MagicMock
 
 import pytest
 import torch
 
+from vllm.config import CacheConfig
 from vllm.v1.core.kv_cache_utils import (
-    _get_kv_cache_config_packed,
-    _get_kv_cache_groups_uniform_groups,
+    _get_kv_cache_bytes_per_block,
+    _pool_bytes_per_block,
     get_kv_cache_config_from_groups,
 )
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheGroupSpec,
-    KVCacheTensor,
+    KVCacheLayout,
     MLAAttentionSpec,
-    SlidingWindowMLASpec,
-    SlidingWindowSpec,
     UniformTypeKVCacheSpecs,
 )
+from vllm.v1.worker.utils import allocate_kv_cache
+
+MEMORY = 8 * 1024 * 1024
 
 
-def _make_mla_spec(page_size: int, block_size: int = 256) -> MLAAttentionSpec:
+def _mla(head_size: int) -> MLAAttentionSpec:
     return MLAAttentionSpec(
-        block_size=block_size,
-        num_kv_heads=1,
-        head_size=512,
-        dtype=torch.uint8,
-        page_size_padded=page_size,
-        cache_dtype_str="fp8_ds_mla",
-        model_version="deepseek_v4",
-        alignment=576,
+        block_size=64, num_kv_heads=1, head_size=head_size, dtype=torch.uint8
     )
 
 
-def _make_full_spec() -> FullAttentionSpec:
+def _full() -> FullAttentionSpec:
     return FullAttentionSpec(
-        block_size=16,
-        num_kv_heads=2,
-        head_size=64,
-        dtype=torch.float16,
+        block_size=16, num_kv_heads=2, head_size=64, dtype=torch.float16
     )
 
 
-def _make_sw_spec() -> SlidingWindowSpec:
-    return SlidingWindowSpec(
-        block_size=16,
-        num_kv_heads=2,
-        head_size=64,
-        dtype=torch.float16,
-        sliding_window=128,
+def _uniform_group(specs: dict) -> KVCacheGroupSpec:
+    return KVCacheGroupSpec(
+        list(specs),
+        UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=specs),
     )
 
 
-def _make_groups(n_c4, n_c128, n_swa):
-    PS_C4_MLA = 37440
-    PS_C4_IDX = 8640
-    PS_C128 = 1728
-    PS_SWA = 37440
-
-    mla_specs = {}
-    for i in range(n_c4):
-        mla_specs[f"c4_mla.{i}"] = _make_mla_spec(PS_C4_MLA)
-        mla_specs[f"c4_idx.{i}"] = _make_mla_spec(PS_C4_IDX)
-    for i in range(n_c128):
-        mla_specs[f"c128_mla.{i}"] = _make_mla_spec(PS_C128)
-
-    mla_group = KVCacheGroupSpec(
-        layer_names=list(mla_specs.keys()),
-        kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=mla_specs),
-    )
-
-    swa_specs = {}
-    for i in range(n_swa):
-        swa_specs[f"swa.{i}"] = _make_mla_spec(PS_SWA)
-
-    swa_group = KVCacheGroupSpec(
-        layer_names=list(swa_specs.keys()),
-        kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=swa_specs),
-    )
-
-    return [mla_group, swa_group]
+def _mixed_page_groups(n_mla=3, n_idx=3, n_swa=5):
+    """A DeepSeek-V4-style hybrid: two groups with different page mixes."""
+    g1 = {f"mla.{i}": _mla(512) for i in range(n_mla)}
+    g1.update({f"idx.{i}": _mla(128) for i in range(n_idx)})
+    g2 = {f"swa.{i}": _mla(512) for i in range(n_swa)}
+    return [_uniform_group(g1), _uniform_group(g2)], g1, g2
 
 
-def _mock_vllm_config(kv_connector_extra_config: dict[str, str] | None = None):
+def _mock_vllm_config(layout: str | None):
     config = MagicMock()
+    config.cache_config = CacheConfig()
     config.cache_config.num_gpu_blocks_override = None
-    config.cache_config.prefix_cache_retention_interval = 0
-    config.kv_transfer_config = None
-    if kv_connector_extra_config is not None:
-        config.kv_transfer_config = MagicMock()
-        config.kv_transfer_config.kv_connector_extra_config = kv_connector_extra_config
+    config.cache_config.kv_cache_layout = layout
     return config
 
 
-def _run(n_c4=3, n_c128=2, n_swa=5, mem=100 * 1024 * 1024):
-    groups = _make_groups(n_c4, n_c128, n_swa)
-    return _get_kv_cache_config_packed(_mock_vllm_config(), groups, mem)
-
-
-def _page_sizes_by_layer(
-    groups: list[KVCacheGroupSpec],
-) -> dict[str, int]:
-    page_sizes = {}
-    for group in groups:
-        specs = group.kv_cache_spec.kv_cache_specs
-        for layer_name in group.layer_names:
-            page_sizes[layer_name] = specs[layer_name].page_size_bytes
-    return page_sizes
-
-
-def _packing_by_layer(
-    tensors: list[KVCacheTensor],
-) -> dict[str, tuple[int, int]]:
+def _pages(groups) -> dict[str, int]:
     return {
-        layer_name: (tensor.offset, tensor.block_stride)
-        for tensor in tensors
-        for layer_name in tensor.shared_by
+        name: group.kv_cache_spec.kv_cache_specs[name].page_size_bytes
+        if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
+        else group.kv_cache_spec.page_size_bytes
+        for group in groups
+        for name in group.layer_names
     }
 
 
-def _make_views(
-    groups: list[KVCacheGroupSpec],
-    num_blocks: int,
-    tensors: list[KVCacheTensor],
-) -> dict[str, torch.Tensor]:
-    page_sizes = _page_sizes_by_layer(groups)
-    packing = _packing_by_layer(tensors)
-    backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
-    return {
-        layer_name: torch.as_strided(
-            backing,
-            size=(num_blocks, page_size),
-            stride=(packing[layer_name][1], 1),
-            storage_offset=packing[layer_name][0],
+def _expected_bytes_per_block(groups) -> int:
+    pages = _pages(groups)
+    return max(sum(pages[n] for n in g.layer_names) for g in groups)
+
+
+def _bind(config, layout: str):
+    return allocate_kv_cache(config, torch.device("cpu"), KVCacheLayout[layout], None)
+
+
+class TestDensePacking:
+    def test_bytes_per_block_is_largest_group(self):
+        groups, g1, g2 = _mixed_page_groups()
+        assert _get_kv_cache_bytes_per_block(groups) == _expected_bytes_per_block(
+            groups
         )
-        for layer_name, page_size in page_sizes.items()
-    }
 
-
-def _make_page_group(prefix: str, page_sizes: list[int]) -> KVCacheGroupSpec:
-    specs = {
-        f"{prefix}.{i}": MagicMock(page_size_bytes=page_size)
-        for i, page_size in enumerate(page_sizes)
-    }
-    return KVCacheGroupSpec(
-        layer_names=list(specs),
-        kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=specs),
-    )
-
-
-class TestInterleavedPacking:
-    def test_compact_cache_overlays_fp32_state_group(self):
-        full_specs = {}
-        state_specs = {}
-        for i in range(2):
-            full_specs[f"mla.{i}"] = MLAAttentionSpec(
-                block_size=256,
-                num_kv_heads=1,
-                head_size=512,
-                dtype=torch.uint8,
-                page_size_padded=32768,
-                indexes_kv_by_block_stride=True,
-                compress_ratio=4,
-            )
-            full_specs[f"indexer.{i}"] = MLAAttentionSpec(
-                block_size=256,
-                num_kv_heads=1,
-                head_size=68,
-                dtype=torch.uint8,
-                page_size_padded=4608,
-                compress_ratio=4,
-            )
-            state_specs[f"mla_state.{i}"] = SlidingWindowMLASpec(
-                block_size=4,
-                num_kv_heads=1,
-                head_size=2048,
-                dtype=torch.float32,
-                sliding_window=8,
-                indexes_kv_by_block_stride=True,
-            )
-            state_specs[f"indexer_state.{i}"] = SlidingWindowMLASpec(
-                block_size=4,
-                num_kv_heads=1,
-                head_size=512,
-                dtype=torch.float32,
-                sliding_window=8,
-                indexes_kv_by_block_stride=True,
-            )
-
-        grouped_specs = [
-            UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=full_specs),
-            UniformTypeKVCacheSpecs(block_size=4, kv_cache_specs=state_specs),
+        config = get_kv_cache_config_from_groups(
+            _mock_vllm_config("BLHNC"), groups, MEMORY
+        )
+        # Groups overlay: both start at offset 0.
+        assert [tensor.offset for tensor in config.kv_cache_tensors].count(0) == 2
+        assert [tensor.layers for tensor in config.kv_cache_tensors] == [
+            list(g1)[:3],
+            list(g1)[3:],
+            list(g2),
         ]
-        groups = _get_kv_cache_groups_uniform_groups(grouped_specs)
 
-        assert len(groups) == 2
-        assert {full_specs[f"indexer.{i}"].page_size_bytes for i in range(2)} == {4608}
-        assert {full_specs[f"indexer.{i}"].real_page_size_bytes for i in range(2)} == {
-            4352
+    def test_layers_within_a_group_are_dense(self):
+        groups, _, _ = _mixed_page_groups()
+        pages = _pages(groups)
+        config = get_kv_cache_config_from_groups(
+            _mock_vllm_config("BLHNC"), groups, MEMORY
+        )
+        offsets = {
+            name: tensor.offset + i * tensor.layer_stride
+            for tensor in config.kv_cache_tensors
+            for i, name in enumerate(tensor.layers)
         }
-        assert {
-            state_specs[f"indexer_state.{i}"].page_size_bytes for i in range(2)
-        } == {8192}
-
-        full_group_bytes = 2 * (32768 + 4608)
-        state_group_bytes = 2 * (32768 + 8192)
-        bytes_per_block = max(full_group_bytes, state_group_bytes)
-        num_blocks, tensors = _get_kv_cache_config_packed(
-            _mock_vllm_config(), groups, bytes_per_block * 32
-        )
-        assert num_blocks == 32
-        assert {tensor.block_stride for tensor in tensors} == {bytes_per_block}
-
-        packing = _packing_by_layer(tensors)
-        assert packing["mla.0"][0] == packing["mla_state.0"][0] == 0
-        assert packing["indexer.0"][0] == 32768
-        assert packing["indexer_state.0"][0] == 32768
-
-    def test_deepseek_v4_pro_stride(self):
-        groups = [
-            _make_page_group("full", [32768, 4608] * 30 + [1024] * 31),
-            _make_page_group("c4_state", [32768, 8192] * 30),
-            _make_page_group("c128_state", [32768] * 31),
-            _make_page_group("swa.0", [32768] * 31),
-            _make_page_group("swa.1", [32768] * 30),
-        ]
-        expected_stride = 1_228_800
-
-        num_blocks, tensors = _get_kv_cache_config_packed(
-            _mock_vllm_config(), groups, expected_stride * 32
-        )
-
-        assert num_blocks == 32
-        assert {tensor.block_stride for tensor in tensors} == {expected_stride}
-        assert {tensor.size for tensor in tensors} == {expected_stride * 32}
-
-    def test_all_tensors_have_block_stride(self):
-        _, tensors = _run()
-        for t in tensors:
-            assert t.block_stride > 0
-
-    def test_all_tensors_share_same_size(self):
-        _, tensors = _run()
-        sizes = set(t.size for t in tensors)
-        assert len(sizes) == 1
-        assert sizes.pop() > 0
-
-    def test_offsets_within_one_block(self):
-        groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
-        _, tensors = _get_kv_cache_config_packed(
-            _mock_vllm_config(), groups, 100 * 1024 * 1024
-        )
-        page_sizes = _page_sizes_by_layer(groups)
-        packing = _packing_by_layer(tensors)
-        for layer_name, page_size in page_sizes.items():
-            offset, block_stride = packing[layer_name]
-            assert offset + page_size <= block_stride
-
-    def test_layouts_are_disjoint_within_each_group(self):
-        groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
-        _, tensors = _get_kv_cache_config_packed(
-            _mock_vllm_config(), groups, 100 * 1024 * 1024
-        )
-        page_sizes = _page_sizes_by_layer(groups)
-        packing = _packing_by_layer(tensors)
-
         for group in groups:
-            ranges = sorted(
-                (packing[name][0], packing[name][0] + page_sizes[name])
-                for name in group.layer_names
+            expected = 0
+            for name in group.layer_names:
+                assert offsets[name] == expected
+                expected += pages[name]
+
+    @pytest.mark.parametrize("layout", ["LBNHC", "BLHNC"])
+    def test_allocation_is_layout_invariant(self, layout):
+        specs = {f"l.{i}": _full() for i in range(4)}
+        groups = [KVCacheGroupSpec(list(specs), _full())]
+        config = get_kv_cache_config_from_groups(
+            _mock_vllm_config(layout), groups, MEMORY
+        )
+        (tensor,) = config.kv_cache_tensors
+        page = _full().page_size_bytes
+        assert tensor.layers == list(specs)
+        assert config.num_blocks == MEMORY // (4 * page)
+        assert tensor.size == 4 * page * config.num_blocks
+        if layout == "LBNHC":
+            assert (tensor.layer_stride, tensor.block_stride) == (
+                page * config.num_blocks,
+                page,
             )
-            assert all(left[1] <= right[0] for left, right in zip(ranges, ranges[1:]))
-
-    def test_all_layers_accounted_for(self):
-        n_c4, n_c128, n_swa = 5, 4, 7
-        _, tensors = _run(n_c4=n_c4, n_c128=n_c128, n_swa=n_swa)
-        all_names = set()
-        for t in tensors:
-            all_names.update(t.shared_by)
-        expected = n_c4 * 2 + n_c128 + n_swa
-        assert len(all_names) == expected
-
-    def test_group_owned_blocks_do_not_alias(self):
-        groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
-        num_blocks, tensors = _get_kv_cache_config_packed(
-            _mock_vllm_config(), groups, 100 * 1024 * 1024
-        )
-        views = _make_views(groups, num_blocks, tensors)
-
-        expected = {}
-        value = 1
-        for block_id, group in enumerate(groups):
-            for layer_name in group.layer_names:
-                views[layer_name][block_id].fill_(value)
-                expected[layer_name] = (block_id, value)
-                value += 1
-
-        for layer_name, (block_id, value) in expected.items():
-            assert (views[layer_name][block_id] == value).all()
-
-        # Once the first group releases its block, another group may reuse it.
-        for layer_name in groups[1].layer_names:
-            views[layer_name][0].fill_(255)
-        for layer_name in groups[1].layer_names:
-            assert (views[layer_name][0] == 255).all()
-
-    def test_hma_attention_groups_keep_default_backing(self):
-        full = _make_full_spec()
-        sw = _make_sw_spec()
-        page_size = full.page_size_bytes
-        groups = [
-            KVCacheGroupSpec(["full.0", "full.1"], full),
-            KVCacheGroupSpec(["sw.0", "sw.2"], sw),
-            KVCacheGroupSpec(["sw.1", "sw.3"], sw),
-        ]
+        else:
+            assert (tensor.layer_stride, tensor.block_stride) == (page, 4 * page)
 
+    @pytest.mark.parametrize("layout", ["LBNHC", "BLHNC"])
+    def test_single_group_mixed_pages_follows_layout(self, layout):
+        specs = {"mla.0": _mla(512), "mla.1": _mla(512), "idx.0": _mla(128)}
+        groups = [_uniform_group(specs)]
         config = get_kv_cache_config_from_groups(
-            _mock_vllm_config(), groups, available_memory=page_size * 2 * 32
+            _mock_vllm_config(layout), groups, MEMORY
         )
+        block_stride = _expected_bytes_per_block(groups)
+        mla_tensor, idx_tensor = config.kv_cache_tensors
+        assert mla_tensor.layers == ["mla.0", "mla.1"]
+        assert idx_tensor.layers == ["idx.0"]
+        assert {t.size for t in config.kv_cache_tensors} == {
+            block_stride * config.num_blocks
+        }
+        if layout == "LBNHC":
+            assert (
+                idx_tensor.offset == 2 * _mla(512).page_size_bytes * config.num_blocks
+            )
+            assert mla_tensor.block_stride == _mla(512).page_size_bytes
+        else:
+            assert idx_tensor.offset == 2 * _mla(512).page_size_bytes
+            assert mla_tensor.block_stride == block_stride
 
-        assert config.num_blocks == 32
-        assert config.prefix_cache_retention_interval == 0
-        assert sum(t.size for t in config.kv_cache_tensors) == page_size * 2 * 32
-        assert config.kv_cache_tensors == [
-            KVCacheTensor(size=page_size * 32, shared_by=["full.0", "sw.0", "sw.1"]),
-            KVCacheTensor(size=page_size * 32, shared_by=["full.1", "sw.2", "sw.3"]),
-        ]
-
-    def test_hma_attention_groups_use_packed_backing_with_enable_cross_layers(self):
-        full = _make_full_spec()
-        sw = _make_sw_spec()
-        page_size = full.page_size_bytes
-        groups = [
-            KVCacheGroupSpec(["full.0", "full.1"], full),
-            KVCacheGroupSpec(["sw.0", "sw.2"], sw),
-            KVCacheGroupSpec(["sw.1", "sw.3"], sw),
-        ]
-
+    def test_overlaid_groups_alias_and_stay_isolated(self):
+        groups, g1, g2 = _mixed_page_groups()
+        # Overlay models resolve to a block-outer layout at backend selection (the
+        # model's backend declares it); mirror that here.
         config = get_kv_cache_config_from_groups(
-            _mock_vllm_config({"enable_cross_layers_blocks": "True"}),
-            groups,
-            available_memory=page_size * 2 * 32,
+            _mock_vllm_config("BLNHC"), groups, MEMORY
         )
+        assert config.num_blocks == MEMORY // _expected_bytes_per_block(groups)
+        assert _pool_bytes_per_block(groups) == _expected_bytes_per_block(groups)
 
-        assert config.num_blocks == 32
-        assert {t.size for t in config.kv_cache_tensors} == {page_size * 2 * 32}
-        assert config.kv_cache_tensors == [
-            KVCacheTensor(
-                size=page_size * 2 * 32,
-                shared_by=["full.0", "sw.0", "sw.1"],
-                offset=0,
-                block_stride=page_size * 2,
-            ),
-            KVCacheTensor(
-                size=page_size * 2 * 32,
-                shared_by=["full.1", "sw.2", "sw.3"],
-                offset=page_size,
-                block_stride=page_size * 2,
-            ),
-        ]
+        views = _bind(config, "BLNHC")
+        assert set(views) == set(g1) | set(g2)
+        assert views["swa.0"].data_ptr() == views["mla.0"].data_ptr()
 
-    def test_single_group_attention_keeps_unpacked_layout(self):
-        spec = _make_full_spec()
-        groups = [KVCacheGroupSpec(["full.0", "full.1"], spec)]
+        # A block is owned by one group at a time: writes to group-owned
+        # blocks never disturb the other group's blocks.
+        for i, name in enumerate(g1):
+            views[name][0].fill_(i + 1)
+            views[name][2].fill_(i + 1)
+        for i, name in enumerate(g2):
+            views[name][1].fill_(100 + i)
+            views[name][3].fill_(100 + i)
+        for i, name in enumerate(g1):
+            assert (views[name][0].to(torch.int32) == i + 1).all()
+            assert (views[name][2].to(torch.int32) == i + 1).all()
+        for i, name in enumerate(g2):
+            assert (views[name][1].to(torch.int32) == 100 + i).all()
+            assert (views[name][3].to(torch.int32) == 100 + i).all()
+        # Layers within a group are disjoint.
+        views["mla.0"][0].fill_(77)
+        for i, name in enumerate(list(g1)[1:], start=1):
+            assert (views[name][0].to(torch.int32) == i + 1).all()
 
+    def test_layer_compact_layout_rejected_for_overlaid_groups(self):
+        # The layout has a single writer (backend-selection resolution); a layer-
+        # compact layout reaching an overlay model's allocation is an error, not a
+        # silent flip.
+        groups, _, _ = _mixed_page_groups()
+        with pytest.raises(
+            ValueError, match="cannot express this model's mixed page sizes"
+        ):
+            get_kv_cache_config_from_groups(_mock_vllm_config("LBNHC"), groups, MEMORY)
+
+    def test_unresolved_layout_rejected(self):
+        groups, _, _ = _mixed_page_groups()
+        with pytest.raises(ValueError, match="has not been resolved"):
+            get_kv_cache_config_from_groups(_mock_vllm_config(None), groups, MEMORY)
+
+    def test_head_outer_layout_rejected_for_mixed_pages(self):
+        groups, _, _ = _mixed_page_groups()
+        with pytest.raises(
+            ValueError, match="cannot express this model's mixed page sizes"
+        ):
+            get_kv_cache_config_from_groups(_mock_vllm_config("LHBNC"), groups, MEMORY)
+
+    @pytest.mark.parametrize("layout", ["LBNHC", "BLHNC"])
+    def test_bound_views_round_trip(self, layout):
+        specs = {"mla.0": _mla(512), "mla.1": _mla(512), "idx.0": _mla(128)}
+        groups = [_uniform_group(specs)]
         config = get_kv_cache_config_from_groups(
-            _mock_vllm_config(), groups, available_memory=spec.page_size_bytes * 2 * 32
+            _mock_vllm_config(layout), groups, MEMORY
         )
-
-        assert sum(t.size for t in config.kv_cache_tensors) == (
-            spec.page_size_bytes * 2 * 32
-        )
-        assert [t.block_stride for t in config.kv_cache_tensors] == [0, 0]
+        views = _bind(config, layout)
+        for i, name in enumerate(specs):
+            views[name].fill_(i + 1)
+        for i, name in enumerate(specs):
+            assert (views[name].to(torch.int32) == i + 1).all(), name
+            assert views[name].shape[0] == config.num_blocks
 
 
 if __name__ == "__main__":

tests/v1/core/test_kv_cache_utils.py

@@ -124,7 +124,6 @@ def new_kv_cache_spec(
     page_size_padded=None,
     sliding_window=None,
     attention_chunk_size=None,
-    indexes_kv_by_block_stride=False,
     kv_quant_mode=KVQuantMode.NONE,
 ):
     return FullAttentionSpec(
@@ -135,7 +134,6 @@ def new_kv_cache_spec(
         page_size_padded=page_size_padded,
         sliding_window=sliding_window,
         attention_chunk_size=attention_chunk_size,
-        indexes_kv_by_block_stride=indexes_kv_by_block_stride,
         kv_quant_mode=kv_quant_mode,
     )
 
@@ -147,7 +145,6 @@ def new_sliding_window_spec(
     dtype=torch.float32,
     page_size_padded=None,
     sliding_window=1,
-    indexes_kv_by_block_stride=False,
 ):
     return SlidingWindowSpec(
         block_size=block_size,
@@ -156,7 +153,6 @@ def new_sliding_window_spec(
         dtype=dtype,
         page_size_padded=page_size_padded,
         sliding_window=sliding_window,
-        indexes_kv_by_block_stride=indexes_kv_by_block_stride,
     )
 
 
@@ -875,6 +871,7 @@ def test_metrics_empty_stats():
 def test_get_kv_cache_configs_multiple_workers():
     model_config = ModelConfig(max_model_len=16)
     vllm_config = VllmConfig(model_config=model_config)
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
     vllm_config.cache_config.prefix_cache_retention_interval = None
 
     ref_kv_cache_spec = new_kv_cache_spec()
@@ -898,36 +895,21 @@ def test_get_kv_cache_configs_multiple_workers():
             ref_kv_cache_spec.page_size_bytes * 2 * 10,
         ],
     )
-    assert kv_cache_configs == [
-        KVCacheConfig(
-            num_blocks=10,
-            kv_cache_tensors=[
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
-                ),
-            ],
-            kv_cache_groups=[
-                KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
-            ],
-        ),
-        KVCacheConfig(
-            num_blocks=10,
-            kv_cache_tensors=[
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
-                ),
-            ],
-            kv_cache_groups=[
-                KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
-            ],
-        ),
-    ]
+    expected = KVCacheConfig(
+        num_blocks=10,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
+                layers=["layer1", "layer2"],
+                layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
+                block_stride=ref_kv_cache_spec.page_size_bytes,
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
+        ],
+    )
+    assert kv_cache_configs == [expected, expected]
 
     # Different available memory. This is the case for TP.
     # Use the smallest memory available.
@@ -939,36 +921,7 @@ def test_get_kv_cache_configs_multiple_workers():
             ref_kv_cache_spec.page_size_bytes * 2 * 20,
         ],
     )
-    assert kv_cache_configs == [
-        KVCacheConfig(
-            num_blocks=10,
-            kv_cache_tensors=[
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
-                ),
-            ],
-            kv_cache_groups=[
-                KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
-            ],
-        ),
-        KVCacheConfig(
-            num_blocks=10,
-            kv_cache_tensors=[
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
-                ),
-            ],
-            kv_cache_groups=[
-                KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
-            ],
-        ),
-    ]
+    assert kv_cache_configs == [expected, expected]
 
     # Different KV cache specs. This is the case for PP.
     different_layer_specs = [
@@ -995,7 +948,10 @@ def test_get_kv_cache_configs_multiple_workers():
             num_blocks=10,
             kv_cache_tensors=[
                 KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
+                    size=ref_kv_cache_spec.page_size_bytes * 10,
+                    layers=["layer1"],
+                    layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
+                    block_stride=ref_kv_cache_spec.page_size_bytes,
                 ),
             ],
             kv_cache_groups=[
@@ -1006,10 +962,10 @@ def test_get_kv_cache_configs_multiple_workers():
             num_blocks=10,
             kv_cache_tensors=[
                 KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer3"]
+                    size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
+                    layers=["layer2", "layer3"],
+                    layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
+                    block_stride=ref_kv_cache_spec.page_size_bytes,
                 ),
             ],
             kv_cache_groups=[
@@ -1046,57 +1002,39 @@ def test_get_kv_cache_configs_multiple_workers():
             ref_kv_cache_spec.page_size_bytes * 2 * 10,
         ],
     )
+    expected_12 = KVCacheConfig(
+        num_blocks=10,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
+                layers=["layer1", "layer2"],
+                layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
+                block_stride=ref_kv_cache_spec.page_size_bytes,
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
+        ],
+    )
+    expected_3 = KVCacheConfig(
+        num_blocks=10,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=ref_kv_cache_spec.page_size_bytes * 10,
+                layers=["layer3"],
+                layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
+                block_stride=ref_kv_cache_spec.page_size_bytes,
+            ),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["layer3"], ref_kv_cache_spec),
+        ],
+    )
     assert kv_cache_configs == [
-        KVCacheConfig(
-            num_blocks=10,
-            kv_cache_tensors=[
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
-                ),
-            ],
-            kv_cache_groups=[
-                KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
-            ],
-        ),
-        KVCacheConfig(
-            num_blocks=10,
-            kv_cache_tensors=[
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
-                ),
-            ],
-            kv_cache_groups=[
-                KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
-            ],
-        ),
-        KVCacheConfig(
-            num_blocks=10,
-            kv_cache_tensors=[
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer3"]
-                ),
-            ],
-            kv_cache_groups=[
-                KVCacheGroupSpec(["layer3"], ref_kv_cache_spec),
-            ],
-        ),
-        KVCacheConfig(
-            num_blocks=10,
-            kv_cache_tensors=[
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer3"]
-                ),
-            ],
-            kv_cache_groups=[
-                KVCacheGroupSpec(["layer3"], ref_kv_cache_spec),
-            ],
-        ),
+        expected_12,
+        expected_12,
+        expected_3,
+        expected_3,
     ]
 
     # Different workers have different types of layers. This is the case for
@@ -1124,10 +1062,10 @@ def test_get_kv_cache_configs_multiple_workers():
             num_blocks=10,
             kv_cache_tensors=[
                 KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer1"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer2"]
+                    size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
+                    layers=["layer1", "layer2"],
+                    layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
+                    block_stride=ref_kv_cache_spec.page_size_bytes,
                 ),
             ],
             kv_cache_groups=[
@@ -1139,10 +1077,10 @@ def test_get_kv_cache_configs_multiple_workers():
             num_blocks=10,
             kv_cache_tensors=[
                 KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer3"]
-                ),
-                KVCacheTensor(
-                    size=ref_kv_cache_spec.page_size_bytes * 10, shared_by=["layer4"]
+                    size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
+                    layers=["layer3", "layer4"],
+                    layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
+                    block_stride=ref_kv_cache_spec.page_size_bytes,
                 ),
             ],
             kv_cache_groups=[
@@ -1180,7 +1118,21 @@ def test_get_kv_cache_configs_multiple_workers():
             kv_cache_tensors=[
                 KVCacheTensor(
                     size=ref_kv_cache_spec.page_size_bytes * 10,
-                    shared_by=["layer1", "layer2", "layer3"],
+                    layers=["layer1"],
+                    layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
+                    block_stride=ref_kv_cache_spec.page_size_bytes,
+                ),
+                KVCacheTensor(
+                    size=ref_kv_cache_spec.page_size_bytes * 10,
+                    layers=["layer2"],
+                    layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
+                    block_stride=ref_kv_cache_spec.page_size_bytes,
+                ),
+                KVCacheTensor(
+                    size=ref_kv_cache_spec.page_size_bytes * 10,
+                    layers=["layer3"],
+                    layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
+                    block_stride=ref_kv_cache_spec.page_size_bytes,
                 ),
             ],
             kv_cache_groups=[
@@ -1194,7 +1146,21 @@ def test_get_kv_cache_configs_multiple_workers():
             kv_cache_tensors=[
                 KVCacheTensor(
                     size=ref_kv_cache_spec.page_size_bytes * 10,
-                    shared_by=["layer4", "layer5", "layer6"],
+                    layers=["layer4"],
+                    layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
+                    block_stride=ref_kv_cache_spec.page_size_bytes,
+                ),
+                KVCacheTensor(
+                    size=ref_kv_cache_spec.page_size_bytes * 10,
+                    layers=["layer5"],
+                    layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
+                    block_stride=ref_kv_cache_spec.page_size_bytes,
+                ),
+                KVCacheTensor(
+                    size=ref_kv_cache_spec.page_size_bytes * 10,
+                    layers=["layer6"],
+                    layer_stride=ref_kv_cache_spec.page_size_bytes * 10,
+                    block_stride=ref_kv_cache_spec.page_size_bytes,
                 ),
             ],
             kv_cache_groups=[
@@ -1233,6 +1199,7 @@ def test_get_kv_cache_configs_multiple_workers():
 def test_get_kv_cache_configs_pp_sharding(asymmetric_memory):
     model_config = ModelConfig(max_model_len=512)
     vllm_config = VllmConfig(model_config=model_config)
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
     vllm_config.cache_config.prefix_cache_retention_interval = None
 
     ref_kv_cache_spec = new_kv_cache_spec()
@@ -1262,7 +1229,10 @@ def test_get_kv_cache_configs_pp_sharding(asymmetric_memory):
             kv_cache_tensors=[
                 KVCacheTensor(
                     size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks,
-                    shared_by=["layer1"],
+                    layers=["layer1"],
+                    layer_stride=ref_kv_cache_spec.page_size_bytes
+                    * expected_num_blocks,
+                    block_stride=ref_kv_cache_spec.page_size_bytes,
                 ),
             ],
             kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)],
@@ -1272,7 +1242,10 @@ def test_get_kv_cache_configs_pp_sharding(asymmetric_memory):
             kv_cache_tensors=[
                 KVCacheTensor(
                     size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks,
-                    shared_by=["layer2"],
+                    layers=["layer2"],
+                    layer_stride=ref_kv_cache_spec.page_size_bytes
+                    * expected_num_blocks,
+                    block_stride=ref_kv_cache_spec.page_size_bytes,
                 ),
             ],
             kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)],
@@ -1628,78 +1601,18 @@ def test_get_max_concurrency_for_kv_cache_config():
     )
 
 
-def test_get_max_concurrency_packed_kv_cache_config():
-    from vllm.v1.core.kv_cache_utils import (
-        _get_kv_cache_config_packed,
-        _use_packed_kv_cache_config,
-    )
-
-    model_config = ModelConfig(
-        "Qwen/Qwen1.5-7B",
-        runner="generate",
-        dtype="float16",
-        max_model_len=16384,
-    )
-    scheduler_config = SchedulerConfig(
-        max_num_batched_tokens=1024,
-        enable_chunked_prefill=True,
-        max_model_len=model_config.max_model_len,
-        is_encoder_decoder=model_config.is_encoder_decoder,
-        async_scheduling=False,
-    )
-    vllm_config = VllmConfig(
-        model_config=model_config,
-        scheduler_config=scheduler_config,
-    )
-
-    # All-UniformTypeKVCacheSpecs groups select the packed layout.
-    mla_specs = {f"layer_{i}": new_mla_spec() for i in range(4)}
-    swa_specs = {
-        f"layer_{i}": SlidingWindowMLASpec(
-            block_size=16,
-            num_kv_heads=1,
-            head_size=576,
-            dtype=torch.float32,
-            sliding_window=128,
-        )
-        for i in range(4, 6)
-    }
-    kv_cache_groups = [
-        KVCacheGroupSpec(
-            list(mla_specs),
-            UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=mla_specs),
-        ),
-        KVCacheGroupSpec(
-            list(swa_specs),
-            UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=swa_specs),
-        ),
-    ]
-    assert _use_packed_kv_cache_config(vllm_config, kv_cache_groups)
-    num_blocks, kv_cache_tensors = _get_kv_cache_config_packed(
-        vllm_config, kv_cache_groups, 2 * GiB_bytes
-    )
-    assert num_blocks > 0
-    kv_cache_config_packed = KVCacheConfig(
-        num_blocks=num_blocks,
-        kv_cache_tensors=kv_cache_tensors,
-        kv_cache_groups=kv_cache_groups,
-    )
-    # Per-request blocks: the MLA group needs cdiv(16384, 16) = 1024 pages;
-    # the SWA group cdiv(min(128 - 1 + 1024, 16384), 16) + 1 = 73. The
-    # previous formula normalized by the first group's page size and gave
-    # 1061 blocks per request instead of 1097.
-    assert get_max_concurrency_for_kv_cache_config(
-        vllm_config, kv_cache_config_packed
-    ) == num_blocks / (1024 + 73)
-
-
 def test_allocate_with_lookahead():
     """Verify that lookahead tokens correctly affect block allocation"""
     block_size = 4
     config = KVCacheConfig(
         num_blocks=10,
         kv_cache_tensors=[
-            KVCacheTensor(size=100, shared_by=["layer1"]),
+            KVCacheTensor(
+                size=100,
+                layers=["layer1"],
+                layer_stride=100,
+                block_stride=10,
+            ),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(["layer1"], new_kv_cache_spec(block_size=block_size)),
@@ -1763,6 +1676,7 @@ def test_get_kv_cache_config_one_worker():
     # pass max_model_len to pass check_enough_kv_cache_memory
     model_config = ModelConfig(max_model_len=16)
     vllm_config = VllmConfig(model_config=model_config)
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
     vllm_config.cache_config.prefix_cache_retention_interval = None
 
     mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
@@ -1778,8 +1692,12 @@ def test_get_kv_cache_config_one_worker():
     assert kv_cache_config_full == KVCacheConfig(
         num_blocks=32,
         kv_cache_tensors=[
-            KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_1"]),
-            KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_2"]),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 32 * 2,
+                layers=["layer_1", "layer_2"],
+                layer_stride=mem_per_block_per_layer * 32,
+                block_stride=mem_per_block_per_layer,
+            ),
         ],
         kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())],
     )
@@ -1795,8 +1713,12 @@ def test_get_kv_cache_config_one_worker():
     assert kv_cache_config_sliding == KVCacheConfig(
         num_blocks=32,
         kv_cache_tensors=[
-            KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_1"]),
-            KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_2"]),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 32 * 2,
+                layers=["layer_1", "layer_2"],
+                layer_stride=mem_per_block_per_layer * 32,
+                block_stride=mem_per_block_per_layer,
+            ),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(["layer_1", "layer_2"], new_sliding_window_spec())
@@ -1815,8 +1737,12 @@ def test_get_kv_cache_config_one_worker():
     assert kv_cache_config_hybrid == KVCacheConfig(
         num_blocks=32,
         kv_cache_tensors=[
-            KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_1"]),
-            KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_2"]),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 32 * 2,
+                layers=["layer_1", "layer_2"],
+                layer_stride=mem_per_block_per_layer * 32,
+                block_stride=mem_per_block_per_layer,
+            ),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(
@@ -1838,7 +1764,16 @@ def test_get_kv_cache_config_one_worker():
         num_blocks=64,
         kv_cache_tensors=[
             KVCacheTensor(
-                size=mem_per_block_per_layer * 64, shared_by=["layer_1", "layer_2"]
+                size=mem_per_block_per_layer * 64,
+                layers=["layer_1"],
+                layer_stride=mem_per_block_per_layer * 64,
+                block_stride=mem_per_block_per_layer,
+            ),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 64,
+                layers=["layer_2"],
+                layer_stride=mem_per_block_per_layer * 64,
+                block_stride=mem_per_block_per_layer,
             ),
         ],
         kv_cache_groups=[
@@ -1863,12 +1798,22 @@ def test_get_kv_cache_config_one_worker():
         num_blocks=32,
         kv_cache_tensors=[
             KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_1", "layer_3", "layer_4"],
+                size=mem_per_block_per_layer * 32 * 2,
+                layers=["layer_1", "layer_2"],
+                layer_stride=mem_per_block_per_layer * 32,
+                block_stride=mem_per_block_per_layer,
             ),
             KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_2", "layer_5", "layer_6"],
+                size=mem_per_block_per_layer * 32 * 2,
+                layers=["layer_3", "layer_5"],
+                layer_stride=mem_per_block_per_layer * 32,
+                block_stride=mem_per_block_per_layer,
+            ),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 32 * 2,
+                layers=["layer_4", "layer_6"],
+                layer_stride=mem_per_block_per_layer * 32,
+                block_stride=mem_per_block_per_layer,
             ),
         ],
         kv_cache_groups=[
@@ -1898,15 +1843,28 @@ def test_get_kv_cache_config_one_worker():
         num_blocks=32,
         kv_cache_tensors=[
             KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_1", "layer_4", "layer_5", "layer_6"],
+                size=mem_per_block_per_layer * 32 * 3,
+                layers=["layer_1", "layer_2", "layer_3"],
+                layer_stride=mem_per_block_per_layer * 32,
+                block_stride=mem_per_block_per_layer,
             ),
             KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_2", "layer_7", "layer_8", "layer_9"],
+                size=mem_per_block_per_layer * 32 * 3,
+                layers=["layer_4", "layer_7", "layer_10"],
+                layer_stride=mem_per_block_per_layer * 32,
+                block_stride=mem_per_block_per_layer,
             ),
             KVCacheTensor(
-                size=mem_per_block_per_layer * 32, shared_by=["layer_3", "layer_10"]
+                size=mem_per_block_per_layer * 32 * 3,
+                layers=["layer_5", "layer_8"],
+                layer_stride=mem_per_block_per_layer * 32,
+                block_stride=mem_per_block_per_layer,
+            ),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 32 * 3,
+                layers=["layer_6", "layer_9"],
+                layer_stride=mem_per_block_per_layer * 32,
+                block_stride=mem_per_block_per_layer,
             ),
         ],
         kv_cache_groups=[
@@ -1943,28 +1901,23 @@ def test_get_kv_cache_config_one_worker():
         num_blocks=32,
         kv_cache_tensors=[
             KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_1", "layer_7"],
+                size=mem_per_block_per_layer * 32 * 6,
+                layers=[
+                    "layer_1",
+                    "layer_2",
+                    "layer_3",
+                    "layer_4",
+                    "layer_5",
+                    "layer_6",
+                ],
+                layer_stride=mem_per_block_per_layer * 32,
+                block_stride=mem_per_block_per_layer,
             ),
             KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_2", "layer_8"],
-            ),
-            KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_3", "layer_9"],
-            ),
-            KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_4", "layer_10"],
-            ),
-            KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_5", "layer_11"],
-            ),
-            KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_6"],
+                size=mem_per_block_per_layer * 32 * 6,
+                layers=["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"],
+                layer_stride=mem_per_block_per_layer * 32,
+                block_stride=mem_per_block_per_layer,
             ),
         ],
         kv_cache_groups=[
@@ -1987,11 +1940,24 @@ def test_get_kv_cache_config_one_worker():
     kv_cache_config_hybrid = get_kv_cache_configs(
         vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 3 * 32]
     )[0]
+    # Layers of different page sizes pack densely into one allocation: the
+    # 2x-larger layer_1 takes the first region, layer_2 the next.
     assert kv_cache_config_hybrid == KVCacheConfig(
         num_blocks=32,
         kv_cache_tensors=[
-            KVCacheTensor(size=mem_per_block_per_layer * 32 * 2, shared_by=["layer_1"]),
-            KVCacheTensor(size=mem_per_block_per_layer * 32, shared_by=["layer_2"]),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 32 * 3,
+                layers=["layer_1"],
+                layer_stride=mem_per_block_per_layer * 2 * 32,
+                block_stride=mem_per_block_per_layer * 2,
+            ),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 32 * 3,
+                layers=["layer_2"],
+                layer_stride=mem_per_block_per_layer * 32,
+                block_stride=mem_per_block_per_layer,
+                offset=mem_per_block_per_layer * 2 * 32,
+            ),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(
@@ -2015,7 +1981,16 @@ def test_get_kv_cache_config_one_worker():
         num_blocks=32,
         kv_cache_tensors=[
             KVCacheTensor(
-                size=mem_per_block_per_layer * 32, shared_by=["layer_1", "layer_2"]
+                size=mem_per_block_per_layer * 32,
+                layers=["layer_1"],
+                layer_stride=mem_per_block_per_layer * 32,
+                block_stride=mem_per_block_per_layer,
+            ),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 32,
+                layers=["layer_2"],
+                layer_stride=mem_per_block_per_layer * 32,
+                block_stride=mem_per_block_per_layer,
             ),
         ],
         kv_cache_groups=[
@@ -2028,9 +2003,9 @@ def test_get_kv_cache_config_one_worker():
 
     # different hidden size that cannot be aligned by using different block size,
     # but can be aligned by padding the smaller physical page.
-    swa_spec = new_sliding_window_spec(head_size=96, indexes_kv_by_block_stride=True)
+    swa_spec = new_sliding_window_spec(head_size=96)
     kv_cache_specs_hybrid = {
-        "layer_1": new_kv_cache_spec(head_size=64, indexes_kv_by_block_stride=True),
+        "layer_1": new_kv_cache_spec(head_size=64),
         "layer_2": swa_spec,
     }
 
@@ -2041,7 +2016,18 @@ def test_get_kv_cache_config_one_worker():
     assert kv_cache_config_hybrid == KVCacheConfig(
         num_blocks=42,
         kv_cache_tensors=[
-            KVCacheTensor(size=padded_page_size * 42, shared_by=["layer_1", "layer_2"]),
+            KVCacheTensor(
+                size=padded_page_size * 42,
+                layers=["layer_1"],
+                layer_stride=padded_page_size * 42,
+                block_stride=padded_page_size,
+            ),
+            KVCacheTensor(
+                size=padded_page_size * 42,
+                layers=["layer_2"],
+                layer_stride=padded_page_size * 42,
+                block_stride=padded_page_size,
+            ),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(
@@ -2049,12 +2035,11 @@ def test_get_kv_cache_config_one_worker():
                 new_kv_cache_spec(
                     head_size=64,
                     page_size_padded=padded_page_size,
-                    indexes_kv_by_block_stride=True,
                 ),
             ),
             KVCacheGroupSpec(
                 ["layer_2"],
-                new_sliding_window_spec(head_size=96, indexes_kv_by_block_stride=True),
+                new_sliding_window_spec(head_size=96),
             ),
         ],
     )
@@ -2067,8 +2052,12 @@ def test_get_kv_cache_config_one_worker():
     assert kv_cache_config_override_blocks == KVCacheConfig(
         num_blocks=16,
         kv_cache_tensors=[
-            KVCacheTensor(size=mem_per_block_per_layer * 16, shared_by=["layer_1"]),
-            KVCacheTensor(size=mem_per_block_per_layer * 16, shared_by=["layer_2"]),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 16 * 2,
+                layers=["layer_1", "layer_2"],
+                layer_stride=mem_per_block_per_layer * 16,
+                block_stride=mem_per_block_per_layer,
+            ),
         ],
         kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())],
     )
@@ -2319,6 +2308,45 @@ def test_hidden_state_group_preserves_hybrid_prefix_cache_granularity():
     ) == (544, 136)
 
 
+def test_multi_run_layer_compact_strides_place_hoisted_heads():
+    """A layer-compact run region is its own dense allocation: under LHBNC the head
+    groups sit between the layers and the blocks, so a run's block stride is one head
+    group's slice, not the whole page (regression test for setStorage out-of-bounds
+    on ROCm hybrid models)."""
+    full = new_kv_cache_spec()
+    swa = new_sliding_window_spec(sliding_window=full.block_size * 2)
+    page = full.page_size_bytes
+    assert swa.page_size_bytes == page
+
+    vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16))
+    vllm_config.cache_config.kv_cache_layout = "LHBNC"
+    config = kv_cache_utils.get_kv_cache_config_from_groups(
+        vllm_config,
+        [
+            KVCacheGroupSpec(["full.0", "full.1"], full),
+            KVCacheGroupSpec(["swa.0", "swa.1"], swa),
+        ],
+        available_memory=4 * page * 8,
+    )
+
+    num_blocks = config.num_blocks
+    head_group = full.block_size * full.state_content_size_bytes
+    for tensor in config.kv_cache_tensors:
+        # One head group per block, head groups between the layers and blocks.
+        assert tensor.block_stride == head_group
+        assert tensor.layer_stride == full.num_heads * num_blocks * head_group
+        # The last block of the last layer's last head group must stay within the
+        # allocation (the old page-based stride overran it).
+        end = (
+            tensor.offset
+            + tensor.layer_stride * (len(tensor.layers) - 1)
+            + (full.num_heads - 1) * num_blocks * head_group
+            + tensor.block_stride * (num_blocks - 1)
+            + head_group
+        )
+        assert end <= tensor.size
+
+
 def test_mla_draft_prefers_standard_layout_when_pages_can_be_unified():
     specs = {
         "target.0.attn": new_mla_spec(),
@@ -2640,6 +2668,7 @@ def test_auto_fit_max_model_len():
     # Simulate the user passing -1 by setting original_max_model_len
     model_config.original_max_model_len = -1
     vllm_config = VllmConfig(model_config=model_config)
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
 
     mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2  # 16KB per block per layer
     kv_cache_specs = {
@@ -2658,6 +2687,7 @@ def test_auto_fit_max_model_len():
     model_config = ModelConfig(max_model_len=1024)
     model_config.original_max_model_len = -1
     vllm_config = VllmConfig(model_config=model_config)
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
 
     # With limited memory, max_model_len should be reduced
     # Need memory for at least max_model_len tokens
@@ -2678,6 +2708,7 @@ def test_auto_fit_max_model_len_with_hybrid():
     # Simulate the user passing -1 by setting original_max_model_len
     model_config.original_max_model_len = -1
     vllm_config = VllmConfig(model_config=model_config)
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
 
     mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2  # 16KB per block per layer
     gamma = 2
@@ -2700,6 +2731,7 @@ def test_auto_fit_max_model_len_not_triggered():
     model_config = ModelConfig(max_model_len=16)
     # original_max_model_len should be None by default, not -1
     vllm_config = VllmConfig(model_config=model_config)
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
 
     mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
     kv_cache_specs = {
@@ -2722,6 +2754,7 @@ def test_auto_fit_max_model_len_respects_num_gpu_blocks_override():
     model_config = ModelConfig(max_model_len=16384)
     model_config.original_max_model_len = -1  # request auto-fit
     vllm_config = VllmConfig(model_config=model_config)
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
     # Cap the cache to 32 blocks regardless of available memory.
     vllm_config.cache_config.num_gpu_blocks_override = 32
 
@@ -2747,6 +2780,7 @@ def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override():
     """
     model_config = ModelConfig(max_model_len=16384)
     vllm_config = VllmConfig(model_config=model_config)
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
     # 32 blocks is far too small for max_model_len=16384 (would need 1024).
     vllm_config.cache_config.num_gpu_blocks_override = 32
 
@@ -2769,14 +2803,13 @@ def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes():
     128-dim KV heads. The resulting page sizes are 3:2 rather than an integer
     block-size multiple, so the smaller page must be padded instead.
     """
-    # Both layers' backends opt into the padded-page strided view (e.g.
-    # FlashAttention / its DiffKV subclass), so padding is allowed.
+    # Attention layers read padded pages through a strided view, so padding is
+    # allowed.
     target_spec = new_kv_cache_spec(
         block_size=16,
         num_kv_heads=1,
         head_size=192,
         dtype=torch.bfloat16,
-        indexes_kv_by_block_stride=True,
     )
     draft_spec = new_sliding_window_spec(
         block_size=16,
@@ -2784,7 +2817,6 @@ def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes():
         head_size=128,
         dtype=torch.bfloat16,
         sliding_window=1024,
-        indexes_kv_by_block_stride=True,
     )
 
     unified_specs = kv_cache_utils.unify_kv_cache_spec_page_size(
@@ -2802,35 +2834,6 @@ def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes():
     assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes
 
 
-def test_unify_kv_cache_page_size_padding_requires_backend_support():
-    """Padding is gated on the backend declaring ``indexes_kv_by_block_stride``.
-
-    A backend that does not support the strided padded-page view must raise
-    rather than silently padding (and misreading KV at runtime).
-    """
-    target_spec = new_kv_cache_spec(
-        block_size=16,
-        num_kv_heads=1,
-        head_size=192,
-        dtype=torch.bfloat16,
-        indexes_kv_by_block_stride=True,
-    )
-    # The non-divisible draft layer needs padding but its backend does not
-    # support the strided padded-page view -> must raise, not silently pad.
-    draft_spec = new_sliding_window_spec(
-        block_size=16,
-        num_kv_heads=1,
-        head_size=128,
-        dtype=torch.bfloat16,
-        sliding_window=1024,
-        indexes_kv_by_block_stride=False,
-    )
-    specs = {"target_attn": target_spec, "draft_attn": draft_spec}
-
-    with pytest.raises(NotImplementedError):
-        kv_cache_utils.unify_kv_cache_spec_page_size(specs)
-
-
 def test_unpadded_page_size_includes_per_token_head_scales():
     # Per-token-head quant carries inline fp32 scales that are carved from the
     # raw KV allocation, so they must be budgeted into the offload width. The
@@ -2850,10 +2853,14 @@ def test_unpadded_page_size_includes_per_token_head_scales():
 
 def test_page_size_padded_wins():
     # An explicit padded page size takes precedence over the unpadded size.
-    spec = new_kv_cache_spec(
-        dtype=torch.uint8,
-        kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD,
-        page_size_padded=65536,
+    from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
+
+    spec = TritonAttentionBackend.customize_spec(
+        new_kv_cache_spec(
+            dtype=torch.uint8,
+            kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD,
+            page_size_padded=65536,
+        )
     )
     assert spec.page_size_bytes == 65536
 
@@ -2996,14 +3003,18 @@ def test_unify_kv_cache_spec_page_size_mamba():
     )
     assert unified["mamba_layer"].page_size_bytes == 32768
 
-    # 4. Attention layers with non-divisible page sizes still raise.
-    with pytest.raises(NotImplementedError):
-        kv_cache_utils.unify_kv_cache_spec_page_size(
-            {
-                "attn_layer": new_kv_cache_spec(block_size=24),  # 24576
-                "draft_attn_layer": draft_attn_spec,  # 32768
-            }
-        )
+    # 4. Attention layers with non-divisible page sizes are padded too: every
+    # backend reads a padded page through the view's block stride, so there is
+    # no longer a case that must raise.
+    unified = kv_cache_utils.unify_kv_cache_spec_page_size(
+        {
+            "attn_layer": new_kv_cache_spec(block_size=24),  # 24576
+            "draft_attn_layer": draft_attn_spec,  # 32768
+        }
+    )
+    assert unified["attn_layer"].page_size_padded == 32768
+    assert unified["attn_layer"].page_size_bytes == 32768
+    assert unified["attn_layer"].block_size == 24
 
     # 5. Uniform page sizes are returned unchanged.
     specs = {
@@ -3135,6 +3146,7 @@ def test_auto_fit_max_model_len_reserves_null_block():
     model_config = ModelConfig(max_model_len=1024)
     model_config.original_max_model_len = -1
     vllm_config = VllmConfig(model_config=model_config)
+    vllm_config.cache_config.kv_cache_layout = "LBNHC"
     spec = new_kv_cache_spec(block_size=block_size)
 
     # Exactly the 1024 / 16 = 64 blocks a full-length request would need.

tests/v1/core/test_prefix_caching.py

@@ -154,7 +154,7 @@ def make_kv_cache_config_hybrid_model(
     elif second_spec_type == "mamba":
         second_spec = MambaSpec(
             block_size=block_size,
-            shapes=(1, 1),
+            shapes=((1, 1),),
             dtypes=(torch.float32,),
         )
 
@@ -189,7 +189,7 @@ def make_kv_cache_config_three_types(
     if third_spec_type == "mamba":
         third_spec = MambaSpec(
             block_size=block_size,
-            shapes=(1, 1),
+            shapes=((1, 1),),
             dtypes=(torch.float32,),
         )
     elif third_spec_type == "sliding_window":
@@ -768,12 +768,12 @@ def _make_hybrid_kv_cache_config(
         ),
         "mamba": lambda: MambaSpec(
             block_size=block_size,
-            shapes=(1, 1),
+            shapes=((1, 1),),
             dtypes=(torch.float32,),
         ),
         "mamba_align": lambda: MambaSpec(
             block_size=block_size,
-            shapes=(1, 1),
+            shapes=((1, 1),),
             dtypes=(torch.float32,),
             mamba_cache_mode="align",
         ),
@@ -3273,7 +3273,7 @@ def test_hybrid_local_kv_retention_interval_survives_recycling():
                     num_kv_heads=1,
                     head_size=1,
                     dtype=torch.uint8,
-                    compress_ratio=4,
+                    tokens_per_state=4,
                 ),
             ),
             KVCacheGroupSpec(

tests/v1/engine/test_engine_core_client.py

@@ -1329,6 +1329,7 @@ def test_engine_core_proc_instantiation_cuda_empty(monkeypatch: pytest.MonkeyPat
         )
 
         mock_executor.get_kv_cache_specs.return_value = [{"default": mock_spec}]
+        mock_executor.get_supported_kv_cache_layouts.return_value = [["LBNHC"]]
         mock_executor.determine_available_memory.return_value = [1024 * 1024 * 1024]
         mock_executor.initialize_from_config.return_value = None
         mock_executor.supports_draft_weight_updates.return_value = False

tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh

@@ -36,7 +36,9 @@ if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
   echo "vLLM serve extra args: $VLLM_SERVE_EXTRA_ARGS"
 fi
 
-DECODER_KV_LAYOUT=${DECODER_KV_LAYOUT:-"HND"} # Default to HND, optional NHD
+# Empty leaves the layout to the backends, for models that constrain it.
+PREFILLER_KV_LAYOUT=${PREFILLER_KV_LAYOUT-"HND"}
+DECODER_KV_LAYOUT=${DECODER_KV_LAYOUT-"HND"} # Default to HND, optional NHD
 if [[ "$DECODER_KV_LAYOUT" == "NHD" ]]; then
   KV_CONFIG_HETERO_LAYOUT=',"enable_permute_local_kv":"True"'
 else
@@ -44,9 +46,9 @@ else
 fi
 
 if [[ "$CROSS_LAYERS_BLOCKS" == "True" ]]; then
-  KV_EXTRA_CONFIG=',"kv_connector_extra_config":{"enable_cross_layers_blocks": "True"}'
-else
-  KV_EXTRA_CONFIG=''
+  # Cross-layer blocks are exercised via the block-outermost layout.
+  PREFILLER_KV_LAYOUT="BLHNC"
+  DECODER_KV_LAYOUT="BLHNC"
 fi
 
 # Connector: default pull NixlConnector; NixlPushConnector enables PP prefill.
@@ -54,11 +56,11 @@ KV_CONNECTOR=${KV_CONNECTOR:-NixlConnector}
 
 # Build the kv-transfer-config for P and D
 if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then
-  KV_CONFIG_P='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_producer"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'
-  KV_CONFIG_D='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_consumer"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'
+  KV_CONFIG_P='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_producer"'${KV_CONFIG_HETERO_LAYOUT}'}'
+  KV_CONFIG_D='{"kv_connector":"'"$KV_CONNECTOR"'","kv_role":"kv_consumer"'${KV_CONFIG_HETERO_LAYOUT}'}'
 else
-  KV_CONFIG_P="{\"kv_connector\":\"$KV_CONNECTOR\",\"kv_role\":\"kv_producer\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"
-  KV_CONFIG_D="{\"kv_connector\":\"$KV_CONNECTOR\",\"kv_role\":\"kv_consumer\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"
+  KV_CONFIG_P="{\"kv_connector\":\"${KV_CONNECTOR}\",\"kv_role\":\"kv_producer\",\"kv_buffer_device\":\"${KV_BUFFER_DEVICE}\"${KV_CONFIG_HETERO_LAYOUT}}"
+  KV_CONFIG_D="{\"kv_connector\":\"${KV_CONNECTOR}\",\"kv_role\":\"kv_consumer\",\"kv_buffer_device\":\"${KV_BUFFER_DEVICE}\"${KV_CONFIG_HETERO_LAYOUT}}"
 fi
 
 # Models to run
@@ -176,7 +178,7 @@ run_tests_for_model() {
 
     # Build the command with or without model-specific args
     BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
-    VLLM_KV_CACHE_LAYOUT='HND' \
+    ${PREFILLER_KV_LAYOUT:+VLLM_KV_CACHE_LAYOUT=$PREFILLER_KV_LAYOUT} \
     VLLM_PORT=$INTERNAL_PORT \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
@@ -239,7 +241,7 @@ run_tests_for_model() {
 
     # Build the command with or without model-specific args
     BASE_CMD="CUDA_VISIBLE_DEVICES=$GPU_ID \
-    VLLM_KV_CACHE_LAYOUT=$DECODER_KV_LAYOUT \
+    ${DECODER_KV_LAYOUT:+VLLM_KV_CACHE_LAYOUT=$DECODER_KV_LAYOUT} \
     $DECODER_INTERNAL_PORT_ENV \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \

tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh

@@ -9,8 +9,8 @@
 # test_accuracy.py.
 #
 # By default runs two configurations:
-#   1. Normal KV layout (NixlConnector without cross-layer blocks)
-#   2. Cross-layer KV layout (NixlConnector with enable_cross_layers_blocks)
+#   1. Normal KV layout (LBHNC)
+#   2. Cross-layer KV layout (BLHNC) via VLLM_KV_CACHE_LAYOUT
 #
 # Usage:
 #   bash tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
@@ -55,8 +55,6 @@ fi
 
 # ── KV transfer configs ─────────────────────────────────────────────────
 
-# Normal layout: OffloadingConnector prefers cross-layer but NixlConnector
-# does not, so MultiConnector.prefer_cross_layer_blocks = False.
 KV_CONFIG_NORMAL='{
   "kv_connector":"MultiConnector",
   "kv_role":"kv_both",
@@ -72,22 +70,6 @@ KV_CONFIG_NORMAL='{
 # Remove whitespace for CLI safety
 KV_CONFIG_NORMAL=$(echo "$KV_CONFIG_NORMAL" | tr -d '[:space:]')
 
-# Cross-layer layout: both connectors prefer cross-layer blocks.
-KV_CONFIG_CROSS_LAYERS='{
-  "kv_connector":"MultiConnector",
-  "kv_role":"kv_both",
-  "kv_connector_extra_config":{
-    "connectors":[
-      {"kv_connector":"NixlConnector","kv_role":"kv_both",
-       "kv_buffer_device":"'$KV_BUFFER_DEVICE'",
-       "kv_connector_extra_config":{"enable_cross_layers_blocks":"True"}},
-      {"kv_connector":"OffloadingConnector","kv_role":"kv_both",
-       "kv_connector_extra_config":{"cpu_bytes_to_use":1000000000}}
-    ]
-  }
-}'
-KV_CONFIG_CROSS_LAYERS=$(echo "$KV_CONFIG_CROSS_LAYERS" | tr -d '[:space:]')
-
 # ── Helpers ──────────────────────────────────────────────────────────────
 
 PROCESS_GROUPS=()
@@ -218,7 +200,7 @@ run_tests_for_model() {
   # ── Start prefill instance ──
   echo "Starting prefill instance on GPU $PREFILL_GPU, port $PREFILL_PORT"
   BASE_CMD="$DEVICE_VISIBILITY_ENV=$PREFILL_GPU \
-    VLLM_KV_CACHE_LAYOUT='HND' \
+    VLLM_KV_CACHE_LAYOUT='${VLLM_KV_CACHE_LAYOUT:-LBHNC}' \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_SIDE_CHANNEL_PORT \
     vllm serve $model_name \
@@ -245,7 +227,7 @@ run_tests_for_model() {
   # ── Start decode instance ──
   echo "Starting decode instance on GPU $DECODE_GPU, port $DECODE_PORT"
   BASE_CMD="$DEVICE_VISIBILITY_ENV=$DECODE_GPU \
-    VLLM_KV_CACHE_LAYOUT='HND' \
+    VLLM_KV_CACHE_LAYOUT='${VLLM_KV_CACHE_LAYOUT:-LBHNC}' \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_SIDE_CHANNEL_PORT \
     vllm serve $model_name \
@@ -305,7 +287,8 @@ for model in "${MODELS[@]}"; do
   fi
 
   if [[ -z "${SKIP_CROSS_LAYERS:-}" ]]; then
-    run_tests_for_model "$model" "$KV_CONFIG_CROSS_LAYERS" "MultiConnector cross-layer layout"
+    VLLM_KV_CACHE_LAYOUT=BLHNC \
+      run_tests_for_model "$model" "$KV_CONFIG_NORMAL" "MultiConnector cross-layer layout"
   fi
 done
 

tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh

@@ -95,7 +95,7 @@ run_tests_for_model() {
   # ── Start prefill instance ──
   echo "Starting prefill instance on GPU $PREFILL_GPU, port $PREFILL_PORT"
   BASE_CMD="CUDA_VISIBLE_DEVICES=$PREFILL_GPU \
-    VLLM_KV_CACHE_LAYOUT='HND' \
+    VLLM_KV_CACHE_LAYOUT='LBHNC' \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_SIDE_CHANNEL_PORT \
     vllm serve \"$model_name\" \
@@ -121,7 +121,7 @@ run_tests_for_model() {
   # ── Start decode instance ──
   echo "Starting decode instance on GPU $DECODE_GPU, port $DECODE_PORT"
   BASE_CMD="CUDA_VISIBLE_DEVICES=$DECODE_GPU \
-    VLLM_KV_CACHE_LAYOUT='HND' \
+    VLLM_KV_CACHE_LAYOUT='LBHNC' \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_SIDE_CHANNEL_PORT \
     vllm serve \"$model_name\" \

tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

@@ -249,7 +249,7 @@ run_test_for_device() {
     echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"
     env \
     ${GPU_DEVICE_VAR}="$GPU_ID" \
-    VLLM_KV_CACHE_LAYOUT='HND' \
+    VLLM_KV_CACHE_LAYOUT='LBHNC' \
     UCX_NET_DEVICES=all \
     ${VLLM_SSM_CONV_STATE_LAYOUT:+VLLM_SSM_CONV_STATE_LAYOUT=$VLLM_SSM_CONV_STATE_LAYOUT} \
     VLLM_NIXL_SIDE_CHANNEL_HOST="$NIXL_SIDE_CHANNEL_HOST" \
@@ -288,7 +288,7 @@ run_test_for_device() {
     echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"
     env \
     ${GPU_DEVICE_VAR}="$GPU_ID" \
-    VLLM_KV_CACHE_LAYOUT='HND' \
+    VLLM_KV_CACHE_LAYOUT='LBHNC' \
     UCX_NET_DEVICES=all \
     ${VLLM_SSM_CONV_STATE_LAYOUT:+VLLM_SSM_CONV_STATE_LAYOUT=$VLLM_SSM_CONV_STATE_LAYOUT} \
     VLLM_NIXL_SIDE_CHANNEL_HOST="$NIXL_SIDE_CHANNEL_HOST" \

tests/v1/kv_connector/unit/offloading_connector/test_canonical_mapping.py

@@ -12,6 +12,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.canonical_mapping i
     _verify_tiling,
     derive_canonical_mappings,
 )
+from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheGroupSpec,
@@ -37,12 +38,14 @@ def _ctx(rank, tp=1, dcp=1, pcp=1, interleave=1, total=None, heads=2):
 
 def _full_spec(num_kv_heads: int = 2, **kwargs) -> FullAttentionSpec:
     # block_size=4, head_dim=64, int8
-    return FullAttentionSpec(
-        block_size=4,
-        num_kv_heads=num_kv_heads,
-        head_size=64,
-        dtype=torch.int8,
-        **kwargs,
+    return TritonAttentionBackend.customize_spec(
+        FullAttentionSpec(
+            block_size=4,
+            num_kv_heads=num_kv_heads,
+            head_size=64,
+            dtype=torch.int8,
+            **kwargs,
+        )
     )
 
 
@@ -381,7 +384,9 @@ def test_fail_closed_cases():
     quant_spec = _full_spec(kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD)
     assert _try_mapping(quant_spec, _split_nhd_cache(quant_spec), _ctx(0, tp=4)) is None
     # Compressed MLA slots are not 1:1 with tokens
-    assert _try_mapping(_mla_spec(compress_ratio=2), None, _ctx(0, tp=2, dcp=2)) is None
+    assert (
+        _try_mapping(_mla_spec(tokens_per_state=2), None, _ctx(0, tp=2, dcp=2)) is None
+    )
     # Unrecognized physical layouts
     swapped = torch.zeros(
         NUM_BLOCKS,

tests/v1/kv_connector/unit/offloading_connector/test_config.py

@@ -72,43 +72,54 @@ def _make_vllm_config(
 
 def _make_kv_cache_config() -> KVCacheConfig:
     num_blocks = 16
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+    )
     kv_tensor = KVCacheTensor(
-        size=num_blocks * 8,
-        shared_by=["layer"],
-        block_stride=0,
+        size=spec.page_size_bytes * num_blocks,
+        layers=["layer"],
+        layer_stride=spec.page_size_bytes * num_blocks,
+        block_stride=spec.page_size_bytes,
     )
     return KVCacheConfig(
         num_blocks=num_blocks,
         kv_cache_tensors=[kv_tensor],
-        kv_cache_groups=[
-            KVCacheGroupSpec(
-                ["layer"],
-                FullAttentionSpec(
-                    block_size=16,
-                    num_kv_heads=1,
-                    head_size=1,
-                    dtype=torch.float32,
-                ),
-            )
-        ],
+        kv_cache_groups=[KVCacheGroupSpec(["layer"], spec)],
     )
 
 
 def _make_sizing_kv_cache_config(packed: bool) -> KVCacheConfig:
+    """One 16 byte-per-block allocation, described two ways.
+
+    Packed: both layers are one dense run. Unpacked: the same bytes as two runs, the
+    second starting after the first layer's region. Either way the connector accounts
+    for 16 KV bytes per block.
+    """
     num_blocks = 4
+    page = 8
+    size = 2 * page * num_blocks
     if packed:
         kv_cache_tensors = [
             KVCacheTensor(
-                size=64,
-                shared_by=[layer_name],
-                block_stride=16,
+                size=size,
+                layers=["layer0", "layer1"],
+                layer_stride=page * num_blocks,
+                block_stride=page,
             )
-            for layer_name in ("layer0", "layer1")
         ]
     else:
         kv_cache_tensors = [
-            KVCacheTensor(size=40, shared_by=["layer0"]),
-            KVCacheTensor(size=24, shared_by=["layer1"]),
+            KVCacheTensor(
+                size=size,
+                layers=[layer],
+                layer_stride=page * num_blocks,
+                block_stride=page,
+                offset=i * page * num_blocks,
+            )
+            for i, layer in enumerate(("layer0", "layer1"))
         ]
 
     return KVCacheConfig(
@@ -150,6 +161,18 @@ def _mla_spec(
     )
 
 
+_MAMBA_SPEC = MambaSpec(
+    block_size=16,
+    shapes=((16, 1),),
+    dtypes=(torch.float32,),
+)
+# Page sizes of the specs the replicated-layout cases below are built from.
+_MLA_PAGE = _mla_spec().page_size_bytes
+_HALF_MLA_PAGE = _mla_spec(head_size=256).page_size_bytes
+_FULL_PAGE = _full_attention_spec().page_size_bytes
+_MAMBA_PAGE = _MAMBA_SPEC.page_size_bytes
+
+
 def _make_mla_kv_cache_config(
     layer_names: list[str] | None = None,
     head_size: int = 512,
@@ -159,12 +182,14 @@ def _make_mla_kv_cache_config(
     if layer_names is None:
         layer_names = ["layer0", "layer1"]
     spec = _mla_spec(head_size=head_size, dtype=dtype)
+    layer_stride = spec.page_size_bytes * num_blocks
     kv_cache_tensors = [
         KVCacheTensor(
-            size=spec.page_size_bytes * num_blocks,
-            shared_by=[layer_name],
+            size=layer_stride * len(layer_names),
+            layers=layer_names,
+            layer_stride=layer_stride,
+            block_stride=spec.page_size_bytes,
         )
-        for layer_name in layer_names
     ]
     return KVCacheConfig(
         num_blocks=num_blocks,
@@ -174,15 +199,31 @@ def _make_mla_kv_cache_config(
 
 
 def _make_hybrid_kv_cache_config() -> KVCacheConfig:
+    num_blocks = 4
+    full_spec = _full_attention_spec(block_size=12)
+    mla_spec = _mla_spec()
+    # Mixed page sizes across overlaying groups: a block is a window of the largest
+    # group's packing, so the layer dim sits inside the block dim.
+    window = max(full_spec.page_size_bytes, mla_spec.page_size_bytes)
     return KVCacheConfig(
-        num_blocks=4,
+        num_blocks=num_blocks,
         kv_cache_tensors=[
-            KVCacheTensor(size=40, shared_by=["full_layer"]),
-            KVCacheTensor(size=24, shared_by=["mla_layer"]),
+            KVCacheTensor(
+                size=window * num_blocks,
+                layers=["full_layer"],
+                layer_stride=full_spec.page_size_bytes,
+                block_stride=window,
+            ),
+            KVCacheTensor(
+                size=window * num_blocks,
+                layers=["mla_layer"],
+                layer_stride=mla_spec.page_size_bytes,
+                block_stride=window,
+            ),
         ],
         kv_cache_groups=[
-            KVCacheGroupSpec(["full_layer"], _full_attention_spec(block_size=12)),
-            KVCacheGroupSpec(["mla_layer"], _mla_spec()),
+            KVCacheGroupSpec(["full_layer"], full_spec),
+            KVCacheGroupSpec(["mla_layer"], mla_spec),
         ],
     )
 
@@ -263,18 +304,9 @@ def test_worker_kv_bytes_preserves_tensor_layout(packed: bool):
     assert offloading_config.cache.blocks_per_chunk == 2
 
 
-def test_rejects_partially_packed_tensor_layout():
-    kv_cache_config = _make_sizing_kv_cache_config(packed=False)
-    kv_cache_config.kv_cache_tensors[0].block_stride = 16
-
-    with pytest.raises(AssertionError):
-        build_offloading_config(_make_vllm_config(), kv_cache_config)
-
-
 def test_zero_blocks_skips_tensor_layout_validation():
     kv_cache_config = _make_sizing_kv_cache_config(packed=False)
     kv_cache_config.num_blocks = 0
-    kv_cache_config.kv_cache_tensors[0].block_stride = 16
 
     offloading_config = build_offloading_config(_make_vllm_config(), kv_cache_config)
 
@@ -374,8 +406,10 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
                 num_blocks=4,
                 kv_cache_tensors=[
                     KVCacheTensor(
-                        size=_mla_spec().page_size_bytes * 4,
-                        shared_by=["layer"],
+                        size=_MLA_PAGE * 4,
+                        layers=["layer"],
+                        layer_stride=_MLA_PAGE * 4,
+                        block_stride=_MLA_PAGE,
                     )
                 ],
                 kv_cache_groups=[
@@ -398,8 +432,10 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
                 num_blocks=4,
                 kv_cache_tensors=[
                     KVCacheTensor(
-                        size=_mla_spec().page_size_bytes * 4,
-                        shared_by=["layer"],
+                        size=_MLA_PAGE * 4,
+                        layers=["layer"],
+                        layer_stride=_MLA_PAGE * 4,
+                        block_stride=_MLA_PAGE,
                     )
                 ],
                 kv_cache_groups=[
@@ -417,16 +453,22 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
             "hidden-state",
         ),
         (
+            # One group, two page sizes: layer1's run starts past layer0's region.
             KVCacheConfig(
                 num_blocks=4,
                 kv_cache_tensors=[
                     KVCacheTensor(
-                        size=_mla_spec().page_size_bytes * 4,
-                        shared_by=["layer0"],
+                        size=(_MLA_PAGE + _HALF_MLA_PAGE) * 4,
+                        layers=["layer0"],
+                        layer_stride=_MLA_PAGE * 4,
+                        block_stride=_MLA_PAGE,
                     ),
                     KVCacheTensor(
-                        size=_mla_spec(head_size=256).page_size_bytes * 4,
-                        shared_by=["layer1"],
+                        size=(_MLA_PAGE + _HALF_MLA_PAGE) * 4,
+                        layers=["layer1"],
+                        layer_stride=_HALF_MLA_PAGE * 4,
+                        block_stride=_HALF_MLA_PAGE,
+                        offset=_MLA_PAGE * 4,
                     ),
                 ],
                 kv_cache_groups=[
@@ -445,16 +487,22 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
             "uniform-wrapper",
         ),
         (
+            # Overlaid groups with different page sizes: a block is a window of
+            # the largest group's packing.
             KVCacheConfig(
                 num_blocks=4,
                 kv_cache_tensors=[
                     KVCacheTensor(
-                        size=_mla_spec().page_size_bytes * 4,
-                        shared_by=["mla"],
+                        size=_FULL_PAGE * 4,
+                        layers=["mla"],
+                        layer_stride=_MLA_PAGE,
+                        block_stride=_FULL_PAGE,
                     ),
                     KVCacheTensor(
-                        size=_full_attention_spec().page_size_bytes * 4,
-                        shared_by=["full"],
+                        size=_FULL_PAGE * 4,
+                        layers=["full"],
+                        layer_stride=_FULL_PAGE,
+                        block_stride=_FULL_PAGE,
                     ),
                 ],
                 kv_cache_groups=[
@@ -469,21 +517,21 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
                 num_blocks=4,
                 kv_cache_tensors=[
                     KVCacheTensor(
-                        size=_mla_spec().page_size_bytes * 4,
-                        shared_by=["mla"],
+                        size=_MLA_PAGE * 4,
+                        layers=["mla"],
+                        layer_stride=_MLA_PAGE,
+                        block_stride=_MLA_PAGE,
+                    ),
+                    KVCacheTensor(
+                        size=_MLA_PAGE * 4,
+                        layers=["mamba"],
+                        layer_stride=_MAMBA_PAGE,
+                        block_stride=_MLA_PAGE,
                     ),
-                    KVCacheTensor(size=64 * 4, shared_by=["mamba"]),
                 ],
                 kv_cache_groups=[
                     KVCacheGroupSpec(["mla"], _mla_spec()),
-                    KVCacheGroupSpec(
-                        ["mamba"],
-                        MambaSpec(
-                            block_size=16,
-                            shapes=((16, 1),),
-                            dtypes=(torch.float32,),
-                        ),
-                    ),
+                    KVCacheGroupSpec(["mamba"], _MAMBA_SPEC),
                 ],
             ),
             "mla-mamba-hybrid",
@@ -493,13 +541,12 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
                 num_blocks=4,
                 kv_cache_tensors=[
                     KVCacheTensor(
-                        size=_mla_spec().page_size_bytes * 4,
-                        shared_by=["layer0"],
-                    ),
-                    KVCacheTensor(
-                        size=_mla_spec().page_size_bytes * 4,
-                        shared_by=["layer1"],
-                    ),
+                        size=_MLA_PAGE * 4,
+                        layers=[layer],
+                        layer_stride=_MLA_PAGE * 4,
+                        block_stride=_MLA_PAGE,
+                    )
+                    for layer in ("layer0", "layer1")
                 ],
                 kv_cache_groups=[
                     KVCacheGroupSpec(["layer0"], _mla_spec()),
@@ -531,16 +578,26 @@ def test_replicated_layout_rejects_bare_mla_with_mixed_page_accounting():
     indexer_spec = _mla_spec(head_size=128, dtype=torch.uint8)
     main_layers = [f"main_{i}" for i in range(61)]
     indexer_layers = [f"indexer_{i}" for i in range(61)]
+    # A DSA-style group: the main pages and the smaller indexer pages are packed one
+    # after the other, so a block holds more than 61 MLA pages.
+    main_bytes = main_spec.page_size_bytes * len(main_layers)
+    indexer_bytes = indexer_spec.page_size_bytes * len(indexer_layers)
+    size = (main_bytes + indexer_bytes) * num_blocks
     kv_cache_config = KVCacheConfig(
         num_blocks=num_blocks,
         kv_cache_tensors=[
             KVCacheTensor(
-                size=main_spec.page_size_bytes * len(main_layers) * num_blocks,
-                shared_by=main_layers,
+                size=size,
+                layers=main_layers,
+                layer_stride=main_spec.page_size_bytes * num_blocks,
+                block_stride=main_spec.page_size_bytes,
             ),
             KVCacheTensor(
-                size=indexer_spec.page_size_bytes * len(indexer_layers) * num_blocks,
-                shared_by=indexer_layers,
+                size=size,
+                layers=indexer_layers,
+                layer_stride=indexer_spec.page_size_bytes * num_blocks,
+                block_stride=indexer_spec.page_size_bytes,
+                offset=main_bytes * num_blocks,
             ),
         ],
         kv_cache_groups=[KVCacheGroupSpec(main_layers + indexer_layers, main_spec)],
@@ -665,34 +722,6 @@ def test_canonical_layout_certifies_v2_model_runner():
     ).parallel.is_parallelism_agnostic
 
 
-def test_prefer_cross_layer_blocks_yields_to_canonical_layout():
-    """The connector must not request cross-layer blocks under
-    canonical_layout: cross-layer slabs have no per-layer refs to certify."""
-    from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
-    from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
-        OffloadingConnector,
-    )
-
-    kv_cache_config = _make_kv_cache_config()
-    connector_module = "vllm.distributed.kv_transfer.kv_connector.v1"
-
-    def make_connector(extra_config: dict[str, Any] | None) -> OffloadingConnector:
-        with (
-            patch(f"{connector_module}.offloading_connector.OffloadingSpecFactory"),
-            patch(
-                f"{connector_module}.offloading_connector.OffloadingConnectorScheduler"
-            ),
-        ):
-            return OffloadingConnector(
-                _make_vllm_config(extra_config=extra_config),
-                KVConnectorRole.SCHEDULER,
-                kv_cache_config,
-            )
-
-    assert make_connector(None).prefer_cross_layer_blocks
-    assert not make_connector({"canonical_layout": True}).prefer_cross_layer_blocks
-
-
 def test_parallelism_agnostic_disabled_on_v2_model_runner():
     config = _make_vllm_config()
     config.use_v2_model_runner = True

tests/v1/kv_connector/unit/offloading_connector/test_worker.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from collections import defaultdict
 from unittest.mock import MagicMock
 
 import pytest
@@ -13,11 +12,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import get_dtype_size
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.v1.attention.backends.utils import set_kv_cache_layout
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
+    KVCacheLayout,
     KVCacheTensor,
     MambaSpec,
     MLAAttentionSpec,
@@ -65,40 +64,18 @@ elif current_platform.is_xpu():
 # ---------------------------------------------------------------------------
 
 
-def _allocate_and_reshape_kv_caches(
+def _allocate_kv_caches(
     kv_cache_config: KVCacheConfig,
     attn_groups: list[list],
     device: torch.device,
 ):
-    """
-    Use the real GPUModelRunner allocation and reshape methods to produce
-    kv_caches, just like the model runner does during initialization.
-    """
-    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+    """Allocate kv_caches exactly as the model runner does at startup."""
+    from vllm.v1.worker.utils import allocate_kv_cache
 
-    # Some backends (e.g. FlashAttention) query the KV cache layout during
-    # reshape, which ultimately calls get_current_vllm_config(). Setting
-    # the layout override avoids needing a full VllmConfig context.
-    set_kv_cache_layout("NHD")
-    try:
-        runner = object.__new__(GPUModelRunner)
-        runner.device = device
-        runner.runner_only_attn_layers = set()
-        runner.attn_groups = attn_groups
-        runner.kv_cache_config = kv_cache_config
-        runner.cache_config = MagicMock(cache_dtype="auto")
-        runner.shared_kv_cache_layers = {}
-        runner.model_config = MagicMock()
-        runner.model_config.hf_config.model_type = ""
-        runner.compilation_config = MagicMock(
-            static_forward_context=defaultdict(MagicMock)
-        )
-        runner.kv_caches = []
-
-        kernel_block_sizes = [BLOCK_SIZE] * len(kv_cache_config.kv_cache_groups)
-        return runner.initialize_kv_cache_tensors(kv_cache_config, kernel_block_sizes)
-    finally:
-        set_kv_cache_layout(None)
+    kernel_block_sizes = [BLOCK_SIZE] * len(kv_cache_config.kv_cache_groups)
+    return allocate_kv_cache(
+        kv_cache_config, device, KVCacheLayout.LBNHC, kernel_block_sizes
+    )
 
 
 def _single_rank_vllm_config(total_kv_heads: int):
@@ -321,10 +298,9 @@ def test_register_kv_caches(backend):
     Creates one FullAttention group, one MLA group, one Mamba group, and
     one Mamba-padded group. Each group has GROUP_SIZE layers.
 
-    KVCacheTensors are shared across all groups mirroring the real allocation
-    in kv_cache_utils.py: tensor i is shared by layer i from every group.
-    The padded-mamba group has a different page size so its layers get their
-    own dedicated tensors.
+    KVCacheTensors mirror the real allocation in kv_cache_utils.py: each group is one
+    run of the shared allocation starting at offset 0, so layer i of every group
+    aliases the same bytes.
 
     Uses the real GPUModelRunner.initialize_kv_cache_tensors to produce
     the raw per-layer kv_caches registered by the connector.
@@ -413,18 +389,15 @@ def test_register_kv_caches(backend):
         aligned_mamba_layer_names,
     ]
 
-    kv_cache_tensors: list[KVCacheTensor] = []
-    for i in range(GROUP_SIZE):
-        shared_by: list[str] = []
-        for group_layer_names in layer_groups:
-            if len(group_layer_names) > i:
-                shared_by.append(group_layer_names[i])
-        kv_cache_tensors.append(
-            KVCacheTensor(
-                size=PAGE_SIZE_BYTES * NUM_BLOCKS,
-                shared_by=shared_by,
-            )
+    kv_cache_tensors: list[KVCacheTensor] = [
+        KVCacheTensor(
+            size=PAGE_SIZE_BYTES * NUM_BLOCKS * GROUP_SIZE,
+            layers=group_layer_names,
+            layer_stride=PAGE_SIZE_BYTES * NUM_BLOCKS,
+            block_stride=PAGE_SIZE_BYTES,
         )
+        for group_layer_names in layer_groups
+    ]
 
     kv_cache_groups = [
         KVCacheGroupSpec(layer_names=attn_layer_names, kv_cache_spec=attn_spec),
@@ -472,7 +445,7 @@ def test_register_kv_caches(backend):
         kv_cache_groups=kv_cache_groups,
     )
 
-    kv_caches = _allocate_and_reshape_kv_caches(
+    kv_caches = _allocate_kv_caches(
         kv_cache_config,
         attn_groups,
         device=torch.device(f"{DEVICE_TYPE}:0"),
@@ -580,16 +553,24 @@ def test_register_kv_caches_uniform_type(backend):
         kv_cache_specs={layer_a: spec_a, layer_b: spec_b},
     )
 
+    # The group packs its two layers densely, one run each as their pages differ:
+    # layer_a occupies the first page of every block, layer_b the rest.
+    window = spec_a.page_size_bytes + spec_b.page_size_bytes
     kv_cache_config = KVCacheConfig(
         num_blocks=NUM_BLOCKS,
         kv_cache_tensors=[
             KVCacheTensor(
-                size=spec_a.page_size_bytes * NUM_BLOCKS,
-                shared_by=[layer_a],
+                size=window * NUM_BLOCKS,
+                layers=[layer_a],
+                layer_stride=spec_a.page_size_bytes * NUM_BLOCKS,
+                block_stride=spec_a.page_size_bytes,
             ),
             KVCacheTensor(
-                size=spec_b.page_size_bytes * NUM_BLOCKS,
-                shared_by=[layer_b],
+                size=window * NUM_BLOCKS,
+                layers=[layer_b],
+                layer_stride=spec_b.page_size_bytes * NUM_BLOCKS,
+                block_stride=spec_b.page_size_bytes,
+                offset=spec_a.page_size_bytes * NUM_BLOCKS,
             ),
         ],
         kv_cache_groups=[
@@ -617,7 +598,7 @@ def test_register_kv_caches_uniform_type(backend):
         ]
     ]
 
-    kv_caches = _allocate_and_reshape_kv_caches(
+    kv_caches = _allocate_kv_caches(
         kv_cache_config,
         attn_groups,
         device=torch.device(f"{DEVICE_TYPE}:0"),

tests/v1/kv_connector/unit/offloading_connector/utils.py

@@ -241,13 +241,19 @@ class RequestRunner:
                 )
             ]
 
+        # Groups overlay one allocation, sized by the largest group.
+        block_stride = max(
+            group.kv_cache_spec.page_size_bytes * len(group.layer_names)
+            for group in kv_cache_groups
+        )
         kv_cache_tensors = [
             KVCacheTensor(
-                size=group.kv_cache_spec.page_size_bytes * num_gpu_blocks,
-                shared_by=[layer_name],
+                size=block_stride * num_gpu_blocks,
+                layers=list(group.layer_names),
+                layer_stride=group.kv_cache_spec.page_size_bytes * num_gpu_blocks,
+                block_stride=group.kv_cache_spec.page_size_bytes,
             )
             for group in kv_cache_groups
-            for layer_name in group.layer_names
         ]
 
         kv_cache_config = KVCacheConfig(
@@ -277,7 +283,6 @@ class RequestRunner:
         )
 
         # register worker kv_caches to enable OffloadingWorker creations
-        # set_current_vllm_config is needed for get_kv_cache_layout() to work
         kv_caches: dict[str, torch.Tensor] = {}
         for group in kv_cache_groups:
             spec = group.kv_cache_spec

tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py

@@ -98,7 +98,7 @@ def _make_connector_with_fake_worker(
     )
     worker = connector.connector_worker
     assert isinstance(worker.nixl_wrapper, FakeNixlWrapper)
-    worker.kv_cache_layout = "HND"
+    worker.kv_cache_layout = "LBHNC"
     if do_handshake:
         remote_agents, _ = worker._nixl_handshake(
             host="localhost",

tests/v1/kv_connector/unit/test_kv_cache_layout.py

@@ -1,75 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-
-from vllm.platforms import current_platform
-
-
-def test_mla_common_backend_rejects_cross_layer_kv_cache():
-    """MLACommonBackend defaults to the identity permutation (layers dim
-    first) so MLA backends whose decode kernels are not verified to honor
-    the cache's block-dim stride stay opted out of cross-layer KV cache."""
-    from vllm.model_executor.layers.attention.mla_attention import (
-        MLACommonBackend,
-    )
-
-    stride_order = MLACommonBackend.get_kv_cache_stride_order(
-        include_num_layers_dimension=True
-    )
-    assert stride_order == (0, 1, 2, 3)
-    assert stride_order[0] == 0  # layers dim first => no cross-layer
-    assert MLACommonBackend.get_kv_cache_stride_order(
-        include_num_layers_dimension=False
-    ) == (0, 1, 2)
-
-
-@pytest.mark.parametrize(
-    "backend_path",
-    # See: https://github.com/vllm-project/vllm/issues/46411
-    [
-        "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend",
-    ]
-    if current_platform.is_rocm() or current_platform.is_xpu()
-    else [
-        "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend",
-        "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend",
-        "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend",
-        "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend",
-        "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend",
-    ],
-)
-def test_verified_mla_backends_support_cross_layer_kv_cache(backend_path):
-    """Backends whose decode kernels honor the cache's block-dim stride opt
-    in to the cross-layer layout with a non-identity permutation placing
-    num_blocks first in physical layout."""
-    module_path, name = backend_path.rsplit(".", 1)
-    backend = getattr(
-        pytest.importorskip(module_path, reason="backend deps unavailable"), name
-    )
-
-    stride_order = backend.get_kv_cache_stride_order(include_num_layers_dimension=True)
-    assert stride_order == (1, 0, 2, 3)
-    assert stride_order[0] != 0  # num_blocks first => cross-layer supported
-    assert backend.get_kv_cache_stride_order(include_num_layers_dimension=False) == (
-        0,
-        1,
-        2,
-    )
-
-
-def test_deepseek_v32_indexer_rejects_cross_layer_kv_cache():
-    """DeepseekV32Indexer returns identity permutation (layers dim first)
-    to signal cross-layer KV cache is unsupported."""
-    from vllm.v1.attention.backends.mla.indexer import (
-        DeepseekV32IndexerBackend,
-    )
-
-    stride_order = DeepseekV32IndexerBackend.get_kv_cache_stride_order(
-        include_num_layers_dimension=True
-    )
-    assert stride_order == (0, 1, 2, 3)
-    assert stride_order[0] == 0  # layers dim first => no cross-layer
-    assert DeepseekV32IndexerBackend.get_kv_cache_stride_order(
-        include_num_layers_dimension=False
-    ) == (0, 1, 2)

tests/v1/kv_connector/unit/test_mooncake_connector.py

@@ -11,6 +11,7 @@ import pytest
 import torch
 import zmq.asyncio
 
+from tests.v1.attention.utils import dense_kv_cache_views
 from vllm import envs
 from vllm.config import set_current_vllm_config
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
@@ -32,11 +33,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import
     MooncakeBootstrapServer,
 )
 from vllm.utils.network_utils import get_open_port
-from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
+    KVCacheLayout,
 )
 from vllm.v1.request import RequestStatus
 
@@ -1093,7 +1094,19 @@ async def test_worker_get_finished_timeout(monkeypatch):
         assert "tx-active" in prefill_worker.reqs_need_send
 
 
-def test_register_kv_caches():
+@pytest.mark.parametrize(
+    ("layout", "separate_kv_head_groups"),
+    [
+        (KVCacheLayout.LBHNC, False),
+        (KVCacheLayout.BLHNC, False),
+        (KVCacheLayout.BHLNC, False),
+        # LHBNC gives each head group its own region; the K/V split doubles the
+        # head count but the registration shape is driven by the layout.
+        (KVCacheLayout.LHBNC, False),
+        (KVCacheLayout.LHBNC, True),
+    ],
+)
+def test_register_kv_caches(layout: KVCacheLayout, separate_kv_head_groups: bool):
     """Tests the memory registration logic with the underlying Mooncake engine."""
 
     vllm_config = create_vllm_config(
@@ -1118,15 +1131,23 @@ def test_register_kv_caches():
         worker = connector.connector_worker
         mock_thread.return_value.is_alive.return_value = False
 
-        kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(
-            num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
+        spec = FullAttentionSpec(
+            block_size=16,
+            num_kv_heads=4,
+            head_size=64,
+            dtype=torch.float16,
+            num_head_slots=2 if separate_kv_head_groups else None,
+            state_content_bytes=4 * 64 * 2 if separate_kv_head_groups else None,
         )
-        tensor1 = torch.zeros(*kv_cache_shape, dtype=torch.float16)
-        tensor2 = torch.zeros(*kv_cache_shape, dtype=torch.float16)
-        kv_caches = {
-            "model.layers.0.self_attn": tensor1,
-            "model.layers.1.self_attn": tensor2,
-        }
+        layer_names = [
+            "model.layers.0.self_attn",
+            "model.layers.1.self_attn",
+        ]
+        for layer_name in layer_names:
+            worker._layer_specs[layer_name] = spec
+        raw = torch.zeros(2 * 2 * spec.page_size_bytes, dtype=torch.int8)
+        tensor1, tensor2 = dense_kv_cache_views(raw, spec, 2, 2, layout)
+        kv_caches = dict(zip(layer_names, (tensor1, tensor2)))
 
         with patch.object(
             worker.engine, "batch_register_memory", return_value=0
@@ -1135,16 +1156,35 @@ def test_register_kv_caches():
 
             mock_batch_register.assert_called_once()
             registered_ptrs, registered_lens = mock_batch_register.call_args[0]
-            expected_ptrs = {tensor.data_ptr() for tensor in kv_caches.values()}
-            assert set(registered_ptrs) == expected_ptrs
-            assert set(registered_lens) == {tensor1.nbytes}
+            assert registered_ptrs == [raw.data_ptr()]
+            assert registered_lens == [raw.nbytes]
 
-            # Verify block_len_per_layer is set correctly.
-            assert len(worker.block_len_per_layer) == len(registered_ptrs)
-            for bl in worker.block_len_per_layer:
-                assert bl == tensor1.nbytes // tensor1.shape[0]
-            assert worker.registered_layer_names == list(kv_caches)
-            assert worker.registered_layer_indices == [0, 1]
+            if not layout.is_block_compact:
+                expected_addrs = [
+                    cache[:, head_idx].data_ptr()
+                    for cache in (tensor1, tensor2)
+                    for head_idx in range(cache.shape[1])
+                ]
+                head_block_bytes = tensor1.stride(0) * tensor1.element_size()
+                assert worker.kv_caches_base_addr == expected_addrs
+                assert worker.block_len_per_layer == [head_block_bytes] * len(
+                    expected_addrs
+                )
+                assert worker.kv_block_len_per_layer == [head_block_bytes] * len(
+                    expected_addrs
+                )
+                assert worker.registered_layer_names == [
+                    layer_name
+                    for layer_name in layer_names
+                    for _ in range(tensor1.shape[1])
+                ]
+            else:
+                assert len(worker.block_len_per_layer) == len(kv_caches)
+                for bl in worker.block_len_per_layer:
+                    assert bl == tensor1.stride(0) * tensor1.element_size()
+                assert worker.kv_block_len_per_layer == [spec.page_size_bytes] * 2
+                assert worker.registered_layer_names == list(kv_caches)
+                assert worker.registered_layer_indices == [0, 1]
 
 
 def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():

tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py

@@ -14,6 +14,7 @@ from unittest.mock import patch
 import pytest
 import torch
 
+from tests.v1.attention.utils import dense_kv_cache_views
 from vllm.config import set_current_vllm_config
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (
     KVConnectorRole,
@@ -30,6 +31,7 @@ from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
+    KVCacheLayout,
     MambaSpec,
 )
 
@@ -138,14 +140,22 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch):
         )
         worker = connector.connector_worker
 
-        fa_cache = torch.empty((2, 2, 11), dtype=torch.float16)
-        gdn_conv_state = torch.empty((2, 22), dtype=torch.float16)
-        gdn_ssm_state = torch.empty((2, 4), dtype=torch.float16)
+        num_blocks = kv_cache_config.num_blocks
+        fa_spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec
+        gdn_spec = kv_cache_config.kv_cache_groups[1].kv_cache_spec
+        fa_raw = torch.empty(num_blocks * fa_spec.page_size_bytes, dtype=torch.int8)
+        gdn_raw = torch.empty(num_blocks * gdn_spec.page_size_bytes, dtype=torch.int8)
+        (fa_cache,) = dense_kv_cache_views(
+            fa_raw, fa_spec, num_blocks, 1, KVCacheLayout.LBHNC
+        )
+        (gdn_cache,) = dense_kv_cache_views(
+            gdn_raw, gdn_spec, num_blocks, 1, KVCacheLayout.LBHNC
+        )
 
         worker.register_kv_caches(
             {
                 "model.layers.0.self_attn": fa_cache,
-                "model.layers.1.linear_attn": (gdn_conv_state, gdn_ssm_state),
+                "model.layers.1.linear_attn": gdn_cache,
             }
         )
 
@@ -154,10 +164,14 @@ def test_register_kv_caches_emits_fa_and_gdn_regions(monkeypatch):
             "model.layers.0.self_attn",
             "model.layers.1.linear_attn",
         ]
+        assert worker.block_len_per_layer == [
+            fa_spec.page_size_bytes,
+            gdn_spec.page_size_bytes,
+        ]
         assert worker.registered_group_indices == [0, 1]
         assert worker.kv_caches_base_addr == [
             fa_cache.data_ptr(),
-            gdn_conv_state.data_ptr(),
+            gdn_cache.data_ptr(),
         ]
 
         worker.shutdown()
@@ -185,8 +199,7 @@ def test_register_kv_caches_deduplicates_shared_backing_memory(monkeypatch):
 
         backing = torch.empty((4, 64), dtype=torch.float16)
         fa_cache = backing[:2, :16]
-        gdn_conv_state = backing[:3]
-        gdn_ssm_state = torch.empty((3, 4), dtype=torch.float16)
+        gdn_cache = backing[:3]
 
         with patch.object(
             worker.engine, "batch_register_memory", return_value=0
@@ -194,13 +207,13 @@ def test_register_kv_caches_deduplicates_shared_backing_memory(monkeypatch):
             worker.register_kv_caches(
                 {
                     "model.layers.0.self_attn": fa_cache,
-                    "model.layers.1.linear_attn": (gdn_conv_state, gdn_ssm_state),
+                    "model.layers.1.linear_attn": gdn_cache,
                 }
             )
 
         assert worker.kv_caches_base_addr == [
             fa_cache.data_ptr(),
-            gdn_conv_state.data_ptr(),
+            gdn_cache.data_ptr(),
         ]
         batch_register_memory.assert_called_once()
         registered_ptrs, registered_lens = batch_register_memory.call_args[0]
@@ -339,7 +352,7 @@ def test_logical_to_kernel_block_ids_expands_fa_not_gdn():
     assert kernel_block_ids == [list(range(34, 51)), [2]]
 
 
-def test_hybrid_gdn_splits_fa_regions_but_keeps_gdn_state_whole(
+def test_hybrid_gdn_keeps_packed_fa_and_gdn_regions_whole(
     monkeypatch,
 ):
     monkeypatch.setenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "5")
@@ -359,7 +372,7 @@ def test_hybrid_gdn_splits_fa_regions_but_keeps_gdn_state_whole(
         )
         worker = connector.connector_worker
 
-        worker.transfer_topo = SimpleNamespace(virtually_split_kv_in_blocks=True)
+        worker.transfer_topo = SimpleNamespace(is_kv_layout_blocks_first=False)
         regions = worker._get_transfer_regions(
             base_addrs=[0x1000, 0x2000],
             block_lens=[0x100, 0x100],
@@ -377,7 +390,6 @@ def test_hybrid_gdn_splits_fa_regions_but_keeps_gdn_state_whole(
             for region in regions
         ] == [
             (0, 0x1000, 0x40),
-            (0, 0x1040, 0x40),
             (1, 0x2000, 0x100),
         ]
 

tests/v1/kv_connector/unit/test_mooncake_store_connector.py

@@ -48,7 +48,14 @@ def _make_kv_cache_config() -> KVCacheConfig:
     spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)
     return KVCacheConfig(
         num_blocks=4,
-        kv_cache_tensors=[KVCacheTensor(size=8192, shared_by=["layer0"])],
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=8192,
+                layers=["layer0"],
+                layer_stride=8192,
+                block_stride=2048,
+            )
+        ],
         kv_cache_groups=[KVCacheGroupSpec(["layer0"], spec)],
     )
 
@@ -210,64 +217,6 @@ def test_get_kv_connector_kv_cache_events_wraps_worker_events():
     assert kv_events.get_all_events() == [event]
 
 
-def test_prefer_cross_layer_blocks_from_config():
-    # Default: disabled
-    vllm_config = _make_vllm_config()
-    kv_cache_config = _make_kv_cache_config()
-    with (
-        set_current_vllm_config(vllm_config),
-        patch(
-            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
-            "connector.MooncakeStoreScheduler"
-        ),
-    ):
-        connector = mooncake_store_connector.MooncakeStoreConnector(
-            vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
-        )
-    assert connector.prefer_cross_layer_blocks is False
-
-    # Enabled via config
-    vllm_config_enabled = create_vllm_config(
-        kv_connector="MooncakeStoreConnector",
-        kv_role="kv_both",
-        kv_connector_extra_config={"enable_cross_layers_blocks": "true"},
-    )
-    with (
-        set_current_vllm_config(vllm_config_enabled),
-        patch(
-            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
-            "connector.MooncakeStoreScheduler"
-        ),
-    ):
-        connector_enabled = mooncake_store_connector.MooncakeStoreConnector(
-            vllm_config_enabled, KVConnectorRole.SCHEDULER, kv_cache_config
-        )
-    assert connector_enabled.prefer_cross_layer_blocks is True
-
-
-def test_register_cross_layers_kv_cache_delegates_to_worker():
-    vllm_config = _make_vllm_config()
-    kv_cache_config = _make_kv_cache_config()
-
-    with (
-        set_current_vllm_config(vllm_config),
-        patch(
-            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
-            "connector.MooncakeStoreWorker"
-        ) as mock_worker_cls,
-    ):
-        connector = mooncake_store_connector.MooncakeStoreConnector(
-            vllm_config, KVConnectorRole.WORKER, kv_cache_config
-        )
-
-    fake_tensor = MagicMock()
-    fake_backend = MagicMock()
-    connector.register_cross_layers_kv_cache(fake_tensor, fake_backend)
-
-    worker = mock_worker_cls.return_value
-    worker.register_cross_layers_kv_caches.assert_called_once_with(fake_tensor)
-
-
 def test_update_connector_output_and_take_events():
     vllm_config = _make_vllm_config()
     kv_cache_config = _make_kv_cache_config()

tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py

@@ -10,6 +10,7 @@ from unittest.mock import MagicMock, patch
 
 import torch
 
+from tests.v1.attention.utils import dense_kv_cache_views
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
     worker as mooncake_store_worker,
 )
@@ -35,6 +36,7 @@ from vllm.v1.kv_cache_interface import (
     MambaSpec,
     SlidingWindowSpec,
 )
+from vllm.v1.kv_cache_layout import KVCacheLayout
 
 
 class _DictStore:
@@ -160,8 +162,18 @@ def test_e2e_swa_plus_full_save_then_lookup_hits():
     cfg = KVCacheConfig(
         num_blocks=4,
         kv_cache_tensors=[
-            KVCacheTensor(size=8192, shared_by=["L0"]),
-            KVCacheTensor(size=8192, shared_by=["L1"]),
+            KVCacheTensor(
+                size=4 * full.page_size_bytes,
+                layers=["L0"],
+                layer_stride=4 * full.page_size_bytes,
+                block_stride=full.page_size_bytes,
+            ),
+            KVCacheTensor(
+                size=4 * swa.page_size_bytes,
+                layers=["L1"],
+                layer_stride=4 * swa.page_size_bytes,
+                block_stride=swa.page_size_bytes,
+            ),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(["L0"], full),
@@ -178,9 +190,11 @@ def test_e2e_swa_plus_full_save_then_lookup_hits():
 
     # Register kv_caches using mocked thread classes so register_kv_caches
     # doesn't try to start real background threads (which set ready_event).
+    raw_full = torch.zeros(4 * full.page_size_bytes, dtype=torch.int8)
+    raw_swa = torch.zeros(4 * swa.page_size_bytes, dtype=torch.int8)
     kv_caches = {
-        "L0": torch.zeros(2, 4, 8, 8, 64),
-        "L1": torch.zeros(2, 4, 8, 8, 64),
+        "L0": dense_kv_cache_views(raw_full, full, 4, 1, KVCacheLayout.LBNHC)[0],
+        "L1": dense_kv_cache_views(raw_swa, swa, 4, 1, KVCacheLayout.LBNHC)[0],
     }
 
     def _fake_thread_init(*args, **kwargs):

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

@@ -16,6 +16,7 @@ from unittest.mock import MagicMock, patch
 import pytest
 import torch
 
+from tests.v1.attention.utils import dense_kv_cache_views
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import (
     rdma_utils,
 )
@@ -38,6 +39,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import
 )
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheGroupSpec,
+    KVCacheLayout,
+)
 
 
 class _RecordingBlockHashes:
@@ -2961,112 +2967,83 @@ def test_putting_consumer_queues_decode_save():
     assert req.current_event is event
 
 
-def test_register_kv_caches_blocks_first_single_segment():
-    """Blocks-first layout (FlashInfer/MLA): one segment per layer."""
+@pytest.mark.parametrize("layout", list(KVCacheLayout))
+def test_register_kv_caches_shared_storage(layout: KVCacheLayout):
     num_blocks = 10
-    page_size_elements = 64
+    num_layers = 2
     worker = _make_bare_worker(num_gpu_blocks=num_blocks)
-
-    # Shape: (num_blocks, page_size_elements) — blocks outermost, no outer_dims
-    tensor = torch.zeros(num_blocks, page_size_elements, dtype=torch.float16)
-    _register_with_mocked_threads(worker, {"layer0": tensor})
-
-    db = worker.token_dbs[0]
-    assert db.kv_caches_base_addr == [tensor.untyped_storage().data_ptr()]
-    assert db.block_len == [tensor.untyped_storage().nbytes() // num_blocks]
-    worker.store.register_buffer.assert_called_once_with(
-        tensor.untyped_storage().data_ptr(),
-        tensor.untyped_storage().nbytes(),
-    )
-
-
-def test_register_kv_caches_kv_first_two_segments():
-    """K/V-first layout (FlashAttn): two segments (K, V) per layer."""
-    num_blocks = 10
-    block_size_tokens = 16
-    num_kv_heads = 4
-    head_size = 8
-    worker = _make_bare_worker(num_gpu_blocks=num_blocks)
-
-    # Shape: (2, num_blocks, block_size, num_kv_heads, head_size) — K/V outermost
-    tensor = torch.zeros(
-        2,
-        num_blocks,
-        block_size_tokens,
-        num_kv_heads,
-        head_size,
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=2,
+        head_size=8,
         dtype=torch.float16,
     )
-    _register_with_mocked_threads(worker, {"layer0": tensor})
-
-    db = worker.token_dbs[0]
-    seg_stride = tensor.stride(0) * tensor.element_size()
-    base = tensor.untyped_storage().data_ptr()
-    assert db.kv_caches_base_addr == [base, base + seg_stride]
-    assert db.block_len == [seg_stride // num_blocks] * 2
-
-
-def test_register_kv_caches_cross_layer_single_segment():
-    """Cross-layer tensor: single segment with block_len = page_size * num_layers."""
-    num_blocks = 10
-    num_layers = 4
-    per_layer_page_elements = 64  # elements per layer per block
-
-    worker = _make_bare_worker(num_gpu_blocks=num_blocks)
-
-    # Cross-layer blocks-first tensor: all layers packed into a single
-    # contiguous block.  Shape (num_blocks, num_layers * per_layer_page)
-    # mimics the physical layout after stride reordering.
-    total_page_elements = num_layers * per_layer_page_elements
-    tensor = torch.zeros(num_blocks, total_page_elements, dtype=torch.float16)
-
-    with (
-        patch(
-            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
-            "worker.KVCacheStoreSendingThread",
-            side_effect=_auto_set_ready_event,
-        ),
-        patch(
-            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
-            "worker.KVCacheStoreRecvingThread",
-            side_effect=_auto_set_ready_event,
-        ),
-    ):
-        # Use the cross-layer wrapper key, same as register_cross_layers_kv_caches
-        worker.register_kv_caches({"__cross_layer__": tensor})
-
-    db = worker.token_dbs[0]
-    assert len(db.kv_caches_base_addr) == 1
-    assert db.kv_caches_base_addr[0] == tensor.untyped_storage().data_ptr()
-
-    expected_block_len = tensor.untyped_storage().nbytes() // num_blocks
-    # block_len should be per_layer_page_size * num_layers
-    assert (
-        expected_block_len
-        == num_layers * per_layer_page_elements * tensor.element_size()
+    raw = torch.zeros(
+        num_blocks * num_layers * spec.page_size_bytes,
+        dtype=torch.int8,
     )
-    assert len(db.block_len) == 1
-    assert db.block_len[0] == expected_block_len
+    caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
 
-    # Also verify via register_cross_layers_kv_caches wrapper
-    worker2 = _make_bare_worker(num_gpu_blocks=num_blocks)
-    with (
-        patch(
-            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
-            "worker.KVCacheStoreSendingThread",
-            side_effect=_auto_set_ready_event,
-        ),
-        patch(
-            "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
-            "worker.KVCacheStoreRecvingThread",
-            side_effect=_auto_set_ready_event,
-        ),
-    ):
-        worker2.register_cross_layers_kv_caches(tensor)
+    _register_with_mocked_threads(
+        worker,
+        {"layer0": caches[0], "__cross_layer__": caches[1]},
+    )
 
-    db2 = worker2.token_dbs[0]
-    assert db2.kv_caches_base_addr == db.kv_caches_base_addr
-    assert db2.block_len == db.block_len
+    db = worker.token_dbs[0]
+    if not layout.is_block_compact:
+        # Each head group is its own region, so a block spans L*H regions.
+        head_cache = caches[0][:, 0]
+        assert db.kv_caches_base_addr == [
+            cache[:, head_idx].data_ptr()
+            for cache in caches
+            for head_idx in range(cache.shape[1])
+        ]
+        assert db.block_len == [head_cache.stride(0) * head_cache.element_size()] * (
+            num_layers * caches[0].shape[1]
+        )
+    elif layout.is_layer_compact:
+        assert db.kv_caches_base_addr == [cache.data_ptr() for cache in caches]
+        assert db.block_len == [spec.page_size_bytes] * num_layers
+    else:
+        assert db.kv_caches_base_addr == [raw.data_ptr()]
+        assert db.block_len == [num_layers * spec.page_size_bytes]
+    worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes)
+
+
+def test_register_kv_caches_separate_head_groups():
+    # LHBNC gives each head group its own region; the K/V split doubles heads.
+    layout = KVCacheLayout.LHBNC
+    num_blocks = 3
+    num_layers = 2
+    worker = _make_bare_worker(num_gpu_blocks=num_blocks)
+    spec = FullAttentionSpec(
+        block_size=4,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.float16,
+        num_head_slots=2,
+        state_content_bytes=2 * 8 * 2,
+    )
+    layer_names = ["layer0", "__cross_layer__"]
+    worker._kv_cache_groups = [KVCacheGroupSpec(layer_names, spec)]
+    raw = torch.zeros(
+        num_blocks * num_layers * spec.page_size_bytes,
+        dtype=torch.int8,
+    )
+    caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
+
+    _register_with_mocked_threads(worker, dict(zip(layer_names, caches)))
+
+    head_block_bytes = caches[0].stride(0) * caches[0].element_size()
+    expected_addrs = [
+        cache[:, head_idx].data_ptr()
+        for cache in caches
+        for head_idx in range(cache.shape[1])
+    ]
+    db = worker.token_dbs[0]
+    assert db.kv_caches_base_addr == expected_addrs
+    assert db.block_len == [head_block_bytes] * len(expected_addrs)
+    worker.store.register_buffer.assert_called_once_with(raw.data_ptr(), raw.nbytes)
 
 
 # ---------------------------------------------------------------------------

tests/v1/kv_connector/unit/test_moriio_connector.py

@@ -43,6 +43,7 @@ from vllm.v1.kv_cache_interface import (
     KVCacheConfig,
     KVCacheGroupSpec,
     KVCacheTensor,
+    compute_layer_kv_cache_shape_bytes,
 )
 
 from .utils import create_request, create_scheduler
@@ -56,20 +57,25 @@ def _find_free_port() -> int:
 
 def _make_test_kv_cache_config() -> KVCacheConfig:
     layer_names = ["layer0", "layer1", "layer2"]
+    num_blocks = 2
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=4,
+        head_size=64,
+        dtype=torch.float16,
+    )
+    page = spec.page_size_bytes
     return KVCacheConfig(
-        num_blocks=2,
-        kv_cache_tensors=[KVCacheTensor(size=0, shared_by=layer_names)],
-        kv_cache_groups=[
-            KVCacheGroupSpec(
-                layer_names=layer_names,
-                kv_cache_spec=FullAttentionSpec(
-                    block_size=16,
-                    num_kv_heads=4,
-                    head_size=64,
-                    dtype=torch.float16,
-                ),
+        num_blocks=num_blocks,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=len(layer_names) * num_blocks * page,
+                layers=layer_names,
+                layer_stride=num_blocks * page,
+                block_stride=page,
             )
         ],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec)],
     )
 
 
@@ -225,7 +231,7 @@ class FakeMoRIIOConnectorWorker(MoRIIOConnectorWorker):
         engine_id,
         *args,
         hand_shake_latency: float = 1.8,
-        kv_cache_layout="HND",
+        kv_cache_layout="LBHNC",
         kv_cache_config=None,
         **kwargs,
     ):
@@ -531,16 +537,15 @@ def test_register_kv_caches(mock_parallel_groups):
     DEFAULT_PORT = 6301
     TP_RANK = 0
     DP_RANK = 0
-    from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
-
-    backend_cls = AiterFlashAttentionBackend
-
-    # Create test kv cache tensors using proper backend shape
-    kv_cache_shape = backend_cls.get_kv_cache_shape(
-        num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
+    # Create test kv cache tensors using KVCacheSpec layout
+    shape = compute_layer_kv_cache_shape_bytes(
+        FullAttentionSpec(
+            block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16
+        ),
+        2,
     )
-    shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
-    unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
+    shared_tensor = torch.zeros(*shape, dtype=torch.int8).view(torch.float16)
+    unique_tensor = torch.zeros(*shape, dtype=torch.int8).view(torch.float16)
     kv_caches = {
         "layer0": shared_tensor,
         "layer1": unique_tensor,
@@ -627,16 +632,15 @@ def test_moriio_handshake_returns_metadata(mock_parallel_groups):
 
     ROLE = "kv_consumer"
     vllm_config = create_vllm_config(role=ROLE)
-    from vllm.v1.attention.backends.rocm_aiter_fa import AiterFlashAttentionBackend
-
-    backend_cls = AiterFlashAttentionBackend
-
-    # Create test kv cache tensors using proper backend shape
-    kv_cache_shape = backend_cls.get_kv_cache_shape(
-        num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
+    # Create test kv cache tensors using KVCacheSpec layout
+    shape = compute_layer_kv_cache_shape_bytes(
+        FullAttentionSpec(
+            block_size=16, num_kv_heads=4, head_size=64, dtype=torch.float16
+        ),
+        2,
     )
-    shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
-    unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
+    shared_tensor = torch.zeros(*shape, dtype=torch.int8).view(torch.float16)
+    unique_tensor = torch.zeros(*shape, dtype=torch.int8).view(torch.float16)
     kv_caches = {
         "layer0": shared_tensor,
         "layer1": unique_tensor,

tests/v1/kv_connector/unit/test_moriio_kv_layout.py

@@ -12,7 +12,10 @@ import pytest
 import torch
 
 from vllm.platforms import current_platform
-from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    MLAAttentionSpec,
+)
 
 aiter_available = importlib.util.find_spec("aiter") is not None
 mori_available = importlib.util.find_spec("mori") is not None
@@ -241,7 +244,7 @@ def _write_task(layer_name: str, transfer_id: str = "xfer") -> Any:
             id="interleaved-kernel-axis-from-spec",
         ),
         pytest.param(
-            (8, 4, 3),
+            (8, 1, 4, 3),
             _mla_spec(),
             16,
             {
@@ -297,7 +300,7 @@ def test_mixed_layers_compute_distinct_offsets_per_layer():
     kv_caches = {
         "separated": torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16),
         "interleaved": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16),
-        "indexer": torch.empty((8, 4, 3), dtype=torch.bfloat16),
+        "indexer": torch.empty((8, 1, 4, 3), dtype=torch.bfloat16),
     }
     worker = _worker(
         kv_caches,
@@ -342,7 +345,7 @@ def test_write_transfer_plan_caches_offsets_per_geometry():
     kv_caches = {
         "dense0": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16),
         "dense1": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16),
-        "indexer": torch.empty((8, 4, 3), dtype=torch.bfloat16),
+        "indexer": torch.empty((8, 1, 4, 3), dtype=torch.bfloat16),
     }
     calls: list[str] = []
 
@@ -692,7 +695,7 @@ def test_empty_local_block_ids_is_free_only_noop():
 def test_registration_regions_do_not_split_interleaved_or_mla_cache():
     separated = torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16)
     interleaved = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
-    indexer = torch.empty((8, 4, 3), dtype=torch.bfloat16)
+    indexer = torch.empty((8, 1, 4, 3), dtype=torch.bfloat16)
     worker = _worker(
         {
             "separated": separated,
@@ -745,7 +748,9 @@ def test_registration_regions_use_layer_num_blocks():
 
 
 def test_unsupported_shape_raises_value_error():
-    cache = torch.empty((8, 4, 2, 3), dtype=torch.bfloat16)
+    # A 4-D view whose head/state/content dims all disagree with the spec (a
+    # standardized [B, H, N, C] view for _full_spec would be (8, 2, 4, 6)).
+    cache = torch.empty((8, 3, 5, 7), dtype=torch.bfloat16)
     worker = _worker({"layer": cache}, {"layer": _full_spec()})
 
     with pytest.raises(ValueError, match="Unsupported MoRIIO K/V cache shape"):

tests/v1/kv_connector/unit/test_multi_connector.py

@@ -889,16 +889,6 @@ Options:
 """)
 
 
-def test_multi_connector_prefer_cross_layer_blocks(mc):
-    mc._connectors[0].prefer_cross_layer_blocks = False
-    mc._connectors[1].prefer_cross_layer_blocks = True
-    assert mc.prefer_cross_layer_blocks is False
-
-    mc._connectors[0].prefer_cross_layer_blocks = True
-    mc._connectors[1].prefer_cross_layer_blocks = True
-    assert mc.prefer_cross_layer_blocks is True
-
-
 def test_multi_connector_worker_metadata(mc):
     class MockConnectorWorkerMetadata(KVConnectorWorkerMetadata):
         def __init__(self, data: set[str]):

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -19,6 +19,7 @@ import pytest
 import ray
 import torch
 
+from tests.v1.attention.utils import dense_kv_cache_views
 from vllm import LLM
 from vllm.config import KVTransferConfig, set_current_vllm_config
 from vllm.distributed.kv_transfer.kv_connector.utils import (
@@ -54,8 +55,6 @@ from vllm.outputs import RequestOutput
 from vllm.platforms import current_platform
 from vllm.platforms.interface import Platform
 from vllm.sampling_params import SamplingParams
-from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
-from vllm.v1.attention.backends.utils import set_kv_cache_layout
 from vllm.v1.engine import EngineCoreRequest
 from vllm.v1.engine.output_processor import OutputProcessor
 from vllm.v1.kv_cache_interface import (
@@ -63,12 +62,12 @@ from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
-    KVCacheTensor,
+    KVCacheLayout,
+    MLAAttentionSpec,
+    compute_layer_kv_cache_shape_bytes,
 )
 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
 from vllm.v1.request import RequestStatus
-from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
-from vllm.v1.worker.utils import AttentionGroup
 
 from .utils import (
     create_request,
@@ -96,9 +95,6 @@ def clear_kv_transfer():
     yield
     if has_kv_transfer_group():
         ensure_kv_transfer_shutdown()
-    # Reset any KV cache layout override set during tests so it doesn't
-    # leak into tests in other modules.
-    set_kv_cache_layout(None)
 
 
 def get_default_xfer_telemetry(
@@ -350,12 +346,12 @@ def test_abort_immediately_remote_prefill_enqueues_empty_recv():
 )
 def test_kv_transfer_handshake(dist_init):
     """Unit test for basic NixlConnector interface functionality."""
-    from vllm.config import set_current_vllm_config
 
     # Test setup, we creates a scheduler that contains a NixlConnector
     # of role SCHEDULER, and expect it to be serving NixlAgentMetadata from
     # all workers of the instance.
     vllm_config = create_vllm_config()
+    vllm_config.cache_config.kv_cache_layout = "BLHNC"
     # in case the test runs on non-GPU machine
     vllm_config.kv_transfer_config.kv_buffer_device = "cpu"
     scheduler = create_scheduler(vllm_config)
@@ -386,18 +382,19 @@ def test_kv_transfer_handshake(dist_init):
         kv_cache_spec = cast(
             AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec
         )
-        kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(
-            num_blocks=kv_cache_config.num_blocks,
-            block_size=kv_cache_spec.block_size,
-            num_kv_heads=kv_cache_spec.num_kv_heads,
-            head_size=kv_cache_spec.head_size,
+        raw = torch.zeros(
+            kv_cache_spec.page_size_bytes * kv_cache_config.num_blocks * 3,
+            dtype=torch.int8,
+        )
+        caches = dense_kv_cache_views(
+            raw,
+            kv_cache_spec,
+            kv_cache_config.num_blocks,
+            num_layers=3,
+            layout=KVCacheLayout.BLHNC,
         )
-        shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
-        unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
         kv_caches = {
-            "layer0": shared_tensor,
-            "layer1": unique_tensor,
-            "layer2": shared_tensor,
+            f"layer{layer_idx}": cache for layer_idx, cache in enumerate(caches)
         }
         prefill_connector.register_kv_caches(kv_caches)
 
@@ -474,7 +471,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
         self,
         *args,
         hand_shake_latency: float = 1.8,
-        kv_cache_layout="HND",
+        kv_cache_layout="LBHNC",
         kv_cache_config=None,
         **kwargs,
     ):
@@ -486,9 +483,8 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
         # Mock register_kv_caches attributes needed for tests that do not call it.
         self.src_xfer_handles_by_block_size = {self.block_size: 1}
         self.src_blocks_data = np.empty((0, 3), dtype=np.uint64)
-        test_shape = self.attn_backends[0].get_kv_cache_shape(
-            num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
-        )
+        rep_spec = self.kv_cache_config.kv_cache_groups[0].kv_cache_spec
+        test_shape = compute_layer_kv_cache_shape_bytes(rep_spec, 1)
         self.transfer_topo = TransferTopology(
             tp_rank=self.tp_rank,
             tp_size=self.world_size,
@@ -502,7 +498,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
         )
 
         self.compat_hash = compute_nixl_compatibility_hash(
-            self.vllm_config, self.backend_name, self.transfer_topo.cross_layers_blocks
+            self.vllm_config, self.backend_name
         )
 
     def _nixl_handshake(
@@ -553,9 +549,8 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
                     device_id=remote_tp_rank,
                     num_blocks=1,
                     block_lens=remote_block_lens,
-                    # `self.kv_cache_layout` is only forced to HND when vllm engine
-                    # is started. We mock HND here.
-                    kv_cache_layout="HND",
+                    block_strides=remote_block_lens,
+                    kv_cache_layout="LBHNC",
                     block_size=self.block_size,
                     ssm_sizes=(0, 0),
                     attn_backend_name=self.backend_name,
@@ -579,6 +574,8 @@ class TestNixlHandshake:
         self, default_vllm_config, dist_init, pcp_rank
     ):
         """Only PCP rank zero publishes and reports sending completion."""
+        from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
+
         vllm_config = create_vllm_config(kv_role="kv_producer")
         vllm_config.parallel_config.prefill_context_parallel_size = 2
         with (
@@ -656,7 +653,7 @@ class TestNixlHandshake:
         worker.dst_xfer_side_handles = {
             FakeNixlConnectorWorker.REMOTE_ENGINE_ID: {0: 1}
         }
-        worker.kv_cache_layout = "HND"
+        worker.kv_cache_layout = "LBHNC"
         num_xfers = 4
         while True:
             # For the same request_id, initiate multiple xfers across different
@@ -1053,7 +1050,9 @@ class TestNixlHandshake:
             worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
 
             # Metadata with different kv_cache_layout than local worker
-            mismatched_layout = "HND" if worker.kv_cache_layout != "HND" else "NHD"
+            mismatched_layout = (
+                "LBHNC" if worker.kv_cache_layout != "LBHNC" else "LBNHC"
+            )
             meta = NixlAgentMetadata(
                 engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
                 agent_metadata=FakeNixlWrapper.AGENT_METADATA,
@@ -1061,6 +1060,7 @@ class TestNixlHandshake:
                 device_id=0,
                 num_blocks=1,
                 block_lens=worker.block_len_per_layer,
+                block_strides=worker.block_len_per_layer,
                 kv_cache_layout=mismatched_layout,
                 block_size=worker.block_size,
                 ssm_sizes=(0, 0),
@@ -1100,7 +1100,7 @@ class TestNixlHandshake:
                 vllm_config,
                 connector.engine_id,
                 hand_shake_latency=0,
-                kv_cache_layout="NHD",
+                kv_cache_layout="LBNHC",
             )
             worker = connector.connector_worker
 
@@ -1111,15 +1111,17 @@ class TestNixlHandshake:
             worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
 
             # Metadata with different kv_cache_layout than local worker
+            # prefill TP=1, decode TP=2, remote block_lens is double to local
+            remote_block_lens = [i * 2 for i in worker.block_len_per_layer]
             meta = NixlAgentMetadata(
                 engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
                 agent_metadata=FakeNixlWrapper.AGENT_METADATA,
                 kv_caches_base_addr=[0],
                 device_id=0,
                 num_blocks=1,
-                # prefill TP=1, decode TP=2, remote block_lens is double to local
-                block_lens=[i * 2 for i in worker.block_len_per_layer],
-                kv_cache_layout="HND",
+                block_lens=remote_block_lens,
+                block_strides=remote_block_lens,
+                kv_cache_layout="LBHNC",
                 block_size=worker.block_size,
                 ssm_sizes=(0, 0),
                 attn_backend_name=worker.backend_name,
@@ -1160,8 +1162,6 @@ class TestNixlHandshake:
             attn_backends=worker.attn_backends,
             tensor_shape=None,
         )
-        assert worker.transfer_topo.virtually_split_kv_in_blocks
-
         plan = MagicMock(
             source_ranks_per_group=((0,), (0,)),
             rank_offset_factor=1,
@@ -1171,6 +1171,7 @@ class TestNixlHandshake:
             device_id=0,
             num_blocks=1,
             block_lens=[remote_block_len],
+            block_strides=[remote_block_len],
         )
 
         assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [
@@ -1226,6 +1227,7 @@ class TestNixlHandshake:
                 device_id=0,
                 num_blocks=1,
                 block_lens=[fa_len * tp_ratio, idx_len],
+                block_strides=[fa_len * tp_ratio, idx_len],
                 kv_cache_layout=worker.kv_cache_layout,
                 block_size=worker.block_size,
                 ssm_sizes=(0, 0),
@@ -1252,6 +1254,7 @@ class TestNixlHandshake:
                 num_blocks=1,
                 # WRONG: MLA region scaled by tp_ratio (it should be replicated).
                 block_lens=[fa_len * tp_ratio, idx_len * tp_ratio],
+                block_strides=[fa_len * tp_ratio, idx_len * tp_ratio],
                 kv_cache_layout=worker2.kv_cache_layout,
                 block_size=worker2.block_size,
                 ssm_sizes=(0, 0),
@@ -1295,7 +1298,7 @@ class TestNixlHandshake:
 
             worker.transfer_topo.total_num_kv_heads = 8
             worker.transfer_topo.local_physical_heads = 1
-            worker.kv_cache_layout = "HND"
+            worker.kv_cache_layout = "LBHNC"
 
             worker.slot_size_per_layer = [4096]
             worker.block_len_per_layer = [4096 * worker.block_size]
@@ -1311,7 +1314,8 @@ class TestNixlHandshake:
                 device_id=0,
                 num_blocks=1,
                 block_lens=list(worker.block_len_per_layer),
-                kv_cache_layout="HND",
+                block_strides=list(worker.block_len_per_layer),
+                kv_cache_layout="LBHNC",
                 block_size=worker.block_size,
                 ssm_sizes=(0, 0),
                 attn_backend_name=worker.backend_name,
@@ -1350,7 +1354,7 @@ class TestNixlHandshake:
 
             worker.transfer_topo.total_num_kv_heads = 32
             worker.transfer_topo.local_physical_heads = 8  # 32 // 4
-            worker.kv_cache_layout = "HND"
+            worker.kv_cache_layout = "LBHNC"
 
             slot_size = 4096
             worker.slot_size_per_layer = [slot_size]
@@ -1368,7 +1372,8 @@ class TestNixlHandshake:
                 device_id=0,
                 num_blocks=1,
                 block_lens=list(worker.block_len_per_layer),
-                kv_cache_layout="HND",
+                block_strides=list(worker.block_len_per_layer),
+                kv_cache_layout="LBHNC",
                 block_size=worker.block_size,
                 ssm_sizes=(0, 0),
                 attn_backend_name=worker.backend_name,
@@ -1839,7 +1844,19 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
     llm.llm_engine.engine_core.shutdown()
 
 
-@pytest.mark.parametrize("enable_cross_layers", ["False", "True"])
+def element_byte_addrs(view: torch.Tensor) -> list[int]:
+    """Absolute byte addresses of every element of a (possibly strided) view."""
+    offsets = torch.zeros(view.shape, dtype=torch.int64)
+    for dim, (size, stride) in enumerate(zip(view.shape, view.stride())):
+        shape = [1] * view.ndim
+        shape[dim] = size
+        offsets = offsets + torch.arange(size, dtype=torch.int64).view(shape) * stride
+    esize = view.element_size()
+    byte_offsets = (offsets.flatten() * esize).tolist()
+    base = view.data_ptr()
+    return [base + offset + i for offset in byte_offsets for i in range(esize)]
+
+
 @pytest.mark.parametrize(
     "attn_backend",
     [
@@ -1853,8 +1870,14 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
         "TRITON_ATTN",
     ],
 )
+@pytest.mark.parametrize("layout", [layout.name for layout in KVCacheLayout])
+@pytest.mark.parametrize("separate_kv_head_groups", [False, True])
 def test_register_kv_caches(
-    default_vllm_config, dist_init, attn_backend, enable_cross_layers
+    default_vllm_config,
+    dist_init,
+    attn_backend,
+    layout,
+    separate_kv_head_groups,
 ):
     """
     Test that register_kv_caches() properly calls nixl_wrapper methods with
@@ -1868,12 +1891,7 @@ def test_register_kv_caches(
     """
 
     vllm_config = create_vllm_config(attention_backend=attn_backend)
-
-    # Enable cross layers blocks
-    vllm_config.kv_transfer_config.kv_connector_extra_config[
-        "enable_cross_layers_blocks"
-    ] = enable_cross_layers
-    set_kv_cache_layout("HND")
+    vllm_config.cache_config.kv_cache_layout = layout
 
     # Import the appropriate backend based on the parameter
     if attn_backend == "FLASH_ATTN":
@@ -1890,59 +1908,37 @@ def test_register_kv_caches(
         backend_cls = TritonAttentionBackend
 
     nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker"
-    nixl_connector = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector"
     with (
         patch(f"{nixl_worker}.NixlWrapper") as mock_nixl_wrapper,
         patch(f"{nixl_worker}.threading.Event"),
         patch(f"{nixl_worker}.threading.Thread") as mock_thread,
-        patch(f"{nixl_connector}.get_current_attn_backend") as mock_get_attn_backend,
         patch(f"{nixl_worker}.get_current_attn_backends") as mock_get_attn_backends,
     ):
-        # Ensure get_attn_backend returns the correct value due to
-        # _cached_get_attn_backend returning the backend from previous
-        # test run if not mocking.
-        mock_get_attn_backend.return_value = backend_cls
         mock_get_attn_backends.return_value = [backend_cls]
-        num_layers = 32
         block_size = 16
         num_blocks = 8
         num_heads = 4
         head_size = 16
 
-        # TODO (NickLucche) the fact that connector depends on kv_cache_config for init
-        # but cross-layer preference cant be inferred prior to creating kv_cache_config
-        # is a bit awkward.
-        dummy_connector = NixlConnector(
-            vllm_config,
-            KVConnectorRole.WORKER,
-            make_kv_cache_config(block_size=block_size),
-        )
         kv_cache_spec = FullAttentionSpec(
             block_size=block_size,
             num_kv_heads=num_heads,
             head_size=head_size,
             dtype=torch.float16,
+            num_head_slots=2 if separate_kv_head_groups else None,
+            state_content_bytes=num_heads * head_size * 2
+            if separate_kv_head_groups
+            else None,
+        )
+        kv_cache_config = KVCacheConfig(
+            num_blocks=num_blocks,
+            kv_cache_tensors=[],
+            kv_cache_groups=[
+                KVCacheGroupSpec(
+                    ["layer0", "layer1", "layer2", "layer3"], kv_cache_spec
+                )
+            ],
         )
-        if dummy_connector.prefer_cross_layer_blocks:
-            kv_cache_config = KVCacheConfig(
-                num_blocks=num_blocks,
-                kv_cache_tensors=[
-                    KVCacheTensor(
-                        size=kv_cache_spec.page_size_bytes * num_blocks,
-                        shared_by=["all-layers"],
-                    )
-                    for _ in range(num_layers)
-                ],
-                kv_cache_groups=[KVCacheGroupSpec(["all-layers"], kv_cache_spec)],
-            )
-        else:
-            kv_cache_config = KVCacheConfig(
-                num_blocks=num_blocks,
-                kv_cache_tensors=[],
-                kv_cache_groups=[
-                    KVCacheGroupSpec(["layer0", "layer1", "layer2"], kv_cache_spec)
-                ],
-            )
         # Create connector
         connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
         connector.connector_worker = FakeNixlConnectorWorker(
@@ -1962,93 +1958,36 @@ def test_register_kv_caches(
         # Reassure the shutdown() check that the thread is terminated
         mock_thread.return_value.is_alive.return_value = False
 
-        expected_tensor_size: int
-        expected_base_addrs: list[int]
-        expected_num_entries: int
-        kv_caches: dict[str, torch.Tensor]
-        if str(enable_cross_layers).lower() == "true":
-            assert connector.prefer_cross_layer_blocks == (
-                attn_backend in ("FLASH_ATTN", "FLASHINFER", "TRITON_ATTN")
-            )
-        else:
-            assert not connector.prefer_cross_layer_blocks
-
-        test_shape = backend_cls.get_kv_cache_shape(
-            num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
+        raw0 = torch.zeros(
+            kv_cache_spec.page_size_bytes * kv_cache_config.num_blocks * 2,
+            dtype=torch.int8,
+            device=current_platform.device_type,
         )
-        is_blocks_first = len(test_shape) == 4 and test_shape[0] == 1
-
-        if connector.prefer_cross_layer_blocks:
-            with set_current_vllm_config(vllm_config):
-                _, cross_layers_kv_cache, _ = (
-                    KVConnectorModelRunnerMixin.allocate_uniform_kv_caches(
-                        kv_cache_config=kv_cache_config,
-                        attn_groups=[
-                            [
-                                AttentionGroup(
-                                    backend=backend_cls,
-                                    layer_names=[],
-                                    kv_cache_spec=kv_cache_spec,
-                                    kv_cache_group_id=0,
-                                )
-                            ]
-                        ],
-                        cache_dtype="bfloat16",
-                        device=torch.accelerator.current_device_index(),
-                        kernel_block_sizes=[block_size],
-                    )
-                )
-            # Store tensor info for validation
-            expected_tensor_size = (
-                cross_layers_kv_cache.element_size() * cross_layers_kv_cache.numel()
-            )
-            expected_base_addrs = [
-                cross_layers_kv_cache.data_ptr(),
-            ]
-            expected_num_entries = 1
-
-            expected_blocks_count = num_blocks
-
-            kv_caches = {"all-layers": cross_layers_kv_cache}
-        else:
-            # Create test kv cache tensors using proper backend shape
-            kv_cache_shape = backend_cls.get_kv_cache_shape(
-                num_blocks=kv_cache_config.num_blocks,
-                block_size=kv_cache_spec.block_size,
-                num_kv_heads=kv_cache_spec.num_kv_heads,
-                head_size=kv_cache_spec.head_size,
-            )
-            shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
-            unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
-            kv_caches = {
-                "layer0": shared_tensor,
-                "layer1": unique_tensor,
-                "layer2": shared_tensor,
-            }
-
-            # Store tensor info for validation
-            if is_blocks_first:
-                expected_tensor_size = (
-                    shared_tensor.element_size() * shared_tensor.numel()
-                )
-                expected_base_addrs = [
-                    shared_tensor.data_ptr(),
-                    unique_tensor.data_ptr(),
-                ]
-                expected_num_entries = 2
-                expected_blocks_count = kv_cache_config.num_blocks * 2
-            else:
-                expected_tensor_size = (
-                    shared_tensor[0].element_size() * shared_tensor[0].numel()
-                )
-                expected_base_addrs = [
-                    shared_tensor[0].data_ptr(),
-                    shared_tensor[1].data_ptr(),
-                    unique_tensor[0].data_ptr(),
-                    unique_tensor[1].data_ptr(),
-                ]
-                expected_num_entries = 4
-                expected_blocks_count = kv_cache_config.num_blocks * 4
+        raw1 = torch.zeros(
+            kv_cache_spec.page_size_bytes * kv_cache_config.num_blocks,
+            dtype=torch.int8,
+            device=current_platform.device_type,
+        )
+        tensor0, tensor1 = dense_kv_cache_views(
+            raw0,
+            kv_cache_spec,
+            kv_cache_config.num_blocks,
+            num_layers=2,
+            layout=KVCacheLayout[layout],
+        )
+        (tensor2,) = dense_kv_cache_views(
+            raw1,
+            kv_cache_spec,
+            kv_cache_config.num_blocks,
+            num_layers=1,
+            layout=KVCacheLayout[layout],
+        )
+        kv_caches = {
+            "layer0": tensor0,
+            "layer1": tensor1,
+            "layer2": tensor2,
+            "layer3": tensor0,
+        }
 
         # Execute register_kv_caches
         connector.register_kv_caches(kv_caches)
@@ -2056,44 +1995,123 @@ def test_register_kv_caches(
         # Verify get_reg_descs was called with caches_data
         assert mock_wrapper_instance.get_reg_descs.called
         caches_data, _ = mock_wrapper_instance.get_reg_descs.call_args[0]
-        assert len(caches_data) == expected_num_entries
+        assert len(caches_data) == 2
 
-        for i, cache_entry in enumerate(caches_data):
+        for cache_entry, raw in zip(caches_data, (raw0, raw1)):
             base_addr, size, _tp_rank, _ = cache_entry
-            assert size == expected_tensor_size, (
-                f"Entry {i}: Expected tensor size {expected_tensor_size}, got {size}"
-            )
-            assert base_addr == expected_base_addrs[i], (
-                f"Entry {i}: Expected base address {expected_base_addrs[i]}, "
-                f"got {base_addr}"
-            )
+            assert size == raw.nbytes
+            assert base_addr == raw.data_ptr()
 
         # Verify get_xfer_descs was called with blocks_data
         assert mock_wrapper_instance.get_xfer_descs.called
         blocks_data, _ = mock_wrapper_instance.get_xfer_descs.call_args[0]
 
-        # Validate blocks_data structure and size
-        assert len(blocks_data) == expected_blocks_count, (
-            f"Expected {expected_blocks_count} blocks, got {len(blocks_data)}"
-        )
+        # Layout-blind contract: whatever regions the worker carves out,
+        # transferring "block b" must move exactly logical block b's bytes for
+        # every layer. Map each registered byte to the block that owns it, then
+        # require every descriptor window to hold bytes of a single block and
+        # the windows of block b to cover exactly block b's bytes.
+        owner: dict[int, int] = {}
+        for cache in (tensor0, tensor1, tensor2):
+            for blk in range(num_blocks):
+                for addr in element_byte_addrs(cache[blk]):
+                    assert owner.setdefault(addr, blk) == blk
+        block_bytes = defaultdict(set)
+        for addr, blk in owner.items():
+            block_bytes[blk].add(addr)
 
-        if connector.prefer_cross_layer_blocks:
-            num_blocks = 8
-        else:
-            num_blocks = kv_cache_config.num_blocks
+        covered: defaultdict[int, set[int]] = defaultdict(set)
+        for block_start_addr, block_len, _tp_rank in blocks_data:
+            window = range(block_start_addr, block_start_addr + block_len)
+            owners = {owner[addr] for addr in window}
+            assert len(owners) == 1, "descriptor window spans logical blocks"
+            covered[owners.pop()].update(window)
+        assert covered == block_bytes
 
-        expected_block_len = expected_tensor_size // num_blocks
-
-        for i, block_entry in enumerate(blocks_data):
-            block_start_addr, block_len, tp_rank = block_entry
-            assert block_len == expected_block_len, (
-                f"Block entry {i}: Expected block len {expected_block_len}, "
-                f"got {block_len}"
-            )
+        # Region bases are exactly the block-0 window starts, one per region.
+        base_addrs = connector.connector_worker.kv_caches_base_addr[
+            connector.connector_worker.engine_id
+        ][0]
+        if layout == "BLHNC":
+            assert len(base_addrs) == 3
+        assert set(base_addrs) == {
+            start for start, _len, _tp in blocks_data if owner[start] == 0
+        }
+        assert len(blocks_data) == num_blocks * len(base_addrs)
 
         assert connector.connector_worker.block_size == 16
 
 
+def test_register_packed_dsv4_mla_cache_as_single_region(
+    default_vllm_config, dist_init
+):
+    from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
+
+    nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker"
+    with (
+        patch(f"{nixl_worker}.NixlWrapper") as mock_nixl_wrapper,
+        patch(f"{nixl_worker}.threading.Event"),
+        patch(f"{nixl_worker}.threading.Thread") as mock_thread,
+        patch(f"{nixl_worker}.get_current_attn_backends") as mock_backends,
+    ):
+        mock_backends.return_value = [TritonAttentionBackend]
+        num_blocks = 2
+        num_layers = 4
+        block_size = 256
+        spec = MLAAttentionSpec(
+            block_size=block_size,
+            num_kv_heads=1,
+            head_size=512,
+            dtype=torch.uint8,
+            cache_dtype_str="fp8_ds_mla",
+            tokens_per_state=4,
+            alignment=576,
+            model_version="deepseek_v4",
+            state_content_bytes=584,
+        )
+        layer_names = [f"layer{idx}" for idx in range(num_layers)]
+        kv_cache_config = KVCacheConfig(
+            num_blocks=num_blocks,
+            kv_cache_tensors=[],
+            kv_cache_groups=[KVCacheGroupSpec(layer_names, spec)],
+        )
+        vllm_config = create_vllm_config(attention_backend="TRITON_ATTN")
+        vllm_config.cache_config.block_size = block_size
+        vllm_config.cache_config.kv_cache_layout = "BLHNC"
+        connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
+        connector.connector_worker = FakeNixlConnectorWorker(
+            vllm_config,
+            connector.engine_id,
+            hand_shake_latency=0,
+            kv_cache_layout="BLHNC",
+            kv_cache_config=kv_cache_config,
+        )
+        wrapper = mock_nixl_wrapper.return_value
+        connector.connector_worker.nixl_wrapper = wrapper
+        wrapper.get_agent_metadata.return_value = b"fake_agent_metadata"
+        mock_thread.return_value.is_alive.return_value = False
+
+        raw = torch.zeros(
+            spec.page_size_bytes * num_blocks * num_layers,
+            dtype=torch.int8,
+            device=current_platform.device_type,
+        )
+        views = dense_kv_cache_views(
+            raw, spec, num_blocks, num_layers, KVCacheLayout.BLHNC
+        )
+        connector.register_kv_caches(dict(zip(layer_names, views)))
+
+        blocks_data, _ = wrapper.get_xfer_descs.call_args[0]
+        packed_block_len = num_layers * spec.page_size_bytes
+        assert connector.connector_worker.kv_caches_base_addr[
+            connector.connector_worker.engine_id
+        ][0] == [raw.data_ptr()]
+        assert blocks_data.tolist() == [
+            [raw.data_ptr() + block_idx * packed_block_len, packed_block_len, 0]
+            for block_idx in range(num_blocks)
+        ]
+
+
 class FakePlatform(Platform):
     device_type: str = "oot"
 
@@ -2318,6 +2336,8 @@ def test_engine_ttl_disabled(default_vllm_config, dist_init):
 
 def test_transfer_topology_unregister():
     """TransferTopology.unregister_remote_engine removes the engine."""
+    from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
+
     topo = TransferTopology(
         tp_rank=0,
         tp_size=1,
@@ -2908,8 +2928,8 @@ def test_speculative_config_compatibility_hash(
     _set_test_speculative_config(local_config)
     _set_test_speculative_config(remote_config, **remote_overrides)
 
-    local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN", False)
-    remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN", False)
+    local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN")
+    remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN")
 
     assert (local_hash == remote_hash) is should_match
 
@@ -2920,10 +2940,8 @@ def test_missing_speculative_config_changes_compatibility_hash():
     speculative_config = create_vllm_config()
     _set_test_speculative_config(speculative_config)
 
-    regular_hash = compute_nixl_compatibility_hash(regular_config, "FLASH_ATTN", False)
-    speculative_hash = compute_nixl_compatibility_hash(
-        speculative_config, "FLASH_ATTN", False
-    )
+    regular_hash = compute_nixl_compatibility_hash(regular_config, "FLASH_ATTN")
+    speculative_hash = compute_nixl_compatibility_hash(speculative_config, "FLASH_ATTN")
 
     assert regular_hash != speculative_hash
 
@@ -2938,8 +2956,8 @@ def test_speculative_kv_cache_dtype_resolves_to_target():
     _set_test_speculative_config(local_config, kv_cache_dtype="fp8")  # explicit
     _set_test_speculative_config(remote_config, kv_cache_dtype=None)  # inherits
 
-    local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN", False)
-    remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN", False)
+    local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN")
+    remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN")
 
     assert local_hash == remote_hash
 
@@ -2955,8 +2973,8 @@ def test_speculative_attention_backend_not_in_compatibility_hash():
     _set_test_speculative_config(local_config, attention_backend=None)
     _set_test_speculative_config(remote_config, attention_backend="FLASHINFER")
 
-    local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN", False)
-    remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN", False)
+    local_hash = compute_nixl_compatibility_hash(local_config, "FLASH_ATTN")
+    remote_hash = compute_nixl_compatibility_hash(remote_config, "FLASH_ATTN")
 
     assert local_hash == remote_hash
 
@@ -2969,17 +2987,17 @@ def test_transfer_mode_changes_compatibility_hash():
     config = create_vllm_config()
 
     pull_hash = compute_nixl_compatibility_hash(
-        config, "FLASH_ATTN", False, transfer_mode="pull"
+        config, "FLASH_ATTN", transfer_mode="pull"
     )
     push_hash = compute_nixl_compatibility_hash(
-        config, "FLASH_ATTN", False, transfer_mode="push"
+        config, "FLASH_ATTN", transfer_mode="push"
     )
 
     assert pull_hash != push_hash
     assert pull_hash == compute_nixl_compatibility_hash(
-        config, "FLASH_ATTN", False, transfer_mode="pull"
+        config, "FLASH_ATTN", transfer_mode="pull"
     )
-    assert compute_nixl_compatibility_hash(config, "FLASH_ATTN", False) == pull_hash
+    assert compute_nixl_compatibility_hash(config, "FLASH_ATTN") == pull_hash
 
 
 @pytest.mark.skip_global_cleanup
@@ -3057,14 +3075,11 @@ def test_compatibility_hash_validation(
     kv_cache_spec = cast(
         AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec
     )
-    kv_cache_shape = decode_worker.attn_backends[0].get_kv_cache_shape(
-        num_blocks=kv_cache_config.num_blocks,
-        block_size=kv_cache_spec.block_size,
-        num_kv_heads=kv_cache_spec.num_kv_heads,
-        head_size=kv_cache_spec.head_size,
+    shape = compute_layer_kv_cache_shape_bytes(
+        kv_cache_spec, kv_cache_config.num_blocks
     )
-    shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
-    unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
+    shared_tensor = torch.zeros(*shape, dtype=torch.int8).view(kv_cache_spec.dtype)
+    unique_tensor = torch.zeros(*shape, dtype=torch.int8).view(kv_cache_spec.dtype)
     # Build kv_caches from the actual layer names in kv_cache_config so that
     # _layer_specs lookups in register_kv_caches always find a matching key.
     layer_names = [
@@ -3099,18 +3114,19 @@ def test_compatibility_hash_validation(
         remote_hash = compute_nixl_compatibility_hash(
             remote_vllm_config,
             decode_worker.backend_name,
-            decode_worker.transfer_topo.cross_layers_blocks,
         )
 
     prefill_block_size = config_overrides.get("block_size", 16)
+    prefill_block_lens = [4096 * prefill_block_size]
     prefill_metadata = NixlAgentMetadata(
         engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
         agent_metadata=FakeNixlWrapper.AGENT_METADATA,
         kv_caches_base_addr=[0],
         device_id=0,
         num_blocks=1,
-        block_lens=[4096 * prefill_block_size],  # slot_size * block_size
-        kv_cache_layout="HND",
+        block_lens=prefill_block_lens,
+        block_strides=prefill_block_lens,
+        kv_cache_layout="LBHNC",
         block_size=prefill_block_size,
         ssm_sizes=(0, 0),
         attn_backend_name=decode_worker.backend_name,
@@ -3187,9 +3203,8 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
     decode_worker = decode_connector.connector_worker
 
     backend = get_current_attn_backend(local_vllm_config)
-    test_shape = backend.get_kv_cache_shape(
-        num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
-    )
+    probe_spec = decode_worker.kv_cache_config.kv_cache_groups[0].kv_cache_spec
+    test_shape = compute_layer_kv_cache_shape_bytes(probe_spec, 1)
     decode_worker.transfer_topo = TransferTopology(
         tp_rank=decode_worker.tp_rank,
         tp_size=decode_worker.world_size,
@@ -3205,7 +3220,6 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
     decode_worker.compat_hash = compute_nixl_compatibility_hash(
         decode_worker.vllm_config,
         decode_worker.backend_name,
-        decode_worker.transfer_topo.cross_layers_blocks,
     )
 
     if error_scenario == "handshake_decode_error":

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -1505,14 +1505,19 @@ def _make_hybrid_mla_kv_cache_config(num_blocks: int = 4):
         mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
     )
     assert kda_spec.page_size_bytes == unified_page
+    # The three groups overlay each other, so layer i of every group aliases the same
+    # region: mla.i, kda_a.i and kda_b.i all live at i * layer_stride.
+    layer_stride = num_blocks * unified_page
     return KVCacheConfig(
         num_blocks=num_blocks,
         kv_cache_tensors=[
             KVCacheTensor(
-                size=num_blocks * unified_page,
-                shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"],
+                size=2 * layer_stride,
+                layers=[f"{prefix}.0", f"{prefix}.1"],
+                layer_stride=layer_stride,
+                block_stride=unified_page,
             )
-            for i in range(2)
+            for prefix in ("mla", "kda_a", "kda_b")
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec),

tests/v1/kv_connector/unit/test_nixl_desc_geometry.py

@@ -129,14 +129,19 @@ def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blo
         page_size_padded=unified_page,
         mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
     )
+    # The three groups overlay each other, so layer i of every group aliases the same
+    # region: mla.i, kda_a.i and kda_b.i all live at i * layer_stride.
+    layer_stride = num_logical_blocks * unified_page
     kv_cache_config = KVCacheConfig(
         num_blocks=num_logical_blocks,
         kv_cache_tensors=[
             KVCacheTensor(
-                size=num_logical_blocks * unified_page,
-                shared_by=[f"mla.{i}", f"kda_a.{i}", f"kda_b.{i}"],
+                size=2 * layer_stride,
+                layers=[f"{prefix}.0", f"{prefix}.1"],
+                layer_stride=layer_stride,
+                block_stride=unified_page,
             )
-            for i in range(2)
+            for prefix in ("mla", "kda_a", "kda_b")
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(["mla.0", "mla.1"], mla_spec),
@@ -230,6 +235,8 @@ def _make_remote_meta(
         kv_caches_base_addr=[0x10_000_000, 0x20_000_000],
         num_blocks=remote_num_logical * remote_ppl,
         block_lens=[kernel_page, kernel_page],
+        # Non-interleaved remote: consecutive blocks abut, so stride == page length.
+        block_strides=[kernel_page, kernel_page],
         kv_cache_layout=worker.kv_cache_layout,
         block_size=remote_kernel_block_size,
         ssm_sizes=remote_ssm_sizes,

tests/v1/kv_connector/unit/test_nixl_push_connector.py

@@ -1081,6 +1081,8 @@ class TestPushPipelineParallel:
             device_id=0,
             num_blocks=4,
             block_lens=[block_len] * 4,
+            # Non-interleaved: consecutive blocks abut, so stride == block length.
+            block_strides=[block_len] * 4,
             kv_cache_layout="HND",
             block_size=16,
             ssm_sizes=(0, 0),

tests/v1/kv_connector/unit/test_nixl_simple_cpu_offload.py

@@ -145,10 +145,15 @@ def _make_kv_cache_config(
             ),
         )
     ]
+    # Each group packs its two layers densely; groups overlay one allocation.
+    layer_stride = _BYTES_PER_BLOCK * num_blocks
+    size = 2 * layer_stride
     tensors = [
         KVCacheTensor(
-            size=_BYTES_PER_BLOCK * num_blocks,
-            shared_by=fa_layers,
+            size=size,
+            layers=fa_layers,
+            layer_stride=layer_stride,
+            block_stride=_BYTES_PER_BLOCK,
         )
     ]
     if swa_enabled:
@@ -167,8 +172,10 @@ def _make_kv_cache_config(
         )
         tensors.append(
             KVCacheTensor(
-                size=_BYTES_PER_BLOCK * num_blocks,
-                shared_by=sw_layers,
+                size=size,
+                layers=sw_layers,
+                layer_stride=layer_stride,
+                block_stride=_BYTES_PER_BLOCK,
             )
         )
     return KVCacheConfig(

tests/v1/kv_connector/unit/test_transfer_topology_sharded.py

@@ -12,14 +12,7 @@ pytestmark = pytest.mark.cpu_test
 
 
 class _FakeAttentionBackend:
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-    ) -> tuple[int, int, int, int]:
-        return (num_blocks, num_kv_heads, block_size, 2 * head_size)
+    pass
 
 
 def _make_topology(

tests/v1/kv_connector/unit/utils.py

@@ -130,6 +130,8 @@ def create_vllm_config(
         cache_dtype=cache_dtype,
         enable_prefix_caching=True,
     )
+    # Connectors are constructed after layout resolution; mirror that here.
+    cache_config.kv_cache_layout = "LBNHC"
     kv_transfer_config = KVTransferConfig(
         kv_connector=kv_connector,
         kv_connector_module_path=kv_connector_module_path,

tests/v1/kv_offload/test_file_mapper.py

@@ -55,6 +55,7 @@ def make_mapper_from_offloading_spec(**kwargs) -> FileMapper:
         ),
         replicated_layout=kwargs.get("replicated_layout", False),
         canonical_layout=kwargs.get("canonical_layout", False),
+        kv_cache_layout=kwargs.get("kv_cache_layout", "LBNHC"),
     )
     spec = MagicMock(spec=OffloadingSpec)
     spec.config = config
@@ -214,14 +215,8 @@ def test_parallel_agnostic_separates_persistent_layouts():
 def test_canonical_layout_changes_storage_namespace():
     # Canonical bytes are not interchangeable with the direct layout, so the
     # format id must fork the storage namespace.
-    from vllm.v1.attention.backends.utils import set_kv_cache_layout
-
-    set_kv_cache_layout("NHD")
-    try:
-        direct = make_mapper_from_offloading_spec()
-        canonical = make_mapper_from_offloading_spec(canonical_layout=True)
-    finally:
-        set_kv_cache_layout(None)
+    direct = make_mapper_from_offloading_spec()
+    canonical = make_mapper_from_offloading_spec(canonical_layout=True)
     assert "canonical_format" not in direct.fields
     assert canonical.fields["canonical_format"] == "v1-nhd"
     assert direct.base_path != canonical.base_path

tests/v1/simple_kv_offload/test_scheduler.py

@@ -91,7 +91,9 @@ def _make_kv_cache_config(
         tensors.append(
             KVCacheTensor(
                 size=_BYTES_PER_BLOCK * num_blocks,
-                shared_by=layer_names,
+                layers=layer_names,
+                layer_stride=_BYTES_PER_BLOCK * num_blocks,
+                block_stride=_BYTES_PER_BLOCK,
             )
         )
     return KVCacheConfig(

tests/v1/simple_kv_offload/test_worker.py

@@ -10,6 +10,7 @@ read partially written / stale blocks and silently corrupt the CPU cache.
 from __future__ import annotations
 
 import time
+from unittest.mock import MagicMock
 
 import pytest
 import torch
@@ -19,6 +20,8 @@ from vllm.platforms import current_platform
 if not current_platform.is_cuda_alike():
     pytest.skip("Requires CUDA or ROCm", allow_module_level=True)
 
+from tests.v1.attention.utils import dense_kv_cache_views
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout
 from vllm.v1.simple_kv_offload.copy_backend import DmaCopyBackend
 from vllm.v1.simple_kv_offload.cuda_mem_ops import (
     CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
@@ -188,3 +191,90 @@ def test_build_params_src_access_order():
         gpu, cpu, stream, src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_STREAM
     )
     assert ordered.attrs.srcAccessOrder == CU_MEMCPY_SRC_ACCESS_ORDER_STREAM
+
+
+@pytest.mark.parametrize("layout", list(KVCacheLayout))
+def test_register_shared_kv_cache_storage(monkeypatch, layout: KVCacheLayout):
+    num_blocks = 4
+    num_layers = 2
+    spec = FullAttentionSpec(
+        block_size=2,
+        num_kv_heads=2,
+        head_size=2,
+        dtype=torch.float16,
+    )
+    raw = torch.zeros(
+        num_blocks * num_layers * spec.page_size_bytes,
+        dtype=torch.int8,
+        device="cuda",
+    )
+    caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
+    cache_config = MagicMock(num_blocks=num_blocks)
+    worker = SimpleCPUOffloadWorker(
+        vllm_config=None,
+        kv_cache_config=cache_config,
+        cpu_capacity_bytes=raw.nbytes,
+    )
+    worker._backend = MagicMock()
+    monkeypatch.setattr("vllm.v1.simple_kv_offload.worker.PIN_MEMORY", False)
+
+    worker.register_kv_caches(
+        {f"layer.{layer_idx}": cache for layer_idx, cache in enumerate(caches)}
+    )
+
+    assert worker.gpu_kv_caches is not None
+    if layout.is_layer_compact and not layout.is_block_compact:
+        expected_regions = num_layers * spec.num_heads
+        expected_block_bytes = spec.page_size_bytes // spec.num_heads
+    elif layout.is_layer_compact:
+        expected_regions = num_layers
+        expected_block_bytes = spec.page_size_bytes
+    else:
+        expected_regions = 1
+        expected_block_bytes = spec.page_size_bytes * num_layers
+    assert len(worker.gpu_kv_caches) == expected_regions
+    assert {cache.shape for cache in worker.gpu_kv_caches.values()} == {
+        (num_blocks, expected_block_bytes)
+    }
+
+
+def test_register_separate_kv_head_groups(monkeypatch):
+    # LHBNC hoists the K/V head groups outside the block dim, so each layer's
+    # blocks are registered as one region per group (K, V).
+    layout = KVCacheLayout.LHBNC
+    num_blocks = 4
+    num_layers = 2
+    spec = FullAttentionSpec(
+        block_size=2,
+        num_kv_heads=2,
+        head_size=2,
+        dtype=torch.float16,
+        num_head_slots=2,
+        state_content_bytes=2 * 2 * 2,
+    )
+    raw = torch.zeros(
+        num_blocks * num_layers * spec.page_size_bytes,
+        dtype=torch.int8,
+        device="cuda",
+    )
+    caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
+    worker = SimpleCPUOffloadWorker(
+        vllm_config=None,
+        kv_cache_config=MagicMock(num_blocks=num_blocks),
+        cpu_capacity_bytes=raw.nbytes,
+    )
+    worker._backend = MagicMock()
+    monkeypatch.setattr("vllm.v1.simple_kv_offload.worker.PIN_MEMORY", False)
+
+    worker.register_kv_caches(
+        {f"layer.{layer_idx}": cache for layer_idx, cache in enumerate(caches)}
+    )
+
+    assert worker.gpu_kv_caches is not None
+    assert len(worker.gpu_kv_caches) == num_layers * spec.num_heads
+    per_group_block_bytes = (
+        spec.num_kv_heads * spec.block_size * spec.head_size * spec.dtype.itemsize
+    )
+    assert {cache.shape for cache in worker.gpu_kv_caches.values()} == {
+        (num_blocks, per_group_block_bytes)
+    }

tests/v1/worker/test_attn_utils.py

@@ -1,48 +1,33 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Padded-page handling in create_kv_cache_views.
 
+Guards that a page_size_padded spec strides the block dimension by the padded page
+while keeping per-block content compact, so padding bytes at the end of each page are
+never addressed by the logical view.
+"""
+
+import pytest
 import torch
 
+from tests.v1.attention.utils import dense_kv_cache_views
+from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheLayout,
     KVCacheTensor,
-    KVQuantMode,
-    MambaSpec,
+    MLAAttentionSpec,
+    compute_layout_strides,
+)
+from vllm.v1.worker.utils import (
+    allocate_kv_cache,
+    copy_kv_cache_blocks_inplace,
 )
-from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache
-from vllm.v1.worker.utils import AttentionGroup
 
 
-class FakeFlashAttentionBackend:
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        return (num_blocks, 2, block_size, num_kv_heads, head_size)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        assert not include_num_layers_dimension
-        return (0, 1, 2, 3, 4)
-
-
-class FakeHNDFlashAttentionBackend(FakeFlashAttentionBackend):
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        assert not include_num_layers_dimension
-        return (0, 1, 3, 2, 4)
-
-
-def test_reshape_padded_flash_attention_kv_cache_strides_by_page():
+def test_reshape_padded_kv_cache_strides_by_padded_page():
     num_blocks = 3
     spec = FullAttentionSpec(
         block_size=16,
@@ -53,292 +38,192 @@ def test_reshape_padded_flash_attention_kv_cache_strides_by_page():
     )
     assert spec.real_page_size_bytes == 256
 
-    raw_tensors = {
-        "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
-    }
-    attn_groups = [
-        AttentionGroup(
-            backend=FakeFlashAttentionBackend,
-            layer_names=["layer"],
-            kv_cache_spec=spec,
-            kv_cache_group_id=0,
-        )
-    ]
+    raw = torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
+    (kv_cache,) = dense_kv_cache_views(raw, spec, num_blocks, 1, KVCacheLayout.LBHNC)
 
-    kv_cache = _reshape_kv_cache(
-        attn_groups,
-        raw_tensors,
-        "auto",
-        [spec.block_size],
-        {},
-    )["layer"]
-
-    assert kv_cache.shape == (num_blocks, 2, 16, 1, 2)
-    assert kv_cache.stride(0) == spec.page_size_bytes // 4
-    assert kv_cache.stride(1) == spec.real_page_size_bytes // 2 // 4
-    assert kv_cache[1, 0].storage_offset() == spec.page_size_bytes // 4
-    assert (
-        kv_cache[1, 1].storage_offset()
-        == (spec.page_size_bytes + spec.real_page_size_bytes // 2) // 4
-    )
+    elem_size = 4  # float32
+    # Content dim packs K and V: 2 * head_size.
+    assert kv_cache.shape == (num_blocks, 1, 16, 2 * spec.head_size)
+    assert kv_cache.dtype == spec.dtype
+    assert kv_cache.stride(0) == spec.page_size_padded // elem_size
+    assert kv_cache[1].storage_offset() == spec.page_size_padded // elem_size
+    # Within one block the (unpadded) content stays compact.
+    assert kv_cache[0].is_contiguous()
 
 
-def test_reshape_padded_hnd_flash_attention_kv_cache_strides_by_page():
-    num_blocks = 3
-    spec = FullAttentionSpec(
-        block_size=16,
-        num_kv_heads=3,
-        head_size=2,
-        dtype=torch.float32,
-        page_size_padded=1024,
-    )
-    assert spec.real_page_size_bytes == 768
-
-    raw_tensors = {
-        "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
-    }
-    attn_groups = [
-        AttentionGroup(
-            backend=FakeHNDFlashAttentionBackend,
-            layer_names=["layer"],
-            kv_cache_spec=spec,
-            kv_cache_group_id=0,
-        )
-    ]
-
-    kv_cache = _reshape_kv_cache(
-        attn_groups,
-        raw_tensors,
-        "auto",
-        [spec.block_size],
-        {},
-    )["layer"]
-
-    assert kv_cache.shape == (num_blocks, 2, 16, 3, 2)
-    assert kv_cache.stride(0) == spec.page_size_bytes // 4
-    assert kv_cache.stride(1) == spec.real_page_size_bytes // 2 // 4
-    assert kv_cache.stride(2) == 2
-    assert kv_cache.stride(3) == spec.block_size * spec.head_size
-    assert kv_cache[1, 0].storage_offset() == spec.page_size_bytes // 4
-    assert (
-        kv_cache[1, 1].storage_offset()
-        == (spec.page_size_bytes + spec.real_page_size_bytes // 2) // 4
-    )
-    assert (
-        kv_cache[1, 1, 3, 2].storage_offset()
-        == (
-            spec.page_size_bytes
-            + spec.real_page_size_bytes // 2
-            + 3 * spec.head_size * 4
-            + 2 * spec.block_size * spec.head_size * 4
-        )
-        // 4
-    )
-
-
-class FakeDiffKVBackend:
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        return (num_blocks, block_size, num_kv_heads, head_size * 2)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        assert not include_num_layers_dimension
-        return (0, 1, 2, 3)
-
-
-def test_reshape_padded_diff_kv_cache_does_not_infer_kv_dim():
-    num_blocks = 3
-    spec = FullAttentionSpec(
-        block_size=16,
+@pytest.mark.parametrize(
+    ("kernel_block_sizes", "expected_num_blocks", "expected_num_states"),
+    [
+        (None, 4, 64),
+        ([256], 4, 64),
+        ([64], 16, 16),
+    ],
+)
+def test_allocate_compressed_mla_cache(
+    kernel_block_sizes: list[int] | None,
+    expected_num_blocks: int,
+    expected_num_states: int,
+):
+    spec = MLAAttentionSpec(
+        block_size=256,
         num_kv_heads=1,
-        head_size=2,
-        dtype=torch.float32,
-        page_size_padded=384,
+        head_size=128,
+        dtype=torch.bfloat16,
+        tokens_per_state=4,
     )
-
-    raw_tensors = {
-        "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
-    }
-    attn_groups = [
-        AttentionGroup(
-            backend=FakeDiffKVBackend,
-            layer_names=["layer"],
-            kv_cache_spec=spec,
-            kv_cache_group_id=0,
-        )
-    ]
-
-    kv_cache = _reshape_kv_cache(
-        attn_groups,
-        raw_tensors,
-        "auto",
-        [spec.block_size],
-        {},
-    )["layer"]
-
-    assert kv_cache.shape == (num_blocks, 16, 1, 4)
-    assert kv_cache.stride(0) == spec.page_size_bytes // 4
-    assert kv_cache.stride(1) == 4
-
-
-class FakePerTokenScaleBackend:
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        return (num_blocks, 2, block_size, num_kv_heads, head_size + 4)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        assert not include_num_layers_dimension
-        return (0, 1, 2, 3, 4)
-
-
-def test_reshape_padded_quantized_kv_cache_preserves_scale_stride():
-    num_blocks = 3
-    spec = FullAttentionSpec(
-        block_size=16,
-        num_kv_heads=1,
-        head_size=4,
-        dtype=torch.int8,
-        kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD,
-        page_size_padded=384,
-    )
-    assert spec.real_page_size_bytes == 128
-    assert spec.page_size_bytes == 384
-
-    raw_tensors = {
-        "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
-    }
-    attn_groups = [
-        AttentionGroup(
-            backend=FakePerTokenScaleBackend,
-            layer_names=["layer"],
-            kv_cache_spec=spec,
-            kv_cache_group_id=0,
-        )
-    ]
-
-    kv_cache = _reshape_kv_cache(
-        attn_groups,
-        raw_tensors,
-        "int8_per_token_head",
-        [spec.block_size],
-        {},
-    )["layer"]
-
-    assert kv_cache.shape == (num_blocks, 2, 16, 1, 8)
-    assert kv_cache.stride(0) == spec.page_size_bytes
-    assert kv_cache.stride(1) == 16 * 1 * 8
-    assert kv_cache[1, 1].storage_offset() == spec.page_size_bytes + 16 * 1 * 8
-
-
-class FakeKVFirstBackend:
-    """ROCm-style backend that puts K and V ahead of the block dim."""
-
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        return (2, num_blocks, block_size, num_kv_heads, head_size)
-
-    @staticmethod
-    def get_kv_cache_block_dim(
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> int:
-        return 1
-
-
-def _kv_first_setup(shared_by: list[str]):
-    num_blocks = 3
-    attn_spec = FullAttentionSpec(
-        block_size=16,
-        num_kv_heads=1,
-        head_size=2,
-        dtype=torch.float32,
-    )
-    mamba_spec = MambaSpec(
-        block_size=16,
-        shapes=((64,),),
-        dtypes=(torch.float32,),
-    )
-    assert attn_spec.page_size_bytes == mamba_spec.page_size_bytes == 256
-
-    raw_tensor = torch.zeros(attn_spec.page_size_bytes * num_blocks, dtype=torch.int8)
-    raw_tensors = {name: raw_tensor for name in shared_by}
-    attn_groups = [
-        AttentionGroup(
-            backend=FakeKVFirstBackend,
-            layer_names=["attn"],
-            kv_cache_spec=attn_spec,
-            kv_cache_group_id=0,
-        )
-    ]
-    if "mamba" in shared_by:
-        attn_groups.append(
-            AttentionGroup(
-                backend=FakeKVFirstBackend,
-                layer_names=["mamba"],
-                kv_cache_spec=mamba_spec,
-                kv_cache_group_id=1,
-            )
-        )
-    kv_cache_config = KVCacheConfig(
-        num_blocks=num_blocks,
+    num_pages = 4
+    config = KVCacheConfig(
+        num_blocks=num_pages,
         kv_cache_tensors=[
-            KVCacheTensor(size=raw_tensor.numel(), shared_by=list(shared_by))
+            KVCacheTensor(
+                size=num_pages * spec.page_size_bytes,
+                layers=["layer.0"],
+                layer_stride=num_pages * spec.page_size_bytes,
+                block_stride=spec.page_size_bytes,
+            )
         ],
-        kv_cache_groups=[],
+        kv_cache_groups=[KVCacheGroupSpec(["layer.0"], spec)],
     )
 
-    kv_caches = _reshape_kv_cache(
-        attn_groups,
-        raw_tensors,
-        "auto",
-        [attn_spec.block_size] * len(attn_groups),
-        {},
-        kv_cache_config,
+    caches = allocate_kv_cache(
+        config, torch.device("cpu"), KVCacheLayout.LBHNC, kernel_block_sizes
     )
-    return num_blocks, kv_caches["attn"]
+
+    assert caches["layer.0"].shape == (expected_num_blocks, 1, expected_num_states, 128)
 
 
-def test_reshape_kv_first_kv_cache_pages_blocks_when_shared_with_mamba():
-    num_blocks, kv_cache = _kv_first_setup(["attn", "mamba"])
+@pytest.mark.parametrize("layout", list(KVCacheLayout))
+def test_copy_kv_cache_blocks_shared_storage(layout: KVCacheLayout):
+    num_blocks = 4
+    num_layers = 2
+    spec = FullAttentionSpec(
+        block_size=2,
+        num_kv_heads=2,
+        head_size=2,
+        dtype=torch.float32,
+    )
+    raw = torch.zeros(num_blocks * num_layers * spec.page_size_bytes, dtype=torch.int8)
+    caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
 
-    assert kv_cache.shape == (2, num_blocks, 16, 1, 2)
-    # Block b has to own page b, so its K and V sit side by side within it.
-    page = 16 * 1 * 2 * 2
-    for block in range(num_blocks):
-        assert kv_cache[0, block].storage_offset() == block * page
-        assert kv_cache[1, block].storage_offset() == block * page + page // 2
+    for layer_idx, cache in enumerate(caches):
+        for block_idx in range(num_blocks):
+            cache[block_idx].fill_(10 * layer_idx + block_idx)
+
+    expected = [[cache[i].clone() for i in range(num_blocks)] for cache in caches]
+    copies = [KVCacheBlockCopy(src_block_id=0, dst_block_id=2)]
+
+    copy_kv_cache_blocks_inplace(caches, num_blocks, copies)
+
+    for layer_idx, cache in enumerate(caches):
+        torch.testing.assert_close(cache[2], expected[layer_idx][0])
+        torch.testing.assert_close(cache[1], expected[layer_idx][1])
 
 
-def test_reshape_kv_first_kv_cache_keeps_layout_without_mamba():
-    num_blocks, kv_cache = _kv_first_setup(["attn"])
+def test_fixed_block_stride_propagates_outward_in_lhbnc():
+    num_blocks = 3
+    num_layers = 2
+    spec = FullAttentionSpec(
+        block_size=2,
+        num_kv_heads=2,
+        head_size=2,
+        dtype=torch.float32,
+    )
+    natural = compute_layout_strides(spec, num_blocks, num_layers, KVCacheLayout.LHBNC)
+    block_stride = natural[1] + 8
 
-    assert kv_cache.shape == (2, num_blocks, 16, 1, 2)
-    # Nothing else indexes this allocation by page, so K and V stay split into
-    # one contiguous half each.
-    assert kv_cache[1, 0].storage_offset() == num_blocks * 16 * 1 * 2
+    strides = compute_layout_strides(
+        spec,
+        num_blocks,
+        num_layers,
+        KVCacheLayout.LHBNC,
+        fixed_strides=(None, block_stride, None, None, None),
+    )
+
+    assert strides[1] == block_stride
+    assert strides[2] == block_stride * num_blocks
+    assert strides[0] == strides[2] * spec.num_heads
+
+
+def test_copy_kv_cache_blocks_separate_head_groups():
+    # LHBNC stores each head group separately, so a block's bytes are scattered
+    # across L*H regions.
+    layout = KVCacheLayout.LHBNC
+    num_blocks = 4
+    num_layers = 2
+    spec = FullAttentionSpec(
+        block_size=2,
+        num_kv_heads=2,
+        head_size=2,
+        dtype=torch.float32,
+        num_head_slots=2,
+        state_content_bytes=2 * 2 * 4,
+    )
+    raw = torch.zeros(num_blocks * num_layers * spec.page_size_bytes, dtype=torch.int8)
+    caches = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
+
+    for layer_idx, cache in enumerate(caches):
+        for block_idx in range(num_blocks):
+            for head_idx in range(cache.shape[1]):
+                cache[block_idx, head_idx].fill_(
+                    100 * layer_idx + 10 * head_idx + block_idx
+                )
+
+    expected = [[cache[i].clone() for i in range(num_blocks)] for cache in caches]
+    copy_kv_cache_blocks_inplace(
+        caches,
+        num_blocks,
+        [KVCacheBlockCopy(src_block_id=0, dst_block_id=2)],
+    )
+
+    for layer_idx, cache in enumerate(caches):
+        torch.testing.assert_close(cache[2], expected[layer_idx][0])
+        torch.testing.assert_close(cache[1], expected[layer_idx][1])
+
+
+@pytest.mark.parametrize(
+    "layout,num_layers",
+    [
+        (KVCacheLayout.LBHNC, 2),
+        # Splitting needs a manager block to be one dense page, which a
+        # block-outermost layout only gives when the block holds one layer.
+        (KVCacheLayout.BLHNC, 1),
+    ],
+)
+def test_copy_kv_cache_blocks_with_virtual_block_splitting(
+    layout: KVCacheLayout, num_layers: int
+):
+    num_blocks = 4
+    physical_per_logical = 2
+    spec = FullAttentionSpec(
+        block_size=4,
+        num_kv_heads=1,
+        head_size=2,
+        dtype=torch.float32,
+    )
+    raw = torch.zeros(num_blocks * num_layers * spec.page_size_bytes, dtype=torch.int8)
+    caches = dense_kv_cache_views(
+        raw,
+        spec,
+        num_blocks,
+        num_layers,
+        layout,
+        kernel_block_size=spec.block_size // physical_per_logical,
+    )
+
+    for layer_idx, cache in enumerate(caches):
+        for block_idx in range(cache.shape[0]):
+            cache[block_idx].fill_(100 * layer_idx + block_idx)
+    expected = [[cache[i].clone() for i in range(cache.shape[0])] for cache in caches]
+
+    copy_kv_cache_blocks_inplace(
+        caches,
+        num_blocks,
+        [KVCacheBlockCopy(src_block_id=0, dst_block_id=2)],
+    )
+
+    dst_start = 2 * physical_per_logical
+    for layer_idx, cache in enumerate(caches):
+        for physical_idx in range(physical_per_logical):
+            torch.testing.assert_close(
+                cache[dst_start + physical_idx], expected[layer_idx][physical_idx]
+            )

tests/v1/worker/test_dsv4_packed_zeroer_geometry.py

@@ -1,35 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""CPU regression for the packed-DSV4 KV zeroer geometry.
+"""CPU regression for the packed-DSV4 KV zeroer geometry (upstream #50276).
 
-Constructor-crossing regression for the packed-KV zeroing stride fix (upstream
-PR #50276): every number below comes from instantiating the REAL production
-code (no arithmetic is reimplemented here):
+Four DSV4 fp8_ds_mla layers share one cross-layer allocation: each block row packs
+the four alignment-padded pages side by side (layout BLHNC). The zeroer must step by
+the full packed row per block while zeroing only each layer's meaningful page — never
+the alignment padding, never an adjacent layer's page.
 
-  - ``vllm.v1.core.kv_cache_utils._get_packed_kv_cache_layout``  (packed
-    offsets / block stride)
-  - ``vllm.v1.worker.gpu.attn_utils._reshape_attention_kv_cache``
-    (per-layer column views)
-  - ``vllm.models.deepseek_v4.sparse_mla.DeepseekV4FlashMLABackend``
-    (``get_kv_cache_shape`` / ``get_kv_cache_block_dim``)
-  - ``vllm.v1.worker.utils.AttentionGroup`` / ``KVBlockZeroer``
-    (``KVBlockZeroer.__init__`` precomputes the segment tables)
-
-CPU only: no Triton launch, no CUDA tensor, no CUDA call is made. The test
-asserts the packed-DSV4 constructor metadata:
-
-  * block stride   = the full packed row (4 aligned pages)
-  * per-layer zero = the meaningful page (the spec's unpadded
-                     ``real_page_size_bytes``), so adjacent packed layers are
-                     never wiped
-  * 4 packed layer views; block 99 of the highest-offset layer is the last
-    written region and stays within the packed backing.
-
-The pre-fix constructor exposed a five-field ``_meta`` with no block stride, so
-its zero width equals the FULL packed row; the extraction below records that
-legacy defect and then asserts the fixed semantics, so failure on the old code
-is a semantic one (wrong zero width / out-of-bounds block 99), not a
-schema-shape accident.
+CPU only: the zeroer's ``__init__`` precomputes the segment tables; no kernel is
+launched.
 """
 
 from types import SimpleNamespace
@@ -37,160 +16,154 @@ from types import SimpleNamespace
 import pytest
 import torch
 
-from vllm.models.deepseek_v4.sparse_mla import DeepseekV4FlashMLABackend
-from vllm.v1.core.kv_cache_utils import _get_packed_kv_cache_layout
-from vllm.v1.kv_cache_interface import MLAAttentionSpec, UniformTypeKVCacheSpecs
-from vllm.v1.worker.gpu.attn_utils import _reshape_attention_kv_cache
-from vllm.v1.worker.utils import AttentionGroup, KVBlockZeroer
+from tests.v1.attention.utils import dense_kv_cache_views
+from vllm.v1.kv_cache_interface import KVCacheLayout, MLAAttentionSpec
+from vllm.v1.worker.utils import (
+    AttentionGroup,
+    KVBlockZeroer,
+    allocate_kv_cache,
+)
 
 pytestmark = pytest.mark.cpu_test
 
 NUM_BLOCKS = 100
-BLOCK_SIZE_CFG = 256  # vllm_config.cache_config.block_size for DSV4 (group 0)
-COMPRESS = 4  # main-MLA compress_ratio
-STORAGE_BS = BLOCK_SIZE_CFG // COMPRESS  # 64
-KERNEL_BS = 256  # DeepseekV4FlashMLABackend supports [256]
-LAST_BLOCK = NUM_BLOCKS - 1  # block 99
+NUM_LAYERS = 4
+ALIGNMENT = 576
 
 
-def _extract_segments(meta, base):
-    """Read (offset_bytes, block_stride_bytes, zero_span_bytes) per segment.
-
-    Accepts both the pre-fix five-field ``_meta`` (no block stride; the zero
-    width is the FULL packed row -- the defect) and the fixed six-field ``_meta``.
-    All values come from the real constructor output; nothing is recomputed. Returns
-    ``(segs, legacy_why)`` where ``legacy_why`` describes the pre-fix defect
-    when detected.
-    """
-    legacy_why = None
-    if len(meta) == 6:
-        seg_addrs, seg_block_strides, seg_page_sizes, _max_chunks, _blk, n_segs = meta
-    else:
-        seg_addrs, seg_page_sizes, _max_chunks, _blk, n_segs = meta
-        seg_block_strides = seg_page_sizes
-        legacy_why = (
-            f"legacy {len(meta)}-field _meta (no block stride): zero width is "
-            f"the FULL packed row ({int(seg_page_sizes[0]) * 4} B) -- the defect"
-        )
-    segs = [
-        (int(a) - base, int(bs) * 4, int(ps) * 4)
-        for a, bs, ps in zip(seg_addrs, seg_block_strides, seg_page_sizes)
-    ]
-    return segs, legacy_why, n_segs
-
-
-def test_dsv4_packed_zeroer_geometry():
-    """Real DSV4 constructor metadata separates block stride from zero span."""
-    layer_names = [f"model.layers.{i}.self_attn" for i in range(4)]
-
-    specs = [
-        MLAAttentionSpec(
-            block_size=BLOCK_SIZE_CFG,
-            num_kv_heads=1,
-            head_size=512,
-            dtype=torch.uint8,
-            compress_ratio=COMPRESS,
-            cache_dtype_str="fp8_ds_mla",
-            alignment=576,
-            model_version="deepseek_v4",
-            state_content_bytes=584,  # >576 to allocate room for scales at back of page
-        )
-        for _ in layer_names
-    ]
-    uniform = UniformTypeKVCacheSpecs.from_specs(
-        {ln: sp for ln, sp in zip(layer_names, specs)}
+def test_packed_dsv4_zeroer_zeroes_only_each_layers_page():
+    spec = MLAAttentionSpec(
+        block_size=256,
+        num_kv_heads=1,
+        head_size=512,
+        dtype=torch.uint8,
+        cache_dtype_str="fp8_ds_mla",
+        model_version="deepseek_v4",
+        tokens_per_state=4,
+        alignment=ALIGNMENT,
+        # DeepseekV4 fp8_ds_mla: 584B per token, published by the layer.
+        state_content_bytes=584,
     )
-    assert uniform is not None
+    # DSV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per stored state.
+    assert spec.state_content_size_bytes == 584
+    assert spec.num_states == 64
+    unpadded_page = spec.unpadded_page_size_bytes
+    padded_page = spec.page_size_bytes
+    assert unpadded_page == 64 * 584
+    assert padded_page > unpadded_page, "alignment must pad the page"
 
-    # --- Real packed layout planner -------------------------------------------
-    block_stride, layers_by_offset = _get_packed_kv_cache_layout(
-        [
-            AttentionGroup(
-                backend=DeepseekV4FlashMLABackend,
-                layer_names=layer_names,
-                kv_cache_spec=uniform,
-                kv_cache_group_id=0,
-            )
-        ]
-    )
-    page_bytes = specs[0].page_size_bytes  # DSV4 fp8_ds_mla page after 576B alignment
-    real_page_bytes = specs[0].real_page_size_bytes  # unpadded meaningful page
-    alignment_gap = page_bytes - real_page_bytes  # trailing planner padding
-    total_size = block_stride * NUM_BLOCKS
-    offsets = sorted(layers_by_offset)
-    assert block_stride == 4 * page_bytes, (block_stride, 4 * page_bytes)
-    assert offsets == [i * page_bytes for i in range(4)]
+    raw = torch.zeros(NUM_BLOCKS * NUM_LAYERS * padded_page, dtype=torch.int8)
+    views = dense_kv_cache_views(raw, spec, NUM_BLOCKS, NUM_LAYERS, KVCacheLayout.BLHNC)
+    base = raw.data_ptr()
+    block_row = NUM_LAYERS * padded_page
+    for i, view in enumerate(views):
+        assert view.data_ptr() - base == i * padded_page
+        assert view.stride(0) * view.element_size() == block_row
 
-    # --- Real per-layer views (packed branch) --------------------------------
-    backing = torch.zeros(total_size, dtype=torch.uint8)
-    base = backing.data_ptr()
-    views = {}
-    for ln, off in zip(layer_names, offsets):
-        packing = (off, block_stride)
-        kv_shape = DeepseekV4FlashMLABackend.get_kv_cache_shape(
-            NUM_BLOCKS, STORAGE_BS, 1, 512, cache_dtype_str="fp8_ds_mla"
-        )
-        stride_order = tuple(range(len(kv_shape)))
-        view = _reshape_attention_kv_cache(
-            backing, specs[0], kv_shape, stride_order, NUM_BLOCKS, packing
-        )
-        views[ln] = view
-        assert view.data_ptr() - base == off, "data_ptr must be base+offset"
-        assert view.stride(0) * view.element_size() == block_stride, (
-            "row stride == block_stride (logical block stride)"
-        )
-
-    # --- Real KVBlockZeroer.__init__ (production arithmetic, CPU) ----------
-    sctx = {ln: SimpleNamespace(kv_cache=v) for ln, v in views.items()}
     zeroer = KVBlockZeroer(
         torch.device("cpu"),
-        attn_groups_iter=[
-            AttentionGroup(
-                backend=DeepseekV4FlashMLABackend,
-                layer_names=layer_names,
-                kv_cache_spec=specs[0],
-                kv_cache_group_id=0,
-            )
-        ],
-        kernel_block_sizes=[KERNEL_BS],
-        cache_dtype="fp8_ds_mla",
-        static_forward_context=sctx,
+        attn_groups_iter=iter(
+            [
+                AttentionGroup(
+                    backend=None,
+                    layer_names=[f"layer.{i}" for i in range(NUM_LAYERS)],
+                    kv_cache_spec=spec,
+                    kv_cache_group_id=0,
+                )
+            ]
+        ),
+        kernel_block_sizes=[spec.block_size],
+        static_forward_context={
+            f"layer.{i}": SimpleNamespace(kv_cache=views[i]) for i in range(NUM_LAYERS)
+        },
     )
-    segs, legacy_why, n_segs = _extract_segments(zeroer._meta, base)
+    seg_addrs, seg_block_strides, seg_page_sizes, _, _, n_segs = zeroer._meta
 
-    # --- Fixed semantics (proven by metadata values, not by schema shape) ----
-    assert n_segs == 4, "all 4 packed layers must register as segments"
-    assert [off for off, _, _ in segs] == offsets, "one segment per layer offset"
-    assert all(bs == block_stride for _, bs, _ in segs), (
-        "block stride must be the full packed row (4 aligned pages)"
+    assert n_segs == NUM_LAYERS
+    # Segments step by the full packed row per block...
+    assert (seg_block_strides * 4 == block_row).all()
+    # ...but zero only the layer's meaningful page: no alignment padding, no
+    # adjacent layer's page.
+    assert (seg_page_sizes * 4 == unpadded_page).all()
+    assert sorted(a - base for a in seg_addrs.tolist()) == [
+        i * padded_page for i in range(NUM_LAYERS)
+    ]
+    # Block 99 of the highest-offset layer stays within the packed backing.
+    last_end = (
+        max(seg_addrs.tolist())
+        + (NUM_BLOCKS - 1) * int(seg_block_strides[0]) * 4
+        + int(seg_page_sizes[0]) * 4
     )
-    assert all(zs == real_page_bytes for _, _, zs in segs), (
-        f"zero span must be the meaningful page ({real_page_bytes} B), got "
-        f"{[zs for _, _, zs in segs]} B" + (f"; {legacy_why}" if legacy_why else "")
+    assert last_end <= base + raw.numel()
+
+
+def test_overlaid_zeroer_dedups_segments_with_max_span():
+    """Two groups overlay one allocation; the zeroer must emit one segment per distinct
+    byte offset, spanning the widest overlaid page, so a newly allocated block is fully
+    zeroed no matter which group owns it."""
+    from unittest.mock import MagicMock
+
+    from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups
+    from vllm.v1.kv_cache_interface import (
+        KVCacheGroupSpec,
+        KVCacheLayout,
+        UniformTypeKVCacheSpecs,
     )
 
-    # --- Block 99 (the real final block) bounds, via constructor metadata ---
-    # Production write for block b of a segment is
-    #   [off + b*block_stride_bytes, off + b*block_stride_bytes + zero_span_bytes).
-    worst = 0
-    for off, bs, zs in segs:
-        start = off + LAST_BLOCK * bs
-        end = start + zs
-        oob = end - total_size
-        worst = max(worst, oob)
-        assert end <= total_size, (
-            f"segment@off={off} block {LAST_BLOCK} writes [{start}, {end}) "
-            f"-> {oob} bytes past backing ({total_size})"
+    def make_spec(head_size):
+        return MLAAttentionSpec(
+            block_size=64, num_kv_heads=1, head_size=head_size, dtype=torch.uint8
         )
-    assert worst == 0, f"highest-offset final block goes {worst} bytes OOB"
 
-    # Highest-offset segment's block 99 is the last written region and ends exactly at
-    # total_size - alignment_gap: the trailing planner padding stays untouched.
-    off_last, bs_last, zs_last = segs[-1]
-    assert off_last + LAST_BLOCK * bs_last + zs_last == total_size - alignment_gap, (
-        "highest-offset final block must end at total_size - alignment_gap "
-        f"({alignment_gap}-byte gap intact)"
+    g1_specs = {"g1.big": make_spec(512), "g1.small": make_spec(128)}
+    g2_specs = {"g2.huge": make_spec(1024)}
+    groups = [
+        KVCacheGroupSpec(
+            list(g1_specs),
+            UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=g1_specs),
+        ),
+        KVCacheGroupSpec(
+            list(g2_specs),
+            UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=g2_specs),
+        ),
+    ]
+    from vllm.config import CacheConfig
+
+    vllm_config = MagicMock()
+    vllm_config.cache_config = CacheConfig()
+    vllm_config.cache_config.num_gpu_blocks_override = None
+    vllm_config.cache_config.kv_cache_layout = "BLHNC"
+    config = get_kv_cache_config_from_groups(vllm_config, groups, 8 * 1024 * 1024)
+    views = allocate_kv_cache(config, torch.device("cpu"), KVCacheLayout.BLHNC, None)
+    buf_ptr = views["g1.big"].data_ptr()
+
+    attn_groups = [
+        AttentionGroup(
+            backend=None,
+            layer_names=list(specs),
+            kv_cache_spec=next(iter(specs.values())),
+            kv_cache_group_id=gid,
+        )
+        for gid, specs in enumerate((g1_specs, g2_specs))
+    ]
+    zeroer = KVBlockZeroer(
+        torch.device("cpu"),
+        attn_groups_iter=iter(attn_groups),
+        kernel_block_sizes=[64, 64],
+        static_forward_context={
+            name: SimpleNamespace(kv_cache=views[name]) for name in views
+        },
     )
+    seg_addrs, seg_block_strides, seg_page_sizes, _, _, n_segs = zeroer._meta
 
-    # Zero span < block stride proves adjacent packed layers are never wiped.
-    assert max(zs for _, _, zs in segs) < block_stride
+    # g1.big and g2.huge overlay at offset 0 -> one segment with g2's wider
+    # span; g1.small keeps its own segment.
+    assert n_segs == 2
+    pages = {n: s.page_size_bytes for n, s in (g1_specs | g2_specs).items()}
+    by_offset = {
+        a - buf_ptr: p * 4 for a, p in zip(seg_addrs.tolist(), seg_page_sizes.tolist())
+    }
+    assert by_offset[0] == max(pages["g1.big"], pages["g2.huge"])
+    assert by_offset[pages["g1.big"]] == pages["g1.small"]
+    packed_block_stride = max(sum(pages[n] for n in g) for g in (g1_specs, g2_specs))
+    assert (seg_block_strides * 4 == packed_block_stride).all()

tests/v1/worker/test_gpu_model_runner.py

@@ -92,7 +92,12 @@ def initialize_kv_cache(runner: GPUModelRunner):
     kv_cache_config = KVCacheConfig(
         num_blocks=NUM_BLOCKS,
         kv_cache_tensors=[
-            KVCacheTensor(size=tensor_size, shared_by=["layer.0"]),
+            KVCacheTensor(
+                size=tensor_size,
+                layers=["layer.0"],
+                layer_stride=tensor_size,
+                block_stride=attn_spec.page_size_bytes,
+            ),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=attn_spec)
@@ -131,6 +136,7 @@ def get_vllm_config():
         gpu_memory_utilization=0.9,
         cache_dtype="auto",
     )
+    cache_config.kv_cache_layout = "LBNHC"
     parallel_config = ParallelConfig()
     vllm_config = VllmConfig(
         model_config=model_config,
@@ -839,56 +845,6 @@ def test_update_states_pp_async_multi_request_keeps_rank_state_consistent(
         )
 
 
-def test_kv_cache_stride_order(monkeypatch, model_runner):
-    # This test checks if GPUModelRunner initializes correctly when an attention
-    # backend enforces a non-default KV cache stride order.
-    n_heads = model_runner.model_config.get_num_kv_heads(model_runner.parallel_config)
-    head_size = model_runner.model_config.get_head_size()
-
-    # Get the expected shape from the backend's get_kv_cache_shape method
-    # to ensure compatibility with different backends (triton vs flexattention)
-    attn_backend = None
-    for attn_group in model_runner._attn_group_iterator():
-        attn_backend = attn_group.backend
-        break
-
-    assert attn_backend is not None, "No attention backend found"
-    expected_kv_cache_shape = list(
-        attn_backend.get_kv_cache_shape(NUM_BLOCKS, BLOCK_SIZE, n_heads, head_size)
-    )
-
-    # TODO mla test
-    default_stride = tuple(range(len(expected_kv_cache_shape)))
-    non_default_stride = (*default_stride[1:], default_stride[0])
-    # Permutation that gets you back to expected kv shape
-    for test_stride in (non_default_stride, default_stride):
-
-        def rnd_stride_order(
-            include_num_layers_dimension: bool = False, test_stride=test_stride
-        ):
-            assert not include_num_layers_dimension
-            return test_stride
-
-        # Patch the attention backend class and re-trigger the KV cache creation
-        for attn_group in model_runner._attn_group_iterator():
-            attn_backend = attn_group.backend
-            monkeypatch.setattr(
-                attn_backend, "get_kv_cache_stride_order", rnd_stride_order
-            )
-
-        model_runner.attn_groups = []
-        model_runner.kv_caches = []
-        model_runner.initialize_kv_cache(model_runner.kv_cache_config)
-
-        # Shape is unchanged, but layout may differ
-        kv_cache_shape = model_runner.kv_caches[0].shape
-        assert list(kv_cache_shape) == expected_kv_cache_shape
-        if default_stride == test_stride:
-            assert all(kv.is_contiguous() for kv in model_runner.kv_caches)
-        else:
-            assert all(not kv.is_contiguous() for kv in model_runner.kv_caches)
-
-
 def test_update_config(model_runner):
     # Simple update
     model_runner.update_config({"load_config": {"load_format": "dummy"}})
@@ -1124,21 +1080,22 @@ def test_init_kv_cache_without_kv_sharing(default_vllm_config):
         vllm_config, [kv_cache_spec], [available_memory]
     )[0]
     assert kv_cache_config.num_blocks == num_expected_blocks
-    assert len(kv_cache_config.kv_cache_tensors) == 2
-    assert kv_cache_config.kv_cache_tensors[0].size == available_memory // 2
-    assert kv_cache_config.kv_cache_tensors[1].size == available_memory // 2
+    assert len(kv_cache_config.kv_cache_tensors) == 1
+    assert kv_cache_config.kv_cache_tensors[0].size == available_memory
 
     max_context_len = estimate_max_model_len(vllm_config, kv_cache_spec, 5 * GiB_bytes)
     # max context len with KV sharing should be 2x as large as without
     assert max_context_len == 1310720
 
     # important: override tensor size to prevent large mem alloc during test
-    # this will only allocate 2 block worth of memory (2 * 32kb)
+    # this will only allocate 1 block worth of memory per layer (2 layers * 32kb)
     kv_cache_config.num_blocks = 1
     for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
-        kv_cache_tensor.size = kv_cache_spec[
-            kv_cache_tensor.shared_by[0]
-        ].page_size_bytes
+        page_size = kv_cache_spec[kv_cache_tensor.layers[0]].page_size_bytes
+        kv_cache_tensor.size = page_size * len(kv_cache_tensor.layers)
+        # One block per layer: a layer's region is exactly one page.
+        kv_cache_tensor.layer_stride = page_size
+        kv_cache_tensor.block_stride = page_size
 
     runner.initialize_kv_cache(kv_cache_config)
 
@@ -1207,7 +1164,9 @@ def test_init_kv_cache_with_kv_sharing_valid(default_vllm_config):
     # important: override tensor size to prevent large mem alloc during test
     # this will only allocate 1 block worth of memory (32kb)
     kv_cache_config.num_blocks = 1
-    kv_cache_config.kv_cache_tensors[0].size = kv_cache_spec[layer_0].page_size_bytes
+    page_size = kv_cache_spec[layer_0].page_size_bytes
+    kv_cache_config.kv_cache_tensors[0].size = page_size
+    kv_cache_config.kv_cache_tensors[0].layer_stride = page_size
 
     runner.initialize_kv_cache(kv_cache_config)
     kv_cache_config_after_init = runner.kv_cache_config
@@ -1232,7 +1191,7 @@ def test_hybrid_attention_mamba_tensor_shapes():
     """
     The GPU model runner creates different views into the
     KVCacheTensors for the attention and mamba layers
-    (via _reshape_kv_cache_tensors function). This test verifies
+    (via _allocate_kv_caches). This test verifies
     that the views are compatible: writing a mamba block
     will not corrupt an attention block and vice versa
     """
@@ -1270,6 +1229,7 @@ def test_hybrid_attention_mamba_tensor_shapes():
         gpu_memory_utilization=0.9,
         cache_dtype="auto",
     )
+    cache_config.kv_cache_layout = "LBNHC"
     parallel_config = ParallelConfig()
     attention_config = AttentionConfig(backend=AttentionBackendEnum.FLASHINFER)
     vllm_config = VllmConfig(
@@ -1561,7 +1521,12 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init):
     kv_cache_config = KVCacheConfig(
         num_blocks=NUM_BLOCKS,
         kv_cache_tensors=[
-            KVCacheTensor(size=tensor_size, shared_by=["layer.0"]),
+            KVCacheTensor(
+                size=tensor_size,
+                layers=["layer.0"],
+                layer_stride=tensor_size,
+                block_stride=attn_spec.page_size_bytes,
+            ),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=attn_spec)
@@ -1726,6 +1691,7 @@ def test_mamba_cache_raises_when_max_num_seqs_exceeds_blocks():
         gpu_memory_utilization=0.9,
         cache_dtype="auto",
     )
+    cache_config.kv_cache_layout = "LBNHC"
     parallel_config = ParallelConfig()
     attention_config = AttentionConfig(backend=AttentionBackendEnum.FLASHINFER)
     vllm_config = VllmConfig(

tests/v1/worker/test_kv_block_zeroer.py

@@ -7,8 +7,11 @@ from types import SimpleNamespace
 import pytest
 import torch
 
+from tests.v1.attention.utils import dense_kv_cache_views
 from vllm.v1.kv_cache_interface import (
     ChunkedLocalAttentionSpec,
+    FullAttentionSpec,
+    KVCacheLayout,
     SlidingWindowSpec,
 )
 from vllm.v1.worker import utils as worker_utils
@@ -19,12 +22,6 @@ from vllm.v1.worker.utils import (
 )
 
 
-class _BlockFirstBackend:
-    @staticmethod
-    def get_kv_cache_block_dim(*args, **kwargs):
-        return 0
-
-
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 @pytest.mark.parametrize(
     "spec",
@@ -52,11 +49,8 @@ def test_attention_blocks_are_zeroed(spec):
     layer_name = "draft.self_attn"
     zeroer = KVBlockZeroer(
         device,
-        attn_groups_iter=[
-            AttentionGroup(_BlockFirstBackend, [layer_name], spec, 0)  # type: ignore[arg-type]
-        ],
+        attn_groups_iter=[AttentionGroup(None, [layer_name], spec, 0)],
         kernel_block_sizes=[2],
-        cache_dtype="fp8",
         static_forward_context={
             layer_name: SimpleNamespace(kv_cache=storage),
         },
@@ -205,16 +199,10 @@ def test_large_dsv4_launch_geometry(monkeypatch):
     zeroer = KVBlockZeroer(
         device,
         attn_groups_iter=[
-            AttentionGroup(
-                _BlockFirstBackend,  # type: ignore[arg-type]
-                [name],
-                spec,
-                group_id,
-            )
+            AttentionGroup(None, [name], spec, group_id)
             for group_id, name in enumerate(layer_names)
         ],
         kernel_block_sizes=[1] * n_segs,
-        cache_dtype="auto",
         static_forward_context={
             name: SimpleNamespace(kv_cache=storage)
             for name, storage in storages.items()
@@ -311,3 +299,46 @@ def test_warmup_respects_available_block_count():
     torch.accelerator.synchronize()
 
     assert torch.all(storage == 1)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("layout", list(KVCacheLayout))
+def test_zeroes_exactly_one_block_per_layer(layout: KVCacheLayout):
+    """The zeroer must zero every byte of the target block in every layer and nothing
+    outside it — per head-group region under LHBNC, and never past the target block's
+    tile under block-major layouts (no out-of-bounds writes, no clobbering)."""
+    device = torch.device("cuda")
+    num_blocks, num_layers = 4, 2
+    spec = FullAttentionSpec(
+        block_size=4, num_kv_heads=2, head_size=8, dtype=torch.float32
+    )
+    raw = torch.empty(
+        num_blocks * num_layers * spec.page_size_bytes,
+        dtype=torch.int8,
+        device=device,
+    ).fill_(1)
+    views = dense_kv_cache_views(raw, spec, num_blocks, num_layers, layout)
+    groups = [
+        AttentionGroup(
+            backend=None,
+            layer_names=[f"layer.{i}" for i in range(num_layers)],
+            kv_cache_spec=spec,
+            kv_cache_group_id=0,
+        )
+    ]
+    ctx = {f"layer.{i}": SimpleNamespace(kv_cache=views[i]) for i in range(num_layers)}
+    zeroer = KVBlockZeroer(
+        device,
+        attn_groups_iter=iter(groups),
+        kernel_block_sizes=[spec.block_size],
+        static_forward_context=ctx,
+    )
+    zeroer.zero_block_ids([2])
+    torch.accelerator.synchronize()
+
+    for view in views:
+        assert (view[2] == 0).all(), layout
+        for b in (0, 1, 3):
+            assert (view[b].view(torch.int8) == 1).all(), layout
+    zero_bytes = int((raw == 0).sum().item())
+    assert zero_bytes == num_layers * spec.page_size_bytes, layout

vllm/compilation/passes/fusion/mla_rope_kvcache_cat_fusion.py

@@ -45,7 +45,7 @@ def fused_rope_unified_mla_kv_cache_update_impl(
             cos_sin_cache,
             is_neox,
             layer_slot_mapping,
-            kv_cache,
+            kv_cache.squeeze(1),
             kv_cache_dtype,
             kv_cache_scale,
         )

vllm/config/cache.py

@@ -3,6 +3,7 @@
 
 from collections.abc import Callable
 from dataclasses import field
+from functools import cache
 from typing import Any, ClassVar, Literal
 
 from pydantic import Field, field_validator, model_validator
@@ -13,9 +14,28 @@ from vllm.utils.torch_utils import (
     is_quantized_kv_cache,
     kv_cache_uses_per_token_head_scales,
 )
+from vllm.v1.kv_cache_layout import KVCacheLayout
 
 logger = init_logger(__name__)
 
+_LAYOUT_COMPAT_ALIASES = {
+    "NHD": "LBNHC",
+    "HND": "LBHNC",
+}
+
+
+@cache
+def _layout_from_name(layout_name: str) -> KVCacheLayout:
+    layout_name = _LAYOUT_COMPAT_ALIASES.get(layout_name, layout_name)
+    try:
+        return KVCacheLayout[layout_name]
+    except KeyError:
+        raise ValueError(
+            f"Unknown KV cache layout {layout_name!r}. "
+            f"Valid layouts: {[m.name for m in KVCacheLayout]}"
+        ) from None
+
+
 CacheDType = Literal[
     "auto",
     "float16",
@@ -65,6 +85,17 @@ class CacheConfig:
     """Whether block_size was explicitly provided. Derived automatically."""
     user_specified_mamba_block_size: bool = field(default=False, init=False)
     """Whether mamba_block_size was explicitly provided. Derived automatically."""
+    kv_cache_layout: str | None = field(default=None, init=False)
+    """Resolved physical KV cache layout name (a ``KVCacheLayout`` member).
+
+    ``None`` means the layout has not been resolved yet. The engine core
+    resolves it once (``resolve_kv_cache_layout``) before memory profiling, and
+    every worker process adopts the resolved name — via the
+    ``set_kv_cache_layout`` RPC, or ``KVCacheConfig.kv_cache_layout`` for
+    workers spawned after resolution — before the KV cache is allocated. Once
+    set the value is final and read with ``get_resolved_kv_cache_layout``,
+    which raises on ``None``. Tests and standalone tools may pre-set a value,
+    which resolution then honors as-is."""
     prefix_match_unit: int | None = Field(default=None, gt=0)
     """The finest token boundary (in tokens) a prefix-cache hit can land on.
 
@@ -309,3 +340,12 @@ class CacheConfig:
                 str(cache_dtype),
             )
         return cache_dtype
+
+    def get_resolved_kv_cache_layout(self) -> KVCacheLayout:
+        if self.kv_cache_layout is None:
+            raise ValueError(
+                "KV cache layout has not been resolved yet; it is resolved once "
+                "by the engine core (resolve_kv_cache_layout) unless explicitly "
+                "set by the user."
+            )
+        return _layout_from_name(self.kv_cache_layout)

vllm/distributed/kv_transfer/kv_connector/utils.py

@@ -12,7 +12,7 @@ import torch
 
 from vllm.config import (
     VllmConfig,
-    get_current_vllm_config,
+    get_current_vllm_config_or_none,
     get_layers_from_vllm_config,
     set_current_vllm_config,
 )
@@ -35,20 +35,20 @@ EngineId = str
 BlockIds = tuple[list[int], ...] | list[list[int]]
 
 
-def get_kv_connector_cache_layout():
-    # NOTE (NickLucche) When running disaggregated PD with NIXL, HND layout is
+def get_kv_connector_cache_layout(vllm_config: VllmConfig | None = None):
+    # NOTE (NickLucche) When running disaggregated PD with NIXL, LBHNC layout is
     # used for faster transfer.
-    vllm_config = get_current_vllm_config()
+    if vllm_config is None:
+        vllm_config = get_current_vllm_config_or_none()
+    if vllm_config is None:
+        return None
     kv_config = vllm_config.kv_transfer_config
     if kv_config is not None:
         connector_cls = KVConnectorFactory.get_connector_class(kv_config)
         required_kvcache_layout = connector_cls.get_required_kvcache_layout(vllm_config)
         if required_kvcache_layout is not None:
             return required_kvcache_layout
-        logger.info_once(
-            "Connectors do not specify a kv cache layout, defaulting to NHD."
-        )
-    return "NHD"
+    return None
 
 
 class KVOutputAggregator:
@@ -288,11 +288,11 @@ def kv_postprocess_layout_on_receive(cache, indices):
 
 def kv_postprocess_blksize_and_layout_on_receive(cache, indices, block_size_ratio):
     """
-    Transforms the layout of received KV cache to the local block_size and HND.
+    Transforms the layout of received KV cache to the local block_size and LBHNC.
     (Only works for local blocksize > remote blocksize)
 
-    prefill is HND, smaller block_size
-    decode(local) is NHD, larger block_size
+    prefill is LBHNC, smaller block_size
+    decode(local) is LBNHC, larger block_size
     """
     blocks_to_update = cache.index_select(0, indices)
 
@@ -427,48 +427,6 @@ class TransferTopology:
 
         self._engines: dict[tuple[EngineId, int], EngineTransferInfo] = {}
 
-        # Figure out whether the first dimension of the cache is K/V
-        # or num_blocks.
-        attn_backend = self.attn_backends[0]
-        if not self.is_mamba:
-            _MOCK_BLOCK_SIZE = 16
-            kv_cache_shape: tuple[int, ...] = attn_backend.get_kv_cache_shape(
-                num_blocks=1,
-                block_size=_MOCK_BLOCK_SIZE,
-                num_kv_heads=1,
-                head_size=1,
-            )
-            logger.debug("Test kv_cache_shape: %s", kv_cache_shape)
-            assert kv_cache_shape[0] == 1, (
-                "KV cache layout must be blocks-first; expected mocked "
-                f"num_blocks=1 in leading dim, got shape {kv_cache_shape}."
-            )
-            if not self.is_mla:
-                assert len(kv_cache_shape) == 4, (
-                    "Attention KV cache layout must be standardized as "
-                    "[num_blocks, num_kv_heads, block_size, content_size], "
-                    f"got shape {kv_cache_shape}."
-                )
-
-        self._cross_layers_blocks = False
-        if self.tensor_shape is not None:
-            self._cross_layers_blocks = (
-                len(self.tensor_shape) == len(kv_cache_shape) + 1
-            )
-
-        if self._cross_layers_blocks:
-            logger.debug("Using cross-layer KV cache")
-            _MOCK_NUM_LAYERS = 80
-            kv_cache_shape = (_MOCK_NUM_LAYERS,) + kv_cache_shape
-            try:
-                kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(
-                    include_num_layers_dimension=self._cross_layers_blocks
-                )
-            except (AttributeError, NotImplementedError):
-                assert self.tensor_shape is not None
-                kv_cache_stride_order = tuple(range(len(self.tensor_shape)))
-            kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
-
     # ============================================================
     # Engine registration
     # ============================================================
@@ -503,24 +461,6 @@ class TransferTopology:
         for key in [k for k in self._engines if k[0] == remote_engine_id]:
             del self._engines[key]
 
-    # ============================================================
-    # Layout properties
-    # ============================================================
-
-    @property
-    def cross_layers_blocks(self) -> bool:
-        return self._cross_layers_blocks
-
-    @property
-    def virtually_split_kv_in_blocks(self) -> bool:
-        # Whether to logically split each block into two separately-indexable
-        # sub-regions. With K and V packed into the content dim, an attention
-        # block transfers as a single unit — no K/V sub-split is needed. Only
-        # Mamba still needs this, to index its two state regions (conv/ssm)
-        # separately. Not applicable to cross-layer blocks (per-layer
-        # interleaving means a simple half-split does not separate the parts).
-        return self.is_mamba and not self._cross_layers_blocks
-
     # ============================================================
     # Common methods
     # ============================================================

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -48,7 +48,7 @@ from typing import TYPE_CHECKING, Any, Literal
 import torch
 
 from vllm.logger import init_logger
-from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
+from vllm.v1.attention.backend import AttentionMetadata
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.outputs import KVConnectorOutput
 
@@ -182,14 +182,6 @@ class KVConnectorBase_V1(ABC):
         """
         return False
 
-    @property
-    def prefer_cross_layer_blocks(self) -> bool:
-        """
-        Indicates whether this connector prefers KV blocks that hold KV data for all
-        layers, which can speed up KV data transfers. Defaults to False.
-        """
-        return False
-
     @property
     def requires_kv_delivery(self) -> bool:
         """Whether this connector hands off KV that must be reliably delivered.
@@ -279,23 +271,6 @@ class KVConnectorBase_V1(ABC):
         """
         return
 
-    def register_cross_layers_kv_cache(
-        self, kv_cache: torch.Tensor, attn_backend: type["AttentionBackend"]
-    ):
-        """
-        Initialize with a single KV cache tensor used by all layers.
-        The first dimension should be num_layers.
-        This function will only be called for models with uniform layers,
-        and only if the prefers_cross_layer_blocks is set to True.
-        Only one of the functions
-        {register_kv_caches, register_cross_layers_kv_cache} will be called.
-
-        Args:
-            kv_cache: a cross-layers kv cache tensor
-            attn_backend: The attention backend that corresponds to all layers
-        """
-        return
-
     def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
         """
         Set the xPU-specific ops for copying KV between host and device.

vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py

@@ -129,20 +129,17 @@ class ExampleConnector(KVConnectorBase_V1):
             """Inject the KV cache into the layer.
 
             Args:
-                dst_kv_cache_layer (torch.Tensor): the destination KV cache
-                    layer. In shape [num_pages, page_size, xxx] for MLA,
-                    [num_pages, 2, page_size, xxx] otherwise.
+                dst_kv_cache_layer (torch.Tensor): the destination KV cache layer,
+                    a standardized [B, H, N, C] per-layer view (H == 1 for MLA).
                 src_kv_cache (torch.Tensor): the source KV cache.
                 slot_mapping (torch.Tensor): the slot mapping. In shape
                     [num_tokens].
             """
             slot_mapping = slot_mapping.to(dst_kv_cache_layer.device, non_blocking=True)
             if isinstance(attn_metadata, MLACommonMetadata):
-                dst_kv_cache_layer_shape = dst_kv_cache_layer.shape
-                num_pages = dst_kv_cache_layer_shape[0]
-                page_size = dst_kv_cache_layer_shape[1]
+                # [B, 1, N, C] -> [B * N, C]; slot_mapping indexes B * N slots.
                 dst_kv_cache_layer = dst_kv_cache_layer.reshape(
-                    num_pages * page_size, -1
+                    -1, dst_kv_cache_layer.shape[-1]
                 )
                 dst_kv_cache_layer[slot_mapping, ...] = src_kv_cache
             else:
@@ -225,13 +222,12 @@ class ExampleConnector(KVConnectorBase_V1):
         ) -> torch.Tensor:
             """Extract the KV cache from the layer.
 
-            Assume the shape of the layer is (num_pages, page_size, xxx)
-            for MLA, and (num_pages, 2, page_size, xxx) otherwise.
+            The layer is a standardized [B, H, N, C] per-layer view (H == 1 for MLA).
             """
             slot_mapping = slot_mapping.to(layer.device, non_blocking=True)
             if isinstance(attn_metadata, MLACommonMetadata):
-                num_pages, page_size = layer.shape[0], layer.shape[1]
-                return layer.reshape(num_pages * page_size, -1)[slot_mapping, ...]
+                # [B, 1, N, C] -> [B * N, C]; slot_mapping indexes B * N slots.
+                return layer.reshape(-1, layer.shape[-1])[slot_mapping, ...]
             block_idxs = slot_mapping // self._block_size
             offsets = slot_mapping % self._block_size
             return layer[block_idxs, :, offsets]

vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py

@@ -38,8 +38,10 @@ def extract_from_kv_cache(
     num_tokens: int,
 ) -> torch.Tensor:
     """Extract data from KV cache."""
-    block_size = kv_cache.shape[1]
-    return kv_cache[slot_mapping // block_size, slot_mapping % block_size][:num_tokens]
+    block_size = kv_cache.shape[2]
+    return kv_cache[slot_mapping // block_size, :, slot_mapping % block_size][
+        :num_tokens
+    ]
 
 
 def load_hidden_states(path: str) -> dict[str, torch.Tensor]:
@@ -100,15 +102,6 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
     Must be used in conjunction with the `extract_hidden_states` spec decoding method.
     """
 
-    @property
-    def prefer_cross_layer_blocks(self) -> bool:
-        """
-        Indicates whether this connector prefers KV blocks that hold KV data for all
-        layers, which can speed up KV data transfers. Defaults to False.
-        """
-        # Must be False so that drafter kv cache isn't merged with verifier's
-        return False
-
     @classmethod
     def _find_cache_kv_group_id(cls, kv_cache_config: "KVCacheConfig | None") -> int:
         """Index of the KV cache group holding the extracted hidden states.
@@ -310,10 +303,12 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
 
         # Block size must match the indexed buffer, else reads hit the wrong
         # slots. Raise (not assert) so the check survives `python -O`.
-        if self._block_size != self._kv_cache.shape[1]:
+        # Views are [num_blocks, num_heads, block_size, head_size], matching what
+        # extract_from_kv_cache() indexes.
+        if self._block_size != self._kv_cache.shape[2]:
             raise ValueError(
                 f"Hidden-states block-size mismatch: derived {self._block_size} "
-                f"but buffer block size is {self._kv_cache.shape[1]}; read slots "
+                f"but buffer block size is {self._kv_cache.shape[2]}; read slots "
                 "would be wrong (likely a hybrid block-size resolution bug)."
             )
 
@@ -605,9 +600,9 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
                 "get_required_kvcache_layout should not be called "
                 "on the abstract base class"
             )
-        # NHD means we have (num_tokens, num_heads)
-        # HND means we have (num_heads, num_tokens)
-        # For now, we only support NHD layout since this keeps the
+        # LBNHC means we have (num_tokens, num_heads)
+        # LBHNC means we have (num_heads, num_tokens)
+        # For now, we only support LBNHC layout since this keeps the
         # hidden states for each token together in memory.
-        # HND is primarily used when sharding heads across devices.
-        return "NHD"
+        # LBHNC is primarily used when sharding heads across devices.
+        return "LBNHC"

vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/hf3fs_connector.py

@@ -541,35 +541,32 @@ class HF3FSKVConnector(KVConnectorBase_V1):
         self._dtype = first_cache.dtype
         element_size = first_cache.element_size()
 
-        if self._use_mla:
-            assert len(first_cache.shape) == 3, "MLA format should have 3 dimensions"
-            # MLA format: [num_blocks, block_size, head_size]
-            num_blocks, block_size, head_size = first_cache.shape
-            num_heads = 1
-        else:
-            # MHA format: [2, num_blocks, block_size, num_heads, head_size]
-            _, num_blocks, block_size, num_heads, head_size = first_cache.shape
+        # Standardized per-layer [B, H, N, C] view; MLA is just H == 1 with the
+        # latent vector as the content dim.
+        assert len(first_cache.shape) == 4, (
+            f"expected a [B, H, N, C] KV cache view, got {tuple(first_cache.shape)}"
+        )
+        num_blocks, num_heads, block_size, content_dim = first_cache.shape
+        cache_strides = first_cache.stride()
+        assert all(
+            cache.shape == first_cache.shape and cache.stride() == cache_strides
+            for cache in self._kv_caches.values()
+        ), "HF3FS requires uniform KV cache shapes and strides."
 
         self._local_total_tokens = num_blocks * block_size
         self._local_block_size = block_size
+        self._num_heads = num_heads
+        self._content_dim = content_dim
+        self._kv_cache_strides = cache_strides
 
-        if self._use_mla:
-            layer_block_size = block_size * head_size * element_size
-            self._bytes_per_page = layer_block_size * len(self._kv_caches)
-            self._shape_per_page = [
-                len(self._kv_caches),
-                block_size,
-                head_size,
-            ]
-        else:
-            layer_block_size = 2 * block_size * num_heads * head_size * element_size
-            self._bytes_per_page = layer_block_size * len(self._kv_caches)
-            self._shape_per_page = [
-                len(self._kv_caches),
-                2,
-                block_size,
-                num_heads * head_size,
-            ]
+        layer_block_size = num_heads * block_size * content_dim * element_size
+        self._bytes_per_page = layer_block_size * len(self._kv_caches)
+        self._shape_per_page = [
+            len(self._kv_caches),
+            num_heads,
+            block_size,
+            content_dim,
+        ]
 
         self._kvcache_ptrs = torch.tensor(
             [cache.data_ptr() for cache in self._kv_caches.values()],
@@ -993,7 +990,10 @@ class HF3FSKVConnector(KVConnectorBase_V1):
                     self._local_total_tokens,
                     buffer_tensor,
                     token_indices,
-                    is_mla=self._use_mla,
+                    self._local_block_size,
+                    self._num_heads,
+                    self._content_dim,
+                    self._kv_cache_strides,
                 )
             else:
                 gather_scatter_helper.scatter_kv_caches(
@@ -1001,7 +1001,10 @@ class HF3FSKVConnector(KVConnectorBase_V1):
                     self._local_total_tokens,
                     buffer_tensor,
                     token_indices,
-                    is_mla=self._use_mla,
+                    self._local_block_size,
+                    self._num_heads,
+                    self._content_dim,
+                    self._kv_cache_strides,
                 )
 
     def _compute_prefix_hash(

vllm/distributed/kv_transfer/kv_connector/v1/hf3fs/utils/gather_scatter_helper.py

@@ -12,10 +12,15 @@ def kv_cache_scatter_kernel(
     source_ptr,
     token_indices_ptr,
     num_tokens_in_block,
-    hidden_size,
     total_token_in_kvcache,
     num_layers,
-    is_mla,
+    tokens_per_block,
+    block_stride,
+    head_stride,
+    token_stride,
+    content_stride,
+    num_heads: tl.constexpr,
+    content_size: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
     layer_idx = tl.program_id(0)
@@ -30,39 +35,25 @@ def kv_cache_scatter_kernel(
     if token_idx >= total_token_in_kvcache:
         return
 
-    if is_mla:
-        # MLA format: source [num_layers, num_tokens_in_block, hidden_size]
-        # MLA format: target [total_token_in_kvcache, hidden_size] (per layer)
-        source_offset = (layer_idx * num_tokens_in_block + token_pos) * hidden_size
-        target_offset = token_idx * hidden_size
-
-        for i in range(0, hidden_size, BLOCK_SIZE):
-            offset = i + tl.arange(0, BLOCK_SIZE)
-            mask = offset < hidden_size
-            val = tl.load(source_ptr + source_offset + offset, mask=mask)
-            tl.store(kv_cache_ptr + target_offset + offset, val, mask=mask)
-    else:
-        # MHA format: source [num_layers, 2, num_tokens_in_block, hidden_size]
-        # MHA format: target [2, total_token_in_kvcache, hidden_size]
-        source_offset_k = (
-            layer_idx * num_tokens_in_block * 2 + token_pos
-        ) * hidden_size
-        source_offset_v = (
-            layer_idx * num_tokens_in_block * 2 + num_tokens_in_block + token_pos
-        ) * hidden_size
-
-        target_offset_k = token_idx * hidden_size
-        target_offset_v = (total_token_in_kvcache + token_idx) * hidden_size
-
-        for i in range(0, hidden_size, BLOCK_SIZE):
-            offset = i + tl.arange(0, BLOCK_SIZE)
-            mask = offset < hidden_size
-
-            val_k = tl.load(source_ptr + source_offset_k + offset, mask=mask)
-            val_v = tl.load(source_ptr + source_offset_v + offset, mask=mask)
-
-            tl.store(kv_cache_ptr + target_offset_k + offset, val_k, mask=mask)
-            tl.store(kv_cache_ptr + target_offset_v + offset, val_v, mask=mask)
+    block_idx = token_idx // tokens_per_block
+    state_idx = token_idx % tokens_per_block
+    hidden_size = num_heads * content_size
+    for i in range(0, hidden_size, BLOCK_SIZE):
+        offset = i + tl.arange(0, BLOCK_SIZE)
+        mask = offset < hidden_size
+        head_idx = offset // content_size
+        content_idx = offset % content_size
+        source_offset = (
+            (layer_idx * num_heads + head_idx) * num_tokens_in_block + token_pos
+        ) * content_size + content_idx
+        target_offset = (
+            block_idx * block_stride
+            + head_idx * head_stride
+            + state_idx * token_stride
+            + content_idx * content_stride
+        )
+        value = tl.load(source_ptr + source_offset, mask=mask)
+        tl.store(kv_cache_ptr + target_offset, value, mask=mask)
 
 
 @triton.jit
@@ -71,10 +62,15 @@ def kv_cache_gather_kernel(
     dst_ptr,
     token_indices_ptr,
     num_tokens_in_block,
-    hidden_size,
     total_token_in_kvcache,
     num_layers,
-    is_mla,
+    tokens_per_block,
+    block_stride,
+    head_stride,
+    token_stride,
+    content_stride,
+    num_heads: tl.constexpr,
+    content_size: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
     layer_idx = tl.program_id(0)
@@ -89,37 +85,25 @@ def kv_cache_gather_kernel(
     if token_idx >= total_token_in_kvcache:
         return
 
-    if is_mla:
-        # MLA format: source [total_token_in_kvcache, hidden_size] (per layer)
-        # MLA format: dst [num_layers, num_tokens_in_block, hidden_size]
-        kvcache_offset = token_idx * hidden_size
-        dst_offset = (layer_idx * num_tokens_in_block + token_pos) * hidden_size
-
-        for i in range(0, hidden_size, BLOCK_SIZE):
-            offset = i + tl.arange(0, BLOCK_SIZE)
-            mask = offset < hidden_size
-            val = tl.load(kv_cache_ptr + kvcache_offset + offset, mask=mask)
-            tl.store(dst_ptr + dst_offset + offset, val, mask=mask)
-    else:
-        # MHA format: source [2, total_token_in_kvcache, hidden_size]
-        # MHA format: dst [num_layers, 2, num_tokens_in_block, hidden_size]
-        dst_offset_k = (layer_idx * num_tokens_in_block * 2 + token_pos) * hidden_size
-        dst_offset_v = (
-            layer_idx * num_tokens_in_block * 2 + num_tokens_in_block + token_pos
-        ) * hidden_size
-
-        kvcache_offset_k = token_idx * hidden_size
-        kvcache_offset_v = (total_token_in_kvcache + token_idx) * hidden_size
-
-        for i in range(0, hidden_size, BLOCK_SIZE):
-            offset = i + tl.arange(0, BLOCK_SIZE)
-            mask = offset < hidden_size
-
-            val_k = tl.load(kv_cache_ptr + kvcache_offset_k + offset, mask=mask)
-            val_v = tl.load(kv_cache_ptr + kvcache_offset_v + offset, mask=mask)
-
-            tl.store(dst_ptr + dst_offset_k + offset, val_k, mask=mask)
-            tl.store(dst_ptr + dst_offset_v + offset, val_v, mask=mask)
+    block_idx = token_idx // tokens_per_block
+    state_idx = token_idx % tokens_per_block
+    hidden_size = num_heads * content_size
+    for i in range(0, hidden_size, BLOCK_SIZE):
+        offset = i + tl.arange(0, BLOCK_SIZE)
+        mask = offset < hidden_size
+        head_idx = offset // content_size
+        content_idx = offset % content_size
+        kvcache_offset = (
+            block_idx * block_stride
+            + head_idx * head_stride
+            + state_idx * token_stride
+            + content_idx * content_stride
+        )
+        dst_offset = (
+            (layer_idx * num_heads + head_idx) * num_tokens_in_block + token_pos
+        ) * content_size + content_idx
+        value = tl.load(kv_cache_ptr + kvcache_offset, mask=mask)
+        tl.store(dst_ptr + dst_offset, value, mask=mask)
 
 
 def scatter_kv_caches(
@@ -127,34 +111,33 @@ def scatter_kv_caches(
     total_token_in_kvcache: int,
     src_tensor: torch.Tensor,
     token_indices: list[int],
-    is_mla: bool = False,
+    tokens_per_block: int,
+    num_heads: int,
+    content_size: int,
+    kv_cache_strides: tuple[int, ...],
 ) -> None:
     """Scatter KV cache data from source tensor to KV cache storage.
 
     Args:
         kv_caches_ptrs: Tensor of KV cache pointers (one per layer)
         total_token_in_kvcache: Total number of tokens in KV cache
-        src_tensor: Source tensor containing data to scatter
-            - MHA format: [num_layers, 2, num_tokens_in_block, hidden_size]
-            - MLA format: [num_layers, num_tokens_in_block, hidden_size]
+        src_tensor: Source ``[L, H, N, C]`` tensor containing data to scatter
         token_indices: List of token positions to update
-        is_mla: Whether using MLA model format
+        tokens_per_block: Number of stored states in each cache block
+        num_heads: Size of the H axis
+        content_size: Size of the C axis
+        kv_cache_strides: Element strides of each ``[B, H, N, C]`` layer view
     """
     num_layers = len(kv_caches_ptrs)
     num_tokens_in_block = len(token_indices)
 
-    if is_mla:
-        # MLA: src_tensor is [num_layers, num_tokens_in_block, hidden_size]
-        assert len(src_tensor.shape) == 3, (
-            f"MLA src_tensor should be 3D, got {src_tensor.shape}"
-        )
-        hidden_size = src_tensor.shape[2]
-    else:
-        # MHA: src_tensor is [num_layers, 2, num_tokens_in_block, hidden_size]
-        assert len(src_tensor.shape) == 4, (
-            f"MHA src_tensor should be 4D, got {src_tensor.shape}"
-        )
-        hidden_size = src_tensor.shape[3]
+    assert src_tensor.shape == (
+        num_layers,
+        num_heads,
+        num_tokens_in_block,
+        content_size,
+    )
+    assert len(kv_cache_strides) == 4
 
     device = src_tensor.device
     token_indices_tensor = torch.tensor(
@@ -169,10 +152,12 @@ def scatter_kv_caches(
         src_tensor,
         token_indices_tensor,
         num_tokens_in_block,
-        hidden_size,
         total_token_in_kvcache,
         num_layers,
-        is_mla,
+        tokens_per_block,
+        *kv_cache_strides,
+        num_heads=num_heads,
+        content_size=content_size,
         BLOCK_SIZE=BLOCK_SIZE,
     )
 
@@ -182,49 +167,33 @@ def gather_kv_caches(
     total_token_in_kvcache: int,
     dst_tensor: torch.Tensor,
     token_indices: list[int],
-    is_mla: bool = False,
+    tokens_per_block: int,
+    num_heads: int,
+    content_size: int,
+    kv_cache_strides: tuple[int, ...],
 ) -> None:
     """Gather KV cache data from KV cache storage to destination tensor.
 
     Args:
         kv_caches_ptrs: Tensor of KV cache pointers (one per layer)
         total_token_in_kvcache: Total number of tokens in KV cache
-        dst_tensor: Destination tensor to store gathered data
-            - MHA format: [num_layers, 2, num_tokens_in_block, hidden_size]
-            - MLA format: [num_layers, num_tokens_in_block, hidden_size]
+        dst_tensor: Destination ``[L, H, N, C]`` tensor
         token_indices: List of token positions to gather
-        is_mla: Whether using MLA model format
+        tokens_per_block: Number of stored states in each cache block
+        num_heads: Size of the H axis
+        content_size: Size of the C axis
+        kv_cache_strides: Element strides of each ``[B, H, N, C]`` layer view
     """
     num_layers = kv_caches_ptrs.shape[0]
     num_tokens_in_block = len(token_indices)
 
-    if is_mla:
-        # MLA: dst_tensor is [num_layers, num_tokens_in_block, hidden_size]
-        assert len(dst_tensor.shape) == 3, (
-            f"MLA dst_tensor should be 3D, got {dst_tensor.shape}"
-        )
-        assert dst_tensor.shape[0] == num_layers, (
-            f"Layer count mismatch: {dst_tensor.shape[0]} vs {num_layers}"
-        )
-        assert dst_tensor.shape[1] == num_tokens_in_block, (
-            f"Token count mismatch: {dst_tensor.shape[1]} vs {num_tokens_in_block}"
-        )
-        hidden_size = dst_tensor.shape[2]
-    else:
-        # MHA: dst_tensor is [num_layers, 2, num_tokens_in_block, hidden_size]
-        assert len(dst_tensor.shape) == 4, (
-            f"MHA dst_tensor should be 4D, got {dst_tensor.shape}"
-        )
-        assert dst_tensor.shape[0] == num_layers, (
-            f"Layer count mismatch: {dst_tensor.shape[0]} vs {num_layers}"
-        )
-        assert dst_tensor.shape[1] == 2, (
-            f"MHA should have 2 (K,V) components, got {dst_tensor.shape[1]}"
-        )
-        assert dst_tensor.shape[2] == num_tokens_in_block, (
-            f"Token count mismatch: {dst_tensor.shape[2]} vs {num_tokens_in_block}"
-        )
-        hidden_size = dst_tensor.shape[3]
+    assert dst_tensor.shape == (
+        num_layers,
+        num_heads,
+        num_tokens_in_block,
+        content_size,
+    )
+    assert len(kv_cache_strides) == 4
 
     device = dst_tensor.device
     token_indices_tensor = torch.tensor(
@@ -239,10 +208,12 @@ def gather_kv_caches(
         dst_tensor,
         token_indices_tensor,
         num_tokens_in_block,
-        hidden_size,
         total_token_in_kvcache,
         num_layers,
-        is_mla,
+        tokens_per_block,
+        *kv_cache_strides,
+        num_heads=num_heads,
+        content_size=content_size,
         BLOCK_SIZE=BLOCK_SIZE,
     )
 

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py

@@ -49,15 +49,15 @@ from vllm.model_executor.models.utils import extract_layer_index
 from vllm.platforms import current_platform
 from vllm.utils.math_utils import cdiv
 from vllm.utils.network_utils import get_ip, make_zmq_path, make_zmq_socket
+from vllm.utils.torch_utils import is_non_overlapping_and_dense
 from vllm.v1.attention.backend import AttentionMetadata
-from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, get_kv_cache_layout
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
     FullAttentionSpec,
     KVCacheSpec,
     MambaSpec,
-    MLAAttentionSpec,
-    SlidingWindowMLASpec,
     SlidingWindowSpec,
 )
 from vllm.v1.request import RequestStatus
@@ -122,9 +122,7 @@ def _expand_transfer_regions(
     kv_block_lens: list[int],
     layer_names: list[str],
     layer_indices: list[int],
-    is_kv_layout_blocks_first: bool,
     group_indices: list[int] | None = None,
-    split_kv_regions: list[bool] | None = None,
 ) -> list[TransferRegion]:
     """Expand registered KV tensors into the regions transferred by Mooncake."""
     assert (
@@ -146,13 +144,6 @@ def _expand_transfer_regions(
         "Mooncake transfer regions require matching group metadata lengths, "
         f"got group_indices={len(group_indices)}, layer_names={len(layer_names)}."
     )
-    if split_kv_regions is None:
-        split_kv_regions = [is_kv_layout_blocks_first] * len(layer_names)
-    assert len(split_kv_regions) == len(layer_names), (
-        "Mooncake transfer regions require matching split metadata, "
-        f"got split_kv_regions={len(split_kv_regions)}, "
-        f"layer_names={len(layer_names)}."
-    )
     regions: list[TransferRegion] = []
     for (
         base_addr,
@@ -161,7 +152,6 @@ def _expand_transfer_regions(
         layer_name,
         layer_index,
         group_index,
-        split_kv_region,
     ) in zip(
         base_addrs,
         block_lens,
@@ -169,7 +159,6 @@ def _expand_transfer_regions(
         layer_names,
         layer_indices,
         group_indices,
-        split_kv_regions,
     ):
         regions.append(
             TransferRegion(
@@ -181,17 +170,6 @@ def _expand_transfer_regions(
                 group_index=group_index,
             )
         )
-        if split_kv_region:
-            regions.append(
-                TransferRegion(
-                    layer_name=layer_name,
-                    layer_index=layer_index,
-                    base_addr=base_addr + kv_block_len,
-                    block_len=block_len,
-                    kv_block_len=kv_block_len,
-                    group_index=group_index,
-                )
-            )
     return regions
 
 
@@ -506,10 +484,10 @@ class MooncakeConnector(KVConnectorBase_V1, SupportsHMA):
         if vllm_config.model_config.use_mla:
             return None
         logger.info_once(
-            "MooncakeConnector setting KV cache layout to HND for "
+            "MooncakeConnector setting KV cache layout to LBHNC for "
             "heterogeneous TP-safe KV transfer."
         )
-        return "HND"
+        return "LBHNC"
 
     ############################################################
     # Scheduler Side Methods
@@ -1023,12 +1001,10 @@ class MooncakeConnectorWorker:
         self._sync_block_size_with_kernel()
 
         self.attn_backends = get_current_attn_backends(vllm_config)
-        self.kv_cache_layout = get_kv_cache_layout()
         logger.debug(
             "Detected attention backends %s",
             [backend.get_name() for backend in self.attn_backends],
         )
-        logger.debug("Detected kv cache layout %s", self.kv_cache_layout)
 
         self._tp_size: dict[EngineId, int] = {self.engine_id: self.tp_size}
         self._layer_specs: dict[str, KVCacheSpec] = {}
@@ -1665,7 +1641,7 @@ class MooncakeConnectorWorker:
         self.registered_layer_indices = []
         self.registered_group_indices = []
 
-        for layer_name, cache_or_caches in kv_caches.items():
+        for layer_name, cache in kv_caches.items():
             layer_index = extract_layer_index(layer_name)
             layer_spec = self._layer_specs.get(layer_name)
             if layer_spec is None:
@@ -1674,34 +1650,30 @@ class MooncakeConnectorWorker:
                     layer_name,
                 )
                 continue
-            if isinstance(layer_spec, MambaSpec):
-                conv, _ = cache_or_caches
-                cache_list = [conv]
+            # One raw page tensor per layer; for Mamba that page holds all the
+            # recurrent states, unpacked only when binding the cache for execution.
+            self._log_debug_cache_registration(layer_name, cache)
+            block_is_contiguous = is_non_overlapping_and_dense(cache[0])
+            if not block_is_contiguous:
+                # Non-block-compact layouts scatter a block across per-head
+                # regions; each region's blocks are contiguous.
+                region_caches = [cache[:, head] for head in range(cache.shape[1])]
+                assert all(
+                    is_non_overlapping_and_dense(region[0]) for region in region_caches
+                )
             else:
-                # K and V are packed into one blocks-first tensor per layer,
-                # so each layer registers as a single region.
-                cache_list = [cache_or_caches]
+                region_caches = [cache]
 
-            logger.debug(
-                "registering layer %s with %d cache tensor(s)",
-                layer_name,
-                len(cache_list),
-            )
-
-            for cache in cache_list:
-                self._log_debug_cache_registration(layer_name, cache)
-                base_addr = cache.data_ptr()
-                block_len = cache.stride(0) * cache.element_size()
+            for region_cache in region_caches:
+                base_addr = region_cache.data_ptr()
+                block_len = region_cache.stride(0) * region_cache.element_size()
                 region_base_addresses.append(base_addr)
 
-                if isinstance(layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)):
-                    kv_block_len = layer_spec.page_size_bytes
-                elif self.transfer_topo.virtually_split_kv_in_blocks and not isinstance(
-                    layer_spec, MambaSpec
-                ):
-                    kv_block_len = block_len // 2
-                else:
-                    kv_block_len = block_len
+                kv_block_len = (
+                    layer_spec.page_size_bytes
+                    if isinstance(layer_spec, AttentionSpec) and block_is_contiguous
+                    else block_len
+                )
                 self.block_len_per_layer.append(block_len)
                 self.kv_block_len_per_layer.append(kv_block_len)
                 self.registered_layer_names.append(layer_name)
@@ -1709,12 +1681,12 @@ class MooncakeConnectorWorker:
                 self.registered_group_indices.append(
                     self._layer_group_indices[layer_name]
                 )
-                storage = cache.untyped_storage()
-                storage_addr = storage.data_ptr()
-                if storage_addr not in seen_storage_ptrs:
-                    seen_storage_ptrs.add(storage_addr)
-                    kv_data_ptrs.append(storage_addr)
-                    kv_data_lens.append(storage.nbytes())
+            storage = cache.untyped_storage()
+            storage_addr = storage.data_ptr()
+            if storage_addr not in seen_storage_ptrs:
+                seen_storage_ptrs.add(storage_addr)
+                kv_data_ptrs.append(storage_addr)
+                kv_data_lens.append(storage.nbytes())
 
         self.kv_caches_base_addr = region_base_addresses
         self.seen_base_addresses = kv_data_ptrs
@@ -2055,24 +2027,13 @@ class MooncakeConnectorWorker:
                 self._layer_group_indices.get(layer_name, 0)
                 for layer_name in layer_names
             ]
-        split_kv_regions = None
-        if self.transfer_topo.virtually_split_kv_in_blocks:
-            split_kv_regions = [
-                not isinstance(
-                    self._layer_specs[layer_name],
-                    (MambaSpec, MLAAttentionSpec, SlidingWindowMLASpec),
-                )
-                for layer_name in layer_names
-            ]
         return _expand_transfer_regions(
             base_addrs=base_addrs,
             block_lens=block_lens,
             kv_block_lens=kv_block_lens,
             layer_names=layer_names,
             layer_indices=layer_indices,
-            is_kv_layout_blocks_first=self.transfer_topo.virtually_split_kv_in_blocks,
             group_indices=group_indices,
-            split_kv_regions=split_kv_regions,
         )
 
     def _get_sender_transfer_plan(

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py

@@ -89,14 +89,6 @@ class MooncakeStoreKVEvents(KVConnectorKVEvents):
 class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
     """KV connector using MooncakeDistributedStore as shared KV pool."""
 
-    @property
-    def prefer_cross_layer_blocks(self) -> bool:
-        extra_config = self._kv_transfer_config.kv_connector_extra_config
-        return (
-            str(extra_config.get("enable_cross_layers_blocks", "False")).lower()
-            == "true"
-        )
-
     @staticmethod
     def _validate_kv_cache_config(
         vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
@@ -290,16 +282,6 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
         assert self.connector_worker is not None
         self.connector_worker.register_kv_caches(kv_caches)
 
-    def register_cross_layers_kv_cache(
-        self, kv_cache: torch.Tensor, attn_backend: type
-    ):
-        assert self.connector_worker is not None
-        assert (
-            self._kv_cache_config is not None
-            and len(self._kv_cache_config.kv_cache_groups) == 1
-        ), "Cross-layer KV cache does not supported with hybrid models"
-        self.connector_worker.register_cross_layers_kv_caches(kv_cache)
-
     def start_load_kv(self, forward_context: ForwardContext, **kwargs: Any) -> None:
         # No-op: loads are issued in get_finished() for compute overlap.
         pass

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

@@ -58,6 +58,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.protocol import
 from vllm.logger import init_logger
 from vllm.utils.math_utils import cdiv
 from vllm.utils.network_utils import get_ip, make_zmq_socket
+from vllm.utils.torch_utils import is_non_overlapping_and_dense
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 from vllm.v1.core.kv_cache_utils import (
     BlockHash,
@@ -72,6 +73,7 @@ from vllm.v1.kv_cache_interface import (
     MLAAttentionSpec,
     SlidingWindowMLASpec,
     UniformTypeKVCacheSpecs,
+    group_kernel_blocks,
 )
 
 from .metrics import MooncakeStoreConnectorStats
@@ -1617,18 +1619,9 @@ class MooncakeStoreWorker:
             for g_idx, db in enumerate(self.token_dbs)
         )
 
-    def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:
-        """Register a cross-layers KV cache tensor.
-
-        Wraps the unified tensor in a single-entry dict so that the
-        existing stride-based logic in register_kv_caches() produces
-        the correct single-segment result (block_len = page_size * num_layers).
-        """
-        self.register_kv_caches({"__cross_layer__": kv_cache})
-
     def register_kv_caches(
         self,
-        kv_caches: dict[str, torch.Tensor | list[torch.Tensor]],
+        kv_caches: dict[str, torch.Tensor],
     ) -> None:
         """Register KV cache tensors and start transfer threads."""
         if self._capacity_only:
@@ -1637,58 +1630,62 @@ class MooncakeStoreWorker:
             logger.warning("No KV caches to offload.")
             return
 
-        # Resolve each entry to a representative tensor for storage
-        # deduplication. For attention layers the value is already a tensor;
-        # for Mamba layers it is a list of tensors that all share the same
-        # underlying raw storage, so we take the first one.
-        def _repr_tensor(v: torch.Tensor | list[torch.Tensor]) -> torch.Tensor:
-            assert isinstance(v, torch.Tensor | list)
-            return v if isinstance(v, torch.Tensor) else v[0]
-
         assert self.cache_config.num_gpu_blocks is not None
         self.num_blocks = self.cache_config.num_gpu_blocks
 
-        seen_ptrs: set[int] = set()
+        seen_storage_ptrs: set[int] = set()
+        seen_region_ptrs: set[int] = set()
         addrs: list[int] = []
         block_lens: list[int] = []
 
-        for value in kv_caches.values():
-            cache = _repr_tensor(value)
+        for cache in kv_caches.values():
+            cache = group_kernel_blocks(cache, self.num_blocks)
             cache_storage = cache.untyped_storage()
             base_addr = cache_storage.data_ptr()
-            if base_addr in seen_ptrs:
-                continue
-            seen_ptrs.add(base_addr)
             region_len = cache_storage.nbytes()
 
-            ret = self.store.register_buffer(base_addr, region_len)
-            if ret != 0:
-                logger.error(
-                    "register_buffer failed for addr %#x len %d: %d",
-                    base_addr,
-                    region_len,
-                    ret,
-                )
+            if base_addr not in seen_storage_ptrs:
+                seen_storage_ptrs.add(base_addr)
+                ret = self.store.register_buffer(base_addr, region_len)
+                if ret != 0:
+                    logger.error(
+                        "register_buffer failed for addr %#x len %d: %d",
+                        base_addr,
+                        region_len,
+                        ret,
+                    )
 
-            # Detect layout via stride: a dim whose byte-stride exceeds
-            # page_size_bytes is an outer segment dim (e.g. the K/V dim of
-            # FlashAttn's (2, num_blocks, ...)). FlashInfer/MLA's blocks-
-            # outermost layout has no such dim and yields a single segment.
-            el = cache.element_size()
-            page_size_bytes = region_len // self.num_blocks
-            outer_dims = [
-                d for d in range(cache.ndim) if cache.stride(d) * el > page_size_bytes
-            ]
-            if not outer_dims:
-                # Blocks-first layout (FlashInfer / MLA): one segment.
+            if not is_non_overlapping_and_dense(cache[0]):
+                # A block is scattered across per-head regions; each region's
+                # blocks are contiguous.
+                for head_idx in range(cache.shape[1]):
+                    head_cache = cache[:, head_idx]
+                    assert is_non_overlapping_and_dense(head_cache[0])
+                    region_addr = head_cache.data_ptr()
+                    if region_addr in seen_region_ptrs:
+                        continue
+                    seen_region_ptrs.add(region_addr)
+                    addrs.append(region_addr)
+                    block_lens.append(head_cache.stride(0) * head_cache.element_size())
+            elif cache.stride(0) * cache.element_size() * self.num_blocks == region_len:
+                # The block stride spans the whole per-block window
+                # (block-outermost and packed layouts, and single-layer
+                # tensors), which may hold other layers' pages at higher
+                # offsets. Register the storage once as one whole-window
+                # region; per-layer regions would copy overlapping windows and
+                # run past the storage for offset layers.
+                if base_addr in seen_region_ptrs:
+                    continue
+                seen_region_ptrs.add(base_addr)
                 addrs.append(base_addr)
-                block_lens.append(page_size_bytes)
+                block_lens.append(region_len // self.num_blocks)
             else:
-                # K/V-first layout (FlashAttn / ROCm): split segments.
-                seg_stride = cache.stride(outer_dims[0]) * el
-                for idx in range(cache.shape[outer_dims[0]]):
-                    addrs.append(base_addr + idx * seg_stride)
-                    block_lens.append(seg_stride // self.num_blocks)
+                region_addr = cache.data_ptr()
+                if region_addr in seen_region_ptrs:
+                    continue
+                seen_region_ptrs.add(region_addr)
+                addrs.append(region_addr)
+                block_lens.append(cache.stride(0) * cache.element_size())
 
         logger.info(
             "Registered KV caches: num_groups=%d, num_segments=%d, num_blocks=%d",

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py

@@ -57,11 +57,6 @@ def is_mla_cache_layer(
     return isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec))
 
 
-def _content_packed_dim(spec: AttentionSpec) -> int:
-    head_size_v = getattr(spec, "head_size_v", spec.head_size)
-    return spec.head_size + head_size_v
-
-
 def _spec_dim_matches(value: int, expected: int | None) -> bool:
     return expected is None or value == expected
 
@@ -110,23 +105,6 @@ def get_layer_transfer_geometry(
     spec = layer_to_spec[layer_name]
     is_mla_cache = is_mla_cache_layer(layer_to_spec, layer_name)
 
-    if is_mla_cache and len(shape) == 3:
-        num_blocks, block_size, latent_dim = shape
-        slot_size_bytes = latent_dim * element_size
-        block_len = block_size * slot_size_bytes
-        return LayerTransferGeometry(
-            num_blocks=num_blocks,
-            block_size=block_size,
-            block_len=block_len,
-            slot_size_bytes=slot_size_bytes,
-            block_stride=stride[0],
-            local_kv_stride=None,
-            remote_kv_stride=None,
-            transfers_per_block=1,
-            regions_per_block=1,
-            split_kv_regions=False,
-        )
-
     if not is_mla_cache and len(shape) == 5 and shape[0] == 2:
         _, num_blocks = shape[:2]
         kernel_blocks_per_block = 1
@@ -230,19 +208,19 @@ def get_layer_transfer_geometry(
         )
 
     if (
-        not is_mla_cache
-        and isinstance(spec, AttentionSpec)
+        isinstance(spec, AttentionSpec)
         and len(shape) == 4
-        and shape[1] == spec.num_kv_heads
-        and shape[2] == spec.block_size
-        and shape[3] == _content_packed_dim(spec)
+        and shape[1] == spec.num_heads
+        and shape[2] == spec.num_states
+        and shape[3] * element_size == spec.state_content_size_bytes
     ):
-        num_blocks, num_kv_heads, block_size, packed_dim = shape
-        slot_size_bytes = num_kv_heads * packed_dim * element_size
-        block_len = block_size * slot_size_bytes
+        # Standardized per-layer [B, H, N, C] view (MLA is just H == 1).
+        num_blocks, num_heads, num_states, content_dim = shape
+        slot_size_bytes = num_heads * content_dim * element_size
+        block_len = num_states * slot_size_bytes
         return LayerTransferGeometry(
             num_blocks=num_blocks,
-            block_size=block_size,
+            block_size=spec.block_size,
             block_len=block_len,
             slot_size_bytes=slot_size_bytes,
             block_stride=stride[0],
@@ -267,6 +245,16 @@ def iter_layer_registration_regions(
 ) -> list[tuple[torch.Tensor, int]]:
     geometry = get_layer_transfer_geometry(layer_name, kv_cache, layer_to_spec)
     region_len = geometry.num_blocks * geometry.regions_per_block * geometry.block_len
+    if geometry.regions_per_block == 1:
+        # With padded or interleaved pages the block stride exceeds the
+        # meaningful block_len; register the strided span. The span ends with
+        # the last block's meaningful bytes, not a whole stride, so it never
+        # runs past the backing allocation of a strided layer view.
+        block_stride_bytes = geometry.block_stride * kv_cache.element_size()
+        region_len = max(
+            region_len,
+            (geometry.num_blocks - 1) * block_stride_bytes + geometry.block_len,
+        )
     if geometry.split_kv_regions:
         return [(cache, region_len) for cache in kv_cache]
     return [(kv_cache, region_len)]

vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py

@@ -27,7 +27,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
     PromMetricT,
 )
 from vllm.logger import init_logger
-from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
+from vllm.v1.attention.backend import AttentionMetadata
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.outputs import KVConnectorOutput
 
@@ -209,12 +209,6 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
             c.supports_divergent_local_hybrid_hits for c in self._connectors
         )
 
-    @property
-    def prefer_cross_layer_blocks(self) -> bool:
-        if not self._connectors:
-            return False
-        return all(c.prefer_cross_layer_blocks for c in self._connectors)
-
     @property
     def requires_kv_delivery(self) -> bool:
         return any(c.requires_kv_delivery for c in self._connectors)
@@ -245,13 +239,6 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
             )
         return ret
 
-    def register_cross_layers_kv_cache(
-        self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
-    ):
-        # Register on all connectors
-        for c in self._connectors:
-            c.register_cross_layers_kv_cache(kv_cache, attn_backend)
-
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         for c in self._connectors:
             c.register_kv_caches(kv_caches)

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

@@ -71,9 +71,9 @@ from vllm.platforms import current_platform
 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
 from vllm.utils.network_utils import make_zmq_path
 from vllm.utils.torch_utils import async_tensor_h2d
-from vllm.v1.attention.backends.utils import get_kv_cache_layout
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
+    KVCacheLayout,
     MambaSpec,
     MLAAttentionSpec,
     SlidingWindowMLASpec,
@@ -89,6 +89,13 @@ if TYPE_CHECKING:
 logger = init_logger(__name__)
 
 
+def _share_storage_and_block_stride(caches: list[torch.Tensor]) -> bool:
+    """Return whether all views share storage and a block stride."""
+    block_strides = {cache.stride(0) * cache.element_size() for cache in caches}
+    storage_ptrs = {cache.untyped_storage().data_ptr() for cache in caches}
+    return len(block_strides) == len(storage_ptrs) == 1
+
+
 class NixlBaseConnectorWorker:
     """Base implementation of Worker side methods shared by pull and push."""
 
@@ -302,7 +309,6 @@ class NixlBaseConnectorWorker:
             for group in kv_cache_config.kv_cache_groups
             for layer in group.layer_names
         }
-        self.hma_group_size = len(kv_cache_config.kv_cache_tensors)
 
         # ---- Model state (derived from model config) ----
         mamba_ssm_size = (0, 0)
@@ -523,7 +529,9 @@ class NixlBaseConnectorWorker:
         self.attn_backends = get_current_attn_backends(vllm_config)
         self.backend_name = self.attn_backends[0].get_name()
 
-        self.kv_cache_layout = get_kv_cache_layout()
+        self.kv_cache_layout = (
+            vllm_config.cache_config.get_resolved_kv_cache_layout().name
+        )
         self.host_buffer_kv_cache_layout = self.kv_cache_layout
         logger.info(
             "Detected attention backend(s) %s",
@@ -558,6 +566,11 @@ class NixlBaseConnectorWorker:
         # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`.
         self.block_len_per_layer = list[int]()
 
+        # Per-region block stride in bytes. Taken from the registered tensor's
+        # stride(0) so it stays correct under layouts that interleave layers
+        # within a block (BLHNC/BHLNC), where stride > block_len.
+        self.block_stride_per_layer = list[int]()
+
         # Per-engine TP mappings. Generated during handshake.
         self.tp_mappings: dict[EngineId, TPMapping] = {}
 
@@ -788,16 +801,16 @@ class NixlBaseConnectorWorker:
                 if not self.use_mla:
                     assert kv_cache.ndim == 4
 
-                    if self.kv_cache_layout == "NHD":
+                    if self.kv_cache_layout == "LBNHC":
                         if self.kv_transfer_config.enable_permute_local_kv:
                             logger.info_once(
                                 "'enable_permute_local_kv' flag is enabled while "
-                                "device KV Layout is NHD. Init host buffer with"
-                                " HND to better support Decode/Prefill TP_ratio > 1."
+                                "device KV Layout is LBNHC. Init host buffer with"
+                                " LBHNC to better support Decode/Prefill TP_ratio > 1."
                             )
-                            # Since NHD will not support Decode/Prefill TP_ratio > 1,
+                            # Since LBNHC will not support Decode/Prefill TP_ratio > 1,
                             # we can leverage host_buffer for permute.
-                            self.host_buffer_kv_cache_layout = "HND"
+                            self.host_buffer_kv_cache_layout = "LBHNC"
                         else:
                             # Packed KV layout is logical (B, H, N, 2*D). Allocate
                             # (B, N, H, 2*D) and view it as logical (B, H, N, 2*D)
@@ -969,120 +982,9 @@ class NixlBaseConnectorWorker:
 
         fut.add_done_callback(request_ready)
 
-    def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:
-        """Register a cross-layers KV cache tensor with NIXL.
-
-        `use_uniform_kv_cache()` guarantees a single KV cache group whose
-        layers all share the same `AttentionSpec`, so any layer name from
-        `_layer_specs` yields the correct per-layer spec for `page_size_bytes`.
-        """
-        first_layer = next(iter(self._layer_specs))
-        # Forwarding a real layer name rather than a synthetic key
-        self.register_kv_caches({first_layer: kv_cache})
-
-    def _register_packed_kv_cache(
-        self,
-        storage: torch.UntypedStorage,
-    ) -> None:
-        """Register a packed KV cache as a single NIXL region.
-
-        The packed allocation interleaves all layers per block, so each
-        block_stride-byte chunk is one logical block.  We register 1
-        NIXL region and create 1 descriptor per block.
-        """
-        self.transfer_topo = TransferTopology(
-            tp_rank=self.tp_rank,
-            tp_size=self.world_size,
-            block_size=self.block_size,
-            engine_id=self.engine_id,
-            is_mla=self.use_mla,
-            total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
-            attn_backends=self.attn_backends,
-            tensor_shape=None,
-            is_mamba=self._has_mamba,
-        )
-        self.compat_hash = compute_nixl_compatibility_hash(
-            self.vllm_config,
-            self.backend_name,
-            self.transfer_topo.cross_layers_blocks,
-            transfer_mode=self._TRANSFER_MODE,
-        )
-
-        total_size = storage.nbytes()
-        block_stride = total_size // self.num_blocks
-        base_addr = storage.data_ptr()
-        device_id = storage.device.index
-        assert device_id is not None
-
-        logger.info(
-            "Registering packed KV cache: total_size=%s, block_stride=%s, "
-            "num_blocks=%s, num_regions=1",
-            total_size,
-            block_stride,
-            self.num_blocks,
-        )
-
-        self.device_id = device_id
-        caches_data = [(base_addr, total_size, self.device_id, "")]
-
-        self.block_len_per_layer = [block_stride]
-        self.num_regions = 1
-        self.num_descs = self.num_blocks
-        self.kv_caches_base_addr[self.engine_id][self.tp_rank] = [base_addr]
-
-        descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type)
-        self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
-        self._registered_descs.append(descs)
-
-        self.dst_num_blocks[self.engine_id] = self.num_blocks
-
-        self.src_xfer_handles_by_block_size[self.block_size], (self.src_blocks_data) = (
-            self.register_local_xfer_handler(self.block_size)
-        )
-
-        agent_metadata = NixlAgentMetadata(
-            engine_id=self.engine_id,
-            agent_metadata=self.nixl_wrapper.get_agent_metadata(),
-            device_id=self.device_id,
-            kv_caches_base_addr=(
-                self.kv_caches_base_addr[self.engine_id][self.tp_rank]
-            ),
-            num_blocks=self.num_blocks,
-            block_lens=self.block_len_per_layer,
-            kv_cache_layout=self.kv_cache_layout,
-            block_size=self.block_size,
-            ssm_sizes=self._mamba_ssm_size,
-            attn_backend_name=self.backend_name,
-            physical_blocks_per_logical_kv_block=(
-                self._physical_blocks_per_logical_kv_block
-            ),
-            dcp_size=self.dcp_size,
-            pcp_size=self.pcp_size,
-        )
-        assert self.compat_hash is not None
-        encoder = msgspec.msgpack.Encoder()
-        self.xfer_handshake_metadata = NixlHandshakePayload(
-            compatibility_hash=self.compat_hash,
-            agent_metadata_bytes=encoder.encode(agent_metadata),
-        )
-
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         """Register the KV Cache data in nixl."""
 
-        # Detect packed allocation: all tensors are strided views into the
-        # same backing storage (different data_ptr but same storage).
-        # This happens with DSv4-style contiguous per-block packing.
-        if len(kv_caches) > 1 and not self._has_mamba:
-            storage = next(iter(kv_caches.values())).untyped_storage()
-            storage_ptrs = {
-                cache.untyped_storage().data_ptr() for cache in kv_caches.values()
-            }
-            data_ptrs = {cache.data_ptr() for cache in kv_caches.values()}
-            if len(storage_ptrs) == 1 and len(data_ptrs) > 1:
-                self._register_packed_kv_cache(storage)
-                self.device_kv_caches = kv_caches
-                return
-
         self.transfer_topo = TransferTopology(
             tp_rank=self.tp_rank,
             tp_size=self.world_size,
@@ -1100,7 +1002,6 @@ class NixlBaseConnectorWorker:
         self.compat_hash = compute_nixl_compatibility_hash(
             self.vllm_config,
             self.backend_name,
-            self.transfer_topo.cross_layers_blocks,
             transfer_mode=self._TRANSFER_MODE,
         )
 
@@ -1127,14 +1028,14 @@ class NixlBaseConnectorWorker:
         )
 
         caches_data = []
-        # With hybrid allocator, layers can share a kv cache tensor
+        seen_storage_addresses: set[int] = set()
         seen_base_addresses: list[int] = []
 
+        packed_storage = _share_storage_and_block_stride(list(xfer_buffers.values()))
+
         # K and V are packed into the content dim, so each attention layer is a
         # single NIXL region whose block transfers as one unit. Mamba layers instead
         # register separate conv/ssm sub-regions (see `_build_mamba_local`).
-        tensor_size_bytes = None
-
         for layer_name, cache in xfer_buffers.items():
             # NOTE (NickLucche) Hybrid SSM mamba/FA physical page_size may differ when
             # kernel requires a specific block size. This leads to SSM and FA layers
@@ -1159,55 +1060,98 @@ class NixlBaseConnectorWorker:
                 else layer_spec.page_size_bytes
                 // self._physical_blocks_per_logical_kv_block
             )
-            if self.transfer_topo._cross_layers_blocks:
-                # When cross-layers blocks are used, multiply by number of layers
-                physical_page_size = physical_page_size * len(
-                    self.kv_cache_config.kv_cache_tensors
-                )
             num_blocks = (
                 self._logical_num_blocks
                 if isinstance(layer_spec, MambaSpec)
                 else self.num_blocks
             )
-            # `page_size` accounts for physical blocks, st KVCache is always
-            # [`num_blocks` * `page_size`]
-            curr_tensor_size_bytes = num_blocks * physical_page_size
-
-            base_addr = cache.data_ptr()
-            is_mla_region = isinstance(
-                layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
-            )
-            if base_addr in seen_base_addresses:
-                # NOTE (NickLucche) HMA employs memory pooling to share tensors
-                # across groups. This results in skipping all tensors but the ones
-                # pointed to by group0. Also, generally we will have more blocks
-                # per tensor but fewer regions.
-                # A shared tensor may back both SSM and attention layers (e.g.
-                # KDA+MLA in KimiLinear); the region's FA view is MLA whichever
-                # layer registered it first.
-                idx = seen_base_addresses.index(base_addr)
-                self._region_is_mla[idx] |= is_mla_region
-                logger.debug("Skipping %s because it's already seen", layer_name)
-                continue
             logger.debug(
                 "Registering layer %s with cache shape: %s", layer_name, cache.shape
             )
-            seen_base_addresses.append(base_addr)
-            # Only record non-Mamba page sizes.
-            if isinstance(layer_spec, MambaSpec):
-                self.block_len_per_layer.append(
-                    physical_page_size // self._physical_blocks_per_logical_kv_block
-                )
-            else:
-                self.block_len_per_layer.append(physical_page_size)
-            self._region_is_mla.append(is_mla_region)
+            storage = cache.untyped_storage()
+            storage_addr = storage.data_ptr()
+            # Memory registration follows allocations, while transfer regions follow
+            # logical layers (or contiguous head segments). This keeps strided
+            # cross-layer views inside their registered allocation.
+            if storage_addr not in seen_storage_addresses:
+                seen_storage_addresses.add(storage_addr)
+                self.device_id = max(cache.get_device(), 0)
+                caches_data.append((storage_addr, storage.nbytes(), self.device_id, ""))
 
-            if not is_mla_region:
-                if tensor_size_bytes is None:
-                    tensor_size_bytes = curr_tensor_size_bytes
-                assert tensor_size_bytes == curr_tensor_size_bytes, (
-                    "All non-MLA kv cache tensors must have the same size"
+            is_mla_region = isinstance(
+                layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
+            )
+
+            if isinstance(layer_spec, MambaSpec):
+                region_specs = [
+                    (
+                        cache.data_ptr(),
+                        physical_page_size
+                        // self._physical_blocks_per_logical_kv_block,
+                        physical_page_size
+                        // self._physical_blocks_per_logical_kv_block,
+                    )
+                ]
+            else:
+                if cache.ndim == 1:
+                    # Flat byte view: HMA tensors shared between layer types carry
+                    # no block dimension, so stride(0) is 1 byte. Blocks abut, so
+                    # the stride is the page size.
+                    block_stride = physical_page_size
+                else:
+                    block_stride = cache.stride(0) * cache.element_size()
+                storage_is_block_major = num_blocks * block_stride == storage.nbytes()
+                # A layer whose [H, N, C] interior is dense addresses its own page
+                # as one chunk, so it can own a region. When it does not, or when
+                # every layer is packed into this one storage, the block's whole
+                # row is the only contiguous unit.
+                hnc_contiguous = (
+                    cache.ndim == 4
+                    and cache.stride(2) == cache.shape[3]
+                    and cache.stride(1) == cache.shape[2] * cache.shape[3]
                 )
+                if storage_is_block_major and (
+                    (packed_storage and is_mla_region) or not hnc_contiguous
+                ):
+                    # TODO(Lucas): handle TP slicing for packed_storage; for now
+                    # restrict to MLA (DSv4) where kv is replicated.
+                    storage_block_len = storage.nbytes() // num_blocks
+                    region_specs = [
+                        (storage_addr, storage_block_len, storage_block_len)
+                    ]
+                elif storage_is_block_major:
+                    region_specs = [
+                        (cache.data_ptr(), physical_page_size, block_stride)
+                    ]
+                else:
+                    segment_bytes = num_blocks * block_stride
+                    assert cache.nbytes % segment_bytes == 0
+                    num_segments = cache.nbytes // segment_bytes
+                    region_block_len = (
+                        block_stride if num_segments > 1 else physical_page_size
+                    )
+                    region_specs = [
+                        (
+                            cache.data_ptr() + segment_idx * segment_bytes,
+                            region_block_len,
+                            block_stride,
+                        )
+                        for segment_idx in range(num_segments)
+                    ]
+
+            for base_addr, block_len, block_stride in region_specs:
+                if base_addr in seen_base_addresses:
+                    if is_mla_region:
+                        # Dual-purpose HMA tensor: an MLA layer shares a region
+                        # that a non-MLA layer registered first. MLA is not
+                        # head-sharded, so the region must be flagged MLA.
+                        idx = seen_base_addresses.index(base_addr)
+                        self._region_is_mla[idx] = True
+                    continue
+                seen_base_addresses.append(base_addr)
+                self.block_len_per_layer.append(block_len)
+                self.block_stride_per_layer.append(block_stride)
+                self._region_is_mla.append(is_mla_region)
 
             # When there's a mismatch between kbs<>bs, we rely on HMA to ensure
             # caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs.
@@ -1228,11 +1172,6 @@ class NixlBaseConnectorWorker:
                     f"kv_cache_layout={self.kv_cache_layout}"
                 )
 
-            # Need to make sure the device ID is non-negative for NIXL,
-            # Torch uses -1 to indicate CPU tensors.
-            self.device_id = max(cache.get_device(), 0)
-            caches_data.append((base_addr, curr_tensor_size_bytes, self.device_id, ""))
-
         logger.debug(
             "Different block lengths collected: %s", set(self.block_len_per_layer)
         )
@@ -1240,10 +1179,11 @@ class NixlBaseConnectorWorker:
             len(self.block_len_per_layer)
             == len(seen_base_addresses)
             == len(self._region_is_mla)
+            == len(self.block_stride_per_layer)
         )
 
         self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses
-        self.num_regions = len(caches_data)
+        self.num_regions = len(seen_base_addresses)
 
         if self.pp_size > 1:
             start_layer, end_layer = self.model_config.get_layers_start_end_indices(
@@ -1293,6 +1233,7 @@ class NixlBaseConnectorWorker:
             kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id][self.tp_rank],
             num_blocks=self.num_blocks,
             block_lens=self.block_len_per_layer,
+            block_strides=self.block_stride_per_layer,
             kv_cache_layout=self.kv_cache_layout
             if not self.use_host_buffer
             else self.host_buffer_kv_cache_layout,
@@ -1431,10 +1372,9 @@ class NixlBaseConnectorWorker:
         block_arange = np.arange(num_blocks, dtype=np.uint64)
         parts: list[np.ndarray] = []
         for i, base_addr in enumerate(base_addresses):
-            # K/V are packed into the content dim, so the whole block transfers
-            # as one unit: desc length equals the block stride.
             block_len = self.block_len_per_layer[i] // block_size_ratio
-            addrs = base_addr + block_arange * block_len
+            block_stride = self.block_stride_per_layer[i] // block_size_ratio
+            addrs = base_addr + block_arange * block_stride
             parts.append(self._stack_descs(addrs, block_len, device_id))
         return np.concatenate(parts)
 
@@ -1476,7 +1416,7 @@ class NixlBaseConnectorWorker:
             )
             local_block_len = local_block_len // num_reads
 
-            page_size = nixl_agent_meta.block_lens[i]
+            page_size = nixl_agent_meta.block_strides[i]
             addrs = base_addr + rank_offset + block_arange * page_size
             parts.append(self._stack_descs(addrs, local_block_len, device_id))
         return np.concatenate(parts)
@@ -1560,7 +1500,7 @@ class NixlBaseConnectorWorker:
                                                  tp_ratio = 4 // 2 = 2
 
         Considering the KV Caches, if P-Worker_i has cache size [2, num_blocksP, kv_heads, block_size, head_dim]
-        then D-Worker_j has [2, num_blocksD, kv_heads//tp_ratio, block_size, head_dim]. Mind the "HND" layout format.
+        then D-Worker_j has [2, num_blocksD, kv_heads//tp_ratio, block_size, head_dim]. Mind the "LBHNC" layout format.
         Assuming num_blocksD >= num_blocksP, D-Worker0 reads from P-Worker0 by preparing the kv_heads//tp_ratio
         first heads from all the slots of all the blocks. D-Worker1 will do the same, but reading the second split
         along the kv_heads dimension, and so forth until "tp_ratio" D TP workers have pulled from P-Worker0.
@@ -1600,6 +1540,7 @@ class NixlBaseConnectorWorker:
                 start:end
             ]
             nixl_agent_meta.block_lens = nixl_agent_meta.block_lens[start:end]
+            nixl_agent_meta.block_strides = nixl_agent_meta.block_strides[start:end]
 
         ### Register remote engine in TransferTopology (idempotent).
         assert self.transfer_topo is not None
@@ -1790,10 +1731,10 @@ class NixlBaseConnectorWorker:
         if not self.use_mla and nixl_agent_meta.kv_cache_layout != kv_cache_layout:
             if (
                 self.kv_transfer_config.enable_permute_local_kv
-                and nixl_agent_meta.kv_cache_layout == "HND"
+                and nixl_agent_meta.kv_cache_layout == "LBHNC"
             ):
                 logger.info(
-                    "Remote is HND and local is NHD, enabled additional permute "
+                    "Remote is LBHNC and local is LBNHC, enabled additional permute "
                     "on local device KV."
                 )
                 assert not self._is_hma_required, (
@@ -1823,18 +1764,18 @@ class NixlBaseConnectorWorker:
             self.enable_heterogeneous_attn_post_process = True
 
         # Heterogeneous TP requires head-splitting, which only works with
-        # HND layout. MLA and replicated-KV cases don't split on heads.
-        # Mamba doesn't support heterogeneous TP.
+        # block-contiguous layouts (e.g. LBHNC). MLA and replicated-KV cases
+        # don't split on heads. Mamba doesn't support heterogeneous TP.
         if (
             abs(tp_ratio) != 1
             and not self.use_mla
             and not self.transfer_topo.is_kv_replicated(remote_engine_id)
-            and kv_cache_layout != "HND"
+            and not KVCacheLayout[kv_cache_layout].is_block_contiguous
             and not self.enable_permute_local_kv
         ):
             raise RuntimeError(
                 "Heterogeneous TP head-dimension splitting requires contiguous heads. "
-                "Use HND layout on the prefill side."
+                "Use a block-contiguous layout (e.g. LBHNC) on the prefill side."
             )
 
         # Per-region block_len validation enforcing the P/D invariant.

vllm/distributed/kv_transfer/kv_connector/v1/nixl/connector.py

@@ -19,7 +19,6 @@ import torch
 from vllm.config import VllmConfig
 from vllm.distributed.kv_transfer.kv_connector.utils import (
     EngineId,
-    get_current_attn_backend,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
     CopyBlocksOp,
@@ -56,10 +55,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import (
 )
 from vllm.forward_context import ForwardContext
 from vllm.logger import init_logger
-from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
-from vllm.v1.attention.backends.utils import get_kv_cache_layout
+from vllm.v1.attention.backend import AttentionMetadata
 from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.kv_cache_interface import MambaSpec
 from vllm.v1.outputs import KVConnectorOutput
 
 if TYPE_CHECKING:
@@ -83,36 +80,6 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
     def supports_divergent_local_hybrid_hits(self) -> bool:
         return True
 
-    @property
-    def prefer_cross_layer_blocks(self) -> bool:
-        if any(
-            [
-                isinstance(group.kv_cache_spec, MambaSpec)
-                for group in self.kv_cache_config.kv_cache_groups
-            ]
-        ):
-            # Hybrid SSM models do not yet support cross-layer layout
-            return False
-
-        backend = get_current_attn_backend(self._vllm_config)
-        if backend.get_name() not in (
-            "FLASH_ATTN",
-            "FLASHINFER",
-            "TRITON_ATTN",
-        ):
-            return False
-
-        # For now there is no benefit to run cross layers when backend
-        # does not support on HND
-        if get_kv_cache_layout() != "HND":
-            return False
-
-        extra_config = self.kv_transfer_config.kv_connector_extra_config
-        return (
-            str(extra_config.get("enable_cross_layers_blocks", "False")).lower()
-            == "true"
-        )
-
     def __init__(
         self,
         vllm_config: VllmConfig,
@@ -179,9 +146,10 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
             # which fallback to the default behavior.
             return None
         logger.info_once(
-            "NixlConnector setting KV cache layout to HND for better xfer performance."
+            "NixlConnector setting KV cache layout to LBHNC for "
+            "better xfer performance."
         )
-        return "HND"
+        return "LBHNC"
 
     ############################################################
     # Scheduler Side Methods
@@ -266,12 +234,6 @@ class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
         assert self.connector_worker is not None
         self.connector_worker.register_kv_caches(kv_caches)
 
-    def register_cross_layers_kv_cache(
-        self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
-    ):
-        assert self.connector_worker is not None
-        self.connector_worker.register_cross_layers_kv_caches(kv_cache)
-
     def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
         assert self.connector_worker is not None
         self.connector_worker.set_host_xfer_buffer_ops(copy_operation)

vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py

@@ -44,8 +44,9 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:"
 #   6: Validate EAGLE/MTP speculative configuration compatibility
 #   7: Include NIXL transfer mode (push vs pull) in the compatibility hash
 #   8: Add dcp_size and pcp_size to NixlAgentMetadata
+#   9: Add block_strides
 #
-NIXL_CONNECTOR_VERSION: int = 8
+NIXL_CONNECTOR_VERSION: int = 9
 
 
 @dataclass
@@ -56,6 +57,7 @@ class NixlAgentMetadata:
     device_id: int
     num_blocks: int
     block_lens: list[int]
+    block_strides: list[int]
     kv_cache_layout: str
     block_size: int
     ssm_sizes: tuple[int, int]
@@ -129,7 +131,6 @@ def _get_speculative_compatibility_factors(
 def compute_nixl_compatibility_hash(
     vllm_config: VllmConfig,
     attn_backend_name: str,
-    cross_layers_blocks: bool,
     transfer_mode: str = "pull",
 ) -> str:
     """
@@ -180,7 +181,6 @@ def compute_nixl_compatibility_hash(
         # Attention backend and KV cache dtype affect memory layout
         "attn_backend_name": attn_backend_name,
         "cache_dtype": str(cache_config.cache_dtype),
-        "cross_layers_blocks": cross_layers_blocks,
         "is_hma_enabled": is_hma_enabled,
         "speculative_config": _get_speculative_compatibility_factors(vllm_config),
         # push (WRITE) and pull (READ) connectors are protocol-incompatible

vllm/distributed/kv_transfer/kv_connector/v1/offloading/canonical_mapping.py

@@ -17,6 +17,7 @@ import torch
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
     KVCacheConfig,
+    KVCacheLayout,
     KVCacheSpec,
     MambaSpec,
     MLAAttentionSpec,
@@ -31,13 +32,16 @@ if TYPE_CHECKING:
 CANONICAL_FORMAT_VERSION = 1
 
 
-def canonical_format_id() -> str:
+def canonical_format_id(kv_cache_layout: str) -> str:
     """Identity of the canonical byte format, for namespacing persisted KV.
     Canonical pages keep the worker's KV layout family, so the id couples the
-    format version with that family; consumers must match it exactly."""
-    from vllm.v1.attention.backends.utils import get_kv_cache_layout
-
-    return f"v{CANONICAL_FORMAT_VERSION}-{get_kv_cache_layout().lower()}"
+    format version with that family; consumers must match it exactly. The family
+    keeps its historical NHD/HND spelling so ids stay stable for KV persisted
+    before the layout enum existed."""
+    layout = KVCacheLayout[kv_cache_layout]
+    legacy = {KVCacheLayout.LBNHC: "nhd", KVCacheLayout.LBHNC: "hnd"}
+    family = legacy.get(layout, layout.name.lower())
+    return f"v{CANONICAL_FORMAT_VERSION}-{family}"
 
 
 @dataclass(frozen=True)
@@ -278,7 +282,7 @@ def _layer_mapping(
     if isinstance(spec, MLAAttentionSpec):
         # TP-replicated latent; CP shards its tokens across the DCP groups
         if (
-            spec.compress_ratio != 1
+            spec.tokens_per_state != 1
             or page % bs
             or ctx.tp_size % ctx.dcp_size
             or spec.kv_quant_mode.is_per_token_head

vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py

@@ -20,12 +20,7 @@ from vllm.v1.kv_offload.config import (
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
-    from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheTensor
-
-
-def is_kv_cache_tensor_packed(kv_cache_tensor: "KVCacheTensor") -> bool:
-    """Return whether a KV cache tensor uses a packed block stride."""
-    return bool(kv_cache_tensor.block_stride)
+    from vllm.v1.kv_cache_interface import KVCacheConfig
 
 
 def build_offloading_config(
@@ -96,18 +91,10 @@ def build_offloading_config(
         blocks_per_chunk = tokens_per_chunk_int // tokens_per_block
 
     worker_kv_bytes_per_block = 0
-    if kv_cache_config.num_blocks > 0:
-        packed_tensors = tuple(
-            is_kv_cache_tensor_packed(tensor)
-            for tensor in kv_cache_config.kv_cache_tensors
-        )
-        is_packed = any(packed_tensors)
-        assert not is_packed or all(packed_tensors)
-        total_gpu_kv_bytes = (
-            kv_cache_config.kv_cache_tensors[0].size
-            if is_packed
-            else sum(tensor.size for tensor in kv_cache_config.kv_cache_tensors)
-        )
+    if kv_cache_config.num_blocks > 0 and kv_cache_config.kv_cache_tensors:
+        # Every KVCacheTensor describes placement within the same backing allocation,
+        # so its size is the total, not a per-tensor share.
+        total_gpu_kv_bytes = kv_cache_config.kv_cache_tensors[0].size
         worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks
 
     single_group_spec = (
@@ -164,7 +151,7 @@ def build_offloading_config(
         if isinstance(single_group_spec, FullAttentionSpec):
             if type(single_group_spec) is MLAAttentionSpec:
                 spec_certifiable = (
-                    single_group_spec.compress_ratio == 1
+                    single_group_spec.tokens_per_state == 1
                     and single_group_spec.real_page_size_bytes
                     % single_group_spec.block_size
                     == 0
@@ -219,4 +206,5 @@ def build_offloading_config(
         ),
         replicated_layout=replicated_layout,
         canonical_layout=canonical_layout,
+        kv_cache_layout=vllm_config.cache_config.kv_cache_layout,
     )

vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py

@@ -14,16 +14,13 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import (
     OffloadingWorkerMetadata,
     ReqId,
 )
-from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
-    is_kv_cache_tensor_packed,
-)
 from vllm.logger import init_logger
-from vllm.v1.attention.backend import AttentionBackend
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
     KVCacheConfig,
     MambaSpec,
     UniformTypeKVCacheSpecs,
+    group_kernel_blocks,
 )
 from vllm.v1.kv_offload.base import (
     CanonicalKVCacheRef,
@@ -73,16 +70,6 @@ class OffloadingConnectorWorker:
             self.vllm_config, kv_cache_config, kv_caches
         )
 
-        # Packed layouts (e.g. DSv4) set block_stride > 0; their tensors use
-        # stride(0) as the manager-block stride (equals total_num_bytes_per_block).
-        # General (non-packed) layouts size the tensor at page_size_bytes per
-        # manager block, so page_size_bytes is the correct offloading stride.
-        layer_is_packed: dict[str, bool] = {
-            ln: is_kv_cache_tensor_packed(kv_tensor)
-            for kv_tensor in kv_cache_config.kv_cache_tensors
-            for ln in kv_tensor.shared_by
-        }
-
         # layer_name -> (num_blocks, page_size_bytes) tensor
         tensors_per_block: dict[str, tuple[torch.Tensor, ...]] = {}
         # layer_name -> size of (un-padded) page in bytes
@@ -100,63 +87,52 @@ class OffloadingConnectorWorker:
                 layer_kv_cache_spec = per_layer_specs.get(
                     layer_name, group_kv_cache_spec
                 )
-                if isinstance(layer_kv_cache_spec, AttentionSpec):
-                    layer_kv_cache = kv_caches[layer_name]
-                    assert isinstance(layer_kv_cache, torch.Tensor)
+                ref = group_kernel_blocks(kv_caches[layer_name], num_blocks)
+                page = layer_kv_cache_spec.page_size_bytes
+                elem_size = ref.element_size()
+                byte_offset = ref.storage_offset() * elem_size
+                # Packed layouts (e.g. DSv4) interleave layers per block, so the
+                # attention tensor's stride(0) (the manager-block stride) exceeds
+                # page_size_bytes; other layouts have stride(0) == page_size_bytes.
+                block_stride_bytes = (
+                    ref.stride(0) * elem_size
+                    if isinstance(layer_kv_cache_spec, AttentionSpec)
+                    else page
+                )
+                tensors_per_block[layer_name] = (
+                    torch.tensor([], dtype=torch.int8, device=ref.device).set_(
+                        ref.untyped_storage(),
+                        byte_offset,
+                        (num_blocks, page),
+                        (block_stride_bytes, 1),
+                    ),
+                )
+                page_size_bytes[layer_name] = page
 
-                    page = layer_kv_cache_spec.page_size_bytes
-                    elem_size = layer_kv_cache.element_size()
-                    byte_offset = layer_kv_cache.storage_offset() * elem_size
-                    block_stride_bytes = (
-                        layer_kv_cache.stride(0) * elem_size
-                        if layer_is_packed[layer_name]
-                        else page
-                    )
-                    raw = torch.empty(
-                        0,
-                        dtype=torch.int8,
-                        device=layer_kv_cache.device,
-                    ).set_(layer_kv_cache.untyped_storage())
-                    tensors_per_block[layer_name] = (
-                        torch.as_strided(
-                            raw,
-                            (num_blocks, page),
-                            (block_stride_bytes, 1),
-                            byte_offset,
-                        ),
-                    )
-                    page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes
+                if isinstance(layer_kv_cache_spec, AttentionSpec):
                     unpadded_page_size_bytes[layer_name] = (
                         layer_kv_cache_spec.unpadded_page_size_bytes
                     )
-
                 elif isinstance(layer_kv_cache_spec, MambaSpec):
-                    layer_kv_cache = kv_caches[layer_name]
-                    assert layer_kv_cache.dtype == torch.int8
-                    tensors_per_block[layer_name] = (
-                        layer_kv_cache.view(
-                            num_blocks, layer_kv_cache_spec.page_size_bytes
-                        ),
-                    )
-
-                    page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes
                     unpadded_page_size_bytes[layer_name] = replace(
                         layer_kv_cache_spec, page_size_padded=None
                     ).page_size_bytes
-
                 else:
                     raise NotImplementedError
 
-        packed_kv_cache_tensor = next(
+        # Packed layouts (e.g. DSv4) interleave all layers within each manager
+        # block: a layer view's block stride exceeds its page size. Offload the
+        # whole packed block as a single transfer region.
+        packed_layer_name = next(
             (
-                t
-                for t in kv_cache_config.kv_cache_tensors
-                if is_kv_cache_tensor_packed(t) and t.shared_by
+                layer_name
+                for layer_name, (tensor,) in tensors_per_block.items()
+                if tensor.stride(0) != tensor.shape[1]
             ),
             None,
         )
-        if packed_kv_cache_tensor is not None:
-            (tensor,) = tensors_per_block[packed_kv_cache_tensor.shared_by[0]]
+        if packed_layer_name is not None:
+            (tensor,) = tensors_per_block[packed_layer_name]
             block_stride = tensor.stride(0)
             packed_tensor = tensor.as_strided(
                 (num_blocks, block_stride),
@@ -176,29 +152,16 @@ class OffloadingConnectorWorker:
 
         block_tensors: list[CanonicalKVCacheTensor] = []
         block_data_refs: dict[str, list[CanonicalKVCacheRef]] = defaultdict(list)
-        for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
-            # Filter to layers that were actually processed above.
-            # Packed KV allocation emits KVCacheTensor entries for
-            # every (tuple_idx, page_size) slot; slots where no group has a
-            # layer at that index produce an empty shared_by (reserved memory
-            # with no corresponding model layer).
-            tensor_layer_names = [
-                n for n in kv_cache_tensor.shared_by if n in tensors_per_block
-            ]
-            if not tensor_layer_names:
-                continue
+        # Layers that alias the same bytes (cache groups overlay each other) share
+        # one transfer region: exactly the views with equal address and strides.
+        aliased_layers: dict[tuple[int, tuple[int, ...]], list[str]] = defaultdict(list)
+        for layer_name, layer_tensors in tensors_per_block.items():
+            view = layer_tensors[0]
+            aliased_layers[(view.data_ptr(), view.stride())].append(layer_name)
 
-            # verify all layers in the group reference the exact same tensors
+        for tensor_layer_names in aliased_layers.values():
             assert len({len(tensors_per_block[n]) for n in tensor_layer_names}) == 1
-            assert (
-                len({tensors_per_block[n][0].data_ptr() for n in tensor_layer_names})
-                == 1
-            )
-            assert (
-                len({tensors_per_block[n][0].stride() for n in tensor_layer_names}) == 1
-            )
 
-            # pick the first layer to represent the group
             first_layer_name = tensor_layer_names[0]
             for tensor in tensors_per_block[first_layer_name]:
                 block_tensors.append(
@@ -242,53 +205,6 @@ class OffloadingConnectorWorker:
 
         self._init_worker(canonical_kv_caches)
 
-    def register_cross_layers_kv_cache(
-        self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
-    ):
-        # verify that num_blocks is at physical position 0 in the cross-layers
-        # tensor layout.
-        test_shape = attn_backend.get_kv_cache_shape(
-            num_blocks=1234, block_size=16, num_kv_heads=1, head_size=256
-        )
-        num_blocks_logical_dim = test_shape.index(1234) + 1
-        physical_to_logical = attn_backend.get_kv_cache_stride_order(
-            include_num_layers_dimension=True
-        )
-        num_blocks_physical_dim = physical_to_logical.index(num_blocks_logical_dim)
-        assert num_blocks_physical_dim == 0
-
-        kv_cache_groups = self.kv_cache_config.kv_cache_groups
-        assert len(kv_cache_groups) == 1
-        kv_cache_spec = kv_cache_groups[0].kv_cache_spec
-        num_layers = len(kv_cache_groups[0].layer_names)
-        page_size_bytes = kv_cache_spec.page_size_bytes * num_layers
-
-        assert kv_cache.storage_offset() == 0
-        storage = kv_cache.untyped_storage()
-        assert len(storage) % page_size_bytes == 0
-        num_blocks = len(storage) // page_size_bytes
-        tensor = (
-            torch.tensor(
-                [],
-                dtype=torch.int8,
-                device=kv_cache.device,
-            )
-            .set_(storage)
-            .view(num_blocks, page_size_bytes)
-        )
-        kv_cache_tensor = CanonicalKVCacheTensor(
-            tensor=tensor, page_size_bytes=page_size_bytes
-        )
-        # in cross layers layout, there's currently only a single group
-        kv_cache_data_ref = CanonicalKVCacheRef(
-            tensor_idx=0, page_size_bytes=page_size_bytes
-        )
-        canonical_kv_caches = CanonicalKVCaches(
-            tensors=[kv_cache_tensor], group_data_refs=[[kv_cache_data_ref]]
-        )
-
-        self._init_worker(canonical_kv_caches)
-
     def handle_preemptions(self, kv_connector_metadata: OffloadingConnectorMetadata):
         assert self.worker is not None
 

vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py

@@ -37,7 +37,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.worker import (
     OffloadingConnectorWorker,
 )
 from vllm.forward_context import ForwardContext
-from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
+from vllm.v1.attention.backend import AttentionMetadata
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.kv_cache_interface import KVCacheConfig
@@ -47,11 +47,6 @@ from vllm.v1.request import Request
 
 
 class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
-    @property
-    def prefer_cross_layer_blocks(self) -> bool:
-        # Cross-layer slabs have no per-layer refs to certify canonical mappings
-        return not self._canonical_layout
-
     @property
     def requires_kv_delivery(self) -> bool:
         # Runs as kv_both, but is a best-effort cache: a dropped save is just a
@@ -91,12 +86,6 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
         assert self.connector_worker is not None
         self.connector_worker.register_kv_caches(kv_caches)
 
-    def register_cross_layers_kv_cache(
-        self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
-    ):
-        assert self.connector_worker is not None
-        self.connector_worker.register_cross_layers_kv_cache(kv_cache, attn_backend)
-
     def handle_preemptions(self, kv_connector_metadata: KVConnectorMetadata):
         assert self.connector_worker is not None
         assert isinstance(kv_connector_metadata, OffloadingConnectorMetadata)
@@ -196,7 +185,7 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
 
     @classmethod
     def get_required_kvcache_layout(cls, vllm_config: VllmConfig) -> str | None:
-        return "HND"
+        return "LBHNC"
 
     def reset_cache(self) -> bool | None:
         assert self.connector_scheduler is not None

vllm/envs.py

@@ -242,7 +242,10 @@ if TYPE_CHECKING:
     VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16
     VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300
     VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS: int = 5
-    VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None
+    VLLM_KV_CACHE_LAYOUT: (
+        Literal["LBNHC", "LBHNC", "LHBNC", "NHD", "HND", "BLHNC", "BLNHC", "BHLNC"]
+        | None
+    ) = None
     VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None
     VLLM_COMPUTE_NANS_IN_LOGITS: bool = False
     VLLM_RAISE_ON_LOGIT_NANS: bool = False
@@ -1784,18 +1787,22 @@ environment_variables: dict[str, Callable[[], Any]] = {
     ),
     # KV Cache layout used throughout vllm.
     # Some common values are:
-    # - NHD
-    # - HND
-    # Where N=num_blocks, H=num_heads and D=head_size. The default value will
+    # - LBNHC
+    # - LBHNC
+    # Where N=num_states, H=num_heads and C=state_content. The default value will
     # leave the layout choice to the backend. Mind that backends may only
     # implement and support a subset of all possible layouts.
+    # LHBNC hoists the head dim outside the block dim; backends must opt in via
+    # AttentionBackend.supported_kv_cache_layouts().
     "VLLM_KV_CACHE_LAYOUT": env_with_choices(
-        "VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"]
+        "VLLM_KV_CACHE_LAYOUT",
+        None,
+        ["LBNHC", "LBHNC", "LHBNC", "NHD", "HND", "BLHNC", "BLNHC", "BHLNC"],
     ),
     # SSM conv state layout used for Mamba models.
     # - SD: (state_len, dim) — dim contiguous (default)
     # - DS: (dim, state_len) — TP-sharded dim on dim1,
-    #   consistent with SSM temporal state and HND KV cache layout.
+    #   consistent with SSM temporal state and LBHNC KV cache layout.
     "VLLM_SSM_CONV_STATE_LAYOUT": env_with_choices(
         "VLLM_SSM_CONV_STATE_LAYOUT", None, ["SD", "DS"]
     ),

vllm/model_executor/layers/attention/mla_attention.py

@@ -649,6 +649,10 @@ class MLAAttention(nn.Module, AttentionLayerBase):
             compile_native=True,
         )
 
+    def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
+        # [B, H=1, N, C] -> [B, N, C]
+        self.kv_cache = kv_cache.squeeze(1)
+
     @property
     def chunked_prefill_workspace_size(self) -> int:
         if self._chunked_prefill_workspace_size is None:
@@ -1451,27 +1455,6 @@ class MLACommonBackend(AttentionBackend):
     def get_builder_cls() -> type["MLACommonMetadataBuilder"]:
         return MLACommonMetadataBuilder
 
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,  # assumed to be 1 for MLA
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        return (num_blocks, block_size, head_size)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        if include_num_layers_dimension:
-            # Default to identity permutation to signal cross-layer allocation
-            # is unsupported. Each MLA backend must opt in to support cross-layer
-            # allocation by overriding this method.
-            return (0, 1, 2, 3)
-        return (0, 1, 2)
-
     @classmethod
     def get_supported_head_sizes(cls) -> list[int]:
         return [320, 576]

vllm/model_executor/layers/hpc/rope_norm.py

@@ -61,16 +61,16 @@ def hpc_rope_norm_forward(
         return
 
     attn_layer = forward_context.no_compile_layers[layer_name]
-    # bind_kv_cache stores the per-layer KV cache as a single 5D tensor
-    # (num_blocks, 2, block_size, num_kv_heads, head_size), so use it directly.
+    # bind_kv_cache stores the per-layer KV cache as a single tensor
+    # (num_blocks, num_kv_heads, block_size, 2 * head_size), so use it directly.
     kv_cache = attn_layer.kv_cache
 
     if kv_cache.numel() == 0:
         output.zero_()
         return
 
-    assert kv_cache.dim() == 5, (
-        f"Expected kv_cache to have 5 dims, got {tuple(kv_cache.shape)}"
+    assert kv_cache.dim() == 4, (
+        f"Expected kv_cache to have 4 dims, got {tuple(kv_cache.shape)}"
     )
 
     rope_norm = _hpc_rope_norm_instances[layer_name]
@@ -295,6 +295,9 @@ class HpcRopeNorm(CustomOp, HpcModule):
         """
         import hpc
 
+        # (B, H, N, 2*hs) -> two (B, N, H, hs) K/V views, as in HpcAttentionImpl.
+        key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_dim, dim=-1)
+
         num_actual_tokens = attn_metadata.num_actual_tokens
         num_prefill_reqs = attn_metadata.num_prefills
         num_decode_reqs = attn_metadata.num_decodes
@@ -333,8 +336,8 @@ class HpcRopeNorm(CustomOp, HpcModule):
 
             if self.use_fp8:
                 _, q_scale, split_k_flag = hpc.rope_norm_store_kv_fp8(
-                    key_cache=kv_cache[:, 0],
-                    value_cache=kv_cache[:, 1],
+                    key_cache=key_cache,
+                    value_cache=value_cache,
                     qkv=qkv_prefill,
                     cos_sin=self.cos_sin_cache,
                     num_seqlen_per_req=seq_lens_prefill,
@@ -353,8 +356,8 @@ class HpcRopeNorm(CustomOp, HpcModule):
                 attn_metadata.hpc_prefill_q_scale = q_scale
             else:
                 hpc.rope_norm_store_kv(
-                    kv_cache[:, 0],
-                    kv_cache[:, 1],
+                    key_cache,
+                    value_cache,
                     qkv_prefill,
                     self.cos_sin_cache,
                     seq_lens_prefill,
@@ -379,8 +382,8 @@ class HpcRopeNorm(CustomOp, HpcModule):
 
             if self.use_fp8:
                 _, q_scale, split_k_flag = hpc.rope_norm_store_kv_fp8(
-                    key_cache=kv_cache[:, 0],
-                    value_cache=kv_cache[:, 1],
+                    key_cache=key_cache,
+                    value_cache=value_cache,
                     qkv=qkv_decode,
                     cos_sin=self.cos_sin_cache,
                     num_seqlen_per_req=num_seq_kvcache,
@@ -401,8 +404,8 @@ class HpcRopeNorm(CustomOp, HpcModule):
                     attn_metadata.hpc_split_k_flag = split_k_flag
             else:
                 hpc.rope_norm_store_kv(
-                    kv_cache[:, 0],
-                    kv_cache[:, 1],
+                    key_cache,
+                    value_cache,
                     qkv_decode,
                     self.cos_sin_cache,
                     num_seq_kvcache,

vllm/model_executor/models/deepseek_v2.py

@@ -645,6 +645,11 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):
             raise ValueError(f"Duplicate layer name: {prefix}")
         compilation_config.static_forward_context[prefix] = self
 
+    def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
+        # [B, H=1, N, C] -> [B, N, C]: the indexer kernels and
+        # kv_cache_as_quant_view index a 3-D block-major cache.
+        self.kv_cache = kv_cache.squeeze(1)
+
     def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
         return MLAAttentionSpec(
             block_size=self.cache_config.block_size,

vllm/model_executor/models/extract_hidden_states.py

@@ -80,14 +80,14 @@ def dummy_attention(layer_name, _placeholder):
 
 def basic_cache(
     to_cache: torch.Tensor,  # shape: [seq_len, num_heads, head_size]
-    kv_cache: torch.Tensor,  # shape: [num_blocks, block_size, num_heads, head_size]
+    kv_cache: torch.Tensor,  # shape: [num_blocks, num_heads, block_size, head_size]
     slot_mapping: torch.Tensor,  # shape: [seq_len]
 ):
     # Padding slots are -1; redirect them to the null block (block 0, never
     # allocated to a request) so the scatter stays branch-free and sync-free.
-    block_size = kv_cache.shape[1]
+    block_size = kv_cache.shape[2]
     slot_mapping = slot_mapping.clamp_min(0)
-    kv_cache[slot_mapping // block_size, slot_mapping % block_size] = to_cache
+    kv_cache[slot_mapping // block_size, :, slot_mapping % block_size] = to_cache
 
 
 ######### CacheOnlyAttentionBackend ########
@@ -123,18 +123,6 @@ class CacheOnlyAttentionBackend(AttentionBackend):
     def get_impl_cls() -> type["CacheOnlyAttentionImpl"]:
         return CacheOnlyAttentionImpl
 
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        # We set `num_kv_heads = num_hidden_layers` and `head_size = hidden_size`
-        # We also don't use a k/v (2) dim
-        return (num_blocks, block_size, num_kv_heads, head_size)
-
     @staticmethod
     def get_builder_cls() -> type["CacheOnlyAttentionMetadataBuilder"]:
         return CacheOnlyAttentionMetadataBuilder

vllm/model_executor/models/whisper_causal.py

@@ -5,6 +5,7 @@ import functools
 import logging
 import math
 from dataclasses import replace
+from fractions import Fraction
 from functools import partial
 
 import torch
@@ -128,14 +129,10 @@ def create_whisper_attention_backend_with_block_pooling(
             vllm_config: VllmConfig,
             device: torch.device,
         ):
-            assert kv_cache_spec.num_kv_heads % block_pool_size == 0
-            # Scale pooled-unit quantities back to unpooled (encoder-token)
-            # units for the underlying attention metadata: the KV cache spec
-            # counts blocks in pooled units, but the kernel runs on the
-            # `block_pool_size`x-expanded sequence built below.
-            spec_overrides = dict(
+            kv_cache_spec = replace(
+                kv_cache_spec,
                 block_size=kv_cache_spec.block_size * block_pool_size,
-                num_kv_heads=kv_cache_spec.num_kv_heads // block_pool_size,
+                tokens_per_state=1,
             )
             if isinstance(kv_cache_spec, SlidingWindowSpec):
                 # The manager keeps `sliding_window` in pooled units; the kernel
@@ -143,8 +140,7 @@ def create_whisper_attention_backend_with_block_pooling(
                 # unpooled units (`get_kv_cache_spec` sizes the pooled window so
                 # this window always stays resident).
                 assert sliding_window is not None
-                spec_overrides["sliding_window"] = sliding_window
-            kv_cache_spec = replace(kv_cache_spec, **spec_overrides)
+                kv_cache_spec = replace(kv_cache_spec, sliding_window=sliding_window)
             super().__init__(kv_cache_spec, layer_names, vllm_config, device)
             # Override model_config-derived values with the actual
             # encoder values from kv_cache_spec
@@ -261,18 +257,6 @@ def create_whisper_attention_backend_with_block_pooling(
         overrides={
             "get_builder_cls": lambda: WhisperCausalAttentionWithBlockPoolingBuilder,
             "get_impl_cls": lambda: WhisperCausalAttentionWithBlockPoolingImpl,
-            "get_kv_cache_shape": lambda num_blocks,
-            block_size,
-            num_kv_heads,
-            head_size,
-            cache_dtype_str: underlying_attn_backend.get_kv_cache_shape(
-                num_blocks,
-                # we stretch each block by `block_pool_size`
-                block_size * block_pool_size,
-                num_kv_heads // block_pool_size,
-                head_size,
-                cache_dtype_str,
-            ),
             "forward_includes_kv_cache_update": True,
         },
     )
@@ -341,7 +325,7 @@ class WhisperCausalAttentionWithBlockPooling(Attention):
         assert isinstance(kv_cache_spec, AttentionSpec)
         kv_cache_spec = replace(
             kv_cache_spec,
-            num_kv_heads=self.block_pool_size * kv_cache_spec.num_kv_heads,
+            tokens_per_state=Fraction(1, self.block_pool_size),
         )
         if isinstance(kv_cache_spec, SlidingWindowSpec):
             # The manager counts blocks in pooled units, so express the window in

vllm/models/deepseek_v4/attention.py

@@ -653,7 +653,8 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
         # per-tensor fp8 writes a separately-allocated fp8 q and quantizes the
         # KV row.
         block_size = swa_metadata.block_size
-        swa_kv_cache_3d = swa_kv_cache.view(-1, block_size, self.head_dim)
+        assert swa_kv_cache.shape[1:] == (block_size, self.head_dim)
+        swa_kv_cache_3d = swa_kv_cache
         if cache_dtype == torch.bfloat16:
             torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
                 q,
@@ -684,6 +685,17 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
         )
         return q_fp8
 
+    def _global_topk_output_buffers(
+        self, topk_indices: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor] | None:
+        if self.compress_ratio != 4 or self.eager_scratch_pool is None:
+            return None
+        return self.eager_scratch_pool.global_topk_outputs(topk_indices)
+
+    def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
+        # [B, H=1, N, C] -> [B, N, C]
+        self.kv_cache = kv_cache.squeeze(1)
+
     def get_attn_backend(self) -> type[AttentionBackend]:
         return self.backend_cls
 
@@ -701,7 +713,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
             num_kv_heads=1,
             head_size=self.head_dim,
             dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,
-            compress_ratio=self.compress_ratio,
+            tokens_per_state=self.compress_ratio,
             cache_dtype_str=self.kv_cache_dtype,
             alignment=576 if uses_fp8_ds_mla_layout else 512,
             model_version="deepseek_v4",
@@ -733,16 +745,20 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
             raise ValueError(f"Duplicate layer name: {prefix}")
         compilation_config.static_forward_context[prefix] = self
 
+    def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
+        # [B, H=1, N, C] -> [B, N, C]
+        self.kv_cache = kv_cache.squeeze(1)
+
     def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
         # head_dim already carries the fp8 scale padding
-        # compress_ratio=1 for V3.2, >1 for DeepseekV4; both use the same cache layout.
+        # tokens_per_state=1 for V3.2, >1 for DeepseekV4; same cache layout.
         uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
         return MLAAttentionSpec(
             block_size=self.cache_config.block_size,
             num_kv_heads=1,
             head_size=self.head_dim,
             dtype=self.dtype,
-            compress_ratio=self.compress_ratio,
+            tokens_per_state=self.compress_ratio,
             # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).
             alignment=576 if uses_fp8_ds_mla_layout else 512,
         )

vllm/models/deepseek_v4/compressor.py

@@ -75,25 +75,6 @@ class CompressorBackend(AttentionBackend):
     def get_builder_cls() -> type["CompressorMetadataBuilder"]:
         return CompressorMetadataBuilder
 
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        assert num_kv_heads == 1
-        return (num_blocks, block_size, head_size)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        if include_num_layers_dimension:
-            return (0, 1, 2, 3)
-        return (0, 1, 2)
-
 
 @dataclass
 class CompressorMetadata:
@@ -185,6 +166,10 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
         else:
             raise ValueError(f"Invalid compress ratio: {compress_ratio}")
 
+    def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
+        # [B, H=1, N, C] -> [B, N, C]
+        self.kv_cache = kv_cache.squeeze(1)
+
     def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
         # fp8_ds_mla is the UE8M0 paged layout and needs 576B alignment. Plain
         # full-cache rows share state pages with contiguous KV pages, so padding

vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py

@@ -22,7 +22,6 @@ from vllm.models.deepseek_v4.sparse_mla import (
     DeepseekV4FlashMLAMetadata,
     DeepseekV4SparseMLABackend,
 )
-from vllm.platforms import current_platform
 from vllm.platforms.interface import DeviceCapability
 from vllm.utils.flashinfer import flashinfer_trtllm_batch_decode_sparse_mla_dsv4
 from vllm.v1.attention.backend import MultipleOf
@@ -165,26 +164,6 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):
             return None
         return "FLASHINFER_MLA_SPARSE_DSV4 requires SM10x or SM12x"
 
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        device_capability = current_platform.get_device_capability()
-        if device_capability is not None and device_capability.major == 12:
-            return DeepseekV4SparseMLABackend.get_kv_cache_shape(
-                num_blocks,
-                block_size,
-                num_kv_heads,
-                head_size,
-                cache_dtype_str,
-            )
-        assert num_kv_heads == 1
-        return (num_blocks, block_size, head_size)
-
 
 class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
     """FlashInfer TRTLLM-gen sparse MLA attention layer for SM100 DeepSeek V4."""

vllm/models/deepseek_v4/sparse_mla.py

@@ -88,21 +88,6 @@ class DeepseekV4SparseMLABackend(AttentionBackend):
     def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
         return capability.major in [9, 10]
 
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        if cache_dtype_str == "fp8_ds_mla":
-            # DeepseekV4 main MLA: 584B per token (448 NoPE + 128 RoPE + 8 fp8 scale).
-            # head_size passed in is the semantic head_dim (512).
-            return (num_blocks, block_size, 584)
-        else:
-            return (num_blocks, block_size, head_size)
-
 
 @dataclass
 class DeepseekV4FlashMLAMetadata(AttentionMetadata):
@@ -151,8 +136,8 @@ class DeepseekV4SparseMLAMetadataBuilder(
             (max_num_batched_tokens,), dtype=torch.int32, device=device
         )
 
-        assert hasattr(self.kv_cache_spec, "compress_ratio")
-        self.compress_ratio = self.kv_cache_spec.compress_ratio
+        assert isinstance(self.kv_cache_spec.tokens_per_state, int)
+        self.compress_ratio = self.kv_cache_spec.tokens_per_state
 
         # Pre-allocate compressed slot mapping buffer for CUDA graph address
         # stability when compress_ratio > 1.
@@ -206,7 +191,7 @@ class DeepseekV4SparseMLAMetadataBuilder(
                 cm.query_start_loc,
                 cm.seq_lens,
                 cm.block_table_tensor.clamp_(min=0),
-                int(self.kv_cache_spec.storage_block_size),
+                int(self.kv_cache_spec.num_states),
                 self.compress_ratio,
                 out=self.compressed_slot_mapping_buffer,
             )

vllm/models/inkling/amd/sconv_swa_attn.py

@@ -107,32 +107,6 @@ class InklingSconvBackend(AttentionBackend):
     def get_name() -> str:
         return "INKLING_SCONV_SWA"
 
-    @classmethod
-    def indexes_kv_by_block_stride(cls) -> bool:
-        # num_blocks is the outermost dim (HND, see get_kv_cache_shape), so the
-        # padded conv page is read through a strided view.
-        return True
-
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        # HND, num-blocks-first, head-major: [num_blocks, H, N, D].
-        return (num_blocks, num_kv_heads, block_size, head_size)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        # Identity: physical layout == logical [num_blocks, H, N, D].
-        if include_num_layers_dimension:
-            return (0, 1, 2, 3, 4)
-        return (0, 1, 2, 3)
-
     @staticmethod
     def get_impl_cls():
         raise NotImplementedError(

vllm/models/inkling/nvidia/sconv_swa_attn.py

@@ -107,32 +107,6 @@ class InklingSconvBackend(AttentionBackend):
     def get_name() -> str:
         return "INKLING_SCONV_SWA"
 
-    @classmethod
-    def indexes_kv_by_block_stride(cls) -> bool:
-        # num_blocks is the outermost dim (HND, see get_kv_cache_shape), so the
-        # padded conv page is read through a strided view.
-        return True
-
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        # HND, num-blocks-first, head-major: [num_blocks, H, N, D].
-        return (num_blocks, num_kv_heads, block_size, head_size)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        # Identity: physical layout == logical [num_blocks, H, N, D].
-        if include_num_layers_dimension:
-            return (0, 1, 2, 3, 4)
-        return (0, 1, 2, 3)
-
     @staticmethod
     def get_impl_cls():
         raise NotImplementedError(

vllm/models/kimi_k3/nvidia/mla.py

@@ -384,6 +384,10 @@ class MultiHeadLatentAttention(nn.Module, AttentionLayerBase):
     # ------------------------------------------------------------------
     # AttentionLayerBase interface
     # ------------------------------------------------------------------
+    def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
+        # [B, H=1, N, C] -> [B, N, C]
+        self.kv_cache = kv_cache.squeeze(1)
+
     def get_attn_backend(self) -> type[AttentionBackend]:
         return self.attn_backend
 

vllm/models/minimax_m3/amd/model.py

@@ -728,6 +728,12 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
 
     def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
         # Main GQA K/V cache. Block size may change after load, refresh it.
+        # AITER sparse PA wants K and V as two separate head groups
+        # (H folded into the content dim).
+        sparse_pa = self.use_aiter_sparse_pa
+        kv_bytes = (
+            self.num_kv_heads * self.head_dim * self.kv_cache_torch_dtype.itemsize
+        )
         return FullAttentionSpec(
             block_size=vllm_config.cache_config.block_size,
             num_kv_heads=self.num_kv_heads,
@@ -735,6 +741,8 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):
             head_size_v=self.head_dim,
             dtype=self.kv_cache_torch_dtype,
             kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
+            num_head_slots=2 if sparse_pa else None,
+            state_content_bytes=kv_bytes if sparse_pa else None,
         )
 
     def _ensure_aiter_sparse_pa_kv_cache(self) -> None:

vllm/models/minimax_m3/common/indexer.py

@@ -97,25 +97,6 @@ class MiniMaxM3IndexerBackend(AttentionBackend):
     def is_sparse(cls) -> bool:
         return True
 
-    @staticmethod
-    def get_kv_cache_shape(
-        num_blocks: int,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> tuple[int, ...]:
-        return (num_blocks, block_size, head_size)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        if include_num_layers_dimension:
-            # M3 does not use cross-layer (per-layer-stacked) KV blocks.
-            raise NotImplementedError
-        return (0, 1, 2)
-
 
 class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase):
     """Side KV cache for the indexer's per-token index keys (key-only).
@@ -156,6 +137,10 @@ class MiniMaxM3IndexerCache(nn.Module, AttentionLayerBase):
             raise ValueError(f"Duplicate layer name: {prefix}")
         compilation_config.static_forward_context[prefix] = self
 
+    def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
+        # [B, H=1, N, C] -> [B, N, C]
+        self.kv_cache = kv_cache.squeeze(1)
+
     def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
         # Key-only: MLAAttentionSpec budgets one vector/token (not 2x for K+V).
         return MLAAttentionSpec(