[BugFix][MLA] Support kv_cache_dtype_skip_layers for MLA attention (#47309)
Signed-off-by: liuruikang <liuruikang.cs@gmail.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
ruikangliu committed
Jul 16, 2026 at 08:06 UTC
0becb7486b3e368bb9c0e56b1ca3ccf8ee53911a
1 file changed
+14
-1
vllm/model_executor/layers/attention/mla_attention.py
+14
-1
@@ -393,6 +393,19 @@ class MLAAttention(nn.Module, AttentionLayerBase):
393
calculate_kv_scales = False
394
self.quant_config = quant_config
395
396
+ if cache_config is not None and cache_config.kv_cache_dtype_skip_layers:
397
+ from vllm.model_executor.models.utils import extract_layer_index
398
+
399
+ layer_idx = extract_layer_index(prefix)
400
+ if str(layer_idx) in cache_config.kv_cache_dtype_skip_layers:
401
+ kv_cache_dtype = "auto"
402
+ calculate_kv_scales = False
403
+ logger.debug(
404
+ "Layer %s: kv_cache_dtype=%s",
405
+ prefix,
406
+ kv_cache_dtype,
407
+ )
408
+
409
dtype = torch.get_default_dtype()
410
if attn_backend is not None:
411
assert attn_backend.is_mla(), (
@@ -1025,7 +1038,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
1038
num_kv_heads=1,
1039
head_size=self.head_size,
1040
dtype=kv_cache_dtype,
1028
- cache_dtype_str=vllm_config.cache_config.cache_dtype,
1041
+ cache_dtype_str=self.kv_cache_dtype,
1042
kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
1043
)
1044