From d8ddb316444e83ce69e5bf20290d1589cbca3db4 Mon Sep 17 00:00:00 2001 From: Yi Liu Date: Mon, 13 Apr 2026 22:50:16 +0800 Subject: [PATCH] [Bugfix][CT] Fix KV cache scale handling (#39418) Signed-off-by: yiliu30 --- .../compressed_tensors/compressed_tensors.py | 11 +++++++++++ 1 file changed, 11 insertions(+) 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