diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py index 8fbf0c3d549..8d16a143b10 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py @@ -1123,6 +1123,17 @@ class CompressedTensorsKVCacheMethod(BaseKVCacheMethod): layer._v_scale = layer.v_scale layer._q_scale = layer.q_scale + # Set the _float variants that the attention backend uses. + def _to_scalar(tensor: torch.Tensor) -> float: + # For n_scales > 1 (e.g., ATTN_HEAD strategy), take max + if tensor.numel() > 1: + return tensor.max().item() + return tensor.item() + + layer._k_scale_float = _to_scalar(layer.k_scale) + layer._v_scale_float = _to_scalar(layer.v_scale) + layer._q_scale_float = _to_scalar(layer.q_scale) + # Discard all placeholders. del layer.k_scale del layer.v_scale