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
Consolidate AWQ quantization into single awq_marlin.py file
Merge awq.py and awq_marlin.py into a single file, eliminating the circular import between them. awq.py becomes a backward-compat shim. Follows the same structure as gptq_marlin.py. Co-authored-by: Claude Signed-off-by: Robert Shaw <robertgshaw2@gmail.com>
Robert Shaw committed
Mar 21, 2026 at 17:09 UTC
8cc700dd6a0d44c3b6983c0305af790b703f6b48
3 files changed
+253
-279
vllm/model_executor/layers/quantization/__init__.py
+1
-2
@@ -107,8 +107,7 @@ 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 import AWQConfig
111
- from .awq_marlin import AWQMarlinConfig
110
+ from .awq_marlin import AWQConfig, AWQMarlinConfig
111
from .bitsandbytes import BitsAndBytesConfig
112
from .compressed_tensors.compressed_tensors import (
113
CompressedTensorsConfig,
vllm/model_executor/layers/quantization/awq.py
+5
-275
@@ -1,278 +1,8 @@
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
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,
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,
8
)
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
+247
-2
@@ -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
4
+from typing import TYPE_CHECKING, Any, Union
5
6
import torch
7
from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
@@ -27,7 +27,6 @@ from vllm.model_executor.layers.linear import (
27
UnquantizedLinearMethod,
28
set_weight_attrs,
29
)
30
-from vllm.model_executor.layers.quantization.awq import AWQConfig
30
from vllm.model_executor.layers.quantization.base_config import (
31
QuantizationConfig,
32
QuantizeMethodBase,
@@ -64,6 +63,252 @@ if TYPE_CHECKING:
63
logger = init_logger(__name__)
64
65
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
+
312
class AWQMarlinConfig(QuantizationConfig):
313
"""Config class for AWQ Marlin"""
314