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 11, 06:15 UTC
[Quantization] Support Quark int4-fp8 w4a8 for MoE (#30071)
Signed-off-by: Bowen Bao <bowenbao@amd.com>
Bowen Bao committed
Dec 17, 2025 at 20:20 UTC
0c738b58bc0e5a5bf2448c95fc2014b83127a4d5
2 files changed
+201
-2
vllm/model_executor/layers/quantization/quark/quark.py
+43
@@ -218,6 +218,49 @@ class QuarkConfig(QuantizationConfig):
218
else:
219
return False
220
221
+ def _is_fp8_w4a8(
222
+ self,
223
+ weight_quant: list[dict[str, Any]] | None,
224
+ input_quant: dict[str, Any] | None,
225
+ ) -> bool:
226
+ # Confirm weights and input quantized.
227
+ if weight_quant is None or input_quant is None:
228
+ return False
229
+
230
+ if not isinstance(weight_quant, list) or len(weight_quant) != 2:
231
+ return False
232
+
233
+ # Confirm weight scheme is supported
234
+ is_w4a8_dtype = (
235
+ weight_quant[0].get("dtype") == "fp8_e4m3"
236
+ and weight_quant[1].get("dtype") == "int4"
237
+ and input_quant.get("dtype") == "fp8_e4m3"
238
+ )
239
+ is_static_weight = not weight_quant[0].get("is_dynamic") and not weight_quant[
240
+ 1
241
+ ].get("is_dynamic")
242
+ is_per_tensor_fp8_and_per_channel_int4_weight = (
243
+ weight_quant[0].get("qscheme") == "per_tensor"
244
+ and weight_quant[1].get("qscheme") == "per_channel"
245
+ and weight_quant[1].get("symmetric") is True
246
+ and weight_quant[1].get("ch_axis") == 0
247
+ )
248
+
249
+ if not (
250
+ is_w4a8_dtype
251
+ and is_static_weight
252
+ and is_per_tensor_fp8_and_per_channel_int4_weight
253
+ ):
254
+ return False
255
+
256
+ # Dynamic quantization is always supported if weights supported.
257
+ if input_quant.get("is_dynamic"):
258
+ return True
259
+
260
+ # Confirm activation scheme is supported.
261
+ is_per_tensor_activation = input_quant.get("qscheme") == "per_tensor"
262
+ return is_per_tensor_activation
263
+
264
def _is_fp8_w8a8(
265
self,
266
weight_quant: dict[str, Any] | None,
vllm/model_executor/layers/quantization/quark/quark_moe.py
+158
-2
@@ -63,8 +63,9 @@ class QuarkMoEMethod(FusedMoEMethodBase):
63
)
64
weight_config = layer_quant_config.get("weight")
65
input_config = layer_quant_config.get("input_tensors")
66
-
67
- if quant_config._is_fp8_w8a8(weight_config, input_config):
66
+ if quant_config._is_fp8_w4a8(weight_config, input_config):
67
+ return QuarkW4A8Fp8MoEMethod(weight_config, input_config, module.moe_config)
68
+ elif quant_config._is_fp8_w8a8(weight_config, input_config):
69
return QuarkW8A8Fp8MoEMethod(weight_config, input_config, module.moe_config)
70
elif quant_config._is_ocp_mx(weight_config, input_config):
71
return QuarkOCP_MX_MoEMethod(weight_config, input_config, module.moe_config)
@@ -396,6 +397,161 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
397
)
398
399
400
+class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
401
+ def __init__(
402
+ self,
403
+ weight_config: dict[str, Any],
404
+ input_config: dict[str, Any],
405
+ moe: FusedMoEConfig,
406
+ ):
407
+ super().__init__(moe)
408
+ self.weight_quant = weight_config
409
+ self.input_quant = input_config
410
+
411
+ assert rocm_aiter_ops.is_fused_moe_enabled(), (
412
+ "W4A8 FP8 MoE requires ROCm AITER fused MoE support."
413
+ )
414
+
415
+ def create_weights(
416
+ self,
417
+ layer: torch.nn.Module,
418
+ num_experts: int,
419
+ hidden_size: int,
420
+ intermediate_size_per_partition: int,
421
+ params_dtype: torch.dtype,
422
+ **extra_weight_attrs,
423
+ ):
424
+ params_dtype = torch.uint32
425
+ w13_weight = torch.nn.Parameter(
426
+ torch.empty(
427
+ num_experts,
428
+ 2 * intermediate_size_per_partition,
429
+ hidden_size // 8, # INT32 packing for W4
430
+ dtype=params_dtype,
431
+ ),
432
+ requires_grad=False,
433
+ )
434
+ w2_weight = torch.nn.Parameter(
435
+ torch.empty(
436
+ num_experts,
437
+ hidden_size,
438
+ intermediate_size_per_partition // 8, # INT32 packing for W4
439
+ dtype=params_dtype,
440
+ ),
441
+ requires_grad=False,
442
+ )
443
+ layer.register_parameter("w13_weight", w13_weight)
444
+ layer.register_parameter("w2_weight", w2_weight)
445
+ set_weight_attrs(w13_weight, extra_weight_attrs)
446
+ set_weight_attrs(w2_weight, extra_weight_attrs)
447
+
448
+ # Per-tensor fp8 weight scales
449
+ w13_weight_scale = torch.nn.Parameter(
450
+ torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False
451
+ )
452
+ w2_weight_scale = torch.nn.Parameter(
453
+ torch.ones(num_experts, dtype=torch.float32), requires_grad=False
454
+ )
455
+ layer.register_parameter("w13_weight_scale", w13_weight_scale)
456
+ layer.register_parameter("w2_weight_scale", w2_weight_scale)
457
+ extra_weight_attrs.update(
458
+ {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
459
+ )
460
+ set_weight_attrs(w13_weight_scale, extra_weight_attrs)
461
+ set_weight_attrs(w2_weight_scale, extra_weight_attrs)
462
+
463
+ # Per-channel int4 weight scales
464
+ w13_weight_scale_2 = torch.nn.Parameter(
465
+ torch.ones(
466
+ num_experts,
467
+ 2 * intermediate_size_per_partition,
468
+ dtype=torch.float32,
469
+ ),
470
+ requires_grad=False,
471
+ )
472
+ w2_weight_scale_2 = torch.nn.Parameter(
473
+ torch.ones(num_experts, hidden_size, dtype=torch.float32),
474
+ requires_grad=False,
475
+ )
476
+ layer.register_parameter("w13_weight_scale_2", w13_weight_scale_2)
477
+ layer.register_parameter("w2_weight_scale_2", w2_weight_scale_2)
478
+ extra_weight_attrs.update(
479
+ {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
480
+ )
481
+ set_weight_attrs(w13_weight_scale_2, extra_weight_attrs)
482
+ set_weight_attrs(w2_weight_scale_2, extra_weight_attrs)
483
+
484
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
485
+ shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
486
+ layer.w13_weight.data, layer.w2_weight.data
487
+ )
488
+ layer.w13_weight = torch.nn.Parameter(shuffled_w13, requires_grad=False)
489
+ layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False)
490
+
491
+ # INT4-FP8 : offset INT4 w13_weight_scale1 to single w13_weight_scale
492
+ # Fp8 moe kernel needs single fp8 w13_weight_scale for w13 per expert.
493
+ # We won't do requant each expert's fp8 weight (not direct available),
494
+ # instead we adjust half of INT4 w13_weight_scale1 numbers
495
+ shard_size = layer.intermediate_size_per_partition
496
+ max_w13_scales = layer.w13_weight_scale.max(dim=1).values
497
+ assert torch.all(max_w13_scales != 0), "fp8 weight scale cannot be zero."
498
+ for expert_id in range(layer.local_num_experts):
499
+ start = 0
500
+ max_w13_scale_fp8 = max_w13_scales[expert_id]
501
+ for shard_id in range(2):
502
+ if layer.w13_weight_scale[expert_id][shard_id] != max_w13_scale_fp8:
503
+ int4_rescale = (
504
+ layer.w13_weight_scale[expert_id][shard_id] / max_w13_scale_fp8
505
+ )
506
+ layer.w13_weight_scale_2[expert_id][start : start + shard_size] *= (
507
+ int4_rescale
508
+ )
509
+ start += shard_size
510
+
511
+ layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales, requires_grad=False)
512
+
513
+ # special hack to asm_moe, which takes (weight_scale1 * weight_scale) as post
514
+ # GEMM scaling optimal design - shall apply per-column weight_scale1 before
515
+ # GEMM, and weight_scale post
516
+ for expert_id in range(layer.local_num_experts):
517
+ layer.w13_weight_scale_2[expert_id] *= max_w13_scales[expert_id]
518
+ layer.w2_weight_scale_2[expert_id] *= layer.w2_weight_scale[expert_id]
519
+
520
+ def get_fused_moe_quant_config(self, layer):
521
+ return fp8_w8a8_moe_quant_config(
522
+ w1_scale=layer.w13_weight_scale_2,
523
+ w2_scale=layer.w2_weight_scale_2,
524
+ per_out_ch_quant=True,
525
+ )
526
+
527
+ def apply(
528
+ self,
529
+ layer: FusedMoE,
530
+ x: torch.Tensor,
531
+ router_logits: torch.Tensor,
532
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
533
+ topk_weights, topk_ids, _ = layer.select_experts(
534
+ hidden_states=x,
535
+ router_logits=router_logits,
536
+ )
537
+
538
+ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
539
+ rocm_aiter_fused_experts,
540
+ )
541
+
542
+ return rocm_aiter_fused_experts(
543
+ hidden_states=x,
544
+ w1=layer.w13_weight,
545
+ w2=layer.w2_weight,
546
+ topk_weights=topk_weights,
547
+ topk_ids=topk_ids,
548
+ activation=layer.activation,
549
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
550
+ quant_config=self.moe_quant_config,
551
+ expert_map=layer.expert_map,
552
+ )
553
+
554
+
555
class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
556
def __init__(
557
self,