[Misc][Platform] Check aligned KV block sizes against every attention backend (#58457)

Signed-off-by: Divy <divy@coralbricks.ai>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

tests/v1/worker/test_gpu_model_runner.py

@@ -486,6 +486,67 @@ def test_preferred_block_size_rejects_backends_with_no_common_size():
         Platform._preferred_block_size_for_backends(classes, 16, None)
 
 
+@pytest.mark.parametrize("cpu", [False, True])
+def test_alignment_rejected_by_a_sibling_backend_raises(monkeypatch, cpu):
+    # Alignment that lands on a size an exact-size sibling rejects must fail
+    # here, not later in select_common_block_size().
+    backends = [_mock_backend([MultipleOf(16)]), _mock_backend([16], exact=True)]
+    monkeypatch.setattr(
+        Platform, "_find_non_ssm_backends", classmethod(lambda cls, c: backends)
+    )
+
+    def bump(cls, vllm_config, backend_cls):
+        vllm_config.cache_config.block_size = 32
+
+    monkeypatch.setattr(Platform, "_align_hybrid_block_size", classmethod(bump))
+    cache_config = SimpleNamespace(
+        block_size=16, user_specified_block_size=True, kv_cache_dtype_skip_layers=None
+    )
+    vllm_config = SimpleNamespace(
+        cache_config=cache_config, model_config=SimpleNamespace(is_hybrid=True)
+    )
+    from vllm.platforms.cpu import CpuPlatform  # overrides the method; must check too
+
+    platform = CpuPlatform if cpu else Platform
+    with pytest.raises(ValueError, match="MOCK_EXACT"):
+        platform.update_block_size_for_backend(vllm_config)
+
+
+def test_xpu_gdn_rounding_rejected_by_a_sibling_backend_raises(monkeypatch):
+    # XPU's GDN rounding runs after super()'s check, so it must check again.
+    try:
+        from vllm.platforms.xpu import XPUPlatform
+    except ImportError:
+        pytest.skip("vllm_xpu_kernels not importable outside an XPU stack")
+
+    backends = [_mock_backend([MultipleOf(16)]), _mock_backend([16], exact=True)]
+    monkeypatch.setattr(
+        Platform, "update_block_size_for_backend", classmethod(lambda cls, c: None)
+    )
+    monkeypatch.setattr(
+        Platform, "_find_non_ssm_backends", classmethod(lambda cls, c: backends)
+    )
+    gdn_layer = SimpleNamespace(
+        get_attn_backend=lambda: SimpleNamespace(get_name=lambda: "GDN_ATTN")
+    )
+    monkeypatch.setattr(
+        "vllm.config.vllm.get_layers_from_vllm_config",
+        lambda vllm_config, layer_type: {"gdn": gdn_layer},
+    )
+    cache_config = SimpleNamespace(
+        block_size=16,
+        mamba_cache_mode="none",
+        mamba_page_size_padded=None,
+        user_specified_block_size=True,
+        kv_cache_dtype_skip_layers=None,
+    )
+    vllm_config = SimpleNamespace(
+        cache_config=cache_config, model_config=SimpleNamespace(is_hybrid=True)
+    )
+    with pytest.raises(ValueError, match="MOCK_EXACT"):
+        XPUPlatform.update_block_size_for_backend(vllm_config)
+
+
 def test_set_active_mm_loras_builds_tower_and_connector_mappings():
     model = Mock()
     model.get_mm_lora_token_counts.side_effect = (

vllm/platforms/cpu.py

@@ -523,7 +523,9 @@ class CpuPlatform(Platform):
         if not backend_classes:
             return
 
+        pre_block_size = vllm_config.cache_config.block_size
         cls._align_hybrid_block_size(vllm_config, backend_classes[0])
+        cls._check_aligned_block_size(vllm_config, backend_classes, pre_block_size)
 
     @classmethod
     def discover_numa_topology(cls) -> list[list[int]]:

vllm/platforms/interface.py

@@ -698,6 +698,7 @@ class Platform:
 
         # Phase 2: Align block/mamba sizes for hybrid models
         # (may override user settings).
+        pre_block_size = cache_config.block_size
         if model_config.is_hybrid:
             cls._align_hybrid_block_size(vllm_config, backend_classes[0])
 
@@ -706,6 +707,36 @@ class Platform:
         # May override the user's --block-size.
         if cache_config.kv_cache_dtype_skip_layers:
             cls._align_heterogeneous_kv_block_size(vllm_config, backend_classes[0])
+        cls._check_aligned_block_size(vllm_config, backend_classes, pre_block_size)
+
+    @classmethod
+    def _check_aligned_block_size(
+        cls,
+        vllm_config: "VllmConfig",
+        backend_classes: "list[type[AttentionBackend]]",
+        pre_block_size: int,
+    ) -> None:
+        """Fail now, not at worker setup, if alignment left a rejected size.
+
+        Alignment rounds by the first backend's kernel block only, so it can
+        land on a size a sibling backend rejects.
+        """
+        from vllm.config.vllm import set_current_vllm_config
+
+        block_size = vllm_config.cache_config.block_size
+        if block_size == pre_block_size:
+            return
+        with set_current_vllm_config(vllm_config):
+            rejecting = [
+                b.get_name()
+                for b in backend_classes
+                if not b.supports_block_size(block_size)
+            ]
+        if rejecting:
+            raise ValueError(
+                f"Aligned KV cache block size {block_size} is not supported by "
+                f"attention backend(s): {', '.join(rejecting)}."
+            )
 
     @classmethod
     def _align_heterogeneous_kv_block_size(

vllm/platforms/xpu.py

@@ -452,6 +452,7 @@ class XPUPlatform(Platform):
         if new_block_size == cache_config.block_size:
             return
 
+        pre_block_size = cache_config.block_size
         if cache_config.mamba_cache_mode == "align":
             cache_config.mamba_block_size = new_block_size
         original_mamba_page_size_padded = cache_config.mamba_page_size_padded
@@ -474,6 +475,10 @@ class XPUPlatform(Platform):
                 original_mamba_page_size_padded,
                 cache_config.mamba_page_size_padded,
             )
+        # This rounding runs after super()'s check, so check its result too.
+        cls._check_aligned_block_size(
+            vllm_config, cls._find_non_ssm_backends(vllm_config), pre_block_size
+        )
 
     @classmethod
     def support_hybrid_kv_cache(cls) -> bool: