@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 62215e72c6
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 1, 21:15 UTC

Remove KV cache scale boilerplate from model weight loading methods (#43167)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Harry Mellor committed Jun 5, 2026 at 13:19 UTC 62215e72c6ff4979009c27aa95f7449213007f8f
56 files changed +88 -731
tests/model_executor/test_eagle_quantization.py
-56
@@ -100,32 +100,6 @@ def test_fc_layer_quant_config_usage(default_vllm_config, dist_init, device) ->
100 assert output.shape == (2, output_size)
101
102
103 -def test_kv_cache_scale_name_handling():
104 - # Mock a quant config that supports cache scales
105 - mock_quant_config = Mock()
106 - mock_quant_config.get_cache_scale = Mock(return_value="layers.0.self_attn.kv_scale")
107 -
108 - # Condition check in load_weights
109 - name = "layers.0.self_attn.k_proj.weight"
110 - scale_name = mock_quant_config.get_cache_scale(name)
111 -
112 - # Check if get_cache_scale is called and returns expected value
113 - mock_quant_config.get_cache_scale.assert_called_once_with(name)
114 - assert scale_name == "layers.0.self_attn.kv_scale"
115 -
116 -
117 -def test_kv_cache_scale_name_no_scale():
118 - # Mock a quant config that returns None for get_cache_scale
119 - mock_quant_config = Mock()
120 - mock_quant_config.get_cache_scale = Mock(return_value=None)
121 -
122 - name = "layers.0.mlp.gate_proj.weight"
123 - scale_name = mock_quant_config.get_cache_scale(name)
124 -
125 - # Should return None for weights that don't have cache scales
126 - assert scale_name is None
127 -
128 -
103 def test_maybe_remap_kv_scale_name():
104 from vllm.model_executor.model_loader.weight_utils import maybe_remap_kv_scale_name
105
@@ -183,33 +157,3 @@ def test_eagle3_lm_head_receives_quant_config():
157 assert call_kwargs["quant_config"] is mock_quant_config, (
158 "ParallelLMHead must receive the draft model's quant_config"
159 )
186 -
187 -
188 -def test_load_weights_kv_scale_handling():
189 - kv_scale_param = Mock()
190 - kv_scale_param.weight_loader = Mock()
191 -
192 - params_dict = {
193 - "layers.0.self_attn.kv_scale": kv_scale_param,
194 - }
195 -
196 - mock_quant_config = Mock()
197 - mock_quant_config.get_cache_scale = Mock(return_value="layers.0.self_attn.kv_scale")
198 -
199 - # Load_weights logic for KV cache scales
200 - name = "layers.0.self_attn.k_proj.weight"
201 - loaded_weight_tensor = torch.tensor([1.0, 2.0])
202 -
203 - if mock_quant_config is not None:
204 - scale_name = mock_quant_config.get_cache_scale(name)
205 - if scale_name:
206 - param = params_dict[scale_name]
207 - assert param is kv_scale_param
208 - weight_to_load = (
209 - loaded_weight_tensor
210 - if loaded_weight_tensor.dim() == 0
211 - else loaded_weight_tensor[0]
212 - )
213 -
214 - assert scale_name == "layers.0.self_attn.kv_scale"
215 - assert weight_to_load == loaded_weight_tensor[0]
vllm/model_executor/layers/quantization/base_config.py
+7 -1
@@ -162,7 +162,13 @@ class QuantizationConfig(ABC):
162 """
163 raise NotImplementedError
164
165 - def get_cache_scale(self, name: str) -> str | None:
165 + def get_cache_scale_mapper(self) -> "WeightsMapper | None":
166 + """Mapping from checkpoint KV-cache scale names to vLLM scale names.
167 +
168 + Returning a mapper here causes `AutoWeightsLoader` to apply it to the
169 + weight stream automatically; individual model `load_weights` methods
170 + do not need to know about KV-cache scales.
171 + """
172 return None
173
174 def apply_vllm_mapper( # noqa: B027
vllm/model_executor/layers/quantization/fp8.py
+12 -19
@@ -207,25 +207,18 @@ class Fp8Config(QuantizationConfig):
207 return Fp8KVCacheMethod(self)
208 return None
209
210 - def get_cache_scale(self, name: str) -> str | None:
211 - """
212 - Check whether the param name matches the format for k/v cache scales
213 - in compressed-tensors. If this is the case, return its equivalent
214 - param name expected by vLLM
215 -
216 - :param name: param name
217 - :return: matching param name for KV cache scale in vLLM
218 - """
219 - if name.endswith(".output_scale") and ".k_proj" in name:
220 - return name.replace(".k_proj.output_scale", ".attn.k_scale")
221 - if name.endswith(".output_scale") and ".v_proj" in name:
222 - return name.replace(".v_proj.output_scale", ".attn.v_scale")
223 - if name.endswith(".output_scale") and ".q_proj" in name:
224 - return name.replace(".q_proj.output_scale", ".attn.q_scale")
225 - if name.endswith("self_attn.prob_output_scale"):
226 - return name.replace(".prob_output_scale", ".attn.prob_scale")
227 - # If no matches, return None
228 - return None
210 + def get_cache_scale_mapper(self) -> "WeightsMapper":
211 + """Map compressed-tensors KV-cache scale names to vLLM names."""
212 + from vllm.model_executor.models.utils import WeightsMapper
213 +
214 + return WeightsMapper(
215 + orig_to_new_suffix={
216 + ".k_proj.output_scale": ".attn.k_scale",
217 + ".v_proj.output_scale": ".attn.v_scale",
218 + ".q_proj.output_scale": ".attn.q_scale",
219 + ".self_attn.prob_output_scale": ".self_attn.attn.prob_scale",
220 + }
221 + )
222
223
224 class CopyNumelCounter(TorchDispatchMode):
vllm/model_executor/layers/quantization/kv_cache.py
+28 -4
@@ -15,6 +15,30 @@ from vllm.v1.kv_cache_interface import kv_cache_uses_per_token_head_scales
15 logger = init_logger(__name__)
16
17
18 +class KVCacheScaleParameter(torch.nn.Parameter):
19 + """Scalar parameter for KV-cache scales.
20 +
21 + Initialized to -1.0 (an invalid sentinel) so call sites just write
22 + `KVCacheScaleParameter()`. The `weight_loader` accepts shape `()` or
23 + `(1,)` and rejects anything else — per-head scales go through a separate
24 + path (compressed-tensors' `_tp_aware_loader`), not this one. Per-instance
25 + overrides still work because instance attribute assignment shadows this
26 + class-level loader.
27 + """
28 +
29 + def __new__(cls) -> "KVCacheScaleParameter":
30 + return super().__new__(cls, torch.tensor(-1.0), requires_grad=False)
31 +
32 + @staticmethod
33 + def weight_loader(param: torch.nn.Parameter, loaded_weight: torch.Tensor) -> None:
34 + if loaded_weight.numel() != 1:
35 + raise ValueError(
36 + f"KV-cache scale expects a scalar weight, got shape "
37 + f"{tuple(loaded_weight.shape)}"
38 + )
39 + param.data.copy_(loaded_weight.reshape(()))
40 +
41 +
42 class BaseKVCacheMethod(QuantizeMethodBase):
43 """
44 Quant method that adds `_k_scale` and `_v_scale` attributes to the
@@ -37,11 +61,11 @@ class BaseKVCacheMethod(QuantizeMethodBase):
61 # Initialize the Q and KV cache scales to -1.0, an invalid value.
62 # If the q and k/v_scales appear in the checkpoint, it will be
63 # overwritten when loading weights.
40 - layer.q_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)
41 - layer.k_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)
42 - layer.v_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)
64 + layer.q_scale = KVCacheScaleParameter()
65 + layer.k_scale = KVCacheScaleParameter()
66 + layer.v_scale = KVCacheScaleParameter()
67 # Initialize P = softmax(QK^T) scales
44 - layer.prob_scale = torch.nn.Parameter(torch.tensor(-1.0), requires_grad=False)
68 + layer.prob_scale = KVCacheScaleParameter()
69
70 def apply(self, layer: torch.nn.Module) -> torch.Tensor:
71 raise RuntimeError(f"{self.__class__.__name__}.apply should not be called.")
vllm/model_executor/layers/quantization/quark/quark.py
+10 -20
@@ -646,26 +646,16 @@ class QuarkConfig(QuantizationConfig):
646
647 return scheme
648
649 - def get_cache_scale(self, name: str) -> str | None:
650 - """
651 - Check whether the param name matches the format for k/v cache scales
652 - in quark. If this is the case, return its equivalent param name
653 - expected by vLLM
654 -
655 - :param name: param name
656 - :return: matching param name for KV cache scale in vLLM
657 - """
658 - if name.endswith(".output_scale") and ".k_proj" in name:
659 - return name.replace(".k_proj.output_scale", ".attn.k_scale")
660 - if name.endswith(".output_scale") and ".v_proj" in name:
661 - return name.replace(".v_proj.output_scale", ".attn.v_scale")
662 - if name.endswith(".output_scale") and ".q_proj" in name:
663 - return name.replace(".q_proj.output_scale", ".attn.q_scale")
664 - if name.endswith("self_attn.prob_output_scale"):
665 - return name.replace(".prob_output_scale", ".attn.prob_scale")
666 -
667 - # If no matches, return None
668 - return None
649 + def get_cache_scale_mapper(self) -> "WeightsMapper":
650 + """Map Quark KV-cache scale names to vLLM names."""
651 + return WeightsMapper(
652 + orig_to_new_suffix={
653 + ".k_proj.output_scale": ".attn.k_scale",
654 + ".v_proj.output_scale": ".attn.v_scale",
655 + ".q_proj.output_scale": ".attn.q_scale",
656 + ".self_attn.prob_output_scale": ".self_attn.attn.prob_scale",
657 + }
658 + )
659
660
661 class QuarkLinearMethod(LinearMethodBase):
vllm/model_executor/model_loader/weight_utils.py
+5
@@ -1541,6 +1541,11 @@ def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> str | None:
1541 if no remapping is needed.
1542 None: If the remapped name is not found in params_dict.
1543 """
1544 + # Already in vLLM's expected form (e.g. weights pre-renamed by a
1545 + # `WeightsMapper` from the quant config). Skip the regex remap, which
1546 + # would otherwise double-apply the `.attn` prefix and drop the weight.
1547 + if name in params_dict:
1548 + return name
1549 if name.endswith(".kv_scale"):
1550 logger.warning_once(
1551 "DEPRECATED. Found kv_scale in the checkpoint. "
vllm/model_executor/models/apertus.py
-12
@@ -430,18 +430,6 @@ class ApertusModel(nn.Module, EagleModelMixin):
430 # Models trained using ColossalAI may include these tensors in
431 # the checkpoint. Skip them.
432 continue
433 - if self.quant_config is not None and (
434 - scale_name := self.quant_config.get_cache_scale(name)
435 - ):
436 - # Loading kv cache quantization scales
437 - param = params_dict[scale_name]
438 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
439 - loaded_weight = (
440 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
441 - )
442 - weight_loader(param, loaded_weight)
443 - loaded_params.add(scale_name)
444 - continue
433 if "scale" in name or "zero_point" in name:
434 # Remapping the name of FP8 kv-scale.
435 name = maybe_remap_kv_scale_name(name, params_dict)
vllm/model_executor/models/arcee.py
-12
@@ -293,18 +293,6 @@ class ArceeModel(nn.Module, EagleModelMixin):
293 if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
294 continue
295
296 - if self.quant_config is not None and (
297 - scale_name := self.quant_config.get_cache_scale(name)
298 - ):
299 - param = params_dict[scale_name]
300 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
301 - loaded_weight = (
302 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
303 - )
304 - weight_loader(param, loaded_weight)
305 - loaded_params.add(scale_name)
306 - continue
307 -
296 if "scale" in name or "zero_point" in name:
297 remapped_name = maybe_remap_kv_scale_name(name, params_dict)
298 if remapped_name is None:
vllm/model_executor/models/aria.py
-12
@@ -363,18 +363,6 @@ class AriaTextModel(LlamaModel, SupportsQuant):
363 # Models trained using ColossalAI may include these tensors in
364 # the checkpoint. Skip them.
365 continue
366 - if self.quant_config is not None and (
367 - scale_name := self.quant_config.get_cache_scale(name)
368 - ):
369 - # Loading kv cache quantization scales
370 - param = params_dict[scale_name]
371 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
372 - loaded_weight = (
373 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
374 - )
375 - weight_loader(param, loaded_weight)
376 - loaded_params.add(scale_name)
377 - continue
366 for param_name, weight_name, shard_id in stacked_params_mapping:
367 if weight_name not in name:
368 continue
vllm/model_executor/models/cohere2_moe.py
-12
@@ -464,18 +464,6 @@ class Cohere2MoeModel(nn.Module):
464 if "rotary_emb.inv_freq" in name:
465 continue
466
467 - if self.quant_config is not None and (
468 - scale_name := self.quant_config.get_cache_scale(name)
469 - ):
470 - param = params_dict[scale_name]
471 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
472 - loaded_weight = (
473 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
474 - )
475 - weight_loader(param, loaded_weight)
476 - loaded_params.add(scale_name)
477 - continue
478 -
467 for param_name, shard_name, shard_id in stacked_params_mapping:
468 if shard_name not in name:
469 continue
vllm/model_executor/models/cohere_eagle.py
-12
@@ -150,18 +150,6 @@ class CohereEagleModel(nn.Module):
150 if "rotary_emb.inv_freq" in name:
151 continue
152
153 - if self.quant_config is not None and (
154 - scale_name := self.quant_config.get_cache_scale(name)
155 - ):
156 - param = params_dict[scale_name]
157 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
158 - loaded_weight = (
159 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
160 - )
161 - weight_loader(param, loaded_weight)
162 - loaded_params.add(scale_name)
163 - continue
164 -
153 for param_name, weight_name, shard_id in stacked_params_mapping:
154 if weight_name not in name:
155 continue
vllm/model_executor/models/commandr.py
-13
@@ -352,19 +352,6 @@ class CohereModel(nn.Module):
352 params_dict = dict(self.named_parameters())
353 loaded_params: set[str] = set()
354 for name, loaded_weight in weights:
355 - if self.quant_config is not None and (
356 - scale_name := self.quant_config.get_cache_scale(name)
357 - ):
358 - # Loading kv cache quantization scales
359 - param = params_dict[scale_name]
360 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
361 - loaded_weight = (
362 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
363 - )
364 - weight_loader(param, loaded_weight)
365 - loaded_params.add(scale_name)
366 - continue
367 -
355 for param_name, shard_name, shard_id in stacked_params_mapping:
356 if shard_name not in name:
357 continue
vllm/model_executor/models/dbrx.py
-13
@@ -394,19 +394,6 @@ class DbrxModel(nn.Module):
394 loaded_params: set[str] = set()
395
396 for name, loaded_weight in weights:
397 - if self.quant_config is not None and (
398 - scale_name := self.quant_config.get_cache_scale(name)
399 - ):
400 - # Loading kv cache quantization scales
401 - param = params_dict[scale_name]
402 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
403 - loaded_weight = (
404 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
405 - )
406 - weight_loader(param, loaded_weight)
407 - loaded_params.add(scale_name)
408 - continue
409 -
397 if name.endswith(("w1", "w2", "v1")):
398 name = name + "_weight"
399 for param_name, weight_name in expert_params_mapping:
vllm/model_executor/models/deepseek_eagle3.py
-13
@@ -284,19 +284,6 @@ class DeepseekV2Eagle3Model(nn.Module):
284 if "midlayer." in name:
285 name = name.replace("midlayer.", "layers.0.")
286
287 - # Handle kv cache quantization scales
288 - if self.quant_config is not None and (
289 - scale_name := self.quant_config.get_cache_scale(name)
290 - ):
291 - param = params_dict[scale_name]
292 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
293 - loaded_weight = (
294 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
295 - )
296 - weight_loader(param, loaded_weight)
297 - loaded_params.add(scale_name)
298 - continue
299 -
287 # Remapping the name FP8 kv-scale
288 if "scale" in name:
289 name = maybe_remap_kv_scale_name(name, params_dict)
vllm/model_executor/models/exaone.py
-12
@@ -391,18 +391,6 @@ class ExaoneModel(nn.Module):
391 # Models trained using ColossalAI may include these tensors in
392 # the checkpoint. Skip them.
393 continue
394 - if self.quant_config is not None and (
395 - scale_name := self.quant_config.get_cache_scale(name)
396 - ):
397 - # Loading kv cache quantization scales
398 - param = params_dict[scale_name]
399 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
400 - loaded_weight = (
401 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
402 - )
403 - weight_loader(param, loaded_weight)
404 - loaded_params.add(scale_name)
405 - continue
394 for param_name, weight_name, shard_id in stacked_params_mapping:
395 if weight_name not in name:
396 continue
vllm/model_executor/models/exaone4.py
-12
@@ -389,18 +389,6 @@ class Exaone4Model(nn.Module):
389 # Models trained using ColossalAI may include these tensors in
390 # the checkpoint. Skip them.
391 continue
392 - if self.quant_config is not None and (
393 - scale_name := self.quant_config.get_cache_scale(name)
394 - ):
395 - # Loading kv cache quantization scales
396 - param = params_dict[scale_name]
397 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
398 - loaded_weight = (
399 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
400 - )
401 - weight_loader(param, loaded_weight)
402 - loaded_params.add(scale_name)
403 - continue
392 for param_name, weight_name, shard_id in stacked_params_mapping:
393 if weight_name not in name:
394 continue
vllm/model_executor/models/exaone_moe.py
-12
@@ -374,18 +374,6 @@ class ExaoneMoeModel(nn.Module):
374 # Models trained using ColossalAI may include these tensors in
375 # the checkpoint. Skip them.
376 continue
377 - if self.quant_config is not None and (
378 - scale_name := self.quant_config.get_cache_scale(name)
379 - ):
380 - # Loading kv cache quantization scales
381 - param = params_dict[scale_name]
382 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
383 - loaded_weight = (
384 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
385 - )
386 - weight_loader(param, loaded_weight)
387 - loaded_params.add(scale_name)
388 - continue
377 for param_name, weight_name, shard_id in stacked_params_mapping:
378 if weight_name not in name:
379 continue
vllm/model_executor/models/gemma2.py
-10
@@ -328,16 +328,6 @@ class Gemma2Model(nn.Module):
328 params_dict = dict(self.named_parameters())
329 loaded_params: set[str] = set()
330 for name, loaded_weight in weights:
331 - if self.quant_config is not None and (
332 - scale_name := self.quant_config.get_cache_scale(name)
333 - ):
334 - # Loading kv cache scales for compressed-tensors quantization
335 - param = params_dict[scale_name]
336 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
337 - loaded_weight = loaded_weight[0]
338 - weight_loader(param, loaded_weight)
339 - loaded_params.add(scale_name)
340 - continue
331 for param_name, shard_name, shard_id in stacked_params_mapping:
332 if shard_name not in name:
333 continue
vllm/model_executor/models/gemma3.py
-11
@@ -386,17 +386,6 @@ class Gemma3Model(nn.Module):
386 ):
387 loaded_weight -= 1
388
389 - if self.quant_config is not None and (
390 - scale_name := self.quant_config.get_cache_scale(name)
391 - ):
392 - # Loading kv cache scales for compressed-tensors quantization
393 - param = params_dict[scale_name]
394 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
395 - loaded_weight = loaded_weight[0]
396 - weight_loader(param, loaded_weight)
397 - loaded_params.add(scale_name)
398 - continue
399 -
389 # Check if this is a scale parameter that needs remapping first
390 if name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")):
391 # Try to remap the scale name first
vllm/model_executor/models/gemma3n.py
-10
@@ -1056,16 +1056,6 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):
1056 ):
1057 name = f"self_decoder.{name}"
1058
1059 - if self.quant_config is not None and (
1060 - scale_name := self.quant_config.get_cache_scale(name)
1061 - ):
1062 - # Loading kv cache scales for compressed-tensors quantization
1063 - param = params_dict[scale_name]
1064 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
1065 - loaded_weight = loaded_weight[0]
1066 - weight_loader(param, loaded_weight)
1067 - loaded_params.add(scale_name)
1068 - continue
1059 for param_name, shard_name, shard_id in stacked_params_mapping:
1060 if shard_name not in name:
1061 continue
vllm/model_executor/models/gemma4.py
-10
@@ -1409,16 +1409,6 @@ class Gemma4Model(nn.Module, EagleModelMixin):
1409 params_dict.update(dict(self.named_buffers()))
1410 loaded_params: set[str] = set()
1411 for name, loaded_weight in weights:
1412 - if self.quant_config is not None and (
1413 - scale_name := self.quant_config.get_cache_scale(name)
1414 - ):
1415 - param = params_dict[scale_name]
1416 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
1417 - loaded_weight = loaded_weight[0]
1418 - weight_loader(param, loaded_weight)
1419 - loaded_params.add(scale_name)
1420 - continue
1421 -
1412 if name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")):
1413 remapped_name = maybe_remap_kv_scale_name(name, params_dict)
1414 if remapped_name is not None and remapped_name in params_dict:
vllm/model_executor/models/glm4.py
-12
@@ -258,18 +258,6 @@ class Glm4Model(LlamaModel):
258 # Models trained using ColossalAI may include these tensors in
259 # the checkpoint. Skip them.
260 continue
261 - if self.quant_config is not None and (
262 - scale_name := self.quant_config.get_cache_scale(name)
263 - ):
264 - # Loading kv cache quantization scales
265 - param = params_dict[scale_name]
266 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
267 - loaded_weight = (
268 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
269 - )
270 - weight_loader(param, loaded_weight)
271 - loaded_params.add(scale_name)
272 - continue
261 if "scale" in name or "zero_point" in name:
262 # Remapping the name of FP8 kv-scale or zero point.
263 name = maybe_remap_kv_scale_name(name, params_dict)
vllm/model_executor/models/glm_ocr_mtp.py
+4 -13
@@ -166,6 +166,10 @@ class GlmOcrMTP(nn.Module, SupportsPP):
166 return self.model.compute_logits(hidden_states, spec_step_idx)
167
168 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
169 + if self.quant_config is not None and (
170 + cache_scale_mapper := self.quant_config.get_cache_scale_mapper()
171 + ):
172 + weights = cache_scale_mapper.apply(weights)
173 stacked_params_mapping = [
174 # (param_name, shard_name, shard_id)
175 (".qkv_proj", ".q_proj", "q"),
@@ -189,19 +193,6 @@ class GlmOcrMTP(nn.Module, SupportsPP):
193
194 name = self._rewrite_spec_layer_name(spec_layer, name)
195
192 - if self.quant_config is not None and (
193 - scale_name := self.quant_config.get_cache_scale(name)
194 - ):
195 - # Loading kv cache quantization scales
196 - param = params_dict[scale_name]
197 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
198 - loaded_weight = (
199 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
200 - )
201 - weight_loader(param, loaded_weight)
202 - loaded_params.add(scale_name)
203 - continue
204 -
196 if "scale" in name or "zero_point" in name:
197 # Remapping the name of FP8 kv-scale or zero point.
198 name = maybe_remap_kv_scale_name(name, params_dict)
vllm/model_executor/models/gpt_j.py
-13
@@ -254,19 +254,6 @@ class GPTJModel(nn.Module):
254 if "attn.bias" in name or "attn.masked_bias" in name:
255 continue
256
257 - if self.quant_config is not None and (
258 - scale_name := self.quant_config.get_cache_scale(name)
259 - ):
260 - # Loading kv cache quantization scales
261 - param = params_dict[scale_name]
262 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
263 - loaded_weight = (
264 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
265 - )
266 - weight_loader(param, loaded_weight)
267 - loaded_params.add(scale_name)
268 - continue
269 -
257 for param_name, weight_name, shard_id in stacked_params_mapping:
258 if weight_name not in name:
259 continue
vllm/model_executor/models/gpt_oss.py
-46
@@ -635,52 +635,6 @@ class GptOssModel(nn.Module, EagleModelMixin):
635
636 moe_quant_method = _get_moe_weight_dtype(layer_id=layer_id)
637
638 - def kv_cache_scale_loader(
639 - quant_config: QuantizationConfig,
640 - name: str,
641 - params_dict: dict[str, typing.Any],
642 - weight: torch.Tensor,
643 - default_weight_loader: Callable[..., None],
644 - loaded_params: set[str],
645 - ) -> tuple[bool, set[str]]:
646 - """
647 - Load KV cache output scales.
648 - Returns:
649 - Tuple of (bool, set):
650 - - bool: True if KV-cache scale was loaded into loaded_params
651 - - set: Updated set of loaded_params if True else the original set
652 - """
653 - # load explicit cached KV output scale from quant_config
654 - if quant_config is not None and (
655 - scale_name := quant_config.get_cache_scale(name)
656 - ):
657 - param = params_dict[scale_name]
658 - weight_loader = getattr(
659 - param, "weight_loader", default_weight_loader
660 - )
661 - if weight.numel() != 1:
662 - raise ValueError(
663 - f"KV cache scale '{scale_name}' is expected to be a "
664 - f"scalar, but got a tensor of shape {weight.shape}."
665 - )
666 - # Ensure weight is a scalar before passing to loader.
667 - weight_loader(param, weight.flatten()[0])
668 - loaded_params.add(scale_name)
669 - return True, loaded_params
670 -
671 - return False, loaded_params
672 -
673 - load_kv_cache_scale_completed, loaded_params = kv_cache_scale_loader(
674 - self.quant_config,
675 - name,
676 - params_dict,
677 - loaded_weight,
678 - default_weight_loader,
679 - loaded_params,
680 - )
681 - if load_kv_cache_scale_completed:
682 - continue
683 -
638 if (
639 all(key in name for key in ["input_scale", "mlp.experts"])
640 and expert_id is not None
vllm/model_executor/models/granite.py
-12
@@ -334,18 +334,6 @@ class GraniteModel(nn.Module):
334 params_dict = dict(self.named_parameters())
335 loaded_params: set[str] = set()
336 for name, loaded_weight in weights:
337 - if self.quant_config is not None and (
338 - scale_name := self.quant_config.get_cache_scale(name)
339 - ):
340 - # Loading kv cache quantization scales
341 - param = params_dict[scale_name]
342 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
343 - loaded_weight = (
344 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
345 - )
346 - weight_loader(param, loaded_weight)
347 - loaded_params.add(scale_name)
348 - continue
337 for param_name, weight_name, shard_id in stacked_params_mapping:
338 if weight_name not in name:
339 continue
vllm/model_executor/models/granitemoe.py
-13
@@ -365,19 +365,6 @@ class GraniteMoeModel(nn.Module):
365 params_dict = dict(self.named_parameters())
366 loaded_params: set[str] = set()
367 for name, loaded_weight in weights:
368 - if self.quant_config is not None and (
369 - scale_name := self.quant_config.get_cache_scale(name)
370 - ):
371 - # Loading kv cache quantization scales
372 - param = params_dict[scale_name]
373 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
374 - loaded_weight = (
375 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
376 - )
377 - weight_loader(param, loaded_weight)
378 - loaded_params.add(scale_name)
379 - continue
380 -
368 for param_name, weight_name, shard_id in stacked_params_mapping:
369 if weight_name not in name:
370 continue
vllm/model_executor/models/granitemoehybrid.py
-12
@@ -495,18 +495,6 @@ class GraniteMoeHybridModel(nn.Module):
495 if "A_log" in n:
496 n = n.replace("A_log", "A")
497
498 - if self.quant_config is not None and (
499 - scale_name := self.quant_config.get_cache_scale(n)
500 - ):
501 - # Loading kv cache quantization scales
502 - loaded_weight = p
503 - loaded_weight = (
504 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
505 - )
506 - _load(scale_name, loaded_weight)
507 - loaded_params.add(scale_name)
508 - continue
509 -
498 if _load_quant_expert(n, p):
499 continue
500
vllm/model_executor/models/grok1.py
-12
@@ -548,18 +548,6 @@ class Grok1Model(nn.Module):
548 for old_pattern, new_pattern in self.weight_name_remapping.items():
549 if old_pattern in name:
550 name = name.replace(old_pattern, new_pattern)
551 - if self.quant_config is not None and (
552 - scale_name := self.quant_config.get_cache_scale(name)
553 - ):
554 - # Loading kv cache quantization scales
555 - param = params_dict[scale_name]
556 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
557 - loaded_weight = (
558 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
559 - )
560 - weight_loader(param, loaded_weight)
561 - loaded_params.add(scale_name)
562 - continue
551
552 for param_name, weight_name, shard_id in stacked_params_mapping:
553 if weight_name not in name:
vllm/model_executor/models/hunyuan_v1.py
-9
@@ -771,15 +771,6 @@ class HunYuanModel(nn.Module, EagleModelMixin):
771 # processed with quantization, LoRA, fine-tuning, etc.
772 if self.config.tie_word_embeddings and "lm_head.weight" in name:
773 continue
774 - if self.quant_config is not None and (
775 - scale_name := self.quant_config.get_cache_scale(name)
776 - ):
777 - # Loading kv cache scales for compressed-tensors quantization
778 - param = params_dict[scale_name]
779 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
780 - loaded_weight = loaded_weight[0]
781 - weight_loader(param, loaded_weight)
782 - continue
774
775 is_found = False
776 for param_name, weight_name, shard_id in stacked_params_mapping:
vllm/model_executor/models/hy_v3.py
-11
@@ -531,17 +531,6 @@ class HYV3Model(nn.Module):
531 for name, loaded_weight in weights:
532 if self.config.tie_word_embeddings and "lm_head.weight" in name:
533 continue
534 - if self.quant_config is not None and (
535 - scale_name := self.quant_config.get_cache_scale(name)
536 - ):
537 - param = params_dict[scale_name]
538 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
539 - loaded_weight = (
540 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
541 - )
542 - weight_loader(param, loaded_weight)
543 - loaded_params.add(scale_name)
544 - continue
534 if "scale" in name:
535 # Remapping the name of FP8 kv-scale.
536 name = maybe_remap_kv_scale_name(name, params_dict)
vllm/model_executor/models/hy_v3_mtp.py
+4 -8
@@ -264,6 +264,10 @@ class HYV3MTP(nn.Module):
264 return torch.concat((q, k, v))
265
266 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
267 + if self.quant_config is not None and (
268 + cache_scale_mapper := self.quant_config.get_cache_scale_mapper()
269 + ):
270 + weights = cache_scale_mapper.apply(weights)
271 cla_factor = _get_cla_factor(self.config)
272 stacked_params_mapping = [
273 # (param_name, shard_name, shard_id)
@@ -336,14 +340,6 @@ class HYV3MTP(nn.Module):
340 continue
341 if self.config.tie_word_embeddings and "lm_head.weight" in name:
342 continue
339 - if self.quant_config is not None and (
340 - scale_name := self.quant_config.get_cache_scale(name)
341 - ):
342 - param = params_dict[scale_name]
343 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
344 - loaded_weight = loaded_weight[0]
345 - weight_loader(param, loaded_weight)
346 - continue
343 spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
344 if spec_layer is None:
345 continue
vllm/model_executor/models/hyperclovax.py
-12
@@ -395,18 +395,6 @@ class HyperCLOVAXModel(nn.Module):
395 # Models trained using ColossalAI may include these tensors in
396 # the checkpoint. Skip them.
397 continue
398 - if self.quant_config is not None and (
399 - scale_name := self.quant_config.get_cache_scale(name)
400 - ):
401 - # Loading kv cache quantization scales
402 - param = params_dict[scale_name]
403 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
404 - loaded_weight = (
405 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
406 - )
407 - weight_loader(param, loaded_weight)
408 - loaded_params.add(scale_name)
409 - continue
398 if "scale" in name or "zero_point" in name:
399 # Remapping the name of FP8 kv-scale or zero point.
400 remapped_name = maybe_remap_kv_scale_name(name, params_dict)
vllm/model_executor/models/iquest_loopcoder.py
-12
@@ -476,18 +476,6 @@ class IQuestLoopCoderModel(nn.Module):
476 for name, loaded_weight in weights:
477 if "rotary_emb.inv_freq" in name:
478 continue
479 - if self.quant_config is not None and (
480 - scale_name := self.quant_config.get_cache_scale(name)
481 - ):
482 - # Loading kv cache quantization scales
483 - param = params_dict[scale_name]
484 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
485 - loaded_weight = (
486 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
487 - )
488 - weight_loader(param, loaded_weight)
489 - loaded_params.add(scale_name)
490 - continue
479 for param_name, weight_name, shard_id in stacked_params_mapping:
480 if "gate_projections" in name:
481 continue
vllm/model_executor/models/jais2.py
-10
@@ -386,16 +386,6 @@ class Jais2Model(nn.Module):
386 # Models trained using ColossalAI may include these tensors in
387 # the checkpoint. Skip them.
388 continue
389 - if self.quant_config is not None and (
390 - scale_name := self.quant_config.get_cache_scale(name)
391 - ):
392 - # Loading kv cache scales for compressed-tensors quantization
393 - param = params_dict[scale_name]
394 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
395 - loaded_weight = loaded_weight[0]
396 - weight_loader(param, loaded_weight)
397 - loaded_params.add(scale_name)
398 - continue
389 if "scale" in name:
390 name = maybe_remap_kv_scale_name(name, params_dict)
391 if name is None:
vllm/model_executor/models/keye.py
-15
@@ -790,21 +790,6 @@ class KeyeSiglipVisionModel(nn.Module):
790 continue
791 if "head.mlp" in name or "head.probe" in name:
792 continue
793 - if self.quant_config is not None and (
794 - scale_name := self.quant_config.get_cache_scale(name)
795 - ):
796 - param = params_dict[scale_name]
797 - weight_loader = getattr(
798 - param,
799 - "weight_loader",
800 - default_weight_loader,
801 - )
802 - loaded_weight = (
803 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
804 - )
805 - weight_loader(param, loaded_weight)
806 - loaded_params.add(scale_name)
807 - continue
793 for (
794 param_name,
795 weight_name,
vllm/model_executor/models/laguna.py
-14
@@ -724,20 +724,6 @@ class LagunaModel(nn.Module, EagleModelMixin):
724 loaded_params.add(name)
725 continue
726
727 - # Handle KV cache quantization scales
728 - if self.quant_config is not None and (
729 - scale_name := self.quant_config.get_cache_scale(name)
730 - ):
731 - param = params_dict[scale_name]
732 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
733 - assert loaded_weight.numel() == 1, (
734 - f"KV scale numel {loaded_weight.numel()} != 1"
735 - )
736 - loaded_weight = loaded_weight.squeeze()
737 - weight_loader(param, loaded_weight)
738 - loaded_params.add(scale_name)
739 - continue
740 -
727 # Handle stacked params (QKV, gate_up for
728 # non-expert layers and shared_expert)
729 for param_name, weight_name, shard_id in stacked_params_mapping:
vllm/model_executor/models/llama.py
-12
@@ -451,18 +451,6 @@ class LlamaModel(nn.Module, EagleModelMixin):
451 # Models trained using ColossalAI may include these tensors in
452 # the checkpoint. Skip them.
453 continue
454 - if self.quant_config is not None and (
455 - scale_name := self.quant_config.get_cache_scale(name)
456 - ):
457 - # Loading kv cache quantization scales
458 - param = params_dict[scale_name]
459 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
460 - loaded_weight = (
461 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
462 - )
463 - weight_loader(param, loaded_weight)
464 - loaded_params.add(scale_name)
465 - continue
454 if "scale" in name or "zero_point" in name:
455 # Remapping the name of FP8 kv-scale or zero point.
456 name = maybe_remap_kv_scale_name(name, params_dict)
vllm/model_executor/models/llama4.py
+3 -18
@@ -588,21 +588,6 @@ class Llama4Model(LlamaModel):
588 fused_experts_params = True
589 expert_params_mapping = expert_params_mapping_fused
590
591 - # If kv cache quantization scales exist and the weight name
592 - # corresponds to one of the kv cache quantization scales, load
593 - # them.
594 - if self.quant_config is not None and (
595 - scale_name := self.quant_config.get_cache_scale(name)
596 - ):
597 - param = params_dict[scale_name]
598 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
599 - loaded_weight = (
600 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
601 - )
602 - weight_loader(param, loaded_weight)
603 - loaded_params.add(scale_name)
604 - continue
605 -
591 # Iterate over stacked_params_mapping to check if the current weight
592 # is one of the stacked parameters. If so, load the weight with the
593 # corresponding shard id. Note that MoE weights are handled
@@ -625,9 +610,9 @@ class Llama4Model(LlamaModel):
610 if is_pp_missing_parameter(name, self):
611 continue
612
628 - # Remap kv cache scale names for ModelOpt checkpoints.
629 - # TODO: ModelOpt should implement get_cache_scale() such that
630 - # kv cache scale name remapping can be done there.
613 + # Remap kv cache scale names for any checkpoint format the
614 + # quant config's `get_cache_scale_mapper` does not cover
615 + # (idempotent for names already renamed by the mapper).
616 if name.endswith("scale"):
617 name = maybe_remap_kv_scale_name(name, params_dict)
618 if name is None:
vllm/model_executor/models/llama_eagle.py
-13
@@ -127,19 +127,6 @@ class LlamaModel(nn.Module):
127 params_dict = dict(self.named_parameters())
128 loaded_params: set[str] = set()
129 for name, loaded_weight in weights:
130 - # Handle kv cache quantization scales
131 - if self.quant_config is not None and (
132 - scale_name := self.quant_config.get_cache_scale(name)
133 - ):
134 - # Loading kv cache quantization scales
135 - param = params_dict[scale_name]
136 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
137 - loaded_weight = (
138 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
139 - )
140 - weight_loader(param, loaded_weight)
141 - loaded_params.add(scale_name)
142 - continue
130 # Remapping the name FP8 kv-scale or zero point.
131 if "scale" in name or "zero_point" in name:
132 name = maybe_remap_kv_scale_name(name, params_dict)
vllm/model_executor/models/llama_eagle3.py
-13
@@ -267,19 +267,6 @@ class LlamaModel(nn.Module):
267 for name, loaded_weight in weights:
268 if "midlayer." in name:
269 name = name.replace("midlayer.", "layers.0.")
270 - # Handle kv cache quantization scales
271 - if self.quant_config is not None and (
272 - scale_name := self.quant_config.get_cache_scale(name)
273 - ):
274 - # Loading kv cache quantization scales
275 - param = params_dict[scale_name]
276 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
277 - loaded_weight = (
278 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
279 - )
280 - weight_loader(param, loaded_weight)
281 - loaded_params.add(scale_name)
282 - continue
270 # Remapping the name FP8 kv-scale or zero point.
271 if "scale" in name or "zero_point" in name:
272 name = maybe_remap_kv_scale_name(name, params_dict)
vllm/model_executor/models/mimo.py
-12
@@ -104,18 +104,6 @@ class MiMoModel(Qwen2Model):
104 continue
105 if "rotary_emb.inv_freq" in name:
106 continue
107 - if self.quant_config is not None and (
108 - scale_name := self.quant_config.get_cache_scale(name)
109 - ):
110 - # Loading kv cache quantization scales
111 - param = params_dict[scale_name]
112 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
113 - loaded_weight = (
114 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
115 - )
116 - weight_loader(param, loaded_weight)
117 - loaded_params.add(scale_name)
118 - continue
107 for param_name, weight_name, shard_id in stacked_params_mapping:
108 if weight_name not in name:
109 continue
vllm/model_executor/models/mimo_v2.py
-16
@@ -555,22 +555,6 @@ class MiMoV2Model(nn.Module):
555 if "mtp" in name:
556 continue
557
558 - if self.quant_config is not None:
559 - cache_scale_name = self.quant_config.get_cache_scale(name)
560 - if cache_scale_name is not None and cache_scale_name in params_dict:
561 - param = params_dict[cache_scale_name]
562 - weight_loader = getattr(
563 - param, "weight_loader", default_weight_loader
564 - )
565 -
566 - kv_scale = loaded_weight
567 - if kv_scale.dim() > 0 and kv_scale.numel() > 1:
568 - kv_scale = kv_scale.view(-1)[0]
569 -
570 - weight_loader(param, kv_scale)
571 - loaded_params.add(cache_scale_name)
572 - continue
573 -
558 expert_matched = False
559 for param_name, weight_name, expert_id, shard_id in expert_params_mapping:
560 if weight_name not in name:
vllm/model_executor/models/mixtral.py
-13
@@ -388,19 +388,6 @@ class MixtralModel(nn.Module):
388 loaded_params: set[str] = set()
389 expert_params_mapping = self.get_expert_mapping()
390 for name, loaded_weight in weights:
391 - if self.quant_config is not None and (
392 - scale_name := self.quant_config.get_cache_scale(name)
393 - ):
394 - # Loading kv cache quantization scales
395 - param = params_dict[scale_name]
396 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
397 - loaded_weight = (
398 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
399 - )
400 - weight_loader(param, loaded_weight)
401 - loaded_params.add(scale_name)
402 - continue
403 -
391 for param_name, weight_name, shard_id in stacked_params_mapping:
392 if weight_name not in name:
393 continue
vllm/model_executor/models/nemotron.py
-12
@@ -375,18 +375,6 @@ class NemotronModel(nn.Module):
375 params_dict = dict(self.named_parameters())
376 loaded_params: set[str] = set()
377 for name, loaded_weight in weights:
378 - if self.quant_config is not None and (
379 - scale_name := self.quant_config.get_cache_scale(name)
380 - ):
381 - # Loading kv cache quantization scales
382 - param = params_dict[scale_name]
383 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
384 - loaded_weight = (
385 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
386 - )
387 - weight_loader(param, loaded_weight)
388 - loaded_params.add(scale_name)
389 - continue
378 for param_name, weight_name, shard_id in stacked_params_mapping:
379 if weight_name not in name:
380 continue
vllm/model_executor/models/nemotron_nas.py
-12
@@ -334,18 +334,6 @@ class DeciModel(nn.Module):
334 # Models trained using ColossalAI may include these tensors in
335 # the checkpoint. Skip them.
336 continue
337 - if self.quant_config is not None and (
338 - scale_name := self.quant_config.get_cache_scale(name)
339 - ):
340 - # Loading kv cache quantization scales
341 - param = params_dict[scale_name]
342 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
343 - loaded_weight = (
344 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
345 - )
346 - weight_loader(param, loaded_weight)
347 - loaded_params.add(scale_name)
348 - continue
337 if "scale" in name or "zero_point" in name:
338 # Remapping the name of FP8 kv-scale.
339 name = maybe_remap_kv_scale_name(name, params_dict)
vllm/model_executor/models/ouro.py
-12
@@ -390,18 +390,6 @@ class OuroModel(nn.Module):
390 for name, loaded_weight in weights:
391 if "rotary_emb.inv_freq" in name:
392 continue
393 - if self.quant_config is not None and (
394 - scale_name := self.quant_config.get_cache_scale(name)
395 - ):
396 - # Loading kv cache quantization scales
397 - param = params_dict[scale_name]
398 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
399 - loaded_weight = (
400 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
401 - )
402 - weight_loader(param, loaded_weight)
403 - loaded_params.add(scale_name)
404 - continue
393 for param_name, weight_name, shard_id in stacked_params_mapping:
394 if weight_name not in name:
395 continue
vllm/model_executor/models/paddleocr_vl.py
-15
@@ -944,21 +944,6 @@ class SiglipVisionModel(nn.Module):
944 continue
945 if "packing_position_embedding" in name:
946 continue
947 - if self.quant_config is not None and (
948 - scale_name := self.quant_config.get_cache_scale(name)
949 - ):
950 - param = params_dict[scale_name]
951 - weight_loader = getattr(
952 - param,
953 - "weight_loader",
954 - default_weight_loader,
955 - )
956 - loaded_weight = (
957 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
958 - )
959 - weight_loader(param, loaded_weight)
960 - loaded_params.add(scale_name)
961 - continue
947 for (
948 param_name,
949 weight_name,
vllm/model_executor/models/phimoe.py
-13
@@ -537,19 +537,6 @@ class PhiMoEModel(nn.Module):
537 loaded_params: set[str] = set()
538 expert_params_mapping = self.get_expert_mapping()
539 for name, loaded_weight in weights:
540 - if self.quant_config is not None and (
541 - scale_name := self.quant_config.get_cache_scale(name)
542 - ):
543 - # Loading kv cache quantization scales
544 - param = params_dict[scale_name]
545 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
546 - loaded_weight = (
547 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
548 - )
549 - weight_loader(param, loaded_weight)
550 - loaded_params.add(scale_name)
551 - continue
552 -
540 for param_name, weight_name, shard_id in stacked_params_mapping:
541 if weight_name not in name:
542 continue
vllm/model_executor/models/qwen2.py
-12
@@ -439,18 +439,6 @@ class Qwen2Model(nn.Module, EagleModelMixin):
439 for name, loaded_weight in weights:
440 if "rotary_emb.inv_freq" in name:
441 continue
442 - if self.quant_config is not None and (
443 - scale_name := self.quant_config.get_cache_scale(name)
444 - ):
445 - # Loading kv cache quantization scales
446 - param = params_dict[scale_name]
447 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
448 - loaded_weight = (
449 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
450 - )
451 - weight_loader(param, loaded_weight)
452 - loaded_params.add(scale_name)
453 - continue
442 for param_name, weight_name, shard_id in stacked_params_mapping:
443 if weight_name not in name:
444 continue
vllm/model_executor/models/qwen3_dflash.py
-11
@@ -469,17 +469,6 @@ class DFlashQwen3Model(nn.Module):
469 for name, loaded_weight in weights:
470 if "midlayer." in name:
471 name = name.replace("midlayer.", "layers.0.")
472 - if self.quant_config is not None and (
473 - scale_name := self.quant_config.get_cache_scale(name)
474 - ):
475 - param = params_dict[scale_name]
476 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
477 - loaded_weight = (
478 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
479 - )
480 - weight_loader(param, loaded_weight)
481 - loaded_params.add(scale_name)
482 - continue
472 if "scale" in name:
473 name = maybe_remap_kv_scale_name(name, params_dict)
474 if name is None:
vllm/model_executor/models/qwen3_moe.py
-13
@@ -552,19 +552,6 @@ class Qwen3MoeModel(nn.Module, EagleModelMixin):
552 loaded_params: set[str] = set()
553 expert_params_mapping = self.get_expert_mapping()
554 for name, loaded_weight in weights:
555 - if self.quant_config is not None and (
556 - scale_name := self.quant_config.get_cache_scale(name)
557 - ):
558 - # Loading kv cache quantization scales
559 - param = params_dict[scale_name]
560 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
561 - assert loaded_weight.numel() == 1, (
562 - f"KV scale numel {loaded_weight.numel()} != 1"
563 - )
564 - loaded_weight = loaded_weight.squeeze()
565 - weight_loader(param, loaded_weight)
566 - loaded_params.add(scale_name)
567 - continue
555 if "scale" in name or "zero_point" in name:
556 name = maybe_remap_kv_scale_name(name, params_dict)
557 if name is None:
vllm/model_executor/models/rnj1.py
-10
@@ -350,16 +350,6 @@ class Rnj1Model(nn.Module):
350 ):
351 loaded_weight -= 1
352
353 - if self.quant_config is not None and (
354 - scale_name := self.quant_config.get_cache_scale(name)
355 - ):
356 - param = params_dict[scale_name]
357 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
358 - loaded_weight = loaded_weight[0]
359 - weight_loader(param, loaded_weight)
360 - loaded_params.add(scale_name)
361 - continue
362 -
353 if name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")):
354 remapped_name = maybe_remap_kv_scale_name(name, params_dict)
355 if remapped_name is not None and remapped_name in params_dict:
vllm/model_executor/models/seed_oss.py
-12
@@ -376,18 +376,6 @@ class SeedOssModel(nn.Module):
376 for name, loaded_weight in weights:
377 if "rotary_emb.inv_freq" in name:
378 continue
379 - if self.quant_config is not None and (
380 - scale_name := self.quant_config.get_cache_scale(name)
381 - ):
382 - # Loading kv cache quantization scales
383 - param = params_dict[scale_name]
384 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
385 - loaded_weight = (
386 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
387 - )
388 - weight_loader(param, loaded_weight)
389 - loaded_params.add(scale_name)
390 - continue
379 for param_name, weight_name, shard_id in stacked_params_mapping:
380 if weight_name not in name:
381 continue
vllm/model_executor/models/solar.py
-12
@@ -360,18 +360,6 @@ class SolarModel(nn.Module):
360 params_dict = dict(self.named_parameters())
361 loaded_params: set[str] = set()
362 for name, loaded_weight in weights:
363 - if self.quant_config is not None and (
364 - scale_name := self.quant_config.get_cache_scale(name)
365 - ):
366 - # Loading kv cache quantization scales
367 - param = params_dict[scale_name]
368 - weight_loader = getattr(param, "weight_loader", default_weight_loader)
369 - loaded_weight = (
370 - loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
371 - )
372 - weight_loader(param, loaded_weight)
373 - loaded_params.add(scale_name)
374 - continue
363 for param_name, weight_name, shard_id in stacked_params_mapping:
364 if weight_name not in name:
365 continue
vllm/model_executor/models/utils.py
+15
@@ -52,6 +52,7 @@ class WeightsMapper:
52 def __or__(self, other: "WeightsMapper") -> "WeightsMapper":
53 """Combine two `WeightsMapper`s by merging their mappings."""
54 return WeightsMapper(
55 + orig_to_new_regex={**self.orig_to_new_regex, **other.orig_to_new_regex},
56 orig_to_new_substr={**self.orig_to_new_substr, **other.orig_to_new_substr},
57 orig_to_new_prefix={**self.orig_to_new_prefix, **other.orig_to_new_prefix},
58 orig_to_new_suffix={**self.orig_to_new_suffix, **other.orig_to_new_suffix},
@@ -343,6 +344,20 @@ class AutoWeightsLoader:
344 *,
345 mapper: WeightsMapper | None = None,
346 ) -> set[str]:
347 + # Many models store quant_config in the base model instead of the causal model.
348 + # We look at the causal model's direct children for this reason.
349 + modules = (self.module, *self.module.children())
350 + iterator = (m.quant_config for m in modules if hasattr(m, "quant_config"))
351 + quant_config = next(iterator, None)
352 + cache_scale_mapper = (
353 + quant_config.get_cache_scale_mapper() if quant_config is not None else None
354 + )
355 + if cache_scale_mapper is not None:
356 + mapper = (
357 + mapper | cache_scale_mapper
358 + if mapper is not None
359 + else cache_scale_mapper
360 + )
361 if mapper is not None:
362 weights = mapper.apply(weights)
363 # filter out weights with first-prefix/substr to skip in name