@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8d4f54966c
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

fix(quantization): Fix AWQ dequantize on Intel XPU and refactor AutoAWQ config (#42727)

Signed-off-by: Alex <alex.tech.lab@outlook.com> Signed-off-by: AlexHuang <jihuihuang@tencent.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

AlexHuang committed Jun 18, 2026 at 20:12 UTC 8d4f54966cdb8f1d0768fbe5319e400047877a3d
20 files changed +579 -428
docs/features/quantization/auto_awq.md
+2 -2
@@ -49,7 +49,7 @@ To run an AWQ model with vLLM, you can use [TheBloke/Llama-2-7b-Chat-AWQ](https:
49 ```bash
50 python examples/deployment/llm_engine_example.py \
51 --model TheBloke/Llama-2-7b-Chat-AWQ \
52 - --quantization awq
52 + --quantization auto_awq
53 ```
54
55 AWQ models are also supported directly through the LLM entrypoint:
@@ -70,7 +70,7 @@ AWQ models are also supported directly through the LLM entrypoint:
70 sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
71
72 # Create an LLM.
73 - llm = LLM(model="TheBloke/Llama-2-7b-Chat-AWQ", quantization="AWQ")
73 + llm = LLM(model="TheBloke/Llama-2-7b-Chat-AWQ", quantization="auto_awq")
74 # Generate texts from the prompts. The output is a list of RequestOutput objects
75 # that contain the prompt, generated text, and other information.
76 outputs = llm.generate(prompts, sampling_params)
tests/quantization/test_auto_awq.py new
+231
@@ -0,0 +1,231 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Tests for AutoAWQConfig behavior after unification.
5 +
6 +These tests verify the bug fixes for:
7 +1. CPU platform override conflict (auto_awq should not override on CPU)
8 +2. MoE fallback compatibility (full_config["quant_method"] should be "awq")
9 +3. Config attribute consistency
10 +4. End-to-end quantization method loading (auto_awq loads and runs correctly)
11 +
12 +Note: Tests that require importing the full auto_awq module (which has GPU-dependent
13 +imports) should use subprocess or be run in a GPU environment.
14 +"""
15 +
16 +from __future__ import annotations
17 +
18 +import pytest
19 +import torch
20 +
21 +from tests.quantization.utils import is_quant_method_supported
22 +
23 +
24 +def _get_auto_awq_config_source() -> str:
25 + """Read the AutoAWQConfig class source code for isolated testing."""
26 + import inspect
27 +
28 + import vllm.model_executor.layers.quantization.auto_awq as auto_awq_module
29 +
30 + return inspect.getsource(auto_awq_module.AutoAWQConfig)
31 +
32 +
33 +class TestAutoAWQConfigFromConfig:
34 + """Tests for AutoAWQConfig.from_config behavior.
35 +
36 + These tests require GPU environment to import the full module.
37 + They are skipped on non-GPU platforms.
38 + """
39 +
40 + def test_full_config_quant_method_is_awq_for_moe_fallback(self):
41 + """full_config should have quant_method='awq' for MoE fallback compatibility.
42 +
43 + MoeWNA16Config only accepts 'gptq' or 'awq' as linear_quant_method.
44 + If full_config has 'auto_awq', the MoE fallback will fail.
45 + """
46 + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
47 +
48 + config = {
49 + "w_bit": 4,
50 + "q_group_size": 128,
51 + "zero_point": True,
52 + "lm_head": False,
53 + }
54 + awq_config = AutoAWQConfig.from_config(config)
55 +
56 + # Verify quant_method is 'awq' for MoE fallback
57 + assert awq_config.full_config["quant_method"] == "awq", (
58 + f"Expected quant_method='awq', got {awq_config.full_config['quant_method']}"
59 + )
60 +
61 + def test_full_config_preserves_other_fields(self):
62 + """full_config should preserve all original config fields."""
63 + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
64 +
65 + config = {
66 + "w_bit": 4,
67 + "q_group_size": 128,
68 + "zero_point": True,
69 + "lm_head": False,
70 + "custom_field": "custom_value",
71 + }
72 + awq_config = AutoAWQConfig.from_config(config)
73 +
74 + assert awq_config.full_config["w_bit"] == 4
75 + assert awq_config.full_config["q_group_size"] == 128
76 + assert awq_config.full_config["zero_point"] is True
77 + assert awq_config.full_config["lm_head"] is False
78 + assert awq_config.full_config["custom_field"] == "custom_value"
79 +
80 + def test_full_config_is_copy_not_original(self):
81 + """full_config should be a copy, not the original dict."""
82 + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
83 +
84 + config = {
85 + "w_bit": 4,
86 + "q_group_size": 128,
87 + "zero_point": True,
88 + "lm_head": False,
89 + }
90 + original_quant_method = config.get("quant_method")
91 +
92 + AutoAWQConfig.from_config(config)
93 +
94 + # Original config should not be modified
95 + assert config.get("quant_method") == original_quant_method
96 +
97 +
98 +class TestAutoAWQConfigAttributes:
99 + """Tests for AutoAWQConfig attribute consistency.
100 +
101 + These tests require GPU environment to import the full module.
102 + They are skipped on non-GPU platforms.
103 + """
104 +
105 + def test_config_attributes_match_input(self):
106 + """Config attributes should match input values."""
107 + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
108 +
109 + awq_config = AutoAWQConfig(
110 + weight_bits=4,
111 + group_size=128,
112 + zero_point=True,
113 + lm_head_quantized=False,
114 + modules_to_not_convert=["lm_head"],
115 + )
116 +
117 + assert awq_config.weight_bits == 4
118 + assert awq_config.group_size == 128
119 + assert awq_config.zero_point is True
120 + assert awq_config.lm_head_quantized is False
121 + assert awq_config.modules_to_not_convert == ["lm_head"]
122 +
123 + def test_pack_factor_for_4bit(self):
124 + """Pack factor should be 8 for 4-bit quantization."""
125 + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
126 +
127 + awq_config = AutoAWQConfig(
128 + weight_bits=4,
129 + group_size=128,
130 + zero_point=True,
131 + lm_head_quantized=False,
132 + )
133 +
134 + assert awq_config.pack_factor == 8 # 32 // 4
135 +
136 +
137 +class TestAutoAWQConfigOverrideLogic:
138 + """Tests for override logic by parsing source code (no GPU import required)."""
139 +
140 + def _get_auto_awq_source(self) -> str:
141 + """Read the auto_awq.py source file."""
142 + import inspect
143 + import pathlib
144 +
145 + import vllm.model_executor.layers.quantization.auto_awq as auto_awq_module
146 +
147 + source_path = inspect.getfile(auto_awq_module)
148 + return pathlib.Path(source_path).read_text()
149 +
150 + def test_cpu_check_in_override_method(self):
151 + """override_quantization_method should check current_platform.is_cpu()."""
152 + source = self._get_auto_awq_source()
153 +
154 + # Verify the CPU check exists in override method
155 + assert "current_platform.is_cpu()" in source, (
156 + "override_quantization_method should check is_cpu()"
157 + )
158 + assert "return None" in source, (
159 + "override_quantization_method should return None on CPU"
160 + )
161 +
162 + def test_quant_method_normalization_in_from_config(self):
163 + """from_config should normalize quant_method to 'awq' for MoE fallback."""
164 + source = self._get_auto_awq_source()
165 +
166 + # Verify the normalization exists
167 + assert (
168 + '"quant_method"] = "awq"' in source or "'quant_method'] = 'awq'" in source
169 + ), "from_config should set quant_method='awq' in full_config"
170 +
171 +
172 +# =============================================================================
173 +# End-to-end integration tests (require GPU environment)
174 +# =============================================================================
175 +
176 +PROMPT = "On the surface of Mars, we found"
177 +
178 +# Small AWQ model for testing - using Qwen2 1.5B which has official AWQ checkpoint
179 +AWQ_MODELS = [
180 + "Qwen/Qwen2-1.5B-Instruct-AWQ",
181 +]
182 +
183 +
184 +@pytest.mark.skipif(
185 + not is_quant_method_supported("auto_awq"),
186 + reason="auto_awq is not supported on this GPU type.",
187 +)
188 +@pytest.mark.parametrize("model_id", AWQ_MODELS)
189 +def test_auto_awq_quantization_method(vllm_runner, model_id: str, monkeypatch):
190 + """Test that quantization='auto_awq' loads and runs correctly."""
191 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
192 +
193 + with vllm_runner(
194 + model_id,
195 + dtype=torch.float16,
196 + quantization="auto_awq",
197 + max_model_len=2048,
198 + enforce_eager=True,
199 + ) as llm:
200 +
201 + def check_model(model):
202 + from vllm.model_executor.layers.quantization.auto_awq import (
203 + AutoAWQLinearMethod,
204 + AutoAWQMarlinLinearMethod,
205 + )
206 +
207 + for name, submodule in model.named_modules():
208 + if name == "model.layers.0.self_attn.qkv_proj":
209 + # Should use either AutoAWQLinearMethod (Triton) or
210 + # AutoAWQMarlinLinearMethod (Marlin) depending on hardware
211 + assert isinstance(
212 + submodule.quant_method,
213 + (AutoAWQLinearMethod, AutoAWQMarlinLinearMethod),
214 + ), (
215 + f"Expected AutoAWQLinearMethod or AutoAWQMarlinLinearMethod "
216 + f"for {name}, got {type(submodule.quant_method)}"
217 + )
218 + break
219 +
220 + llm.apply_model(check_model)
221 +
222 + outputs = llm.generate_greedy([PROMPT], max_tokens=8)
223 + assert outputs
224 + assert len(outputs[0][1]) > 0
225 +
226 +
227 +def test_auto_awq_config_get_name():
228 + """Test that AutoAWQConfig.get_name() returns 'auto_awq'."""
229 + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
230 +
231 + assert AutoAWQConfig.get_name() == "auto_awq"
tests/quantization/test_auto_round.py
+2 -2
@@ -645,11 +645,11 @@ def test_resolve_awq_moe_uses_marlin_when_supported(monkeypatch) -> None:
645 lambda *args, **kwargs: True,
646 )
647 monkeypatch.setattr(
648 - "vllm.model_executor.layers.quantization.awq_marlin.verify_marlin_supported",
648 + "vllm.model_executor.layers.quantization.auto_awq.verify_marlin_supported",
649 lambda *args, **kwargs: None,
650 )
651 monkeypatch.setattr(
652 - "vllm.model_executor.layers.quantization.awq_marlin.AWQMarlinMoEMethod",
652 + "vllm.model_executor.layers.quantization.auto_awq.AutoAWQMoEMethod",
653 DummyMethod,
654 )
655
tests/quantization/test_configs.py
+5 -3
@@ -43,16 +43,18 @@ MODEL_ARG_EXPTYPES = [
43 ("LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit", "gptq", "auto_gptq"),
44 ("LnL-AI/TinyLlama-1.1B-Chat-v1.0-GPTQ-4bit", "awq", "ERROR"),
45 # AUTOAWQ
46 + # AutoAWQConfig.override_quantization_method() returns "auto_awq" for AWQ models
47 + # when user_quant is None, "awq", "awq_marlin", "marlin", or "auto_awq"
48 (
49 "TheBloke/OpenHermes-2.5-Mistral-7B-AWQ",
50 None,
49 - "awq_marlin" if current_platform.is_cuda_alike() else "awq",
51 + "auto_awq",
52 ),
51 - ("TheBloke/OpenHermes-2.5-Mistral-7B-AWQ", "awq", "awq"),
53 + ("TheBloke/OpenHermes-2.5-Mistral-7B-AWQ", "awq", "auto_awq"),
54 (
55 "TheBloke/OpenHermes-2.5-Mistral-7B-AWQ",
56 "marlin",
55 - "awq_marlin" if current_platform.is_cuda_alike() else "ERROR",
57 + "auto_awq" if current_platform.is_cuda_alike() else "ERROR",
58 ),
59 ("TheBloke/OpenHermes-2.5-Mistral-7B-AWQ", "gptq", "ERROR"),
60 ]
vllm/config/model.py
+2
@@ -984,6 +984,8 @@ class ModelConfig:
984 "auto_gptq",
985 "gptq",
986 "gptq_marlin",
987 + "auto_awq",
988 + "awq",
989 "awq_marlin",
990 "inc",
991 "moe_wna16",
vllm/model_executor/kernels/linear/mixed_precision/cpu.py
+1 -1
@@ -175,7 +175,7 @@ class CPUWNA16LinearKernel(MPLinearKernel):
175 and torch.cpu._is_amx_tile_supported()
176 )
177 # layer.use_w4a8 = False
178 - # AWQ format will be converted to GPTQ format in `AWQMarlinLinearMethod`
178 + # AWQ format will be converted to GPTQ format in `AutoAWQMarlinLinearMethod`
179 if layer.use_w4a8:
180 self._process_gptq_weights_w4a8(layer)
181 else:
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+10 -10
@@ -728,18 +728,18 @@ def _process_weights_cpu(
728 from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
729 prepare_int4_moe_layer_for_cpu,
730 )
731 + from vllm.model_executor.layers.quantization.auto_awq import (
732 + AutoAWQConfig,
733 + )
734 from vllm.model_executor.layers.quantization.auto_gptq import (
735 AutoGPTQConfig,
736 )
734 - from vllm.model_executor.layers.quantization.awq_marlin import (
735 - AWQMarlinConfig,
736 - )
737
738 # Detect packing format.
739 # AWQ: qweight is [E, K, 2*N//8] (packed along output/N dim).
740 # GPTQ: qweight is [E, K//8, 2*N] (packed along input/K dim).
741 # compressed-tensors: qweight is [E, K//8, 2*N] (packed along input/K dim).
742 - if isinstance(quant_config, AWQMarlinConfig):
742 + if isinstance(quant_config, AutoAWQConfig):
743 # AWQ: K is stored unpacked in dim 1.
744 cpu_quant_algo = ops.CPUQuantAlgo.AWQ
745 elif isinstance(quant_config, (AutoGPTQConfig, QuantizationArgs)):
@@ -753,7 +753,7 @@ def _process_weights_cpu(
753 cpu_quant_algo = ops.CPUQuantAlgo.GPTQ
754 else:
755 raise TypeError(
756 - "CPU WNA16 MoE backend requires AWQMarlinConfig, AutoGPTQConfig "
756 + "CPU WNA16 MoE backend requires AutoAWQConfig, AutoGPTQConfig "
757 f"or QuantizationArgs, got {type(quant_config).__name__}."
758 )
759
@@ -916,14 +916,14 @@ def convert_to_wna16_moe_kernel_format(
916 WNA16MoEBackend.MARLIN,
917 WNA16MoEBackend.BATCHED_MARLIN,
918 ):
919 + from vllm.model_executor.layers.quantization.auto_awq import (
920 + AutoAWQConfig,
921 + )
922 from vllm.model_executor.layers.quantization.auto_gptq import (
923 AutoGPTQConfig,
924 )
922 - from vllm.model_executor.layers.quantization.awq_marlin import (
923 - AWQMarlinConfig,
924 - )
925
926 - if isinstance(quant_config, AWQMarlinConfig):
926 + if isinstance(quant_config, AutoAWQConfig):
927 if w13_qzeros is None or w2_qzeros is None:
928 raise ValueError("AWQ Marlin MoE requires zero-point tensors.")
929
@@ -958,7 +958,7 @@ def convert_to_wna16_moe_kernel_format(
958 actorder = quant_config.actorder
959 else:
960 raise TypeError(
961 - "Marlin WNA16 MoE backend requires AutoGPTQConfig, AWQMarlinConfig or "
961 + "Marlin WNA16 MoE backend requires AutoAWQConfig, AutoGPTQConfig or "
962 f"QuantizationArgs, got {type(quant_config).__name__}."
963 )
964 if w13_g_idx is None or w2_g_idx is None:
vllm/model_executor/layers/linear.py
+2 -2
@@ -46,8 +46,8 @@ WEIGHT_LOADER_V2_SUPPORTED = [
46 "UnquantizedLinearMethod",
47 "CompressedTensorsLinearMethod",
48 "CompressedTensorsLinearTransformMethod",
49 - "AWQMarlinLinearMethod",
50 - "AWQLinearMethod",
49 + "AutoAWQMarlinLinearMethod",
50 + "AutoAWQLinearMethod",
51 "AutoGPTQLinearMethod",
52 "Fp8LinearMethod",
53 "FBGEMMFp8LinearMethod",
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+2 -2
@@ -48,8 +48,8 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
48 causal_conv1d_update,
49 )
50 from vllm.model_executor.layers.quantization import QuantizationConfig
51 +from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
52 from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
52 -from vllm.model_executor.layers.quantization.awq_marlin import AWQMarlinConfig
53 from vllm.model_executor.layers.quantization.inc import INCConfig
54 from vllm.model_executor.model_loader.weight_utils import (
55 sharded_weight_loader,
@@ -628,7 +628,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
628 return (
629 current_platform.is_cuda()
630 and not self.gqa_interleaved_layout
631 - and isinstance(quant_config, (AWQMarlinConfig, AutoGPTQConfig, INCConfig))
631 + and isinstance(quant_config, (AutoAWQConfig, AutoGPTQConfig, INCConfig))
632 )
633
634 def split_ba(self, ba: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
vllm/model_executor/layers/quantization/__init__.py
+5 -4
@@ -11,6 +11,7 @@ logger = init_logger(__name__)
11
12 QuantizationMethods = Literal[
13 "awq",
14 + "auto_awq",
15 "fp8",
16 "fbgemm_fp8",
17 "fp_quant",
@@ -113,9 +114,8 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
114 from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig
115 from vllm.models.deepseek_v4 import DeepseekV4FP8Config
116
117 + from .auto_awq import AutoAWQConfig
118 from .auto_gptq import AutoGPTQConfig
117 - from .awq import AWQConfig
118 - from .awq_marlin import AWQMarlinConfig
119 from .bitsandbytes import BitsAndBytesConfig
120 from .compressed_tensors.compressed_tensors import (
121 CompressedTensorsConfig,
@@ -138,7 +138,9 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
138 from .torchao import TorchAOConfig
139
140 method_to_config: dict[str, type[QuantizationConfig]] = {
141 - "awq": AWQConfig,
141 + "awq": AutoAWQConfig,
142 + "awq_marlin": AutoAWQConfig,
143 + "auto_awq": AutoAWQConfig,
144 "fp8": Fp8Config,
145 "fbgemm_fp8": FBGEMMFp8Config,
146 "fp_quant": FPQuantConfig,
@@ -149,7 +151,6 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
151 "auto_gptq": AutoGPTQConfig,
152 "gptq": AutoGPTQConfig,
153 "gptq_marlin": AutoGPTQConfig,
152 - "awq_marlin": AWQMarlinConfig,
154 "compressed-tensors": CompressedTensorsConfig,
155 "bitsandbytes": BitsAndBytesConfig,
156 "experts_int8": ExpertsInt8Config,
vllm/model_executor/layers/quantization/auto_awq.py renamed
+285 -71
@@ -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
@@ -9,6 +9,7 @@ from torch.nn import Parameter
9 from transformers import PretrainedConfig
10
11 import vllm.model_executor.layers.fused_moe # noqa
12 +from vllm import _custom_ops as ops
13 from vllm import envs
14 from vllm.logger import init_logger
15 from vllm.model_executor.kernels.linear import (
@@ -36,7 +37,6 @@ from vllm.model_executor.layers.linear import (
37 UnquantizedLinearMethod,
38 set_weight_attrs,
39 )
39 -from vllm.model_executor.layers.quantization.awq import AWQConfig
40 from vllm.model_executor.layers.quantization.base_config import (
41 QuantizationConfig,
42 QuantizeMethodBase,
@@ -55,7 +55,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
55 kInt4Static,
56 )
57 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
58 -from vllm.model_executor.parameter import GroupQuantScaleParameter, PackedvLLMParameter
58 +from vllm.model_executor.parameter import (
59 + GroupQuantScaleParameter,
60 + PackedvLLMParameter,
61 +)
62 from vllm.platforms import current_platform
63 from vllm.scalar_type import scalar_types
64 from vllm.transformers_utils.config import get_safetensors_params_metadata
@@ -164,8 +167,12 @@ def _convert_awq_to_standard_format(
167 setattr(layer, w_zp_name, new_zp_param)
168
169
167 -class AWQMarlinConfig(QuantizationConfig):
168 - """Config class for AWQ Marlin"""
170 +class AutoAWQConfig(QuantizationConfig):
171 + """Config class for AutoAWQ quantization.
172 +
173 + Unified config that supports multiple backends: Triton, Marlin, and XPU.
174 + Reference: https://arxiv.org/abs/2306.00978
175 + """
176
177 # num_bits -> type
178 TYPE_MAP = {
@@ -178,8 +185,8 @@ class AWQMarlinConfig(QuantizationConfig):
185 group_size: int,
186 zero_point: bool,
187 lm_head_quantized: bool,
181 - modules_to_not_convert: list[str] | None,
182 - full_config: dict[str, Any],
188 + modules_to_not_convert: list[str] | None = None,
189 + full_config: dict[str, Any] | None = None,
190 ) -> None:
191 super().__init__()
192 self.pack_factor = 32 // weight_bits # packed into int32
@@ -188,23 +195,22 @@ class AWQMarlinConfig(QuantizationConfig):
195 self.lm_head_quantized = lm_head_quantized
196 self.weight_bits = weight_bits
197 self.modules_to_not_convert = modules_to_not_convert or []
191 - self.full_config = full_config
198 + self.full_config = full_config or {}
199
200 if self.weight_bits not in self.TYPE_MAP:
201 + supported = ", ".join(str(k) for k in self.TYPE_MAP)
202 raise ValueError(
203 f"Unsupported num_bits = {self.weight_bits}. "
196 - f"Supported num_bits = {self.TYPE_MAP.keys()}"
204 + f"Supported: {supported}. "
205 + f"For 8-bit AWQ, use Marlin backend by setting "
206 + f"backend='awq:marlin' or backend='marlin'."
207 )
208
209 self.quant_type = self.TYPE_MAP[self.weight_bits]
210
201 - verify_marlin_supported(
202 - self.quant_type, group_size=self.group_size, has_zp=self.zero_point
203 - )
204 -
211 def __repr__(self) -> str:
212 return (
207 - f"AWQMarlinConfig(quant_type={self.quant_type}, "
213 + f"AutoAWQConfig(quant_type={self.quant_type}, "
214 f"group_size={self.group_size}, "
215 f"zero_point={self.zero_point}, "
216 f"lm_head_quantized={self.lm_head_quantized}, "
@@ -213,7 +219,7 @@ class AWQMarlinConfig(QuantizationConfig):
219
220 @classmethod
221 def get_name(cls) -> "QuantizationMethods":
216 - return "awq_marlin"
222 + return "auto_awq"
223
224 @classmethod
225 def get_supported_act_dtypes(cls) -> list[torch.dtype]:
@@ -225,60 +231,59 @@ class AWQMarlinConfig(QuantizationConfig):
231
232 @classmethod
233 def get_config_filenames(cls) -> list[str]:
228 - return ["quantize_config.json"]
234 + return ["quantize_config.json", "quant_config.json"]
235
236 @classmethod
231 - def from_config(cls, config: dict[str, Any]) -> "AWQMarlinConfig":
232 - weight_bits = cls.get_from_keys(config, ["bits"])
233 - group_size = cls.get_from_keys(config, ["group_size"])
237 + def from_config(cls, config: dict[str, Any]) -> "AutoAWQConfig":
238 + weight_bits = cls.get_from_keys(config, ["w_bit", "bits"])
239 + group_size = cls.get_from_keys(config, ["q_group_size", "group_size"])
240 zero_point = cls.get_from_keys(config, ["zero_point"])
241 lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False)
242 modules_to_not_convert = cls.get_from_keys_or(
243 config, ["modules_to_not_convert"], None
244 )
245 + # Ensure full_config uses "awq" as quant_method for MoE fallback compatibility.
246 + # MoeWNA16Config only accepts "gptq" or "awq", so we normalize here.
247 + full_config = config.copy()
248 + full_config["quant_method"] = "awq"
249 return cls(
250 weight_bits,
251 group_size,
252 zero_point,
253 lm_head_quantized,
254 modules_to_not_convert,
245 - config,
255 + full_config,
256 )
257
258 @classmethod
259 def override_quantization_method(
260 cls, hf_quant_cfg, user_quant, hf_config=None
261 ) -> "QuantizationMethods | None":
252 - # Skip override to marlin kernels, as they are not
253 - # batch invariant
254 - if envs.VLLM_BATCH_INVARIANT:
262 + """Override to use AutoAWQ for compatible AWQ models."""
263 + # Don't override on CPU - let cpu_awq handle it
264 + if current_platform.is_cpu():
265 + return None
266 +
267 + quant_method = hf_quant_cfg.get("quant_method", "").lower()
268 +
269 + if quant_method != "awq":
270 return None
271
257 - can_convert = cls.is_awq_marlin_compatible(hf_quant_cfg)
258 - is_valid_user_quant = (
259 - user_quant is None or user_quant == "marlin" or user_quant == "awq_marlin"
272 + is_valid_user_quant = user_quant is None or user_quant in (
273 + "awq",
274 + "awq_marlin",
275 + "auto_awq",
276 + "marlin",
277 )
278
262 - if can_convert and is_valid_user_quant:
263 - msg = (
264 - "The model is convertible to {} during runtime."
265 - " Using {} kernel.".format(cls.get_name(), cls.get_name())
266 - )
267 - logger.info(msg)
279 + if is_valid_user_quant:
280 return cls.get_name()
281
270 - if can_convert and user_quant == "awq":
271 - logger.info(
272 - "Detected that the model can run with awq_marlin"
273 - ", however you specified quantization=awq explicitly,"
274 - " so forcing awq. Use quantization=awq_marlin for"
275 - " faster inference"
276 - )
282 return None
283
284 def get_quant_method(
285 self, layer: torch.nn.Module, prefix: str
281 - ) -> "QuantizeMethodBase | None":
286 + ) -> Union["LinearMethodBase", "QuantizeMethodBase"] | None:
287 if isinstance(layer, LinearBase) or (
288 isinstance(layer, ParallelLMHead) and self.lm_head_quantized
289 ):
@@ -289,41 +294,66 @@ class AWQMarlinConfig(QuantizationConfig):
294 skip_with_substr=True,
295 ):
296 return UnquantizedLinearMethod()
292 - # Check if the layer is supported by AWQMarlin; tile-misaligned
293 - # shapes are fixed by padding at weight prep.
294 - if not check_marlin_supports_layer(
295 - layer, self.group_size, allow_tile_padding=True
296 - ):
297 - logger.warning_once(
298 - "Layer '%s' is not supported by AWQMarlin. Falling back to unoptimized AWQ kernels.", # noqa: E501
299 - prefix,
300 - )
301 - return AWQConfig.from_config(self.full_config).get_quant_method(
302 - layer, prefix
297 +
298 + # Check if XPU - use XPU-specific linear method
299 + if current_platform.is_xpu():
300 + return AutoAWQXPULinearMethod(self)
301 +
302 + # On CPU, use Marlin linear method which uses choose_mp_linear_kernel
303 + # to select the best available kernel (CPUWNA16LinearKernel on CPU)
304 + if current_platform.is_cpu():
305 + return AutoAWQMarlinLinearMethod(self)
306 +
307 + # Check if Marlin is supported and not using batch invariant mode
308 + # (Marlin kernels are not batch invariant)
309 + use_marlin = (
310 + not envs.VLLM_BATCH_INVARIANT
311 + and current_platform.is_cuda()
312 + and check_marlin_supported(
313 + self.quant_type, self.group_size, self.zero_point
314 )
304 - quant_method = AWQMarlinLinearMethod(self)
305 - quant_method.input_dtype = get_marlin_input_dtype(prefix)
306 - return quant_method
307 - elif isinstance(layer, RoutedExperts):
308 - from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
315 + )
316 +
317 + if use_marlin:
318 + # tile-misaligned shapes are fixed by padding at weight prep
319 + if not check_marlin_supports_layer(
320 + layer, self.group_size, allow_tile_padding=True
321 + ):
322 + logger.warning_once(
323 + "Layer '%s' is not supported by AutoAWQMarlin. "
324 + "Falling back to unoptimized AWQ kernels.",
325 + prefix,
326 + )
327 + return AutoAWQLinearMethod(self)
328 + quant_method = AutoAWQMarlinLinearMethod(self)
329 + quant_method.input_dtype = get_marlin_input_dtype(prefix)
330 + return quant_method
331 +
332 + return AutoAWQLinearMethod(self)
333
334 + elif isinstance(layer, RoutedExperts):
335 if is_layer_skipped(
336 prefix,
337 getattr(self, "modules_to_not_convert", []),
338 skip_with_substr=True,
339 ):
340 return UnquantizedFusedMoEMethod(layer.moe_config)
341 +
342 if not check_moe_marlin_supports_layer(layer, self.group_size):
343 logger.warning_once(
318 - f"Layer '{prefix}' is not supported by AWQMoeMarlin. "
344 + f"Layer '{prefix}' is not supported by AutoAWQMoEMarlin. "
345 "Falling back to Moe WNA16 kernels."
346 )
347 + from vllm.model_executor.layers.quantization.moe_wna16 import (
348 + MoeWNA16Config,
349 + )
350 +
351 return MoeWNA16Config.from_config(self.full_config).get_quant_method(
352 layer, prefix
353 )
324 - moe_quant_method = AWQMarlinMoEMethod(self, layer.moe_config)
325 - moe_quant_method.input_dtype = get_marlin_input_dtype(prefix)
326 - return moe_quant_method
354 +
355 + return AutoAWQMoEMethod(self, layer.moe_config)
356 +
357 return None
358
359 @classmethod
@@ -378,7 +408,7 @@ class AWQMarlinConfig(QuantizationConfig):
408 self.modules_to_not_convert = list(layers - quant_layers)
409
410
381 -class AWQMarlinLinearMethod(LinearMethodBase):
411 +class AutoAWQMarlinLinearMethod(LinearMethodBase):
412 """Linear method for AWQ Marlin.
413
414 Uses choose_mp_linear_kernel to select the best available kernel
@@ -390,16 +420,18 @@ class AWQMarlinLinearMethod(LinearMethodBase):
420
421 _kernel_backends_being_used: set[str] = set()
422
393 - def __init__(self, quant_config: AWQMarlinConfig) -> None:
423 + def __init__(self, quant_config: AutoAWQConfig) -> None:
424 self.quant_config = quant_config
425 self.quant_type = scalar_types.uint4
426 self.input_dtype = None
427
398 - verify_marlin_supported(
399 - quant_type=self.quant_config.quant_type,
400 - group_size=self.quant_config.group_size,
401 - has_zp=self.quant_config.zero_point,
402 - )
428 + # Skip Marlin verification on CPU - it will use CPUWNA16LinearKernel
429 + if not current_platform.is_cpu():
430 + verify_marlin_supported(
431 + quant_type=self.quant_config.quant_type,
432 + group_size=self.quant_config.group_size,
433 + has_zp=self.quant_config.zero_point,
434 + )
435
436 def create_weights(
437 self,
@@ -435,7 +467,7 @@ class AWQMarlinLinearMethod(LinearMethodBase):
467 kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config)
468
469 if kernel_type.__name__ not in self._kernel_backends_being_used:
438 - logger.info("Using %s for AWQMarlinLinearMethod", kernel_type.__name__)
470 + logger.info("Using %s for AutoAWQMarlinLinearMethod", kernel_type.__name__)
471 self._kernel_backends_being_used.add(kernel_type.__name__)
472
473 # Weights are loaded in AWQ checkpoint format (packed along output dim).
@@ -509,16 +541,16 @@ class AWQMarlinLinearMethod(LinearMethodBase):
541 return self.kernel.apply_weights(layer, x, bias)
542
543
512 -class AWQMarlinMoEMethod(FusedMoEMethodBase):
544 +class AutoAWQMoEMethod(FusedMoEMethodBase):
545 def __init__(
546 self,
515 - quant_config: AWQMarlinConfig,
547 + quant_config: AutoAWQConfig,
548 moe: FusedMoEConfig,
549 ):
550 super().__init__(moe)
551 self.quant_config = quant_config
552 if self.quant_config.weight_bits != 4:
521 - raise ValueError("AWQMarlinMoEMethod only supports 4bit now.")
553 + raise ValueError("AutoAWQMoEMethod only supports 4bit now.")
554 self.quant_type = scalar_types.uint4
555 self.input_dtype = None
556 self.use_marlin = True
@@ -784,3 +816,185 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
816 e_score_correction_bias=layer.e_score_correction_bias,
817 routed_scaling_factor=layer.routed_scaling_factor,
818 )
819 +
820 +
821 +class BaseAWQLinearMethod(LinearMethodBase):
822 + """Base class for AWQ linear methods with shared weight creation logic."""
823 +
824 + def __init__(self, quant_config: AutoAWQConfig):
825 + self.quant_config = quant_config
826 +
827 + def create_weights(
828 + self,
829 + layer: torch.nn.Module,
830 + input_size_per_partition: int,
831 + output_partition_sizes: list[int],
832 + input_size: int,
833 + output_size: int,
834 + params_dtype: torch.dtype,
835 + **extra_weight_attrs,
836 + ):
837 + # Normalize group_size
838 + if self.quant_config.group_size != -1:
839 + group_size = self.quant_config.group_size
840 + else:
841 + group_size = input_size
842 +
843 + if input_size_per_partition % group_size != 0:
844 + raise ValueError(
845 + "The input size is not aligned with the quantized "
846 + "weight shape. This can be caused by too large "
847 + "tensor parallel size."
848 + )
849 +
850 + output_size_per_partition = sum(output_partition_sizes)
851 + if output_size_per_partition % self.quant_config.pack_factor != 0:
852 + raise ValueError(
853 + "The output size is not aligned with the quantized "
854 + "weight shape. This can be caused by too large "
855 + "tensor parallel size."
856 + )
857 +
858 + weight_loader = extra_weight_attrs.get("weight_loader")
859 + qweight = PackedvLLMParameter(
860 + data=torch.empty(
861 + input_size_per_partition,
862 + output_size_per_partition // self.quant_config.pack_factor,
863 + dtype=torch.int32,
864 + ),
865 + input_dim=0,
866 + output_dim=1,
867 + packed_dim=1,
868 + packed_factor=self.quant_config.pack_factor,
869 + weight_loader=weight_loader,
870 + )
871 +
872 + num_groups = input_size_per_partition // group_size
873 +
874 + qzeros = PackedvLLMParameter(
875 + data=torch.empty(
876 + num_groups,
877 + output_size_per_partition // self.quant_config.pack_factor,
878 + dtype=torch.int32,
879 + ),
880 + input_dim=0,
881 + output_dim=1,
882 + packed_dim=1,
883 + packed_factor=self.quant_config.pack_factor,
884 + weight_loader=weight_loader,
885 + )
886 +
887 + scales = GroupQuantScaleParameter(
888 + data=torch.empty(
889 + num_groups,
890 + output_size_per_partition,
891 + dtype=params_dtype,
892 + ),
893 + input_dim=0,
894 + output_dim=1,
895 + weight_loader=weight_loader,
896 + )
897 +
898 + layer.register_parameter("qweight", qweight)
899 + layer.register_parameter("qzeros", qzeros)
900 + layer.register_parameter("scales", scales)
901 +
902 +
903 +class AutoAWQLinearMethod(BaseAWQLinearMethod):
904 + """Linear method for AWQ using Triton kernels.
905 +
906 + Args:
907 + quant_config: The AWQ quantization config.
908 + """
909 +
910 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
911 + layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False)
912 + layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False)
913 + layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
914 +
915 + def apply(
916 + self,
917 + layer: torch.nn.Module,
918 + x: torch.Tensor,
919 + bias: torch.Tensor | None = None,
920 + ) -> torch.Tensor:
921 + qweight = layer.qweight
922 + scales = layer.scales
923 + qzeros = layer.qzeros
924 + pack_factor = self.quant_config.pack_factor
925 + out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,)
926 + reshaped_x = x.reshape(-1, x.shape[-1])
927 +
928 + # num_tokens >= threshold
929 + FP16_MATMUL_HEURISTIC_CONDITION = x.shape[:-1].numel() >= 256
930 + # Batch invariant mode requires torch.matmul path
931 + # for Triton override
932 + if FP16_MATMUL_HEURISTIC_CONDITION or envs.VLLM_BATCH_INVARIANT:
933 + out = ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
934 + out = torch.matmul(reshaped_x, out)
935 + else:
936 + out = ops.awq_gemm(reshaped_x, qweight, scales, qzeros, pack_factor)
937 + if bias is not None:
938 + out.add_(bias)
939 + return out.reshape(out_shape)
940 +
941 +
942 +class AutoAWQXPULinearMethod(BaseAWQLinearMethod):
943 + """Linear method for AWQ on XPU using int4 GEMM kernel.
944 +
945 + Args:
946 + quant_config: The AWQ quantization config.
947 + """
948 +
949 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
950 + layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False)
951 + layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False)
952 + layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
953 +
954 + try:
955 + from vllm_xpu_kernels.quantization._quantize_convert import (
956 + AWQUtils,
957 + transpose_onednn_woq_format,
958 + )
959 + except ImportError as e:
960 + raise ImportError(
961 + "XPU AWQ requires vllm-xpu-kernels. "
962 + "Please install it with: pip install vllm-xpu-kernels"
963 + ) from e
964 +
965 + layer.xpu_output_size = layer.qweight.size(1) * self.quant_config.pack_factor
966 + qweight_new, qzeros_new = AWQUtils.repack(layer.qweight, layer.qzeros)
967 + if qweight_new.shape != layer.qweight.data.shape:
968 + layer.qweight.data = layer.qweight.data.view_as(qweight_new)
969 + if qzeros_new.shape != layer.qzeros.data.shape:
970 + layer.qzeros.data = layer.qzeros.data.view_as(qzeros_new)
971 + layer.qweight.data.copy_(qweight_new)
972 + layer.qzeros.data.copy_(qzeros_new)
973 + transpose_onednn_woq_format(layer, "awq", False)
974 +
975 + def _get_group_size(self, layer: torch.nn.Module) -> int:
976 + """Get the effective group size for kernel computation."""
977 + if self.quant_config.group_size != -1:
978 + return self.quant_config.group_size
979 + return layer.qweight.shape[0] # input_size_per_partition
980 +
981 + def apply(
982 + self,
983 + layer: torch.nn.Module,
984 + x: torch.Tensor,
985 + bias: torch.Tensor | None = None,
986 + ) -> torch.Tensor:
987 + reshaped_x = x.reshape(-1, x.shape[-1])
988 + group_size = self._get_group_size(layer)
989 +
990 + out = torch.ops._xpu_C.int4_gemm_w4a16(
991 + reshaped_x,
992 + layer.qweight,
993 + bias,
994 + layer.scales,
995 + layer.qzeros,
996 + group_size,
997 + None,
998 + )
999 + out_shape = x.shape[:-1] + (layer.xpu_output_size,)
1000 + return out.reshape(out_shape)
vllm/model_executor/layers/quantization/awq.py deleted
-286
@@ -1,286 +0,0 @@
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 -from transformers import PretrainedConfig
9 -
10 -from vllm import _custom_ops as ops
11 -from vllm import envs
12 -from vllm.logger import init_logger
13 -from vllm.model_executor.layers.fused_moe import RoutedExperts
14 -from vllm.model_executor.layers.linear import (
15 - LinearBase,
16 - LinearMethodBase,
17 - UnquantizedLinearMethod,
18 -)
19 -from vllm.model_executor.layers.quantization.base_config import (
20 - QuantizationConfig,
21 - QuantizeMethodBase,
22 -)
23 -from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
24 -from vllm.model_executor.parameter import GroupQuantScaleParameter, PackedvLLMParameter
25 -from vllm.transformers_utils.config import get_safetensors_params_metadata
26 -
27 -if TYPE_CHECKING:
28 - from vllm.model_executor.layers.quantization import QuantizationMethods
29 - from vllm.model_executor.models.utils import WeightsMapper
30 -
31 -logger = init_logger(__name__)
32 -
33 -
34 -class AWQConfig(QuantizationConfig):
35 - """Config class for AWQ.
36 -
37 - Reference: https://arxiv.org/abs/2306.00978
38 - """
39 -
40 - def __init__(
41 - self,
42 - weight_bits: int,
43 - group_size: int,
44 - zero_point: bool,
45 - modules_to_not_convert: list[str] | None = None,
46 - ) -> None:
47 - super().__init__()
48 - self.weight_bits = weight_bits
49 - self.group_size = group_size
50 - self.zero_point = zero_point
51 - self.modules_to_not_convert = modules_to_not_convert or []
52 -
53 - if self.weight_bits != 4:
54 - raise ValueError(
55 - "Currently, only 4-bit weight quantization is supported for "
56 - f"AWQ, but got {self.weight_bits} bits."
57 - )
58 - self.pack_factor = 32 // self.weight_bits
59 -
60 - def __repr__(self) -> str:
61 - return (
62 - f"AWQConfig(weight_bits={self.weight_bits}, "
63 - f"group_size={self.group_size}, "
64 - f"zero_point={self.zero_point}, "
65 - f"modules_to_not_convert={self.modules_to_not_convert})"
66 - )
67 -
68 - def get_name(self) -> "QuantizationMethods":
69 - return "awq"
70 -
71 - def get_supported_act_dtypes(self) -> list[torch.dtype]:
72 - return [torch.half]
73 -
74 - @classmethod
75 - def get_min_capability(cls) -> int:
76 - # The AWQ kernel only supports Turing or newer GPUs.
77 - return 75
78 -
79 - @staticmethod
80 - def get_config_filenames() -> list[str]:
81 - return [
82 - "quant_config.json", # E.g., casperhansen/vicuna-7b-v1.5-awq
83 - # E.g., abhinavkulkarni/mosaicml-mpt-7b-instruct-w4-g128-awq
84 - "quantize_config.json",
85 - ]
86 -
87 - @classmethod
88 - def from_config(cls, config: dict[str, Any]) -> "AWQConfig":
89 - weight_bits = cls.get_from_keys(config, ["w_bit", "bits"])
90 - group_size = cls.get_from_keys(config, ["q_group_size", "group_size"])
91 - zero_point = cls.get_from_keys(config, ["zero_point"])
92 - modules_to_not_convert = cls.get_from_keys_or(
93 - config, ["modules_to_not_convert"], None
94 - )
95 - return cls(weight_bits, group_size, zero_point, modules_to_not_convert)
96 -
97 - def get_quant_method(
98 - self, layer: torch.nn.Module, prefix: str
99 - ) -> Union["LinearMethodBase", "QuantizeMethodBase"] | None:
100 - if isinstance(layer, LinearBase):
101 - if is_layer_skipped(
102 - prefix,
103 - self.modules_to_not_convert,
104 - self.packed_modules_mapping,
105 - skip_with_substr=True,
106 - ):
107 - return UnquantizedLinearMethod()
108 - return AWQLinearMethod(self)
109 - elif isinstance(layer, RoutedExperts):
110 - # Lazy import to avoid circular import.
111 - from .awq_marlin import AWQMarlinConfig
112 - from .moe_wna16 import MoeWNA16Config
113 - from .utils.marlin_utils import check_moe_marlin_supports_layer
114 -
115 - if not check_moe_marlin_supports_layer(layer, self.group_size):
116 - logger.warning_once(
117 - f"Layer '{prefix}' is not supported by AWQMoeMarlin. "
118 - "Falling back to Moe WNA16 kernels."
119 - )
120 - config = {
121 - "quant_method": "awq",
122 - "bits": self.weight_bits,
123 - "group_size": self.group_size,
124 - "zero_point": self.zero_point,
125 - "lm_head": False,
126 - "modules_to_not_convert": self.modules_to_not_convert,
127 - }
128 - return MoeWNA16Config.from_config(config).get_quant_method(
129 - layer, prefix
130 - )
131 - marlin_compatible_config_dict = {
132 - "quant_method": "awq",
133 - "bits": self.weight_bits,
134 - "group_size": self.group_size,
135 - "zero_point": self.zero_point,
136 - "lm_head": False,
137 - "modules_to_not_convert": self.modules_to_not_convert,
138 - }
139 - awq_marlin_config = AWQMarlinConfig.from_config(
140 - marlin_compatible_config_dict
141 - )
142 - return awq_marlin_config.get_quant_method(layer, prefix)
143 - return None
144 -
145 - def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
146 - if self.modules_to_not_convert:
147 - self.modules_to_not_convert = hf_to_vllm_mapper.apply_list(
148 - self.modules_to_not_convert
149 - )
150 -
151 - def maybe_update_config(
152 - self,
153 - model_name: str,
154 - hf_config: PretrainedConfig | None = None,
155 - revision: str | None = None,
156 - ):
157 - if self.modules_to_not_convert:
158 - return
159 -
160 - unquant_dtypes = [torch.float16, torch.bfloat16, torch.float32]
161 - metadata = get_safetensors_params_metadata(model_name, revision=revision)
162 - layers = {param_name.rsplit(".", 1)[0] for param_name in metadata}
163 - quant_layers: set[str] = {
164 - param_name.rsplit(".", 1)[0]
165 - for param_name, info in metadata.items()
166 - if (dtype := info.get("dtype", None))
167 - and _SAFETENSORS_TO_TORCH_DTYPE[dtype] not in unquant_dtypes
168 - }
169 - self.modules_to_not_convert = list(layers - quant_layers)
170 -
171 -
172 -class AWQLinearMethod(LinearMethodBase):
173 - """Linear method for AWQ.
174 -
175 - Args:
176 - quant_config: The AWQ quantization config.
177 - """
178 -
179 - def __init__(self, quant_config: AWQConfig):
180 - self.quant_config = quant_config
181 -
182 - def create_weights(
183 - self,
184 - layer: torch.nn.Module,
185 - input_size_per_partition: int,
186 - output_partition_sizes: list[int],
187 - input_size: int,
188 - output_size: int,
189 - params_dtype: torch.dtype,
190 - **extra_weight_attrs,
191 - ):
192 - # Normalize group_size
193 - if self.quant_config.group_size != -1:
194 - group_size = self.quant_config.group_size
195 - else:
196 - group_size = input_size
197 -
198 - if input_size_per_partition % group_size != 0:
199 - raise ValueError(
200 - "The input size is not aligned with the quantized "
201 - "weight shape. This can be caused by too large "
202 - "tensor parallel size."
203 - )
204 -
205 - output_size_per_partition = sum(output_partition_sizes)
206 - if output_size_per_partition % self.quant_config.pack_factor != 0:
207 - raise ValueError(
208 - "The output size is not aligned with the quantized "
209 - "weight shape. This can be caused by too large "
210 - "tensor parallel size."
211 - )
212 -
213 - weight_loader = extra_weight_attrs.get("weight_loader")
214 - qweight = PackedvLLMParameter(
215 - data=torch.empty(
216 - input_size_per_partition,
217 - output_size_per_partition // self.quant_config.pack_factor,
218 - dtype=torch.int32,
219 - ),
220 - input_dim=0,
221 - output_dim=1,
222 - packed_dim=1,
223 - packed_factor=self.quant_config.pack_factor,
224 - weight_loader=weight_loader,
225 - )
226 -
227 - num_groups = input_size_per_partition // group_size
228 -
229 - qzeros = PackedvLLMParameter(
230 - data=torch.empty(
231 - num_groups,
232 - output_size_per_partition // self.quant_config.pack_factor,
233 - dtype=torch.int32,
234 - ),
235 - input_dim=0,
236 - output_dim=1,
237 - packed_dim=1,
238 - packed_factor=self.quant_config.pack_factor,
239 - weight_loader=weight_loader,
240 - )
241 -
242 - scales = GroupQuantScaleParameter(
243 - data=torch.empty(
244 - num_groups,
245 - output_size_per_partition,
246 - dtype=params_dtype,
247 - ),
248 - input_dim=0,
249 - output_dim=1,
250 - weight_loader=weight_loader,
251 - )
252 -
253 - layer.register_parameter("qweight", qweight)
254 - layer.register_parameter("qzeros", qzeros)
255 - layer.register_parameter("scales", scales)
256 -
257 - def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
258 - layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False)
259 - layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False)
260 - layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
261 -
262 - def apply(
263 - self,
264 - layer: torch.nn.Module,
265 - x: torch.Tensor,
266 - bias: torch.Tensor | None = None,
267 - ) -> torch.Tensor:
268 - qweight = layer.qweight
269 - scales = layer.scales
270 - qzeros = layer.qzeros
271 - pack_factor = self.quant_config.pack_factor
272 - out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,)
273 - reshaped_x = x.reshape(-1, x.shape[-1])
274 -
275 - # num_tokens >= threshold
276 - FP16_MATMUL_HEURISTIC_CONDITION = x.shape[:-1].numel() >= 256
277 - # Batch invariant mode requires torch.matmul path
278 - # for Triton override
279 - if FP16_MATMUL_HEURISTIC_CONDITION or envs.VLLM_BATCH_INVARIANT:
280 - out = ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
281 - out = torch.matmul(reshaped_x, out)
282 - else:
283 - out = ops.awq_gemm(reshaped_x, qweight, scales, qzeros, pack_factor)
284 - if bias is not None:
285 - out.add_(bias)
286 - return out.reshape(out_shape)
vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py
+11 -9
@@ -8,9 +8,8 @@ import torch
8 from torch.nn.parameter import Parameter
9
10 from vllm.logger import init_logger
11 +from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
12 from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
12 -from vllm.model_executor.layers.quantization.awq import AWQConfig
13 -from vllm.model_executor.layers.quantization.awq_marlin import AWQMarlinConfig
13 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
14 check_marlin_supported,
15 )
@@ -125,12 +124,12 @@ class INCWNA16LinearScheme(INCLinearScheme):
124 )
125
126 if use_marlin:
128 - from vllm.model_executor.layers.quantization.awq_marlin import (
129 - AWQMarlinLinearMethod,
127 + from vllm.model_executor.layers.quantization.auto_awq import (
128 + AutoAWQMarlinLinearMethod,
129 )
130
132 - return AWQMarlinLinearMethod(
133 - AWQMarlinConfig(
131 + return AutoAWQMarlinLinearMethod(
132 + AutoAWQConfig(
133 weight_bits=self.layer_config.bits,
134 group_size=self.layer_config.group_size,
135 zero_point=not self.layer_config.sym,
@@ -140,13 +139,16 @@ class INCWNA16LinearScheme(INCLinearScheme):
139 )
140 )
141
143 - from vllm.model_executor.layers.quantization.awq import AWQLinearMethod
142 + from vllm.model_executor.layers.quantization.auto_awq import (
143 + AutoAWQLinearMethod,
144 + )
145
145 - return AWQLinearMethod(
146 - AWQConfig(
146 + return AutoAWQLinearMethod(
147 + AutoAWQConfig(
148 weight_bits=self.layer_config.bits,
149 group_size=self.layer_config.group_size,
150 zero_point=not self.layer_config.sym,
151 + lm_head_quantized=False,
152 )
153 )
154
vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py
+4 -4
@@ -4,8 +4,8 @@
4 from typing import TYPE_CHECKING
5
6 from vllm.logger import init_logger
7 +from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
8 from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
8 -from vllm.model_executor.layers.quantization.awq_marlin import AWQMarlinConfig
9 from vllm.platforms import current_platform
10 from vllm.scalar_type import scalar_types
11
@@ -154,7 +154,7 @@ def _resolve_gptq_moe(layer: "torch.nn.Module", layer_config: "INCLayerConfig"):
154
155
156 def _resolve_awq_moe(layer: "torch.nn.Module", layer_config: "INCLayerConfig"):
157 - from vllm.model_executor.layers.quantization.awq_marlin import AWQMarlinMoEMethod
157 + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQMoEMethod
158 from vllm.model_executor.layers.quantization.moe_wna16 import (
159 MoeWNA16Config,
160 MoeWNA16Method,
@@ -177,8 +177,8 @@ def _resolve_awq_moe(layer: "torch.nn.Module", layer_config: "INCLayerConfig"):
177 ) and check_moe_marlin_supports_layer(layer, layer_config.group_size)
178
179 if use_marlin:
180 - return AWQMarlinMoEMethod(
181 - AWQMarlinConfig(
180 + return AutoAWQMoEMethod(
181 + AutoAWQConfig(
182 weight_bits=layer_config.bits,
183 group_size=layer_config.group_size,
184 zero_point=not layer_config.sym,
vllm/model_executor/layers/quantization/moe_wna16.py
+8 -24
@@ -27,9 +27,6 @@ from vllm.model_executor.layers.quantization.base_config import (
27 QuantizationConfig,
28 QuantizeMethodBase,
29 )
30 -from vllm.model_executor.layers.quantization.utils.marlin_utils import (
31 - check_marlin_supports_layer,
32 -)
30 from vllm.model_executor.utils import set_weight_attrs
31 from vllm.platforms import current_platform
32
@@ -55,10 +52,8 @@ class MoeWNA16Config(QuantizationConfig):
52 self.lm_head_quantized = lm_head_quantized
53 self.linear_quant_method = linear_quant_method
54 self.full_config = full_config
58 - self.use_marlin = False
55 # Avoid circular import
60 - from vllm.model_executor.layers.quantization.awq import AWQConfig
61 - from vllm.model_executor.layers.quantization.awq_marlin import AWQMarlinConfig
56 + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
57
58 if self.linear_quant_method == "gptq":
59 pass
@@ -67,7 +62,7 @@ class MoeWNA16Config(QuantizationConfig):
62 device_capability = (
63 -1 if capability_tuple is None else capability_tuple.to_int()
64 )
70 - awq_min_capability = AWQConfig.get_min_capability()
65 + awq_min_capability = AutoAWQConfig.get_min_capability()
66 if device_capability < awq_min_capability:
67 raise ValueError(
68 "The quantization method moe_wna16 + awq is not supported "
@@ -75,7 +70,6 @@ class MoeWNA16Config(QuantizationConfig):
70 f"Minimum capability: {awq_min_capability}. "
71 f"Current capability: {device_capability}."
72 )
78 - self.use_marlin = AWQMarlinConfig.is_awq_marlin_compatible(full_config)
73 else:
74 raise ValueError("moe_wna16 only support gptq and awq.")
75
@@ -148,9 +142,9 @@ class MoeWNA16Config(QuantizationConfig):
142 -1 if capability_tuple is None else capability_tuple.to_int()
143 )
144 # Avoid circular import
151 - from vllm.model_executor.layers.quantization.awq import AWQConfig
145 + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
146
153 - awq_min_capability = AWQConfig.get_min_capability()
147 + awq_min_capability = AutoAWQConfig.get_min_capability()
148
149 gptq_compatible = quant_method == "gptq" and not desc_act and num_bits in [4, 8]
150 awq_compatible = (
@@ -170,29 +164,19 @@ class MoeWNA16Config(QuantizationConfig):
164 return UnquantizedLinearMethod()
165 elif isinstance(layer, LinearBase):
166 # Avoid circular import
167 + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
168 from vllm.model_executor.layers.quantization.auto_gptq import (
169 AutoGPTQConfig,
170 )
176 - from vllm.model_executor.layers.quantization.awq import AWQConfig
177 - from vllm.model_executor.layers.quantization.awq_marlin import (
178 - AWQMarlinConfig,
179 - )
171
172 if self.linear_quant_method == "gptq":
173 return AutoGPTQConfig.from_config(self.full_config).get_quant_method(
174 layer, prefix
175 )
176 elif self.linear_quant_method in ("awq", "awq_marlin"):
186 - if self.use_marlin and check_marlin_supports_layer(
187 - layer, self.group_size
188 - ):
189 - return AWQMarlinConfig.from_config(
190 - self.full_config
191 - ).get_quant_method(layer, prefix)
192 - else:
193 - return AWQConfig.from_config(self.full_config).get_quant_method(
194 - layer, prefix
195 - )
177 + return AutoAWQConfig.from_config(self.full_config).get_quant_method(
178 + layer, prefix
179 + )
180 else:
181 raise ValueError("moe_wna16 only support gptq and awq.")
182 elif isinstance(layer, RoutedExperts):
vllm/model_executor/models/cohere2_vision.py
+2 -2
@@ -26,7 +26,7 @@ from vllm.model_executor.layers.linear import (
26 RowParallelLinear,
27 )
28 from vllm.model_executor.layers.quantization import QuantizationConfig
29 -from vllm.model_executor.layers.quantization.awq import AWQConfig
29 +from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
30 from vllm.multimodal import MULTIMODAL_REGISTRY
31 from vllm.multimodal.inputs import (
32 MultiModalFieldConfig,
@@ -420,7 +420,7 @@ class Cohere2VisionForConditionalGeneration(
420 ):
421 # the awq models from OpenGVLab missing `modules_to_not_convert`
422 # patch the quant_config to add `modules_to_not_convert` back
423 - if isinstance(quant_config, AWQConfig):
423 + if isinstance(quant_config, AutoAWQConfig):
424 text_config = config.text_config
425 llm_quant_config = getattr(text_config, "quantization_config", None)
426 if (not quant_config.modules_to_not_convert) and (
vllm/model_executor/models/internvl.py
+2 -2
@@ -20,7 +20,7 @@ from vllm.config import VllmConfig
20 from vllm.config.multimodal import BaseDummyOptions
21 from vllm.inputs import MultiModalDataDict
22 from vllm.model_executor.layers.quantization import QuantizationConfig
23 -from vllm.model_executor.layers.quantization.awq import AWQConfig
23 +from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
24 from vllm.model_executor.models.intern_vit import (
25 InternVisionModel,
26 )
@@ -608,7 +608,7 @@ class InternVLChatModel(
608 ):
609 # the awq models from OpenGVLab missing `modules_to_not_convert`
610 # patch the quant_config to add `modules_to_not_convert` back
611 - if isinstance(quant_config, AWQConfig):
611 + if isinstance(quant_config, AutoAWQConfig):
612 text_config = config.text_config
613 llm_quant_config = getattr(text_config, "quantization_config", None)
614 if (not quant_config.modules_to_not_convert) and (
vllm/model_executor/models/nemotron_vl.py
+2 -2
@@ -11,7 +11,7 @@ from vllm.config import VllmConfig
11 from vllm.model_executor.layers.linear import ReplicatedLinear
12 from vllm.model_executor.layers.pooler import DispatchPooler
13 from vllm.model_executor.layers.quantization import QuantizationConfig
14 -from vllm.model_executor.layers.quantization.awq import AWQConfig
14 +from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
15 from vllm.model_executor.models.internvl import (
16 BaseInternVLDummyInputsBuilder,
17 BaseInternVLMultiModalProcessor,
@@ -144,7 +144,7 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor
144 ):
145 # the awq models from OpenGVLab missing `modules_to_not_convert`
146 # patch the quant_config to add `modules_to_not_convert` back
147 - if isinstance(quant_config, AWQConfig):
147 + if isinstance(quant_config, AutoAWQConfig):
148 text_config = config.get_text_config()
149 llm_quant_config = getattr(text_config, "quantization_config", None)
150 if (not quant_config.modules_to_not_convert) and (
vllm/model_executor/models/skyworkr1v.py
+2 -2
@@ -19,7 +19,7 @@ from vllm.config.multimodal import BaseDummyOptions
19 from vllm.inputs import MultiModalDataDict
20 from vllm.model_executor.layers.linear import ReplicatedLinear
21 from vllm.model_executor.layers.quantization import QuantizationConfig
22 -from vllm.model_executor.layers.quantization.awq import AWQConfig
22 +from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
23 from vllm.model_executor.models.intern_vit import (
24 InternVisionModel,
25 )
@@ -205,7 +205,7 @@ class SkyworkR1VChatModel(nn.Module, SupportsMultiModal, SupportsPP):
205 ):
206 # the awq models from OpenGVLab missing `modules_to_not_convert`
207 # patch the quant_config to add `modules_to_not_convert` back
208 - if isinstance(quant_config, AWQConfig):
208 + if isinstance(quant_config, AutoAWQConfig):
209 text_config = config.text_config
210 llm_quant_config = getattr(text_config, "quantization_config", None)
211 if (not quant_config.modules_to_not_convert) and (
vllm/platforms/rocm.py
+1
@@ -440,6 +440,7 @@ class RocmPlatform(Platform):
440
441 supported_quantization: list[str] = [
442 "awq",
443 + "auto_awq",
444 "awq_marlin", # will be overwritten with awq
445 "gptq",
446 "gptq_marlin",