| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import Annotated, Any |
| 5 | |
| 6 | import regex as re |
| 7 | from pydantic import ( |
| 8 | Field, |
| 9 | GetPydanticSchema, |
| 10 | ValidationInfo, |
| 11 | field_validator, |
| 12 | model_validator, |
| 13 | ) |
| 14 | from pydantic_core import core_schema |
| 15 | |
| 16 | from vllm.config.utils import config |
| 17 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 18 | QuantKey, |
| 19 | kFp8Dynamic128Sym, |
| 20 | kFp8DynamicTensorSym, |
| 21 | kFp8DynamicTokenSym, |
| 22 | kFp8Static128BlockSym, |
| 23 | kFp8StaticChannelSym, |
| 24 | kFp8StaticTensorSym, |
| 25 | kInt8StaticChannelSym, |
| 26 | kMxfp4Dynamic, |
| 27 | kMxfp4Static, |
| 28 | kMxfp8Dynamic, |
| 29 | kNvfp4Static, |
| 30 | ) |
| 31 | |
| 32 | # User-facing names addressable from quantization_config. |
| 33 | QUANT_KEY_NAMES: dict[str, QuantKey] = { |
| 34 | "fp8_per_tensor_static": kFp8StaticTensorSym, |
| 35 | "fp8_per_tensor_dynamic": kFp8DynamicTensorSym, |
| 36 | "fp8_per_token": kFp8DynamicTokenSym, |
| 37 | "fp8_per_channel_static": kFp8StaticChannelSym, |
| 38 | "fp8_per_block_static": kFp8Static128BlockSym, |
| 39 | "fp8_per_block_dynamic": kFp8Dynamic128Sym, |
| 40 | "mxfp8": kMxfp8Dynamic, |
| 41 | "mxfp4": kMxfp4Dynamic, |
| 42 | "int8_per_channel_static": kInt8StaticChannelSym, |
| 43 | } |
| 44 | |
| 45 | |
| 46 | def _coerce_quant_key(v: Any) -> QuantKey | None: |
| 47 | if v is None or isinstance(v, QuantKey): |
| 48 | return v |
| 49 | if not isinstance(v, str): |
| 50 | raise TypeError(f"expected str or QuantKey, got {type(v).__name__}") |
| 51 | try: |
| 52 | return QUANT_KEY_NAMES[v] |
| 53 | except KeyError: |
| 54 | raise ValueError( |
| 55 | f"unknown quantization name {v!r}; " |
| 56 | f"expected one of {sorted(QUANT_KEY_NAMES)}" |
| 57 | ) from None |
| 58 | |
| 59 | |
| 60 | # Stop pydantic from introspecting QuantKey: it transitively contains a |
| 61 | # NamedTuple with `ClassVar[GroupShape]` declarations that pydantic refuses. |
| 62 | QuantKeyField = Annotated[ |
| 63 | QuantKey | None, |
| 64 | GetPydanticSchema( |
| 65 | lambda _src, _handler: core_schema.no_info_plain_validator_function( |
| 66 | _coerce_quant_key |
| 67 | ) |
| 68 | ), |
| 69 | ] |
| 70 | |
| 71 | |
| 72 | @config |
| 73 | class QuantSpec: |
| 74 | """Quantization spec for one layer kind (linear or MoE). |
| 75 | |
| 76 | `None` on either side means the method class falls back to its own default |
| 77 | (typically inherited from the checkpoint, or unquantized for online). |
| 78 | """ |
| 79 | |
| 80 | weight: QuantKeyField = None |
| 81 | """Weight quantization key, or a name from QUANT_KEY_NAMES.""" |
| 82 | |
| 83 | activation: QuantKeyField = None |
| 84 | """Activation quantization key, or a name from QUANT_KEY_NAMES.""" |
| 85 | |
| 86 | def __str__(self) -> str: |
| 87 | def quant_key_str(quant_key: QuantKey | None) -> str: |
| 88 | if quant_key is None: |
| 89 | return "None" |
| 90 | return next( |
| 91 | ( |
| 92 | name |
| 93 | for name, known_quant_key in QUANT_KEY_NAMES.items() |
| 94 | if known_quant_key == quant_key |
| 95 | ), |
| 96 | str(quant_key), |
| 97 | ) |
| 98 | |
| 99 | return quant_key_str(self.weight) |
| 100 | |
| 101 | |
| 102 | @config |
| 103 | class QuantizationConfigArgs: |
| 104 | """User-facing quantization configuration. |
| 105 | |
| 106 | See `docs/features/quantization/online.md` for the schema and shorthand |
| 107 | string forms accepted on `linear` and `moe`. |
| 108 | """ |
| 109 | |
| 110 | linear: QuantSpec | None = None |
| 111 | """Spec applied to ``LinearBase`` layers.""" |
| 112 | |
| 113 | moe: QuantSpec | None = None |
| 114 | """Spec applied to ``FusedMoEFactory`` layers.""" |
| 115 | |
| 116 | ignore: list[str] = Field(default_factory=list) |
| 117 | """Layers to skip quantization for. Online quantization also supports |
| 118 | fnmatch-style patterns.""" |
| 119 | |
| 120 | targets: dict[str, str] | None = None |
| 121 | """Per-layer online quantization overrides, keyed by exact layer name or |
| 122 | regex patterns with a `re:`, or fnmatch-style patterns for online |
| 123 | quantization, mapping to an online shorthand name (see |
| 124 | `_ONLINE_SHORTHANDS`). A layer that matches no pattern is left unquantized. |
| 125 | Mutually exclusive with `linear` and `moe`. |
| 126 | """ |
| 127 | |
| 128 | @field_validator("linear", "moe", mode="before") |
| 129 | @classmethod |
| 130 | def _coerce_spec(cls, v: Any, info: ValidationInfo) -> Any: |
| 131 | if not isinstance(v, str): |
| 132 | return v |
| 133 | field_name = info.field_name |
| 134 | assert field_name is not None |
| 135 | if v in _ONLINE_SHORTHANDS: |
| 136 | spec = getattr(_ONLINE_SHORTHANDS[v], field_name) |
| 137 | if spec is None: |
| 138 | raise ValueError( |
| 139 | f"online shorthand {v!r} does not define a {field_name} spec" |
| 140 | ) |
| 141 | return spec |
| 142 | return QuantSpec(weight=_coerce_quant_key(v)) |
| 143 | |
| 144 | @field_validator("targets", mode="before") |
| 145 | @classmethod |
| 146 | def _validate_targets(cls, v: Any) -> Any: |
| 147 | if v is None: |
| 148 | return v |
| 149 | if not isinstance(v, dict): |
| 150 | raise TypeError(f"targets must be a dict, got {type(v).__name__}") |
| 151 | for pattern, shorthand in v.items(): |
| 152 | if not isinstance(pattern, str): |
| 153 | raise ValueError( |
| 154 | f"targets keys must be strings, got {type(pattern).__name__}" |
| 155 | ) |
| 156 | if not isinstance(shorthand, str) or shorthand not in _ONLINE_SHORTHANDS: |
| 157 | raise ValueError( |
| 158 | f"targets[{pattern}] = {shorthand} is not a valid " |
| 159 | f"online shorthand name; expected one of " |
| 160 | f"{sorted(_ONLINE_SHORTHANDS)}" |
| 161 | ) |
| 162 | if pattern.startswith("re:"): |
| 163 | try: |
| 164 | re.compile(pattern[3:]) |
| 165 | except re.error as e: |
| 166 | raise ValueError( |
| 167 | f"targets key {pattern} is not a valid regex: {e}" |
| 168 | ) from e |
| 169 | return v |
| 170 | |
| 171 | @model_validator(mode="after") |
| 172 | def _validate_targets_exclusivity(self) -> "QuantizationConfigArgs": |
| 173 | if self.targets is None: |
| 174 | return self |
| 175 | if self.linear is not None or self.moe is not None: |
| 176 | raise ValueError( |
| 177 | "quantization_config.targets is mutually exclusive with " |
| 178 | f"quantization_config.linear/moe, got " |
| 179 | f"targets={self.targets}, linear={self.linear}, " |
| 180 | f"moe={self.moe}." |
| 181 | ) |
| 182 | return self |
| 183 | |
| 184 | |
| 185 | # CLI shorthands accepted by `--quantization`. Each desugars to a full |
| 186 | # QuantizationConfigArgs; activation overrides go through quantization_config. |
| 187 | _ONLINE_SHORTHANDS: dict[str, QuantizationConfigArgs] = { |
| 188 | "fp8_per_tensor": QuantizationConfigArgs( |
| 189 | linear=QuantSpec(weight=kFp8StaticTensorSym), |
| 190 | moe=QuantSpec(weight=kFp8StaticTensorSym), |
| 191 | ), |
| 192 | "fp8_per_block": QuantizationConfigArgs( |
| 193 | linear=QuantSpec(weight=kFp8Static128BlockSym), |
| 194 | moe=QuantSpec(weight=kFp8Static128BlockSym), |
| 195 | ), |
| 196 | # Per-output-channel weight scale + dynamic per-token activation. |
| 197 | # Same shape as llmcompressor's FP8_DYNAMIC recipe. |
| 198 | "fp8_per_channel": QuantizationConfigArgs( |
| 199 | linear=QuantSpec(weight=kFp8StaticChannelSym), |
| 200 | moe=QuantSpec(weight=kFp8StaticChannelSym), |
| 201 | ), |
| 202 | "mxfp8": QuantizationConfigArgs( |
| 203 | linear=QuantSpec(weight=kMxfp8Dynamic), |
| 204 | moe=QuantSpec(weight=kMxfp8Dynamic), |
| 205 | ), |
| 206 | "mxfp4": QuantizationConfigArgs( |
| 207 | linear=QuantSpec(weight=kMxfp4Static), |
| 208 | moe=QuantSpec(weight=kMxfp4Static), |
| 209 | ), |
| 210 | # INT8 weight-only on MoE; linear stays unquantized (no `linear` field). |
| 211 | "int8_per_channel_weight_only": QuantizationConfigArgs( |
| 212 | moe=QuantSpec(weight=kInt8StaticChannelSym), |
| 213 | ), |
| 214 | # Online NVFP4 on MoE with per-token dynamic activation scales (Blackwell + |
| 215 | # FlashInfer TRTLLM only); linear stays unquantized (no `linear` field). |
| 216 | "nvfp4_per_token": QuantizationConfigArgs( |
| 217 | moe=QuantSpec(weight=kNvfp4Static), |
| 218 | ), |
| 219 | } |
| 220 | |
| 221 | |
| 222 | # Names accepted by `--quantization`; "online" means "use quantization_config". |
| 223 | ONLINE_QUANT_SHORTHAND_NAMES: tuple[str, ...] = ( |
| 224 | *_ONLINE_SHORTHANDS.keys(), |
| 225 | "online", |
| 226 | ) |
| 227 | |
| 228 | # These names are also checkpoint quantization methods. Their online configs |
| 229 | # are resolved only when checkpoint quantization metadata is absent. |
| 230 | _DEFERRED_ONLINE_SHORTHANDS = frozenset(("mxfp4", "mxfp8")) |
| 231 | |
| 232 | |
| 233 | def resolve_quantization_config( |
| 234 | quantization: str | None, |
| 235 | quantization_config: dict[str, Any] | QuantizationConfigArgs | None, |
| 236 | ) -> QuantizationConfigArgs | None: |
| 237 | """Resolve `--quantization` shorthand and `--quantization-config` into a |
| 238 | QuantizationConfigArgs. |
| 239 | |
| 240 | `quantization` is a CLI shorthand that desugars into a base config via |
| 241 | `_ONLINE_SHORTHANDS`. `quantization_config` is a dict or pre-built args |
| 242 | object. When both are given, fields explicitly set in `quantization_config` |
| 243 | take precedence over the shorthand. |
| 244 | """ |
| 245 | if quantization is not None and quantization not in ONLINE_QUANT_SHORTHAND_NAMES: |
| 246 | # Pre-quantized checkpoints can be composed with online quantization |
| 247 | # for layers that the base quant_method leaves unquantized. The |
| 248 | # checkpoint quant_method remains the primary quantization method; composition |
| 249 | # is performed after its config has been loaded. |
| 250 | if quantization_config is None: |
| 251 | return None |
| 252 | |
| 253 | # `quantization_config` may hold both: |
| 254 | # 1. Base quantization method activation key override, |
| 255 | # 2. online quantization config to apply on top of the base quant_method. |
| 256 | if isinstance(quantization_config, dict): |
| 257 | return QuantizationConfigArgs(**quantization_config) |
| 258 | return quantization_config |
| 259 | |
| 260 | base = _ONLINE_SHORTHANDS.get(quantization) if quantization else None |
| 261 | |
| 262 | if quantization_config is None: |
| 263 | if quantization in _DEFERRED_ONLINE_SHORTHANDS: |
| 264 | return None |
| 265 | return base |
| 266 | |
| 267 | if isinstance(quantization_config, dict): |
| 268 | quantization_config = QuantizationConfigArgs(**quantization_config) |
| 269 | |
| 270 | if base is None: |
| 271 | return quantization_config |
| 272 | |
| 273 | return QuantizationConfigArgs( |
| 274 | linear=quantization_config.linear or base.linear, |
| 275 | moe=quantization_config.moe or base.moe, |
| 276 | ignore=quantization_config.ignore or base.ignore, |
| 277 | targets=quantization_config.targets or base.targets, |
| 278 | ) |