| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import Literal, get_args |
| 5 | |
| 6 | from vllm.logger import init_logger |
| 7 | from vllm.model_executor.layers.quantization.base_config import ( |
| 8 | QuantizationConfig, |
| 9 | resolve_quant_method, |
| 10 | ) |
| 11 | from vllm.platforms import current_platform |
| 12 | |
| 13 | logger = init_logger(__name__) |
| 14 | |
| 15 | QuantizationMethods = Literal[ |
| 16 | "awq", |
| 17 | "auto_awq", |
| 18 | "fp8", |
| 19 | "fbgemm_fp8", |
| 20 | "fp_quant", |
| 21 | "modelopt", |
| 22 | "modelopt_fp4", |
| 23 | "modelopt_mxfp8", |
| 24 | "modelopt_mixed", |
| 25 | "auto_gptq", |
| 26 | "gptq", |
| 27 | "gptq_marlin", |
| 28 | "awq_marlin", |
| 29 | "humming", |
| 30 | "compressed-tensors", |
| 31 | "experts_int8", |
| 32 | "quark", |
| 33 | "moe_wna16", |
| 34 | "torchao", |
| 35 | "inc", |
| 36 | "mxfp4", |
| 37 | "gpt_oss_mxfp4", |
| 38 | "deepseek_v4_fp8", |
| 39 | "online", |
| 40 | # Below are online quant shorthand names (see vllm.config.quantization). |
| 41 | # Listed here as strings to avoid a circular import; kept in sync with |
| 42 | # _ONLINE_SHORTHANDS by the assertion in get_quantization_config(). |
| 43 | "fp8_per_tensor", |
| 44 | "fp8_per_block", |
| 45 | "fp8_per_channel", |
| 46 | "int8_per_channel_weight_only", |
| 47 | "nvfp4_per_token", |
| 48 | "mxfp8", |
| 49 | ] |
| 50 | QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) |
| 51 | |
| 52 | DEPRECATED_QUANTIZATION_METHODS = [ |
| 53 | "fbgemm_fp8", |
| 54 | "fp_quant", |
| 55 | ] |
| 56 | |
| 57 | # The customized quantization methods which will be added to this dict. |
| 58 | _CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {} |
| 59 | |
| 60 | |
| 61 | def register_quantization_config(quantization: str): |
| 62 | """Register a customized vllm quantization config. |
| 63 | |
| 64 | When a quantization method is not supported by vllm, you can register a customized |
| 65 | quantization config to support it. |
| 66 | |
| 67 | Args: |
| 68 | quantization (str): The quantization method name. |
| 69 | |
| 70 | Examples: |
| 71 | >>> from vllm.model_executor.layers.quantization import ( |
| 72 | ... register_quantization_config, |
| 73 | ... ) |
| 74 | >>> from vllm.model_executor.layers.quantization import get_quantization_config |
| 75 | >>> from vllm.model_executor.layers.quantization.base_config import ( |
| 76 | ... QuantizationConfig, |
| 77 | ... ) |
| 78 | >>> |
| 79 | >>> @register_quantization_config("my_quant") |
| 80 | ... class MyQuantConfig(QuantizationConfig): |
| 81 | ... pass |
| 82 | >>> |
| 83 | >>> get_quantization_config("my_quant") |
| 84 | <class 'MyQuantConfig'> |
| 85 | |
| 86 | """ # noqa: E501 |
| 87 | |
| 88 | def _wrapper(quant_config_cls): |
| 89 | if quantization in QUANTIZATION_METHODS: |
| 90 | logger.debug( |
| 91 | "The quantization method '%s' already exists and will be " |
| 92 | "overwritten by the quantization config %s.", |
| 93 | quantization, |
| 94 | quant_config_cls, |
| 95 | ) |
| 96 | else: |
| 97 | QUANTIZATION_METHODS.append(quantization) |
| 98 | # Automatically assume the custom quantization config is supported |
| 99 | if sq := current_platform.supported_quantization: |
| 100 | sq.append(quantization) |
| 101 | |
| 102 | if not issubclass(quant_config_cls, QuantizationConfig): |
| 103 | raise ValueError( |
| 104 | "The quantization config must be a subclass of `QuantizationConfig`." |
| 105 | ) |
| 106 | _CUSTOMIZED_METHOD_TO_QUANT_CONFIG[quantization] = quant_config_cls |
| 107 | return quant_config_cls |
| 108 | |
| 109 | return _wrapper |
| 110 | |
| 111 | |
| 112 | def get_quantization_config(quantization: str) -> type[QuantizationConfig]: |
| 113 | if quantization not in QUANTIZATION_METHODS: |
| 114 | raise ValueError(f"Invalid quantization method: {quantization}") |
| 115 | |
| 116 | # lazy import to avoid triggering `torch.compile` too early |
| 117 | from vllm.config.quantization import _ONLINE_SHORTHANDS |
| 118 | from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig |
| 119 | |
| 120 | if current_platform.is_cuda() or current_platform.is_rocm(): |
| 121 | # The v4.1 class is the v4 one extended to accept model_type |
| 122 | # "deepseek_v41" and its 32x32 MXFP8 linear layout. V4.1 has |
| 123 | # platform-specific implementations for both CUDA and ROCm. |
| 124 | from vllm.models.deepseek_v41 import ( |
| 125 | DeepseekV4FP8Config as DeepseekV41FP8Config, |
| 126 | ) |
| 127 | |
| 128 | deepseek_config: type[QuantizationConfig] = DeepseekV41FP8Config |
| 129 | else: |
| 130 | from vllm.models.deepseek_v4 import DeepseekV4FP8Config |
| 131 | |
| 132 | deepseek_config = DeepseekV4FP8Config |
| 133 | |
| 134 | from .auto_awq import AutoAWQConfig |
| 135 | from .auto_gptq import AutoGPTQConfig |
| 136 | from .compressed_tensors.compressed_tensors import ( |
| 137 | CompressedTensorsConfig, |
| 138 | ) |
| 139 | from .experts_int8 import ExpertsInt8Config |
| 140 | from .fbgemm_fp8 import FBGEMMFp8Config |
| 141 | from .fp8 import Fp8Config |
| 142 | from .fp_quant import FPQuantConfig |
| 143 | from .humming import HummingConfig |
| 144 | from .inc import INCConfig |
| 145 | from .modelopt import ( |
| 146 | ModelOptFp8Config, |
| 147 | ModelOptMixedPrecisionConfig, |
| 148 | ModelOptMxFp8Config, |
| 149 | ModelOptNvFp4Config, |
| 150 | ) |
| 151 | from .moe_wna16 import MoeWNA16Config |
| 152 | from .mxfp4 import GptOssMxfp4Config, Mxfp4Config |
| 153 | from .online.base import OnlineQuantizationConfig |
| 154 | from .torchao import TorchAOConfig |
| 155 | |
| 156 | method_to_config: dict[str, type[QuantizationConfig]] = { |
| 157 | "awq": AutoAWQConfig, |
| 158 | "awq_marlin": AutoAWQConfig, |
| 159 | "auto_awq": AutoAWQConfig, |
| 160 | "fp8": Fp8Config, |
| 161 | "fbgemm_fp8": FBGEMMFp8Config, |
| 162 | "fp_quant": FPQuantConfig, |
| 163 | "modelopt": ModelOptFp8Config, |
| 164 | "modelopt_fp4": ModelOptNvFp4Config, |
| 165 | "modelopt_mxfp8": ModelOptMxFp8Config, |
| 166 | "modelopt_mixed": ModelOptMixedPrecisionConfig, |
| 167 | "auto_gptq": AutoGPTQConfig, |
| 168 | "gptq": AutoGPTQConfig, |
| 169 | "gptq_marlin": AutoGPTQConfig, |
| 170 | "compressed-tensors": CompressedTensorsConfig, |
| 171 | "experts_int8": ExpertsInt8Config, |
| 172 | "quark": QuarkConfig, |
| 173 | "moe_wna16": MoeWNA16Config, |
| 174 | "torchao": TorchAOConfig, |
| 175 | "inc": INCConfig, |
| 176 | "mxfp4": Mxfp4Config, |
| 177 | "gpt_oss_mxfp4": GptOssMxfp4Config, |
| 178 | "deepseek_v4_fp8": deepseek_config, |
| 179 | "humming": HummingConfig, |
| 180 | "online": OnlineQuantizationConfig, |
| 181 | # MiniMax-style checkpoints tag `quant_method: "mxfp8"`; load with the |
| 182 | # ModelOpt MXFP8 config (same format). The "mxfp8" online shorthand |
| 183 | # below only applies to the `--quantization mxfp8` CLI path. |
| 184 | "mxfp8": ModelOptMxFp8Config, |
| 185 | } |
| 186 | |
| 187 | # Register online shorthands (e.g. "fp8_per_tensor") as quant methods. |
| 188 | # setdefault so a shorthand that is also a checkpoint method (e.g. "mxfp8") |
| 189 | # keeps its checkpoint config; the shorthand still works via the |
| 190 | # `--quantization` CLI path in `resolve_quantization_config`. |
| 191 | for shorthand in _ONLINE_SHORTHANDS: |
| 192 | method_to_config.setdefault(shorthand, OnlineQuantizationConfig) |
| 193 | |
| 194 | # Update the `method_to_config` with customized quantization methods. |
| 195 | method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG) |
| 196 | |
| 197 | return method_to_config[quantization] |
| 198 | |
| 199 | |
| 200 | __all__ = [ |
| 201 | "QuantizationConfig", |
| 202 | "resolve_quant_method", |
| 203 | "QuantizationMethods", |
| 204 | "get_quantization_config", |
| 205 | "register_quantization_config", |
| 206 | "QUANTIZATION_METHODS", |
| 207 | ] |