[Bugfix] Disable cross-layer KV cache for MLA attention backends

tests/v1/kv_offload/test_cpu_gpu.py

@@ -212,3 +212,33 @@ def test_transfer(
             else:
                 expected_value = orig_dst_cache[dst_block]
             torch.testing.assert_close(dst_cache[dst_block].cpu(), expected_value.cpu())
+
+
+def test_mla_backend_rejects_cross_layer_kv_cache():
+    """MLA backends must not use cross-layer KV cache layout because
+    their kernels assume contiguous per-layer block layout."""
+    from vllm.model_executor.layers.attention.mla_attention import (
+        MLACommonBackend,
+    )
+
+    with pytest.raises(NotImplementedError):
+        MLACommonBackend.get_kv_cache_stride_order(include_num_layers_dimension=True)
+    # Normal (per-layer) path still works
+    assert MLACommonBackend.get_kv_cache_stride_order(
+        include_num_layers_dimension=False
+    ) == (0, 1, 2)
+
+
+def test_deepseek_v32_indexer_rejects_cross_layer_kv_cache():
+    """DeepseekV32Indexer backend must not use cross-layer KV cache layout."""
+    from vllm.v1.attention.backends.mla.indexer import (
+        DeepseekV32IndexerBackend,
+    )
+
+    with pytest.raises(NotImplementedError):
+        DeepseekV32IndexerBackend.get_kv_cache_stride_order(
+            include_num_layers_dimension=True
+        )
+    assert DeepseekV32IndexerBackend.get_kv_cache_stride_order(
+        include_num_layers_dimension=False
+    ) == (0, 1, 2)

vllm/model_executor/layers/attention/mla_attention.py

@@ -1141,10 +1141,13 @@ def get_kv_cache_shape(
     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` to the actual memory layout we want.
-        # (num_blocks, num_layers, block_size, head_size)
-        return (1, 0, 2, 3) if include_num_layers_dimension else (0, 1, 2)
+        if include_num_layers_dimension:
+            # MLA attention kernels require contiguous per-layer KV cache
+            # views. Cross-layer block allocation produces non-contiguous
+            # views (stride(0) includes a num_layers factor) that cause
+            # incorrect memory access in paged decode kernels.
+            raise NotImplementedError
+        return (0, 1, 2)
 
     @classmethod
     def get_supported_head_sizes(cls) -> list[int]:

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

@@ -63,7 +63,9 @@ def get_kv_cache_stride_order(
         include_num_layers_dimension: bool = False,
     ) -> tuple[int, ...]:
         if include_num_layers_dimension:
-            return (0, 1, 2, 3)
+            # DeepseekV32Indexer kernels do not support cross-layer
+            # KV cache layout.
+            raise NotImplementedError
         return (0, 1, 2)