[Bugfix][Quantization] Don't reject fp8_e5m2 KV cache for non-fp8 quantized checkpoints (#45040)
Signed-off-by: Ting Sun <suntcrick@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Ting SUN committed
Jun 19, 2026 at 02:18 UTC
79ca54d2215b22d9a4fc17378eb7aa2b2eb9dbd1
1 file changed
+16
-2
vllm/model_executor/layers/attention/attention.py
+16
-2
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import TYPE_CHECKING, Any
4
+from typing import TYPE_CHECKING, Any, cast
5
6
import torch
7
import torch.nn as nn
@@ -166,7 +166,21 @@ def _init_kv_cache_quant(
166
# TODO (mgoin): kv cache dtype should be specified in the FP8
167
# checkpoint config and become the "auto" behavior
168
if layer.kv_cache_dtype == "fp8_e5m2":
169
- raise ValueError("fp8_e5m2 kv-cache is not supported with fp8 checkpoints.")
169
+ # A compressed-tensors checkpoint stores fp8 KV scales only when it
170
+ # declares a kv_cache_scheme; weight-only ones declare none and must
171
+ # keep fp8_e5m2, the only fp8 KV dtype usable on Ampere.
172
+ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
173
+ CompressedTensorsConfig,
174
+ CompressedTensorsKVCacheMethod,
175
+ )
176
+
177
+ if not isinstance(quant_method, CompressedTensorsKVCacheMethod) or (
178
+ cast(CompressedTensorsConfig, quant_method.quant_config).kv_cache_scheme
179
+ is not None
180
+ ):
181
+ raise ValueError(
182
+ "fp8_e5m2 kv-cache is not supported with fp8 checkpoints."
183
+ )
184
# If quantization is enabled, we make "k_scale" and "v_scale"
185
# parameters so that it can be loaded from the model checkpoint.
186
# The k/v_scale will then be converted back to native float32