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