Fix AMD standardized KV cache regressions

Co-authored-by: OpenAI Codex <noreply@openai.com>

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

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
@@ -37,10 +38,13 @@ from vllm.config import (
 )
 from vllm.v1.attention.backends.utils import (
     CommonAttentionMetadata,
-    get_kv_cache_layout,
-    set_kv_cache_layout,
+    resolve_kv_cache_layout,
+)
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    compute_layer_kv_cache_shape_bytes,
+    reshape_kv_cache,
 )
-from vllm.v1.kv_cache_interface import FullAttentionSpec
 
 # ============================================================================
 # Backend Configuration
@@ -337,52 +341,23 @@ def _create_input_tensors(
 def _create_kv_cache(
     config: BenchmarkConfig,
     max_num_blocks: int,
-    backend_class,
     device: torch.device,
     dtype: torch.dtype,
 ) -> list:
-    """Create KV cache tensors for all layers using the backend's methods.
-
-    Uses the backend's get_kv_cache_shape() and get_kv_cache_stride_order()
-    to create the cache with the correct shape and memory layout.
-    """
-    # Get the logical shape from the backend
-    cache_shape = backend_class.get_kv_cache_shape(
-        num_blocks=max_num_blocks,
+    """Create KV cache tensors for all layers using the standard allocator."""
+    spec = FullAttentionSpec(
         block_size=config.block_size,
         num_kv_heads=config.num_kv_heads,
         head_size=config.head_dim,
+        dtype=dtype,
     )
-
-    # 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
+    layout = resolve_kv_cache_layout()
+    total_bytes = (
+        prod(compute_layer_kv_cache_shape_bytes(spec, max_num_blocks))
+        * config.num_layers
+    )
+    buf = torch.zeros(total_bytes, device=device, dtype=torch.int8)
+    return reshape_kv_cache(buf, spec, max_num_blocks, config.num_layers, layout)
 
 
 # ============================================================================
@@ -500,13 +475,6 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
                 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()
-
             common_metadata = _build_common_attn_metadata(
                 q_lens, kv_lens, config.block_size, device
             )
@@ -541,9 +509,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
                 config, total_q, device, dtype, quantize_query=quantize_query
             )
 
-            cache_list = _create_kv_cache(
-                config, max_num_blocks, backend_class, device, dtype
-            )
+            cache_list = _create_kv_cache(config, max_num_blocks, device, dtype)
 
             timing_stats, mem_stats = _run_single_benchmark(
                 config,

docs/features/nixl_connector_compatibility.md

@@ -59,7 +59,7 @@ th:not(:first-child) {
 
 <sup>1</sup> P and D instances must use the same speculation configuration.
 
-<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` or `--enable-cross-layers`).
 
 <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.
 

docs/features/nixl_connector_usage.md

@@ -414,15 +414,6 @@ Support use case: Prefill with 'HND' and decode with 'NHD' with experimental con
 --kv-transfer-config '{..., "enable_permute_local_kv":"True"}'
 ```
 
-### Cross layers blocks
-
-By default, this feature is disabled. On attention backends that support this feature, each logical block is contiguous in physical memory. This reduces the number of buffers that need to be transferred.
-To enable this feature:
-
-```bash
---kv-transfer-config '{..., "kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}'
-```
-
 ## Metrics Reference
 
 vLLM periodically logs a `KV Transfer metrics` line summarising NIXL transfer

tests/compile/passes/test_fusion_attn.py

@@ -39,7 +39,12 @@ 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 resolve_kv_cache_layout
+from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
+    get_kv_quant_mode,
+    reshape_kv_cache,
+)
 
 DEVICE_TYPE = current_platform.device_type
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -108,32 +113,27 @@ 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),
+        )
+        layout = resolve_kv_cache_layout()
+        num_layer_slots = 1 if layout.is_layer_compact else 2
+        raw_tensor = torch.zeros(
+            num_layer_slots * num_blocks * spec.page_size_bytes,
+            dtype=torch.int8,
             device=self.device,
         )
-        kv_cache = raw_tensor.permute(*inv_order)
+        kv_cache = reshape_kv_cache(
+            raw_tensor,
+            spec,
+            num_blocks,
+            num_layer_slots,
+            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,15 @@ 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 uses a 4D KV cache: (num_blocks, num_heads=1, block_size, head_size).
+        kv_cache_shape = (num_blocks, 1, self.block_size, self.head_size)
+        kv_cache = torch.zeros(
+            kv_cache_shape,
             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

@@ -38,7 +38,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, reshape_kv_cache
 
 INDEX_SELECT_OP = torch.ops.aten.index.Tensor
 FP8_DTYPE = current_platform.fp8_dtype()
@@ -128,20 +128,21 @@ 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.kv_cache_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=AttentionSpec(
-                block_size=self.block_size,
-                num_kv_heads=self.num_kv_heads,
-                head_size=head_size,
-                dtype=self.kv_cache_dtype,
-            ),
+            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(
@@ -151,32 +152,22 @@ 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 = reshape_kv_cache(
+            raw_tensor,
+            self.kv_cache_spec,
+            num_blocks,
+            num_layer_slots=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
-        # attribute directly to the fused/unfused `do_kv_cache_update` impls,
-        # which call `kv_cache.unbind(0)` and therefore require a tensor.
+        # `bind_kv_cache` behavior. `get_attention_context` returns this
+        # attribute directly to the fused/unfused cache update implementations.
         self.attn.kv_cache = kv_cache
 
         attn_metadata = self.builder.build(
@@ -253,7 +244,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,
@@ -326,7 +317,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
             }
@@ -341,7 +332,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
             }
@@ -419,10 +410,10 @@ _FUSION_CONFIGS = [
 @pytest.mark.parametrize("num_tokens", [5, 16, 2048])
 @pytest.mark.parametrize("use_shuffle_kv_layout", ["1", "0"])
 @pytest.mark.parametrize(
-    "kv_stride_order",
+    "kv_layout",
     [
-        pytest.param((0, 1, 2, 3, 4), id="block_first"),
-        pytest.param((1, 0, 2, 3, 4), id="kv_first"),
+        pytest.param(KVCacheLayout.LBHNC, id="head_major"),
+        pytest.param(KVCacheLayout.LBNHC, id="token_major"),
     ],
 )
 @pytest.mark.parametrize("enable_aiter_triton_rope", [True, False])
@@ -435,6 +426,7 @@ _FUSION_CONFIGS = [
     not is_aiter_found_and_supported(),
     reason="Only test on ROCm with AITER installed and supported",
 )
+@pytest.mark.skip(reason="AITER fusion does not support packed standardized K/V caches")
 def test_qk_norm_rope_kvcache_fusion(
     num_tokens: int,
     num_heads: int,
@@ -445,7 +437,7 @@ def test_qk_norm_rope_kvcache_fusion(
     attn_backend: AttentionBackendEnum,
     enable_aiter_triton_rope: bool,
     use_shuffle_kv_layout: str,
-    kv_stride_order: tuple[int, ...],
+    kv_layout: KVCacheLayout,
     block_size: int,
     dtype: torch.dtype,
     kv_cache_dtype: str,
@@ -469,7 +461,7 @@ def test_qk_norm_rope_kvcache_fusion(
         block_size=block_size,
         is_neox=is_neox,
         use_shuffle_kv_layout=use_shuffle_kv_layout,
-        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

@@ -37,6 +37,10 @@ from vllm.v1.attention.backend import (
     CommonAttentionMetadata,
 )
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    compute_layer_kv_cache_shape_bytes,
+)
 
 INDEX_SELECT_OP = torch.ops.aten.index.Tensor
 VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update
@@ -136,28 +140,21 @@ 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
+        kv_cache_shape = compute_layer_kv_cache_shape_bytes(
+            FullAttentionSpec(
+                block_size=self.block_size,
+                num_kv_heads=self.num_kv_heads,
+                head_size=self.head_size,
+                dtype=self.kv_cache_dtype,
+            ),
+            num_blocks,
         )
-        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))
-        ]
-
-        # 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,
+        kv_cache = torch.zeros(
+            kv_cache_shape,
+            dtype=torch.int8,
             device=self.device,
-        )
-        raw_tensor = raw_tensor.view(kv_cache_shape)
-        kv_cache = raw_tensor.permute(*inv_order)
+        ).view(self.kv_cache_dtype)
 
         self.attn.kv_cache = kv_cache
 

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

@@ -19,7 +19,7 @@ NUM_LAYERS = [1]  # Arbitrary values for testing
 NUM_HEADS = [8]  # Arbitrary values for testing
 HEAD_SIZES = [64, 80, 256]
 BLOCK_SIZES = [8, 16, 32]
-CACHE_LAYOUTS = ["NHD", "HND"]
+CACHE_LAYOUTS = ["LBNHC", "LBHNC"]
 KV_SCALE_TYPES = ["tensor", "attn_head"]
 
 # Parameters for MLA tests.
@@ -196,8 +196,8 @@ def test_reshape_and_cache_flash(
     torch.set_default_device(device)
     torch.accelerator.set_device_index(device)
     assert implementation in ["cuda", "triton"]
-    if implementation == "triton" and kv_cache_layout == "HND":
-        pytest.skip("Triton implementation only supports NHD layout.")
+    if implementation == "triton" and kv_cache_layout == "LBHNC":
+        pytest.skip("Triton implementation only supports LBNHC layout.")
 
     if kv_scale_type == "attn_head" and implementation != "cuda":
         pytest.skip("Only CUDA implementation supports attn_head scaling.")
@@ -270,7 +270,7 @@ def test_reshape_and_cache_flash(
         v_scale = (value.amax(dim=(0, 2)) / 64.0).to(torch.float32)
 
     def permute_and_compact(x):
-        y = x if kv_cache_layout == "NHD" else x.permute(0, 2, 1, 3)
+        y = x if kv_cache_layout == "LBNHC" else x.permute(0, 2, 1, 3)
         return y.contiguous()
 
     if kv_cache_dtype != "nvfp4":
@@ -284,8 +284,8 @@ def test_reshape_and_cache_flash(
                 fp8_input.flatten(0, 2), scale, group_shape=None, out_dtype=output.dtype
             ).reshape(*input.shape)
         else:  # per-head: broadcast scale along the head dimension
-            # Original code uses dim 2 for NHD, dim 1 for HND
-            if kv_cache_layout == "NHD":
+            # Original code uses dim 2 for LBNHC, dim 1 for LBHNC
+            if kv_cache_layout == "LBNHC":
                 result = fp8_input.to(output.dtype) * scale.view(1, 1, -1, 1)
             else:
                 result = fp8_input.to(output.dtype) * scale.view(1, -1, 1, 1)
@@ -354,28 +354,29 @@ 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.
-            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
+        def dequant_nvfp4_cache_hnc(data_cache, scale_cache, global_scale):
+            # data_cache:  [H, N, T, data_dim]  HNC layout
+            # scale_cache: [H, N, T, scale_dim] HNC layout
+            return dequant_nvfp4_kv_cache(
+                data_cache, scale_cache, global_scale, head_size, block_size
             )
-            return result_hnd.permute(0, 2, 1, 3)  # back to [N, T, H, D]
 
-        result_key_cache = dequant_nvfp4_cache_nhd(
+        result_key_cache = dequant_nvfp4_cache_hnc(
             nvfp4_key_data, key_scale_cache, k_scale.item()
         )
-        result_value_cache = dequant_nvfp4_cache_nhd(
+        result_value_cache = dequant_nvfp4_cache_hnc(
             nvfp4_value_data, value_scale_cache, v_scale.item()
         )
 
-        # Flatten [num_blocks, block_size] → [num_slots] and index by slot_mapping.
+        # Result is HNC: (num_blocks, num_heads, block_size, head_size).
+        # Flatten to (num_slots, num_heads, head_size) for comparison.
         num_slots = num_blocks * block_size
-        result_key_flat = result_key_cache.reshape(num_slots, num_heads, head_size)
-        result_value_flat = result_value_cache.reshape(num_slots, num_heads, head_size)
+        result_key_flat = result_key_cache.permute(0, 2, 1, 3).reshape(
+            num_slots, num_heads, head_size
+        )
+        result_value_flat = result_value_cache.permute(0, 2, 1, 3).reshape(
+            num_slots, num_heads, head_size
+        )
 
         torch.testing.assert_close(
             result_key_flat[slot_mapping], key.float(), atol=1.5, rtol=0.5
@@ -407,7 +408,7 @@ def test_reshape_and_cache_flash(
     for i in range(num_tokens):
         block_idx = block_indices_lst[i]
         block_offset = block_offsets_lst[i]
-        if kv_cache_layout == "NHD":
+        if kv_cache_layout == "LBNHC":
             cloned_key_cache[block_idx, block_offset, :, :] = key[i]
             cloned_value_cache[block_idx, block_offset, :, :] = value[i]
         else:

tests/kernels/attention/test_minimax_m3.py

@@ -6,9 +6,6 @@ import pytest
 import torch
 
 from vllm import _custom_ops as ops
-from vllm.models.minimax_m3.common.indexer import (
-    MiniMaxM3IndexerBackend,
-)
 from vllm.models.minimax_m3.common.ops.index_topk import (
     minimax_m3_index_decode,
     minimax_m3_index_score,
@@ -20,14 +17,21 @@ from vllm.models.minimax_m3.common.ops.sparse_attn import (
     minimax_m3_sparse_attn_decode,
 )
 from vllm.models.minimax_m3.common.sparse_attention import (
-    MiniMaxM3SparseBackend,
     MiniMaxM3SparseTritonImpl,
+    minimax_m3_use_aiter_sparse_pa,
 )
 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 (
+    resolve_kv_cache_layout,
+    set_kv_cache_layout,
+)
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheLayout,
+    MLAAttentionSpec,
+    compute_layer_kv_cache_shape_bytes,
+    reshape_kv_cache,
+)
 
 if not (current_platform.is_cuda() or current_platform.is_rocm()):
     pytest.skip(
@@ -46,26 +50,41 @@ def kv_layout(request):
         set_kv_cache_layout(None)
 
 
-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(ndim: int) -> tuple[int, ...]:
+    """Per-layer physical stride order for the active layout; the 3-dim
+    indexer side cache (H=1) is contiguous, so identity."""
+    if ndim == 3:
+        return (0, 1, 2)
+    stride_order = resolve_kv_cache_layout().layer_stride_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"
 ) -> 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))
+    expose the inverse-permuted logical [B, H, N, C] view the kernels see."""
+    logical_shape = _main_kv_logical_shape(num_pages)
+    stride_order = _layer_stride_order(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)
@@ -897,41 +916,48 @@ 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)
 
+    assert KVCacheLayout.LBHNC.layer_stride_order == (0, 1, 2, 3)
+    assert KVCacheLayout.LBNHC.layer_stride_order == (0, 2, 1, 3)
+
     try:
         set_kv_cache_layout("HND")
-        assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 1, 2, 3)
+        assert resolve_kv_cache_layout() is KVCacheLayout.LBHNC
         set_kv_cache_layout("NHD")
-        assert MiniMaxM3SparseBackend.get_kv_cache_stride_order() == (0, 2, 1, 3)
+        assert resolve_kv_cache_layout() is KVCacheLayout.LBNHC
     finally:
         set_kv_cache_layout(None)
 
     for layout in ("NHD", "HND"):
         try:
             set_kv_cache_layout(layout)
-            order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
+            order = resolve_kv_cache_layout().layer_stride_order
         finally:
             set_kv_cache_layout(None)
         # 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_unknown_layout_raises():
+    """An unrecognized layout override is rejected at resolution time."""
+    try:
+        set_kv_cache_layout("BOGUS")
+        with pytest.raises(ValueError, match="Unknown KV cache layout"):
+            resolve_kv_cache_layout()
+    finally:
+        set_kv_cache_layout(None)
 
 
-def test_aiter_sparse_pa_layout_contract(monkeypatch):
-    """The shuffle-only AITER path retains separately contiguous K/V storage."""
+def test_aiter_sparse_pa_cache_uses_separate_head_groups(monkeypatch):
     import vllm.models.minimax_m3.common.sparse_attention as sparse_attn_mod
 
     monkeypatch.setattr(sparse_attn_mod.rocm_aiter_ops, "is_enabled", lambda: True)
@@ -941,73 +967,63 @@ 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)
+    assert minimax_m3_use_aiter_sparse_pa(1)
+    with pytest.raises(ValueError, match="num_kv_heads == 1"):
+        minimax_m3_use_aiter_sparse_pa(2)
 
-    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)
+    spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=1,
+        head_size=HEAD_DIM,
+        head_size_v=HEAD_DIM,
+        dtype=DTYPE,
+        separate_kv_head_groups=True,
+    )
+    num_blocks = 7
+    raw = torch.empty(num_blocks * spec.page_size_bytes, dtype=torch.int8)
+    kv_cache = reshape_kv_cache(
+        raw,
+        spec,
+        num_blocks,
+        num_layer_slots=1,
+        layout=KVCacheLayout.LBHNC,
+    )[0]
+    assert kv_cache.shape == (num_blocks, 2, BLOCK_SIZE, HEAD_DIM)
+    key_cache, value_cache = kv_cache.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(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 = reshape_kv_cache(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)
+    logical = _main_kv_logical_shape(nb)
     try:
         set_kv_cache_layout("HND")
-        stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
+        stride_order = resolve_kv_cache_layout().layer_stride_order
     finally:
         set_kv_cache_layout(None)
 
@@ -1035,25 +1051,15 @@ 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
-    )
+    logical = _main_kv_logical_shape(4)
     for layout in ("NHD", "HND"):
         try:
             set_kv_cache_layout(layout)
-            order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
+            order = resolve_kv_cache_layout().layer_stride_order
         finally:
             set_kv_cache_layout(None)
         inv_order = [order.index(i) for i in range(len(order))]
@@ -1357,93 +1363,25 @@ 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
-    `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."""
+    """AC-2: drive the real allocator (`reshape_kv_cache`) for the M3 main
+    `FullAttentionSpec` under HND and assert the kernel-visible view has the
+    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], {})
+        layout = resolve_kv_cache_layout()
     finally:
         set_kv_cache_layout(None)
-    view = kv_caches["main"]
+    view = reshape_kv_cache(raw, spec, nb, 1, layout, 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
@@ -1460,7 +1398,7 @@ def test_padded_main_cache_is_flagged():
 
     try:
         set_kv_cache_layout("HND")
-        stride_order = MiniMaxM3SparseBackend.get_kv_cache_stride_order()
+        stride_order = resolve_kv_cache_layout().layer_stride_order
     finally:
         set_kv_cache_layout(None)
 

tests/kernels/attention/test_trtllm_kvfp8_dequant.py

@@ -3,14 +3,14 @@
 """
 Standalone unit tests for trtllm_prefill_attn_kvfp8_dequant.
 
-Tests both contiguous and non-contiguous (cross-layer unified) KV cache
-layouts against a pure-PyTorch reference implementation.
+Tests KV cache layouts against a pure-PyTorch reference implementation.
 """
 
 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 +34,32 @@ 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(
-    num_blocks,
-    num_kv_heads,
-    block_size,
-    head_size,
-    num_layers=4,
+def make_random_kv_cache(
+    num_blocks, num_kv_heads, block_size, head_size, layout=KVCacheLayout.LBHNC
 ):
-    """
-    Create a non-contiguous per-layer view mimicking cross-layer allocation.
+    """Create a random fp8 KV cache in 5D ``(B, 2, H, N, hs)`` format.
 
-    Physical layout: (num_blocks, 2, num_kv_heads, num_layers, block_size, head_size)
-    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).
+    The cache is allocated in the physical 5D layout, then one logical layer
+    is selected and reshaped. Cross-layer layouts therefore retain their
+    inter-layer stride gaps, matching the actual forward path.
     """
-    raw = torch.randn(
+    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)]
+
+    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 +95,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 +103,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 +111,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 +161,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 +208,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 +240,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 +273,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 +306,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 +343,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,
@@ -395,46 +369,3 @@ def test_large_block_size():
     ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
 
     torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)
-
-
-@torch.inference_mode()
-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).
-    """
-    from vllm.v1.attention.backends.flashinfer import (
-        trtllm_prefill_attn_kvfp8_dequant,
-    )
-
-    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(
-        NUM_BLOCKS,
-        num_kv_heads,
-        block_size,
-        head_size,
-        num_layers=num_layers,
-    )
-    k_scale = v_scale = scale.clone()
-
-    block_tables = torch.randint(
-        1,
-        NUM_BLOCKS,
-        (4, 6),
-        dtype=torch.int32,
-        device="cuda",
-    )
-
-    mock_kv_cache, _ = trtllm_prefill_attn_kvfp8_dequant(
-        kv_cache,
-        block_tables,
-        k_scale,
-        v_scale,
-        torch.bfloat16,
-    )
-    ref = ref_dequant(kv_cache, block_tables, k_scale, v_scale, torch.bfloat16)
-
-    torch.testing.assert_close(mock_kv_cache[1:], ref[1:], atol=1e-3, rtol=1e-3)

tests/kernels/attention/test_xpu_mla_sparse.py

@@ -4,6 +4,7 @@
 import pytest
 import torch
 
+from vllm.v1.attention.backends.mla.xpu_mla_sparse import XPUMLASparseImpl
 from vllm.v1.attention.ops.xpu_mla_sparse import triton_bf16_mla_sparse_interface
 
 
@@ -75,6 +76,39 @@ def reference_mla_sparse_prefill(
     return (out.to(kv.dtype), out, max_logits, orig_lse)
 
 
+def test_xpu_sparse_backend_flattens_standard_cache_to_three_dims(monkeypatch):
+    captured = {}
+
+    def fake_sparse_interface(q, kv, indices, sm_scale):
+        captured["kv"] = kv
+        output = torch.zeros(q.shape[0], q.shape[1], 512, dtype=q.dtype)
+        return output, None, None
+
+    monkeypatch.setattr(
+        "vllm.v1.attention.backends.mla.xpu_mla_sparse."
+        "triton_bf16_mla_sparse_interface",
+        fake_sparse_interface,
+    )
+    impl = type("StubImpl", (), {"num_heads": 4, "softmax_scale": 1.0})()
+    q = torch.zeros(2, 4, 576, dtype=torch.bfloat16)
+    kv_cache = (
+        torch.arange(3 * 8 * 576, dtype=torch.int32)
+        .remainder(127)
+        .to(torch.bfloat16)
+        .view(3, 8, 576)
+    )
+    topk_indices = torch.zeros(2, 128, dtype=torch.int32)
+
+    output = XPUMLASparseImpl._forward_bf16_kv(
+        impl, q, kv_cache, topk_indices, attn_metadata=None
+    )
+
+    expected_kv = kv_cache.reshape(24, 1, 576)
+    assert captured["kv"].shape == expected_kv.shape
+    assert torch.equal(captured["kv"], expected_kv)
+    assert output.shape == (2, 4, 512)
+
+
 @pytest.mark.parametrize("device_str", ["xpu"])
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 @pytest.mark.skipif(

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

@@ -32,16 +32,16 @@ from vllm.v1.attention.backend import (
 )
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.attention.backends.utils import (
+    resolve_kv_cache_layout,
     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,
     AttentionBackendEnum.FLASHINFER,
     AttentionBackendEnum.FLEX_ATTENTION,
     AttentionBackendEnum.TRITON_ATTN,
-    "FLEX_ATTENTION_SLOW",
 ]
 
 DEVICE_TYPE = current_platform.device_type
@@ -115,27 +115,19 @@ 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:
     """Create and prepopulate a KV cache with context data.
 
-    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
-        randomize_blocks: Whether to randomly permute blocks
-                          or use sequential order
+    Mirrors production's ``reshape_kv_cache``: allocates a flat buffer in
+    the physical order dictated by *layout*, then permutes to the logical
+    ``[B, H, N, C]`` shape that every backend expects.
 
     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()
@@ -152,41 +144,43 @@ 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)
+    # --- allocate ---------------------------------------------------------
+    # 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
-    # Start from block_id=1 since block_id=0 is considered the null block
-    start_block_idx = 1
+    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
+    # reshape_kv_cache and retains cross-layer strides in the 4D view.
+    kv_cache = kv_cache_physical.permute(*inv_order)[0]
+
+    # --- populate ---------------------------------------------------------
+    # Write context tokens into the cache via the logical view:
+    # kv_cache[block, :, token_in_block, :] routes correctly regardless
+    # of physical layout.
+    start_block_idx = 1  # block 0 is the null block
     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
-
-        # Stay block aligned and allocate enough blocks for the new tokens
+        for t in range(k_context.shape[0]):
+            blk = start_block_idx + t // block_size
+            off = t % block_size
+            kv_cache[blk, :, off, :head_size] = k_context[t]
+            kv_cache[blk, :, off, head_size:] = v_context[t]
         start_block_idx += cdiv(int(seq_lens[i]), block_size)
 
     blocks_end = start_block_idx
 
     # Permute the context blocks (excluding block 0 which is null)
     if randomize_blocks:
-        # Random permutation starting from block 1
         perm = torch.randperm(blocks_end - 1) + 1
     else:
-        # Sequential order starting from block 1
         perm = torch.arange(1, blocks_end)
 
     inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
-    # Add 1 to account for starting from block 1
     inv_perm[1:] = torch.argsort(perm) + 1
     kv_cache[1:blocks_end, ...] = kv_cache[perm, ...]
 
@@ -211,9 +205,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)
 
@@ -250,18 +241,14 @@ 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 = not current_platform.is_rocm() and is_torch_equal_or_newer(
+        "2.9.0.dev0"
+    )
 
-    use_direct_block_mask = is_torch_equal_or_newer("2.9.0.dev0")
-    if backend == "FLEX_ATTENTION_SLOW":
-        actual_backend = AttentionBackendEnum.FLEX_ATTENTION
-        use_direct_block_mask = False
-
-    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
@@ -290,7 +277,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,
@@ -327,7 +314,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
         )
@@ -341,7 +328,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,
@@ -502,6 +489,8 @@ def _test_backend_correctness(
     common_attn_metadata.causal = causal
 
     # 3. Simulate Paged KV Cache and a realistic slot_mapping
+    attn_backends = tuple(backend.get_class() for backend in backend_to_test)
+    layout = resolve_kv_cache_layout(attn_backends)
     kv_cache = create_and_prepopulate_kv_cache(
         k_contexts=k_contexts,
         v_contexts=v_contexts,
@@ -512,6 +501,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,
     )
@@ -520,41 +510,18 @@ 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 = 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)
-                )
+
+        # FlashInfer reads the layout at plan time; override to match
+        # the physical order of the test cache.
+        set_kv_cache_layout(layout.name)
 
         try:
             backend_output = run_attention_backend(
@@ -573,8 +540,7 @@ def _test_backend_correctness(
                 kv_cache_dtype=kv_cache_dtype,
             )
         finally:
-            if reset_kv_cache_layout:
-                set_kv_cache_layout(None)
+            set_kv_cache_layout(None)
 
         # Check shape and dtype consistency
         assert backend_output.shape == sdpa_output.shape, (
@@ -603,6 +569,41 @@ 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
+
+    set_kv_cache_layout(layout)
+    try:
+        _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,
+        )
+    finally:
+        set_kv_cache_layout(None)
+
+
 @pytest.mark.parametrize(
     "batch_spec_name",
     [
@@ -827,14 +828,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",
     ]
 
 
@@ -852,7 +851,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."""
 
@@ -914,7 +916,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."""
 
@@ -950,7 +955,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():
@@ -971,7 +975,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_cpu_kv_cache_layout.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+
+import pytest
+import torch
+
+from vllm.v1.attention.backends.cpu_attn import _split_cpu_kv_cache
+from vllm.v1.attention.backends.utils import set_kv_cache_layout
+from vllm.v1.kv_cache_interface import KVCacheLayout
+
+
+def _make_cache_with_layout(layout: KVCacheLayout) -> torch.Tensor:
+    logical_shape = (2, 3, 2, 4, 10)
+    physical_shape = tuple(logical_shape[i] for i in layout.stride_order)
+    physical = torch.arange(math.prod(logical_shape)).view(physical_shape)
+    inverse_order = tuple(layout.stride_order.index(i) for i in range(5))
+    return physical.permute(*inverse_order)[0]
+
+
+@pytest.mark.parametrize(
+    "layout", [KVCacheLayout.LBHNC, KVCacheLayout.BLHNC, KVCacheLayout.BHLNC]
+)
+def test_split_cpu_kv_cache_supports_hnd_layouts(layout: KVCacheLayout):
+    set_kv_cache_layout(layout.name)
+    try:
+        kv_cache = _make_cache_with_layout(layout)
+        key_cache, value_cache = _split_cpu_kv_cache(kv_cache)
+    finally:
+        set_kv_cache_layout(None)
+
+    assert key_cache.shape == value_cache.shape == (3, 2, 4, 5)
+    assert key_cache.stride() == value_cache.stride()
+    assert key_cache.stride(-2) == 5
+    assert value_cache.storage_offset() - key_cache.storage_offset() == 20
+
+
+def test_split_cpu_kv_cache_rejects_nhd_layout():
+    set_kv_cache_layout(KVCacheLayout.LBNHC.name)
+    try:
+        kv_cache = _make_cache_with_layout(KVCacheLayout.LBNHC)
+        with pytest.raises(ValueError, match="does not support KV cache layout LBNHC"):
+            _split_cpu_kv_cache(kv_cache)
+    finally:
+        set_kv_cache_layout(None)
+
+
+def test_split_cpu_kv_cache_rejects_incompatible_strides():
+    set_kv_cache_layout(KVCacheLayout.LBHNC.name)
+    try:
+        kv_cache = torch.empty(3, 4, 2, 10).transpose(1, 2)
+        with pytest.raises(ValueError, match="contiguous token and content"):
+            _split_cpu_kv_cache(kv_cache)
+    finally:
+        set_kv_cache_layout(None)

tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py

@@ -17,13 +17,13 @@ def test_indexer_builder_deepseek_v4_compressed_slot_mapping_uses_storage_block_
     """
     device = torch.device("cuda")
 
-    # storage_block_size = block_size // compress_ratio = 256 // 4 = 64
+    # storage_block_size = 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)
     builder = DeepseekV32IndexerMetadataBuilder(

tests/v1/attention/test_mla_backends.py

@@ -279,8 +279,9 @@ def create_and_prepopulate_kv_cache(
         else:
             kv_entry_size = head_size
 
+        # Create MLA KV cache: (num_blocks, num_heads=1, block_size, kv_entry_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
@@ -289,9 +290,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)
 
@@ -312,7 +313,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,
@@ -435,7 +436,7 @@ class MockSparseMLAAttentionLayer:
             ops.concat_and_cache_mla(
                 kv_c,
                 k_pe.squeeze(1),
-                kv_cache,
+                kv_cache.squeeze(1),
                 attn_metadata.slot_mapping.flatten(),
                 kv_cache_dtype=kv_cache_dtype,
                 scale=self._k_scale,
@@ -571,7 +572,7 @@ class MockMLAAttentionLayer(MLAAttention):
             ops.concat_and_cache_mla(
                 kv_c,
                 k_pe.squeeze(1),
-                kv_cache,
+                kv_cache.squeeze(1),
                 attn_metadata.slot_mapping.flatten(),
                 kv_cache_dtype=kv_cache_dtype,
                 scale=self._k_scale,

tests/v1/attention/test_rocm_attention_backends_selection.py

@@ -54,7 +54,7 @@ def mock_on_mi3xx():
         (
             {},
             None,
-            AttentionBackendEnum.ROCM_ATTN.get_path(),
+            AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path(),
         ),
         # Test Case 2: Explicit TRITON_ATTN backend
         (
@@ -66,7 +66,7 @@ def mock_on_mi3xx():
         (
             {},
             "ROCM_ATTN",
-            AttentionBackendEnum.ROCM_ATTN.get_path(),
+            None,
         ),
         # Test Case 4: Explicit ROCM_AITER_FA backend
         (
@@ -84,7 +84,7 @@ def mock_on_mi3xx():
         (
             {"VLLM_ROCM_USE_AITER": "1"},
             None,
-            AttentionBackendEnum.ROCM_ATTN.get_path(),
+            AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path(),
         ),
         # Test Case 7: VLLM_ROCM_USE_AITER=1 + explicit TRITON_ATTN
         (
@@ -96,13 +96,13 @@ def mock_on_mi3xx():
         (
             {"VLLM_ROCM_USE_AITER": "1", "VLLM_ROCM_USE_AITER_MHA": "0"},
             None,
-            AttentionBackendEnum.ROCM_ATTN.get_path(),
+            AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path(),
         ),
         # Test Case 9: VLLM_ROCM_USE_AITER=1 + explicit ROCM_ATTN
         (
             {"VLLM_ROCM_USE_AITER": "1"},
             "ROCM_ATTN",
-            AttentionBackendEnum.ROCM_ATTN.get_path(),
+            None,
         ),
     ],
 )
@@ -146,6 +146,16 @@ def test_standard_attention_backend_selection(
         use_sparse=False,
     )
 
+    if expected_backend_path is None:
+        with pytest.raises(
+            ValueError, match="does not support standardized packed KV caches"
+        ):
+            RocmPlatform.get_attn_backend_cls(
+                selected_backend=backend_enum,
+                attn_selector_config=attn_selector_config,
+            )
+        return
+
     backend_path = RocmPlatform.get_attn_backend_cls(
         selected_backend=backend_enum, attn_selector_config=attn_selector_config
     )

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,13 @@ 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,
+    resolve_kv_cache_layout,
     set_kv_cache_layout,
 )
-from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheLayout, KVQuantMode
 
 if not current_platform.is_device_capability_family(100):
     pytest.skip(
@@ -87,96 +87,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,
@@ -189,39 +99,15 @@ def _create_nvfp4_hnd_kv_cache(
     common_attn_metadata,
     kv_scale_val,
 ):
-    """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 an nvfp4 KV cache with 2H head layout.
 
-    The returned tensor is dtype ``uint8`` with head-group layout
-      ``(num_blocks, 2 * num_kv_heads, block_size, full_dim)``
-    where K heads occupy the first ``num_kv_heads`` heads and V heads the second.
-    Each ``full_dim = head_size // 2 + head_size // 16`` block packs two regions:
-      - **FP4 data** (``head_size // 2`` bytes): pairs of E2M1 values,
-        two per byte.
-      - **FP8 block scales** (``head_size // 16`` bytes): one E4M3
-        scale per 16-element block.
-
-    Args:
-        k_contexts: List of key context tensors, one per sequence.
-        v_contexts: List of value context tensors, one per sequence.
-        block_size: Number of tokens per cache block.
-        num_kv_heads: Number of key/value heads.
-        head_size: Head dimension (must be divisible by 16).
-        dtype: Source data type for the bf16 intermediate cache.
-        device: Target device.
-        num_blocks: Total number of blocks to allocate.
-        common_attn_metadata: Metadata containing block tables and
-            sequence lengths.
-        kv_scale_val: Scalar float used as both k_scale and v_scale
-            during quantization.
-
-    Returns:
-        ``torch.Tensor``: The nvfp4 kv_cache tensor (uint8, HND-strided).
+    The returned tensor is dtype ``uint8`` with logical shape
+    ``(num_blocks, 2 * num_kv_heads, block_size, full_dim)`` where K occupies
+    the first H heads and V occupies the next H heads, and
+    ``full_dim = head_size // 2 + head_size // 16`` packs FP4 data and
+    FP8 block scales per head.
     """
-    # 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(
+    bf16_cache = create_and_prepopulate_kv_cache(
         k_contexts,
         v_contexts,
         block_size,
@@ -231,10 +117,9 @@ 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
     full_dim = nvfp4_kv_cache_full_dim(head_size)
     nvfp4_cache = torch.zeros(
         (num_blocks, 2 * num_kv_heads, block_size, full_dim),
@@ -243,8 +128,6 @@ 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.
     block_table = common_attn_metadata.block_table_tensor
     seq_lens = common_attn_metadata.seq_lens.cpu()
     query_lens = (
@@ -257,19 +140,24 @@ def _create_nvfp4_hnd_kv_cache(
         ctx_len = int(seq_lens[i]) - int(query_lens[i])
         if ctx_len == 0:
             continue
-        # 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]
-        # Build slot mapping for these context tokens.
+        # bf16_cache is (B, H, N, 2*head_size); extract K and V from last dim.
+        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]
+        )
         token_offsets = torch.arange(ctx_len, device=device)
         block_indices = token_offsets // block_size
         intra_offsets = token_offsets % block_size
         slots = block_table[i, block_indices] * block_size + intra_offsets
-        # reshape_and_cache_flash expects (B, N, H, D) cache views.
+        # reshape_and_cache_flash expects (B, N, H, D) cache views
         torch.ops._C_cache_ops.reshape_and_cache_flash(
             k_ctx,
             v_ctx,
@@ -363,7 +251,7 @@ def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL)
 
     common_attn_metadata = create_common_attn_metadata(batch_spec, BLOCK_SIZE, device)
 
-    # 2. Create HND KV cache
+    # 2. Create HNC KV cache
     is_nvfp4 = kv_cache_dtype == "nvfp4"
     if is_nvfp4:
         # Compute a global scale from the context data.
@@ -383,7 +271,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,11 +281,12 @@ 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()
+    set_kv_cache_layout("LBHNC")
+    resolve_kv_cache_layout.cache_clear()
 
     try:
         is_nvfp4 = kv_cache_dtype == "nvfp4"
@@ -513,7 +402,7 @@ def _run_trtllm_integration(batch_spec, kv_cache_dtype="auto", model_name=MODEL)
 
     finally:
         set_kv_cache_layout(None)
-        get_kv_cache_layout.cache_clear()
+        resolve_kv_cache_layout.cache_clear()
 
 
 @pytest.mark.parametrize(

tests/v1/core/test_contiguous_kv_packing.py

@@ -1,231 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for contiguous KV cache packing."""
-
-from unittest.mock import MagicMock
-
-import pytest
-import torch
-
-from vllm.v1.core.kv_cache_utils import (
-    _get_kv_cache_config_packed,
-    get_kv_cache_config_from_groups,
-)
-from vllm.v1.kv_cache_interface import (
-    FullAttentionSpec,
-    KVCacheGroupSpec,
-    KVCacheTensor,
-    MLAAttentionSpec,
-    SlidingWindowSpec,
-    UniformTypeKVCacheSpecs,
-)
-
-
-def _make_mla_spec(page_size: int, block_size: int = 256) -> 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,
-    )
-
-
-def _make_full_spec() -> FullAttentionSpec:
-    return FullAttentionSpec(
-        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 _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 _mock_vllm_config(kv_connector_extra_config: dict[str, str] | None = None):
-    config = MagicMock()
-    config.cache_config.num_gpu_blocks_override = None
-    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
-    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
-
-
-class TestInterleavedPacking:
-    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):
-        _, tensors = _run()
-        for t in tensors:
-            assert t.offset < t.block_stride
-
-    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_strided_views_are_independent(self):
-        groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
-        page_sizes = _page_sizes_by_layer(groups)
-        num_blocks, tensors = _get_kv_cache_config_packed(
-            _mock_vllm_config(), groups, 100 * 1024 * 1024
-        )
-        backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
-        views = []
-        for t in tensors:
-            page_size = page_sizes[t.shared_by[0]]
-            v = torch.as_strided(
-                backing,
-                size=(num_blocks, page_size),
-                stride=(t.block_stride, 1),
-                storage_offset=t.offset,
-            )
-            views.append(v)
-
-        for i, v in enumerate(views):
-            v.fill_(i + 1)
-
-        for i, v in enumerate(views):
-            assert (v == i + 1).all(), f"View {i} was corrupted"
-
-    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),
-        ]
-
-        config = get_kv_cache_config_from_groups(
-            _mock_vllm_config(), groups, available_memory=page_size * 2 * 32
-        )
-
-        assert config.num_blocks == 32
-        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),
-        ]
-
-        config = get_kv_cache_config_from_groups(
-            _mock_vllm_config({"enable_cross_layers_blocks": "True"}),
-            groups,
-            available_memory=page_size * 2 * 32,
-        )
-
-        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,
-            ),
-        ]
-
-    def test_single_group_attention_keeps_unpacked_layout(self):
-        spec = _make_full_spec()
-        groups = [KVCacheGroupSpec(["full.0", "full.1"], spec)]
-
-        config = get_kv_cache_config_from_groups(
-            _mock_vllm_config(), groups, available_memory=spec.page_size_bytes * 2 * 32
-        )
-
-        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]
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-v"])

tests/v1/core/test_kv_cache_utils.py

@@ -835,36 +835,19 @@ 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,
+                shared_by=[["layer1"], ["layer2"]],
+            ),
+        ],
+        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.
@@ -876,36 +859,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 = [
@@ -932,7 +886,8 @@ 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,
+                    shared_by=[["layer1"]],
                 ),
             ],
             kv_cache_groups=[
@@ -943,10 +898,8 @@ 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,
+                    shared_by=[["layer2"], ["layer3"]],
                 ),
             ],
             kv_cache_groups=[
@@ -976,64 +929,37 @@ def test_get_kv_cache_configs_multiple_workers():
     kv_cache_configs = get_kv_cache_configs(
         vllm_config,
         tp_pp_kv_cache_specs,
-        [
-            ref_kv_cache_spec.page_size_bytes * 2 * 10,
-            ref_kv_cache_spec.page_size_bytes * 2 * 10,
-            ref_kv_cache_spec.page_size_bytes * 2 * 10,
-            ref_kv_cache_spec.page_size_bytes * 2 * 10,
+        [ref_kv_cache_spec.page_size_bytes * 2 * 10] * 4,
+    )
+    expected_12 = KVCacheConfig(
+        num_blocks=10,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=ref_kv_cache_spec.page_size_bytes * 10 * 2,
+                shared_by=[["layer1"], ["layer2"]],
+            ),
+        ],
+        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,
+                shared_by=[["layer3"]],
+            ),
+        ],
+        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
@@ -1061,10 +987,8 @@ 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,
+                    shared_by=[["layer1"], ["layer2"]],
                 ),
             ],
             kv_cache_groups=[
@@ -1076,10 +1000,8 @@ 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,
+                    shared_by=[["layer3"], ["layer4"]],
                 ),
             ],
             kv_cache_groups=[
@@ -1106,10 +1028,7 @@ def test_get_kv_cache_configs_multiple_workers():
     kv_cache_configs = get_kv_cache_configs(
         vllm_config,
         different_type_layer_specs,
-        [
-            ref_kv_cache_spec.page_size_bytes * 10,
-            ref_kv_cache_spec.page_size_bytes * 10,
-        ],
+        [ref_kv_cache_spec.page_size_bytes * 10] * 2,
     )
     assert kv_cache_configs == [
         KVCacheConfig(
@@ -1117,7 +1036,7 @@ 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"],
+                    shared_by=[["layer1", "layer2", "layer3"]],
                 ),
             ],
             kv_cache_groups=[
@@ -1131,7 +1050,7 @@ 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"],
+                    shared_by=[["layer4", "layer5", "layer6"]],
                 ),
             ],
             kv_cache_groups=[
@@ -1198,7 +1117,7 @@ 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"],
+                    shared_by=[["layer1"]],
                 ),
             ],
             kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)],
@@ -1208,7 +1127,7 @@ 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"],
+                    shared_by=[["layer2"]],
                 ),
             ],
             kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)],
@@ -1536,78 +1455,13 @@ 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, shared_by=[["layer1"]]),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(["layer1"], new_kv_cache_spec(block_size=block_size)),
@@ -1682,11 +1536,14 @@ def test_get_kv_cache_config_one_worker():
         vllm_config, [kv_cache_specs_full], [mem_per_block_per_layer * 2 * 32]
     )[0]
     print(kv_cache_config_full)
+
     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,
+                shared_by=[["layer_1"], ["layer_2"]],
+            ),
         ],
         kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())],
     )
@@ -1702,8 +1559,10 @@ 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,
+                shared_by=[["layer_1"], ["layer_2"]],
+            ),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(["layer_1", "layer_2"], new_sliding_window_spec())
@@ -1722,8 +1581,10 @@ 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,
+                shared_by=[["layer_1"], ["layer_2"]],
+            ),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(
@@ -1745,7 +1606,8 @@ 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,
+                shared_by=[["layer_1", "layer_2"]],
             ),
         ],
         kv_cache_groups=[
@@ -1770,12 +1632,11 @@ 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"],
-            ),
-            KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_2", "layer_5", "layer_6"],
+                size=mem_per_block_per_layer * 32 * 2,
+                shared_by=[
+                    ["layer_1", "layer_3", "layer_4"],
+                    ["layer_2", "layer_5", "layer_6"],
+                ],
             ),
         ],
         kv_cache_groups=[
@@ -1805,15 +1666,12 @@ 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"],
-            ),
-            KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_2", "layer_7", "layer_8", "layer_9"],
-            ),
-            KVCacheTensor(
-                size=mem_per_block_per_layer * 32, shared_by=["layer_3", "layer_10"]
+                size=mem_per_block_per_layer * 32 * 3,
+                shared_by=[
+                    ["layer_1", "layer_4", "layer_5", "layer_6"],
+                    ["layer_2", "layer_7", "layer_8", "layer_9"],
+                    ["layer_3", "layer_10"],
+                ],
             ),
         ],
         kv_cache_groups=[
@@ -1826,8 +1684,7 @@ def test_get_kv_cache_config_one_worker():
         ],
     )
 
-    # 6 full + 5 sliding, pad to 6 full + 6 sliding. This is a typical case for gpt-oss
-    # eagle where there is only one more full attention layer than sliding window layers
+    # 6 full + 5 sliding
     kv_cache_specs_hybrid = {
         "layer_1": new_kv_cache_spec(),
         "layer_2": new_kv_cache_spec(),
@@ -1845,33 +1702,19 @@ 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 * 6 * 32]
     )[0]
-    print(kv_cache_config_hybrid)
     assert kv_cache_config_hybrid == KVCacheConfig(
         num_blocks=32,
         kv_cache_tensors=[
             KVCacheTensor(
-                size=mem_per_block_per_layer * 32,
-                shared_by=["layer_1", "layer_7"],
-            ),
-            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,
+                shared_by=[
+                    ["layer_1", "layer_7"],
+                    ["layer_2", "layer_8"],
+                    ["layer_3", "layer_9"],
+                    ["layer_4", "layer_10"],
+                    ["layer_5", "layer_11"],
+                    ["layer_6"],
+                ],
             ),
         ],
         kv_cache_groups=[
@@ -1897,8 +1740,14 @@ 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 * 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 * 2,
+                shared_by=[["layer_1"]],
+            ),
+            KVCacheTensor(
+                size=mem_per_block_per_layer * 32,
+                shared_by=[["layer_2"]],
+            ),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(
@@ -1922,7 +1771,8 @@ 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,
+                shared_by=[["layer_1", "layer_2"]],
             ),
         ],
         kv_cache_groups=[
@@ -1948,7 +1798,10 @@ 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,
+                shared_by=[["layer_1", "layer_2"]],
+            ),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(
@@ -1974,8 +1827,10 @@ 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,
+                shared_by=[["layer_1"], ["layer_2"]],
+            ),
         ],
         kv_cache_groups=[KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())],
     )

tests/v1/core/test_prefix_caching.py

@@ -148,7 +148,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,),
         )
 
@@ -183,7 +183,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":
@@ -762,12 +762,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",
         ),
@@ -3244,7 +3244,7 @@ def test_hybrid_local_kv_retention_interval_survives_recycling(monkeypatch):
                     num_kv_heads=1,
                     head_size=1,
                     dtype=torch.uint8,
-                    compress_ratio=4,
+                    tokens_per_state=4,
                 ),
             ),
             KVCacheGroupSpec(

tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

@@ -94,7 +94,6 @@ else
   echo "running with default attention backend"
 fi
 
-# Check if cross-layers is enabled (non-empty)
 if [[ -n "${CROSS_LAYERS_BLOCKS:-}" ]]; then
   echo "CROSS_LAYERS_BLOCKS is set, running with --enable-cross-layers"
   label+=" - CROSS_LAYERS_BLOCKS enabled"

tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh

@@ -4,7 +4,11 @@ set -xe
 # Parse command line arguments
 KV_BUFFER_DEVICE="cuda"  # Default to cuda
 ATTENTION_BACKEND=""  # Default to empty (use vllm default)
-CROSS_LAYERS_BLOCKS="False"
+ENABLE_HMA_VAR=""  # Default to empty (HMA disabled by default for kv connector)
+# Check for ENABLE_HMA_FLAG environment variable
+if [[ -n "${ENABLE_HMA_FLAG:-}" ]]; then
+  ENABLE_HMA_VAR="--no-disable-hybrid-kv-cache-manager"
+fi
 
 while [[ $# -gt 0 ]]; do
   case $1 in
@@ -17,12 +21,12 @@ while [[ $# -gt 0 ]]; do
       shift 2
       ;;
     --enable-cross-layers)
-      CROSS_LAYERS_BLOCKS="True"
+      export VLLM_KV_CACHE_LAYOUT="BLHNC"
       shift 1
       ;;
     *)
       echo "Unknown option $1"
-      echo "Usage: $0 [--kv_buffer_device <cuda|cpu>] [--attention-backend <backend>]"
+      echo "Usage: $0 [--kv_buffer_device <cuda|cpu>] [--attention-backend <backend>] [--enable-cross-layers]"
       exit 1
       ;;
   esac
@@ -36,29 +40,24 @@ 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
-if [[ "$DECODER_KV_LAYOUT" == "NHD" ]]; then
+PREFILLER_KV_LAYOUT=${VLLM_KV_CACHE_LAYOUT:-"LBHNC"}
+DECODER_KV_LAYOUT=${DECODER_KV_LAYOUT:-"$PREFILLER_KV_LAYOUT"}
+if [[ "$DECODER_KV_LAYOUT" == "LBNHC" ]]; then
   KV_CONFIG_HETERO_LAYOUT=',"enable_permute_local_kv":"True"'
 else
   KV_CONFIG_HETERO_LAYOUT=''
 fi
 
-if [[ "$CROSS_LAYERS_BLOCKS" == "True" ]]; then
-  KV_EXTRA_CONFIG=',"kv_connector_extra_config":{"enable_cross_layers_blocks": "True"}'
-else
-  KV_EXTRA_CONFIG=''
-fi
-
 # Connector: default pull NixlConnector; NixlPushConnector enables PP prefill.
 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
@@ -165,7 +164,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' \
+    VLLM_KV_CACHE_LAYOUT='$PREFILLER_KV_LAYOUT' \
     VLLM_PORT=$INTERNAL_PORT \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \

tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh

@@ -8,9 +8,9 @@
 # wrapping NixlConnector and OffloadingConnector, then runs gsm8k accuracy via
 # 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)
+# Runs two configurations:
+#   1. Standard 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
@@ -44,8 +44,6 @@ SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "")
 
 # ── 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",
@@ -60,21 +58,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_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 ──────────────────────────────────────────────────────────────
 
 trap 'kill $(jobs -pr) 2>/dev/null' SIGINT SIGTERM EXIT
@@ -125,7 +108,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='${VLLM_KV_CACHE_LAYOUT:-LBHNC}' \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_SIDE_CHANNEL_PORT \
     vllm serve $model_name \
@@ -150,7 +133,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='${VLLM_KV_CACHE_LAYOUT:-LBHNC}' \
     UCX_NET_DEVICES=all \
     VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_SIDE_CHANNEL_PORT \
     vllm serve $model_name \
@@ -207,7 +190,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

@@ -247,7 +247,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 \
@@ -286,7 +286,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_worker.py

@@ -9,7 +9,6 @@ import torch
 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,
@@ -60,31 +59,25 @@ def _allocate_and_reshape_kv_caches(
     Use the real GPUModelRunner allocation and reshape methods to produce
     kv_caches, just like the model runner does during initialization.
     """
+    from vllm.v1.kv_cache_interface import KVCacheLayout
     from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
-    # 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 = []
+    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 runner._allocate_and_reshape_kv_cache(
+        kv_cache_config, kernel_block_sizes, layout=KVCacheLayout.LBNHC
+    )
 
 
 def _make_worker(kv_cache_config: KVCacheConfig):
@@ -211,18 +204,19 @@ def test_register_kv_caches(backend):
         aligned_mamba_layer_names,
     ]
 
-    kv_cache_tensors: list[KVCacheTensor] = []
+    shared_by: list[list[str]] = []
     for i in range(GROUP_SIZE):
-        shared_by: list[str] = []
+        slot_layers: 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,
-            )
+                slot_layers.append(group_layer_names[i])
+        shared_by.append(slot_layers)
+    kv_cache_tensors: list[KVCacheTensor] = [
+        KVCacheTensor(
+            size=PAGE_SIZE_BYTES * NUM_BLOCKS * GROUP_SIZE,
+            shared_by=shared_by,
         )
+    ]
 
     kv_cache_groups = [
         KVCacheGroupSpec(layer_names=attn_layer_names, kv_cache_spec=attn_spec),
@@ -381,11 +375,11 @@ def test_register_kv_caches_uniform_type(backend):
         kv_cache_tensors=[
             KVCacheTensor(
                 size=spec_a.page_size_bytes * NUM_BLOCKS,
-                shared_by=[layer_a],
+                shared_by=[[layer_a]],
             ),
             KVCacheTensor(
                 size=spec_b.page_size_bytes * NUM_BLOCKS,
-                shared_by=[layer_b],
+                shared_by=[[layer_b]],
             ),
         ],
         kv_cache_groups=[

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

@@ -275,7 +275,7 @@ class RequestRunner:
         )
 
         # register worker kv_caches to enable OffloadingWorker creations
-        # set_current_vllm_config is needed for get_kv_cache_layout() to work
+        # set_current_vllm_config is needed for resolve_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 +1,79 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for reshape_kv_cache."""
 
 import pytest
+import torch
 
-from vllm.platforms import current_platform
+from vllm.v1.attention.backends.utils import (
+    get_flashinfer_layout_string,
+    set_kv_cache_layout,
+)
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheLayout,
+    compute_layer_kv_cache_shape_bytes,
+    reshape_kv_cache,
+)
 
-
-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)
+NUM_BLOCKS = 4
+BLOCK_SIZE = 4
+NUM_KV_HEADS = 2
+HEAD_SIZE = 8
+DTYPE = torch.bfloat16
 
 
 @pytest.mark.parametrize(
-    "backend_path",
-    # See: https://github.com/vllm-project/vllm/issues/46411
+    ("layout", "expected"),
     [
-        "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",
+        ("LBHNC", "HND"),
+        ("LBNHC", "NHD"),
+        ("BLHNC", "HND"),
+        ("BLNHC", "NHD"),
+        ("BHLNC", "HND"),
     ],
 )
-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_flashinfer_layout_string(layout: str, expected: str):
+    set_kv_cache_layout(layout)
+    try:
+        assert get_flashinfer_layout_string() == expected
+    finally:
+        set_kv_cache_layout(None)
 
 
-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,
+@pytest.mark.parametrize("layout", list(KVCacheLayout))
+def test_reshape_kv_cache(layout):
+    spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=NUM_KV_HEADS,
+        head_size=HEAD_SIZE,
+        dtype=DTYPE,
     )
+    num_slots = 2
+    total_bytes = spec.page_size_bytes * NUM_BLOCKS * num_slots
+    raw = torch.zeros(total_bytes, dtype=torch.int8, device="cuda")
+    views = reshape_kv_cache(raw, spec, NUM_BLOCKS, num_slots, layout)
 
-    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)
+    byte_4d = compute_layer_kv_cache_shape_bytes(spec, NUM_BLOCKS)
+    dtype_size = torch.tensor([], dtype=spec.dtype).element_size()
+    expected_shape = (*byte_4d[:3], byte_4d[3] // dtype_size)
+
+    assert len(views) == num_slots
+    for v in views:
+        assert v.shape == expected_shape
+        assert v.dtype == spec.dtype
+
+    # The per-layer view preserves the physical order of B, H, N, and C
+    # after the layer dimension is selected. Dimensions later in that order
+    # have smaller strides, including when layer interleaving creates gaps.
+    stride_order = layout.layer_view_order
+    strides = views[0].stride()
+    for i in range(3):
+        for j in range(i + 1, 4):
+            if stride_order[i] < stride_order[j]:
+                assert strides[i] >= strides[j], (
+                    f"layout {layout.name}: dim {i} (physical pos "
+                    f"{stride_order[i]}) should have >= stride than "
+                    f"dim {j} (physical pos {stride_order[j]}), got "
+                    f"strides={strides}"
+                )

tests/v1/kv_connector/unit/test_mooncake_connector.py

@@ -32,17 +32,25 @@ 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.attention.backends.utils import set_kv_cache_layout
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
+    KVCacheLayout,
+    reshape_kv_cache,
 )
 from vllm.v1.request import RequestStatus
 
 from .utils import create_request, create_scheduler, create_vllm_config
 
 
+@pytest.fixture(autouse=True)
+def reset_kv_cache_layout():
+    yield
+    set_kv_cache_layout(None)
+
+
 def _make_test_kv_cache_config() -> KVCacheConfig:
     return KVCacheConfig(
         num_blocks=0,
@@ -1093,13 +1101,23 @@ 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.LBHNC, True),
+        (KVCacheLayout.BHLNC, 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(
         kv_connector="MooncakeConnector", kv_role="kv_consumer"
     )
 
+    set_kv_cache_layout(layout.name)
     with (
         set_current_vllm_config(vllm_config),
         patch_worker_dependencies(),
@@ -1118,15 +1136,22 @@ 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,
+            separate_kv_head_groups=separate_kv_head_groups,
         )
-        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 = reshape_kv_cache(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 +1160,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 separate_kv_head_groups:
+                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

@@ -30,7 +30,9 @@ from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
+    KVCacheLayout,
     MambaSpec,
+    reshape_kv_cache,
 )
 
 from .test_mooncake_connector import patch_worker_dependencies
@@ -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,) = reshape_kv_cache(
+            fa_raw, fa_spec, num_blocks, 1, KVCacheLayout.LBHNC
+        )
+        (gdn_cache,) = reshape_kv_cache(
+            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,7 @@ 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, shared_by=[["layer0"]])],
         kv_cache_groups=[KVCacheGroupSpec(["layer0"], spec)],
     )
 
@@ -210,64 +210,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

@@ -159,8 +159,8 @@ 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=8192, shared_by=[["L0"]]),
+            KVCacheTensor(size=8192, shared_by=[["L1"]]),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(["L0"], full),

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

@@ -34,7 +34,20 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import (
     MooncakeStoreConnectorStats,
 )
+from vllm.v1.attention.backends.utils import set_kv_cache_layout
 from vllm.v1.core.kv_cache_utils import BlockHash
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheGroupSpec,
+    KVCacheLayout,
+    reshape_kv_cache,
+)
+
+
+@pytest.fixture(autouse=True)
+def reset_kv_cache_layout():
+    yield
+    set_kv_cache_layout(None)
 
 
 class _RecordingBlockHashes:
@@ -1818,112 +1831,72 @@ def test_lookup_applies_swa_mask_before_accessing_hashes():
 # ---------------------------------------------------------------------------
 
 
-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 = reshape_kv_cache(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)
+    set_kv_cache_layout(layout.name)
+    _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 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)
+
+
+@pytest.mark.parametrize("layout", list(KVCacheLayout))
+def test_register_kv_caches_separate_head_groups(layout: KVCacheLayout):
+    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,
+        separate_kv_head_groups=True,
+    )
+    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 = reshape_kv_cache(raw, spec, num_blocks, num_layers, layout)
+
+    set_kv_cache_layout(layout.name)
+    _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
@@ -58,7 +59,9 @@ def _make_test_kv_cache_config() -> KVCacheConfig:
     layer_names = ["layer0", "layer1", "layer2"]
     return KVCacheConfig(
         num_blocks=2,
-        kv_cache_tensors=[KVCacheTensor(size=0, shared_by=layer_names)],
+        kv_cache_tensors=[
+            KVCacheTensor(size=0, shared_by=[[name] for name in layer_names])
+        ],
         kv_cache_groups=[
             KVCacheGroupSpec(
                 layer_names=layer_names,
@@ -225,7 +228,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,
     ):
@@ -525,16 +528,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,
@@ -621,16 +623,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_multi_connector.py

@@ -869,16 +869,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

@@ -53,7 +53,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
@@ -62,12 +61,12 @@ from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
-    KVCacheTensor,
+    KVCacheLayout,
+    compute_layer_kv_cache_shape_bytes,
+    reshape_kv_cache,
 )
 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,
@@ -100,6 +99,12 @@ def clear_kv_transfer():
     set_kv_cache_layout(None)
 
 
+@pytest.fixture(autouse=True)
+def reset_kv_cache_layout():
+    yield
+    set_kv_cache_layout(None)
+
+
 def get_default_xfer_telemetry(
     xferDurationS: float = 1,
     postDurationS: float = 1,
@@ -350,8 +355,8 @@ 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
 
+    set_kv_cache_layout("BLHNC")
     # 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.
@@ -386,18 +391,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 = reshape_kv_cache(
+            raw,
+            kv_cache_spec,
+            kv_cache_config.num_blocks,
+            num_layer_slots=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)
 
@@ -471,7 +477,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,
     ):
@@ -482,9 +488,8 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
         self.kv_cache_layout = kv_cache_layout
         # Mock register_kv_caches attribute needed for tests that do not call it.
         self.src_xfer_handles_by_block_size = {self.block_size: 1}
-        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,
@@ -498,7 +503,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(
@@ -549,9 +554,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,
@@ -600,7 +604,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
@@ -996,7 +1000,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,
@@ -1004,6 +1010,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),
@@ -1043,7 +1050,7 @@ class TestNixlHandshake:
                 vllm_config,
                 connector.engine_id,
                 hand_shake_latency=0,
-                kv_cache_layout="NHD",
+                kv_cache_layout="LBNHC",
             )
             worker = connector.connector_worker
 
@@ -1054,15 +1061,16 @@ class TestNixlHandshake:
             worker.dst_num_blocks[worker.engine_id] = worker.num_blocks
 
             # Metadata with different kv_cache_layout than local worker
+            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,
@@ -1073,65 +1081,6 @@ class TestNixlHandshake:
             # whole block is moved.
             worker.add_remote_agent(meta, remote_tp_size=1)
 
-    @patch(
-        "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
-        FakeNixlWrapper,
-    )
-    def test_hybrid_mamba_attention_remote_descs_use_packed_head_slices(
-        self, default_vllm_config, dist_init
-    ):
-        worker = FakeNixlConnectorWorker(
-            create_vllm_config(), "engine", hand_shake_latency=0
-        )
-
-        remote_block_len = 2048
-        local_block_len = remote_block_len // 2
-        worker.block_len_per_layer = [local_block_len]
-        worker._region_is_mla = [False]
-        worker.num_blocks = 1
-        worker.num_regions = 1
-        worker._has_mamba = True
-        worker._mamba_ssm_size = (128, 256)
-        worker.transfer_topo = TransferTopology(
-            tp_rank=1,
-            tp_size=2,
-            block_size=worker.block_size,
-            engine_id=worker.engine_id,
-            is_mla=False,
-            is_mamba=True,
-            total_num_kv_heads=2,
-            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,
-        )
-        meta = MagicMock(
-            kv_caches_base_addr=[0x1000],
-            device_id=0,
-            num_blocks=1,
-            block_lens=[remote_block_len],
-        )
-
-        assert worker.get_backend_aware_kv_block_len(0, mamba_view=False) == (
-            local_block_len
-        )
-        assert (
-            worker.get_backend_aware_kv_block_len(0, first_split=True, mamba_view=True)
-            == worker._mamba_ssm_size[0]
-        )
-        assert (
-            worker.get_backend_aware_kv_block_len(0, first_split=False, mamba_view=True)
-            == worker._mamba_ssm_size[1]
-        )
-
-        assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [
-            [0x1000 + local_block_len, local_block_len, 0]
-        ]
-
     @patch(
         "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
         FakeNixlWrapper,
@@ -1181,6 +1130,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),
@@ -1207,6 +1157,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),
@@ -1250,7 +1201,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]
@@ -1266,7 +1217,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,
@@ -1305,7 +1257,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]
@@ -1323,7 +1275,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,
@@ -1747,7 +1700,6 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
     llm.llm_engine.engine_core.shutdown()
 
 
-@pytest.mark.parametrize("enable_cross_layers", ["False", "True"])
 @pytest.mark.parametrize(
     "attn_backend",
     [
@@ -1761,8 +1713,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
@@ -1776,12 +1734,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")
+    set_kv_cache_layout(layout)
 
     # Import the appropriate backend based on the parameter
     if attn_backend == "FLASH_ATTN":
@@ -1798,59 +1751,34 @@ 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
+        block_size = 16
         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,
+            separate_kv_head_groups=separate_kv_head_groups,
+        )
+        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(
@@ -1870,93 +1798,60 @@ 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
+        raw1 = torch.zeros(
+            kv_cache_spec.page_size_bytes * kv_cache_config.num_blocks,
+            dtype=torch.int8,
+            device=current_platform.device_type,
+        )
+        tensor0, tensor1 = reshape_kv_cache(
+            raw0,
+            kv_cache_spec,
+            kv_cache_config.num_blocks,
+            num_layer_slots=2,
+            layout=KVCacheLayout[layout],
+        )
+        (tensor2,) = reshape_kv_cache(
+            raw1,
+            kv_cache_spec,
+            kv_cache_config.num_blocks,
+            num_layer_slots=1,
+            layout=KVCacheLayout[layout],
+        )
+        kv_caches = {
+            "layer0": tensor0,
+            "layer1": tensor1,
+            "layer2": tensor2,
+            "layer3": tensor0,
+        }
 
-        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()
-            )
+        if separate_kv_head_groups:
             expected_base_addrs = [
-                cross_layers_kv_cache.data_ptr(),
+                cache[:, head_idx].data_ptr()
+                for cache in (tensor0, tensor1, tensor2)
+                for head_idx in range(cache.shape[1])
             ]
-            expected_num_entries = 1
-
-            expected_blocks_count = num_blocks
-
-            kv_caches = {"all-layers": cross_layers_kv_cache}
+            expected_block_len = block_size * head_size * torch.float16.itemsize
+            expected_blocks_count = num_blocks * len(expected_base_addrs)
+        elif layout in ("LBHNC", "LBNHC", "BLHNC"):
+            expected_base_addrs = [
+                tensor0.data_ptr(),
+                tensor1.data_ptr(),
+                tensor2.data_ptr(),
+            ]
+            expected_block_len = kv_cache_spec.page_size_bytes
+            expected_blocks_count = kv_cache_config.num_blocks * 3
         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,
+            expected_base_addrs = [raw0.data_ptr(), raw1.data_ptr()]
+            expected_block_len = {
+                raw0.nbytes // num_blocks,
+                raw1.nbytes // num_blocks,
             }
-
-            # 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
+            expected_blocks_count = kv_cache_config.num_blocks * 2
 
         # Execute register_kv_caches
         connector.register_kv_caches(kv_caches)
@@ -1964,17 +1859,12 @@ 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
@@ -1985,19 +1875,19 @@ def test_register_kv_caches(
             f"Expected {expected_blocks_count} blocks, got {len(blocks_data)}"
         )
 
-        if connector.prefer_cross_layer_blocks:
-            num_blocks = 8
-        else:
-            num_blocks = kv_cache_config.num_blocks
-
-        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}"
-            )
+            if isinstance(expected_block_len, set):
+                assert block_len in expected_block_len
+            else:
+                assert block_len == expected_block_len
+
+        assert (
+            connector.connector_worker.kv_caches_base_addr[
+                connector.connector_worker.engine_id
+            ][0]
+            == expected_base_addrs
+        )
 
         assert connector.connector_worker.block_size == 16
 
@@ -2226,6 +2116,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,
@@ -2820,14 +2712,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 = [
@@ -2862,18 +2751,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,
@@ -2950,9 +2840,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,
@@ -2968,7 +2857,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_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_offload/test_factory.py

@@ -169,9 +169,7 @@ def _make_kv_cache_config():
     head_size = 1
     dtype = torch.float32
     page_size = 2 * num_kv_heads * head_size * torch.finfo(dtype).bits // 8
-    kv_tensor = KVCacheTensor(
-        size=num_blocks * page_size, shared_by=["layer"], block_stride=0
-    )
+    kv_tensor = KVCacheTensor(size=num_blocks * page_size, shared_by=[["layer"]])
     return KVCacheConfig(
         num_blocks=num_blocks,
         kv_cache_tensors=[kv_tensor],
@@ -189,22 +187,12 @@ def _make_kv_cache_config():
     )
 
 
-def _make_sizing_kv_cache_config(packed: bool) -> KVCacheConfig:
+def _make_sizing_kv_cache_config() -> KVCacheConfig:
     num_blocks = 4
-    if packed:
-        kv_cache_tensors = [
-            KVCacheTensor(
-                size=64,
-                shared_by=[layer_name],
-                block_stride=16,
-            )
-            for layer_name in ("layer0", "layer1")
-        ]
-    else:
-        kv_cache_tensors = [
-            KVCacheTensor(size=40, shared_by=["layer0"]),
-            KVCacheTensor(size=24, shared_by=["layer1"]),
-        ]
+    kv_cache_tensors = [
+        KVCacheTensor(size=40, shared_by=[["layer0"]]),
+        KVCacheTensor(size=24, shared_by=[["layer1"]]),
+    ]
 
     return KVCacheConfig(
         num_blocks=num_blocks,
@@ -227,8 +215,8 @@ def _make_hybrid_kv_cache_config() -> KVCacheConfig:
     return KVCacheConfig(
         num_blocks=4,
         kv_cache_tensors=[
-            KVCacheTensor(size=40, shared_by=["full_layer"]),
-            KVCacheTensor(size=24, shared_by=["mla_layer"]),
+            KVCacheTensor(size=40, shared_by=[["full_layer"]]),
+            KVCacheTensor(size=24, shared_by=[["mla_layer"]]),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(
@@ -325,8 +313,7 @@ def test_create_cpu_offloading_spec_end_to_end():
     assert spec.num_blocks > 0
 
 
-@pytest.mark.parametrize("packed", [False, True])
-def test_cpu_spec_sizing_preserves_tensor_layout(packed: bool):
+def test_cpu_spec_sizing_preserves_tensor_layout():
     cpu_bytes_to_use = 1920
     config = _make_layout_vllm_config(
         cpu_bytes_to_use=cpu_bytes_to_use,
@@ -335,7 +322,7 @@ def test_cpu_spec_sizing_preserves_tensor_layout(packed: bool):
         pipeline_parallel_size=2,
     )
 
-    spec = _create_spec(config, _make_sizing_kv_cache_config(packed))
+    spec = _create_spec(config, _make_sizing_kv_cache_config())
 
     assert isinstance(spec, CPUOffloadingSpec)
     assert spec.cpu_page_size_per_worker == 32
@@ -343,29 +330,6 @@ def test_cpu_spec_sizing_preserves_tensor_layout(packed: bool):
     assert spec.num_blocks == cpu_bytes_to_use // 192
 
 
-def test_cpu_spec_rejects_partially_packed_tensor_layout():
-    config = _make_layout_vllm_config(cpu_bytes_to_use=65536)
-    kv_cache_config = _make_sizing_kv_cache_config(packed=False)
-    kv_cache_config.kv_cache_tensors[0].block_stride = 16
-
-    with pytest.raises(AssertionError):
-        _create_spec(config, kv_cache_config)
-
-
-def test_cpu_spec_zero_blocks_skips_tensor_layout_validation():
-    config = _make_layout_vllm_config(cpu_bytes_to_use=65536)
-    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
-
-    spec = _create_spec(config, kv_cache_config)
-
-    assert isinstance(spec, CPUOffloadingSpec)
-    assert spec.cpu_page_size_per_worker == 0
-    assert spec.kv_bytes_per_chunk == 0
-    assert spec.num_blocks == 0
-
-
 def test_tiering_spec_aligns_row_size():
     alignment = SharedOffloadRegion.BLOCK_SIZE_ALIGNMENT
     cpu_bytes_to_use = alignment * 3
@@ -377,7 +341,7 @@ def test_tiering_spec_aligns_row_size():
         pipeline_parallel_size=2,
     )
 
-    spec = _create_spec(config, _make_sizing_kv_cache_config(packed=False))
+    spec = _create_spec(config, _make_sizing_kv_cache_config())
 
     assert isinstance(spec, TieringOffloadingSpec)
     assert spec.cpu_page_size_per_worker == 32

tests/v1/simple_kv_offload/test_scheduler.py

@@ -91,7 +91,7 @@ def _make_kv_cache_config(
         tensors.append(
             KVCacheTensor(
                 size=_BYTES_PER_BLOCK * num_blocks,
-                shared_by=layer_names,
+                shared_by=[layer_names],
             )
         )
     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,12 @@ from vllm.platforms import current_platform
 if not current_platform.is_cuda_alike():
     pytest.skip("Requires CUDA or ROCm", allow_module_level=True)
 
+from vllm.v1.attention.backends.utils import set_kv_cache_layout
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheLayout,
+    reshape_kv_cache,
+)
 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,
@@ -181,3 +188,88 @@ 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 = reshape_kv_cache(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)
+
+    set_kv_cache_layout(layout.name)
+    try:
+        worker.register_kv_caches(
+            {f"layer.{layer_idx}": cache for layer_idx, cache in enumerate(caches)}
+        )
+    finally:
+        set_kv_cache_layout(None)
+
+    assert worker.gpu_kv_caches is not None
+    expected_regions = num_layers if layout.is_layer_compact else 1
+    assert len(worker.gpu_kv_caches) == expected_regions
+    expected_block_bytes = spec.page_size_bytes * (
+        1 if layout.is_layer_compact else num_layers
+    )
+    assert {cache.shape for cache in worker.gpu_kv_caches.values()} == {
+        (num_blocks, expected_block_bytes)
+    }
+
+
+@pytest.mark.parametrize("layout", [KVCacheLayout.BLHNC, KVCacheLayout.BHLNC])
+def test_register_separate_kv_head_groups(monkeypatch, layout: KVCacheLayout):
+    num_blocks = 4
+    num_layers = 2
+    spec = FullAttentionSpec(
+        block_size=2,
+        num_kv_heads=2,
+        head_size=2,
+        dtype=torch.float16,
+        separate_kv_head_groups=True,
+    )
+    raw = torch.zeros(
+        num_blocks * num_layers * spec.page_size_bytes,
+        dtype=torch.int8,
+        device="cuda",
+    )
+    caches = reshape_kv_cache(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)
+
+    set_kv_cache_layout(layout.name)
+    try:
+        worker.register_kv_caches(
+            {f"layer.{layer_idx}": cache for layer_idx, cache in enumerate(caches)}
+        )
+    finally:
+        set_kv_cache_layout(None)
+
+    assert worker.gpu_kv_caches is not None
+    assert len(worker.gpu_kv_caches) == num_layers * spec.num_heads
+    per_head_block_bytes = spec.block_size * spec.head_size * spec.dtype.itemsize
+    assert {cache.shape for cache in worker.gpu_kv_caches.values()} == {
+        (num_blocks, per_head_block_bytes)
+    }

tests/v1/worker/test_attn_utils.py

@@ -1,42 +1,32 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Padded-page handling in reshape_kv_cache.
 
+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 vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode
-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 set_kv_cache_layout
+from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheLayout,
+    KVCacheTensor,
+    KVQuantMode,
+    MLAAttentionSpec,
+    reshape_kv_cache,
+)
+from vllm.v1.worker.gpu.attn_utils import _allocate_and_reshape_kv_cache
+from vllm.v1.worker.utils import copy_kv_cache_blocks_inplace
 
 
-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,
@@ -47,163 +37,20 @@ 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,) = reshape_kv_cache(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,
-        num_kv_heads=1,
-        head_size=2,
-        dtype=torch.float32,
-        page_size_padded=384,
-    )
-
-    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():
+def test_reshape_padded_quantized_kv_cache_budgets_scale_bytes():
     num_blocks = 3
     spec = FullAttentionSpec(
         block_size=16,
@@ -213,30 +60,182 @@ def test_reshape_padded_quantized_kv_cache_preserves_scale_stride():
         kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD,
         page_size_padded=384,
     )
+    # Per-token-head scales are budgeted into the page but live past the
+    # real content, so the logical view must stride by the padded page.
     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,
+    raw = torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
+    (kv_cache,) = reshape_kv_cache(raw, spec, num_blocks, 1, KVCacheLayout.LBHNC)
+
+    assert kv_cache.shape == (num_blocks, 1, 16, 2 * spec.head_size)
+    assert kv_cache.stride(0) == spec.page_size_padded
+    assert kv_cache[1].storage_offset() == spec.page_size_padded
+
+
+@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=128,
+        dtype=torch.bfloat16,
+        tokens_per_state=4,
+    )
+    num_pages = 4
+    config = KVCacheConfig(
+        num_blocks=num_pages,
+        kv_cache_tensors=[
+            KVCacheTensor(
+                size=num_pages * spec.page_size_bytes,
+                shared_by=[["layer.0"]],
+            )
+        ],
+        kv_cache_groups=[KVCacheGroupSpec(["layer.0"], spec)],
+    )
+
+    caches = _allocate_and_reshape_kv_cache(
+        config,
+        torch.device("cpu"),
+        layout=KVCacheLayout.LBHNC,
+        kernel_block_sizes=kernel_block_sizes,
+    )
+
+    assert caches["layer.0"].shape == (
+        expected_num_blocks,
+        1,
+        expected_num_states,
+        128,
+    )
+
+
+@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 = reshape_kv_cache(raw, spec, num_blocks, num_layers, layout)
+
+    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)]
+
+    set_kv_cache_layout(layout.name)
+    try:
+        copy_kv_cache_blocks_inplace(caches, num_blocks, copies)
+    finally:
+        set_kv_cache_layout(None)
+
+    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", [KVCacheLayout.BLHNC, KVCacheLayout.BHLNC])
+def test_copy_kv_cache_blocks_separate_head_groups(layout: KVCacheLayout):
+    num_blocks = 4
+    num_layers = 2
+    spec = FullAttentionSpec(
+        block_size=2,
+        num_kv_heads=2,
+        head_size=2,
+        dtype=torch.float32,
+        separate_kv_head_groups=True,
+    )
+    raw = torch.zeros(
+        num_blocks * num_layers * spec.page_size_bytes,
+        dtype=torch.int8,
+    )
+    caches = reshape_kv_cache(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]
+    set_kv_cache_layout(layout.name)
+    try:
+        copy_kv_cache_blocks_inplace(
+            caches,
+            num_blocks,
+            [KVCacheBlockCopy(src_block_id=0, dst_block_id=2)],
         )
-    ]
+    finally:
+        set_kv_cache_layout(None)
 
-    kv_cache = _reshape_kv_cache(
-        attn_groups,
-        raw_tensors,
-        "int8_per_token_head",
-        [spec.block_size],
-        {},
-    )["layer"]
+    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])
 
-    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
+
+@pytest.mark.parametrize("layout", [KVCacheLayout.LBHNC, KVCacheLayout.BLHNC])
+def test_copy_kv_cache_blocks_with_virtual_block_splitting(layout: KVCacheLayout):
+    num_blocks = 4
+    num_layers = 2
+    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 = reshape_kv_cache(
+        raw,
+        spec,
+        num_blocks * physical_per_logical,
+        num_layers,
+        layout,
+        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]
+
+    set_kv_cache_layout(layout.name)
+    try:
+        copy_kv_cache_blocks_inplace(
+            caches,
+            num_blocks,
+            [KVCacheBlockCopy(src_block_id=0, dst_block_id=2)],
+        )
+    finally:
+        set_kv_cache_layout(None)
+
+    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_gpu_model_runner.py

@@ -72,7 +72,7 @@ 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, shared_by=[["layer.0"]]),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=attn_spec)
@@ -776,56 +776,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"}})
@@ -1030,21 +980,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 slot (2 slots * 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
+        num_layer_slots = len(kv_cache_tensor.shared_by)
+        kv_cache_tensor.size = (
+            kv_cache_spec[kv_cache_tensor.shared_by[0][0]].page_size_bytes
+            * num_layer_slots
+        )
 
     runner.initialize_kv_cache(kv_cache_config)
 
@@ -1138,7 +1089,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
     """
@@ -1301,10 +1252,9 @@ def test_hybrid_attention_mamba_tensor_shapes():
             actual_kv = vllm_ctx[layer].kv_cache[kernel_block, :]
             expected = attn_blocks_constant[i]
 
-            # Packed layout: (num_kv_heads, block_size, 2*head_size). Every
-            # head in the block was filled with the same constant.
-            for head_idx in range(actual_kv.shape[0]):
-                assert torch.equal(actual_kv[head_idx], expected)
+            # Check K and V separately
+            assert torch.equal(actual_kv[0], expected)
+            assert torch.equal(actual_kv[1], expected)
 
     for layer in [layer_2, layer_3, layer_4, layer_5]:
         for i, kv_block in enumerate(kv_blocks_for_mamba):
@@ -1444,7 +1394,7 @@ 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, shared_by=[["layer.0"]]),
         ],
         kv_cache_groups=[
             KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=attn_spec)

vllm/_aiter_ops.py

@@ -2541,10 +2541,9 @@ class rocm_aiter_ops:
     ) -> None:
         """Run the fused QK-norm+RoPE+KV-cache op on already-split k/v caches.
 
-        Shared by the AITER FA and unified-attention impls. The caller splits
-        kv_cache, since the unbind dim depends on the layout (e.g. the unified
-        encoder-decoder path is K/V-first), and passes use_shuffle_layout
-        (unified reads NHD and must pass False).
+        Shared by the AITER FA and unified-attention impls. The caller converts
+        the standardized cache view to NHD and splits its packed K/V content,
+        then passes use_shuffle_layout (unified reads NHD and must pass False).
         """
         if kv_cache_dtype.startswith("fp8"):
             key_cache = key_cache.view(current_platform.fp8_dtype())

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/kv_transfer.py

@@ -64,7 +64,7 @@ class KVTransferConfig:
     Only supported in V1."""
 
     enable_permute_local_kv: bool = False
-    """Experiment feature flag to enable HND to NHD KV Transfer"""
+    """Experiment feature flag to enable HNC to NHC KV Transfer"""
 
     kv_load_failure_policy: Literal["recompute", "fail"] = "fail"
     """Policy for handling KV cache load failures.

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,19 +35,18 @@ BlockIds = tuple[list[int], ...] | list[list[int]]
 
 
 def get_kv_connector_cache_layout():
-    # NOTE (NickLucche) When running disaggregated PD with NIXL, HND layout is
-    # used for faster transfer.
-    vllm_config = get_current_vllm_config()
+    # NOTE (NickLucche) When running disaggregated PD with NIXL, LBHNC layout
+    # is used for faster transfer.
+    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:
@@ -279,11 +278,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.
-    (Only works for local blocksize > remote blocksize)
+    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)
 
@@ -418,47 +417,12 @@ 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}."
-                )
+        # Cross-layer layouts (BLHNC) have B outermost, so all layers
+        # for a block are contiguous — transfers can coalesce multiple
+        # layers into one operation.
+        from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
 
-        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)
+        self._is_kv_layout_blocks_first = not resolve_kv_cache_layout().is_layer_compact
 
     # ============================================================
     # Engine registration
@@ -499,18 +463,8 @@ class TransferTopology:
     # ============================================================
 
     @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
+    def is_kv_layout_blocks_first(self) -> bool:
+        return self._is_kv_layout_blocks_first
 
     # ============================================================
     # 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
 
@@ -173,14 +173,6 @@ class KVConnectorBase_V1(ABC):
     Base class for KV connectors.
     """
 
-    @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
-
     def __init__(
         self,
         vllm_config: "VllmConfig",
@@ -258,23 +250,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.
@@ -594,7 +569,7 @@ class KVConnectorBase_V1(ABC):
             vllm_config (VllmConfig): the vllm config.
 
         Returns:
-            str: the required KV cache layout. e.g. HND, or NHD.
+            str: the required KV cache layout. e.g. HNC, or NHC.
             None if the connector does not require a specific layout.
         """
 

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.
@@ -596,7 +589,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
             vllm_config (VllmConfig): the vllm config.
 
         Returns:
-            str: the required KV cache layout. e.g. HND, or NHD.
+            str: the required KV cache layout. e.g. HNC, or NHC.
             None if the connector does not require a specific layout.
         """
 
@@ -605,9 +598,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/lmcache_mp_connector.py

@@ -974,7 +974,7 @@ class LMCacheMPConnectorUpstream(KVConnectorBase_V1):
             vllm_config (VllmConfig): the vllm config.
 
         Returns:
-            str: the required KV cache layout. e.g. HND, or NHD.
+            str: the required KV cache layout. e.g. HNC, or NHC.
             None if the connector does not require a specific layout.
         """
 

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

@@ -50,14 +50,13 @@ 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.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, resolve_kv_cache_layout
 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 +121,8 @@ def _expand_transfer_regions(
     kv_block_lens: list[int],
     layer_names: list[str],
     layer_indices: list[int],
-    is_kv_layout_blocks_first: bool,
+    is_kv_layout_blocks_first: bool,  # kept for API compat, unused
     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,7 +1001,7 @@ 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()
+        self.kv_cache_layout = resolve_kv_cache_layout().name
         logger.debug(
             "Detected attention backends %s",
             [backend.get_name() for backend in self.attn_backends],
@@ -1674,13 +1652,10 @@ class MooncakeConnectorWorker:
                     layer_name,
                 )
                 continue
-            if isinstance(layer_spec, MambaSpec):
-                conv, _ = cache_or_caches
-                cache_list = [conv]
-            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]
+            # Standardized allocation exposes one raw page tensor per layer.
+            # For Mamba that page contains all recurrent states; the layer
+            # unpacks it only when binding the cache for model execution.
+            cache_list = [cache_or_caches]
 
             logger.debug(
                 "registering layer %s with %d cache tensor(s)",
@@ -1690,25 +1665,32 @@ class MooncakeConnectorWorker:
 
             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()
-                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
+                if (
+                    isinstance(layer_spec, AttentionSpec)
+                    and layer_spec.separate_kv_head_groups
                 ):
-                    kv_block_len = block_len // 2
+                    region_caches = [cache[:, head] for head in range(cache.shape[1])]
                 else:
-                    kv_block_len = 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)
-                self.registered_layer_indices.append(layer_index)
-                self.registered_group_indices.append(
-                    self._layer_group_indices[layer_name]
-                )
+                    region_caches = [cache]
+
+                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)
+
+                    kv_block_len = (
+                        layer_spec.page_size_bytes
+                        if isinstance(layer_spec, AttentionSpec)
+                        and not layer_spec.separate_kv_head_groups
+                        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)
+                    self.registered_layer_indices.append(layer_index)
+                    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:
@@ -2055,24 +2037,14 @@ 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,
+            is_kv_layout_blocks_first=self.transfer_topo.is_kv_layout_blocks_first,
             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

@@ -87,14 +87,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
@@ -262,16 +254,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

@@ -61,12 +61,18 @@ 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.v1.attention.backends.utils import resolve_kv_cache_layout
 from vllm.v1.core.kv_cache_utils import (
     BlockHash,
     maybe_convert_block_hash,
     resolve_kv_cache_block_sizes,
 )
-from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec
+from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    UniformTypeKVCacheSpecs,
+)
 
 from .metrics import MooncakeStoreConnectorStats
 
@@ -1193,15 +1199,6 @@ class MooncakeStoreWorker:
         )
         self._lookup_expected_per_key = len(rank_namespaces)
 
-    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]],
@@ -1222,47 +1219,65 @@ class MooncakeStoreWorker:
         assert self.cache_config.num_gpu_blocks is not None
         self.num_blocks = self.cache_config.num_gpu_blocks
 
-        seen_ptrs: set[int] = set()
+        layout = resolve_kv_cache_layout()
+        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)
-            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,
+        layer_specs = {}
+        for group in self._kv_cache_groups:
+            group_spec = group.kv_cache_spec
+            for layer_name in group.layer_names:
+                layer_specs[layer_name] = (
+                    group_spec.kv_cache_specs[layer_name]
+                    if isinstance(group_spec, UniformTypeKVCacheSpecs)
+                    else group_spec
                 )
 
-            # 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.
+        for layer_name, value in kv_caches.items():
+            cache = _repr_tensor(value)
+            layer_spec = layer_specs.get(layer_name)
+            cache_storage = cache.untyped_storage()
+            base_addr = cache_storage.data_ptr()
+            region_len = cache_storage.nbytes()
+
+            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,
+                    )
+
+            if (
+                isinstance(layer_spec, AttentionSpec)
+                and layer_spec.separate_kv_head_groups
+            ):
+                for head_idx in range(cache.shape[1]):
+                    head_cache = cache[:, head_idx]
+                    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 not layout.is_layer_compact:
+                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/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
 
@@ -203,12 +203,6 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
         # Propagated from scheduler to worker side via the connector metadata.
         self._extra_async_saves: dict[str, int] = {}
 
-    @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)
-
     @classmethod
     def _get_connector_classes_and_configs(
         cls, vllm_config: "VllmConfig"
@@ -235,13 +229,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)
@@ -551,7 +538,7 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
             vllm_config (VllmConfig): the vllm config.
 
         Returns:
-            str: the required KV cache layout. e.g. HND, or NHD.
+            str: the required KV cache layout. e.g. HNC, or NHC.
             None if the connector does not require a specific layout.
         """
         assert vllm_config.kv_transfer_config is not None

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

@@ -67,9 +67,10 @@ from vllm.distributed.parallel_state import (
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.utils.network_utils import make_zmq_path
-from vllm.v1.attention.backends.utils import get_kv_cache_layout
+from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
+    KVCacheLayout,
     MambaSpec,
     MLAAttentionSpec,
     SlidingWindowMLASpec,
@@ -286,7 +287,9 @@ 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)
+        self.hma_group_size = sum(
+            len(t.shared_by) for t in kv_cache_config.kv_cache_tensors
+        )
 
         # ---- Model state (derived from model config) ----
         mamba_ssm_size = (0, 0)
@@ -501,7 +504,7 @@ 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 = resolve_kv_cache_layout().name
         self.host_buffer_kv_cache_layout = self.kv_cache_layout
         logger.info(
             "Detected attention backend(s) %s",
@@ -536,6 +539,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-block physical stride per layer (bytes). Read from each
+        # registered tensor's stride(0) so it stays correct under layouts
+        # that interleave layers within a block (BLHNC/BHLNC).
+        self.block_stride_per_layer = list[int]()
+
         # Per-engine TP mappings. Generated during handshake.
         self.tp_mappings: dict[EngineId, TPMapping] = {}
 
@@ -746,20 +754,22 @@ 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,
-                            # we can leverage host_buffer for permute.
-                            self.host_buffer_kv_cache_layout = "HND"
+                            # Since LBNHC will not support Decode/Prefill
+                            # TP_ratio > 1, we can leverage host_buffer for
+                            # permute.
+                            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)
-                            # so raw NIXL transfers see NHD physical strides.
+                            # so raw NIXL transfers see NHC physical strides.
                             kv_shape = tuple(kv_shape[i] for i in inv_order)
                             permute_shape = True
 
@@ -927,117 +937,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,
-        )
-
-        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
-            ),
-        )
-        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,
@@ -1053,7 +955,7 @@ class NixlBaseConnectorWorker:
             is_mamba=self._has_mamba,
         )
         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
         )
 
         if self.use_host_buffer:
@@ -1079,8 +981,8 @@ class NixlBaseConnectorWorker:
         )
 
         caches_data = []
-        # With hybrid allocator, layers can share a kv cache tensor
-        seen_base_addresses = []
+        seen_storage_addresses: set[int] = set()
+        seen_base_addresses: list[int] = []
 
         # Note(tms): I modified this from the original region setup code.
         # K and V are now in different regions. Advantage is that we can
@@ -1090,8 +992,6 @@ class NixlBaseConnectorWorker:
         # (roughly 8KB vs 5KB).
         # Conversely for FlashInfer, K and V are registered in the same region
         # to better exploit the memory layout (ie num_blocks is the first dim).
-        tensor_size_bytes = None
-
         for layer_name, cache in xfer_buffers.items():
             # NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to
             # that of FI, with block laid out as in `get_backend_aware_kv_block_len`.
@@ -1117,50 +1017,78 @@ 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()
-            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.
-                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)
+            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, ""))
+
             is_mla_region = isinstance(
                 layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
             )
-            self._region_is_mla.append(is_mla_region)
 
-            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"
+            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:
+                block_stride = cache.stride(0) * cache.element_size()
+                storage_is_block_major = num_blocks * block_stride == storage.nbytes()
+                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 not hnc_contiguous:
+                    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:
+                    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.
@@ -1181,11 +1109,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)
         )
@@ -1193,10 +1116,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(
@@ -1246,6 +1170,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,
@@ -1368,7 +1293,7 @@ class NixlBaseConnectorWorker:
                 )
                 // block_size_ratio
             )
-            page_stride = self.block_len_per_layer[i] // block_size_ratio
+            page_stride = self.block_stride_per_layer[i] // block_size_ratio
             addrs = base_addr + block_arange * page_stride
             parts.append(self._stack_descs(addrs, kv_block_len, device_id))
         return np.concatenate(parts)
@@ -1413,7 +1338,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)
@@ -1497,7 +1422,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.
@@ -1538,6 +1463,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
@@ -1714,10 +1640,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, (
@@ -1727,7 +1653,7 @@ class NixlBaseConnectorWorker:
             else:
                 raise RuntimeError(
                     "Heterogeneous TP expects same kv_cache_layout. "
-                    "Or enable experimental feature to use HND to NHD support by "
+                    "Or enable experimental feature to use HNC to NHC support by "
                     "setting 'enable_permute_local_kv'=True in --kv-transfer-config."
                 )
         # if remote_agent used attn is not same as local,
@@ -1747,18 +1673,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 (H before N, e.g. HNC / BHLNC).
+        # MLA and replicated-KV cases don't split on heads.
         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 HNC layout on the prefill side."
             )
 
         # Per-region block_len validation enforcing the P/D invariant.
@@ -1870,11 +1796,11 @@ class NixlBaseConnectorWorker:
         Post process device kv cache after receiving from remote.
 
         3 types of post processing supported:
-            * kv_cache_postprocess_layout => convert from HND to NHD
+            * kv_cache_postprocess_layout => convert from HNC to NHC
             * kv_cache_postprocess_blksize => convert from small block size
               to large block size
             * kv_cache_postprocess_blksize_and_layout => convert from small
-              block size to large block size and convert from HND to NHD
+              block size to large block size and convert from HNC to NHC
 
         """
         if len(self.device_kv_caches) == 0:
@@ -1884,14 +1810,14 @@ class NixlBaseConnectorWorker:
         if self.enable_permute_local_kv and block_size_ratio > 1:
             logger.debug(
                 "Post-processing device kv cache on receive by converting "
-                "block_size with %sx bigger and permuting layout from HND"
-                " to NHD.",
+                "block_size with %sx bigger and permuting layout from HNC"
+                " to NHC.",
                 block_size_ratio,
             )
         elif self.enable_permute_local_kv:
             logger.debug(
                 "Post-processing device kv cache on receive by permuting layout"
-                "from HND to NHD."
+                "from HNC to NHC."
             )
         else:
             logger.debug(
@@ -2333,7 +2259,7 @@ class NixlBaseConnectorWorker:
            |1st_split-2nd_split|         |1st_split-2nd_split |
         """
         assert self.transfer_topo is not None
-        if self.transfer_topo.virtually_split_kv_in_blocks and mamba_view:
+        if self.transfer_topo.is_kv_layout_blocks_first and mamba_view:
             block_len = self._mamba_ssm_size[not first_split]
         else:
             block_len = self.block_len_per_layer[layer_idx]

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:
@@ -79,36 +76,6 @@ logger = init_logger(__name__)
 class NixlBaseConnector(KVConnectorBase_V1, SupportsHMA):
     """Base connector with common logic shared by pull and push modes."""
 
-    @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,
@@ -152,9 +119,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
@@ -227,12 +195,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

@@ -41,8 +41,9 @@ PUSH_REG_NOTIF_PREFIX = b"PUSH_REG:"
 #   4: Add KV block lease renewal through heartbeats
 #   5: Add remote_blocks_expiry_time to kv_transfer_params + handshake
 #      clock-sync timestamp
+#   6: Add block_strides
 #
-NIXL_CONNECTOR_VERSION: int = 5
+NIXL_CONNECTOR_VERSION: int = 6
 
 
 @dataclass
@@ -53,6 +54,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]
@@ -79,7 +81,7 @@ class NixlHandshakePayload(KVConnectorHandshakeMetadata):
 
 
 def compute_nixl_compatibility_hash(
-    vllm_config: VllmConfig, attn_backend_name: str, cross_layers_blocks: bool
+    vllm_config: VllmConfig, attn_backend_name: str
 ) -> str:
     """
     Compute compatibility hash for NIXL KV transfer.
@@ -123,7 +125,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,
     }
 

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

@@ -16,12 +16,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(
@@ -89,16 +84,8 @@ def build_offloading_config(
 
     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)
+        total_gpu_kv_bytes = sum(
+            tensor.size for tensor in kv_cache_config.kv_cache_tensors
         )
         worker_kv_bytes_per_block = total_gpu_kv_bytes // kv_cache_config.num_blocks
 

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

@@ -10,11 +10,7 @@ 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,
@@ -56,21 +52,14 @@ class OffloadingConnectorWorker:
     def _init_worker(self, kv_caches: CanonicalKVCaches) -> None:
         self.worker = self.spec.get_worker(kv_caches)
 
-    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
+    def register_kv_caches(
+        self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]]
+    ):
         kv_cache_config = self.kv_cache_config
         num_blocks = kv_cache_config.num_blocks
 
-        # 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
+        # layer_name -> (num_blocks, page_size_bytes) int8 view.
+        # Standardized layouts always have num_blocks as the leading dim.
         tensors_per_block: dict[str, tuple[torch.Tensor, ...]] = {}
         # layer_name -> size of (un-padded) page in bytes
         unpadded_page_size_bytes: dict[str, int] = {}
@@ -87,63 +76,62 @@ 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)
+                layer_kv_cache = kv_caches[layer_name]
+                # AttentionSpec yields a single tensor; MambaSpec yields a
+                # list of typed state tensors that share one underlying
+                # buffer. Either way, the first tensor's storage_offset
+                # marks the start of this layer's region.
+                ref = (
+                    layer_kv_cache[0]
+                    if isinstance(layer_kv_cache, list)
+                    else layer_kv_cache
+                )
+                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. Non-packed 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),
@@ -164,46 +152,43 @@ 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
+            for slot_layers in kv_cache_tensor.shared_by:
+                # Filter to layers that were actually processed above.
+                # Some slots may have no corresponding model layer (reserved
+                # memory with no group layer at that index).
+                tensor_layer_names = [n for n in slot_layers if n in tensors_per_block]
+                if not tensor_layer_names:
+                    continue
 
-            # verify all layers in the group reference the exact same tensors
-            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(
-                    CanonicalKVCacheTensor(
-                        tensor=tensor,
-                        page_size_bytes=page_size_bytes[first_layer_name],
-                    )
+                # Verify all layers in the slot reference the same tensors.
+                assert len({len(tensors_per_block[n]) for n in tensor_layer_names}) == 1
+                data_ptrs = {
+                    n: tensors_per_block[n][0].data_ptr() for n in tensor_layer_names
+                }
+                assert len(set(data_ptrs.values())) == 1, data_ptrs
+                assert (
+                    len({tensors_per_block[n][0].stride() for n in tensor_layer_names})
+                    == 1
                 )
 
-                curr_tensor_idx = len(block_tensors) - 1
-                for layer_name in tensor_layer_names:
-                    block_data_refs[layer_name].append(
-                        CanonicalKVCacheRef(
-                            tensor_idx=curr_tensor_idx,
-                            page_size_bytes=(unpadded_page_size_bytes[layer_name]),
+                first_layer_name = tensor_layer_names[0]
+                for tensor in tensors_per_block[first_layer_name]:
+                    block_tensors.append(
+                        CanonicalKVCacheTensor(
+                            tensor=tensor,
+                            page_size_bytes=page_size_bytes[first_layer_name],
                         )
                     )
 
+                    curr_tensor_idx = len(block_tensors) - 1
+                    for layer_name in tensor_layer_names:
+                        block_data_refs[layer_name].append(
+                            CanonicalKVCacheRef(
+                                tensor_idx=curr_tensor_idx,
+                                page_size_bytes=(unpadded_page_size_bytes[layer_name]),
+                            )
+                        )
+
         group_data_refs: list[list[CanonicalKVCacheRef]] = []
         for kv_cache_group in kv_cache_config.kv_cache_groups:
             group_refs: list[CanonicalKVCacheRef] = []
@@ -218,53 +203,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,10 +47,6 @@ from vllm.v1.request import Request
 
 
 class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
-    @property
-    def prefer_cross_layer_blocks(self) -> bool:
-        return True
-
     def __init__(
         self,
         vllm_config: VllmConfig,
@@ -81,12 +77,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)
@@ -186,7 +176,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

@@ -223,7 +223,20 @@ 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",
+            "NHC",
+            "HNC",
+            "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_ROCM_QUICK_REDUCE_QUANTIZATION: Literal[
@@ -1684,18 +1697,30 @@ 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
-    # leave the layout choice to the backend. Mind that backends may only
+    # - 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.
     "VLLM_KV_CACHE_LAYOUT": env_with_choices(
-        "VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"]
+        "VLLM_KV_CACHE_LAYOUT",
+        None,
+        [
+            "LBNHC",
+            "LBHNC",
+            "NHC",
+            "HNC",
+            "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

@@ -577,6 +577,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:
@@ -1337,27 +1341,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]
@@ -2219,7 +2202,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
                 )
             elif not use_fp8_prefill:
                 ops.gather_and_maybe_dequant_cache(
-                    src_cache=kv_c_and_k_pe_cache,
+                    src_cache=kv_c_and_k_pe_cache.squeeze(1),
                     dst=workspace,
                     block_table=prefill_metadata.block_table,
                     cu_seq_lens=prefill_metadata.chunked_context.cu_seq_lens[i],
@@ -2232,7 +2215,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
             else:
                 # FP8 path: gather cache without dequantization
                 ops.cp_gather_cache(
-                    src_cache=kv_c_and_k_pe_cache,
+                    src_cache=kv_c_and_k_pe_cache.squeeze(1),
                     dst=workspace,
                     block_table=prefill_metadata.block_table,
                     cu_seq_lens=prefill_metadata.chunked_context.cu_seq_lens[i],
@@ -2344,7 +2327,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
             elif is_quantized_kv_cache(self.kv_cache_dtype):
                 assert k_scale is not None
                 ops.gather_and_maybe_dequant_cache(
-                    src_cache=kv_c_and_k_pe_cache,
+                    src_cache=kv_c_and_k_pe_cache.squeeze(1),
                     dst=workspace,
                     block_table=prefill_metadata.block_table,
                     cu_seq_lens=padded_local_cu_seq_lens,
@@ -2358,7 +2341,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
                 )
             else:
                 ops.cp_gather_cache(
-                    src_cache=kv_c_and_k_pe_cache,
+                    src_cache=kv_c_and_k_pe_cache.squeeze(1),
                     dst=workspace,
                     block_table=prefill_metadata.block_table,
                     cu_seq_lens=padded_local_cu_seq_lens,

vllm/model_executor/layers/attention_layer_base.py

@@ -24,11 +24,7 @@ class AttentionLayerBase(ABC):
     supports_dcp: bool = True
 
     def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
-        """Bind the allocated KV cache tensor to this layer.
-
-        The default stores the cache view as-is; subclasses (e.g. Mamba)
-        override this to unpack the raw buffer into per-state views.
-        """
+        """Bind a ``[B, H, N, C]`` cache view; override to reshape."""
         self.kv_cache = kv_cache
 
     @abstractmethod

vllm/model_executor/layers/mamba/abstract.py

@@ -27,11 +27,7 @@ class MambaBase(AttentionLayerBase):
     supports_dcp: bool = False
 
     def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
-        """Unpack a raw ``[B, 1, 1, C]`` int8 page view into per-state views.
-
-        Each block's ``C`` bytes hold the layer's states (e.g. conv, ssm)
-        packed contiguously; slice them out and reinterpret per dtype/shape.
-        """
+        """Unpack a raw 4D ``[B, 1, 1, C]`` int8 view into per-state views."""
         pages = kv_cache.squeeze(dim=(1, 2))
         states: list[torch.Tensor] = []
         offset = 0

vllm/model_executor/layers/mamba/mamba_utils.py

@@ -28,7 +28,7 @@ def get_conv_state_layout() -> ConvStateLayoutType:
     """Return the SSM conv state layout.
 
     SD = (state_len, dim) — dim is the innermost contiguous dimension.
-    DS = (dim, state_len) — TP-sharded dim is on dim-1 (like HND for KV
+    DS = (dim, state_len) — TP-sharded dim is on dim-1 (like HNC for KV
          cache), consistent with SSM temporal state layout.
     """
     layout: ConvStateLayoutType | None = envs.VLLM_SSM_CONV_STATE_LAYOUT

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

@@ -128,23 +128,13 @@ 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(
-                block_size=kv_cache_spec.block_size * block_pool_size,
-                num_kv_heads=kv_cache_spec.num_kv_heads // block_pool_size,
-            )
             if isinstance(kv_cache_spec, SlidingWindowSpec):
                 # The manager keeps `sliding_window` in pooled units; the kernel
                 # runs on the expanded sequence, so give it the model's window in
                 # 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 +251,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,
         },
     )
@@ -339,9 +317,11 @@ class WhisperCausalAttentionWithBlockPooling(Attention):
     def get_kv_cache_spec(self, vllm_config: VllmConfig):
         kv_cache_spec = super().get_kv_cache_spec(vllm_config)
         assert isinstance(kv_cache_spec, AttentionSpec)
+        assert kv_cache_spec.num_kv_heads % self.block_pool_size == 0
         kv_cache_spec = replace(
             kv_cache_spec,
-            num_kv_heads=self.block_pool_size * kv_cache_spec.num_kv_heads,
+            block_size=kv_cache_spec.block_size * self.block_pool_size,
+            num_kv_heads=kv_cache_spec.num_kv_heads // 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

@@ -599,6 +599,10 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
         )
         return q_fp8
 
+    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
 
@@ -616,7 +620,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",
@@ -645,16 +649,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

@@ -62,25 +62,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:
@@ -166,6 +147,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

@@ -21,7 +21,6 @@ from vllm.models.deepseek_v4.sparse_mla import (
     DeepseekV4FlashMLABackend,
     DeepseekV4FlashMLAMetadata,
 )
-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
@@ -135,26 +134,6 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4FlashMLABackend):
             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 DeepseekV4FlashMLABackend.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/nvidia/ops/sparse_attn_compress_cutedsl.py

@@ -2098,6 +2098,8 @@ def compress_norm_rope_store_cutedsl(
     store_full_fp8: bool = False,
     fp8_scale: torch.Tensor | None = None,
 ) -> None:
+    # (B, H=1, N, C) -> (B, N, C)
+    kv_cache = kv_cache.squeeze(1)
     if compress_ratio == 4:
         # For C4A, the single fused kernel is faster than the two-kernel version.
         fused_kv_compress_norm_rope_insert_sparse_attn_cutedsl(

vllm/models/deepseek_v4/sparse_mla.py

@@ -92,21 +92,6 @@ class DeepseekV4FlashMLABackend(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):
@@ -155,8 +140,7 @@ class DeepseekV4FlashMLAMetadataBuilder(
             (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
+        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.

vllm/models/inkling/nvidia/sconv_swa_attn.py

@@ -109,30 +109,10 @@ class InklingSconvBackend(AttentionBackend):
 
     @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.
+        # The standardized layout keeps num_blocks outermost for this cache,
+        # 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/minimax_m3/amd/model.py

@@ -705,6 +705,7 @@ 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),
+            separate_kv_head_groups=self.use_aiter_sparse_pa,
         )
 
     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(

vllm/models/minimax_m3/common/sparse_attention.py

@@ -50,10 +50,7 @@ from vllm.v1.attention.backend import (
     CommonAttentionMetadata,
     MultipleOf,
 )
-from vllm.v1.attention.backends.utils import (
-    get_kv_cache_layout,
-    split_decodes_and_prefills,
-)
+from vllm.v1.attention.backends.utils import split_decodes_and_prefills
 from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache
 
 logger = init_logger(__name__)
@@ -112,48 +109,6 @@ class MiniMaxM3SparseBackend(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, ...]:
-        if minimax_m3_use_aiter_sparse_pa(num_kv_heads):
-            # AITER's assembly paged-attention kernels require independently
-            # contiguous K and V storage. Keep that specialized layout behind
-            # the shuffle flag while every other implementation uses the
-            # packed-content contract introduced by #44455.
-            return (num_blocks, 2, block_size, num_kv_heads, head_size)
-        # K and V are packed into the content dim: logical (B, H, N, 2*hs).
-        return (num_blocks, num_kv_heads, block_size, 2 * head_size)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        # `stride_order` indicates the permutation that gets us from
-        # `get_kv_cache_shape` (logical (B, H, N, 2*hs)) to the actual memory
-        # layout we want.
-        if include_num_layers_dimension:
-            raise NotImplementedError  # no cross-layer KV blocks in M3
-        if _minimax_m3_aiter_sparse_pa_requested():
-            # The AITER page-16 sparse PA path reinterprets the K and V slices
-            # as separate SHUFFLE caches. Keep K/V physically separated so
-            # kv_cache[:, 0] and kv_cache[:, 1] are contiguous byte ranges.
-            return (1, 0, 2, 3, 4)
-        cache_layout = get_kv_cache_layout()
-        if cache_layout == "NHD":
-            # (num_blocks, block_size, num_kv_heads, 2*head_size)
-            stride_order = (0, 2, 1, 3)
-        elif cache_layout == "HND":
-            # (num_blocks, num_kv_heads, block_size, 2*head_size)
-            stride_order = (0, 1, 2, 3)
-        else:
-            raise ValueError(f"Unknown cache layout format {cache_layout}.")
-        return stride_order
-
 
 @dataclass
 class MiniMaxM3SparsePrefillMetadata:

vllm/platforms/rocm.py

@@ -427,10 +427,6 @@ def _get_backend_priorities(
             ]
 
     backends = []
-    # Keep ROCM_ATTN disabled for KV connectors until connector transfer
-    # semantics are validated for its asymmetric native K/V cache views.
-    if not use_kv_connector:
-        backends.append(AttentionBackendEnum.ROCM_ATTN)
     if rocm_aiter_ops.is_mha_enabled():
         backends.append(AttentionBackendEnum.ROCM_AITER_FA)
     if is_aiter_found_and_supported():
@@ -511,20 +507,6 @@ class RocmPlatform(Platform):
             attn_selector_config.use_sparse,
             attn_selector_config.use_kv_connector,
         )
-        from vllm.config import get_current_vllm_config_or_none
-
-        vllm_config = get_current_vllm_config_or_none()
-        is_encoder_decoder = (
-            getattr(getattr(vllm_config, "model_config", None), "attn_type", None)
-            == "encoder_decoder"
-        )
-        # ROCM_ATTN still uses a legacy attention layout (KV is the outer
-        # dimension) that is incompatible with the encoder backend layouts. The
-        # encoder and decoder need the layouts to match. This is currently
-        # enforced implicitly.
-        # TODO: Make this explicit in the selector in a future PR.
-        if is_encoder_decoder and AttentionBackendEnum.ROCM_ATTN in backend_priorities:
-            backend_priorities.remove(AttentionBackendEnum.ROCM_ATTN)
         for priority, backend in enumerate(backend_priorities):
             try:
                 backend_class = backend.get_class()
@@ -553,14 +535,19 @@ class RocmPlatform(Platform):
 
         # First try checking just the selected backend, if there is one.
         if selected_backend is not None:
-            try:
-                backend_class = selected_backend.get_class()
-                invalid_reasons = backend_class.validate_configuration(
-                    device_capability=device_capability,
-                    **attn_selector_config._asdict(),
-                )
-            except ImportError:
-                invalid_reasons = ["ImportError"]
+            if selected_backend == AttentionBackendEnum.ROCM_ATTN:
+                invalid_reasons = [
+                    "ROCM_ATTN does not support standardized packed KV caches"
+                ]
+            else:
+                try:
+                    backend_class = selected_backend.get_class()
+                    invalid_reasons = backend_class.validate_configuration(
+                        device_capability=device_capability,
+                        **attn_selector_config._asdict(),
+                    )
+                except ImportError:
+                    invalid_reasons = ["ImportError"]
             if invalid_reasons:
                 raise ValueError(
                     f"Selected backend {selected_backend} is not valid for "

vllm/platforms/xpu.py

@@ -124,6 +124,14 @@ class XPUPlatform(Platform):
         attn_selector_config: "AttentionSelectorConfig",
         num_heads: int | None = None,
     ) -> str:
+        from vllm.v1.attention.backends.utils import set_kv_cache_layout
+
+        set_kv_cache_layout("LBNHC")
+        logger.info_once(
+            "Setting VLLM_KV_CACHE_LAYOUT to 'LBNHC' for XPU; "
+            "only LBNHC layout is supported by XPU attention kernels."
+        )
+
         # TurboQuant KV cache: route directly to TQ backend
         kv_cache_dtype = attn_selector_config.kv_cache_dtype
         if kv_cache_dtype is not None and kv_cache_dtype.startswith("turboquant_"):

vllm/utils/torch_utils.py

@@ -422,11 +422,11 @@ def nvfp4_split_data_scale(
     """Split one side (K or V) of an NVFP4 KV cache into data and scale.
 
     The input is a 4D uint8 tensor whose last dimension is
-    ``full_dim = data_dim + scale_dim``.  The physical layout within each
-    side is ``[data | scale]``, both packed contiguously.
+    ``full_dim = data_dim + scale_dim``.  The physical layout within
+    each side is ``[data | scale]``, both packed contiguously.
 
-    The caller is responsible for slicing K and V from the combined cache
-    first (e.g. ``kv_cache.split(num_kv_heads, dim=1)``).
+    The caller is responsible for slicing K and V from the combined
+    cache first (e.g. ``kv_cache.split(num_kv_heads, dim=1)``).
 
     Args:
         kv_side: 4D uint8 tensor ``(B, H, N, full_dim)``.
@@ -444,9 +444,6 @@ def nvfp4_split_data_scale(
     data_per_kv = dim_1 * dim_2 * data_dim
     page_bytes = kv_side.stride(0)
 
-    # Derive inner strides from the kv_side strides, scaling by the
-    # ratio of the target dim to full_dim.  This preserves the physical
-    # layout (NHD vs HND) encoded in the input tensor's strides.
     s1 = kv_side.stride(1) * data_dim // full_dim
     s2 = kv_side.stride(2) * data_dim // full_dim
     data_shape = (num_pages, dim_1, dim_2, data_dim)
@@ -460,7 +457,10 @@ def nvfp4_split_data_scale(
     base = kv_side.storage_offset()
     data = torch.as_strided(kv_side, data_shape, data_strides, storage_offset=base)
     scale = torch.as_strided(
-        kv_side, scale_shape, scale_strides, storage_offset=base + data_per_kv
+        kv_side,
+        scale_shape,
+        scale_strides,
+        storage_offset=base + data_per_kv,
     ).view(torch.float8_e4m3fn)
 
     return data, scale
@@ -476,14 +476,14 @@ def create_kv_caches_with_random_flash(
     model_dtype: str | torch.dtype | None = None,
     seed: int | None = None,
     device: str | None = "cuda",
-    cache_layout: str | None = "NHD",
+    cache_layout: str | None = "LBNHC",
 ) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
     set_random_seed(seed)
 
     dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
     generic_kv_cache_shape = (num_blocks, 2, block_size, num_heads, head_size)
-    assert cache_layout in ("NHD", "HND")
-    stride_order = (0, 1, 2, 3, 4) if cache_layout == "NHD" else (0, 1, 3, 2, 4)
+    assert cache_layout in ("LBNHC", "LBHNC")
+    stride_order = (0, 1, 2, 3, 4) if cache_layout == "LBNHC" else (0, 1, 3, 2, 4)
 
     kv_cache_allocation_shape = tuple(generic_kv_cache_shape[i] for i in stride_order)
     scale = head_size**-0.5

vllm/v1/attention/backend.py

@@ -24,7 +24,6 @@ if TYPE_CHECKING:
     from vllm.model_executor.layers.linear import ColumnParallelLinear
     from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
     from vllm.platforms.interface import DeviceCapability
-    from vllm.v1.attention.backends.utils import KVCacheLayoutType
     from vllm.v1.kv_cache_interface import AttentionSpec, KVQuantMode
 
 from vllm.v1.kv_cache_interface import get_kv_quant_mode
@@ -85,72 +84,14 @@ class AttentionBackend(ABC):
     def get_builder_cls():  # -> Type["AttentionMetadataBuilder"]:
         raise NotImplementedError
 
-    @staticmethod
-    @abstractmethod
-    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, ...]:
-        raise NotImplementedError
-
-    @classmethod
-    def get_kv_cache_block_dim(
-        cls,
-        block_size: int,
-        num_kv_heads: int,
-        head_size: int,
-        cache_dtype_str: str = "auto",
-    ) -> int:
-        """Discover which tensor dim is the block index, since different
-        backends lay out dims differently."""
-        _S = 1234567
-        shape = cls.get_kv_cache_shape(
-            _S,
-            block_size,
-            num_kv_heads,
-            head_size,
-            cache_dtype_str=cache_dtype_str,
-        )
-        return shape.index(_S)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        """
-        Get the physical (memory layout) ordering of the kv cache dimensions.
-        Standard attention backends pack K and V into the content dim, giving
-        the logical shape [num_blocks, num_heads, block_size, 2 * head_size].
-        e.g. if get_kv_cache_stride_order returns (0, 2, 1, 3) then the physical
-        ordering of dimensions is
-        [num_blocks, block_size, num_heads, 2 * head_size].
-
-        If this function is unimplemented / raises NotImplementedError,
-        the physical layout of the KV cache will match the logical shape.
-
-        Args:
-            include_num_layers_dimension: if True, includes an additional
-                num_layers dimension, which is assumed to be prepended
-                to the logical KV cache shape.
-                With the above example, a return value (1, 0, 3, 2, 4)
-                corresponds to
-                [num_blocks, num_layers, block_size, num_heads, 2 * head_size].
-
-                If an additional dimension is NOT included in the returned
-                tuple, the physical layout will not include a layers dimension.
-
-        Returns:
-            A tuple of ints which is a permutation of range(len(shape)).
-        """
-        raise NotImplementedError
-
     @classmethod
     def full_cls_name(cls) -> tuple[str, str]:
         return (cls.__module__, cls.__qualname__)
 
+    @classmethod
+    def get_required_kv_cache_layout(cls) -> str | None:
+        return None
+
     @classmethod
     def get_supported_head_sizes(cls) -> list[int]:
         return []
@@ -206,33 +147,12 @@ class AttentionBackend(ABC):
     def indexes_kv_by_block_stride(cls) -> bool:
         """Whether the backend reads KV pages by the runtime block stride.
 
-        True when ``num_blocks`` is the outermost physical dimension of the KV
-        cache, so the backend tolerates a non-contiguous block dim. This gates
-        page size padding and cross-layer uniform KV layout.
-
-        Returns:
-            True if the backend's physical KV layout is num-blocks-first. False
-            otherwise, including when the backend does not define a layered
-            stride order.
+        Under the standardized ``[L, B, H, N, C]`` layouts (RFC #42082) the
+        per-layer KV view is always blocks-first, so pages are addressed
+        through the view's block stride and page-size padding is tolerated.
+        This gates page size padding and cross-layer uniform KV layout.
         """
-        try:
-            kv_cache_stride_order = cls.get_kv_cache_stride_order(
-                include_num_layers_dimension=False
-            )
-            layered_kv_cache_stride_order = cls.get_kv_cache_stride_order(
-                include_num_layers_dimension=True
-            )
-        except (AttributeError, NotImplementedError):
-            return False
-
-        # Check that attention backend includes a layers dimension.
-        if len(layered_kv_cache_stride_order) != len(kv_cache_stride_order) + 1:
-            return False
-
-        # stride_order[0] == 0 means num_layers stays first in physical
-        # layout (identity permutation), so indexing by block stride is
-        # not supported.
-        return layered_kv_cache_stride_order[0] != 0
+        return True
 
     @classmethod
     def is_mla(cls) -> bool:
@@ -392,10 +312,6 @@ class AttentionBackend(ABC):
             invalid_reasons.append(combination_reason)
         return invalid_reasons
 
-    @classmethod
-    def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
-        return None
-
     @classmethod
     def is_ssm(cls) -> bool:
         return False

vllm/v1/attention/backends/cpu_attn.py

@@ -24,9 +24,7 @@ from vllm.v1.attention.backend import (
     CommonAttentionMetadata,
     MultipleOf,
 )
-from vllm.v1.attention.backends.utils import (
-    KVCacheLayoutType,
-)
+from vllm.v1.attention.backends.utils import resolve_kv_cache_layout
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
     CrossAttentionSpec,
@@ -35,6 +33,36 @@ from vllm.v1.kv_cache_interface import (
 
 logger = init_logger(__name__)
 
+_CPU_KV_CACHE_LAYOUTS = frozenset(("LBHNC", "BLHNC", "BHLNC"))
+
+
+def _split_cpu_kv_cache(
+    kv_cache: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Split a logical CPU KV cache into contiguous-token K and V views."""
+    layout = resolve_kv_cache_layout()
+    if layout.name not in _CPU_KV_CACHE_LAYOUTS:
+        raise ValueError(
+            f"CPU attention does not support KV cache layout {layout.name}. "
+            f"Supported layouts: {sorted(_CPU_KV_CACHE_LAYOUTS)}."
+        )
+    if kv_cache.ndim != 4:
+        raise ValueError(f"CPU KV cache must be 4D, got {kv_cache.ndim}D")
+
+    num_blocks, num_kv_heads, block_size, content_size = kv_cache.shape
+    if content_size % 2 != 0:
+        raise ValueError(f"CPU KV cache content size must be even, got {content_size}")
+    if kv_cache.stride(-1) != 1 or kv_cache.stride(-2) != content_size:
+        raise ValueError(
+            "CPU attention requires contiguous token and content dimensions; "
+            f"got shape {tuple(kv_cache.shape)} and strides {kv_cache.stride()}"
+        )
+
+    kv_cache = kv_cache.view(
+        num_blocks, num_kv_heads, block_size * 2, content_size // 2
+    )
+    return kv_cache.chunk(2, dim=2)
+
 
 class CPUAttentionBackend(AttentionBackend):
     forward_includes_kv_cache_update: bool = False
@@ -90,20 +118,6 @@ class CPUAttentionBackend(AttentionBackend):
     def get_builder_cls() -> type["CPUAttentionMetadataBuilder"]:
         return CPUAttentionMetadataBuilder
 
-    @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, num_kv_heads, block_size, 2 * head_size
-
-    @classmethod
-    def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
-        return "HND"
-
     @staticmethod
     def use_cascade_attention(*args, **kwargs) -> bool:
         return False
@@ -354,12 +368,7 @@ class CPUAttentionBackendImpl(AttentionImpl):
             # For encoder attention,
             kv_cache = attn_metadata.encoder_cache
 
-        # KV cache size are [num_blocks, num_kv_heads, block_size,
-        # 2 * head_size]. Make a view [num_blocks, num_kv_heads,
-        # block_size * 2, head_size]. Then slice KV at dim 2
-        num_blocks, num_kv_heads, block_size, _ = kv_cache.size()
-        kv_cache = kv_cache.view((num_blocks, num_kv_heads, block_size * 2, -1))
-        key_cache, value_cache = kv_cache.chunk(2, dim=2)
+        key_cache, value_cache = _split_cpu_kv_cache(kv_cache)
 
         # key and value may be None in the case of cross attention. They are
         # calculated once based on the output from the encoder and then cached
@@ -415,9 +424,7 @@ class CPUAttentionBackendImpl(AttentionImpl):
         if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
             return
 
-        num_blocks, num_kv_heads, block_size, _ = kv_cache.size()
-        kv_cache = kv_cache.view((num_blocks, num_kv_heads, block_size * 2, -1))
-        key_cache, value_cache = kv_cache.chunk(2, dim=2)
+        key_cache, value_cache = _split_cpu_kv_cache(kv_cache)
         ops.cpu_attn_reshape_and_cache(
             key,
             value,

vllm/v1/attention/backends/flash_attn.py

@@ -57,7 +57,6 @@ from vllm.v1.attention.backend import (
     AttentionMetadataBuilder,
     CommonAttentionMetadata,
 )
-from vllm.v1.attention.backends.utils import get_kv_cache_layout
 from vllm.v1.kv_cache_interface import AttentionSpec
 from vllm.v1.worker.cp_utils import (
     run_split_fa2_dcp_context_attention,
@@ -130,43 +129,6 @@ class FlashAttentionBackend(AttentionBackend):
     def get_builder_cls() -> type["FlashAttentionMetadataBuilder"]:
         return FlashAttentionMetadataBuilder
 
-    @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 block_size % 16 != 0:
-            raise ValueError("Block size must be a multiple of 16.")
-        # K and V are packed into the content dim: logical (B, H, N, 2*D).
-        return (num_blocks, num_kv_heads, block_size, 2 * head_size)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        # `stride_order` indicates the permutation that gets us from
-        # `get_kv_cache_shape` (logical (B, H, N, 2*D)) to the actual memory
-        # layout we want.
-        cache_layout = get_kv_cache_layout()
-        if cache_layout == "NHD" and include_num_layers_dimension:
-            # (num_blocks, num_layers, block_size, num_kv_heads, 2*head_size)
-            return (1, 0, 3, 2, 4)
-        elif cache_layout == "NHD":
-            # (num_blocks, block_size, num_kv_heads, 2*head_size)
-            stride_order = (0, 2, 1, 3)
-        elif cache_layout == "HND" and include_num_layers_dimension:
-            # (num_blocks, num_kv_heads, num_layers, block_size, 2*head_size)
-            return (1, 2, 0, 3, 4)
-        elif cache_layout == "HND":
-            # (num_blocks, num_kv_heads, block_size, 2*head_size)
-            stride_order = (0, 1, 2, 3)
-        else:
-            raise ValueError(f"Unknown cache layout format {cache_layout}.")
-        return stride_order
-
     @classmethod
     def supports_head_size(cls, head_size: int) -> bool:
         if head_size % 8 != 0:
@@ -853,7 +815,7 @@ class FlashAttentionImpl(AttentionImpl):
             key: shape = [num_tokens, num_kv_heads, head_size]
             value: shape = [num_tokens, num_kv_heads, head_size]
             kv_cache: shape =
-                [num_blocks, num_kv_heads, block_size, 2 * head_size]
+                [num_blocks, num_kv_heads, block_size, 2*head_size]
             attn_metadata: Metadata for attention.
         Returns:
             shape = [num_tokens, num_heads * head_size]
@@ -900,7 +862,7 @@ class FlashAttentionImpl(AttentionImpl):
                 layer,
             )
 
-        # (B, H, N, 2*D) -> ((B, N, H, D), (B, N, H, D))
+        # (B, H, N, 2*head_size) -> ((B, N, H, head_size), (B, N, H, head_size))
         key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
         # Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP).
         # FA3/4 on H100+ uses TMA, which requires ≥16-byte stride alignment.
@@ -1108,9 +1070,7 @@ class FlashAttentionImpl(AttentionImpl):
             return
 
         # Scatter write into the KV cache using slot_mapping indices.
-        # No TMA kernel is invoked here, so stride canonicalization is not needed.
-        # (B, H, N, 2*D) -> ((B, N, H, D), (B, N, H, D))
-        key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
+        k_cache, v_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
 
         # Reshape the input keys and values and store them in the cache.
         # Skip this if sharing KV cache with an earlier attention layer.
@@ -1122,8 +1082,8 @@ class FlashAttentionImpl(AttentionImpl):
         reshape_and_cache_flash(
             key,
             value,
-            key_cache,
-            value_cache,
+            k_cache,
+            v_cache,
             slot_mapping,
             self.kv_cache_dtype,
             layer._k_scale,
@@ -1629,10 +1589,11 @@ def cascade_attention(
 
     num_tokens = query.shape[0]
     block_size = key_cache.shape[-3]
+    num_kv_heads = key_cache.shape[-2]
     assert common_prefix_len % block_size == 0
     num_common_kv_blocks = common_prefix_len // block_size
     assert num_common_kv_blocks > 0
-    descale_shape = (cu_prefix_query_lens.shape[0] - 1, key_cache.shape[-2])
+    descale_shape = (cu_prefix_query_lens.shape[0] - 1, num_kv_heads)
 
     # Process shared prefix.
     prefix_output, prefix_lse = flash_attn_varlen_func(
@@ -1660,7 +1621,7 @@ def cascade_attention(
         num_splits=1 if envs.VLLM_BATCH_INVARIANT else max_num_splits,
     )
 
-    descale_shape = (cu_query_lens.shape[0] - 1, key_cache.shape[-2])
+    descale_shape = (cu_query_lens.shape[0] - 1, num_kv_heads)
 
     # Process suffix per query.
     suffix_output, suffix_lse = flash_attn_varlen_func(

vllm/v1/attention/backends/flash_attn_diffkv.py

@@ -21,8 +21,6 @@ from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
 
 if is_flash_attn_varlen_func_available():
     from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func
-from vllm.v1.attention.backends.utils import get_kv_cache_layout
-
 from .flash_attn import (
     FlashAttentionBackend,
     FlashAttentionImpl,
@@ -73,49 +71,6 @@ class FlashAttentionDiffKVBackend(FlashAttentionBackend):
     def get_impl_cls() -> type["FlashAttentionImpl"]:
         return FlashAttentionDiffKVImpl
 
-    @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 block_size % 16 != 0:
-            raise ValueError("Block size must be a multiple of 16.")
-        # Logical (blocks-first, head-major) layout: K and V (with their
-        # different head sizes) packed in the content dim.
-        return (
-            num_blocks,
-            num_kv_heads,
-            block_size,
-            head_size + FlashAttentionDiffKVBackend.head_size_v,
-        )
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        # `stride_order` indicates the permutation that gets us from
-        # `get_kv_cache_shape` (logical (B, H, N, C_k+C_v)) to the actual
-        # memory layout we want.
-        cache_layout = get_kv_cache_layout()
-        if cache_layout == "NHD" and include_num_layers_dimension:
-            # (num_blocks, num_layers, block_size, num_kv_heads, C_k+C_v)
-            return (1, 0, 3, 2, 4)
-        elif cache_layout == "NHD":
-            # (num_blocks, block_size, num_kv_heads, C_k+C_v)
-            stride_order = (0, 2, 1, 3)
-        elif cache_layout == "HND" and include_num_layers_dimension:
-            # (num_blocks, num_kv_heads, num_layers, block_size, C_k+C_v)
-            return (1, 2, 0, 3, 4)
-        elif cache_layout == "HND":
-            # (num_blocks, num_kv_heads, block_size, C_k+C_v)
-            stride_order = (0, 1, 2, 3)
-        else:
-            raise ValueError(f"Unknown cache layout format {cache_layout}.")
-        return stride_order
-
 
 class FlashAttentionDiffKVImpl(FlashAttentionImpl):
     vllm_flash_attn_version: int | None
@@ -225,6 +180,7 @@ class FlashAttentionDiffKVImpl(FlashAttentionImpl):
                 layer,
             )
 
+        # (B, H, N, C) -> (B, N, H, C) for kernel compatibility.
         # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
         key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
         # Fix degenerate strides on size-1 dims (e.g. num_kv_heads=1 with TP).

vllm/v1/attention/backends/flashinfer.py

@@ -50,7 +50,6 @@ from vllm.utils.torch_utils import (
     canonicalize_singleton_dim_strides,
     is_quantized_kv_cache,
     is_strictly_contiguous,
-    nvfp4_kv_cache_full_dim,
     nvfp4_split_data_scale,
 )
 from vllm.v1.attention.backend import (
@@ -63,12 +62,12 @@ from vllm.v1.attention.backend import (
     MultipleOf,
 )
 from vllm.v1.attention.backends.utils import (
-    KVCacheLayoutType,
     get_dcp_local_seq_lens,
-    get_kv_cache_layout,
+    get_flashinfer_layout_string,
     get_num_attention_heads_from_layers,
     get_per_layer_parameters,
     infer_global_hyperparameters,
+    resolve_kv_cache_layout,
     split_decodes_and_prefills,
 )
 from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
@@ -238,7 +237,7 @@ class BatchDCPPrefillWrapper:
         else:
             self._dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False)
         self._context = BatchPrefillWithPagedKVCacheWrapper(
-            workspace_buffer, get_kv_cache_layout()
+            workspace_buffer, get_flashinfer_layout_string()
         )
         self._new_tokens = BatchPrefillWithRaggedKVCacheWrapper(workspace_buffer)
 
@@ -393,44 +392,6 @@ class FlashInferBackend(AttentionBackend):
     def get_builder_cls() -> type["FlashInferMetadataBuilder"]:
         return FlashInferMetadataBuilder
 
-    @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 == "nvfp4":
-            full_dim = nvfp4_kv_cache_full_dim(head_size)
-            return (num_blocks, 2 * num_kv_heads, block_size, full_dim)
-        # Pack K and V in the content dim (B, H, N, 2*hs).
-        return (num_blocks, num_kv_heads, block_size, 2 * head_size)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        # `stride_order` indicates the permutation that gets us from
-        # `get_kv_cache_shape` (logical (B, H, N, 2*hs)) to the actual memory
-        # layout we want.
-        cache_layout = get_kv_cache_layout()
-        if cache_layout == "NHD" and include_num_layers_dimension:
-            # (num_blocks, num_layers, block_size, num_kv_heads, 2*head_size)
-            return (1, 0, 3, 2, 4)
-        elif cache_layout == "NHD":
-            # (num_blocks, block_size, num_kv_heads, 2*head_size)
-            stride_order = (0, 2, 1, 3)
-        elif cache_layout == "HND" and include_num_layers_dimension:
-            # (num_blocks, num_kv_heads, num_layers, block_size, 2*head_size)
-            return (1, 2, 0, 3, 4)
-        elif cache_layout == "HND":
-            # (num_blocks, num_kv_heads, block_size, 2*head_size)
-            stride_order = (0, 1, 2, 3)
-        else:
-            raise ValueError(f"Unknown cache layout format {cache_layout}.")
-        return stride_order
-
     @staticmethod
     def get_dtype_for_flashinfer(kv_cache_dtype: str) -> torch.dtype:
         if kv_cache_dtype in ("fp8", "fp8_e4m3"):
@@ -489,7 +450,7 @@ class FlashInferBackend(AttentionBackend):
         ) and supports_trtllm_attention(is_prefill=True)
 
     @classmethod
-    def get_required_kv_cache_layout(cls) -> KVCacheLayoutType | None:
+    def get_required_kv_cache_layout(cls) -> str | None:
         capability = current_platform.get_device_capability()
         if capability is not None and capability.major == 10:
             return "HND"
@@ -985,7 +946,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
             if self._noncausal_prefill_wrapper is None:
                 self._noncausal_prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
                     self._get_workspace_buffer(),
-                    get_kv_cache_layout(),
+                    get_flashinfer_layout_string(),
                     backend="auto",
                 )
             return self._noncausal_prefill_wrapper
@@ -1002,7 +963,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
                 backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto"
                 self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
                     self._get_workspace_buffer(),
-                    get_kv_cache_layout(),
+                    get_flashinfer_layout_string(),
                     backend=backend,
                 )
         assert self._prefill_wrapper is not None
@@ -1028,7 +989,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
             backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto"
             decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
                 self._get_workspace_buffer(),
-                get_kv_cache_layout(),
+                get_flashinfer_layout_string(),
                 use_cuda_graph=use_cudagraph,
                 paged_kv_indptr_buffer=paged_kv_indptr,
                 paged_kv_indices_buffer=paged_kv_indices,
@@ -1051,7 +1012,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
     def _get_cascade_wrapper(self):
         if self._cascade_wrapper is None:
             self._cascade_wrapper = MultiLevelCascadeAttentionWrapper(
-                2, self._get_workspace_buffer(), get_kv_cache_layout()
+                2, self._get_workspace_buffer(), get_flashinfer_layout_string()
             )
         return self._cascade_wrapper
 
@@ -1694,9 +1655,7 @@ class FlashInferImpl(AttentionImpl):
             query: shape = [num_tokens, num_heads, head_size]
             key: shape = [num_tokens, num_kv_heads, head_size]
             value: shape = [num_tokens, num_kv_heads, head_size]
-            kv_cache: KV cache tensor with different possible shapes:
-                - NHD: [num_blocks, 2, block_size, num_kv_heads, head_size]
-                - HND: [num_blocks, 2, num_kv_heads, block_size, head_size]
+            kv_cache: [num_blocks, num_kv_heads, block_size, 2*head_size]
             attn_metadata: Metadata for attention.
         Returns:
             shape = [num_tokens, num_heads * head_size]
@@ -1791,31 +1750,9 @@ class FlashInferImpl(AttentionImpl):
         output_padded = output
         output = output[:num_actual_tokens]
 
-        if attn_metadata.use_cascade:
-            # Cascade attention (rare case).
-            assert attn_metadata.cascade_wrapper is not None
-            stride_order = FlashInferBackend.get_kv_cache_stride_order()
-            if self.is_kvcache_nvfp4:
-                kv_cache_views = tuple(
-                    cache.permute(*stride_order)
-                    for cache in kv_cache.split(self.num_kv_heads, dim=1)
-                )
-            else:
-                kv_perm = kv_cache.permute(*stride_order)
-                kv_cache_views = kv_perm.split(self.head_size, dim=-1)
-            kv_tuple = tuple(
-                canonicalize_singleton_dim_strides(cache) for cache in kv_cache_views
-            )
-            output.copy_(attn_metadata.cascade_wrapper.run(query, kv_tuple))
-            return output
-
-        # When using spec decoding, num_decodes can be < num_decode_tokens
-        # because some decode requests may have more than one query token.
-        num_decode_tokens = attn_metadata.num_decode_tokens
-        num_prefill_tokens = attn_metadata.num_prefill_tokens
-
-        stride_order = FlashInferBackend.get_kv_cache_stride_order()
-        kv_cache_permute = kv_cache.permute(*stride_order)  # HND and contiguous
+        # Permute to FlashInfer's expected layout (metadata-only).
+        stride_order = resolve_kv_cache_layout().layer_view_order
+        kv_cache_permute = kv_cache.permute(*stride_order)
         # Fix degenerate strides on any size-1 dimension (e.g. num_kv_heads=1
         # with TP=8).  PyTorch permits non-canonical strides on size-1 dims;
         # CUDA TMA requires ≥16-byte alignment on all non-outermost strides.
@@ -1850,6 +1787,14 @@ class FlashInferImpl(AttentionImpl):
         else:
             kv_cache_tuple = kv_cache_permute.split(hs, dim=-1)
 
+        if attn_metadata.use_cascade:
+            assert attn_metadata.cascade_wrapper is not None
+            output.copy_(attn_metadata.cascade_wrapper.run(query, kv_cache_tuple))
+            return output
+
+        num_decode_tokens = attn_metadata.num_decode_tokens
+        num_prefill_tokens = attn_metadata.num_prefill_tokens
+
         use_dcp = self.dcp_world_size > 1
 
         # Regular attention (common case).
@@ -1949,7 +1894,7 @@ class FlashInferImpl(AttentionImpl):
                 seq_lens_prefill = attn_metadata.prefill.seq_lens
 
                 # This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND
-                assert get_kv_cache_layout() == "HND"
+                assert get_flashinfer_layout_string() == "HND"
                 assert is_strictly_contiguous(prefill_query)
                 assert is_strictly_contiguous(workspace_buffer)
                 assert is_strictly_contiguous(block_tables_prefill)
@@ -2132,7 +2077,7 @@ class FlashInferImpl(AttentionImpl):
                 # trtllm-gen needs HND layout on SM100. XQA is selected
                 # separately on SM90 and does not use this SM100 layout gate.
                 if decode_with_trtllm_gen:
-                    assert get_kv_cache_layout() == "HND"
+                    assert get_flashinfer_layout_string() == "HND"
                 else:
                     assert decode_with_xqa
                 assert is_strictly_contiguous(decode_query)
@@ -2226,7 +2171,7 @@ class FlashInferImpl(AttentionImpl):
                     sinks=self.sinks,
                     o_sf_scale=self.o_sf_scale,
                     out=out,
-                    kv_layout=get_kv_cache_layout(),
+                    kv_layout=get_flashinfer_layout_string(),
                     backend=attn_metadata.decode.kernel.value,
                     q_len_per_req=q_len_per_req,
                     kv_cache_sf=(

vllm/v1/attention/backends/flex_attention.py

@@ -100,6 +100,12 @@ class FlexAttentionBackend(AttentionBackend):
     def get_name() -> str:
         return "FLEX_ATTENTION"
 
+    @classmethod
+    def get_required_kv_cache_layout(cls) -> str | None:
+        if current_platform.is_rocm():
+            return "LBNHC"
+        return "BLNHC"
+
     @classmethod
     def supports_sliding_window(cls) -> bool:
         return True
@@ -126,25 +132,6 @@ class FlexAttentionBackend(AttentionBackend):
     def get_impl_cls() -> type["FlexAttentionImpl"]:
         return FlexAttentionImpl
 
-    @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, ...]:
-        # K and V are packed into the content dim: logical (B, H, N, 2*hs).
-        return (num_blocks, num_kv_heads, block_size, 2 * head_size)
-
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        if include_num_layers_dimension:
-            return (1, 0, 3, 2, 4)
-        return (0, 2, 1, 3)
-
     @staticmethod
     def get_builder_cls() -> type["FlexAttentionMetadataBuilder"]:
         return FlexAttentionMetadataBuilder
@@ -1279,7 +1266,7 @@ class FlexAttentionImpl(AttentionImpl):
             key: shape = [num_tokens, num_kv_heads, head_size]
             value: shape = [num_tokens, num_kv_heads, head_size]
             kv_cache: shape =
-                [num_blocks, num_kv_heads, block_size, 2 * head_size]
+                [num_blocks, 2, block_size, num_kv_heads, head_size]
             attn_metadata: Metadata for attention.
         Returns:
             shape = [num_tokens, num_heads * head_size]
@@ -1354,12 +1341,12 @@ class FlexAttentionImpl(AttentionImpl):
         else:
             assert self.attn_type == AttentionType.DECODER
             kv_cache = kv_cache.transpose(1, 2)
-            hs = self.head_size
-            key_cache, value_cache = kv_cache.split(hs, dim=-1)
+            key_cache, value_cache = kv_cache.split(self.head_size, dim=-1)
 
             # Flatten (num_blocks, block_size) into a single token dim.
-            key_cache = key_cache.view(-1, self.num_kv_heads, self.head_size)
-            value_cache = value_cache.view(-1, self.num_kv_heads, self.head_size)
+            # The transposed views are non-contiguous, so use reshape.
+            key_cache = key_cache.reshape(-1, self.num_kv_heads, self.head_size)
+            value_cache = value_cache.reshape(-1, self.num_kv_heads, self.head_size)
             query, key_tensor, value_tensor = map(
                 lambda x: self.view_as_4d(x).permute(0, 2, 1, 3),
                 (query, key_cache, value_cache),

vllm/v1/attention/backends/hpc_attn.py

@@ -249,7 +249,9 @@ class HpcAttnMetadataBuilder(AttentionMetadataBuilder[HpcAttnMetadata]):
 class HpcAttentionBackend(AttentionBackend):
     """HPC attention backend (pure attention, no RoPE/Norm).
 
-    KV cache layout: NHD (num_blocks, 2, block_size, num_kv_heads, head_size).
+    KV cache layout: NHD; the logical per-layer view is
+    (num_blocks, num_kv_heads, block_size, 2 * head_size) with K and V
+    packed along the content dim.
     """
 
     accept_output_buffer: bool = True
@@ -282,24 +284,6 @@ class HpcAttentionBackend(AttentionBackend):
     def get_builder_cls() -> type["HpcAttnMetadataBuilder"]:
         return HpcAttnMetadataBuilder
 
-    @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, ...]:
-        if include_num_layers_dimension:
-            return (1, 0, 2, 3, 4, 5)
-        return (0, 1, 2, 3, 4)
-
     @classmethod
     def get_supported_head_sizes(cls) -> list[int]:
         return [128]
@@ -316,7 +300,7 @@ class HpcAttentionBackend(AttentionBackend):
 
     @classmethod
     def get_required_kv_cache_layout(cls) -> KVCacheLayoutType | None:
-        return "NHD"
+        return "LBNHC"
 
 
 class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
@@ -436,13 +420,16 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
         num_decode_reqs = attn_metadata.num_decodes
         num_decode_tokens = attn_metadata.num_decode_tokens
 
+        # Logical (B, H, N, 2*head_size) -> (B, N, H, head_size) K/V views.
+        key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
+
         # Write KV cache if not already done by HpcRopeNorm.
         if self.kv_sharing_target_layer_name is None and not hpc_kv_written:
             torch.ops._C_cache_ops.reshape_and_cache_flash(
                 key,
                 value,
-                kv_cache[:, 0],
-                kv_cache[:, 1],
+                key_cache,
+                value_cache,
                 attn_metadata.slot_mapping,
                 self.kv_cache_dtype,
                 layer._k_scale,
@@ -451,7 +438,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
 
         if self.use_fp8:
             torch_dtype = _get_fp8_dtype_for_kv_cache(self.kv_cache_dtype)
-            kv_cache = kv_cache.view(torch_dtype)
+            key_cache = key_cache.view(torch_dtype)
+            value_cache = value_cache.view(torch_dtype)
 
         if self.use_fp8:
             if not hpc_kv_written:
@@ -483,8 +471,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
             if self.use_fp8:
                 hpc.attention_with_kvcache_prefill_fp8(
                     q_prefill,
-                    kv_cache[:, 0],
-                    kv_cache[:, 1],
+                    key_cache,
+                    value_cache,
                     hpc_prefill_q_scale,
                     k_scale,
                     v_scale,
@@ -498,8 +486,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
             else:
                 hpc.attention_with_kvcache_prefill_bf16(
                     q_prefill,
-                    kv_cache[:, 0],
-                    kv_cache[:, 1],
+                    key_cache,
+                    value_cache,
                     cu_seqlens_prefill,
                     block_table_prefill,
                     seq_lens_prefill,
@@ -520,8 +508,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
             if self.use_fp8:
                 hpc.attention_decode_fp8(
                     q_decode,
-                    kv_cache[:, 0],
-                    kv_cache[:, 1],
+                    key_cache,
+                    value_cache,
                     block_table_decode,
                     num_seq_kvcache,
                     hpc_decode_q_scale,
@@ -541,8 +529,8 @@ class HpcAttentionImpl(AttentionImpl[HpcAttnMetadata]):
             else:
                 hpc.attention_decode_bf16(
                     q_decode,
-                    kv_cache[:, 0],
-                    kv_cache[:, 1],
+                    key_cache,
+                    value_cache,
                     block_table_decode,
                     num_seq_kvcache,
                     mtp=mtp,

vllm/v1/attention/backends/mla/cutlass_mla.py

@@ -49,14 +49,6 @@ class CutlassMLABackend(MLACommonBackend):
     def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
         return [128]
 
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        if include_num_layers_dimension:
-            return (1, 0, 2, 3)
-        return (0, 1, 2)
-
     @staticmethod
     def get_name() -> str:
         return "CUTLASS_MLA"

vllm/v1/attention/backends/mla/flashattn_mla.py

@@ -52,14 +52,6 @@ class FlashAttnMLABackend(MLACommonBackend):
     def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
         return [MultipleOf(16)]
 
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        if include_num_layers_dimension:
-            return (1, 0, 2, 3)
-        return (0, 1, 2)
-
     @staticmethod
     def get_name() -> str:
         return "FLASH_ATTN_MLA"
@@ -335,6 +327,7 @@ class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):
         if is_quantized_kv_cache(self.kv_cache_dtype):
             raise NotImplementedError("FP8 FlashAttention MLA not yet supported")
 
+        kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.unsqueeze(2)
         kv_c_cache = kv_c_and_k_pe_cache[..., : self.kv_lora_rank]
         k_pe_cache = kv_c_and_k_pe_cache[..., self.kv_lora_rank :]
 
@@ -345,8 +338,8 @@ class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):
 
         attn_out = flash_attn_varlen_func(
             q=q_pe,
-            k=k_pe_cache.unsqueeze(-2),  # Add head dim of 1
-            v=kv_c_cache.unsqueeze(-2),  # Add head dim of 1
+            k=k_pe_cache,
+            v=kv_c_cache,
             q_v=q_nope,
             max_seqlen_q=max_seqlen_q,
             cu_seqlens_q=attn_metadata.decode.query_start_loc,

vllm/v1/attention/backends/mla/flashattn_mla_sparse.py

@@ -103,16 +103,6 @@ class FlashAttnMLASparseBackend(AttentionBackend):
                 return "FlashAttention MLA Sparse requires model with index_topk"
         return None
 
-    @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)
-
 
 @dataclass
 class FlashAttnMLASparseMetadata(AttentionMetadata):

vllm/v1/attention/backends/mla/flashinfer_mla.py

@@ -23,7 +23,6 @@ from vllm.v1.attention.backend import (
     AttentionType,
     MultipleOf,
 )
-from vllm.v1.attention.backends.utils import KVCacheLayoutType
 
 logger = init_logger(__name__)
 
@@ -67,14 +66,6 @@ class FlashInferMLABackend(MLACommonBackend):
     def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
         return [32, 64]
 
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        if include_num_layers_dimension:
-            return (1, 0, 2, 3)
-        return (0, 1, 2)
-
     @staticmethod
     def get_name() -> str:
         return "FLASHINFER_MLA"
@@ -118,10 +109,6 @@ class FlashInferMLABackend(MLACommonBackend):
                 )
         return None
 
-    @classmethod
-    def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
-        return "HND"
-
 
 class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
     can_return_lse_for_decode: bool = True
@@ -219,7 +206,7 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
         workspace_buffer = _get_workspace_buffer(return_lse)
         kernel_out = trtllm_batch_decode_with_kv_cache_mla(
             query=q,
-            kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
+            kv_cache=kv_c_and_k_pe_cache,
             workspace_buffer=workspace_buffer,
             qk_nope_head_dim=self.qk_nope_head_dim,
             kv_lora_rank=self.kv_lora_rank,

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

@@ -31,7 +31,6 @@ from vllm.v1.attention.backends.mla.sparse_utils import (
     triton_convert_req_index_to_global_index,
     triton_filter_and_convert_dcp_index,
 )
-from vllm.v1.attention.backends.utils import KVCacheLayoutType
 from vllm.v1.kv_cache_interface import AttentionSpec
 
 if TYPE_CHECKING:
@@ -123,20 +122,6 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
                 return "FlashInfer MLA Sparse requires model with index_topk config"
         return None
 
-    @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)
-
-    @classmethod
-    def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
-        return "HND"
-
 
 class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase):
     """FlashInfer sparse MLA backend for SM120."""
@@ -216,23 +201,6 @@ class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase):
                 )
         return None
 
-    @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, ...]:
-        if cache_dtype_str in ("auto", "fp8", "fp8_e4m3", "fp8_ds_mla"):
-            # fp8_ds_mla packed layout: 512 NoPE + 16 scales + 128 RoPE.
-            return (num_blocks, block_size, 656)
-        return (num_blocks, block_size, head_size)
-
-    @classmethod
-    def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
-        return None
-
 
 @dataclass
 class FlashInferMLASparseMetadata(AttentionMetadata):
@@ -439,7 +407,7 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
 
         kernel_out = trtllm_batch_decode_with_kv_cache_mla(
             query=query,
-            kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
+            kv_cache=kv_c_and_k_pe_cache,
             workspace_buffer=self._workspace_buffer,
             qk_nope_head_dim=self.qk_nope_head_dim,
             kv_lora_rank=self.kv_lora_rank,

vllm/v1/attention/backends/mla/flashmla.py

@@ -58,14 +58,6 @@ class FlashMLABackend(MLACommonBackend):
     def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
         return [64]
 
-    @staticmethod
-    def get_kv_cache_stride_order(
-        include_num_layers_dimension: bool = False,
-    ) -> tuple[int, ...]:
-        if include_num_layers_dimension:
-            return (1, 0, 2, 3)
-        return (0, 1, 2)
-
     @staticmethod
     def get_name() -> str:
         return "FLASHMLA"
@@ -317,7 +309,7 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
         if is_quantized_kv_cache(self.kv_cache_dtype):
             o, lse = flash_mla_with_kvcache_fp8(
                 q=q,
-                k_cache=kv_c_and_k_pe_cache.unsqueeze(-2),  # Add head dim of 1
+                k_cache=kv_c_and_k_pe_cache.unsqueeze(2),
                 block_table=attn_metadata.decode.block_table,
                 cache_seqlens=attn_metadata.decode.seq_lens,
                 head_dim_v=self.kv_lora_rank,
@@ -331,7 +323,7 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
         else:
             o, lse = flash_mla_with_kvcache(
                 q=q,
-                k_cache=kv_c_and_k_pe_cache.unsqueeze(-2),  # Add head dim of 1
+                k_cache=kv_c_and_k_pe_cache.unsqueeze(2),
                 block_table=attn_metadata.decode.block_table,
                 cache_seqlens=attn_metadata.decode.seq_lens,
                 head_dim_v=self.kv_lora_rank,

vllm/v1/attention/backends/mla/flashmla_sparse.py

@@ -16,6 +16,7 @@ from vllm.model_executor.layers.attention.sparse_mla_attention import (
 )
 from vllm.platforms import current_platform
 from vllm.platforms.interface import DeviceCapability
+from vllm.utils.math_utils import cdiv
 from vllm.utils.platform_utils import num_compute_units
 from vllm.utils.torch_utils import is_quantized_kv_cache
 from vllm.v1.attention.backend import (
@@ -127,19 +128,15 @@ class FlashMLASparseBackend(AttentionBackend):
     def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
         return capability.major in [9, 10]
 
+
+class DeepseekV4FlashMLASparseBackend(FlashMLASparseBackend):
     @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, ...]:
-        if cache_dtype_str == "fp8_ds_mla":
-            # V3.2 main MLA: 656-byte custom storage format. See module docstring.
-            return (num_blocks, block_size, 656)
-        else:
-            return (num_blocks, block_size, head_size)
+    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+        return [256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "V4_FLASHMLA_SPARSE"
 
 
 @dataclass
@@ -313,6 +310,67 @@ class FlashMLASparseMetadataBuilder(
             device=device,
         )
 
+        # DeepseekV4: has compress_ratios in hf_config.
+        hf_config = vllm_config.model_config.hf_config
+        self.is_deepseek_v4 = (
+            hasattr(hf_config, "compress_ratios") and len(hf_config.compress_ratios) > 0
+        )
+        self.compress_ratio = 1
+        if self.is_deepseek_v4:
+            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.
+            if self.compress_ratio > 1:
+                max_num_batched_tokens = (
+                    vllm_config.scheduler_config.max_num_batched_tokens
+                )
+                self.compressed_slot_mapping_buffer = torch.empty(
+                    max_num_batched_tokens,
+                    dtype=torch.int64,
+                    device=self.device,
+                )
+
+            # Pre-allocate C128A topk buffers for CUDA graph address stability.
+            if self.compress_ratio == 128:
+                max_num_batched_tokens = (
+                    vllm_config.scheduler_config.max_num_batched_tokens
+                )
+                # Pad to B_TOPK alignment (128 covers both h_q=64 B_TOPK=64 and
+                # h_q=128 B_TOPK=128). FlashMLA decode asserts extra_topk % B_TOPK
+                # == 0; unaligned widths (e.g. 17 = ceil(2136/128)) crash the
+                # sm100 head64 kernel. Padded slots stay -1 and decode_lens caps
+                # them via topk_length, so the pad is a no-op at kernel level.
+                # Mirrors _SPARSE_PREFILL_TOPK_ALIGNMENT in cache_utils.py.
+                _C128A_TOPK_ALIGNMENT = 128
+                c128a_max_compressed = cdiv(
+                    self.model_config.max_model_len, self.compress_ratio
+                )
+                c128a_max_compressed = (
+                    cdiv(c128a_max_compressed, _C128A_TOPK_ALIGNMENT)
+                    * _C128A_TOPK_ALIGNMENT
+                )
+                # Stored so _build_c128a_metadata passes it as the kernel's
+                # max_compressed_tokens, matching the buffer stride. Otherwise
+                # the kernel's default 8192 iterates past row width and spills
+                # writes into adjacent rows (present in both decode and prefill
+                # branches of _build_c128a_topk_metadata_kernel).
+                self.c128a_max_compressed = c128a_max_compressed
+                self.c128a_global_decode_buffer = torch.empty(
+                    (max_num_batched_tokens, c128a_max_compressed),
+                    dtype=torch.int32,
+                    device=self.device,
+                )
+                self.c128a_decode_lens_buffer = torch.empty(
+                    max_num_batched_tokens,
+                    dtype=torch.int32,
+                    device=self.device,
+                )
+                self.c128a_prefill_buffer = torch.empty(
+                    (max_num_batched_tokens, c128a_max_compressed),
+                    dtype=torch.int32,
+                    device=self.device,
+                )
+
     def _build_fp8_mixed_decode_prefill(
         self,
         common_attn_metadata: CommonAttentionMetadata,
@@ -783,7 +841,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
 
         out, lse = flash_mla_with_kvcache(
             q=q,
-            k_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(-2),
+            k_cache=kv_c_and_k_pe_cache.view(torch.uint8).unsqueeze(2),
             block_table=kernel_metadata.dummy_block_table,
             head_dim_v=512,
             cache_seqlens=kernel_metadata.cache_lens,
@@ -808,7 +866,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
     ) -> torch.Tensor:
         num_tokens = q.shape[0]
         kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.view(
-            -1, 1, kv_c_and_k_pe_cache.shape[-1]
+            -1, 1, 1, kv_c_and_k_pe_cache.shape[-1]
         )
 
         # NOTE(Chen): kernel requires num_local_head to be a multiple of

vllm/v1/attention/backends/mla/indexer.py

@@ -146,28 +146,6 @@ class DeepseekV32IndexerBackend(AttentionBackend):
     def get_builder_cls() -> type["DeepseekV32IndexerMetadataBuilder"]:
         return DeepseekV32IndexerMetadataBuilder
 
-    @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:
-            # DeepseekV32Indexer kernels do not support cross-layer
-            # KV cache layout. Identity permutation keeps num_layers
-            # first, signaling incompatibility.
-            return (0, 1, 2, 3)
-        return (0, 1, 2)
-
 
 class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
     @staticmethod
@@ -578,7 +556,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
         self.compress_ratio = 1
         # Get compress_ratio for DeepseekV4 support
         if isinstance(self.kv_cache_spec, MLAAttentionSpec):
-            self.compress_ratio = self.kv_cache_spec.compress_ratio
+            self.compress_ratio = self.kv_cache_spec.tokens_per_state
         if self.dcp_world_size > 1 and self.compress_ratio > 1:
             raise NotImplementedError(
                 "DCP is not supported with sparse indexer KV compression "

vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py

@@ -292,16 +292,6 @@ class ROCMAiterMLASparseBackend(AttentionBackend):
     def get_impl_cls() -> type["ROCMAiterMLASparseImpl"]:
         return ROCMAiterMLASparseImpl
 
-    @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)
-
     @classmethod
     def is_mla(cls) -> bool:
         return True

vllm/v1/attention/backends/mla/sparse_swa.py

@@ -84,6 +84,10 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
         # contiguous full-cache layout.
         assert self.dtype in (torch.uint8, torch.bfloat16, torch.float8_e4m3fn)
 
+    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's UE8M0 paged layout needs 576B alignment; contiguous
         # bf16/fp8 cache uses the natural element-size page.
@@ -134,30 +138,6 @@ class DeepseekSparseSWABackend(AttentionBackend):
             return DeepseekV4ROCMAiterSparseSWAMetadataBuilder
         return DeepseekSparseSWAMetadataBuilder
 
-    @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
-        if cache_dtype_str == "fp8_ds_mla":
-            # DeepseekV4 SWA: 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)
-
-    @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 DeepseekSparseSWAMetadata:
@@ -406,7 +386,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
         assert isinstance(self.kv_cache_spec, SlidingWindowMLASpec | MLAAttentionSpec)
         mla_spec = cast(SlidingWindowMLASpec | MLAAttentionSpec, self.kv_cache_spec)
         self.head_size = mla_spec.head_size  # Already considered quantization.
-        self.compress_ratio = mla_spec.compress_ratio
+        self.compress_ratio = mla_spec.tokens_per_state
         self.block_size = mla_spec.block_size
         self.max_model_len = self.vllm_config.model_config.max_model_len
         self.max_num_batched_tokens = (

vllm/v1/attention/backends/mla/tokenspeed_mla.py

@@ -23,7 +23,6 @@ from vllm.v1.attention.backend import (
     AttentionType,
     MultipleOf,
 )
-from vllm.v1.attention.backends.utils import KVCacheLayoutType
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
@@ -144,10 +143,6 @@ class TokenspeedMLABackend(MLACommonBackend):
                 )
         return None
 
-    @classmethod
-    def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":
-        return "HND"
-
 
 class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):
     can_return_lse_for_decode: bool = True