[Platform] Add check_runner_kv_caches_multi_layer (#51633)

Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>

vllm/platforms/cpu.py

@@ -70,6 +70,10 @@ class CpuPlatform(Platform):
         # x86/aarch64 CPU has supported both bf16 and fp16 natively.
         return [torch.bfloat16, torch.float16, torch.float32]
 
+    @classmethod
+    def check_runner_kv_caches_multi_layer(cls) -> None:
+        pass
+
     @classmethod
     def get_device_name(cls, device_id: int = 0) -> str:
         return "cpu"

vllm/platforms/cuda.py

@@ -229,6 +229,10 @@ class CudaPlatformBase(Platform):
         with contextlib.suppress(ImportError):
             import vllm._qutlass_C  # noqa: F401
 
+    @classmethod
+    def check_runner_kv_caches_multi_layer(cls) -> None:
+        pass
+
     @property
     def supported_dtypes(self) -> list[torch.dtype]:
         if self.has_device_capability(80):

vllm/platforms/interface.py

@@ -1200,6 +1200,23 @@ class Platform:
         """
         return False
 
+    @classmethod
+    def check_runner_kv_caches_multi_layer(cls) -> None:
+        """
+        Check whether the platform's ModelRunner can handle multiple attention
+        layers that share the same layer index (e.g. cross attention and self
+        attention in the same decoder block of an encoder-decoder model such as
+        BART).
+
+        Platforms that have verified that their ``runner_kv_caches`` is not
+        impacted by this case should override this to a no-op. Otherwise the
+        default implementation raises ``NotImplementedError``.
+        """
+        raise NotImplementedError(
+            "Multiple attention layers with the same layer index are not "
+            "supported on the current platform."
+        )
+
     @classmethod
     def support_deep_gemm(cls) -> bool:
         """

vllm/platforms/rocm.py

@@ -547,6 +547,10 @@ class RocmPlatform(Platform):
         with contextlib.suppress(ImportError):
             import vllm._rocm_C  # noqa: F401
 
+    @classmethod
+    def check_runner_kv_caches_multi_layer(cls) -> None:
+        pass
+
     @classmethod
     def is_pin_memory_available(cls) -> bool:
         if in_wsl():

vllm/platforms/xpu.py

@@ -134,6 +134,10 @@ class XPUPlatform(Platform):
         with contextlib.suppress(ImportError):
             import vllm._moe_C  # noqa: F401
 
+    @classmethod
+    def check_runner_kv_caches_multi_layer(cls) -> None:
+        pass
+
     @classmethod
     def get_attn_backend_cls(
         cls,

vllm/v1/worker/utils.py

@@ -538,17 +538,7 @@ def bind_kv_cache(
             # TODO - analyze where runner_kv_caches is used and the right
             # way to ensure it properly reflects multiple attention layers
             # in the same decoder block.
-            if (
-                current_platform.is_cuda_alike()
-                or current_platform.is_xpu()
-                or current_platform.is_cpu()
-            ):
-                # We know that the GPU / CPU runner is not impacted by this
-                # case. Some test code depends on runner_kv_caches, but
-                # not in a way that's impacted by ignoring this.
-                pass
-            else:
-                raise NotImplementedError
+            current_platform.check_runner_kv_caches_multi_layer()
         for layer_name in layer_names:
             runner_kv_caches.append(kv_caches[layer_name])