From 79ca54d2215b22d9a4fc17378eb7aa2b2eb9dbd1 Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Fri, 19 Jun 2026 02:18:25 +0800 Subject: [PATCH] [Bugfix][Quantization] Don't reject fp8_e5m2 KV cache for non-fp8 quantized checkpoints (#45040) Signed-off-by: Ting Sun Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../layers/attention/attention.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 5974e09624..cdfe9fa1bc 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import torch import torch.nn as nn @@ -166,7 +166,21 @@ def _init_kv_cache_quant( # TODO (mgoin): kv cache dtype should be specified in the FP8 # checkpoint config and become the "auto" behavior if layer.kv_cache_dtype == "fp8_e5m2": - raise ValueError("fp8_e5m2 kv-cache is not supported with fp8 checkpoints.") + # A compressed-tensors checkpoint stores fp8 KV scales only when it + # declares a kv_cache_scheme; weight-only ones declare none and must + # keep fp8_e5m2, the only fp8 KV dtype usable on Ampere. + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 + CompressedTensorsConfig, + CompressedTensorsKVCacheMethod, + ) + + if not isinstance(quant_method, CompressedTensorsKVCacheMethod) or ( + cast(CompressedTensorsConfig, quant_method.quant_config).kv_cache_scheme + is not None + ): + raise ValueError( + "fp8_e5m2 kv-cache is not supported with fp8 checkpoints." + ) # If quantization is enabled, we make "k_scale" and "v_scale" # parameters so that it can be loaded from the model checkpoint. # The k/v_scale will then be converted back to native float32