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 3, 19:00 UTC
[Bugfix][Quark] Fix W8A8 INT8 garbage outputs on Step-3.5-Flash (and other 3-key fused-MoE Quark exports) (#41892)
Signed-off-by: kangletian <kangletian@hotmail.com>
PikaPikachu committed
May 13, 2026 at 19:59 UTC
3b1ef03be4a3c7f11b801ec3f2aa806b8c37573b
3 files changed
+46
-4
vllm/model_executor/layers/quantization/quark/quark_moe.py
+25
@@ -23,6 +23,7 @@ from vllm.model_executor.layers.fused_moe.config import (
23
FusedMoEParallelConfig,
24
FusedMoEQuantConfig,
25
fp8_w8a8_moe_quant_config,
26
+ int8_w8a8_moe_quant_config,
27
mxfp4_w4a8_moe_quant_config,
28
mxfp4_w4a16_moe_quant_config,
29
ocp_mx_moe_quant_config,
@@ -708,6 +709,20 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
709
layer.w2_input_scale.max(), requires_grad=False
710
)
711
712
+ # Per-channel scales: 2D [E, N] -> 3D [E, N, 1] for the int8 MoE kernel.
713
+ if self.weight_qscheme == "per_channel":
714
+ for attr in ("w13_weight_scale", "w2_weight_scale"):
715
+ param = getattr(layer, attr, None)
716
+ if param is not None and param.dim() == 2:
717
+ replace_parameter(
718
+ layer,
719
+ attr,
720
+ torch.nn.Parameter(
721
+ param.data.unsqueeze(-1).contiguous(),
722
+ requires_grad=False,
723
+ ),
724
+ )
725
+
726
# For per-tensor weights, merge w1/w3 scales into single per-expert
727
if self.weight_qscheme == "per_tensor":
728
assert layer.w13_weight_scale is not None
@@ -736,6 +751,16 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
751
def get_fused_moe_quant_config(
752
self, layer: torch.nn.Module
753
) -> FusedMoEQuantConfig | None:
754
+ if self.weight_qscheme == "per_channel" and not self.static_input_scales:
755
+ return int8_w8a8_moe_quant_config(
756
+ w1_scale=layer.w13_weight_scale,
757
+ w2_scale=layer.w2_weight_scale,
758
+ a1_scale=layer.w13_input_scale,
759
+ a2_scale=layer.w2_input_scale,
760
+ w1_bias=getattr(layer, "w13_bias", None),
761
+ w2_bias=getattr(layer, "w2_bias", None),
762
+ per_act_token_quant=True,
763
+ )
764
is_dynamic = not self.static_input_scales
765
is_per_channel = self.weight_qscheme == "per_channel"
766
return FusedMoEQuantConfig.make(
vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py
+15
-4
@@ -47,6 +47,17 @@ class QuarkW8A8Int8(QuarkScheme):
47
):
48
layer.logical_widths = output_partition_sizes
49
50
+ # Quark stores per-channel weight_scale as 1D [N]; reshape to [N, 1].
51
+ def _scale_weight_loader(
52
+ param: torch.nn.Parameter,
53
+ loaded_weight: torch.Tensor,
54
+ *args,
55
+ **kwargs,
56
+ ):
57
+ if loaded_weight.dim() == 1:
58
+ loaded_weight = loaded_weight.unsqueeze(-1)
59
+ return weight_loader(param, loaded_weight, *args, **kwargs)
60
+
61
self.kernel = init_int8_linear_kernel(
62
is_channelwise=(self.qscheme == "per_channel"),
63
is_static_input_scheme=(self.is_static_input_scheme is True),
@@ -69,15 +80,15 @@ class QuarkW8A8Int8(QuarkScheme):
80
# WEIGHT SCALE
81
if self.qscheme == "per_channel":
82
weight_scale = ChannelQuantScaleParameter(
72
- data=torch.empty((sum(output_partition_sizes)), dtype=torch.float32),
83
+ data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32),
84
output_dim=0,
74
- weight_loader=weight_loader,
85
+ weight_loader=_scale_weight_loader,
86
)
87
ChannelQuantZPParameter = ChannelQuantScaleParameter
88
weight_zero_point = ChannelQuantZPParameter(
78
- data=torch.empty((sum(output_partition_sizes)), dtype=torch.int8),
89
+ data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.int8),
90
output_dim=0,
80
- weight_loader=weight_loader,
91
+ weight_loader=_scale_weight_loader,
92
)
93
else:
94
assert self.qscheme == "per_tensor"
vllm/model_executor/models/step3p5.py
+6
@@ -817,6 +817,12 @@ class Step3p5Model(nn.Module):
817
818
819
class Step3p5ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
820
+ # Required so quantization exclude lists match fused module prefixes.
821
+ packed_modules_mapping = {
822
+ "qkv_proj": ["q_proj", "k_proj", "v_proj"],
823
+ "gate_up_proj": ["gate_proj", "up_proj"],
824
+ }
825
+
826
hf_to_vllm_mapper = WeightsMapper(
827
orig_to_new_substr={".share_expert.": ".moe.share_expert."}
828
)