From 04adc8843bbe0711fed8edf50d6d4cd4fca400e7 Mon Sep 17 00:00:00 2001 From: Zhiyi Hong <118732786+ACEEE-1222@users.noreply.github.com> Date: Tue, 7 Jul 2026 03:56:44 +0800 Subject: [PATCH] [Bugfix]Fix DeepSeek-V4 fp8_ds_mla KV cache reshape (#47716) Co-authored-by: yy-fighting <23518844576@qq.com> Co-authored-by: Michael Goin --- vllm/models/deepseek_v4/attention.py | 7 ++++++- vllm/v1/attention/backends/mla/sparse_swa.py | 2 ++ vllm/v1/worker/gpu_model_runner.py | 7 ++++++- 3 files changed, 14 insertions(+), 2 deletions(-) diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 29a19d9026..906d594ac5 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/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), ) diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 46b55cc7b2..ca5be3a302 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/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): ... diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 5a267a8bb8..2cf7d2db54 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/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,