[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>
Zhiyi Hong committed
Jul 7, 2026 at 03:56 UTC
04adc8843bbe0711fed8edf50d6d4cd4fca400e7
3 files changed
+14
-2
vllm/models/deepseek_v4/attention.py
+6
-1
@@ -56,7 +56,11 @@ from vllm.v1.attention.backends.mla.indexer import (
56
get_max_prefill_buffer_size,
57
)
58
from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache
59
-from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
59
+from vllm.v1.kv_cache_interface import (
60
+ KVCacheSpec,
61
+ MLAAttentionSpec,
62
+ get_kv_quant_mode,
63
+)
64
65
logger = init_logger(__name__)
66
@@ -616,6 +620,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
620
cache_dtype_str=self.kv_cache_dtype,
621
alignment=576 if uses_fp8_ds_mla_layout else None,
622
model_version="deepseek_v4",
623
+ kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
624
)
625
626
vllm/v1/attention/backends/mla/sparse_swa.py
+2
@@ -23,6 +23,7 @@ from vllm.v1.kv_cache_interface import (
23
KVCacheSpec,
24
MLAAttentionSpec,
25
SlidingWindowMLASpec,
26
+ get_kv_quant_mode,
27
)
28
29
# DeepseekV4 decode layer types, keyed by compress_ratio. Each type has a distinct
@@ -91,6 +92,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
92
cache_dtype_str=self.cache_config.cache_dtype,
93
alignment=576 if uses_fp8_ds_mla_layout else None,
94
model_version="deepseek_v4",
95
+ kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype),
96
)
97
98
def forward(self): ...
vllm/v1/worker/gpu_model_runner.py
+6
-1
@@ -7164,7 +7164,12 @@ class GPUModelRunner(
7164
layer_cache_dtype_str = (
7165
"auto"
7166
if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
7167
- else self.cache_config.cache_dtype
7167
+ else getattr(
7168
+ kv_cache_spec,
7169
+ "cache_dtype_str",
7170
+ None,
7171
+ )
7172
+ or self.cache_config.cache_dtype
7173
)
7174
kv_cache_shape = attn_backend.get_kv_cache_shape(
7175
kernel_num_blocks,