@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 827268e98d
Fork 0

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 5, 09:45 UTC

[Quantization] Support Quark W8A8 INT8 MoE inference (#36320)

Signed-off-by: kangletian <Letian.Kang@amd.com>

PikaPikachu committed Apr 10, 2026 at 01:24 UTC 827268e98d92761e9189c15baec6a452cf3ac945
4 files changed +360 -2
tests/quantization/test_quark.py
+31
@@ -22,6 +22,9 @@ from vllm.model_executor.layers.quantization.quark.quark import ( # noqa: E501
22 QuarkW8A8Fp8,
23 QuarkW8A8Int8,
24 )
25 +from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E501
26 + QuarkW8A8Int8MoEMethod,
27 +)
28 from vllm.platforms import current_platform
29
30 from .reference_mxfp4 import dq_mxfp4_torch, qdq_mxfp4_torch
@@ -126,6 +129,34 @@ def test_quark_int8_w_per_tensor_a_per_tensor(vllm_runner, tp):
129 assert output
130
131
132 +@pytest.mark.parametrize("tp", [1])
133 +def test_quark_int8_w8a8_moe(vllm_runner, tp):
134 + """Test W8A8 INT8 MoE quantization with a tiny Qwen3 MoE model."""
135 + model_path = "nameistoken/tiny-qwen3-moe-w8a8-int8-quark"
136 + with vllm_runner(
137 + model_path,
138 + enforce_eager=True,
139 + tensor_parallel_size=tp,
140 + gpu_memory_utilization=0.1,
141 + ) as llm:
142 +
143 + def check_model(model):
144 + layer = model.model.layers[0]
145 + # MoE experts should use QuarkW8A8Int8MoEMethod
146 + moe = layer.mlp.experts
147 + assert isinstance(moe.quant_method, QuarkW8A8Int8MoEMethod), (
148 + f"Expected QuarkW8A8Int8MoEMethod, got {type(moe.quant_method)}"
149 + )
150 + # Non-MoE linear layers should use QuarkW8A8Int8
151 + qkv_proj = layer.self_attn.qkv_proj
152 + assert isinstance(qkv_proj.scheme, QuarkW8A8Int8)
153 +
154 + llm.apply_model(check_model)
155 +
156 + output = llm.generate_greedy("Hello", max_tokens=4)
157 + assert output
158 +
159 +
160 def test_quark_fp8_parity(vllm_runner):
161 quark_model_id = "amd-quark/llama-tiny-fp8-quark-quant-method"
162 fp8_model_id = "amd-quark/llama-tiny-fp8-quant-method"
vllm/model_executor/layers/fused_moe/utils.py
+9 -2
@@ -163,8 +163,15 @@ def _int8_quantize(
163 # activations apply per-token quantization. Otherwise, assume
164 # activation tensor-wise fp8/int8 quantization, dynamic or static
165 if block_shape is None:
166 - assert per_act_token, "int8 quantization only supports block or channel-wise"
167 - A, A_scale = per_token_quant_int8(A)
166 + if per_act_token:
167 + A, A_scale = per_token_quant_int8(A)
168 + elif A_scale is not None:
169 + # Static per-tensor: use the optimized CUDA kernel
170 + A, A_scale, _ = ops.scaled_int8_quant(A, scale=A_scale)
171 + elif A_scale is None:
172 + # Dynamic per-tensor: compute scale then quantize via kernel
173 + A_scale = torch.clamp(A.abs().max() / 127.0, min=1e-10)
174 + A, A_scale, _ = ops.scaled_int8_quant(A, scale=A_scale)
175 else:
176 assert not per_act_token
177 assert len(block_shape) == 2
vllm/model_executor/layers/quantization/quark/quark.py
+38
@@ -389,6 +389,37 @@ class QuarkConfig(QuantizationConfig):
389
390 return is_weight_mxfp4 and is_input_fp8
391
392 + def _is_dynamic_per_token_w8a8(
393 + self,
394 + weight_quant: dict[str, Any] | None,
395 + input_quant: dict[str, Any] | None,
396 + ) -> bool:
397 + """Detect W8A8 INT8 with per-tensor or per-channel
398 + weights and dynamic per-token input."""
399 + if weight_quant is None or input_quant is None:
400 + return False
401 +
402 + is_int8_dtype = (
403 + weight_quant.get("dtype") == "int8" and input_quant.get("dtype") == "int8"
404 + )
405 +
406 + is_valid_weight_scheme = weight_quant.get("qscheme") in [
407 + "per_tensor",
408 + "per_channel",
409 + ]
410 + is_per_token_input = input_quant.get("qscheme") == "per_channel"
411 +
412 + is_dynamic_input = input_quant.get("is_dynamic") is True
413 + is_weight_symmetric = weight_quant.get("symmetric") is True
414 +
415 + return (
416 + is_int8_dtype
417 + and is_valid_weight_scheme
418 + and is_per_token_input
419 + and is_dynamic_input
420 + and is_weight_symmetric
421 + )
422 +
423 def _is_w_ocp_mx_a_x(
424 self, weight_quant: dict[str, Any] | None, input_quant: dict[str, Any] | None
425 ) -> bool:
@@ -556,6 +587,13 @@ class QuarkConfig(QuantizationConfig):
587 )
588 if is_w4a8_supported:
589 return QuarkW4A8_MXFP4_FP8(weight_config, input_config)
590 + elif self._is_dynamic_per_token_w8a8(weight_config, input_config):
591 + weight_qscheme = cast(str, weight_config.get("qscheme"))
592 + return QuarkW8A8Int8(
593 + qscheme=weight_qscheme,
594 + is_static_input_scheme=False,
595 + input_symmetric=input_config.get("symmetric"),
596 + )
597 elif self._is_w_ocp_mx_a_x(weight_config, input_config):
598 return QuarkOCP_MX(
599 weight_config, input_config, dynamic_mxfp4_quant=dynamic_mxfp4_quant
vllm/model_executor/layers/quantization/quark/quark_moe.py
+282
@@ -109,6 +109,12 @@ class QuarkMoEMethod(FusedMoEMethodBase):
109 return QuarkOCP_MX_MoEMethod(
110 weight_config, input_config, module.moe_config
111 )
112 + elif quant_config._is_static_tensor_w8a8(
113 + weight_config, input_config
114 + ) or quant_config._is_dynamic_per_token_w8a8(weight_config, input_config):
115 + return QuarkW8A8Int8MoEMethod(
116 + weight_config, input_config, module.moe_config
117 + )
118 else:
119 raise RuntimeError("Unsupported FusedMoe scheme")
120
@@ -505,6 +511,282 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
511 )
512
513
514 +class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):
515 + """Quark W8A8 INT8 MoE method."""
516 +
517 + def __init__(
518 + self,
519 + weight_config: dict[str, Any],
520 + input_config: dict[str, Any],
521 + moe: FusedMoEConfig,
522 + ):
523 + super().__init__(moe)
524 + self.weight_quant = weight_config
525 + self.input_quant = input_config
526 + self.weight_qscheme = self.weight_quant.get("qscheme", "per_tensor")
527 + self.static_input_scales = not self.input_quant.get("is_dynamic", False)
528 +
529 + def create_weights(
530 + self,
531 + layer: torch.nn.Module,
532 + num_experts: int,
533 + hidden_size: int,
534 + intermediate_size_per_partition: int,
535 + params_dtype: torch.dtype,
536 + **extra_weight_attrs,
537 + ):
538 + layer.num_experts = num_experts
539 + layer.orig_dtype = params_dtype
540 + layer.weight_block_size = None
541 + params_dtype = torch.int8
542 +
543 + # WEIGHTS
544 + w13_weight = torch.nn.Parameter(
545 + torch.empty(
546 + num_experts,
547 + 2 * intermediate_size_per_partition,
548 + hidden_size,
549 + dtype=params_dtype,
550 + ),
551 + requires_grad=False,
552 + )
553 + layer.register_parameter("w13_weight", w13_weight)
554 + set_weight_attrs(w13_weight, extra_weight_attrs)
555 +
556 + w2_weight = torch.nn.Parameter(
557 + torch.empty(
558 + num_experts,
559 + hidden_size,
560 + intermediate_size_per_partition,
561 + dtype=params_dtype,
562 + ),
563 + requires_grad=False,
564 + )
565 + layer.register_parameter("w2_weight", w2_weight)
566 + set_weight_attrs(w2_weight, extra_weight_attrs)
567 +
568 + # WEIGHT_SCALES
569 + if self.weight_qscheme == "per_channel":
570 + w13_weight_scale = torch.nn.Parameter(
571 + torch.ones(
572 + num_experts,
573 + 2 * intermediate_size_per_partition,
574 + dtype=torch.float32,
575 + ),
576 + requires_grad=False,
577 + )
578 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
579 + w2_weight_scale = torch.nn.Parameter(
580 + torch.ones(num_experts, hidden_size, dtype=torch.float32),
581 + requires_grad=False,
582 + )
583 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
584 + extra_weight_attrs.update(
585 + {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
586 + )
587 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)
588 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)
589 + else:
590 + # per-tensor: one scalar per expert
591 + w13_weight_scale = torch.nn.Parameter(
592 + torch.ones(num_experts, 2, dtype=torch.float32),
593 + requires_grad=False,
594 + )
595 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
596 + w2_weight_scale = torch.nn.Parameter(
597 + torch.ones(num_experts, dtype=torch.float32),
598 + requires_grad=False,
599 + )
600 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
601 + extra_weight_attrs.update(
602 + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
603 + )
604 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)
605 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)
606 +
607 + # INPUT_SCALES
608 + if self.static_input_scales:
609 + w13_input_scale = torch.nn.Parameter(
610 + torch.ones(num_experts, dtype=torch.float32),
611 + requires_grad=False,
612 + )
613 + layer.register_parameter("w13_input_scale", w13_input_scale)
614 + set_weight_attrs(w13_input_scale, extra_weight_attrs)
615 +
616 + w2_input_scale = torch.nn.Parameter(
617 + torch.ones(num_experts, dtype=torch.float32),
618 + requires_grad=False,
619 + )
620 + layer.register_parameter("w2_input_scale", w2_input_scale)
621 + set_weight_attrs(w2_input_scale, extra_weight_attrs)
622 + else:
623 + layer.w13_input_scale = None
624 + layer.w2_input_scale = None
625 +
626 + # ZERO POINTS (loaded but discarded after loading; kernel uses symmetric)
627 + w13_input_zero_point = torch.nn.Parameter(
628 + torch.zeros(num_experts, 2, dtype=torch.int8),
629 + requires_grad=False,
630 + )
631 + layer.register_parameter("w13_input_zero_point", w13_input_zero_point)
632 + set_weight_attrs(w13_input_zero_point, extra_weight_attrs)
633 +
634 + w2_input_zero_point = torch.nn.Parameter(
635 + torch.zeros(num_experts, dtype=torch.int8),
636 + requires_grad=False,
637 + )
638 + layer.register_parameter("w2_input_zero_point", w2_input_zero_point)
639 + set_weight_attrs(w2_input_zero_point, extra_weight_attrs)
640 +
641 + if self.weight_qscheme == "per_channel":
642 + w13_weight_zero_point = torch.nn.Parameter(
643 + torch.zeros(
644 + num_experts,
645 + 2 * intermediate_size_per_partition,
646 + dtype=torch.int8,
647 + ),
648 + requires_grad=False,
649 + )
650 + w2_weight_zero_point = torch.nn.Parameter(
651 + torch.zeros(num_experts, hidden_size, dtype=torch.int8),
652 + requires_grad=False,
653 + )
654 + else:
655 + w13_weight_zero_point = torch.nn.Parameter(
656 + torch.zeros(num_experts, 2, dtype=torch.int8),
657 + requires_grad=False,
658 + )
659 + w2_weight_zero_point = torch.nn.Parameter(
660 + torch.zeros(num_experts, dtype=torch.int8),
661 + requires_grad=False,
662 + )
663 + layer.register_parameter("w13_weight_zero_point", w13_weight_zero_point)
664 + set_weight_attrs(w13_weight_zero_point, extra_weight_attrs)
665 + layer.register_parameter("w2_weight_zero_point", w2_weight_zero_point)
666 + set_weight_attrs(w2_weight_zero_point, extra_weight_attrs)
667 +
668 + # BIAS
669 + if self.has_bias:
670 + w13_bias = torch.nn.Parameter(
671 + torch.zeros(
672 + num_experts,
673 + 2 * intermediate_size_per_partition,
674 + dtype=torch.float32,
675 + ),
676 + requires_grad=False,
677 + )
678 + layer.register_parameter("w13_bias", w13_bias)
679 + set_weight_attrs(w13_bias, extra_weight_attrs)
680 + w2_bias = torch.nn.Parameter(
681 + torch.zeros(num_experts, hidden_size, dtype=torch.float32),
682 + requires_grad=False,
683 + )
684 + layer.register_parameter("w2_bias", w2_bias)
685 + set_weight_attrs(w2_bias, extra_weight_attrs)
686 + else:
687 + layer.w13_bias, layer.w2_bias = None, None
688 +
689 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
690 + # Discard zero points (INT8 fused MoE kernel uses symmetric quant)
691 + for attr in (
692 + "w13_input_zero_point",
693 + "w2_input_zero_point",
694 + "w13_weight_zero_point",
695 + "w2_weight_zero_point",
696 + ):
697 + if hasattr(layer, attr):
698 + delattr(layer, attr)
699 +
700 + # For static input scales, collapse per-expert scales to single max
701 + if self.static_input_scales:
702 + if layer.w13_input_scale is None or layer.w2_input_scale is None:
703 + raise ValueError(
704 + "QuantConfig has static quantization, but found "
705 + "activation scales are None."
706 + )
707 + if not all_close_1d(layer.w13_input_scale) or not all_close_1d(
708 + layer.w2_input_scale
709 + ):
710 + logger.warning_once(
711 + "Found input_scales that are not equal for "
712 + "INT8 MoE layer. Using the maximum across experts "
713 + "for each layer."
714 + )
715 + layer.w13_input_scale = torch.nn.Parameter(
716 + layer.w13_input_scale.max(), requires_grad=False
717 + )
718 + layer.w2_input_scale = torch.nn.Parameter(
719 + layer.w2_input_scale.max(), requires_grad=False
720 + )
721 +
722 + # For per-tensor weights, merge w1/w3 scales into single per-expert
723 + if self.weight_qscheme == "per_tensor":
724 + assert layer.w13_weight_scale is not None
725 + shard_size = layer.intermediate_size_per_partition
726 + max_w13_scales = layer.w13_weight_scale.max(dim=1).values
727 +
728 + for expert_id in range(layer.local_num_experts):
729 + start = 0
730 + for shard_id in range(2):
731 + dq_weight = per_tensor_dequantize(
732 + layer.w13_weight[expert_id][start : start + shard_size, :],
733 + layer.w13_weight_scale[expert_id][shard_id],
734 + )
735 + layer.w13_weight[expert_id][start : start + shard_size, :], _, _ = (
736 + ops.scaled_int8_quant(
737 + dq_weight,
738 + scale=max_w13_scales[expert_id],
739 + )
740 + )
741 + start += shard_size
742 +
743 + layer.w13_weight_scale = torch.nn.Parameter(
744 + max_w13_scales, requires_grad=False
745 + )
746 +
747 + def get_fused_moe_quant_config(
748 + self, layer: torch.nn.Module
749 + ) -> FusedMoEQuantConfig | None:
750 + is_dynamic = not self.static_input_scales
751 + is_per_channel = self.weight_qscheme == "per_channel"
752 + return FusedMoEQuantConfig.make(
753 + torch.int8,
754 + w1_scale=layer.w13_weight_scale,
755 + w2_scale=layer.w2_weight_scale,
756 + a1_scale=layer.w13_input_scale,
757 + a2_scale=layer.w2_input_scale,
758 + w1_bias=getattr(layer, "w13_bias", None),
759 + w2_bias=getattr(layer, "w2_bias", None),
760 + per_act_token_quant=is_dynamic,
761 + per_out_ch_quant=is_per_channel,
762 + block_shape=None,
763 + )
764 +
765 + def apply(
766 + self,
767 + layer: FusedMoE,
768 + x: torch.Tensor,
769 + topk_weights: torch.Tensor,
770 + topk_ids: torch.Tensor,
771 + shared_experts_input: torch.Tensor | None,
772 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
773 + from vllm.model_executor.layers.fused_moe import fused_experts
774 +
775 + return fused_experts(
776 + hidden_states=x,
777 + w1=layer.w13_weight,
778 + w2=layer.w2_weight,
779 + topk_weights=topk_weights,
780 + topk_ids=topk_ids,
781 + inplace=not self.moe.disable_inplace,
782 + activation=layer.activation,
783 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
784 + global_num_experts=layer.global_num_experts,
785 + expert_map=layer.expert_map,
786 + quant_config=self.moe_quant_config,
787 + )
788 +
789 +
790 class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):
791 def __init__(
792 self,