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, 15:00 UTC
[Feature] AWQ marlin quantization support for fused moe with lora (#30442)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
汪志鹏 committed
Dec 12, 2025 at 02:03 UTC
0e71eaa6447d99e76de8e03213ec22bc1d3b07df
2 files changed
+130
-1
vllm/model_executor/layers/fused_moe/config.py
+36
@@ -700,6 +700,42 @@ def int4_w4afp8_moe_quant_config(
700
)
701
702
703
+def awq_marlin_moe_quant_config(
704
+ w1_scale: torch.Tensor,
705
+ w2_scale: torch.Tensor,
706
+ w1_zp: torch.Tensor | None,
707
+ w2_zp: torch.Tensor | None,
708
+ weight_bits: int,
709
+ group_size: int,
710
+ w1_bias: torch.Tensor | None = None,
711
+ w2_bias: torch.Tensor | None = None,
712
+) -> FusedMoEQuantConfig:
713
+ """
714
+ Construct a quant config for awq marlin quantization.
715
+ """
716
+ from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
717
+
718
+ w_shape = None if group_size == -1 else GroupShape(row=1, col=group_size)
719
+
720
+ # Activations are NOT quantized for AWQ (fp16/bf16)
721
+ a_shape = w_shape # Same as weight shape for alignment
722
+
723
+ # Determine weight dtype
724
+ if weight_bits == 4:
725
+ weight_dtype = "int4"
726
+ elif weight_bits == 8:
727
+ weight_dtype = torch.int8
728
+ else:
729
+ raise ValueError(f"Unsupported weight_bits: {weight_bits}")
730
+
731
+ return FusedMoEQuantConfig(
732
+ _a1=FusedMoEQuantDesc(dtype=None, shape=a_shape),
733
+ _a2=FusedMoEQuantDesc(dtype=None, shape=a_shape),
734
+ _w1=FusedMoEQuantDesc(weight_dtype, w_shape, w1_scale, None, w1_zp, w1_bias),
735
+ _w2=FusedMoEQuantDesc(weight_dtype, w_shape, w2_scale, None, w2_zp, w2_bias),
736
+ )
737
+
738
+
739
def biased_moe_quant_config(
740
w1_bias: torch.Tensor | None,
741
w2_bias: torch.Tensor | None,
vllm/model_executor/layers/quantization/awq_marlin.py
+94
-1
@@ -470,6 +470,11 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
470
}
471
)
472
473
+ intermediate_size_full = extra_weight_attrs.pop(
474
+ "intermediate_size_full", intermediate_size_per_partition
475
+ )
476
+ self.is_k_full = intermediate_size_per_partition == intermediate_size_full
477
+
478
w13_qweight = Parameter(
479
torch.empty(
480
num_experts,
@@ -597,6 +602,13 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
602
)
603
replace_parameter(layer, "w2_qweight", marlin_w2_qweight)
604
605
+ # The modular kernel expects w13_weight and w2_weight,
606
+ # but AWQ uses w13_qweight and w2_qweight
607
+ # Alias for modular kernel
608
+ layer.w13_weight = layer.w13_qweight
609
+ # Alias for modular kernel
610
+ layer.w2_weight = layer.w2_qweight
611
+
612
# Why does this take the intermediate size for size_k?
613
marlin_w13_scales = marlin_moe_permute_scales(
614
s=layer.w13_scales,
@@ -661,7 +673,88 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
673
def get_fused_moe_quant_config(
674
self, layer: torch.nn.Module
675
) -> FusedMoEQuantConfig | None:
664
- return None
676
+ from vllm.model_executor.layers.fused_moe.config import (
677
+ awq_marlin_moe_quant_config,
678
+ )
679
+
680
+ return awq_marlin_moe_quant_config(
681
+ w1_scale=layer.w13_scales,
682
+ w2_scale=layer.w2_scales,
683
+ weight_bits=self.quant_config.weight_bits,
684
+ group_size=self.quant_config.group_size,
685
+ w1_zp=getattr(layer, "w13_qzeros", None)
686
+ if self.quant_config.zero_point
687
+ else None,
688
+ w2_zp=getattr(layer, "w2_qzeros", None)
689
+ if self.quant_config.zero_point
690
+ else None,
691
+ w1_bias=getattr(layer, "w13_bias", None),
692
+ w2_bias=getattr(layer, "w2_bias", None),
693
+ )
694
+
695
+ def select_gemm_impl(
696
+ self,
697
+ prepare_finalize,
698
+ layer: torch.nn.Module,
699
+ ):
700
+ """
701
+ Select the GEMM implementation for AWQ-Marlin MoE.
702
+ Returns MarlinExperts configured for AWQ quantization.
703
+ This is ONLY used when LoRA is enabled.
704
+ Without LoRA, AWQ uses its own apply() method.
705
+ """
706
+ # Only use modular kernels when LoRA is enabled
707
+ # Without LoRA, AWQ's own apply() method works fine and is more efficient
708
+ if not self.moe.is_lora_enabled:
709
+ raise NotImplementedError(
710
+ "AWQ-Marlin uses its own apply() method when LoRA is not enabled. "
711
+ "Modular kernels are only used for LoRA support."
712
+ )
713
+
714
+ from vllm.model_executor.layers.fused_moe import modular_kernel as mk
715
+ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
716
+ BatchedMarlinExperts,
717
+ MarlinExperts,
718
+ )
719
+
720
+ # Ensure quant config is initialized
721
+ assert self.moe_quant_config is not None, (
722
+ "moe_quant_config must be initialized before select_gemm_impl"
723
+ )
724
+
725
+ w13_g_idx = getattr(layer, "w13_g_idx", None)
726
+ w2_g_idx = getattr(layer, "w2_g_idx", None)
727
+ w13_g_idx_sort_indices = getattr(layer, "w13_g_idx_sort_indices", None)
728
+ w2_g_idx_sort_indices = getattr(layer, "w2_g_idx_sort_indices", None)
729
+
730
+ # Check if using batched expert format (for Expert Parallelism)
731
+ if (
732
+ prepare_finalize.activation_format
733
+ == mk.FusedMoEActivationFormat.BatchedExperts
734
+ ):
735
+ # For batched format, use BatchedMarlinExperts
736
+ max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()
737
+ assert max_num_tokens_per_rank is not None
738
+ return BatchedMarlinExperts(
739
+ max_num_tokens=max_num_tokens_per_rank,
740
+ num_dispatchers=prepare_finalize.num_dispatchers(),
741
+ quant_config=self.moe_quant_config,
742
+ w13_g_idx=w13_g_idx,
743
+ w2_g_idx=w2_g_idx,
744
+ w13_g_idx_sort_indices=w13_g_idx_sort_indices,
745
+ w2_g_idx_sort_indices=w2_g_idx_sort_indices,
746
+ is_k_full=self.is_k_full,
747
+ )
748
+ else:
749
+ # Standard Marlin experts for AWQ
750
+ return MarlinExperts(
751
+ quant_config=self.moe_quant_config,
752
+ w13_g_idx=w13_g_idx,
753
+ w2_g_idx=w2_g_idx,
754
+ w13_g_idx_sort_indices=w13_g_idx_sort_indices,
755
+ w2_g_idx_sort_indices=w2_g_idx_sort_indices,
756
+ is_k_full=self.is_k_full,
757
+ )
758
759
def apply(
760
self,