[Platform] Add check_runner_kv_caches_multi_layer (#51633)
Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>
wangxiyuan committed
Aug 14, 2026 at 03:25 UTC
73b83949f76bfb806f3beca1be11b1dfb120145a
6 files changed
+34
-11
vllm/platforms/cpu.py
+4
@@ -70,6 +70,10 @@ class CpuPlatform(Platform):
70
# x86/aarch64 CPU has supported both bf16 and fp16 natively.
71
return [torch.bfloat16, torch.float16, torch.float32]
72
73
+ @classmethod
74
+ def check_runner_kv_caches_multi_layer(cls) -> None:
75
+ pass
76
+
77
@classmethod
78
def get_device_name(cls, device_id: int = 0) -> str:
79
return "cpu"
vllm/platforms/cuda.py
+4
@@ -229,6 +229,10 @@ class CudaPlatformBase(Platform):
229
with contextlib.suppress(ImportError):
230
import vllm._qutlass_C # noqa: F401
231
232
+ @classmethod
233
+ def check_runner_kv_caches_multi_layer(cls) -> None:
234
+ pass
235
+
236
@property
237
def supported_dtypes(self) -> list[torch.dtype]:
238
if self.has_device_capability(80):
vllm/platforms/interface.py
+17
@@ -1200,6 +1200,23 @@ class Platform:
1200
"""
1201
return False
1202
1203
+ @classmethod
1204
+ def check_runner_kv_caches_multi_layer(cls) -> None:
1205
+ """
1206
+ Check whether the platform's ModelRunner can handle multiple attention
1207
+ layers that share the same layer index (e.g. cross attention and self
1208
+ attention in the same decoder block of an encoder-decoder model such as
1209
+ BART).
1210
+
1211
+ Platforms that have verified that their ``runner_kv_caches`` is not
1212
+ impacted by this case should override this to a no-op. Otherwise the
1213
+ default implementation raises ``NotImplementedError``.
1214
+ """
1215
+ raise NotImplementedError(
1216
+ "Multiple attention layers with the same layer index are not "
1217
+ "supported on the current platform."
1218
+ )
1219
+
1220
@classmethod
1221
def support_deep_gemm(cls) -> bool:
1222
"""
vllm/platforms/rocm.py
+4
@@ -547,6 +547,10 @@ class RocmPlatform(Platform):
547
with contextlib.suppress(ImportError):
548
import vllm._rocm_C # noqa: F401
549
550
+ @classmethod
551
+ def check_runner_kv_caches_multi_layer(cls) -> None:
552
+ pass
553
+
554
@classmethod
555
def is_pin_memory_available(cls) -> bool:
556
if in_wsl():
vllm/platforms/xpu.py
+4
@@ -134,6 +134,10 @@ class XPUPlatform(Platform):
134
with contextlib.suppress(ImportError):
135
import vllm._moe_C # noqa: F401
136
137
+ @classmethod
138
+ def check_runner_kv_caches_multi_layer(cls) -> None:
139
+ pass
140
+
141
@classmethod
142
def get_attn_backend_cls(
143
cls,
vllm/v1/worker/utils.py
+1
-11
@@ -538,17 +538,7 @@ def bind_kv_cache(
538
# TODO - analyze where runner_kv_caches is used and the right
539
# way to ensure it properly reflects multiple attention layers
540
# in the same decoder block.
541
- if (
542
- current_platform.is_cuda_alike()
543
- or current_platform.is_xpu()
544
- or current_platform.is_cpu()
545
- ):
546
- # We know that the GPU / CPU runner is not impacted by this
547
- # case. Some test code depends on runner_kv_caches, but
548
- # not in a way that's impacted by ignoring this.
549
- pass
550
- else:
551
- raise NotImplementedError
541
+ current_platform.check_runner_kv_caches_multi_layer()
542
for layer_name in layer_names:
543
runner_kv_caches.append(kv_caches[layer_name])
544