[Bugfix][Quantization][XPU] Fix moe_wna16 linear weight loading (#52651)
Signed-off-by: Artur Fierka <artur.fierka@intel.com>
Artur Fierka committed
Sep 7, 2026 at 14:20 UTC
5e6f6a8ed42fd420ef206b938e3bb48739a95d6c
3 files changed
+65
-6
tests/quantization/test_moe_wna16.py
+53
@@ -313,3 +313,56 @@ def test_compressed_tensors_wna16_moe_converts_and_sets_up_humming_kernel():
313
assert not hasattr(layer, "w2_weight_packed")
314
assert layer.w13_weight.dtype is torch.int32
315
assert layer.w2_weight.dtype is torch.int32
316
+
317
+
318
+def test_moe_wna16_forwards_packed_modules_mapping_to_linear_delegate(monkeypatch):
319
+ """The linear delegate must receive packed_modules_mapping.
320
+
321
+ It is rebuilt from the raw HF quantization dict, which lists shard names and
322
+ never fused ones, so without the mapping a fused layer resolves to
323
+ `UnquantizedLinearMethod` and the checkpoint's qweight has nowhere to load.
324
+ """
325
+ from vllm.model_executor.layers.linear import ColumnParallelLinear
326
+ from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
327
+
328
+ config = MoeWNA16Config(
329
+ linear_quant_method="gptq",
330
+ weight_bits=4,
331
+ group_size=128,
332
+ has_zp=False,
333
+ lm_head_quantized=False,
334
+ modules_to_not_convert=None,
335
+ full_config={
336
+ "bits": 4,
337
+ "group_size": 128,
338
+ "desc_act": False,
339
+ "sym": True,
340
+ "quant_method": "gptq",
341
+ # As emitted by AutoGPTQ: shard names, never the fused name.
342
+ "modules_in_block_to_quantize": [["mlp.gate_proj", "mlp.up_proj"]],
343
+ },
344
+ )
345
+ config.packed_modules_mapping = {"gate_up_proj": ["gate_proj", "up_proj"]}
346
+
347
+ seen: dict[str, dict[str, list[str]]] = {}
348
+ monkeypatch.setattr(
349
+ AutoGPTQConfig,
350
+ "get_quant_method",
351
+ lambda self, layer, prefix: seen.setdefault(
352
+ "mapping", self.packed_modules_mapping
353
+ ),
354
+ )
355
+ layer = ColumnParallelLinear.__new__(ColumnParallelLinear)
356
+ config.get_quant_method(layer, "model.layers.0.mlp.gate_up_proj")
357
+
358
+ assert seen["mapping"] == {"gate_up_proj": ["gate_proj", "up_proj"]}
359
+
360
+
361
+def test_xpu_platform_supports_moe_wna16():
362
+ """Regression guard for the XPU quantization allowlist."""
363
+ try:
364
+ from vllm.platforms.xpu import XPUPlatform
365
+ except ImportError:
366
+ pytest.skip("vllm_xpu_kernels not importable outside an XPU stack")
367
+
368
+ assert "moe_wna16" in XPUPlatform.supported_quantization
vllm/model_executor/layers/quantization/moe_wna16.py
+11
-6
@@ -182,16 +182,21 @@ class MoeWNA16Config(QuantizationConfig):
182
AutoGPTQConfig,
183
)
184
185
+ linear_config: QuantizationConfig
186
if self.linear_quant_method == "gptq":
186
- return AutoGPTQConfig.from_config(self.full_config).get_quant_method(
187
- layer, prefix
188
- )
187
+ linear_config = AutoGPTQConfig.from_config(self.full_config)
188
elif self.linear_quant_method in ("awq", "awq_marlin"):
190
- return AutoAWQConfig.from_config(self.full_config).get_quant_method(
191
- layer, prefix
192
- )
189
+ linear_config = AutoAWQConfig.from_config(self.full_config)
190
else:
191
raise ValueError("moe_wna16 only support gptq and awq.")
192
+
193
+ # The delegate is rebuilt from the raw HF quantization dict, which
194
+ # lists shard names and never fused ones. Without the mapping it
195
+ # resolves fused layers (qkv_proj, gate_up_proj) to
196
+ # UnquantizedLinearMethod and the checkpoint's qweight has nowhere
197
+ # to load.
198
+ linear_config.packed_modules_mapping = self.packed_modules_mapping
199
+ return linear_config.get_quant_method(layer, prefix)
200
elif isinstance(layer, RoutedExperts):
201
return MoeWNA16Method(self, layer.moe_config)
202
return None
vllm/platforms/xpu.py
+1
@@ -113,6 +113,7 @@ class XPUPlatform(Platform):
113
supported_quantization: list[str] = [
114
"awq",
115
"gptq",
116
+ "moe_wna16",
117
"auto_awq",
118
"auto_gptq",
119
"inc",