@tduncan / Nvidia-Optimization-Stack-VLLM / commits / dd11df04f3

[Misc] Remove deprecated calculate_kv_scales runtime KV scale calculation (#49389)

Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>

wangxiyuan committed Aug 3, 2026 at 23:15 UTC dd11df04f3b7046c40f13e586ac38a3725bc3c03

19 files changed +16 -263

.buildkite/test-amd.yaml

+1 -1

@@ -636,7 +636,7 @@ steps:

636 - tests/compile

637 - vllm/platforms/rocm.py

638 commands:

639 - - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'

639 + - pytest -v -s compile/fullgraph/test_full_graph.py

640

641 - label: Pytorch Nightly Dependency Override Check # TBD

642 timeout_in_minutes: 180

.buildkite/test_areas/compile.yaml

-3

@@ -90,9 +90,6 @@ steps:

90 - pytest -v -s tests/compile/passes/test_silu_mul_quant_fusion.py

91 # this runner has 2 GPUs available even though num_devices=2 is not set

92 - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py

93 - # test_fp8_kv_scale_compile requires FlashAttention (not supported on default L4/L40)

94 - # TODO(luka) move to H100 once pass tests run on H100

95 - - pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile

94 - label: Fusion E2E Quick (H100)

95 key: fusion-e2e-quick-h100

.buildkite/test_areas/pytorch.yaml

+1 -2

@@ -211,8 +211,7 @@ steps:

211 - vllm/v1/

212 - tests/compile

213 commands:

214 - # fp8 kv scales not supported on sm89, tested on Blackwell instead

215 - - pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'

214 + - pytest -v -s compile/fullgraph/test_full_graph.py

215

216 - label: Pytorch Nightly Dependency Override Check # 2min

217 key: pytorch-nightly-dependency-override-check

docs/design/metrics.md

+1 -1

@@ -367,7 +367,7 @@ We use this concept for the `vllm:cache_config_info` metric:

367 ```text

368 # HELP vllm:cache_config_info Information of the LLMEngine CacheConfig

369 # TYPE vllm:cache_config_info gauge

370 -vllm:cache_config_info{block_size="16",cache_dtype="auto",calculate_kv_scales="False",cpu_offload_gb="0",enable_prefix_caching="False",gpu_memory_utilization="0.9",...} 1.0

370 +vllm:cache_config_info{block_size="16",cache_dtype="auto",cpu_offload_gb="0",enable_prefix_caching="False",gpu_memory_utilization="0.9",...} 1.0

371 ```

372

373 However, `prometheus_client` has

docs/features/quantization/quantized_kvcache.md

+3 -33

@@ -27,18 +27,9 @@ You can configure how the quantization scales are computed in vLLM using three d

27 _Configure with:_

28 ```python

29 kv_cache_dtype="fp8"

30 - calculate_kv_scales=False

30 ```

33 -2. **Random token calibration (on-the-fly):**

34 - Scales are automatically estimated from a single batch of random tokens during warmup and then fixed.

35 - _Configure with:_

36 - ```python

37 - kv_cache_dtype="fp8"

38 - calculate_kv_scales=True

39 - ```

40 -

41 -3. **[Recommended] Calibration with a dataset (via llm-compressor):**

32 +2. **[Recommended] Calibration with a dataset (via llm-compressor):**

33 Scales are estimated using a curated calibration dataset for maximum accuracy.

34 This requires the [llm-compressor](https://github.com/vllm-project/llm-compressor) library.

35 _See example below!_

@@ -79,7 +70,7 @@ llm = LLM(

71 ## Examples

82 -### 1. No Calibration (`kv_cache_dtype="fp8"`, `calculate_kv_scales=False`)

73 +### 1. No Calibration (`kv_cache_dtype="fp8"`)

75 All quantization scales are set to 1.0.

@@ -90,27 +81,6 @@ sampling_params = SamplingParams(temperature=0.7, top_p=0.8)

81 llm = LLM(

82 model="meta-llama/Llama-2-7b-chat-hf",

83 kv_cache_dtype="fp8",

93 - calculate_kv_scales=False,

94 -)

95 -prompt = "London is the capital of"

96 -out = llm.generate(prompt, sampling_params)[0].outputs[0].text

97 -print(out)

98 -```

99 -

100 ----

101 -

102 -### 2. Random Token Calibration (`kv_cache_dtype="fp8"`, `calculate_kv_scales=True`)

103 -

104 -Scales are automatically estimated from a single batch of tokens during warmup.

105 -

106 -```python

107 -from vllm import LLM, SamplingParams

108 -

109 -sampling_params = SamplingParams(temperature=0.7, top_p=0.8)

110 -llm = LLM(

111 - model="meta-llama/Llama-2-7b-chat-hf",

112 - kv_cache_dtype="fp8",

113 - calculate_kv_scales=True,

84 )

85 prompt = "London is the capital of"

86 out = llm.generate(prompt, sampling_params)[0].outputs[0].text

@@ -119,7 +89,7 @@ print(out)

90 ---

122 -### 3. **[Recommended] Calibration Using a Dataset (with `llm-compressor`)**

92 +### 2. **[Recommended] Calibration Using a Dataset (with `llm-compressor`)**

94 For the highest-quality quantization, we recommend calibrating against a dataset using `llm-compressor`. This enables advanced strategies such as per-attention-head quantization.

tests/compile/fullgraph/test_full_graph.py

-32

@@ -13,7 +13,6 @@ from vllm import LLM, SamplingParams

13 from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode, PassConfig

14 from vllm.platforms import current_platform

15 from vllm.utils.torch_utils import is_torch_equal_or_newer

16 -from vllm.v1.attention.backends.registry import AttentionBackendEnum

17 from ...utils import create_new_process_for_each_test

@@ -182,37 +181,6 @@ def test_custom_compile_config(

181 run_model(compilation_config, model, **model_kwargs)

182

183

185 -@pytest.mark.parametrize(

186 - "compilation_mode",

187 - [CompilationMode.NONE, CompilationMode.VLLM_COMPILE],

188 -)

189 -@pytest.mark.parametrize(

190 - "model, backend",

191 - [

192 - ("Qwen/Qwen2-0.5B", None), # Standard attention model

193 - (

194 - "deepseek-ai/DeepSeek-V2-Lite",

195 - AttentionBackendEnum.FLASHINFER_MLA,

196 - ), # MLA (Multi-head Latent Attention) model

197 - ],

198 -)

199 -def test_fp8_kv_scale_compile(

200 - compilation_mode: CompilationMode,

201 - model: str,

202 - backend: AttentionBackendEnum | None,

203 -):

204 - model_kwargs = {

205 - "quantization": "fp8",

206 - "kv_cache_dtype": "fp8_e4m3",

207 - "calculate_kv_scales": True,

208 - "max_model_len": 512,

209 - }

210 - if backend:

211 - model_kwargs["attention_config"] = {"backend": backend.name}

212 -

213 - run_model(compilation_mode, model, **model_kwargs)

214 -

215 -

184 def run_model(

185 compile_config: CompilationMode | CompilationConfig, model: str, **model_kwargs

186 ):

tests/models/quantization/test_per_token_kv_cache.py

-4

@@ -52,9 +52,6 @@ def test_per_token_head_kv_cache_accuracy(

52 ) -> None:

53 """Compare logprobs between bf16 baseline and per-token-head quantized KV

54 cache.

55 -

56 - Uses calculate_kv_scales (dynamic scale computation) since there are

57 - no per-token-head calibrated checkpoints available yet.

55 """

56 with monkeypatch.context() as m:

57 m.setenv("TOKENIZERS_PARALLELISM", "true")

@@ -80,7 +77,6 @@ def test_per_token_head_kv_cache_accuracy(

77 tensor_parallel_size=tensor_parallel_size,

78 enforce_eager=enforce_eager,

79 kv_cache_dtype=kv_cache_dtype,

83 - calculate_kv_scales=True,

80 attention_config={"backend": backend},

81 ) as vllm_model:

82 test_outputs = vllm_model.generate_greedy_logprobs(

tests/quantization/test_fp8.py

+8 -20

@@ -18,7 +18,6 @@ from vllm.model_executor.kernels.linear.scaled_mm import (

18 MarlinFP8ScaledMMLinearKernel,

19 )

20 from vllm.model_executor.layers.attention.attention import (

21 - Attention,

21 set_default_quant_scales,

22 )

23 from vllm.model_executor.layers.fused_moe import FusedMoEFactory

@@ -433,8 +432,7 @@ def test_fp8_reloading(

432 method.process_weights_after_loading(layer)

433

434

436 -@pytest.mark.parametrize("source", ["checkpoint", "runtime_calc"])

437 -def test_kv_cache_scale_sync_to_host_copies(source):

435 +def test_kv_cache_scale_sync_to_host_copies():

436 """Test device-to-host sync of the k/v quantization scales, for both the

437 checkpoint-load and runtime-calc paths that produce them.

438 """

@@ -442,23 +440,13 @@ def test_kv_cache_scale_sync_to_host_copies(source):

440 set_default_quant_scales(layer, register_buffer=True)

441 layer.kv_cache_dtype = "fp8"

442

445 - if source == "checkpoint":

446 - # Scales come from the checkpoint, so runtime calc is disabled.

447 - layer.calculate_kv_scales = False

448 - method = BaseKVCacheMethod(quant_config=None)

449 - method.create_weights(layer)

450 - # 0.3 stays != 1.0 even after the fp8_fnuz x2 rescale.

451 - checkpoint_scale = torch.tensor(0.3, dtype=torch.float32)

452 - layer.k_scale.weight_loader(layer.k_scale, checkpoint_scale)

453 - layer.v_scale.weight_loader(layer.v_scale, checkpoint_scale)

454 - method.process_weights_after_loading(layer)

455 - else:

456 - # First forward computes distinct, non-unity scales from live k/v.

457 - layer.calculate_kv_scales = True

458 - query = torch.full((4, 8), 10.0)

459 - key = torch.full((4, 8), 60.0)

460 - value = torch.full((4, 8), 50.0)

461 - Attention.calc_kv_scales(layer, query, key, value)

443 + method = BaseKVCacheMethod(quant_config=None)

444 + method.create_weights(layer)

445 + # 0.3 stays != 1.0 even after the fp8_fnuz x2 rescale.

446 + checkpoint_scale = torch.tensor(0.3, dtype=torch.float32)

447 + layer.k_scale.weight_loader(layer.k_scale, checkpoint_scale)

448 + layer.v_scale.weight_loader(layer.v_scale, checkpoint_scale)

449 + method.process_weights_after_loading(layer)

450

451 assert layer._k_scale_float != 1.0

452 assert layer._v_scale_float != 1.0

tests/quantization/test_per_token_kv_cache.py

-1

@@ -497,7 +497,6 @@ def test_process_weights_sets_placeholder_scales(kv_cache_dtype: str):

497

498 layer = MagicMock()

499 layer.kv_cache_dtype = kv_cache_dtype

500 - layer.calculate_kv_scales = False

500 layer.k_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)

501 layer.v_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)

502 layer.q_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)

vllm/config/cache.py

-17

@@ -108,11 +108,6 @@ class CacheConfig:

108 security risk tolerance against the performance benefits before turning this on.

109 - "xxhash_cbor" combines canonical CBOR serialization with xxHash for

110 reproducible hashing. Requires the optional ``xxhash`` package."""

111 - calculate_kv_scales: bool = False

112 - """Deprecated: This option is deprecated and will be removed in v0.19.

113 - It enables dynamic calculation of `k_scale` and `v_scale` when

114 - kv_cache_dtype is fp8. If `False`, the scales will be loaded from the model

115 - checkpoint if available. Otherwise, the scales will default to 1.0."""

111 kv_cache_dtype_skip_layers: list[str] = field(default_factory=list)

112 """Layer patterns to skip KV cache quantization. Accepts layer indices

113 (e.g., '0', '2', '4') or attention type names (e.g., 'sliding_window')."""

@@ -270,18 +265,6 @@ class CacheConfig:

265 self.user_specified_mamba_block_size = True

266 return self

267

273 - @field_validator("calculate_kv_scales", mode="after")

274 - @classmethod

275 - def _warn_deprecated_calculate_kv_scales(cls, calculate_kv_scales: bool) -> bool:

276 - if calculate_kv_scales:

277 - logger.warning(

278 - "The `--calculate-kv-scales` option is deprecated and will "

279 - "be removed in v0.19. The scales will be loaded from the "

280 - "model checkpoint if available, otherwise they default to "

281 - "1.0."

282 - )

283 - return calculate_kv_scales

284 -

268 @field_validator("cache_dtype", mode="after")

269 @classmethod

270 def _validate_cache_dtype(cls, cache_dtype: CacheDType) -> CacheDType:

vllm/engine/arg_utils.py

-5

@@ -695,7 +695,6 @@ class EngineArgs:

695 override_attention_dtype: str | None = ModelConfig.override_attention_dtype

696 attention_backend: AttentionBackendEnum | None = AttentionConfig.backend

697

698 - calculate_kv_scales: bool = CacheConfig.calculate_kv_scales

698 kv_cache_dtype_skip_layers: list[str] = get_field(

699 CacheConfig, "kv_cache_dtype_skip_layers"

700 )

@@ -1210,9 +1209,6 @@ class EngineArgs:

1209 cache_group.add_argument(

1210 "--prefix-caching-hash-algo", **cache_kwargs["prefix_caching_hash_algo"]

1211 )

1213 - cache_group.add_argument(

1214 - "--calculate-kv-scales", **cache_kwargs["calculate_kv_scales"]

1215 - )

1212 cache_group.add_argument(

1213 "--kv-cache-dtype-skip-layers", **cache_kwargs["kv_cache_dtype_skip_layers"]

1214 )

@@ -1957,7 +1953,6 @@ class EngineArgs:

1953 sliding_window=sliding_window,

1954 enable_prefix_caching=self.enable_prefix_caching,

1955 prefix_caching_hash_algo=self.prefix_caching_hash_algo,

1960 - calculate_kv_scales=self.calculate_kv_scales,

1956 kv_cache_dtype_skip_layers=self.kv_cache_dtype_skip_layers,

1957 kv_sharing_fast_prefill=self.kv_sharing_fast_prefill,

1958 mamba_cache_dtype=self.mamba_cache_dtype,

vllm/envs.py

-6

@@ -1325,12 +1325,6 @@ environment_variables: dict[str, Callable[[], Any]] = {

1325 "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB": lambda: maybe_convert_int(

1326 os.environ.get("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB", None)

1327 ),

1328 - # Divisor for dynamic query scale factor calculation for FP8 KV Cache

1329 - "Q_SCALE_CONSTANT": lambda: int(os.getenv("Q_SCALE_CONSTANT", "200")),

1330 - # Divisor for dynamic key scale factor calculation for FP8 KV Cache

1331 - "K_SCALE_CONSTANT": lambda: int(os.getenv("K_SCALE_CONSTANT", "200")),

1332 - # Divisor for dynamic value scale factor calculation for FP8 KV Cache

1333 - "V_SCALE_CONSTANT": lambda: int(os.getenv("V_SCALE_CONSTANT", "100")),

1328 # If set, enable multiprocessing in LLM for the V1 code path.

1329 "VLLM_ENABLE_V1_MULTIPROCESSING": lambda: bool(

1330 int(os.getenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1"))

vllm/model_executor/layers/attention/attention.py

-62

@@ -144,11 +144,6 @@ def set_default_quant_scales(layer: nn.Module, register_buffer: bool = False) ->

144 layer._v_scale_cpu = torch.tensor(1.0, dtype=torch.float32)

145 layer._prob_scale_float = 1.0

146

147 - # Initialize q/k/v range constants used by calc_kv_scales

148 - layer.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)

149 - layer.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)

150 - layer.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32)

151 -

147

148 def _init_kv_cache_quant(

149 layer: nn.Module,

@@ -270,10 +265,8 @@ class Attention(nn.Module, AttentionLayerBase):

265 vllm_config = get_current_vllm_config()

266 if cache_config is not None:

267 kv_cache_dtype = cache_config.cache_dtype

273 - calculate_kv_scales = cache_config.calculate_kv_scales

268 else:

269 kv_cache_dtype = "auto"

276 - calculate_kv_scales = False

270

271 # llm-compressor models declare an FP8 KV-cache scheme in their

272 # checkpoint config. Honor it only when the user did not explicitly

@@ -284,10 +277,8 @@ class Attention(nn.Module, AttentionLayerBase):

277 kv_cache_scheme = getattr(quant_config, "kv_cache_scheme", None)

278 if kv_cache_scheme is not None and kv_cache_dtype == "auto":

279 kv_cache_dtype = "fp8"

287 - calculate_kv_scales = False

280 if cache_config is not None:

281 cache_config.cache_dtype = "fp8"

290 - cache_config.calculate_kv_scales = False

282

283 # Check if per-head quant scales are required based on kv_cache_scheme

284 use_per_head_quant_scales = (

@@ -312,7 +303,6 @@ class Attention(nn.Module, AttentionLayerBase):

303 skip = True

304 if skip:

305 kv_cache_dtype = "auto"

315 - calculate_kv_scales = False

306 logger.debug(

307 "Layer %s: kv_cache_dtype=%s, sliding_window=%s",

308 prefix,

@@ -324,7 +314,6 @@ class Attention(nn.Module, AttentionLayerBase):

314 kv_cache_dtype, vllm_config.model_config

315 )

316 self.kv_cache_dtype = kv_cache_dtype

327 - self.calculate_kv_scales = calculate_kv_scales

317 if num_kv_heads is None:

318 num_kv_heads = num_heads

319 assert num_heads % num_kv_heads == 0, (

@@ -505,10 +494,6 @@ class Attention(nn.Module, AttentionLayerBase):

494 context using

495 `vllm.forward_context.get_forward_context().attn_metadata`.

496 """

508 - if self.calculate_kv_scales:

509 - torch.ops.vllm.maybe_calc_kv_scales(

510 - query, key, value, _encode_layer_name(self.layer_name)

511 - )

497 if output_dtype is None:

498 output_dtype = query.dtype

499 if self.query_quant is not None:

@@ -581,18 +566,6 @@ class Attention(nn.Module, AttentionLayerBase):

566 )

567 return output.view(-1, hidden_size)

568

584 - def calc_kv_scales(self, query, key, value):

585 - self._q_scale.copy_(torch.abs(query).max() / self.q_range)

586 - self._k_scale.copy_(torch.abs(key).max() / self.k_range)

587 - self._v_scale.copy_(torch.abs(value).max() / self.v_range)

588 - self._q_scale_float = self._q_scale.item()

589 - self._k_scale_float = self._k_scale.item()

590 - self._v_scale_float = self._v_scale.item()

591 - self._k_scale_cpu.fill_(self._k_scale_float)

592 - self._v_scale_cpu.fill_(self._v_scale_float)

593 - # We only calculate the scales once

594 - self.calculate_kv_scales = False

595 -

569 def extra_repr(self) -> str:

570 s = f"head_size={self.impl.head_size}" # type: ignore

571 s += f", num_heads={self.impl.num_heads}" # type: ignore

@@ -694,41 +667,6 @@ class Attention(nn.Module, AttentionLayerBase):

667 )

668

669

697 -def maybe_calc_kv_scales(

698 - query: torch.Tensor,

699 - key: torch.Tensor,

700 - value: torch.Tensor,

701 - layer_name: LayerNameType,

702 -) -> None:

703 - layer_name = _resolve_layer_name(layer_name)

704 - forward_context: ForwardContext = get_forward_context()

705 - self = forward_context.no_compile_layers[layer_name]

706 -

707 - # Only calculate if the layer's calculate_kv_scales flag is True

708 - # This flag gets set to False after the first forward pass

709 - if not self.calculate_kv_scales:

710 - return

711 -

712 - self.calc_kv_scales(query, key, value)

713 -

714 -

715 -def maybe_calc_kv_scales_fake(

716 - query: torch.Tensor,

717 - key: torch.Tensor,

718 - value: torch.Tensor,

719 - layer_name: LayerNameType,

720 -) -> None:

721 - return

722 -

723 -

724 -direct_register_custom_op(

725 - op_name="maybe_calc_kv_scales",

726 - op_func=maybe_calc_kv_scales,

727 - mutates_args=["query", "key", "value"],

728 - fake_impl=maybe_calc_kv_scales_fake,

729 -)

730 -

731 -

670 def get_attention_context(

671 layer_name: str,

672 ) -> tuple[Any, "Attention | MLAAttention", torch.Tensor, torch.Tensor]:

vllm/model_executor/layers/attention/mla_attention.py

-41

@@ -399,10 +399,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):

399

400 if cache_config is not None:

401 kv_cache_dtype: CacheDType = cache_config.cache_dtype

402 - calculate_kv_scales = cache_config.calculate_kv_scales

402 else:

403 kv_cache_dtype = "auto"

405 - calculate_kv_scales = False

404 self.quant_config = quant_config

405

406 if cache_config is not None and cache_config.kv_cache_dtype_skip_layers:

@@ -411,7 +409,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):

409 layer_idx = extract_layer_index(prefix)

410 if str(layer_idx) in cache_config.kv_cache_dtype_skip_layers:

411 kv_cache_dtype = "auto"

414 - calculate_kv_scales = False

412 logger.debug(

413 "Layer %s: kv_cache_dtype=%s",

414 prefix,

@@ -460,7 +457,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):

457

458 # Initialize KV cache quantization attributes

459 self.kv_cache_dtype = kv_cache_dtype

463 - self.calculate_kv_scales = calculate_kv_scales

460 _init_kv_cache_quant(self, quant_config, prefix)

461

462 if (

@@ -560,11 +556,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):

556 and _vllm_config.parallel_config.dcp_comm_backend == "a2a"

557 )

558

563 - # Initialize q/k/v range constants.

564 - self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)

565 - self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)

566 - self.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32)

567 -

559 self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()

560

561 # If kv_b_proj_weight is unquantized, quantize it to mxfp4 if supported

@@ -606,14 +597,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):

597 output_shape: torch.Size | None = None,

598 q_dcp_replicated: torch.Tensor | None = None,

599 ) -> torch.Tensor:

609 - if self.calculate_kv_scales:

610 - torch.ops.vllm.maybe_calc_kv_scales(

611 - q,

612 - kv_c_normed,

613 - k_pe,

614 - _encode_layer_name(self.layer_name),

615 - )

616 -

600 if self.use_direct_call:

601 forward_context: ForwardContext = get_forward_context()

602 attn_metadata_raw = forward_context.attn_metadata

@@ -1110,30 +1093,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):

1093 if not should_load_quant_weights(quant_method):

1094 set_default_quant_scales(self, register_buffer=False)

1095

1113 - def calc_kv_scales(

1114 - self, q: torch.Tensor, kv_c_normed: torch.Tensor, k_pe: torch.Tensor

1115 - ) -> None:

1116 - """Optional scale calculation for MLA inputs.

1117 -

1118 - Mirrors Attention.calc_kv_scales. Not all MLA backends require this

1119 - """

1120 - # Use safe defaults if ranges are not present

1121 - q_range = getattr(self, "q_range", torch.tensor(1.0))

1122 - k_range = getattr(self, "k_range", torch.tensor(1.0))

1123 - v_range = getattr(self, "v_range", torch.tensor(1.0))

1124 -

1125 - self._q_scale.copy_(torch.abs(q).max() / q_range)

1126 - # kv_c_normed is the compressed KV representation; use it for k/v

1127 - kv_abs_max = torch.abs(kv_c_normed).max()

1128 - self._k_scale.copy_(kv_abs_max / k_range)

1129 - self._v_scale.copy_(kv_abs_max / v_range)

1130 - self._q_scale_float = self._q_scale.item()

1131 - self._k_scale_float = self._k_scale.item()

1132 - self._v_scale_float = self._v_scale.item()

1133 - self._k_scale_cpu.fill_(self._k_scale_float)

1134 - self._v_scale_cpu.fill_(self._v_scale_float)

1135 - self.calculate_kv_scales = False

1136 -

1096 def get_attn_backend(self) -> type[AttentionBackend]:

1097 return self.attn_backend

1098

vllm/model_executor/layers/quantization/kv_cache.py

+2 -7

@@ -80,8 +80,7 @@ class BaseKVCacheMethod(QuantizeMethodBase):

80 return

82 # Per-token-head quantized KV cache: scales are computed dynamically

83 - # per (token, head) in the kernel at cache-write time. Checkpoint

84 - # scales are never used regardless of calculate_kv_scales.

83 + # per (token, head) in the kernel at cache-write time.

84 if kv_cache_uses_per_token_head_scales(layer.kv_cache_dtype):

85 layer._k_scale.copy_(1.0)

86 layer._v_scale.copy_(1.0)

@@ -97,10 +96,7 @@ class BaseKVCacheMethod(QuantizeMethodBase):

96 # regardless whether the kv-scale is available in the checkpoint.

97 # No need to process kv scales after loading if we are going to

98 # calculate them on the fly.

100 - if (

101 - is_quantized_kv_cache(layer.kv_cache_dtype)

102 - and not layer.calculate_kv_scales

103 - ):

99 + if is_quantized_kv_cache(layer.kv_cache_dtype):

100 if layer.k_scale > 0.0 and layer.v_scale > 0.0:

101 # We prefer to use separate k_scale and v_scale if present

102 k_scale = layer.k_scale.to("cpu").tolist()

@@ -158,7 +154,6 @@ class BaseKVCacheMethod(QuantizeMethodBase):

154 q_scale = layer.q_scale

155 if current_platform.is_fp8_fnuz():

156 q_scale *= 2

161 - layer.calculate_kv_scales = False

157 else:

158 q_scale = 1.0

159 if layer.prob_scale > 0.0:

vllm/model_executor/models/config.py

-17

@@ -445,23 +445,6 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):

445 Args:

446 vllm_config: vLLM Config

447 """

448 - cache_config = vllm_config.cache_config

449 -

450 - # Disable calculate_kv_scales for hybrid models: uninitialized

451 - # recurrent state corrupts scales during the calibration pass.

452 - # See issue: https://github.com/vllm-project/vllm/issues/37554

453 -

454 - if cache_config.calculate_kv_scales:

455 - logger.warning(

456 - "Disabling calculate_kv_scales for hybrid model '%s'. "

457 - "Hybrid models with recurrent layers (GDN, Mamba, SSM) "

458 - "produce unreliable KV cache scales during the "

459 - "calibration pass because recurrent state is "

460 - "uninitialized. Using default scale of 1.0 instead.",

461 - vllm_config.model_config.model,

462 - )

463 - cache_config.calculate_kv_scales = False

464 -

448 # Enable FULL_AND_PIECEWISE by default

449 MambaModelConfig.verify_and_update_config(vllm_config)

450

vllm/models/minimax_m3/amd/model.py

-1

@@ -642,7 +642,6 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):

642 # MiniMax-M3 sparse attention owns its KV-cache insert/read path instead

643 # of wrapping the generic Attention module. Keep the same runtime scale

644 # attributes so FP8 KV reads can honor vLLM's per-layer descale contract.

645 - self.calculate_kv_scales = False

645 set_default_quant_scales(self, register_buffer=True)

646

647 # Shared top-k buffer: the indexer writes the selected blocks into it and

vllm/models/minimax_m3/nvidia/model.py

-1

@@ -495,7 +495,6 @@ class MiniMaxM3SparseAttention(nn.Module, AttentionLayerBase):

495 # MiniMax-M3 sparse attention owns its KV-cache insert/read path instead

496 # of wrapping the generic Attention module. Keep the same runtime scale

497 # attributes so FP8 KV reads can honor vLLM's per-layer descale contract.

498 - self.calculate_kv_scales = False

498 set_default_quant_scales(self, register_buffer=True)

499 # Indexer side-cache dtype, mirroring --kv-cache-dtype for the main

500 # cache (--attention-config '{"indexer_kv_dtype": ...}').

vllm/v1/worker/gpu_model_runner.py

-9

@@ -499,7 +499,6 @@ class GPUModelRunner(

499 self.max_model_len = model_config.max_model_len

500

501 # Always set to false after the first forward pass

502 - self.calculate_kv_scales = self.cache_config.calculate_kv_scales

502 self.dcp_world_size = self.parallel_config.decode_context_parallel_size

503 self.dcp_rank = 0 if self.dcp_world_size <= 1 else get_dcp_group().rank_in_group

504 self.max_num_tokens = scheduler_config.max_num_batched_tokens

@@ -4402,14 +4401,6 @@ class GPUModelRunner(

4401 scheduler_output, num_tokens_padded, intermediate_tensors

4402 )

4403

4405 - # Set cudagraph mode to none if calc_kv_scales is true.

4406 - # KV scales calculation involves dynamic operations that are incompatible

4407 - # with CUDA graph capture.

4408 - if self.calculate_kv_scales:

4409 - cudagraph_mode = CUDAGraphMode.NONE

4410 - # Mark KV scales as calculated after the first forward pass

4411 - self.calculate_kv_scales = False

4412 -

4404 # Encoder-decoder models can only compile the pure decode steps where no

4405 # encoder inputs are present. Use eager for the first pass.

4406 num_encoder_reqs = len(scheduler_output.scheduled_encoder_inputs)