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
gptq marlin quantization support for fused moe with lora (#30254)
Signed-off-by: Bhanu068 <voutharoja.bhanu06@gmail.com>
Bhanu Prakash Voutharoja committed
Dec 12, 2025 at 13:27 UTC
6a6fc41c799916521b1fa2914f72e108352e1bf6
3 files changed
+146
-2
csrc/moe/marlin_moe_wna16/ops.cu
+1
-1
@@ -860,4 +860,4 @@ torch::Tensor moe_wna16_marlin_gemm(
860
861
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
862
m.impl("moe_wna16_marlin_gemm", &moe_wna16_marlin_gemm);
863
-}
863
+}
\ No newline at end of file
vllm/model_executor/layers/fused_moe/config.py
+36
@@ -543,6 +543,42 @@ def int8_w8a8_moe_quant_config(
543
)
544
545
546
+def gptq_marlin_moe_quant_config(
547
+ w1_scale: torch.Tensor,
548
+ w2_scale: torch.Tensor,
549
+ weight_bits: int,
550
+ group_size: int,
551
+ w1_zp: torch.Tensor | None = None,
552
+ w2_zp: torch.Tensor | None = None,
553
+ w1_bias: torch.Tensor | None = None,
554
+ w2_bias: torch.Tensor | None = None,
555
+):
556
+ """
557
+ Construct a quant config for gptq marlin quantization.
558
+ """
559
+ from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
560
+
561
+ w_shape = None if group_size == -1 else GroupShape(row=1, col=group_size)
562
+
563
+ # Activations are NOT quantized for GPTQ (fp16/bf16)
564
+ a_shape = w_shape # Same as weight shape for alignment
565
+
566
+ # Determine weight dtype
567
+ if weight_bits == 4:
568
+ weight_dtype = "int4"
569
+ elif weight_bits == 8:
570
+ weight_dtype = torch.int8
571
+ else:
572
+ raise ValueError(f"Unsupported weight_bits: {weight_bits}")
573
+
574
+ return FusedMoEQuantConfig(
575
+ _a1=FusedMoEQuantDesc(dtype=None, shape=a_shape),
576
+ _a2=FusedMoEQuantDesc(dtype=None, shape=a_shape),
577
+ _w1=FusedMoEQuantDesc(weight_dtype, w_shape, w1_scale, None, w1_zp, w1_bias),
578
+ _w2=FusedMoEQuantDesc(weight_dtype, w_shape, w2_scale, None, w2_zp, w2_bias),
579
+ )
580
+
581
+
582
def mxfp4_w4a16_moe_quant_config(
583
w1_scale: Union[torch.Tensor, "PrecisionConfig"],
584
w2_scale: Union[torch.Tensor, "PrecisionConfig"],
vllm/model_executor/layers/quantization/gptq_marlin.py
+109
-1
@@ -732,6 +732,14 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
732
is_a_8bit=is_a_8bit,
733
)
734
replace_parameter(layer, "w2_qweight", marlin_w2_qweight)
735
+
736
+ # The modular kernel expects w13_weight and w2_weight,
737
+ # but GPTQ uses w13_qweight and w2_qweight
738
+ # Alias for modular kernel
739
+ layer.w13_weight = layer.w13_qweight
740
+ # Alias for modular kernel
741
+ layer.w2_weight = layer.w2_qweight
742
+
743
# Repack scales
744
marlin_w13_scales = marlin_moe_permute_scales(
745
s=layer.w13_scales,
@@ -782,7 +790,107 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
790
def get_fused_moe_quant_config(
791
self, layer: torch.nn.Module
792
) -> FusedMoEQuantConfig | None:
785
- return None
793
+ from vllm.model_executor.layers.fused_moe.config import (
794
+ gptq_marlin_moe_quant_config,
795
+ )
796
+
797
+ return gptq_marlin_moe_quant_config(
798
+ w1_scale=layer.w13_scales,
799
+ w2_scale=layer.w2_scales,
800
+ weight_bits=self.quant_config.weight_bits,
801
+ group_size=self.quant_config.group_size,
802
+ w1_zp=getattr(layer, "w13_qzeros", None)
803
+ if not self.quant_config.is_sym
804
+ else None,
805
+ w2_zp=getattr(layer, "w2_qzeros", None)
806
+ if not self.quant_config.is_sym
807
+ else None,
808
+ w1_bias=getattr(layer, "w13_bias", None),
809
+ w2_bias=getattr(layer, "w2_bias", None),
810
+ )
811
+
812
+ def select_gemm_impl(
813
+ self,
814
+ prepare_finalize,
815
+ layer: torch.nn.Module,
816
+ ):
817
+ """
818
+ Select the GEMM implementation for GPTQ-Marlin MoE.
819
+
820
+ Returns MarlinExperts configured for GPTQ quantization.
821
+ This is ONLY used when LoRA is enabled.
822
+ Without LoRA, GPTQ uses its own apply() method.
823
+ """
824
+ # Only use modular kernels when LoRA is enabled
825
+ # Without LoRA, GPTQ's own apply() method works fine and is more efficient
826
+ if not self.moe.is_lora_enabled:
827
+ raise NotImplementedError(
828
+ "GPTQ-Marlin uses its own apply() method when LoRA is not enabled. "
829
+ "Modular kernels are only used for LoRA support."
830
+ )
831
+
832
+ # The modular marlin kernels do not support 8-bit weights.
833
+ if self.quant_config.weight_bits == 8:
834
+ raise NotImplementedError(
835
+ "GPTQ-Marlin kernel does not support 8-bit weights."
836
+ )
837
+
838
+ from vllm.model_executor.layers.fused_moe import modular_kernel as mk
839
+ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
840
+ BatchedMarlinExperts,
841
+ MarlinExperts,
842
+ )
843
+
844
+ # Ensure quant config is initialized
845
+ assert self.moe_quant_config is not None, (
846
+ "moe_quant_config must be initialized before select_gemm_impl"
847
+ )
848
+
849
+ w13_g_idx = (
850
+ getattr(layer, "w13_g_idx", None) if self.quant_config.desc_act else None
851
+ )
852
+ w2_g_idx = (
853
+ getattr(layer, "w2_g_idx", None) if self.quant_config.desc_act else None
854
+ )
855
+ w13_g_idx_sort_indices = (
856
+ getattr(layer, "w13_g_idx_sort_indices", None)
857
+ if self.quant_config.desc_act
858
+ else None
859
+ )
860
+ w2_g_idx_sort_indices = (
861
+ getattr(layer, "w2_g_idx_sort_indices", None)
862
+ if self.quant_config.desc_act
863
+ else None
864
+ )
865
+
866
+ # Check if using batched expert format (for Expert Parallelism)
867
+ if (
868
+ prepare_finalize.activation_format
869
+ == mk.FusedMoEActivationFormat.BatchedExperts
870
+ ):
871
+ # For batched format, use BatchedMarlinExperts
872
+ max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()
873
+ assert max_num_tokens_per_rank is not None
874
+ return BatchedMarlinExperts(
875
+ max_num_tokens=max_num_tokens_per_rank,
876
+ num_dispatchers=prepare_finalize.num_dispatchers(),
877
+ quant_config=self.moe_quant_config,
878
+ w13_g_idx=w13_g_idx,
879
+ w2_g_idx=w2_g_idx,
880
+ w13_g_idx_sort_indices=w13_g_idx_sort_indices,
881
+ w2_g_idx_sort_indices=w2_g_idx_sort_indices,
882
+ is_k_full=self.is_k_full,
883
+ )
884
+ else:
885
+ # Standard Marlin experts for GPTQ
886
+ return MarlinExperts(
887
+ quant_config=self.moe_quant_config,
888
+ w13_g_idx=w13_g_idx,
889
+ w2_g_idx=w2_g_idx,
890
+ w13_g_idx_sort_indices=w13_g_idx_sort_indices,
891
+ w2_g_idx_sort_indices=w2_g_idx_sort_indices,
892
+ is_k_full=self.is_k_full,
893
+ )
894
895
def apply(
896
self,