[Bugfix][Model] DeepseekV4: resolve kv_cache_dtype="auto" to fp8_ds_mla

tests/models/test_deepseek_v4_kv_cache_dtype.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.models.deepseek_v4.attention import _resolve_dsv4_kv_cache_dtype
+
+
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3", "fp8_ds_mla"])
+def test_fp8_ds_mla_layout_resolves_to_fp8_ds_mla(kv_cache_dtype: str):
+    # The fp8_ds_mla layout architecturally requires fp8 storage; "auto" and
+    # fp8 aliases resolve to the canonical string and are written back to the
+    # cache config so page-size specs pick the 576B per-token slot.
+    cache_config = SimpleNamespace(cache_dtype=kv_cache_dtype)
+    resolved, torch_dtype = _resolve_dsv4_kv_cache_dtype(
+        True, kv_cache_dtype, cache_config
+    )
+    assert resolved == "fp8_ds_mla"
+    assert torch_dtype is torch.uint8
+    assert cache_config.cache_dtype == "fp8_ds_mla"
+
+
+def test_fp8_ds_mla_layout_rejects_explicit_non_fp8():
+    with pytest.raises(AssertionError, match="only supports fp8"):
+        _resolve_dsv4_kv_cache_dtype(True, "bfloat16", None)
+
+
+def test_plain_layout_keeps_auto_as_bf16():
+    assert _resolve_dsv4_kv_cache_dtype(False, "auto", None) == (
+        "auto",
+        torch.bfloat16,
+    )
+    assert _resolve_dsv4_kv_cache_dtype(False, "fp8", None) == (
+        "fp8",
+        torch.float8_e4m3fn,
+    )

vllm/models/deepseek_v4/attention.py

@@ -102,7 +102,10 @@ def _resolve_dsv4_kv_cache_dtype(
     """
     if use_fp8_ds_mla_layout:
         # fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8.
-        assert kv_cache_dtype.startswith("fp8"), (
+        # The layout requires fp8 storage, so "auto" also resolves to
+        # fp8_ds_mla, like _canonicalize_sparse_mla_kv_cache_dtype does for
+        # FLASHMLA_SPARSE / FLASHINFER_MLA_SPARSE_SM120.
+        assert kv_cache_dtype == "auto" or kv_cache_dtype.startswith("fp8"), (
             f"DeepseekV4 fp8_ds_mla layout only supports fp8 kv-cache, "
             f"got {kv_cache_dtype}"
         )

vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py

@@ -157,7 +157,9 @@ def supports_combination(
                 return "kv_cache_dtype not supported"
             return None
         if device_capability.major == 12:
-            if kv_cache_dtype not in ("fp8", "fp8_e4m3", "fp8_ds_mla"):
+            # "auto" is accepted and resolved to fp8_ds_mla at layer
+            # construction (see _resolve_dsv4_kv_cache_dtype).
+            if kv_cache_dtype not in ("auto", "fp8", "fp8_e4m3", "fp8_ds_mla"):
                 return "kv_cache_dtype not supported"
             from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120