A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 06:00 UTC
[Bugfix][CT] Fix KV cache scale handling (#39418)
Signed-off-by: yiliu30 <yi4.liu@intel.com>
Yi Liu committed
Apr 13, 2026 at 22:50 UTC
d8ddb316444e83ce69e5bf20290d1589cbca3db4
1 file changed
+11
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+11
@@ -1123,6 +1123,17 @@ class CompressedTensorsKVCacheMethod(BaseKVCacheMethod):
1123
layer._v_scale = layer.v_scale
1124
layer._q_scale = layer.q_scale
1125
1126
+ # Set the _float variants that the attention backend uses.
1127
+ def _to_scalar(tensor: torch.Tensor) -> float:
1128
+ # For n_scales > 1 (e.g., ATTN_HEAD strategy), take max
1129
+ if tensor.numel() > 1:
1130
+ return tensor.max().item()
1131
+ return tensor.item()
1132
+
1133
+ layer._k_scale_float = _to_scalar(layer.k_scale)
1134
+ layer._v_scale_float = _to_scalar(layer.v_scale)
1135
+ layer._q_scale_float = _to_scalar(layer.q_scale)
1136
+
1137
# Discard all placeholders.
1138
del layer.k_scale
1139
del layer.v_scale