[Bugfix] Pick KV cache block size from all attention backends

tests/v1/worker/test_gpu_model_runner.py

@@ -28,14 +28,16 @@
 from vllm.lora.layers import LoRAMappingType
 from vllm.lora.request import LoRARequest
 from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
 from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
 from vllm.platforms import current_platform
+from vllm.platforms.interface import Platform
 from vllm.sampling_params import SamplingParams
 from vllm.utils.mem_constants import GiB_bytes
 from vllm.utils.system_utils import update_environment_variables
 from vllm.utils.torch_utils import set_random_seed
-from vllm.v1.attention.backend import MultipleOf
+from vllm.v1.attention.backend import AttentionBackend, MultipleOf
 from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend
 from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse import (
     ROCMAiterMLASparseBackend,
@@ -322,6 +324,53 @@ def test_select_common_block_size_accepts_rocm_sparse_block_size_16(monkeypatch)
     assert selected_size == 16
 
 
+def test_update_block_size_for_backend_honors_strictest_backend():
+    # A backend with a fixed block size requirement (e.g. the DeepSeek V3.2
+    # indexer, which only supports 64) must win over more permissive sibling
+    # backends when picking the cache block size.
+    class _FlexibleBackend(AttentionBackend):
+        @staticmethod
+        def get_name():
+            return "FLEXIBLE"
+
+        @staticmethod
+        def get_supported_kernel_block_sizes():
+            return [MultipleOf(16)]
+
+    class _StrictBackend(AttentionBackend):
+        @staticmethod
+        def get_name():
+            return "STRICT"
+
+        @staticmethod
+        def get_supported_kernel_block_sizes():
+            return [64]
+
+    class _FakeLayer(AttentionLayerBase):
+        def __init__(self, backend):
+            self._backend = backend
+
+        def get_attn_backend(self):
+            return self._backend
+
+        def get_kv_cache_spec(self, vllm_config):
+            return None
+
+    vllm_config = SimpleNamespace(
+        cache_config=CacheConfig(),
+        model_config=SimpleNamespace(is_hybrid=False),
+        compilation_config=SimpleNamespace(
+            static_forward_context={
+                "layers.0.attn": _FakeLayer(_FlexibleBackend),
+                "layers.0.indexer": _FakeLayer(_StrictBackend),
+            }
+        ),
+    )
+
+    Platform.update_block_size_for_backend(vllm_config)
+    assert vllm_config.cache_config.block_size == 64
+
+
 def test_reasoning_config_without_custom_logitsprocs_does_not_need_output_token_ids(
     dist_init,
 ):

vllm/platforms/cpu.py

@@ -458,9 +458,10 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
             return
 
         # reconcile attention and mamba page sizes
-        backend_cls = cls._find_non_ssm_backend(vllm_config)
-        if backend_cls is None:
+        backend_classes = cls._find_non_ssm_backends(vllm_config)
+        if not backend_classes:
             return
+        backend_cls = backend_classes[0]
 
         cls._align_hybrid_block_size(vllm_config, backend_cls)
 

vllm/platforms/interface.py

@@ -586,10 +586,10 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
         pass
 
     @classmethod
-    def _find_non_ssm_backend(
+    def _find_non_ssm_backends(
         cls, vllm_config: "VllmConfig"
-    ) -> "type[AttentionBackend] | None":
-        """Find the first non-SSM attention backend from model layers."""
+    ) -> "list[type[AttentionBackend]]":
+        """Find all distinct non-SSM attention backends from model layers."""
         from vllm.config.vllm import get_layers_from_vllm_config
         from vllm.model_executor.layers.attention_layer_base import (
             AttentionLayerBase,
@@ -599,11 +599,12 @@ def _find_non_ssm_backend(
             vllm_config,
             AttentionLayerBase,  # type: ignore[type-abstract]
         )
+        backends: list[type[AttentionBackend]] = []
         for layer in attn_layers.values():
             b = layer.get_attn_backend()
-            if not b.is_ssm():
-                return b
-        return None
+            if not b.is_ssm() and b not in backends:
+                backends.append(b)
+        return backends
 
     @classmethod
     def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
@@ -621,21 +622,31 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
         if not model_config:
             return
 
-        backend_cls = cls._find_non_ssm_backend(vllm_config)
-        if backend_cls is None:
+        backend_classes = cls._find_non_ssm_backends(vllm_config)
+        if not backend_classes:
             return
+        backend_cls = backend_classes[0]
 
-        # Phase 1: Pick block size from backend (skip if user set --block-size)
+        # Phase 1: Pick block size from backends (skip if user set --block-size).
+        # Honor the strictest backend: e.g. the DeepSeek V3.2 indexer only
+        # supports block size 64 while the main sparse attention backend may
+        # accept the default.
         if not cache_config.user_specified_block_size:
+            preferred = CacheConfig.DEFAULT_BLOCK_SIZE
+            preferred_backend_cls = backend_cls
             with set_current_vllm_config(vllm_config):
-                preferred = backend_cls.get_preferred_block_size(
-                    CacheConfig.DEFAULT_BLOCK_SIZE
-                )
+                for backend in backend_classes:
+                    backend_preferred = backend.get_preferred_block_size(
+                        CacheConfig.DEFAULT_BLOCK_SIZE
+                    )
+                    if backend_preferred > preferred:
+                        preferred = backend_preferred
+                        preferred_backend_cls = backend
             if preferred != CacheConfig.DEFAULT_BLOCK_SIZE:
                 logger.info(
                     "Setting kv cache block size to %d for %s backend.",
                     preferred,
-                    backend_cls.get_name(),
+                    preferred_backend_cls.get_name(),
                 )
             cache_config.block_size = preferred