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 7, 06:30 UTC
[Bugfix][Model] Fix FP8 k_scale/v_scale not loaded for Qwen3-MoE (#35656)
Signed-off-by: raghavan <oneraghavan@gmail.com>
Raghavan committed
Mar 4, 2026 at 18:45 UTC
c8c3935b701380aee7494f7f5021e8e06de2d9b7
3 files changed
+129
-36
tests/model_executor/test_weight_utils.py
+116
@@ -11,6 +11,7 @@ from huggingface_hub.utils import LocalEntryNotFoundError
11
from vllm.model_executor.model_loader.weight_utils import (
12
download_weights_from_hf,
13
enable_hf_transfer,
14
+ maybe_remap_kv_scale_name,
15
)
16
17
@@ -61,6 +62,121 @@ def test_download_weights_from_hf():
62
)
63
64
65
+class TestMaybeRemapKvScaleName:
66
+ """Tests for maybe_remap_kv_scale_name covering all checkpoint formats."""
67
+
68
+ PARAMS_DICT = {
69
+ "model.layers.0.self_attn.attn.k_scale": None,
70
+ "model.layers.0.self_attn.attn.v_scale": None,
71
+ "model.layers.0.self_attn.attn.q_scale": None,
72
+ "model.layers.0.self_attn.qkv_proj.weight": None,
73
+ }
74
+
75
+ def test_qkv_proj_k_scale(self):
76
+ """Qwen3-MoE / llm-compressor format: qkv_proj.k_scale -> attn.k_scale
77
+ Regression test for https://github.com/vllm-project/vllm/issues/25047"""
78
+ result = maybe_remap_kv_scale_name(
79
+ "model.layers.0.self_attn.qkv_proj.k_scale", self.PARAMS_DICT
80
+ )
81
+ assert result == "model.layers.0.self_attn.attn.k_scale"
82
+
83
+ def test_qkv_proj_v_scale(self):
84
+ """Qwen3-MoE / llm-compressor format: qkv_proj.v_scale -> attn.v_scale
85
+ Regression test for https://github.com/vllm-project/vllm/issues/25047"""
86
+ result = maybe_remap_kv_scale_name(
87
+ "model.layers.0.self_attn.qkv_proj.v_scale", self.PARAMS_DICT
88
+ )
89
+ assert result == "model.layers.0.self_attn.attn.v_scale"
90
+
91
+ def test_modelopt_k_proj_k_scale(self):
92
+ """ModelOpt format: k_proj.k_scale -> attn.k_scale"""
93
+ result = maybe_remap_kv_scale_name(
94
+ "model.layers.0.self_attn.k_proj.k_scale", self.PARAMS_DICT
95
+ )
96
+ assert result == "model.layers.0.self_attn.attn.k_scale"
97
+
98
+ def test_modelopt_v_proj_v_scale(self):
99
+ """ModelOpt format: v_proj.v_scale -> attn.v_scale"""
100
+ result = maybe_remap_kv_scale_name(
101
+ "model.layers.0.self_attn.v_proj.v_scale", self.PARAMS_DICT
102
+ )
103
+ assert result == "model.layers.0.self_attn.attn.v_scale"
104
+
105
+ def test_deprecated_kv_scale(self):
106
+ """Old format: kv_scale -> attn.k_scale (deprecated)"""
107
+ result = maybe_remap_kv_scale_name(
108
+ "model.layers.0.self_attn.kv_scale", self.PARAMS_DICT
109
+ )
110
+ assert result == "model.layers.0.self_attn.attn.k_scale"
111
+
112
+ def test_default_bare_k_scale(self):
113
+ """Default format: .k_scale -> .attn.k_scale"""
114
+ result = maybe_remap_kv_scale_name(
115
+ "model.layers.0.self_attn.k_scale", self.PARAMS_DICT
116
+ )
117
+ assert result == "model.layers.0.self_attn.attn.k_scale"
118
+
119
+ def test_non_scale_name_unchanged(self):
120
+ """Non-scale names should be returned unchanged."""
121
+ name = "model.layers.0.self_attn.qkv_proj.weight"
122
+ result = maybe_remap_kv_scale_name(name, self.PARAMS_DICT)
123
+ assert result == name
124
+
125
+ def test_nvfp4_modelopt_k_proj_k_scale(self):
126
+ """ModelOpt NVFP4 format (e.g. nvidia/Qwen3-30B-A3B-NVFP4):
127
+ k_proj.k_scale -> attn.k_scale.
128
+ Validates that NVFP4 checkpoints are not broken by this change."""
129
+ result = maybe_remap_kv_scale_name(
130
+ "model.layers.0.self_attn.k_proj.k_scale", self.PARAMS_DICT
131
+ )
132
+ assert result == "model.layers.0.self_attn.attn.k_scale"
133
+
134
+ def test_nvfp4_modelopt_v_proj_v_scale(self):
135
+ """ModelOpt NVFP4 format (e.g. nvidia/Qwen3-30B-A3B-NVFP4):
136
+ v_proj.v_scale -> attn.v_scale.
137
+ Validates that NVFP4 checkpoints are not broken by this change."""
138
+ result = maybe_remap_kv_scale_name(
139
+ "model.layers.0.self_attn.v_proj.v_scale", self.PARAMS_DICT
140
+ )
141
+ assert result == "model.layers.0.self_attn.attn.v_scale"
142
+
143
+ def test_qwen3_vl_moe_qkv_proj_k_scale(self):
144
+ """Qwen3-VL-MoE uses the same fused qkv_proj naming as Qwen3-MoE.
145
+ Regression test for qwen3_vl_moe.py fix (same bug as #25047)."""
146
+ result = maybe_remap_kv_scale_name(
147
+ "model.layers.0.self_attn.qkv_proj.k_scale", self.PARAMS_DICT
148
+ )
149
+ assert result == "model.layers.0.self_attn.attn.k_scale"
150
+
151
+ def test_qwen3_vl_moe_qkv_proj_v_scale(self):
152
+ """Qwen3-VL-MoE uses the same fused qkv_proj naming as Qwen3-MoE.
153
+ Regression test for qwen3_vl_moe.py fix (same bug as #25047)."""
154
+ result = maybe_remap_kv_scale_name(
155
+ "model.layers.0.self_attn.qkv_proj.v_scale", self.PARAMS_DICT
156
+ )
157
+ assert result == "model.layers.0.self_attn.attn.v_scale"
158
+
159
+ def test_nvfp4_weight_scale_not_remapped(self):
160
+ """NVFP4 weight_scale should not be touched by remap (not a kv scale)."""
161
+ name = "model.layers.0.self_attn.k_proj.weight_scale"
162
+ result = maybe_remap_kv_scale_name(name, self.PARAMS_DICT)
163
+ assert result == name
164
+
165
+ def test_nvfp4_input_scale_not_remapped(self):
166
+ """NVFP4 input_scale should not be touched by remap (not a kv scale)."""
167
+ name = "model.layers.0.self_attn.k_proj.input_scale"
168
+ result = maybe_remap_kv_scale_name(name, self.PARAMS_DICT)
169
+ assert result == name
170
+
171
+ def test_missing_target_returns_none(self):
172
+ """If remapped name not in params_dict, return None."""
173
+ empty_params: dict[str, None] = {}
174
+ result = maybe_remap_kv_scale_name(
175
+ "model.layers.0.self_attn.qkv_proj.k_scale", empty_params
176
+ )
177
+ assert result is None
178
+
179
+
180
if __name__ == "__main__":
181
test_hf_transfer_auto_activation()
182
test_download_weights_from_hf()
vllm/model_executor/models/qwen3_moe.py
+6
-18
@@ -535,10 +535,6 @@ class Qwen3MoeModel(nn.Module):
535
ignore_suffixes = (
536
".bias",
537
"_bias",
538
- ".k_scale",
539
- "_k_scale",
540
- ".v_scale",
541
- "_v_scale",
538
".weight_scale",
539
"_weight_scale",
540
".input_scale",
@@ -562,6 +558,10 @@ class Qwen3MoeModel(nn.Module):
558
weight_loader(param, loaded_weight)
559
loaded_params.add(scale_name)
560
continue
561
+ if "scale" in name or "zero_point" in name:
562
+ name = maybe_remap_kv_scale_name(name, params_dict)
563
+ if name is None:
564
+ continue
565
for param_name, weight_name, shard_id in stacked_params_mapping:
566
# Skip non-stacked layers and experts (experts handled below).
567
if weight_name not in name:
@@ -654,20 +654,8 @@ class Qwen3MoeModel(nn.Module):
654
# Skip layers on other devices.
655
if is_pp_missing_parameter(name, self):
656
continue
657
- # Remapping the name of FP8 kv-scale.
658
- if name.endswith("kv_scale"):
659
- remapped_kv_scale_name = name.replace(
660
- ".kv_scale", ".attn.kv_scale"
661
- )
662
- if remapped_kv_scale_name not in params_dict:
663
- logger.warning_once(
664
- "Found kv scale in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). kv-scale is not loaded.", # noqa: E501
665
- name,
666
- remapped_kv_scale_name,
667
- )
668
- continue
669
- else:
670
- name = remapped_kv_scale_name
657
+ if name not in params_dict:
658
+ continue
659
param = params_dict[name]
660
weight_loader = getattr(
661
param, "weight_loader", default_weight_loader
vllm/model_executor/models/qwen3_vl_moe.py
+7
-18
@@ -172,10 +172,6 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
172
ignore_suffixes = (
173
".bias",
174
"_bias",
175
- ".k_scale",
176
- "_k_scale",
177
- ".v_scale",
178
- "_v_scale",
175
".weight_scale",
176
"_weight_scale",
177
".input_scale",
@@ -191,6 +187,11 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
187
]
188
num_experts = self.config.num_experts
189
for name, loaded_weight in weights:
190
+ if "scale" in name or "zero_point" in name:
191
+ name = maybe_remap_kv_scale_name(name, params_dict)
192
+ if name is None:
193
+ continue
194
+
195
for param_name, weight_name, shard_id in stacked_params_mapping:
196
if "experts.gate_up_proj" in name or "experts.down_proj" in name:
197
is_fused_expert = True
@@ -305,20 +306,8 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
306
# Skip layers on other devices.
307
if is_pp_missing_parameter(name, self):
308
continue
308
- # Remapping the name of FP8 kv-scale.
309
- if name.endswith("kv_scale"):
310
- remapped_kv_scale_name = name.replace(
311
- ".kv_scale", ".attn.kv_scale"
312
- )
313
- if remapped_kv_scale_name not in params_dict:
314
- logger.warning_once(
315
- "Found kv scale in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). kv-scale is not loaded.", # noqa: E501
316
- name,
317
- remapped_kv_scale_name,
318
- )
319
- continue
320
- else:
321
- name = remapped_kv_scale_name
309
+ if name not in params_dict:
310
+ continue
311
param = params_dict[name]
312
weight_loader = getattr(
313
param, "weight_loader", default_weight_loader