[Bugfix] Fix moe_wna16 w13 zero-point shard split for 8-bit asym GPTQ MoE (#58950)
Signed-off-by: Spandan Tiwari <sptiwari@amd.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Spandan Tiwari committed
Sep 28, 2026 at 16:38 UTC
9b29370524cfd84b3a85bae8c2db8112fd52003b
2 files changed
+80
-2
tests/quantization/test_moe_wna16.py
+74
@@ -546,3 +546,77 @@ def test_compressed_tensors_wna16_moe_marlin_prep_with_unset_group_size():
546
)
547
assert layer.w13_weight_scale.shape == (num_experts, 1, 2 * intermediate_size)
548
assert layer.w2_weight_scale.shape == (num_experts, 1, hidden_size)
549
+
550
+
551
+def test_moe_wna16_w13_qzeros_shard_split_int8_asym_gptq(monkeypatch):
552
+ """8-bit asymmetric GPTQ MoE must split w13 zero-points by pack factor.
553
+
554
+ ``w13_qzeros`` is packed along the intermediate dim by ``bit8_pack_factor``,
555
+ so the gate/up (w1/w3) boundary in the weight loader is
556
+ ``shard_size // bit8_pack_factor``. Hardcoding ``shard_size // 2`` only
557
+ holds for 4-bit and raises a shape mismatch when loading an 8-bit
558
+ (unpacked) checkpoint.
559
+ """
560
+ from tests.kernels.moe.utils import make_dummy_moe_config
561
+
562
+ monkeypatch.setattr(
563
+ moe_wna16, "get_tp_group", lambda: SimpleNamespace(device="cpu")
564
+ )
565
+ monkeypatch.setattr(moe_wna16, "get_tensor_model_parallel_rank", lambda: 0)
566
+
567
+ num_experts, hidden, intermediate, group_size = 2, 256, 512, 128
568
+ moe_config = make_dummy_moe_config(
569
+ num_experts=num_experts,
570
+ hidden_dim=hidden,
571
+ intermediate_size=intermediate,
572
+ )
573
+ quant_config = MoeWNA16Config(
574
+ linear_quant_method="gptq",
575
+ weight_bits=8,
576
+ group_size=group_size,
577
+ has_zp=True,
578
+ lm_head_quantized=False,
579
+ modules_to_not_convert=[],
580
+ full_config={},
581
+ )
582
+ method = object.__new__(MoeWNA16Method)
583
+ method.quant_config = quant_config
584
+ method.moe = moe_config
585
+
586
+ layer = torch.nn.Module()
587
+ layer.intermediate_size_per_partition = intermediate
588
+ layer.moe_config = moe_config
589
+ method.create_weights(
590
+ layer,
591
+ num_experts=num_experts,
592
+ hidden_size=hidden,
593
+ intermediate_size_per_partition=intermediate,
594
+ params_dtype=torch.float16,
595
+ weight_loader=lambda *a, **k: True,
596
+ )
597
+
598
+ # 8-bit is unpacked along the intermediate dim: dim1 == 2 * intermediate.
599
+ assert layer.w13_qzeros.shape == (
600
+ num_experts,
601
+ 2 * intermediate,
602
+ hidden // group_size,
603
+ )
604
+ loader = layer.w13_qzeros.weight_loader
605
+
606
+ # Raw GPTQ 8-bit qzeros per shard: int32 (num_groups, intermediate // 4);
607
+ # the loader views them as uint8 -> (num_groups, intermediate), then
608
+ # transposes and offsets by 1.
609
+ num_groups = hidden // group_size
610
+ raw_w1 = torch.arange(num_groups * (intermediate // 4), dtype=torch.int32).reshape(
611
+ num_groups, intermediate // 4
612
+ )
613
+ raw_w3 = raw_w1 + 1000
614
+
615
+ # Must not raise (old code raised a shape mismatch here for 8-bit).
616
+ loader(layer.w13_qzeros, raw_w1, "w13_qzeros", "w1", 0, False)
617
+ loader(layer.w13_qzeros, raw_w3, "w13_qzeros", "w3", 0, False)
618
+
619
+ expected_w1 = raw_w1.view(torch.uint8).T + 1
620
+ expected_w3 = raw_w3.view(torch.uint8).T + 1
621
+ assert torch.equal(layer.w13_qzeros[0, :intermediate].cpu(), expected_w1)
622
+ assert torch.equal(layer.w13_qzeros[0, intermediate:].cpu(), expected_w3)
vllm/model_executor/layers/quantization/moe_wna16.py
+6
-2
@@ -635,10 +635,14 @@ class MoeWNA16Method(FusedMoEMethodBase):
635
tensor = loaded_weight.view(
636
layer.moe_config.tp_size, -1, loaded_weight.size(1)
637
)[tp_rank]
638
+ # qzeros are packed along the intermediate dim by
639
+ # bit8_pack_factor, so the w1/w3 boundary is in packed units
640
+ # (equals shard_size // 2 only for 4-bit).
641
+ shard_zp_size = shard_size // layer.quant_config.bit8_pack_factor
642
if shard_id == "w1":
639
- param.data[expert_id, : shard_size // 2] = tensor
643
+ param.data[expert_id, :shard_zp_size] = tensor
644
else:
641
- param.data[expert_id, shard_size // 2 :] = tensor
645
+ param.data[expert_id, shard_zp_size:] = tensor
646
return True if return_success else None
647
elif "w2_qzeros" in weight_name:
648
param.data[expert_id] = loaded_weight.view(