diff --git a/tests/config/test_checkpoint_kv_quant_policy.py b/tests/config/test_checkpoint_kv_quant_policy.py new file mode 100644 index 0000000000..d573dbab95 --- /dev/null +++ b/tests/config/test_checkpoint_kv_quant_policy.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""A checkpoint's KV-cache quantization directive is honored only on Ampere+. + +The directive describes how the weights were made; on Volta and Turing there +is no FP8 hardware, so under ``--kv-cache-dtype auto`` the KV cache keeps the +model dtype. An explicit ``--kv-cache-dtype`` is never touched. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm import platforms +from vllm.config import CacheConfig, VllmConfig +from vllm.config import vllm as vllm_config_module +from vllm.config.vllm import checkpoint_kv_quant_allowed + +SM70 = (7, 0) +SM75 = (7, 5) +SM80 = (8, 0) +SM90 = (9, 0) + + +def _fake_platform(capabilities: list[tuple[int, int]]): + return SimpleNamespace( + is_cuda=lambda: True, + device_count=lambda: len(capabilities), + is_device_capability=lambda capability, device_id=0: ( + capabilities[device_id] == capability + ), + ) + + +def _placement_config(world_size: int): + return SimpleNamespace( + parallel_config=SimpleNamespace( + distributed_executor_backend="mp" if world_size > 1 else "uni", + data_parallel_backend="mp", + world_size=world_size, + local_world_size=world_size, + nnodes_within_dp=1, + data_parallel_rank_local=0, + data_parallel_index=0, + tensor_parallel_size=world_size, + pipeline_parallel_size=1, + ), + device_config=SimpleNamespace(device=torch.device("cuda")), + ) + + +@pytest.mark.parametrize( + ("capabilities", "world_size", "expected"), + [ + pytest.param([SM80], 1, True, id="ampere"), + pytest.param([SM90, SM90], 2, True, id="hopper-pair"), + pytest.param([SM70], 1, False, id="volta"), + pytest.param([SM75], 1, False, id="turing"), + pytest.param([SM75, SM70], 2, False, id="mixed-pre-ampere"), + pytest.param([SM80, SM75], 2, False, id="ampere-with-turing"), + # Only participating devices count: the Turing card is visible but + # not part of this single-GPU engine. + pytest.param([SM80, SM75], 1, True, id="turing-visible-not-used"), + ], +) +def test_policy_follows_participating_devices( + monkeypatch, capabilities, world_size, expected +): + monkeypatch.setattr(platforms, "current_platform", _fake_platform(capabilities)) + assert checkpoint_kv_quant_allowed(_placement_config(world_size)) is expected + + +def _checkpoint_resolved_cache_config() -> CacheConfig: + cache_config = CacheConfig(cache_dtype="fp8_e4m3") + cache_config.cache_dtype_from_checkpoint = True + return cache_config + + +def test_checkpoint_directive_dropped_on_pre_ampere(monkeypatch): + monkeypatch.setattr( + vllm_config_module, "_any_participating_device_is_pre_ampere", lambda cfg: True + ) + config = VllmConfig(cache_config=_checkpoint_resolved_cache_config()) + assert config.cache_config.cache_dtype == "auto" + assert config.cache_config.cache_dtype_from_checkpoint is False + + +def test_checkpoint_directive_kept_on_ampere(monkeypatch): + monkeypatch.setattr( + vllm_config_module, "_any_participating_device_is_pre_ampere", lambda cfg: False + ) + config = VllmConfig(cache_config=_checkpoint_resolved_cache_config()) + assert config.cache_config.cache_dtype == "fp8_e4m3" + assert config.cache_config.cache_dtype_from_checkpoint is True + + +def test_explicit_request_is_never_touched(monkeypatch): + monkeypatch.setattr( + vllm_config_module, "_any_participating_device_is_pre_ampere", lambda cfg: True + ) + config = VllmConfig(cache_config=CacheConfig(cache_dtype="fp8_e4m3")) + assert config.cache_config.cache_dtype == "fp8_e4m3" diff --git a/vllm/config/cache.py b/vllm/config/cache.py index c3861d0dea..1023385136 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -52,6 +52,10 @@ class CacheConfig: """Whether block_size was explicitly provided. Derived automatically.""" user_specified_mamba_block_size: bool = field(default=False, init=False) """Whether mamba_block_size was explicitly provided. Derived automatically.""" + cache_dtype_from_checkpoint: bool = field(default=False, init=False) + """Whether cache_dtype was resolved from the checkpoint's KV-cache + quantization metadata rather than requested by the user. Derived + automatically; a request other than "auto" never sets it.""" hash_block_size: int | None = Field(default=None, gt=0) """Block size (in tokens) used for computing Request's block_hashes. @@ -73,7 +77,10 @@ class CacheConfig: example, if you have two vLLM instances running on the same GPU, you can set the GPU memory utilization to 0.5 for each instance.""" cache_dtype: CacheDType = "auto" - """Data type for kv cache storage. If "auto", will use model data type. + """Data type for kv cache storage. If "auto", will use model data type, + unless the checkpoint declares a KV-cache quantization algorithm; that + declaration is honored on Ampere and newer and ignored on Volta and + Turing, which have no FP8 hardware (pass the dtype explicitly to force it). CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. ROCm (AMD GPU) supports fp8 (=fp8_e4m3). Intel Gaudi (HPU) supports fp8 (using fp8_inc). On SM70 with the 1Cat Flash-V100 backend enabled, the user-facing ``fp8`` diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index c0933fe43b..8b38aa5b2b 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -312,6 +312,23 @@ def _any_participating_device_is_pre_ampere(cfg: "VllmConfig") -> bool: ) or _any_participating_device_is_capability(cfg, (7, 5)) +def checkpoint_kv_quant_allowed(cfg: "VllmConfig") -> bool: + """May the checkpoint's own metadata select a quantized KV cache here? + + A checkpoint that declares ``kv_cache_quant_algo`` or ``kv_cache_scheme`` + describes how its weights were produced. With ``--kv-cache-dtype auto`` + vLLM reads that as permission to also store the KV cache in FP8. On + Volta and Turing there is no FP8 hardware: the cache is unpacked in + software and decode attention loses its tensor-core route (measured on + 4x V100 with Qwen3.8-27B: +4.82 ms per decode round, 4.5x the cost of + the FP8 weights the checkpoint ships with), and on Turing the FP8 cast + is not compiled at all. So the directive is honored only when every + participating device is Ampere or newer. An explicit ``--kv-cache-dtype`` + never reaches this policy. + """ + return not _any_participating_device_is_pre_ampere(cfg) + + def _apply_sm70_dflash2_verifier_defaults() -> tuple[str, ...]: """Set quality-audited defaults while preserving every explicit override.""" applied = [] @@ -1480,6 +1497,20 @@ def __post_init__(self): if self.performance_mode != "balanced": logger.info_once("Performance mode set to '%s'.", self.performance_mode) + if ( + self.cache_config.cache_dtype_from_checkpoint + and not checkpoint_kv_quant_allowed(self) + ): + logger.info_once( + "Ignoring the checkpoint's KV-cache quantization directive (%s): " + "a participating device is Volta or Turing, which has no FP8 " + "hardware, so the KV cache keeps the model dtype. Pass " + "--kv-cache-dtype explicitly to override.", + self.cache_config.cache_dtype, + ) + self.cache_config.cache_dtype = "auto" + self.cache_config.cache_dtype_from_checkpoint = False + self.try_verify_and_update_config() if self.model_config is not None: diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index ac032debed..64df1cdf3d 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -2006,6 +2006,10 @@ def create_engine_config( kv_offloading_backend=self.kv_offloading_backend, ) + cache_config.cache_dtype_from_checkpoint = ( + self.kv_cache_dtype == "auto" and resolved_cache_dtype != "auto" + ) + if resolved_cache_dtype.startswith("turboquant_"): from vllm.model_executor.layers.quantization.turboquant.config import ( TurboQuantConfig, diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 1e4e3996d2..2af74effa3 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/vllm/model_executor/layers/attention/attention.py @@ -8,7 +8,7 @@ import vllm.envs as envs from vllm.config import CacheConfig, get_current_vllm_config -from vllm.config.vllm import VllmConfig +from vllm.config.vllm import VllmConfig, checkpoint_kv_quant_allowed from vllm.forward_context import ForwardContext, get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.attention.kv_transfer_utils import ( @@ -273,8 +273,15 @@ def __init__( # The "auto" case is normally resolved upstream in # resolve_kv_cache_dtype_string, but we re-apply here defensively in # case anything bypassed that path. + # The same pre-Ampere policy as VllmConfig applies here, so a + # compressed-tensors checkpoint cannot quantize the cache on a device + # where the resolve path just refused to. kv_cache_scheme = getattr(quant_config, "kv_cache_scheme", None) - if kv_cache_scheme is not None and kv_cache_dtype == "auto": + if ( + kv_cache_scheme is not None + and kv_cache_dtype == "auto" + and checkpoint_kv_quant_allowed(vllm_config) + ): kv_cache_dtype = "fp8" calculate_kv_scales = False if cache_config is not None: