[Quantization][Autoround][XPU] Support AutoRound MXFP8 MoE models (#51248)
Signed-off-by: jl9876 <jia.li@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
jl9876 committed
Aug 31, 2026 at 16:37 UTC
1b9539d37c90f179886a462d2149189e3fe38045
3 files changed
+341
tests/quantization/test_auto_round.py
+133
@@ -104,11 +104,20 @@ QWEN3_AUTOROUND_MODELS = [
104
),
105
id="auto_round:mxfp4:qwen3-30b-a3b",
106
),
107
+ pytest.param(
108
+ "INCModel/Qwen3-30B-A3B-12L-MXFP8-test",
109
+ marks=pytest.mark.skipif(
110
+ not (current_platform.is_cuda() or current_platform.is_xpu()),
111
+ reason="Qwen3-30B-A3B MXFP8 AutoRound model requires CUDA/XPU.",
112
+ ),
113
+ id="auto_round:mxfp8:qwen3-30b-a3b",
114
+ ),
115
]
116
117
MODEL_RUNNER_KWARGS: dict[str, dict[str, Any]] = {
118
"INCModel/Qwen3-1.7B-AutoRound-MXFP4-W4A4": {"enforce_eager": True},
119
"INCModel/Qwen3-30B-A3B-12L-MXFP4-test": {"enforce_eager": True},
120
+ "INCModel/Qwen3-30B-A3B-12L-MXFP8-test": {"enforce_eager": True},
121
"Intel/Qwen3-8B-w2g64-for-ut": {
122
"block_size": 64,
123
"gpu_memory_utilization": 0.8,
@@ -232,6 +241,23 @@ def make_qwen3_autoround_config(kind: str) -> INCConfig:
241
},
242
},
243
},
244
+ "qwen3_30b_a3b_mxfp8": {
245
+ "quant_method": "auto-round",
246
+ "bits": 8,
247
+ "group_size": 32,
248
+ "sym": True,
249
+ "packing_format": "auto_round:llm_compressor",
250
+ "data_type": "mx_fp",
251
+ "act_bits": 8,
252
+ "act_group_size": 32,
253
+ "act_data_type": "mx_fp",
254
+ "extra_config": {
255
+ "model.layers.0.mlp.gate": {
256
+ "bits": 16,
257
+ "data_type": "float",
258
+ },
259
+ },
260
+ },
261
}
262
try:
263
config = configs[kind]
@@ -606,6 +632,113 @@ def test_qwen3_30b_a3b_mxfp4_autoround_routes_to_mxfp4_moe(
632
assert method.moe_config is layer.moe_config
633
634
635
+def test_qwen3_30b_a3b_mxfp8_autoround_routes_to_mxfp8_moe(
636
+ monkeypatch,
637
+) -> None:
638
+ expected_backend = object()
639
+ expected_experts_cls = object()
640
+ captured = {}
641
+
642
+ def fake_select_mxfp8_moe_backend(*, config):
643
+ captured["moe_config"] = config
644
+ return expected_backend, expected_experts_cls
645
+
646
+ monkeypatch.setattr(
647
+ "vllm.model_executor.layers.quantization.inc.schemes."
648
+ "inc_mxfp8_moe.select_mxfp8_moe_backend",
649
+ fake_select_mxfp8_moe_backend,
650
+ )
651
+
652
+ config = make_qwen3_autoround_config("qwen3_30b_a3b_mxfp8")
653
+ layer = object.__new__(RoutedExperts)
654
+ layer.moe_config = cast(Any, "moe-config")
655
+
656
+ linear_method = config.get_quant_method(
657
+ object.__new__(LinearBase), "model.layers.0.self_attn.q_proj"
658
+ )
659
+ moe_method = config.get_quant_method(layer, "model.layers.0.mlp")
660
+ gate_method = config.get_quant_method(
661
+ object.__new__(LinearBase), "model.layers.0.mlp.gate"
662
+ )
663
+
664
+ from vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_moe import (
665
+ INCMxfp8MoEMethod,
666
+ )
667
+
668
+ assert isinstance(linear_method, INCLinearMethod)
669
+ assert isinstance(linear_method.scheme, INCMxfp8LinearScheme)
670
+ assert isinstance(moe_method, INCMxfp8MoEMethod)
671
+ assert moe_method.mxfp8_backend is expected_backend
672
+ assert moe_method.experts_cls is expected_experts_cls
673
+ assert captured["moe_config"] == layer.moe_config
674
+ assert isinstance(gate_method, UnquantizedLinearMethod)
675
+
676
+
677
+def test_inc_mxfp8_moe_method_registers_block_scaled_weights(monkeypatch) -> None:
678
+ expected_backend = object()
679
+ expected_experts_cls = object()
680
+
681
+ monkeypatch.setattr(
682
+ "vllm.model_executor.layers.quantization.inc.schemes."
683
+ "inc_mxfp8_moe.select_mxfp8_moe_backend",
684
+ lambda *, config: (expected_backend, expected_experts_cls),
685
+ )
686
+ monkeypatch.setattr(
687
+ "vllm.model_executor.parameter.get_tensor_model_parallel_rank",
688
+ lambda: 0,
689
+ )
690
+ monkeypatch.setattr(
691
+ "vllm.model_executor.parameter.get_tensor_model_parallel_world_size",
692
+ lambda: 1,
693
+ )
694
+
695
+ from vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_moe import (
696
+ INCMxfp8MoEMethod,
697
+ )
698
+
699
+ class DummyMoeConfig:
700
+ is_act_and_mul = True
701
+
702
+ method = INCMxfp8MoEMethod(cast(Any, DummyMoeConfig()))
703
+ layer = torch.nn.Module()
704
+ layer.hidden_size = 64
705
+ layer.intermediate_size_per_partition = 32
706
+
707
+ method.create_weights(
708
+ layer,
709
+ num_experts=2,
710
+ hidden_size=64,
711
+ intermediate_size_per_partition=32,
712
+ params_dtype=torch.bfloat16,
713
+ )
714
+
715
+ assert method.mxfp8_backend is expected_backend
716
+ assert method.experts_cls is expected_experts_cls
717
+ assert method.weight_block_size == [1, 32]
718
+ assert layer.w13_weight.shape == (2, 64, 64)
719
+ assert layer.w2_weight.shape == (2, 64, 32)
720
+ assert layer.w13_weight.dtype is torch.float8_e4m3fn
721
+ assert layer.w2_weight.dtype is torch.float8_e4m3fn
722
+ assert layer.w13_weight_scale.shape == (2, 64, 2)
723
+ assert layer.w2_weight_scale.shape == (2, 64, 1)
724
+ assert layer.w13_weight_scale.dtype is torch.uint8
725
+ assert layer.w2_weight_scale.dtype is torch.uint8
726
+ assert layer.w13_weight_scale.quant_method == "block"
727
+ assert layer.w2_weight_scale.quant_method == "block"
728
+
729
+ misaligned_layer = torch.nn.Module()
730
+ misaligned_layer.hidden_size = 65
731
+ misaligned_layer.intermediate_size_per_partition = 32
732
+ with pytest.raises(ValueError, match="hidden_size divisible by 32"):
733
+ method.create_weights(
734
+ misaligned_layer,
735
+ num_experts=2,
736
+ hidden_size=65,
737
+ intermediate_size_per_partition=32,
738
+ params_dtype=torch.bfloat16,
739
+ )
740
+
741
+
742
def test_inc_mxfp4_linear_method_registers_and_processes_weights(
743
monkeypatch,
744
) -> None:
vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_moe.py
new
+196
@@ -0,0 +1,196 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm.model_executor.layers.fused_moe import (
7
+ FusedMoEQuantConfig,
8
+ FusedMoeWeightScaleSupported,
9
+ RoutedExperts,
10
+ SharedExperts,
11
+)
12
+from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
13
+ FusedMoEMethodBase,
14
+)
15
+from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
16
+ convert_to_fp8_moe_kernel_format,
17
+ make_fp8_moe_kernel,
18
+ make_fp8_moe_quant_config,
19
+)
20
+from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
21
+ select_mxfp8_moe_backend,
22
+)
23
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
24
+ MXFP8_BLOCK_SIZE,
25
+ MXFP8_SCALE_DTYPE,
26
+ MXFP8_VALUE_DTYPE,
27
+)
28
+from vllm.model_executor.parameter import ModelWeightParameter
29
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
30
+
31
+
32
+class INCMxfp8MoEMethod(FusedMoEMethodBase):
33
+ """W8A8 MXFP8 MoE method for serialized AutoRound checkpoints."""
34
+
35
+ def __init__(self, moe) -> None:
36
+ super().__init__(moe)
37
+ self.weight_block_size = [1, MXFP8_BLOCK_SIZE]
38
+ self.mxfp8_backend, self.experts_cls = select_mxfp8_moe_backend(config=self.moe)
39
+
40
+ def create_weights(
41
+ self,
42
+ layer: RoutedExperts,
43
+ num_experts: int,
44
+ hidden_size: int,
45
+ intermediate_size_per_partition: int,
46
+ params_dtype: torch.dtype,
47
+ **extra_weight_attrs,
48
+ ) -> None:
49
+ assert layer.intermediate_size_per_partition == intermediate_size_per_partition
50
+ assert layer.hidden_size == hidden_size
51
+ layer.orig_dtype = params_dtype
52
+
53
+ if hidden_size % MXFP8_BLOCK_SIZE != 0:
54
+ raise ValueError(
55
+ f"MXFP8 MoE requires hidden_size divisible by {MXFP8_BLOCK_SIZE}, "
56
+ f"got {hidden_size}."
57
+ )
58
+ if intermediate_size_per_partition % MXFP8_BLOCK_SIZE != 0:
59
+ raise ValueError(
60
+ "MXFP8 MoE requires intermediate_size_per_partition divisible by "
61
+ f"{MXFP8_BLOCK_SIZE}, got {intermediate_size_per_partition}."
62
+ )
63
+
64
+ layer.num_experts = num_experts
65
+ weight_loader = extra_weight_attrs.get("weight_loader")
66
+ w13_num_shards = 2 if self.moe.is_act_and_mul else 1
67
+
68
+ w13_weight = ModelWeightParameter(
69
+ data=torch.empty(
70
+ num_experts,
71
+ w13_num_shards * intermediate_size_per_partition,
72
+ hidden_size,
73
+ dtype=MXFP8_VALUE_DTYPE,
74
+ ),
75
+ input_dim=2,
76
+ output_dim=1,
77
+ weight_loader=weight_loader,
78
+ )
79
+ layer.register_parameter("w13_weight", w13_weight)
80
+
81
+ w2_weight = ModelWeightParameter(
82
+ data=torch.empty(
83
+ num_experts,
84
+ hidden_size,
85
+ intermediate_size_per_partition,
86
+ dtype=MXFP8_VALUE_DTYPE,
87
+ ),
88
+ input_dim=2,
89
+ output_dim=1,
90
+ weight_loader=weight_loader,
91
+ )
92
+ layer.register_parameter("w2_weight", w2_weight)
93
+
94
+ w13_weight_scale = ModelWeightParameter(
95
+ data=torch.empty(
96
+ num_experts,
97
+ w13_num_shards * intermediate_size_per_partition,
98
+ hidden_size // MXFP8_BLOCK_SIZE,
99
+ dtype=MXFP8_SCALE_DTYPE,
100
+ ),
101
+ input_dim=2,
102
+ output_dim=1,
103
+ weight_loader=weight_loader,
104
+ )
105
+ layer.register_parameter("w13_weight_scale", w13_weight_scale)
106
+
107
+ w2_weight_scale = ModelWeightParameter(
108
+ data=torch.empty(
109
+ num_experts,
110
+ hidden_size,
111
+ intermediate_size_per_partition // MXFP8_BLOCK_SIZE,
112
+ dtype=MXFP8_SCALE_DTYPE,
113
+ ),
114
+ input_dim=2,
115
+ output_dim=1,
116
+ weight_loader=weight_loader,
117
+ )
118
+ layer.register_parameter("w2_weight_scale", w2_weight_scale)
119
+
120
+ set_weight_attrs(
121
+ layer.w13_weight_scale,
122
+ {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value},
123
+ )
124
+ set_weight_attrs(
125
+ layer.w2_weight_scale,
126
+ {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value},
127
+ )
128
+
129
+ def process_weights_after_loading(self, layer: RoutedExperts) -> None:
130
+ layer.weight_block_size = self.weight_block_size
131
+
132
+ w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
133
+ fp8_backend=self.mxfp8_backend,
134
+ layer=layer,
135
+ w13=layer.w13_weight,
136
+ w2=layer.w2_weight,
137
+ w13_scale=layer.w13_weight_scale,
138
+ w2_scale=layer.w2_weight_scale,
139
+ w13_input_scale=None,
140
+ w2_input_scale=None,
141
+ )
142
+
143
+ replace_parameter(layer, "w13_weight", w13)
144
+ replace_parameter(layer, "w2_weight", w2)
145
+ replace_parameter(layer, "w13_weight_scale", w13_scale)
146
+ replace_parameter(layer, "w2_weight_scale", w2_scale)
147
+
148
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
149
+ assert self.moe_quant_config is not None
150
+ self.moe_kernel = make_fp8_moe_kernel(
151
+ moe_quant_config=self.moe_quant_config,
152
+ moe_config=self.moe,
153
+ fp8_backend=self.mxfp8_backend,
154
+ experts_cls=self.experts_cls,
155
+ routing_tables=layer._expert_routing_tables(),
156
+ )
157
+
158
+ def get_fused_moe_quant_config(
159
+ self, layer: RoutedExperts
160
+ ) -> FusedMoEQuantConfig | None:
161
+ return make_fp8_moe_quant_config(
162
+ fp8_backend=self.mxfp8_backend,
163
+ w1_scale=layer.w13_weight_scale,
164
+ w2_scale=layer.w2_weight_scale,
165
+ a1_scale=None,
166
+ a2_scale=None,
167
+ block_shape=self.weight_block_size,
168
+ swiglu_limit=getattr(layer, "swiglu_limit", None),
169
+ gemm1_alpha=getattr(layer, "swiglu_alpha", None),
170
+ gemm1_beta=getattr(layer, "swiglu_beta", None),
171
+ layer=layer,
172
+ )
173
+
174
+ def apply(
175
+ self,
176
+ layer: RoutedExperts,
177
+ x: torch.Tensor,
178
+ topk_weights: torch.Tensor,
179
+ topk_ids: torch.Tensor,
180
+ shared_experts: SharedExperts | None,
181
+ shared_experts_input: torch.Tensor | None,
182
+ ) -> torch.Tensor:
183
+ assert self.moe_kernel is not None
184
+ return self.moe_kernel.apply(
185
+ x,
186
+ layer.w13_weight,
187
+ layer.w2_weight,
188
+ topk_weights,
189
+ topk_ids,
190
+ activation=layer.activation,
191
+ global_num_experts=layer.global_num_experts,
192
+ expert_map=layer.expert_map,
193
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
194
+ shared_experts=shared_experts,
195
+ shared_experts_input=shared_experts_input,
196
+ )
vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp8_scheme.py
+12
@@ -29,3 +29,15 @@ class INCMxfp8Scheme(INCScheme):
29
from .inc_mxfp8_linear import INCMxfp8LinearScheme
30
31
return INCLinearMethod(INCMxfp8LinearScheme())
32
+
33
+ def get_moe_method(
34
+ self,
35
+ config: "INCConfig",
36
+ layer: "torch.nn.Module",
37
+ prefix: str,
38
+ layer_config: "INCLayerConfig",
39
+ ):
40
+ del config, prefix, layer_config
41
+ from .inc_mxfp8_moe import INCMxfp8MoEMethod
42
+
43
+ return INCMxfp8MoEMethod(layer.moe_config)