A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 15:45 UTC
[Bugfix] Disable cross-layer KV cache for MLA attention backends (#37090)
Signed-off-by: haosdent <haosdent@gmail.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>
haosdent committed
Mar 17, 2026 at 01:03 UTC
ca1954d58c49e3a3209ec86d743a99f3a605028b
5 files changed
+56
-8
tests/v1/kv_connector/unit/test_kv_cache_layout.py
new
+36
@@ -0,0 +1,36 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+
5
+def test_mla_backend_rejects_cross_layer_kv_cache():
6
+ """MLA backends return identity permutation (layers dim first)
7
+ to signal cross-layer KV cache is unsupported."""
8
+ from vllm.model_executor.layers.attention.mla_attention import (
9
+ MLACommonBackend,
10
+ )
11
+
12
+ stride_order = MLACommonBackend.get_kv_cache_stride_order(
13
+ include_num_layers_dimension=True
14
+ )
15
+ assert stride_order == (0, 1, 2, 3)
16
+ assert stride_order[0] == 0 # layers dim first => no cross-layer
17
+ assert MLACommonBackend.get_kv_cache_stride_order(
18
+ include_num_layers_dimension=False
19
+ ) == (0, 1, 2)
20
+
21
+
22
+def test_deepseek_v32_indexer_rejects_cross_layer_kv_cache():
23
+ """DeepseekV32Indexer returns identity permutation (layers dim first)
24
+ to signal cross-layer KV cache is unsupported."""
25
+ from vllm.v1.attention.backends.mla.indexer import (
26
+ DeepseekV32IndexerBackend,
27
+ )
28
+
29
+ stride_order = DeepseekV32IndexerBackend.get_kv_cache_stride_order(
30
+ include_num_layers_dimension=True
31
+ )
32
+ assert stride_order == (0, 1, 2, 3)
33
+ assert stride_order[0] == 0 # layers dim first => no cross-layer
34
+ assert DeepseekV32IndexerBackend.get_kv_cache_stride_order(
35
+ include_num_layers_dimension=False
36
+ ) == (0, 1, 2)
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+4
-2
@@ -24,7 +24,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
24
)
25
from vllm.forward_context import ForwardContext
26
from vllm.logger import init_logger
27
-from vllm.model_executor.layers.attention import Attention
27
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
28
from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
29
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
30
from vllm.v1.core.kv_cache_utils import BlockHash
@@ -601,7 +601,9 @@ class OffloadingConnectorWorker:
601
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
602
layer_names = list(kv_caches.keys())
603
layers = get_layers_from_vllm_config(
604
- self.spec.vllm_config, Attention, layer_names
604
+ self.spec.vllm_config,
605
+ AttentionLayerBase, # type: ignore[type-abstract]
606
+ layer_names,
607
)
608
attn_backends = {
609
layer_name: layers[layer_name].get_attn_backend()
vllm/model_executor/layers/attention/mla_attention.py
+6
-4
@@ -1142,10 +1142,12 @@ class MLACommonBackend(AttentionBackend):
1142
def get_kv_cache_stride_order(
1143
include_num_layers_dimension: bool = False,
1144
) -> tuple[int, ...]:
1145
- # `stride_order` indicates the permutation that gets
1146
- # us from `get_kv_cache_shape` to the actual memory layout we want.
1147
- # (num_blocks, num_layers, block_size, head_size)
1148
- return (1, 0, 2, 3) if include_num_layers_dimension else (0, 1, 2)
1145
+ if include_num_layers_dimension:
1146
+ # MLA kernels require contiguous per-layer KV cache views.
1147
+ # Identity permutation keeps num_layers first in physical
1148
+ # layout, signaling cross-layer allocation is unsupported.
1149
+ return (0, 1, 2, 3)
1150
+ return (0, 1, 2)
1151
1152
@classmethod
1153
def get_supported_head_sizes(cls) -> list[int]:
vllm/v1/attention/backends/mla/indexer.py
+3
@@ -63,6 +63,9 @@ class DeepseekV32IndexerBackend(AttentionBackend):
63
include_num_layers_dimension: bool = False,
64
) -> tuple[int, ...]:
65
if include_num_layers_dimension:
66
+ # DeepseekV32Indexer kernels do not support cross-layer
67
+ # KV cache layout. Identity permutation keeps num_layers
68
+ # first, signaling incompatibility.
69
return (0, 1, 2, 3)
70
return (0, 1, 2)
71
vllm/v1/worker/kv_connector_model_runner_mixin.py
+7
-2
@@ -191,8 +191,13 @@ class KVConnectorModelRunnerMixin:
191
except (AttributeError, NotImplementedError):
192
return False
193
194
- # check that attention backend include a layers dimension
195
- return len(kv_cache_stride_order) == len(kv_cache_shape) + 1
194
+ # check that attention backend includes a layers dimension
195
+ if len(kv_cache_stride_order) != len(kv_cache_shape) + 1:
196
+ return False
197
+
198
+ # stride_order[0] == 0 means num_layers stays first in physical
199
+ # layout (identity permutation), so cross-layer is unsupported.
200
+ return kv_cache_stride_order[0] != 0
201
202
@staticmethod
203
def allocate_uniform_kv_caches(