diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 032f849841..108e6ff152 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -393,6 +393,19 @@ class MLAAttention(nn.Module, AttentionLayerBase): calculate_kv_scales = False self.quant_config = quant_config + if cache_config is not None and cache_config.kv_cache_dtype_skip_layers: + from vllm.model_executor.models.utils import extract_layer_index + + layer_idx = extract_layer_index(prefix) + if str(layer_idx) in cache_config.kv_cache_dtype_skip_layers: + kv_cache_dtype = "auto" + calculate_kv_scales = False + logger.debug( + "Layer %s: kv_cache_dtype=%s", + prefix, + kv_cache_dtype, + ) + dtype = torch.get_default_dtype() if attn_backend is not None: assert attn_backend.is_mla(), ( @@ -1025,7 +1038,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): num_kv_heads=1, head_size=self.head_size, dtype=kv_cache_dtype, - cache_dtype_str=vllm_config.cache_config.cache_dtype, + cache_dtype_str=self.kv_cache_dtype, kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), )