[Bugfix][CPU] Reject unsupported --kv-cache-dtype at CPU attention backend selection

tests/kernels/attention/test_cpu_attn.py

@@ -10,13 +10,15 @@
 
 from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
 from vllm.platforms import CpuArchEnum, current_platform
+from vllm.platforms.cpu import CpuPlatform
 from vllm.utils.torch_utils import set_default_torch_num_threads, set_random_seed
 from vllm.v1.attention.backend import AttentionType
 from vllm.v1.attention.backends.cpu_attn import _get_attn_isa
 from vllm.v1.attention.backends.zentorch_sdpa import (
     should_use_zentorch_sdpa,
     zentorch_sdpa_attn,
 )
+from vllm.v1.attention.selector import AttentionSelectorConfig
 
 if not current_platform.is_cpu():
     pytest.skip("skipping CPU-only tests", allow_module_level=True)
@@ -1589,3 +1591,38 @@ def test_amx_fp8_qk_covers_full_64_token_group() -> None:
     )
 
     torch.testing.assert_close(output, torch.ones_like(output), atol=0.05, rtol=0)
+
+
+def _select_attn_backend(kv_cache_dtype: str, use_mla: bool = False) -> str:
+    config = AttentionSelectorConfig(
+        head_size=576 if use_mla else 64,
+        dtype=torch.bfloat16,
+        kv_cache_dtype=kv_cache_dtype,
+        block_size=128,
+        use_mla=use_mla,
+    )
+    return CpuPlatform.get_attn_backend_cls(None, config)
+
+
+@pytest.mark.parametrize("fp8_isa", [True, False])
+@pytest.mark.parametrize(
+    "kv_cache_dtype",
+    ["auto", "bfloat16", "fp8", "turboquant_k8v4", "nvfp4"],
+)
+def test_cpu_attn_kv_cache_dtype_validation(
+    monkeypatch: pytest.MonkeyPatch, kv_cache_dtype: str, fp8_isa: bool
+) -> None:
+    monkeypatch.setattr("torch.cpu._is_amx_tile_supported", lambda: False)
+    monkeypatch.setattr("torch.cpu._is_avx512_supported", lambda: fp8_isa)
+
+    if kv_cache_dtype in ("auto", "bfloat16") or (kv_cache_dtype == "fp8" and fp8_isa):
+        assert _select_attn_backend(kv_cache_dtype).endswith("CPUAttentionBackend")
+    else:
+        with pytest.raises(ValueError, match=kv_cache_dtype):
+            _select_attn_backend(kv_cache_dtype)
+
+
+def test_cpu_mla_rejects_quantized_kv_cache() -> None:
+    assert "MLA" in _select_attn_backend("auto", use_mla=True)
+    with pytest.raises(ValueError, match="fp8"):
+        _select_attn_backend("fp8", use_mla=True)

vllm/platforms/cpu.py

@@ -148,7 +148,13 @@ def get_attn_backend_cls(
     ) -> str:
         if attn_selector_config.use_sparse:
             raise NotImplementedError("Sparse Attention is not supported on CPU.")
+        kv_cache_dtype = attn_selector_config.kv_cache_dtype
         if attn_selector_config.use_mla:
+            if kv_cache_dtype not in (None, "auto", "bfloat16", "float16"):
+                raise ValueError(
+                    f"--kv-cache-dtype {kv_cache_dtype} is not supported by "
+                    "MLA attention on CPU."
+                )
             amx_available = (
                 cls.get_cpu_architecture() == CpuArchEnum.X86
                 and torch.cpu._is_amx_tile_supported()
@@ -176,6 +182,13 @@ def get_attn_backend_cls(
             return AttentionBackendEnum.CPU_MLA.get_path()
         if selected_backend and selected_backend != AttentionBackendEnum.CPU_ATTN:
             logger.info("Cannot use %s backend on CPU.", selected_backend)
+        if not AttentionBackendEnum.CPU_ATTN.get_class().supports_kv_cache_dtype(
+            kv_cache_dtype
+        ):
+            raise ValueError(
+                f"--kv-cache-dtype {kv_cache_dtype} is not supported by the CPU "
+                "attention backend on this host."
+            )
         return AttentionBackendEnum.CPU_ATTN.get_path()
 
     @classmethod

vllm/v1/attention/backends/cpu_attn.py

@@ -65,6 +65,15 @@ class CPUAttentionBackend(AttentionBackend):
     def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(32)]
 
+    @classmethod
+    def supports_kv_cache_dtype(cls, kv_cache_dtype: "CacheDType | None") -> bool:
+        if kv_cache_dtype in (None, "auto", "bfloat16", "float16"):
+            return True
+        if kv_cache_dtype not in cls.supported_kv_cache_dtypes:
+            return False
+        # Mirrors the FP8 ISA requirement enforced in _get_attn_isa.
+        return torch.cpu._is_amx_tile_supported() or torch.cpu._is_avx512_supported()
+
     @classmethod
     def get_supported_head_sizes(cls) -> list[int]:
         return [32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512]