[BugFix] Fix ModelOpt mixed-precision quantization for sparse `quantized_layers` configs. (#47318)
Signed-off-by: Daniel Afrimi <dafrimi@nvidia.com> Signed-off-by: <dafrimi@nvidia.com>
danielafrimi committed
Jul 7, 2026 at 14:45 UTC
0a2965b1b35f34fde29c9a653696bac0ef1e499d
2 files changed
+66
-10
tests/quantization/test_modelopt.py
+43
@@ -133,6 +133,49 @@ def test_modelopt_mixed_precision_quantizes_parallel_lm_head():
133
assert isinstance(method, ModelOptNvFp4LinearMethod)
134
135
136
+def test_modelopt_mixed_precision_infers_fused_gate_up_projection():
137
+ from vllm.model_executor.layers.linear import LinearBase
138
+
139
+ config = _mixed_precision_config(
140
+ {
141
+ "model.layers.0.mlp.gate_proj": {"quant_algo": "NVFP4"},
142
+ "model.layers.0.mlp.up_proj": {"quant_algo": "NVFP4"},
143
+ }
144
+ )
145
+
146
+ fake_layer = MagicMock(spec=LinearBase)
147
+ with patch(
148
+ "vllm.model_executor.layers.quantization.modelopt.init_nvfp4_linear_kernel"
149
+ ):
150
+ method = config.get_quant_method(fake_layer, "model.layers.0.mlp.gate_up_proj")
151
+
152
+ assert isinstance(method, ModelOptNvFp4LinearMethod)
153
+
154
+
155
+@pytest.mark.parametrize(
156
+ ("quantized_prefix", "missing_prefix"),
157
+ [
158
+ ("model.layers.0.mlp.gate_proj", "model.layers.0.mlp.down_proj"),
159
+ ("model.layers.0.self_attn.o_proj", "model.layers.0.self_attn.qkv_proj"),
160
+ ],
161
+)
162
+def test_modelopt_mixed_precision_does_not_infer_missing_sibling_linear(
163
+ quantized_prefix, missing_prefix
164
+):
165
+ from vllm.model_executor.layers.linear import LinearBase
166
+
167
+ config = _mixed_precision_config(
168
+ {
169
+ quantized_prefix: {"quant_algo": "NVFP4"},
170
+ }
171
+ )
172
+
173
+ fake_layer = MagicMock(spec=LinearBase)
174
+ method = config.get_quant_method(fake_layer, missing_prefix)
175
+
176
+ assert isinstance(method, UnquantizedLinearMethod)
177
+
178
+
179
def test_vocab_parallel_embedding_weight_loader_accepts_scalar_scale():
180
holder = Mock()
181
scale = torch.nn.Parameter(torch.empty(1))
vllm/model_executor/layers/quantization/modelopt.py
+23
-10
@@ -2461,16 +2461,29 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
2461
if key.startswith(parent_dot):
2462
return info["quant_algo"].upper()
2463
2464
- # 4. Parent-prefix fallback for fused projections (qkv_proj, gate_up_proj).
2465
- for candidate in self._quantized_layer_prefix_candidates(prefix):
2466
- parent_dot = candidate.rsplit(".", 1)[0] + "."
2467
- algos = {
2468
- info["quant_algo"].upper()
2469
- for key, info in self.quantized_layers.items()
2470
- if key.startswith(parent_dot) and "." not in key[len(parent_dot) :]
2471
- }
2472
- if len(algos) == 1:
2473
- return algos.pop()
2464
+ # 4. Parent-prefix fallback for fused projections whose config lists
2465
+ # shard names instead of vLLM's packed module name.
2466
+ fused_projection_shards = {
2467
+ "qkv_proj": ("q_proj", "k_proj", "v_proj"),
2468
+ "gate_up_proj": ("gate_proj", "up_proj"),
2469
+ }
2470
+ shard_names = fused_projection_shards.get(proj_name)
2471
+ if shard_names is not None:
2472
+ for candidate in self._quantized_layer_prefix_candidates(prefix):
2473
+ parent_dot = candidate.rsplit(".", 1)[0] + "."
2474
+ shard_algos: set[str] = set()
2475
+ for shard_name in shard_names:
2476
+ shard_prefix = f"{parent_dot}{shard_name}"
2477
+ if shard_prefix in self.quantized_layers:
2478
+ algo = self.quantized_layers[shard_prefix]["quant_algo"].upper()
2479
+ shard_algos.add(algo)
2480
+ if len(shard_algos) == 1:
2481
+ return shard_algos.pop()
2482
+ if len(shard_algos) > 1:
2483
+ raise ValueError(
2484
+ f"Mixed quant_algo within fused layer {prefix}: "
2485
+ f"{shard_algos}. All shards must use the same quantization."
2486
+ )
2487
2488
return None
2489