[Bugfix][ROCm] Fix a few int4/int8 quantization errors (#52112)
Signed-off-by: Qiang Li <qiang.li2@amd.com>
qli88 committed
Aug 18, 2026 at 03:31 UTC
e8ad2855e7f5d40665250eb468bfd567e3a4b3c1
4 files changed
+69
-7
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+7
-2
@@ -48,6 +48,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
48
kFp8StaticTensorSym,
49
kInt4Static,
50
kInt4Static32,
51
+ kInt4Static32Asym,
52
+ kInt4StaticAsym,
53
kInt8DynamicTensorSym,
54
kInt8DynamicTokenSym,
55
kInt8Static,
@@ -186,8 +188,8 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
188
swiglu_limit_func(output, input, activation_config.clamp_limit)
189
return
190
189
- # SWIGLUOAI_UNINTERLEAVE routes to the silu_and_mul_with_clamp kernel and
190
- # requires a clamp limit. Other activations ignore these parameters.
191
+ # SWIGLUOAI_UNINTERLEAVE routes to torch.ops._C.silu_and_mul_with_clamp
192
+ # via apply_moe_activation() and requires clamp_limit to be set.
193
if activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
194
assert activation_config.clamp_limit is not None, (
195
"SWIGLUOAI_UNINTERLEAVE requires gemm1_clamp_limit"
@@ -580,6 +582,8 @@ class TritonWNA16Experts(TritonExperts):
582
kInt4Static,
583
kInt8Static,
584
kInt4Static32,
585
+ kInt4StaticAsym,
586
+ kInt4Static32Asym,
587
# other group sizes?
588
]
589
return weight_key in SUPPORTED_W
@@ -591,6 +595,7 @@ class TritonWNA16Experts(TritonExperts):
595
MoEActivation.GELU,
596
MoEActivation.GELU_TANH,
597
MoEActivation.SWIGLUOAI,
598
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE,
599
MoEActivation.SWIGLUSTEP,
600
MoEActivation.SILU_NO_MUL,
601
MoEActivation.GELU_NO_MUL,
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+36
@@ -1649,6 +1649,42 @@ def convert_to_wna16_moe_kernel_format(
1649
w2_uint8 = w2.transpose(1, 2).contiguous().view(torch.uint8)
1650
w13_scale = w13_scale.transpose(1, 2).contiguous()
1651
w2_scale = w2_scale.transpose(1, 2).contiguous()
1652
+ # Zero points from compressed-tensors checkpoints are K-first int32
1653
+ # with 8 int4 ZPs packed per element: shape (E, K//gs, N//8).
1654
+ # fused_moe_kernel_gptq_awq expects N-first uint8 with 2 int4 ZPs
1655
+ # per byte: shape (E, N//2, K//gs), indexed as
1656
+ # (offs_bn // 2) * stride_bzn + offs_k_group * stride_bzk.
1657
+ # Conversion steps:
1658
+ # (E, K//gs, N//8) int32
1659
+ # → transpose(1,2) → (E, N//8, K//gs) int32
1660
+ # → view(uint8) → (E, N//8, K//gs*4) [each int32 → 4 bytes]
1661
+ # → reshape(…,4) → (E, N//8, K//gs, 4) [isolate byte index]
1662
+ # → permute(0,1,3,2) → (E, N//8, 4, K//gs) [byte index before K]
1663
+ # → reshape → (E, N//2, K//gs) [kernel expected layout]
1664
+ # After this, element [e, offs_bn//2, k_group] is the uint8 byte
1665
+ # holding the two int4 ZPs for output channels offs_bn and offs_bn+1.
1666
+ if w13_qzeros is not None:
1667
+ E13, Kg13, Np13 = w13_qzeros.shape
1668
+ w13_qzeros = (
1669
+ w13_qzeros.transpose(1, 2)
1670
+ .contiguous()
1671
+ .view(torch.uint8)
1672
+ .reshape(E13, Np13, Kg13, 4)
1673
+ .permute(0, 1, 3, 2)
1674
+ .reshape(E13, Np13 * 4, Kg13)
1675
+ .contiguous()
1676
+ )
1677
+ if w2_qzeros is not None:
1678
+ E2, Kg2, Np2 = w2_qzeros.shape
1679
+ w2_qzeros = (
1680
+ w2_qzeros.transpose(1, 2)
1681
+ .contiguous()
1682
+ .view(torch.uint8)
1683
+ .reshape(E2, Np2, Kg2, 4)
1684
+ .permute(0, 1, 3, 2)
1685
+ .reshape(E2, Np2 * 4, Kg2)
1686
+ .contiguous()
1687
+ )
1688
else:
1689
# MoeWNA16 uses N-first uint8 weights and scales.
1690
w13_uint8 = w13.view(torch.uint8)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
+26
-4
@@ -123,8 +123,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
123
# grouped actorder isn't supported by this kernel
124
assert weight_quant.actorder != "group"
125
126
- assert self.symmetric, "Only symmetric quantization is supported for MoE"
127
-
126
# Non-Marlin WNA16 always uses bf16/fp16 inputs
127
self.input_dtype = torch.bfloat16
128
@@ -317,6 +315,27 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
315
num_groups_w2 = w2_scales_size // self.group_size
316
num_groups_w13 = hidden_size // self.group_size
317
318
+ if not self.symmetric:
319
+ # For asymmetric int4 quantization, packed_factor (=8) int4 ZPs are
320
+ # stored per int32. Each TP shard must contain a whole number of
321
+ # int32 elements so the packed ZP can be sliced without straddling
322
+ # an int32 boundary.
323
+ w13_n = 2 * intermediate_size_per_partition # gate+up output channels
324
+ if w13_n % self.packed_factor != 0:
325
+ raise ValueError(
326
+ f"CompressedTensors WNA16 MoE: gate+up output channels per "
327
+ f"TP rank (2 * intermediate_size_per_partition = {w13_n}) "
328
+ f"must be divisible by packed_factor ({self.packed_factor}). "
329
+ f"Use a TP size where 2 * intermediate_size is divisible by "
330
+ f"{self.packed_factor}."
331
+ )
332
+ if hidden_size % self.packed_factor != 0:
333
+ raise ValueError(
334
+ f"CompressedTensors WNA16 MoE: hidden_size ({hidden_size}) "
335
+ f"must be divisible by packed_factor ({self.packed_factor}) "
336
+ f"for correct ZP unpacking."
337
+ )
338
+
339
layer.num_groups_w13 = num_groups_w13
340
layer.num_groups_w2 = num_groups_w2
341
@@ -513,8 +532,11 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
532
replace_parameter(layer, "w13_weight_scale", w13_scales)
533
replace_parameter(layer, "w2_weight_scale", w2_scales)
534
516
- # CPU fused_experts_cpu requires zero points even for symmetric quant
517
- if not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU:
535
+ # CPU fused_experts_cpu requires zero points even for symmetric quant.
536
+ # EMULATION bakes ZP into the dequantized bf16 weights — ZP is None.
537
+ if (
538
+ not self.symmetric or self.wna16_backend == WNA16MoEBackend.CPU
539
+ ) and self.wna16_backend != WNA16MoEBackend.EMULATION:
540
assert w13_qzeros is not None and w2_qzeros is not None
541
replace_parameter(layer, "w13_weight_zero_point", w13_qzeros)
542
replace_parameter(layer, "w2_weight_zero_point", w2_qzeros)
vllm/model_executor/layers/quantization/moe_wna16.py
-1
@@ -222,7 +222,6 @@ class MoeWNA16Method(FusedMoEMethodBase):
222
else:
223
scale = kInt4StaticGroupScale
224
elif num_bits == 8:
225
- assert group_size == -1
225
quant_type = INT8_DTYPE
226
scale = kInt8StaticGroupScale
227
else: