[Bugfix][Attention] Resolve kernel block-size constraints per KV cache spec

tests/v1/worker/test_attn_utils.py

@@ -79,7 +79,7 @@ def test_get_kv_cache_spec_resolves_hisparse_block_size(
         backend = SimpleNamespace(
             get_name=lambda name=name: name,
             customize_spec=AttentionBackend.customize_spec,
-            get_supported_kernel_block_sizes=lambda sizes=sizes: sizes,
+            get_supported_kernel_block_sizes=lambda *_, sizes=sizes: sizes,
         )
         layers[name] = SimpleNamespace(
             get_kv_cache_spec=lambda _, spec=specs[name]: spec,

tests/v1/worker/test_gpu_model_runner.py

@@ -36,7 +36,7 @@
 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,
@@ -303,7 +303,7 @@ def get_name():
             return "MOCK"
 
         @staticmethod
-        def get_supported_kernel_block_sizes():
+        def get_supported_kernel_block_sizes(kv_cache_spec=None):
             return supported_sizes
 
     return _MockBackend()
@@ -441,13 +441,13 @@ def get_name() -> str:
             return "MOCK_EXACT" if exact else "MOCK"
 
         @staticmethod
-        def get_supported_kernel_block_sizes():
+        def get_supported_kernel_block_sizes(kv_cache_spec=None):
             return list(supported)
 
         if exact:
 
             @classmethod
-            def supports_block_size(cls, block_size: int | None) -> bool:
+            def supports_block_size(cls, block_size, kv_cache_spec=None):
                 return block_size is None or block_size in supported
 
     return _MockBackendCls
@@ -469,21 +469,61 @@ def supports_block_size(cls, block_size: int | None) -> bool:
     ],
 )
 def test_preferred_block_size_satisfies_every_backend(backends, expected):
-    classes = [_mock_backend(s) for s in backends]
-    assert Platform._preferred_block_size_for_backends(classes, 16, None) == expected
+    specs = [(_mock_backend(s), None) for s in backends]
+    assert Platform._preferred_block_size_for_backends(specs, 16, None) == expected
 
 
 def test_preferred_block_size_searches_past_an_exact_size_backend():
     # Extending greedily picks lcm(16, 32) = 32, which the exact backend
     # rejects; 96 is accepted by both.
-    classes = [_mock_backend([16, 96], exact=True), _mock_backend([MultipleOf(32)])]
-    assert Platform._preferred_block_size_for_backends(classes, 16, None) == 96
+    specs = [(_mock_backend([16, 96], exact=True), None), (_mock_backend([32]), None)]
+    assert Platform._preferred_block_size_for_backends(specs, 16, None) == 96
 
 
 def test_preferred_block_size_rejects_backends_with_no_common_size():
-    classes = [_mock_backend([16], exact=True), _mock_backend([MultipleOf(64)])]
+    specs = [(_mock_backend([16], exact=True), None), (_mock_backend([64]), None)]
     with pytest.raises(ValueError, match="share no supported KV cache block size"):
-        Platform._preferred_block_size_for_backends(classes, 16, None)
+        Platform._preferred_block_size_for_backends(specs, 16, None)
+
+
+class _SpecAwareMockBackend(AttentionBackend):
+    """FlashAttention-like: a 128-token page for head size 256 (and model-wide),
+    multiples of 16 for other head sizes."""
+
+    @staticmethod
+    def get_name() -> str:
+        return "MOCK_SPEC_AWARE"
+
+    @staticmethod
+    def get_supported_kernel_block_sizes(kv_cache_spec=None):
+        if kv_cache_spec is None or kv_cache_spec.head_size == 256:
+            return [128]
+        return [MultipleOf(16)]
+
+
+def _spec(head_size: int) -> FullAttentionSpec:
+    return FullAttentionSpec(
+        block_size=16, num_kv_heads=1, head_size=head_size, dtype=torch.bfloat16
+    )
+
+
+def test_preferred_block_size_resolves_constraints_per_spec():
+    target = _mock_backend([16, 32, 64])  # FlashInfer-like head-size-256 target
+    drafter = _SpecAwareMockBackend  # FlashAttention drafter, head size 128
+    # Without specs, the drafter is held to the target's 128-token page.
+    backend_specs = [(target, None), (drafter, None)]
+    assert Platform._preferred_block_size_for_backends(backend_specs, 16, None) == 128
+    backend_specs = [(target, _spec(256)), (drafter, _spec(128))]
+    assert Platform._preferred_block_size_for_backends(backend_specs, 16, None) == 16
+
+
+def test_select_common_block_size_resolves_sizes_per_spec():
+    backend = _SpecAwareMockBackend
+    # 832 is a hybrid head-size-256 target's mamba-aligned manager block; its
+    # drafter's head-size-128 layers must not be held to the 128-token page.
+    with pytest.raises(ValueError):
+        select_common_block_size(832, [backend])
+    assert select_common_block_size(832, [backend], [_spec(128)]) == 832
 
 
 def test_set_active_mm_loras_builds_tower_and_connector_mappings():

vllm/platforms/interface.py

@@ -27,6 +27,7 @@
     from vllm.utils.argparse_utils import FlexibleArgumentParser
     from vllm.v1.attention.backend import AttentionBackend
     from vllm.v1.attention.selector import AttentionSelectorConfig
+    from vllm.v1.kv_cache_interface import KVCacheSpec
 else:
     FlexibleArgumentParser = object
 
@@ -592,10 +593,13 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
         pass
 
     @classmethod
-    def _find_non_ssm_backends(
+    def _find_non_ssm_backend_specs(
         cls, vllm_config: "VllmConfig"
-    ) -> "list[type[AttentionBackend]]":
-        """Distinct non-SSM attention backends, in layer order."""
+    ) -> "list[tuple[type[AttentionBackend], KVCacheSpec | None]]":
+        """Distinct (non-SSM attention backend, KV cache spec) pairs, in layer
+        order. A backend's kernel block sizes may depend on the spec it serves
+        (e.g. a drafter's head size), so one backend can appear with several.
+        """
         from vllm.config.vllm import get_layers_from_vllm_config
         from vllm.model_executor.layers.attention_layer_base import (
             AttentionLayerBase,
@@ -605,21 +609,33 @@ def _find_non_ssm_backends(
             vllm_config,
             AttentionLayerBase,  # type: ignore[type-abstract]
         )
-        backends: list[type[AttentionBackend]] = []
+        backend_specs: list[tuple[type[AttentionBackend], KVCacheSpec | None]] = []
         for layer in attn_layers.values():
             b = layer.get_attn_backend()
-            if not b.is_ssm() and b not in backends:
-                backends.append(b)
-        return backends
+            if b.is_ssm():
+                continue
+            backend_spec = (b, layer.get_kv_cache_spec(vllm_config))
+            if backend_spec not in backend_specs:
+                backend_specs.append(backend_spec)
+        return backend_specs
+
+    @classmethod
+    def _find_non_ssm_backends(
+        cls, vllm_config: "VllmConfig"
+    ) -> "list[type[AttentionBackend]]":
+        """Distinct non-SSM attention backends, in layer order."""
+        return list(
+            dict.fromkeys(b for b, _ in cls._find_non_ssm_backend_specs(vllm_config))
+        )
 
     @classmethod
     def _preferred_block_size_for_backends(
         cls,
-        backend_classes: "list[type[AttentionBackend]]",
+        backend_specs: "list[tuple[type[AttentionBackend], KVCacheSpec | None]]",
         default_block_size: int,
         vllm_config: "VllmConfig",
     ) -> int:
-        """Smallest block size every backend accepts.
+        """Smallest block size every backend accepts for the spec it serves.
 
         ``supports_block_size`` may be overridden to accept exact sizes only
         (CPU_MLA takes 16 and no multiple of it), so candidates are the LCMs of
@@ -634,30 +650,36 @@ def _preferred_block_size_for_backends(
 
         # Backends may read the current config to decide a preference.
         with set_current_vllm_config(vllm_config):
-            if len(backend_classes) == 1:
-                return backend_classes[0].get_preferred_block_size(default_block_size)
-            if all(b.supports_block_size(default_block_size) for b in backend_classes):
+            # A lone backend keeps its preference, whichever specs it serves.
+            if len({b for b, _ in backend_specs}) == 1:
+                return backend_specs[0][0].get_preferred_block_size(default_block_size)
+            if all(
+                b.supports_block_size(default_block_size, spec)
+                for b, spec in backend_specs
+            ):
                 return default_block_size
             # A backend declaring no sizes accepts any, so it contributes 1.
             per_backend_sizes = [
                 [
                     s.base if isinstance(s, MultipleOf) else s
-                    for s in b.get_supported_kernel_block_sizes()
+                    for s in b.get_supported_kernel_block_sizes(spec)
                 ]
                 or [1]
-                for b in backend_classes
+                for b, spec in backend_specs
             ]
             candidates = sorted(
                 {math.lcm(*sizes) for sizes in itertools.product(*per_backend_sizes)}
             )
             for candidate in candidates:
-                if all(b.supports_block_size(candidate) for b in backend_classes):
+                if all(
+                    b.supports_block_size(candidate, spec) for b, spec in backend_specs
+                ):
                     return candidate
         raise ValueError(
             "The attention backends share no supported KV cache block size ("
             + "; ".join(
-                f"{b.get_name()}: {b.get_supported_kernel_block_sizes()}"
-                for b in backend_classes
+                f"{b.get_name()}: {b.get_supported_kernel_block_sizes(spec)}"
+                for b, spec in backend_specs
             )
             + ")."
         )
@@ -676,7 +698,8 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
         if not model_config:
             return
 
-        backend_classes = cls._find_non_ssm_backends(vllm_config)
+        backend_specs = cls._find_non_ssm_backend_specs(vllm_config)
+        backend_classes = list(dict.fromkeys(b for b, _ in backend_specs))
         if not backend_classes:
             return
 
@@ -686,7 +709,7 @@ def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
         # select_common_block_size().
         if not cache_config.user_specified_block_size:
             preferred = cls._preferred_block_size_for_backends(
-                backend_classes, CacheConfig.DEFAULT_BLOCK_SIZE, vllm_config
+                backend_specs, CacheConfig.DEFAULT_BLOCK_SIZE, vllm_config
             )
             if preferred != CacheConfig.DEFAULT_BLOCK_SIZE:
                 logger.info(

vllm/v1/attention/backend.py

@@ -118,11 +118,15 @@ def supports_kv_cache_dtype(cls, kv_cache_dtype: "CacheDType | None") -> bool:
         )
 
     @classmethod
-    def supports_block_size(cls, block_size: int | None) -> bool:
+    def supports_block_size(
+        cls, block_size: int | None, kv_cache_spec: "KVCacheSpec | None" = None
+    ) -> bool:
         if block_size is None:
             return True
 
-        supported_kernel_block_sizes = cls.get_supported_kernel_block_sizes()
+        supported_kernel_block_sizes = cls.get_supported_kernel_block_sizes(
+            kv_cache_spec
+        )
         if not supported_kernel_block_sizes:
             return True
 

vllm/v1/attention/backends/b12x.py

@@ -187,7 +187,7 @@ def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleO
         return list(_B12X_SUPPORTED_PAGE_SIZES)
 
     @classmethod
-    def supports_block_size(cls, block_size: int | None) -> bool:
+    def supports_block_size(cls, block_size: int | None, kv_cache_spec=None) -> bool:
         return block_size is None or int(block_size) in _B12X_SUPPORTED_PAGE_SIZES
 
     @classmethod

vllm/v1/attention/backends/composite.py

@@ -397,10 +397,10 @@ def get_supported_kernel_block_sizes(
             )
 
         @classmethod
-        def supports_block_size(cls, block_size):
+        def supports_block_size(cls, block_size, kv_cache_spec=None):
             return bool(
-                cls.get_supported_kernel_block_sizes()
-            ) and super().supports_block_size(block_size)
+                cls.get_supported_kernel_block_sizes(kv_cache_spec)
+            ) and super().supports_block_size(block_size, kv_cache_spec)
 
         @classmethod
         def get_supported_head_sizes(cls):

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

@@ -134,7 +134,7 @@ def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleO
         return [MultipleOf(32)]
 
     @classmethod
-    def supports_block_size(cls, block_size: int | None) -> bool:
+    def supports_block_size(cls, block_size: int | None, kv_cache_spec=None) -> bool:
         if block_size is None:
             return True
         return block_size % 32 == 0

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

@@ -68,7 +68,7 @@ def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleO
         return [16]
 
     @classmethod
-    def supports_block_size(cls, block_size: int | None) -> bool:
+    def supports_block_size(cls, block_size: int | None, kv_cache_spec=None) -> bool:
         return block_size is None or block_size == 16
 
     @staticmethod

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

@@ -111,7 +111,7 @@ def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleO
         return [MultipleOf(16)]
 
     @classmethod
-    def supports_block_size(cls, block_size: int | None) -> bool:
+    def supports_block_size(cls, block_size: int | None, kv_cache_spec=None) -> bool:
         if block_size is None:
             return True
         return block_size % 16 == 0

vllm/v1/attention/backends/rocm_aiter_unified_attn.py

@@ -72,7 +72,7 @@ def get_preferred_block_size(cls, default_block_size: int) -> int:
         return 64
 
     @classmethod
-    def supports_block_size(cls, block_size: int | None) -> bool:
+    def supports_block_size(cls, block_size: int | None, kv_cache_spec=None) -> bool:
         if block_size is None:
             return True
         return block_size % 16 == 0

vllm/v1/attention/backends/triton_attn.py

@@ -344,7 +344,7 @@ def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleO
         return [MultipleOf(16)]
 
     @classmethod
-    def supports_block_size(cls, block_size: int | None) -> bool:
+    def supports_block_size(cls, block_size: int | None, kv_cache_spec=None) -> bool:
         if block_size is None:
             return True
         return block_size % 16 == 0