[Bugfix]Fix DeepSeek-V4 fp8_ds_mla KV cache reshape (#47716)

Co-authored-by: yy-fighting <23518844576@qq.com>
Co-authored-by: Michael Goin <mgoin64@gmail.com>

vllm/models/deepseek_v4/attention.py

@@ -56,7 +56,11 @@ from vllm.v1.attention.backends.mla.indexer import (
     get_max_prefill_buffer_size,
 )
 from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache
-from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
+from vllm.v1.kv_cache_interface import (
+    KVCacheSpec,
+    MLAAttentionSpec,
+    get_kv_quant_mode,
+)
 
 logger = init_logger(__name__)
 
@@ -616,6 +620,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
             cache_dtype_str=self.kv_cache_dtype,
             alignment=576 if uses_fp8_ds_mla_layout else None,
             model_version="deepseek_v4",
+            kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
         )
 
 

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

@@ -23,6 +23,7 @@ from vllm.v1.kv_cache_interface import (
     KVCacheSpec,
     MLAAttentionSpec,
     SlidingWindowMLASpec,
+    get_kv_quant_mode,
 )
 
 # DeepseekV4 decode layer types, keyed by compress_ratio. Each type has a distinct
@@ -91,6 +92,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
             cache_dtype_str=self.cache_config.cache_dtype,
             alignment=576 if uses_fp8_ds_mla_layout else None,
             model_version="deepseek_v4",
+            kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype),
         )
 
     def forward(self): ...

vllm/v1/worker/gpu_model_runner.py

@@ -7164,7 +7164,12 @@ class GPUModelRunner(
                     layer_cache_dtype_str = (
                         "auto"
                         if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
-                        else self.cache_config.cache_dtype
+                        else getattr(
+                            kv_cache_spec,
+                            "cache_dtype_str",
+                            None,
+                        )
+                        or self.cache_config.cache_dtype
                     )
                     kv_cache_shape = attn_backend.get_kv_cache_shape(
                         kernel_num_blocks,