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 7, 21:30 UTC
[kv-cache, ct] Use compressed-tensors as a source of ground-truth for quant strategies (#34254)
Signed-off-by: Your Name <you@example.com> Co-authored-by: Your Name <you@example.com>
Eldar Kurtić committed
Feb 23, 2026 at 15:37 UTC
b95bb6927f7ea5f00f60b76032eb23b0128f3c48
1 file changed
+10
-8
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+10
-8
@@ -951,11 +951,11 @@ class CompressedTensorsKVCacheMethod(BaseKVCacheMethod):
951
f"received num_bits={num_bits}, type={type_}"
952
)
953
954
- # TODO: delegate validation to compressed-tensors library so that we have a
955
- # single source of truth. Right now this is not possible until the next release
956
- # of compressed-tensors.
957
- strategy = kv_cache_scheme.get("strategy")
958
- supported_strategies = ("tensor", "attn_head")
954
+ strategy = QuantizationStrategy(kv_cache_scheme.get("strategy"))
955
+ supported_strategies = (
956
+ QuantizationStrategy.TENSOR,
957
+ QuantizationStrategy.ATTN_HEAD,
958
+ )
959
if strategy not in supported_strategies:
960
raise NotImplementedError(
961
"Invalid strategy for compressed-tensors KV cache. "
@@ -981,9 +981,11 @@ class CompressedTensorsKVCacheMethod(BaseKVCacheMethod):
981
hasattr(self.quant_config, "kv_cache_scheme")
982
and self.quant_config.kv_cache_scheme is not None
983
):
984
- strategy = self.quant_config.kv_cache_scheme["strategy"]
984
+ strategy = QuantizationStrategy(
985
+ self.quant_config.kv_cache_scheme["strategy"]
986
+ )
987
986
- if strategy == "attn_head":
988
+ if strategy == QuantizationStrategy.ATTN_HEAD:
989
assert layer.impl.supports_per_head_quant_scales, (
990
f"Layer {layer.__class__.__name__} with implementation "
991
f"{layer.impl.__class__.__name__} does not support per-head scales."
@@ -1020,7 +1022,7 @@ class CompressedTensorsKVCacheMethod(BaseKVCacheMethod):
1022
# - q_scale is partitioned over query heads.
1023
# - k/v_scale is partitioned over kv heads when total_kv_heads >= tp_size,
1024
# and replicated when total_kv_heads < tp_size.
1023
- if strategy == "attn_head":
1025
+ if strategy == QuantizationStrategy.ATTN_HEAD:
1026
1027
def _tp_aware_loader(
1028
param: torch.Tensor,