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 10, 02:45 UTC
[Quant] Support MXFP4 W4A16 for compressed-tensors MoE models (#32285)
Signed-off-by: Dipika Sikka <dipikasikka1@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Dipika Sikka committed
Jan 15, 2026 at 10:25 UTC
361dfdc9d81ef16f00cf9f51dc104b55e382ee61
4 files changed
+156
tests/evals/gsm8k/configs/Qwen3-30B-A3B-MXFP4A16.yaml
new
+5
@@ -0,0 +1,5 @@
1
+model_name: nm-testing/Qwen3-30B-A3B-MXFP4A16
2
+accuracy_threshold: 0.88
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 4096"
tests/evals/gsm8k/configs/models-small.txt
+1
@@ -4,3 +4,4 @@ Llama-3-8B-Instruct-nonuniform-CT.yaml
4
Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
5
Qwen1.5-MoE-W4A16-CT.yaml
6
DeepSeek-V2-Lite-Instruct-FP8.yaml
7
+Qwen3-30B-A3B-MXFP4A16.yaml
\ No newline at end of file
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+11
@@ -10,6 +10,7 @@ from vllm.logger import init_logger
10
from vllm.model_executor.layers.fused_moe.config import (
11
FusedMoEConfig,
12
FusedMoEQuantConfig,
13
+ mxfp4_w4a16_moe_quant_config,
14
nvfp4_moe_quant_config,
15
nvfp4_w4a16_moe_quant_config,
16
)
@@ -193,6 +194,16 @@ def convert_to_nvfp4_moe_kernel_format(
194
)
195
196
197
+def make_mxfp4_moe_quant_config(
198
+ w13_scale: torch.Tensor,
199
+ w2_scale: torch.Tensor,
200
+) -> FusedMoEQuantConfig:
201
+ return mxfp4_w4a16_moe_quant_config(
202
+ w1_scale=w13_scale,
203
+ w2_scale=w2_scale,
204
+ )
205
+
206
+
207
def make_nvfp4_moe_quant_config(
208
backend: NvFp4MoeBackend,
209
w13_scale: torch.Tensor,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+139
@@ -52,6 +52,7 @@ from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
52
NvFp4MoeBackend,
53
convert_to_nvfp4_moe_kernel_format,
54
is_global_sf_supported_for_nvfp4_backend,
55
+ make_mxfp4_moe_quant_config,
56
make_nvfp4_moe_kernel,
57
make_nvfp4_moe_quant_config,
58
select_nvfp4_moe_backend,
@@ -79,6 +80,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
80
)
81
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
82
is_fp4_marlin_supported,
83
+ prepare_moe_fp4_layer_for_marlin,
84
)
85
from vllm.model_executor.layers.quantization.utils.quant_utils import (
86
convert_bf16_scales_to_fp8,
@@ -145,6 +147,9 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
147
input_quant = scheme_dict.get("input_activations")
148
format = scheme_dict.get("format")
149
150
+ if quant_config._is_mxfp4(weight_quant):
151
+ return CompressedTensorsW4A4Mxfp4MoEMethod(layer.moe_config)
152
+
153
if quant_config._is_wNa16_group_channel(weight_quant, input_quant):
154
# group_size=None means channelwise
155
group_size = weight_quant.group_size or -1
@@ -224,6 +229,140 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
229
)
230
231
232
+class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
233
+ def __init__(self, moe):
234
+ super().__init__(moe)
235
+ self.group_size = 32
236
+ self.mxfp4_backend = NvFp4MoeBackend.MARLIN
237
+ self.kernel: mk.FusedMoEModularKernel | None = None
238
+
239
+ def create_weights(
240
+ self,
241
+ layer: torch.nn.Module,
242
+ num_experts: int,
243
+ hidden_size: int,
244
+ intermediate_size_per_partition: int,
245
+ params_dtype: torch.dtype,
246
+ **extra_weight_attrs,
247
+ ):
248
+ layer.num_experts = num_experts
249
+ layer.params_dtype = params_dtype
250
+
251
+ w13_weight = torch.nn.Parameter(
252
+ torch.empty(
253
+ num_experts,
254
+ 2 * intermediate_size_per_partition,
255
+ # 2 fp4 items are packed in the input dimension
256
+ hidden_size // 2,
257
+ requires_grad=False,
258
+ dtype=torch.uint8,
259
+ ),
260
+ requires_grad=False,
261
+ )
262
+ layer.register_parameter("w13_weight_packed", w13_weight)
263
+ set_weight_attrs(w13_weight, extra_weight_attrs)
264
+
265
+ w2_weight = torch.nn.Parameter(
266
+ torch.empty(
267
+ num_experts,
268
+ hidden_size,
269
+ # 2 fp4 items are packed in the input dimension
270
+ intermediate_size_per_partition // 2,
271
+ dtype=torch.uint8,
272
+ ),
273
+ requires_grad=False,
274
+ )
275
+ layer.register_parameter("w2_weight_packed", w2_weight)
276
+ set_weight_attrs(w2_weight, extra_weight_attrs)
277
+
278
+ w13_weight_scale = torch.nn.Parameter(
279
+ torch.empty(
280
+ num_experts,
281
+ 2 * intermediate_size_per_partition,
282
+ # 2 fp4 items are packed in the input dimension
283
+ hidden_size // self.group_size,
284
+ dtype=torch.uint8,
285
+ ),
286
+ requires_grad=False,
287
+ )
288
+ layer.register_parameter("w13_weight_scale", w13_weight_scale)
289
+ extra_weight_attrs.update(
290
+ {"quant_method": FusedMoeWeightScaleSupported.GROUP.value}
291
+ )
292
+ set_weight_attrs(w13_weight_scale, extra_weight_attrs)
293
+
294
+ w2_weight_scale = torch.nn.Parameter(
295
+ torch.empty(
296
+ num_experts,
297
+ hidden_size,
298
+ # 2 fp4 items are packed in the input dimension
299
+ intermediate_size_per_partition // self.group_size,
300
+ dtype=torch.uint8,
301
+ ),
302
+ requires_grad=False,
303
+ )
304
+ layer.register_parameter("w2_weight_scale", w2_weight_scale)
305
+ set_weight_attrs(w2_weight_scale, extra_weight_attrs)
306
+
307
+ def get_fused_moe_quant_config(
308
+ self, layer: torch.nn.Module
309
+ ) -> FusedMoEQuantConfig | None:
310
+ return make_mxfp4_moe_quant_config(
311
+ w13_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale
312
+ )
313
+
314
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
315
+ layer.w13_weight = torch.nn.Parameter(
316
+ layer.w13_weight_packed.data, requires_grad=False
317
+ )
318
+ delattr(layer, "w13_weight_packed")
319
+
320
+ layer.w2_weight = torch.nn.Parameter(
321
+ layer.w2_weight_packed.data, requires_grad=False
322
+ )
323
+ delattr(layer, "w2_weight_packed")
324
+
325
+ prepare_moe_fp4_layer_for_marlin(layer)
326
+
327
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
328
+ if self.moe_quant_config is not None:
329
+ self.kernel = make_nvfp4_moe_kernel(
330
+ backend=self.mxfp4_backend,
331
+ quant_config=self.moe_quant_config,
332
+ moe_config=self.moe,
333
+ )
334
+
335
+ def apply(
336
+ self,
337
+ layer: FusedMoE,
338
+ router: FusedMoERouter,
339
+ x: torch.Tensor,
340
+ router_logits: torch.Tensor,
341
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
342
+ if isinstance(x, tuple):
343
+ x_routing, _ = x
344
+ else:
345
+ x_routing = x
346
+
347
+ topk_weights, topk_ids = router.select_experts(
348
+ hidden_states=x_routing,
349
+ router_logits=router_logits,
350
+ )
351
+ assert self.kernel is not None
352
+ return self.kernel(
353
+ x,
354
+ layer.w13_weight,
355
+ layer.w2_weight,
356
+ topk_weights,
357
+ topk_ids,
358
+ inplace=False,
359
+ activation=layer.activation,
360
+ global_num_experts=layer.global_num_experts,
361
+ expert_map=layer.expert_map,
362
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
363
+ )
364
+
365
+
366
class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
367
def __init__(
368
self,