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 15, 06:00 UTC
[Quantization] Improve AWQ logic (#19431)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Jun 12, 2025 at 19:02 UTC
73e2e0118f208aa9fd1cb72925e2ca58a37e7210
1 file changed
+38
-4
vllm/model_executor/layers/quantization/awq.py
+38
-4
@@ -1,19 +1,23 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import Any, Optional
4
+from typing import Any, Optional, Union
5
6
import torch
7
8
from vllm import _custom_ops as ops
9
+from vllm.logger import init_logger
10
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE
11
from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
12
UnquantizedLinearMethod)
13
from vllm.model_executor.layers.quantization import QuantizationMethods
14
from vllm.model_executor.layers.quantization.base_config import (
13
- QuantizationConfig)
15
+ QuantizationConfig, QuantizeMethodBase)
16
from vllm.model_executor.parameter import (GroupQuantScaleParameter,
17
PackedvLLMParameter)
18
19
+logger = init_logger(__name__)
20
+
21
22
class AWQConfig(QuantizationConfig):
23
"""Config class for AWQ.
@@ -74,12 +78,42 @@ class AWQConfig(QuantizationConfig):
78
config, ["modules_to_not_convert"], None)
79
return cls(weight_bits, group_size, zero_point, modules_to_not_convert)
80
77
- def get_quant_method(self, layer: torch.nn.Module,
78
- prefix: str) -> Optional["LinearMethodBase"]:
81
+ def get_quant_method(
82
+ self, layer: torch.nn.Module, prefix: str
83
+ ) -> Optional[Union["LinearMethodBase", "QuantizeMethodBase"]]:
84
if isinstance(layer, LinearBase):
85
if is_layer_skipped_awq(prefix, self.modules_to_not_convert):
86
return UnquantizedLinearMethod()
87
return AWQLinearMethod(self)
88
+ elif isinstance(layer, FusedMoE):
89
+ # Lazy import to avoid circular import.
90
+ from .awq_marlin import AWQMarlinConfig, AWQMoEMethod
91
+ from .moe_wna16 import MoeWNA16Config
92
+ from .utils.marlin_utils import check_moe_marlin_supports_layer
93
+ if not check_moe_marlin_supports_layer(layer, self.group_size):
94
+ logger.warning_once(
95
+ f"Layer '{prefix}' is not supported by AWQMoeMarlin. "
96
+ "Falling back to Moe WNA16 kernels.")
97
+ config = {
98
+ "quant_method": "awq",
99
+ "bits": self.weight_bits,
100
+ "group_size": self.group_size,
101
+ "zero_point": self.zero_point,
102
+ "lm_head": False,
103
+ }
104
+ return MoeWNA16Config.from_config(config).get_quant_method(
105
+ layer, prefix)
106
+ marlin_compatible_config_dict = {
107
+ "quant_method": "awq",
108
+ "bits": self.weight_bits,
109
+ "group_size": self.group_size,
110
+ "zero_point": self.zero_point,
111
+ "lm_head": False,
112
+ "modules_to_not_convert": self.modules_to_not_convert,
113
+ }
114
+ awq_marlin_config = AWQMarlinConfig.from_config(
115
+ marlin_compatible_config_dict)
116
+ return AWQMoEMethod(awq_marlin_config)
117
return None
118
119