| 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 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 ( |
| 16 | MPLinearLayerConfig, |
| 17 | choose_mp_linear_kernel, |
| 18 | ) |
| 19 | from vllm.model_executor.layers.fused_moe import ( |
| 20 | FusedMoEConfig, |
| 21 | FusedMoEMethodBase, |
| 22 | FusedMoEQuantConfig, |
| 23 | FusedMoeWeightScaleSupported, |
| 24 | RoutedExperts, |
| 25 | SharedExperts, |
| 26 | UnquantizedFusedMoEMethod, |
| 27 | ) |
| 28 | from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( |
| 29 | WNA16MoEBackend, |
| 30 | convert_to_wna16_moe_kernel_format, |
| 31 | make_wna16_moe_kernel, |
| 32 | make_wna16_moe_quant_config, |
| 33 | select_wna16_moe_backend, |
| 34 | ) |
| 35 | from vllm.model_executor.layers.linear import ( |
| 36 | LinearBase, |
| 37 | LinearMethodBase, |
| 38 | UnquantizedLinearMethod, |
| 39 | set_weight_attrs, |
| 40 | ) |
| 41 | from vllm.model_executor.layers.quantization.base_config import ( |
| 42 | QuantizationConfig, |
| 43 | QuantizeMethodBase, |
| 44 | ) |
| 45 | from vllm.model_executor.layers.quantization.utils import replace_parameter |
| 46 | from vllm.model_executor.layers.quantization.utils.marlin_utils import ( |
| 47 | check_marlin_supported, |
| 48 | check_marlin_supports_layer, |
| 49 | check_moe_marlin_supports_layer, |
| 50 | get_marlin_input_dtype, |
| 51 | verify_marlin_supported, |
| 52 | ) |
| 53 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 54 | is_layer_skipped, |
| 55 | kInt4Static, |
| 56 | ) |
| 57 | from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead |
| 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 |
| 65 | |
| 66 | if TYPE_CHECKING: |
| 67 | from vllm.model_executor.layers.quantization import QuantizationMethods |
| 68 | from vllm.model_executor.models.utils import WeightsMapper |
| 69 | |
| 70 | logger = init_logger(__name__) |
| 71 | |
| 72 | # AWQ uses a non-standard packing order within int32 values. |
| 73 | # For 4-bit: standard order stores values at bit positions [0,4,8,12,16,20,24,28] |
| 74 | # for indices [0,1,2,3,4,5,6,7], while AWQ stores them for indices |
| 75 | # [0,4,1,5,2,6,3,7]. This permutation reverses that ordering. |
| 76 | _REVERSE_AWQ_PACK_ORDER = [0, 4, 1, 5, 2, 6, 3, 7] |
| 77 | |
| 78 | |
| 79 | def _replace_or_register_parameter( |
| 80 | layer: torch.nn.Module, |
| 81 | name: str, |
| 82 | value: torch.Tensor | None, |
| 83 | ) -> None: |
| 84 | if value is None: |
| 85 | return |
| 86 | if hasattr(layer, name): |
| 87 | replace_parameter(layer, name, value) |
| 88 | else: |
| 89 | layer.register_parameter(name, Parameter(value, requires_grad=False)) |
| 90 | |
| 91 | |
| 92 | def _convert_awq_to_standard_format( |
| 93 | layer: torch.nn.Module, |
| 94 | w_q_name: str, |
| 95 | w_zp_name: str | None, |
| 96 | size_bits: int, |
| 97 | ) -> None: |
| 98 | """Convert AWQ weight and zero-point tensors to standard GPTQ-like format. |
| 99 | |
| 100 | AWQ packs qweight along the output dim with a non-standard bit order. |
| 101 | This converts to standard bit order and repacks qweight along the input |
| 102 | dim, matching the format expected by the MPLinearKernel framework. |
| 103 | If w_zp_name is None (symmetric quantization), only the weight is converted. |
| 104 | """ |
| 105 | pack_factor = 32 // size_bits |
| 106 | mask = (1 << size_bits) - 1 |
| 107 | device = getattr(layer, w_q_name).device |
| 108 | reverse_order = torch.tensor( |
| 109 | _REVERSE_AWQ_PACK_ORDER, dtype=torch.long, device=device |
| 110 | ) |
| 111 | shifts = torch.arange(0, 32, size_bits, dtype=torch.int32, device=device) |
| 112 | |
| 113 | # --- Convert qweight: (K, N // pack) packed_dim=1 → (K // pack, N) packed_dim=0 |
| 114 | qw = getattr(layer, w_q_name).data |
| 115 | K, N_packed = qw.shape |
| 116 | N = N_packed * pack_factor |
| 117 | |
| 118 | # Unpack int32 → individual values, fix AWQ ordering |
| 119 | unpacked = (qw.unsqueeze(-1) >> shifts) & mask # (K, N_packed, pack_factor) |
| 120 | unpacked = unpacked[:, :, reverse_order] |
| 121 | unpacked = unpacked.reshape(K, N) # (K, N) |
| 122 | |
| 123 | # Repack along input dim (dim 0) |
| 124 | unpacked = unpacked.reshape(K // pack_factor, pack_factor, N) |
| 125 | new_qw = (unpacked.to(torch.int32) << shifts[None, :, None]).sum( |
| 126 | dim=1, dtype=torch.int32 |
| 127 | ) |
| 128 | |
| 129 | def _noop_loader(*args, **kwargs): |
| 130 | pass |
| 131 | |
| 132 | new_param = PackedvLLMParameter( |
| 133 | data=new_qw.contiguous(), |
| 134 | input_dim=0, |
| 135 | output_dim=1, |
| 136 | packed_dim=0, |
| 137 | packed_factor=pack_factor, |
| 138 | weight_loader=_noop_loader, |
| 139 | ) |
| 140 | setattr(layer, w_q_name, new_param) |
| 141 | |
| 142 | if w_zp_name is None: |
| 143 | return |
| 144 | |
| 145 | # --- Convert qzeros: fix AWQ bit ordering and repack |
| 146 | # AWQ qzeros: (G, N // pack) packed along dim 1, AWQ bit order |
| 147 | # Target: (N // pack, G) packed along dim 0, standard bit order |
| 148 | # This matches the CompressedTensors layout expected by the kernels. |
| 149 | qz = getattr(layer, w_zp_name).data |
| 150 | G, _ = qz.shape |
| 151 | |
| 152 | unpacked_zp = (qz.unsqueeze(-1) >> shifts) & mask # (G, N_packed, pack_factor) |
| 153 | unpacked_zp = unpacked_zp[:, :, reverse_order] |
| 154 | unpacked_zp = unpacked_zp.reshape(G, N) # (G, N) individual values |
| 155 | |
| 156 | # Transpose and repack along dim 0 (output dim) |
| 157 | unpacked_zp = unpacked_zp.T # (N, G) |
| 158 | unpacked_zp = unpacked_zp.reshape(N // pack_factor, pack_factor, G) |
| 159 | new_qz = (unpacked_zp.to(torch.int32) << shifts[None, :, None]).sum( |
| 160 | dim=1, dtype=torch.int32 |
| 161 | ) |
| 162 | |
| 163 | new_zp_param = PackedvLLMParameter( |
| 164 | data=new_qz.contiguous(), |
| 165 | output_dim=0, |
| 166 | input_dim=1, |
| 167 | packed_dim=0, |
| 168 | packed_factor=pack_factor, |
| 169 | weight_loader=_noop_loader, |
| 170 | ) |
| 171 | setattr(layer, w_zp_name, new_zp_param) |
| 172 | |
| 173 | |
| 174 | class AutoAWQConfig(QuantizationConfig): |
| 175 | """Config class for AutoAWQ quantization. |
| 176 | |
| 177 | Unified config that supports multiple backends: Triton, Marlin, and XPU. |
| 178 | Reference: https://arxiv.org/abs/2306.00978 |
| 179 | """ |
| 180 | |
| 181 | # num_bits -> type |
| 182 | TYPE_MAP = { |
| 183 | 4: scalar_types.uint4, |
| 184 | } |
| 185 | |
| 186 | def __init__( |
| 187 | self, |
| 188 | weight_bits: int, |
| 189 | group_size: int, |
| 190 | zero_point: bool, |
| 191 | lm_head_quantized: bool, |
| 192 | modules_to_not_convert: list[str] | None = None, |
| 193 | full_config: dict[str, Any] | None = None, |
| 194 | ) -> None: |
| 195 | super().__init__() |
| 196 | self.pack_factor = 32 // weight_bits # packed into int32 |
| 197 | self.group_size = group_size |
| 198 | self.zero_point = zero_point |
| 199 | self.lm_head_quantized = lm_head_quantized |
| 200 | self.weight_bits = weight_bits |
| 201 | self.modules_to_not_convert = modules_to_not_convert or [] |
| 202 | self.full_config = full_config or {} |
| 203 | |
| 204 | if self.weight_bits not in self.TYPE_MAP: |
| 205 | supported = ", ".join(str(k) for k in self.TYPE_MAP) |
| 206 | raise ValueError( |
| 207 | f"Unsupported num_bits = {self.weight_bits}. " |
| 208 | f"Supported: {supported}. " |
| 209 | f"For 8-bit AWQ, use Marlin backend by setting " |
| 210 | f"backend='awq:marlin' or backend='marlin'." |
| 211 | ) |
| 212 | |
| 213 | self.quant_type = self.TYPE_MAP[self.weight_bits] |
| 214 | |
| 215 | def __repr__(self) -> str: |
| 216 | return ( |
| 217 | f"AutoAWQConfig(quant_type={self.quant_type}, " |
| 218 | f"group_size={self.group_size}, " |
| 219 | f"zero_point={self.zero_point}, " |
| 220 | f"lm_head_quantized={self.lm_head_quantized}, " |
| 221 | f"modules_to_not_convert={self.modules_to_not_convert})" |
| 222 | ) |
| 223 | |
| 224 | @classmethod |
| 225 | def get_name(cls) -> "QuantizationMethods": |
| 226 | return "auto_awq" |
| 227 | |
| 228 | @classmethod |
| 229 | def get_supported_act_dtypes(cls) -> list[torch.dtype]: |
| 230 | return [torch.half, torch.bfloat16] |
| 231 | |
| 232 | @classmethod |
| 233 | def get_min_capability(cls) -> int: |
| 234 | return 75 |
| 235 | |
| 236 | @classmethod |
| 237 | def get_config_filenames(cls) -> list[str]: |
| 238 | return ["quantize_config.json", "quant_config.json"] |
| 239 | |
| 240 | @classmethod |
| 241 | def from_config(cls, config: dict[str, Any]) -> "AutoAWQConfig": |
| 242 | weight_bits = cls.get_from_keys(config, ["w_bit", "bits"]) |
| 243 | group_size = cls.get_from_keys(config, ["q_group_size", "group_size"]) |
| 244 | zero_point = cls.get_from_keys(config, ["zero_point"]) |
| 245 | lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False) |
| 246 | modules_to_not_convert = cls.get_from_keys_or( |
| 247 | config, ["modules_to_not_convert"], None |
| 248 | ) |
| 249 | # Ensure full_config uses "awq" as quant_method for MoE fallback compatibility. |
| 250 | # MoeWNA16Config only accepts "gptq" or "awq", so we normalize here. |
| 251 | full_config = config.copy() |
| 252 | full_config["quant_method"] = "awq" |
| 253 | return cls( |
| 254 | weight_bits, |
| 255 | group_size, |
| 256 | zero_point, |
| 257 | lm_head_quantized, |
| 258 | modules_to_not_convert, |
| 259 | full_config, |
| 260 | ) |
| 261 | |
| 262 | @classmethod |
| 263 | def override_quantization_method( |
| 264 | cls, hf_quant_cfg, user_quant, hf_config=None |
| 265 | ) -> "QuantizationMethods | None": |
| 266 | """Override to use AutoAWQ for compatible AWQ models.""" |
| 267 | # Don't override on CPU - let cpu_awq handle it |
| 268 | if current_platform.is_cpu(): |
| 269 | return None |
| 270 | |
| 271 | quant_method = hf_quant_cfg.get("quant_method", "").lower() |
| 272 | |
| 273 | if quant_method != "awq": |
| 274 | return None |
| 275 | |
| 276 | is_valid_user_quant = user_quant is None or user_quant in ( |
| 277 | "awq", |
| 278 | "awq_marlin", |
| 279 | "auto_awq", |
| 280 | "marlin", |
| 281 | ) |
| 282 | |
| 283 | if is_valid_user_quant: |
| 284 | return cls.get_name() |
| 285 | |
| 286 | return None |
| 287 | |
| 288 | def get_quant_method( |
| 289 | self, layer: torch.nn.Module, prefix: str |
| 290 | ) -> Union["LinearMethodBase", "QuantizeMethodBase"] | None: |
| 291 | if isinstance(layer, LinearBase) or ( |
| 292 | isinstance(layer, ParallelLMHead) and self.lm_head_quantized |
| 293 | ): |
| 294 | if is_layer_skipped( |
| 295 | prefix, |
| 296 | self.modules_to_not_convert, |
| 297 | self.packed_modules_mapping, |
| 298 | match_mode="substring", |
| 299 | ): |
| 300 | return UnquantizedLinearMethod() |
| 301 | |
| 302 | # On CPU and XPU, use the Marlin linear method which uses |
| 303 | # choose_mp_linear_kernel to select the best available kernel |
| 304 | # (CPUWNA16LinearKernel on CPU, XPUwNa16LinearKernel on XPU). |
| 305 | # The AWQ checkpoint is converted to the standard GPTQ-like format |
| 306 | # in process_weights_after_loading before being handed to the kernel. |
| 307 | if current_platform.is_cpu() or current_platform.is_xpu(): |
| 308 | return AutoAWQMarlinLinearMethod(self) |
| 309 | |
| 310 | # Check if Marlin is supported and not using batch invariant mode |
| 311 | # (Marlin kernels are not batch invariant) |
| 312 | use_marlin = ( |
| 313 | not envs.VLLM_BATCH_INVARIANT |
| 314 | and current_platform.is_cuda() |
| 315 | and check_marlin_supported( |
| 316 | self.quant_type, self.group_size, self.zero_point |
| 317 | ) |
| 318 | ) |
| 319 | |
| 320 | if use_marlin: |
| 321 | # tile-misaligned shapes are fixed by padding at weight prep |
| 322 | if not check_marlin_supports_layer( |
| 323 | layer, self.group_size, allow_tile_padding=True |
| 324 | ): |
| 325 | logger.warning_once( |
| 326 | "Layer '%s' is not supported by AutoAWQMarlin. " |
| 327 | "Falling back to unoptimized AWQ kernels.", |
| 328 | prefix, |
| 329 | ) |
| 330 | return AutoAWQLinearMethod(self) |
| 331 | quant_method = AutoAWQMarlinLinearMethod(self) |
| 332 | quant_method.input_dtype = get_marlin_input_dtype(prefix) |
| 333 | return quant_method |
| 334 | |
| 335 | return AutoAWQLinearMethod(self) |
| 336 | |
| 337 | elif isinstance(layer, RoutedExperts): |
| 338 | if is_layer_skipped( |
| 339 | prefix, |
| 340 | getattr(self, "modules_to_not_convert", []), |
| 341 | match_mode="substring", |
| 342 | ): |
| 343 | return UnquantizedFusedMoEMethod(layer.moe_config) |
| 344 | |
| 345 | if not check_moe_marlin_supports_layer( |
| 346 | layer, self.group_size, allow_tile_padding=True |
| 347 | ): |
| 348 | logger.warning_once( |
| 349 | f"Layer '{prefix}' is not supported by AutoAWQMoEMarlin. " |
| 350 | "Falling back to Moe WNA16 kernels." |
| 351 | ) |
| 352 | from vllm.model_executor.layers.quantization.moe_wna16 import ( |
| 353 | MoeWNA16Config, |
| 354 | ) |
| 355 | |
| 356 | return MoeWNA16Config.from_config(self.full_config).get_quant_method( |
| 357 | layer, prefix |
| 358 | ) |
| 359 | |
| 360 | return AutoAWQMoEMethod(self, layer.moe_config) |
| 361 | |
| 362 | return None |
| 363 | |
| 364 | def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"): |
| 365 | if self.modules_to_not_convert: |
| 366 | self.modules_to_not_convert = hf_to_vllm_mapper.apply_list( |
| 367 | self.modules_to_not_convert |
| 368 | ) |
| 369 | |
| 370 | def maybe_update_config( |
| 371 | self, |
| 372 | model_name: str, |
| 373 | hf_config: PreTrainedConfig | None = None, |
| 374 | revision: str | None = None, |
| 375 | ): |
| 376 | if self.modules_to_not_convert: |
| 377 | return |
| 378 | |
| 379 | unquant_dtypes = [torch.float16, torch.bfloat16, torch.float32] |
| 380 | metadata = get_safetensors_params_metadata(model_name, revision=revision) |
| 381 | layers = {param_name.rsplit(".", 1)[0] for param_name in metadata} |
| 382 | quant_layers: set[str] = { |
| 383 | param_name.rsplit(".", 1)[0] |
| 384 | for param_name, info in metadata.items() |
| 385 | if (dtype := info.get("dtype", None)) |
| 386 | and _SAFETENSORS_TO_TORCH_DTYPE[dtype] not in unquant_dtypes |
| 387 | } |
| 388 | self.modules_to_not_convert = list(layers - quant_layers) |
| 389 | |
| 390 | |
| 391 | class AutoAWQMarlinLinearMethod(LinearMethodBase): |
| 392 | """Linear method for AWQ Marlin. |
| 393 | |
| 394 | Uses choose_mp_linear_kernel to select the best available kernel |
| 395 | (Conch, Exllama, or Marlin) for the current platform. |
| 396 | |
| 397 | Args: |
| 398 | quant_config: The AWQ Marlin quantization config. |
| 399 | |
| 400 | """ |
| 401 | |
| 402 | _kernel_backends_being_used: set[str] = set() |
| 403 | |
| 404 | def __init__(self, quant_config: AutoAWQConfig) -> None: |
| 405 | self.quant_config = quant_config |
| 406 | self.quant_type = scalar_types.uint4 |
| 407 | self.input_dtype = None |
| 408 | |
| 409 | # Skip Marlin verification on CPU/XPU - they use dedicated WNA16 |
| 410 | # kernels (CPUWNA16LinearKernel / XPUwNa16LinearKernel) instead. |
| 411 | if not (current_platform.is_cpu() or current_platform.is_xpu()): |
| 412 | verify_marlin_supported( |
| 413 | quant_type=self.quant_config.quant_type, |
| 414 | group_size=self.quant_config.group_size, |
| 415 | has_zp=self.quant_config.zero_point, |
| 416 | ) |
| 417 | |
| 418 | def create_weights( |
| 419 | self, |
| 420 | layer: torch.nn.Module, |
| 421 | input_size_per_partition: int, |
| 422 | output_partition_sizes: list[int], |
| 423 | input_size: int, |
| 424 | output_size: int, |
| 425 | params_dtype: torch.dtype, |
| 426 | **extra_weight_attrs, |
| 427 | ) -> None: |
| 428 | output_size_per_partition = sum(output_partition_sizes) |
| 429 | weight_loader = extra_weight_attrs.get("weight_loader") |
| 430 | |
| 431 | if self.quant_config.group_size != -1: |
| 432 | group_size = self.quant_config.group_size |
| 433 | else: |
| 434 | group_size = input_size |
| 435 | |
| 436 | mp_linear_kernel_config = MPLinearLayerConfig( |
| 437 | full_weight_shape=(input_size, output_size), |
| 438 | partition_weight_shape=( |
| 439 | input_size_per_partition, |
| 440 | output_size_per_partition, |
| 441 | ), |
| 442 | weight_type=self.quant_config.quant_type, |
| 443 | act_type=params_dtype if self.input_dtype is None else self.input_dtype, |
| 444 | group_size=self.quant_config.group_size, |
| 445 | zero_points=self.quant_config.zero_point, |
| 446 | ) |
| 447 | |
| 448 | kernel_type = choose_mp_linear_kernel(mp_linear_kernel_config) |
| 449 | |
| 450 | if kernel_type.__name__ not in self._kernel_backends_being_used: |
| 451 | logger.info("Using %s for AutoAWQMarlinLinearMethod", kernel_type.__name__) |
| 452 | self._kernel_backends_being_used.add(kernel_type.__name__) |
| 453 | |
| 454 | # Weights are loaded in AWQ checkpoint format (packed along output dim). |
| 455 | # Conversion to GPTQ-like format happens in process_weights_after_loading. |
| 456 | qweight = PackedvLLMParameter( |
| 457 | data=torch.empty( |
| 458 | input_size_per_partition, |
| 459 | output_size_per_partition // self.quant_config.pack_factor, |
| 460 | dtype=torch.int32, |
| 461 | ), |
| 462 | input_dim=0, |
| 463 | output_dim=1, |
| 464 | packed_dim=1, |
| 465 | packed_factor=self.quant_config.pack_factor, |
| 466 | weight_loader=weight_loader, |
| 467 | ) |
| 468 | |
| 469 | num_groups = input_size_per_partition // group_size |
| 470 | |
| 471 | qzeros = PackedvLLMParameter( |
| 472 | data=torch.empty( |
| 473 | num_groups, |
| 474 | output_size_per_partition // self.quant_config.pack_factor, |
| 475 | dtype=torch.int32, |
| 476 | ), |
| 477 | input_dim=0, |
| 478 | output_dim=1, |
| 479 | packed_dim=1, |
| 480 | packed_factor=self.quant_config.pack_factor, |
| 481 | weight_loader=weight_loader, |
| 482 | ) |
| 483 | |
| 484 | scales = GroupQuantScaleParameter( |
| 485 | data=torch.empty( |
| 486 | num_groups, |
| 487 | output_size_per_partition, |
| 488 | dtype=params_dtype, |
| 489 | ), |
| 490 | input_dim=0, |
| 491 | output_dim=1, |
| 492 | weight_loader=weight_loader, |
| 493 | ) |
| 494 | |
| 495 | layer.register_parameter("qweight", qweight) |
| 496 | layer.register_parameter("qzeros", qzeros) |
| 497 | layer.register_parameter("scales", scales) |
| 498 | |
| 499 | self.kernel = kernel_type( |
| 500 | mp_linear_kernel_config, |
| 501 | w_q_param_name="qweight", |
| 502 | w_s_param_name="scales", |
| 503 | w_zp_param_name="qzeros", |
| 504 | ) |
| 505 | |
| 506 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 507 | # AWQ checkpoints use a non-standard packing order and pack qweight |
| 508 | # along the output dimension. Convert to the standard format |
| 509 | # (GPTQ-like: standard bit order, qweight packed along input dim) |
| 510 | # before handing off to the kernel. |
| 511 | _convert_awq_to_standard_format( |
| 512 | layer, "qweight", "qzeros", self.quant_config.quant_type.size_bits |
| 513 | ) |
| 514 | self.kernel.process_weights_after_loading(layer) |
| 515 | |
| 516 | def apply( |
| 517 | self, |
| 518 | layer: torch.nn.Module, |
| 519 | x: torch.Tensor, |
| 520 | bias: torch.Tensor | None = None, |
| 521 | ) -> torch.Tensor: |
| 522 | return self.kernel.apply_weights(layer, x, bias) |
| 523 | |
| 524 | |
| 525 | class AutoAWQMoEMethod(FusedMoEMethodBase): |
| 526 | def __init__( |
| 527 | self, |
| 528 | quant_config: AutoAWQConfig, |
| 529 | moe: FusedMoEConfig, |
| 530 | ): |
| 531 | super().__init__(moe) |
| 532 | self.quant_config = quant_config |
| 533 | if self.quant_config.weight_bits != 4: |
| 534 | raise ValueError("AutoAWQMoEMethod only supports 4bit now.") |
| 535 | self.quant_type = scalar_types.uint4 |
| 536 | self.input_dtype = None |
| 537 | self.use_marlin = True |
| 538 | self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend( |
| 539 | moe, |
| 540 | kInt4Static, |
| 541 | quant_config=self.quant_config, |
| 542 | may_have_zp=self.quant_config.zero_point, |
| 543 | may_have_bias=True, |
| 544 | allow_tile_padding=True, |
| 545 | ) |
| 546 | |
| 547 | def create_weights( |
| 548 | self, |
| 549 | layer: RoutedExperts, |
| 550 | num_experts: int, |
| 551 | hidden_size: int, |
| 552 | intermediate_size_per_partition: int, |
| 553 | params_dtype: torch.dtype, |
| 554 | **extra_weight_attrs, |
| 555 | ): |
| 556 | layer.input_dtype = self.input_dtype |
| 557 | extra_weight_attrs.update( |
| 558 | { |
| 559 | "is_transposed": True, |
| 560 | "quant_method": FusedMoeWeightScaleSupported.GROUP.value, |
| 561 | } |
| 562 | ) |
| 563 | |
| 564 | w13_qweight = Parameter( |
| 565 | torch.empty( |
| 566 | num_experts, |
| 567 | hidden_size, |
| 568 | self.moe.w13_num_shards |
| 569 | * intermediate_size_per_partition |
| 570 | // self.quant_config.pack_factor, |
| 571 | dtype=torch.int32, |
| 572 | ), |
| 573 | requires_grad=False, |
| 574 | ) |
| 575 | layer.register_parameter("w13_qweight", w13_qweight) |
| 576 | set_weight_attrs(w13_qweight, extra_weight_attrs) |
| 577 | |
| 578 | w2_qweight = Parameter( |
| 579 | torch.empty( |
| 580 | num_experts, |
| 581 | intermediate_size_per_partition, |
| 582 | hidden_size // self.quant_config.pack_factor, |
| 583 | dtype=torch.int32, |
| 584 | ), |
| 585 | requires_grad=False, |
| 586 | ) |
| 587 | layer.register_parameter("w2_qweight", w2_qweight) |
| 588 | set_weight_attrs(w2_qweight, extra_weight_attrs) |
| 589 | |
| 590 | num_groups_w13 = hidden_size // self.quant_config.group_size |
| 591 | num_groups_w2 = intermediate_size_per_partition // self.quant_config.group_size |
| 592 | layer.num_groups_w13 = num_groups_w13 |
| 593 | layer.num_groups_w2 = num_groups_w2 |
| 594 | |
| 595 | # WEIGHT_SCALES |
| 596 | # Allocate 2 scales for w1 and w3 respectively. |
| 597 | w13_scales = Parameter( |
| 598 | torch.empty( |
| 599 | num_experts, |
| 600 | num_groups_w13, |
| 601 | intermediate_size_per_partition * self.moe.w13_num_shards, |
| 602 | dtype=params_dtype, |
| 603 | ), |
| 604 | requires_grad=False, |
| 605 | ) |
| 606 | layer.register_parameter("w13_scales", w13_scales) |
| 607 | set_weight_attrs(w13_scales, extra_weight_attrs) |
| 608 | |
| 609 | w2_scales = Parameter( |
| 610 | torch.empty(num_experts, num_groups_w2, hidden_size, dtype=params_dtype), |
| 611 | requires_grad=False, |
| 612 | ) |
| 613 | layer.register_parameter("w2_scales", w2_scales) |
| 614 | set_weight_attrs(w2_scales, extra_weight_attrs) |
| 615 | |
| 616 | # WEIGHT_ZERO_POINT |
| 617 | # Allocate 2 zero points for w1 and w3 respectively. |
| 618 | w13_qzeros = Parameter( |
| 619 | torch.empty( |
| 620 | num_experts, |
| 621 | num_groups_w13, |
| 622 | self.moe.w13_num_shards |
| 623 | * intermediate_size_per_partition |
| 624 | // self.quant_config.pack_factor, |
| 625 | dtype=torch.int32, |
| 626 | ), |
| 627 | requires_grad=False, |
| 628 | ) |
| 629 | layer.register_parameter("w13_qzeros", w13_qzeros) |
| 630 | set_weight_attrs(w13_qzeros, extra_weight_attrs) |
| 631 | |
| 632 | w2_qzeros = Parameter( |
| 633 | torch.empty( |
| 634 | num_experts, |
| 635 | num_groups_w2, |
| 636 | hidden_size // self.quant_config.pack_factor, |
| 637 | dtype=torch.int32, |
| 638 | ), |
| 639 | requires_grad=False, |
| 640 | ) |
| 641 | layer.register_parameter("w2_qzeros", w2_qzeros) |
| 642 | set_weight_attrs(w2_qzeros, extra_weight_attrs) |
| 643 | |
| 644 | def process_weights_after_loading(self, layer: RoutedExperts) -> None: |
| 645 | converted = convert_to_wna16_moe_kernel_format( |
| 646 | backend=self.wna16_moe_backend, |
| 647 | layer=layer, |
| 648 | quant_config=self.quant_config, |
| 649 | input_dtype=self.input_dtype, |
| 650 | w13=layer.w13_qweight, |
| 651 | w2=layer.w2_qweight, |
| 652 | w13_scale=layer.w13_scales, |
| 653 | w2_scale=layer.w2_scales, |
| 654 | w13_qzeros=layer.w13_qzeros, |
| 655 | w2_qzeros=layer.w2_qzeros, |
| 656 | w13_bias=getattr(layer, "w13_bias", None), |
| 657 | w2_bias=getattr(layer, "w2_bias", None), |
| 658 | ) |
| 659 | |
| 660 | if converted is None: |
| 661 | # Backend rewrote the layer's params in place (e.g. Humming). |
| 662 | self._setup_kernel(layer) |
| 663 | return |
| 664 | |
| 665 | ( |
| 666 | w13, |
| 667 | w2, |
| 668 | w13_scale, |
| 669 | w2_scale, |
| 670 | w13_qzeros, |
| 671 | w2_qzeros, |
| 672 | w13_input_global_scale, |
| 673 | w2_input_global_scale, |
| 674 | w13_bias, |
| 675 | w2_bias, |
| 676 | ) = converted |
| 677 | |
| 678 | replace_parameter(layer, "w13_qweight", w13) |
| 679 | replace_parameter(layer, "w2_qweight", w2) |
| 680 | |
| 681 | # The modular kernel expects w13_weight and w2_weight, |
| 682 | # but AWQ uses w13_qweight and w2_qweight |
| 683 | # Alias for modular kernel |
| 684 | layer.w13_weight = layer.w13_qweight |
| 685 | # Alias for modular kernel |
| 686 | layer.w2_weight = layer.w2_qweight |
| 687 | |
| 688 | replace_parameter(layer, "w13_scales", w13_scale) |
| 689 | replace_parameter(layer, "w2_scales", w2_scale) |
| 690 | _replace_or_register_parameter(layer, "w13_qzeros", w13_qzeros) |
| 691 | _replace_or_register_parameter(layer, "w2_qzeros", w2_qzeros) |
| 692 | _replace_or_register_parameter( |
| 693 | layer, "w13_input_global_scale", w13_input_global_scale |
| 694 | ) |
| 695 | _replace_or_register_parameter( |
| 696 | layer, "w2_input_global_scale", w2_input_global_scale |
| 697 | ) |
| 698 | _replace_or_register_parameter(layer, "w13_bias", w13_bias) |
| 699 | _replace_or_register_parameter(layer, "w2_bias", w2_bias) |
| 700 | |
| 701 | self._setup_kernel(layer) |
| 702 | |
| 703 | def _setup_kernel(self, layer: RoutedExperts) -> None: |
| 704 | """Build the FusedMoEKernel for this layer.""" |
| 705 | self.moe_quant_config = self.get_fused_moe_quant_config(layer) |
| 706 | self.moe_kernel = make_wna16_moe_kernel( |
| 707 | moe_quant_config=self.moe_quant_config, |
| 708 | moe_config=self.moe, |
| 709 | experts_cls=self.experts_cls, |
| 710 | backend=self.wna16_moe_backend, |
| 711 | routing_tables=layer._expert_routing_tables(), |
| 712 | ) |
| 713 | |
| 714 | def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: |
| 715 | if self.wna16_moe_backend == WNA16MoEBackend.HUMMING: |
| 716 | from vllm.model_executor.layers.quantization.utils.humming import ( |
| 717 | get_humming_moe_quant_config, |
| 718 | ) |
| 719 | |
| 720 | return get_humming_moe_quant_config( |
| 721 | layer, |
| 722 | gemm1_alpha=getattr(layer, "swiglu_alpha", None), |
| 723 | gemm1_beta=getattr(layer, "swiglu_beta", None), |
| 724 | gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), |
| 725 | ) |
| 726 | return make_wna16_moe_quant_config( |
| 727 | w1_scale=layer.w13_scales, |
| 728 | w2_scale=layer.w2_scales, |
| 729 | group_size=self.quant_config.group_size, |
| 730 | num_bits=self.quant_config.weight_bits, |
| 731 | w1_zp=getattr(layer, "w13_qzeros", None) |
| 732 | if self.quant_config.zero_point |
| 733 | else None, |
| 734 | w2_zp=getattr(layer, "w2_qzeros", None) |
| 735 | if self.quant_config.zero_point |
| 736 | else None, |
| 737 | w1_bias=getattr(layer, "w13_bias", None), |
| 738 | w2_bias=getattr(layer, "w2_bias", None), |
| 739 | a1_gscale=getattr(layer, "w13_input_global_scale", None), |
| 740 | a2_gscale=getattr(layer, "w2_input_global_scale", None), |
| 741 | gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), |
| 742 | gemm1_alpha=getattr(layer, "swiglu_alpha", None), |
| 743 | gemm1_beta=getattr(layer, "swiglu_beta", None), |
| 744 | ) |
| 745 | |
| 746 | def apply( |
| 747 | self, |
| 748 | layer: RoutedExperts, |
| 749 | x: torch.Tensor, |
| 750 | topk_weights: torch.Tensor, |
| 751 | topk_ids: torch.Tensor, |
| 752 | shared_experts: SharedExperts | None, |
| 753 | shared_experts_input: torch.Tensor | None, |
| 754 | ) -> torch.Tensor: |
| 755 | assert not self.is_monolithic |
| 756 | assert self.moe_kernel is not None |
| 757 | return self.moe_kernel.apply( |
| 758 | hidden_states=x, |
| 759 | w1=layer.w13_weight, |
| 760 | w2=layer.w2_weight, |
| 761 | topk_weights=topk_weights, |
| 762 | topk_ids=topk_ids, |
| 763 | activation=layer.activation, |
| 764 | global_num_experts=layer.global_num_experts, |
| 765 | apply_router_weight_on_input=layer.apply_router_weight_on_input, |
| 766 | expert_map=layer.expert_map, |
| 767 | shared_experts=shared_experts, |
| 768 | shared_experts_input=shared_experts_input, |
| 769 | ) |
| 770 | |
| 771 | def apply_monolithic( |
| 772 | self, |
| 773 | layer: RoutedExperts, |
| 774 | x: torch.Tensor, |
| 775 | router_logits: torch.Tensor, |
| 776 | input_ids: torch.Tensor | None = None, |
| 777 | ) -> torch.Tensor: |
| 778 | assert self.is_monolithic |
| 779 | assert self.moe_kernel is not None |
| 780 | return self.moe_kernel.apply_monolithic( |
| 781 | hidden_states=x, |
| 782 | w1=layer.w13_weight, |
| 783 | w2=layer.w2_weight, |
| 784 | router_logits=router_logits, |
| 785 | activation=layer.activation, |
| 786 | global_num_experts=layer.global_num_experts, |
| 787 | expert_map=layer.expert_map, |
| 788 | apply_router_weight_on_input=layer.apply_router_weight_on_input, |
| 789 | num_expert_group=layer.num_expert_group, |
| 790 | topk_group=layer.topk_group, |
| 791 | e_score_correction_bias=layer.e_score_correction_bias, |
| 792 | routed_scaling_factor=layer.routed_scaling_factor, |
| 793 | ) |
| 794 | |
| 795 | |
| 796 | class BaseAWQLinearMethod(LinearMethodBase): |
| 797 | """Base class for AWQ linear methods with shared weight creation logic.""" |
| 798 | |
| 799 | def __init__(self, quant_config: AutoAWQConfig): |
| 800 | self.quant_config = quant_config |
| 801 | |
| 802 | def create_weights( |
| 803 | self, |
| 804 | layer: torch.nn.Module, |
| 805 | input_size_per_partition: int, |
| 806 | output_partition_sizes: list[int], |
| 807 | input_size: int, |
| 808 | output_size: int, |
| 809 | params_dtype: torch.dtype, |
| 810 | **extra_weight_attrs, |
| 811 | ): |
| 812 | # Normalize group_size |
| 813 | if self.quant_config.group_size != -1: |
| 814 | group_size = self.quant_config.group_size |
| 815 | else: |
| 816 | group_size = input_size |
| 817 | |
| 818 | if input_size_per_partition % group_size != 0: |
| 819 | raise ValueError( |
| 820 | "The input size is not aligned with the quantized " |
| 821 | "weight shape. This can be caused by too large " |
| 822 | "tensor parallel size." |
| 823 | ) |
| 824 | |
| 825 | output_size_per_partition = sum(output_partition_sizes) |
| 826 | if output_size_per_partition % self.quant_config.pack_factor != 0: |
| 827 | raise ValueError( |
| 828 | "The output size is not aligned with the quantized " |
| 829 | "weight shape. This can be caused by too large " |
| 830 | "tensor parallel size." |
| 831 | ) |
| 832 | |
| 833 | weight_loader = extra_weight_attrs.get("weight_loader") |
| 834 | qweight = PackedvLLMParameter( |
| 835 | data=torch.empty( |
| 836 | input_size_per_partition, |
| 837 | output_size_per_partition // self.quant_config.pack_factor, |
| 838 | dtype=torch.int32, |
| 839 | ), |
| 840 | input_dim=0, |
| 841 | output_dim=1, |
| 842 | packed_dim=1, |
| 843 | packed_factor=self.quant_config.pack_factor, |
| 844 | weight_loader=weight_loader, |
| 845 | ) |
| 846 | |
| 847 | num_groups = input_size_per_partition // group_size |
| 848 | |
| 849 | qzeros = PackedvLLMParameter( |
| 850 | data=torch.empty( |
| 851 | num_groups, |
| 852 | output_size_per_partition // self.quant_config.pack_factor, |
| 853 | dtype=torch.int32, |
| 854 | ), |
| 855 | input_dim=0, |
| 856 | output_dim=1, |
| 857 | packed_dim=1, |
| 858 | packed_factor=self.quant_config.pack_factor, |
| 859 | weight_loader=weight_loader, |
| 860 | ) |
| 861 | |
| 862 | scales = GroupQuantScaleParameter( |
| 863 | data=torch.empty( |
| 864 | num_groups, |
| 865 | output_size_per_partition, |
| 866 | dtype=params_dtype, |
| 867 | ), |
| 868 | input_dim=0, |
| 869 | output_dim=1, |
| 870 | weight_loader=weight_loader, |
| 871 | ) |
| 872 | |
| 873 | layer.register_parameter("qweight", qweight) |
| 874 | layer.register_parameter("qzeros", qzeros) |
| 875 | layer.register_parameter("scales", scales) |
| 876 | |
| 877 | |
| 878 | class AutoAWQLinearMethod(BaseAWQLinearMethod): |
| 879 | """Linear method for AWQ using Triton kernels. |
| 880 | |
| 881 | Args: |
| 882 | quant_config: The AWQ quantization config. |
| 883 | |
| 884 | """ |
| 885 | |
| 886 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 887 | layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False) |
| 888 | layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False) |
| 889 | layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False) |
| 890 | |
| 891 | def apply( |
| 892 | self, |
| 893 | layer: torch.nn.Module, |
| 894 | x: torch.Tensor, |
| 895 | bias: torch.Tensor | None = None, |
| 896 | ) -> torch.Tensor: |
| 897 | qweight = layer.qweight |
| 898 | scales = layer.scales |
| 899 | qzeros = layer.qzeros |
| 900 | pack_factor = self.quant_config.pack_factor |
| 901 | out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,) |
| 902 | reshaped_x = x.reshape(-1, x.shape[-1]) |
| 903 | |
| 904 | # num_tokens >= threshold |
| 905 | FP16_MATMUL_HEURISTIC_CONDITION = x.shape[:-1].numel() >= 256 |
| 906 | # Batch invariant mode requires torch.matmul path |
| 907 | # for Triton override |
| 908 | if FP16_MATMUL_HEURISTIC_CONDITION or envs.VLLM_BATCH_INVARIANT: |
| 909 | out = ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0) |
| 910 | out = torch.matmul(reshaped_x, out) |
| 911 | else: |
| 912 | out = ops.awq_gemm(reshaped_x, qweight, scales, qzeros, pack_factor) |
| 913 | if bias is not None: |
| 914 | out.add_(bias) |
| 915 | return out.reshape(out_shape) |