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 14, 22:35 UTC
feat: Add Support GPTQ Quantization MOE on ROCM vllm serve (#21733)
JartX committed
Aug 2, 2025 at 03:12 UTC
3654847db5a9b9a0955f8416292d94fa1c827f77
2 files changed
+21
-5
vllm/model_executor/layers/fused_moe/fused_moe.py
+2
-2
@@ -761,8 +761,8 @@ def get_moe_wna16_block_config(config: dict[str,
761
762
def should_moe_wna16_use_cuda(num_valid_tokens: int, group_size: int,
763
num_experts: int, bit: int):
764
- return bit == 4 and group_size in [32, 64, 128] and \
765
- num_valid_tokens / num_experts <= 6
764
+ return current_platform.is_cuda() and bit == 4 and \
765
+ group_size in [32, 64, 128] and num_valid_tokens / num_experts <= 6
766
767
768
def get_default_config(
vllm/model_executor/layers/quantization/gptq.py
+19
-3
@@ -10,10 +10,11 @@ import torch
10
from torch.nn.parameter import Parameter
11
12
from vllm import _custom_ops as ops
13
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE
14
from vllm.model_executor.layers.linear import LinearMethodBase
15
from vllm.model_executor.layers.quantization import QuantizationMethods
16
from vllm.model_executor.layers.quantization.base_config import (
16
- QuantizationConfig)
17
+ QuantizationConfig, QuantizeMethodBase)
18
from vllm.model_executor.layers.quantization.utils.gptq_utils import (
19
get_linear_quant_method)
20
from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
@@ -110,8 +111,23 @@ class GPTQConfig(QuantizationConfig):
111
return cls(weight_bits, group_size, desc_act, lm_head_quantized,
112
dynamic)
113
113
- def get_quant_method(self, layer: torch.nn.Module,
114
- prefix: str) -> Optional["GPTQLinearMethod"]:
114
+ def get_quant_method(
115
+ self, layer: torch.nn.Module, prefix: str
116
+ ) -> Optional[Union["GPTQLinearMethod", "QuantizeMethodBase"]]:
117
+ if isinstance(layer, FusedMoE):
118
+ # GPTQ MoE support: fall back to MoeWNA16 for broad compatibility
119
+ from .moe_wna16 import MoeWNA16Config
120
+
121
+ config = {
122
+ "quant_method": "gptq",
123
+ "bits": self.weight_bits,
124
+ "group_size": self.group_size,
125
+ "sym": True, # GPTQ typically uses symmetric quantization
126
+ "lm_head": False,
127
+ }
128
+ return MoeWNA16Config.from_config(config).get_quant_method(
129
+ layer, prefix)
130
+
131
return get_linear_quant_method(self, layer, prefix, GPTQLinearMethod)
132
133