From 8c98b85cf08de6c6c7c80af6abb0b0b349e131f9 Mon Sep 17 00:00:00 2001 From: Peuqui Date: Sat, 12 Sep 2026 21:16:58 +0200 Subject: [PATCH] [Bugfix][SM70/SM75] Honor a checkpoint's KV-cache quantization directive only on Ampere and newer With --kv-cache-dtype auto, a ModelOpt kv_cache_quant_algo or a compressed-tensors kv_cache_scheme turned the KV cache into FP8 without the user asking. That metadata describes the weights; on Volta and Turing there is no FP8 hardware. Volta unpacks the cache in software and loses the tensor-core decode route (4x V100, Qwen3.8-27B: +4.82 ms per decode round), Turing does not boot at all (FlashAttention rejects fp8 below FA3, the compiled Triton cache write fails with "fp8e4nv not supported"). Honor the directive only when every participating CUDA device is Ampere or newer, reusing the participation helpers from #579 so a visible but unused card does not decide and a mixed rig decides once for all stages. An explicit --kv-cache-dtype is never touched. The compressed-tensors re-apply path in attention.py follows the same policy. Signed-off-by: Peuqui Co-authored-by: Claude Fable 5.1 --- .../config/test_checkpoint_kv_quant_policy.py | 103 ++++++++++++++++++ vllm/config/cache.py | 9 +- vllm/config/vllm.py | 31 ++++++ vllm/engine/arg_utils.py | 4 + .../layers/attention/attention.py | 11 +- 5 files changed, 155 insertions(+), 3 deletions(-) create mode 100644 tests/config/test_checkpoint_kv_quant_policy.py 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: