[Bugfix][CT] Fix KV cache scale handling (#39418)

Signed-off-by: yiliu30 <yi4.liu@intel.com>
This commit is contained in:
Yi Liu
2026-04-13 10:50:16 -04:00
committed by GitHub
parent 1ce0318c68
commit d8ddb31644
@@ -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