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 6, 05:45 UTC
Revert "Consolidate AWQ quantization into single awq_marlin.py file" (#37768)
Robert Shaw committed
Mar 21, 2026 at 17:20 UTC
5ad0446572dcec6bc5af2519ea4a8a54b315f37c
3 files changed
+279
-253
vllm/model_executor/layers/quantization/__init__.py
+2
-1
@@ -107,7 +107,8 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
107
# lazy import to avoid triggering `torch.compile` too early
108
from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
109
110
- from .awq_marlin import AWQConfig, AWQMarlinConfig
110
+ from .awq import AWQConfig
111
+ from .awq_marlin import AWQMarlinConfig
112
from .bitsandbytes import BitsAndBytesConfig
113
from .compressed_tensors.compressed_tensors import (
114
CompressedTensorsConfig,
vllm/model_executor/layers/quantization/awq.py
+275
-5
@@ -1,8 +1,278 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-# Backward compatibility: AWQConfig and AWQLinearMethod have been consolidated
4
-# into awq_marlin.py
5
-from vllm.model_executor.layers.quantization.awq_marlin import ( # noqa: F401
6
- AWQConfig,
7
- AWQLinearMethod,
3
+
4
+from typing import TYPE_CHECKING, Any, Union
5
+
6
+import torch
7
+from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
8
+
9
+from vllm import _custom_ops as ops
10
+from vllm.logger import init_logger
11
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE
12
+from vllm.model_executor.layers.linear import (
13
+ LinearBase,
14
+ LinearMethodBase,
15
+ UnquantizedLinearMethod,
16
)
17
+from vllm.model_executor.layers.quantization.base_config import (
18
+ QuantizationConfig,
19
+ QuantizeMethodBase,
20
+)
21
+from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
22
+from vllm.model_executor.parameter import GroupQuantScaleParameter, PackedvLLMParameter
23
+from vllm.transformers_utils.config import get_safetensors_params_metadata
24
+
25
+if TYPE_CHECKING:
26
+ from vllm.model_executor.layers.quantization import QuantizationMethods
27
+ from vllm.model_executor.models.utils import WeightsMapper
28
+
29
+logger = init_logger(__name__)
30
+
31
+
32
+class AWQConfig(QuantizationConfig):
33
+ """Config class for AWQ.
34
+
35
+ Reference: https://arxiv.org/abs/2306.00978
36
+ """
37
+
38
+ def __init__(
39
+ self,
40
+ weight_bits: int,
41
+ group_size: int,
42
+ zero_point: bool,
43
+ modules_to_not_convert: list[str] | None = None,
44
+ ) -> None:
45
+ super().__init__()
46
+ self.weight_bits = weight_bits
47
+ self.group_size = group_size
48
+ self.zero_point = zero_point
49
+ self.modules_to_not_convert = modules_to_not_convert or []
50
+
51
+ if self.weight_bits != 4:
52
+ raise ValueError(
53
+ "Currently, only 4-bit weight quantization is supported for "
54
+ f"AWQ, but got {self.weight_bits} bits."
55
+ )
56
+ self.pack_factor = 32 // self.weight_bits
57
+
58
+ def __repr__(self) -> str:
59
+ return (
60
+ f"AWQConfig(weight_bits={self.weight_bits}, "
61
+ f"group_size={self.group_size}, "
62
+ f"zero_point={self.zero_point}, "
63
+ f"modules_to_not_convert={self.modules_to_not_convert})"
64
+ )
65
+
66
+ def get_name(self) -> "QuantizationMethods":
67
+ return "awq"
68
+
69
+ def get_supported_act_dtypes(self) -> list[torch.dtype]:
70
+ return [torch.half]
71
+
72
+ @classmethod
73
+ def get_min_capability(cls) -> int:
74
+ # The AWQ kernel only supports Turing or newer GPUs.
75
+ return 75
76
+
77
+ @staticmethod
78
+ def get_config_filenames() -> list[str]:
79
+ return [
80
+ "quant_config.json", # E.g., casperhansen/vicuna-7b-v1.5-awq
81
+ # E.g., abhinavkulkarni/mosaicml-mpt-7b-instruct-w4-g128-awq
82
+ "quantize_config.json",
83
+ ]
84
+
85
+ @classmethod
86
+ def from_config(cls, config: dict[str, Any]) -> "AWQConfig":
87
+ weight_bits = cls.get_from_keys(config, ["w_bit", "bits"])
88
+ group_size = cls.get_from_keys(config, ["q_group_size", "group_size"])
89
+ zero_point = cls.get_from_keys(config, ["zero_point"])
90
+ modules_to_not_convert = cls.get_from_keys_or(
91
+ config, ["modules_to_not_convert"], None
92
+ )
93
+ return cls(weight_bits, group_size, zero_point, modules_to_not_convert)
94
+
95
+ def get_quant_method(
96
+ self, layer: torch.nn.Module, prefix: str
97
+ ) -> Union["LinearMethodBase", "QuantizeMethodBase"] | None:
98
+ if isinstance(layer, LinearBase):
99
+ if is_layer_skipped(
100
+ prefix,
101
+ self.modules_to_not_convert,
102
+ self.packed_modules_mapping,
103
+ skip_with_substr=True,
104
+ ):
105
+ return UnquantizedLinearMethod()
106
+ return AWQLinearMethod(self)
107
+ elif isinstance(layer, FusedMoE):
108
+ # Lazy import to avoid circular import.
109
+ from .awq_marlin import AWQMarlinConfig
110
+ from .moe_wna16 import MoeWNA16Config
111
+ from .utils.marlin_utils import check_moe_marlin_supports_layer
112
+
113
+ if not check_moe_marlin_supports_layer(layer, self.group_size):
114
+ logger.warning_once(
115
+ f"Layer '{prefix}' is not supported by AWQMoeMarlin. "
116
+ "Falling back to Moe WNA16 kernels."
117
+ )
118
+ config = {
119
+ "quant_method": "awq",
120
+ "bits": self.weight_bits,
121
+ "group_size": self.group_size,
122
+ "zero_point": self.zero_point,
123
+ "lm_head": False,
124
+ "modules_to_not_convert": self.modules_to_not_convert,
125
+ }
126
+ return MoeWNA16Config.from_config(config).get_quant_method(
127
+ layer, prefix
128
+ )
129
+ marlin_compatible_config_dict = {
130
+ "quant_method": "awq",
131
+ "bits": self.weight_bits,
132
+ "group_size": self.group_size,
133
+ "zero_point": self.zero_point,
134
+ "lm_head": False,
135
+ "modules_to_not_convert": self.modules_to_not_convert,
136
+ }
137
+ awq_marlin_config = AWQMarlinConfig.from_config(
138
+ marlin_compatible_config_dict
139
+ )
140
+ return awq_marlin_config.get_quant_method(layer, prefix)
141
+ return None
142
+
143
+ def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
144
+ if self.modules_to_not_convert:
145
+ self.modules_to_not_convert = hf_to_vllm_mapper.apply_list(
146
+ self.modules_to_not_convert
147
+ )
148
+
149
+ def maybe_update_config(self, model_name: str, revision: str | None = None):
150
+ if self.modules_to_not_convert:
151
+ return
152
+
153
+ unquant_dtypes = [torch.float16, torch.bfloat16, torch.float32]
154
+ metadata = get_safetensors_params_metadata(model_name, revision=revision)
155
+ layers = {param_name.rsplit(".", 1)[0] for param_name in metadata}
156
+ quant_layers: set[str] = {
157
+ param_name.rsplit(".", 1)[0]
158
+ for param_name, info in metadata.items()
159
+ if (dtype := info.get("dtype", None))
160
+ and _SAFETENSORS_TO_TORCH_DTYPE[dtype] not in unquant_dtypes
161
+ }
162
+ self.modules_to_not_convert = list(layers - quant_layers)
163
+
164
+
165
+class AWQLinearMethod(LinearMethodBase):
166
+ """Linear method for AWQ.
167
+
168
+ Args:
169
+ quant_config: The AWQ quantization config.
170
+ """
171
+
172
+ def __init__(self, quant_config: AWQConfig):
173
+ self.quant_config = quant_config
174
+
175
+ def create_weights(
176
+ self,
177
+ layer: torch.nn.Module,
178
+ input_size_per_partition: int,
179
+ output_partition_sizes: list[int],
180
+ input_size: int,
181
+ output_size: int,
182
+ params_dtype: torch.dtype,
183
+ **extra_weight_attrs,
184
+ ):
185
+ # Normalize group_size
186
+ if self.quant_config.group_size != -1:
187
+ group_size = self.quant_config.group_size
188
+ else:
189
+ group_size = input_size
190
+
191
+ if input_size_per_partition % group_size != 0:
192
+ raise ValueError(
193
+ "The input size is not aligned with the quantized "
194
+ "weight shape. This can be caused by too large "
195
+ "tensor parallel size."
196
+ )
197
+
198
+ output_size_per_partition = sum(output_partition_sizes)
199
+ if output_size_per_partition % self.quant_config.pack_factor != 0:
200
+ raise ValueError(
201
+ "The output size is not aligned with the quantized "
202
+ "weight shape. This can be caused by too large "
203
+ "tensor parallel size."
204
+ )
205
+
206
+ weight_loader = extra_weight_attrs.get("weight_loader")
207
+ qweight = PackedvLLMParameter(
208
+ data=torch.empty(
209
+ input_size_per_partition,
210
+ output_size_per_partition // self.quant_config.pack_factor,
211
+ dtype=torch.int32,
212
+ ),
213
+ input_dim=0,
214
+ output_dim=1,
215
+ packed_dim=1,
216
+ packed_factor=self.quant_config.pack_factor,
217
+ weight_loader=weight_loader,
218
+ )
219
+
220
+ num_groups = input_size_per_partition // group_size
221
+
222
+ qzeros = PackedvLLMParameter(
223
+ data=torch.empty(
224
+ num_groups,
225
+ output_size_per_partition // self.quant_config.pack_factor,
226
+ dtype=torch.int32,
227
+ ),
228
+ input_dim=0,
229
+ output_dim=1,
230
+ packed_dim=1,
231
+ packed_factor=self.quant_config.pack_factor,
232
+ weight_loader=weight_loader,
233
+ )
234
+
235
+ scales = GroupQuantScaleParameter(
236
+ data=torch.empty(
237
+ num_groups,
238
+ output_size_per_partition,
239
+ dtype=params_dtype,
240
+ ),
241
+ input_dim=0,
242
+ output_dim=1,
243
+ weight_loader=weight_loader,
244
+ )
245
+
246
+ layer.register_parameter("qweight", qweight)
247
+ layer.register_parameter("qzeros", qzeros)
248
+ layer.register_parameter("scales", scales)
249
+
250
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
251
+ layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False)
252
+ layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False)
253
+ layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
254
+
255
+ def apply(
256
+ self,
257
+ layer: torch.nn.Module,
258
+ x: torch.Tensor,
259
+ bias: torch.Tensor | None = None,
260
+ ) -> torch.Tensor:
261
+ qweight = layer.qweight
262
+ scales = layer.scales
263
+ qzeros = layer.qzeros
264
+ pack_factor = self.quant_config.pack_factor
265
+ out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,)
266
+ reshaped_x = x.reshape(-1, x.shape[-1])
267
+
268
+ # num_tokens >= threshold
269
+ FP16_MATMUL_HEURISTIC_CONDITION = x.shape[:-1].numel() >= 256
270
+
271
+ if FP16_MATMUL_HEURISTIC_CONDITION:
272
+ out = ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
273
+ out = torch.matmul(reshaped_x, out)
274
+ else:
275
+ out = ops.awq_gemm(reshaped_x, qweight, scales, qzeros, pack_factor)
276
+ if bias is not None:
277
+ out.add_(bias)
278
+ return out.reshape(out_shape)
vllm/model_executor/layers/quantization/awq_marlin.py
+2
-247
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import TYPE_CHECKING, Any, Union
4
+from typing import TYPE_CHECKING, Any
5
6
import torch
7
from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
@@ -27,6 +27,7 @@ from vllm.model_executor.layers.linear import (
27
UnquantizedLinearMethod,
28
set_weight_attrs,
29
)
30
+from vllm.model_executor.layers.quantization.awq import AWQConfig
31
from vllm.model_executor.layers.quantization.base_config import (
32
QuantizationConfig,
33
QuantizeMethodBase,
@@ -63,252 +64,6 @@ if TYPE_CHECKING:
64
logger = init_logger(__name__)
65
66
66
-class AWQConfig(QuantizationConfig):
67
- """Config class for AWQ.
68
-
69
- Reference: https://arxiv.org/abs/2306.00978
70
- """
71
-
72
- def __init__(
73
- self,
74
- weight_bits: int,
75
- group_size: int,
76
- zero_point: bool,
77
- modules_to_not_convert: list[str] | None = None,
78
- ) -> None:
79
- super().__init__()
80
- self.weight_bits = weight_bits
81
- self.group_size = group_size
82
- self.zero_point = zero_point
83
- self.modules_to_not_convert = modules_to_not_convert or []
84
-
85
- if self.weight_bits != 4:
86
- raise ValueError(
87
- "Currently, only 4-bit weight quantization is supported for "
88
- f"AWQ, but got {self.weight_bits} bits."
89
- )
90
- self.pack_factor = 32 // self.weight_bits
91
-
92
- def __repr__(self) -> str:
93
- return (
94
- f"AWQConfig(weight_bits={self.weight_bits}, "
95
- f"group_size={self.group_size}, "
96
- f"zero_point={self.zero_point}, "
97
- f"modules_to_not_convert={self.modules_to_not_convert})"
98
- )
99
-
100
- def get_name(self) -> "QuantizationMethods":
101
- return "awq"
102
-
103
- def get_supported_act_dtypes(self) -> list[torch.dtype]:
104
- return [torch.half]
105
-
106
- @classmethod
107
- def get_min_capability(cls) -> int:
108
- # The AWQ kernel only supports Turing or newer GPUs.
109
- return 75
110
-
111
- @staticmethod
112
- def get_config_filenames() -> list[str]:
113
- return [
114
- "quant_config.json", # E.g., casperhansen/vicuna-7b-v1.5-awq
115
- # E.g., abhinavkulkarni/mosaicml-mpt-7b-instruct-w4-g128-awq
116
- "quantize_config.json",
117
- ]
118
-
119
- @classmethod
120
- def from_config(cls, config: dict[str, Any]) -> "AWQConfig":
121
- weight_bits = cls.get_from_keys(config, ["w_bit", "bits"])
122
- group_size = cls.get_from_keys(config, ["q_group_size", "group_size"])
123
- zero_point = cls.get_from_keys(config, ["zero_point"])
124
- modules_to_not_convert = cls.get_from_keys_or(
125
- config, ["modules_to_not_convert"], None
126
- )
127
- return cls(weight_bits, group_size, zero_point, modules_to_not_convert)
128
-
129
- def get_quant_method(
130
- self, layer: torch.nn.Module, prefix: str
131
- ) -> Union["LinearMethodBase", "QuantizeMethodBase"] | None:
132
- if isinstance(layer, LinearBase):
133
- if is_layer_skipped(
134
- prefix,
135
- self.modules_to_not_convert,
136
- self.packed_modules_mapping,
137
- skip_with_substr=True,
138
- ):
139
- return UnquantizedLinearMethod()
140
- return AWQLinearMethod(self)
141
- elif isinstance(layer, FusedMoE):
142
- from .moe_wna16 import MoeWNA16Config
143
-
144
- if not check_moe_marlin_supports_layer(layer, self.group_size):
145
- logger.warning_once(
146
- f"Layer '{prefix}' is not supported by AWQMoeMarlin. "
147
- "Falling back to Moe WNA16 kernels."
148
- )
149
- config = {
150
- "quant_method": "awq",
151
- "bits": self.weight_bits,
152
- "group_size": self.group_size,
153
- "zero_point": self.zero_point,
154
- "lm_head": False,
155
- "modules_to_not_convert": self.modules_to_not_convert,
156
- }
157
- return MoeWNA16Config.from_config(config).get_quant_method(
158
- layer, prefix
159
- )
160
- marlin_compatible_config_dict = {
161
- "quant_method": "awq",
162
- "bits": self.weight_bits,
163
- "group_size": self.group_size,
164
- "zero_point": self.zero_point,
165
- "lm_head": False,
166
- "modules_to_not_convert": self.modules_to_not_convert,
167
- }
168
- awq_marlin_config = AWQMarlinConfig.from_config(
169
- marlin_compatible_config_dict
170
- )
171
- return awq_marlin_config.get_quant_method(layer, prefix)
172
- return None
173
-
174
- def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
175
- if self.modules_to_not_convert:
176
- self.modules_to_not_convert = hf_to_vllm_mapper.apply_list(
177
- self.modules_to_not_convert
178
- )
179
-
180
- def maybe_update_config(self, model_name: str, revision: str | None = None):
181
- if self.modules_to_not_convert:
182
- return
183
-
184
- unquant_dtypes = [torch.float16, torch.bfloat16, torch.float32]
185
- metadata = get_safetensors_params_metadata(model_name, revision=revision)
186
- layers = {param_name.rsplit(".", 1)[0] for param_name in metadata}
187
- quant_layers: set[str] = {
188
- param_name.rsplit(".", 1)[0]
189
- for param_name, info in metadata.items()
190
- if (dtype := info.get("dtype", None))
191
- and _SAFETENSORS_TO_TORCH_DTYPE[dtype] not in unquant_dtypes
192
- }
193
- self.modules_to_not_convert = list(layers - quant_layers)
194
-
195
-
196
-class AWQLinearMethod(LinearMethodBase):
197
- """Linear method for AWQ.
198
-
199
- Args:
200
- quant_config: The AWQ quantization config.
201
- """
202
-
203
- def __init__(self, quant_config: AWQConfig):
204
- self.quant_config = quant_config
205
-
206
- def create_weights(
207
- self,
208
- layer: torch.nn.Module,
209
- input_size_per_partition: int,
210
- output_partition_sizes: list[int],
211
- input_size: int,
212
- output_size: int,
213
- params_dtype: torch.dtype,
214
- **extra_weight_attrs,
215
- ):
216
- # Normalize group_size
217
- if self.quant_config.group_size != -1:
218
- group_size = self.quant_config.group_size
219
- else:
220
- group_size = input_size
221
-
222
- if input_size_per_partition % group_size != 0:
223
- raise ValueError(
224
- "The input size is not aligned with the quantized "
225
- "weight shape. This can be caused by too large "
226
- "tensor parallel size."
227
- )
228
-
229
- output_size_per_partition = sum(output_partition_sizes)
230
- if output_size_per_partition % self.quant_config.pack_factor != 0:
231
- raise ValueError(
232
- "The output size is not aligned with the quantized "
233
- "weight shape. This can be caused by too large "
234
- "tensor parallel size."
235
- )
236
-
237
- weight_loader = extra_weight_attrs.get("weight_loader")
238
- qweight = PackedvLLMParameter(
239
- data=torch.empty(
240
- input_size_per_partition,
241
- output_size_per_partition // self.quant_config.pack_factor,
242
- dtype=torch.int32,
243
- ),
244
- input_dim=0,
245
- output_dim=1,
246
- packed_dim=1,
247
- packed_factor=self.quant_config.pack_factor,
248
- weight_loader=weight_loader,
249
- )
250
-
251
- num_groups = input_size_per_partition // group_size
252
-
253
- qzeros = PackedvLLMParameter(
254
- data=torch.empty(
255
- num_groups,
256
- output_size_per_partition // self.quant_config.pack_factor,
257
- dtype=torch.int32,
258
- ),
259
- input_dim=0,
260
- output_dim=1,
261
- packed_dim=1,
262
- packed_factor=self.quant_config.pack_factor,
263
- weight_loader=weight_loader,
264
- )
265
-
266
- scales = GroupQuantScaleParameter(
267
- data=torch.empty(
268
- num_groups,
269
- output_size_per_partition,
270
- dtype=params_dtype,
271
- ),
272
- input_dim=0,
273
- output_dim=1,
274
- weight_loader=weight_loader,
275
- )
276
-
277
- layer.register_parameter("qweight", qweight)
278
- layer.register_parameter("qzeros", qzeros)
279
- layer.register_parameter("scales", scales)
280
-
281
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
282
- layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False)
283
- layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False)
284
- layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
285
-
286
- def apply(
287
- self,
288
- layer: torch.nn.Module,
289
- x: torch.Tensor,
290
- bias: torch.Tensor | None = None,
291
- ) -> torch.Tensor:
292
- qweight = layer.qweight
293
- scales = layer.scales
294
- qzeros = layer.qzeros
295
- pack_factor = self.quant_config.pack_factor
296
- out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,)
297
- reshaped_x = x.reshape(-1, x.shape[-1])
298
-
299
- # num_tokens >= threshold
300
- FP16_MATMUL_HEURISTIC_CONDITION = x.shape[:-1].numel() >= 256
301
-
302
- if FP16_MATMUL_HEURISTIC_CONDITION:
303
- out = ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
304
- out = torch.matmul(reshaped_x, out)
305
- else:
306
- out = ops.awq_gemm(reshaped_x, qweight, scales, qzeros, pack_factor)
307
- if bias is not None:
308
- out.add_(bias)
309
- return out.reshape(out_shape)
310
-
311
-
67
class AWQMarlinConfig(QuantizationConfig):
68
"""Config class for AWQ Marlin"""
69