| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Mapping |
| 5 | from enum import Enum |
| 6 | from types import MappingProxyType |
| 7 | from typing import Any |
| 8 | |
| 9 | import torch |
| 10 | |
| 11 | from vllm.config.quantization import ( |
| 12 | _ONLINE_SHORTHANDS, |
| 13 | QuantizationConfigArgs, |
| 14 | QuantSpec, |
| 15 | ) |
| 16 | from vllm.logger import init_logger |
| 17 | from vllm.model_executor.layers.fused_moe import ( |
| 18 | FusedMoEMethodBase, |
| 19 | RoutedExperts, |
| 20 | ) |
| 21 | from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( |
| 22 | UnquantizedFusedMoEMethod, |
| 23 | ) |
| 24 | from vllm.model_executor.layers.linear import ( |
| 25 | LinearBase, |
| 26 | UnquantizedLinearMethod, |
| 27 | ) |
| 28 | from vllm.model_executor.layers.quantization import QuantizationMethods |
| 29 | from vllm.model_executor.layers.quantization.base_config import ( |
| 30 | QuantizationConfig, |
| 31 | QuantizeMethodBase, |
| 32 | ) |
| 33 | from vllm.model_executor.layers.quantization.compressed_tensors.utils import ( |
| 34 | should_ignore_layer, |
| 35 | ) |
| 36 | from vllm.model_executor.layers.quantization.online.fp8 import ( |
| 37 | Fp8PerBlockOnlineLinearMethod, |
| 38 | Fp8PerBlockOnlineMoEMethod, |
| 39 | Fp8PerTensorOnlineLinearMethod, |
| 40 | Fp8PerTensorOnlineMoEMethod, |
| 41 | Fp8PtpcOnlineLinearMethod, |
| 42 | Fp8PtpcOnlineMoEMethod, |
| 43 | OnlineLinearBase, |
| 44 | ) |
| 45 | from vllm.model_executor.layers.quantization.online.int8 import ( |
| 46 | Int8OnlineMoEMethod, |
| 47 | ) |
| 48 | from vllm.model_executor.layers.quantization.online.mxfp4 import ( |
| 49 | Mxfp4OnlineLinearMethod, |
| 50 | Mxfp4OnlineMoEMethod, |
| 51 | ) |
| 52 | from vllm.model_executor.layers.quantization.online.mxfp8 import ( |
| 53 | Mxfp8OnlineLinearMethod, |
| 54 | Mxfp8OnlineMoEMethod, |
| 55 | ) |
| 56 | from vllm.model_executor.layers.quantization.online.nvfp4 import ( |
| 57 | Nvfp4OnlineMoEMethod, |
| 58 | ) |
| 59 | from vllm.model_executor.layers.quantization.utils.config_utils import ( |
| 60 | find_matching_patterns, |
| 61 | get_layer_name_after_index, |
| 62 | ) |
| 63 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 64 | QuantKey, |
| 65 | kFp8Static128BlockSym, |
| 66 | kFp8StaticChannelSym, |
| 67 | kFp8StaticTensorSym, |
| 68 | kInt8StaticChannelSym, |
| 69 | kMxfp4Static, |
| 70 | kMxfp8Dynamic, |
| 71 | kNvfp4Static, |
| 72 | ) |
| 73 | |
| 74 | logger = init_logger(__name__) |
| 75 | |
| 76 | |
| 77 | class OnlineQuantizationSource(str, Enum): |
| 78 | """Supported online quantization configuration sources.""" |
| 79 | |
| 80 | linear = "linear" # LinearBase |
| 81 | moe = "moe" # RoutedExperts |
| 82 | targets = "targets" |
| 83 | |
| 84 | |
| 85 | # Online dispatch tables, keyed by the QuantSpec.weight QuantKey. The |
| 86 | # corresponding method class handles the activation choice via its |
| 87 | # `supported_activation_quant` set. |
| 88 | _ONLINE_LINEAR_METHODS: dict[QuantKey, type] = { |
| 89 | kFp8StaticTensorSym: Fp8PerTensorOnlineLinearMethod, |
| 90 | kFp8Static128BlockSym: Fp8PerBlockOnlineLinearMethod, |
| 91 | kFp8StaticChannelSym: Fp8PtpcOnlineLinearMethod, |
| 92 | kMxfp8Dynamic: Mxfp8OnlineLinearMethod, |
| 93 | kMxfp4Static: Mxfp4OnlineLinearMethod, |
| 94 | } |
| 95 | |
| 96 | _ONLINE_MOE_METHODS: dict[QuantKey, type] = { |
| 97 | kFp8StaticTensorSym: Fp8PerTensorOnlineMoEMethod, |
| 98 | kFp8Static128BlockSym: Fp8PerBlockOnlineMoEMethod, |
| 99 | kFp8StaticChannelSym: Fp8PtpcOnlineMoEMethod, |
| 100 | kMxfp8Dynamic: Mxfp8OnlineMoEMethod, |
| 101 | kMxfp4Static: Mxfp4OnlineMoEMethod, |
| 102 | kInt8StaticChannelSym: Int8OnlineMoEMethod, |
| 103 | kNvfp4Static: Nvfp4OnlineMoEMethod, |
| 104 | } |
| 105 | |
| 106 | |
| 107 | def _find_matching_targets( |
| 108 | prefix: str, |
| 109 | targets: Mapping[str, str], |
| 110 | fused_mapping: Mapping[str, list[str]] = MappingProxyType({}), |
| 111 | ) -> list[str]: |
| 112 | per_shard_matches = find_matching_patterns( |
| 113 | prefix, targets, fused_mapping, use_fnmatch=True |
| 114 | ) |
| 115 | if all(len(matches) == 0 for matches in per_shard_matches): |
| 116 | return [] |
| 117 | if any(len(matches) == 0 for matches in per_shard_matches): |
| 118 | raise ValueError( |
| 119 | f"Found unmatched shards for {prefix}: {per_shard_matches}. vLLM " |
| 120 | "requires all shards of a fused layer to match a target." |
| 121 | ) |
| 122 | if any(len(matches) > 1 for matches in per_shard_matches): |
| 123 | raise ValueError( |
| 124 | f"Found multiple quantization_config.targets matches for the " |
| 125 | f"shards of {prefix}: {per_shard_matches}. Each shard may match " |
| 126 | "at most one target." |
| 127 | ) |
| 128 | |
| 129 | matched_patterns = [next(iter(matches)) for matches in per_shard_matches] |
| 130 | quant_key_strs = {targets[pattern] for pattern in matched_patterns} |
| 131 | if len(quant_key_strs) > 1: |
| 132 | raise ValueError( |
| 133 | f"Found different quantization_config.targets values for the " |
| 134 | f"shards of {prefix}: {matched_patterns}. vLLM requires all " |
| 135 | "shards of a fused layer to use the same target." |
| 136 | ) |
| 137 | return [matched_patterns[0]] |
| 138 | |
| 139 | |
| 140 | class OnlineQuantizationConfig(QuantizationConfig): |
| 141 | """Model-level config for online quantization (quantize fp16/bf16 weights |
| 142 | during model loading, without requiring a pre-quantized checkpoint).""" |
| 143 | |
| 144 | def __init__( |
| 145 | self, |
| 146 | args: QuantizationConfigArgs, |
| 147 | ) -> None: |
| 148 | super().__init__() |
| 149 | if args.linear is None and args.moe is None and args.targets is None: |
| 150 | raise ValueError( |
| 151 | "OnlineQuantizationConfig requires at least one of " |
| 152 | "quantization_config.linear, quantization_config.moe, or " |
| 153 | "quantization_config.targets to be set." |
| 154 | ) |
| 155 | self.args = args |
| 156 | self.ignored_layers: list[str] = args.ignore |
| 157 | self.quantized_layers: dict[str, tuple[str, str, str | None]] = {} |
| 158 | |
| 159 | @property |
| 160 | def quantized_layer_summaries(self) -> list[str]: |
| 161 | counts: dict[tuple[str, str, str | None, str], int] = {} |
| 162 | for layer_name, ( |
| 163 | source, |
| 164 | quant_key_str, |
| 165 | target_pattern, |
| 166 | ) in self.quantized_layers.items(): |
| 167 | key = ( |
| 168 | get_layer_name_after_index(layer_name), |
| 169 | source, |
| 170 | target_pattern, |
| 171 | quant_key_str, |
| 172 | ) |
| 173 | counts[key] = counts.get(key, 0) + 1 |
| 174 | |
| 175 | summaries = [] |
| 176 | # Build summary entries as |
| 177 | # `self_attn.o_proj: 24 (from targets: re:.*self_attn\.o_proj, mxfp4` |
| 178 | for (layer_type, source, target_pattern, quant_key_str), count in sorted( |
| 179 | counts.items() |
| 180 | ): |
| 181 | pattern_prefix = f"{target_pattern}, " if target_pattern else "" |
| 182 | summaries.append( |
| 183 | f"{layer_type}: {count} " |
| 184 | f"(from {source}: {pattern_prefix}{quant_key_str})" |
| 185 | ) |
| 186 | return summaries |
| 187 | |
| 188 | @classmethod |
| 189 | def get_name(cls) -> QuantizationMethods: |
| 190 | return "online" |
| 191 | |
| 192 | @classmethod |
| 193 | def get_supported_act_dtypes(cls) -> list[torch.dtype]: |
| 194 | return [torch.bfloat16, torch.half] |
| 195 | |
| 196 | @classmethod |
| 197 | def get_min_capability(cls) -> int: |
| 198 | # Note: as more online quant schemes will be added, this |
| 199 | # value will become the minimum across all supported schemes. |
| 200 | return 75 |
| 201 | |
| 202 | @classmethod |
| 203 | def get_config_filenames(cls) -> list[str]: |
| 204 | return [] |
| 205 | |
| 206 | @classmethod |
| 207 | def from_config(cls, config: dict[str, Any]) -> "OnlineQuantizationConfig": |
| 208 | raise NotImplementedError( |
| 209 | "OnlineQuantizationConfig does not support loading from a " |
| 210 | "checkpoint config. Use quantization_config or " |
| 211 | "quantization='fp8_per_tensor'/'fp8_per_block' instead." |
| 212 | ) |
| 213 | |
| 214 | def _get_method_cls( |
| 215 | self, |
| 216 | spec: QuantSpec | None, |
| 217 | table: dict[QuantKey, type], |
| 218 | layer: torch.nn.Module, |
| 219 | ) -> type | None: |
| 220 | """Resolve the online method class for a layer's quantization spec. |
| 221 | |
| 222 | Args: |
| 223 | spec: Quantization specification to resolve. |
| 224 | table: Mapping from weight quantization keys to method classes. |
| 225 | layer: Layer that will use the resolved method. |
| 226 | |
| 227 | Returns: |
| 228 | The matching method class, or None when ``spec`` has no weight |
| 229 | quantization. |
| 230 | |
| 231 | """ |
| 232 | if spec is None or spec.weight is None: |
| 233 | return None |
| 234 | cls = table.get(spec.weight) |
| 235 | if cls is None: |
| 236 | raise ValueError( |
| 237 | f"online quantization for {type(layer).__name__} with " |
| 238 | f"weight={spec.weight} is not supported; supported weight " |
| 239 | f"keys: {sorted(str(k) for k in table)}" |
| 240 | ) |
| 241 | # Online method classes pick their own activation format internally. |
| 242 | # Per-class activation overrides are not yet wired through; reject |
| 243 | # explicit overrides until the relevant method class opts in. |
| 244 | if spec.activation is not None: |
| 245 | raise ValueError( |
| 246 | f"activation override (activation={spec.activation}) is not " |
| 247 | f"yet supported for online {cls.__name__}" |
| 248 | ) |
| 249 | return cls |
| 250 | |
| 251 | def resolve_quant_method_cls( |
| 252 | self, layer: torch.nn.Module, prefix: str |
| 253 | ) -> tuple[OnlineQuantizationSource, str, str | None, QuantSpec, type] | None: |
| 254 | """Resolve quantization metadata and method class without instantiating it. |
| 255 | |
| 256 | Args: |
| 257 | layer: Layer for which to resolve online quantization. |
| 258 | prefix: Fully qualified layer name. |
| 259 | |
| 260 | Returns: |
| 261 | A tuple of source, quantization key string, target pattern, spec, |
| 262 | and method class. Returns None when online quantization does not |
| 263 | apply to the layer. |
| 264 | |
| 265 | """ |
| 266 | quant_spec: QuantSpec | None |
| 267 | if self.args.targets is not None: |
| 268 | resolved_pattern = self._resolve_targets_quant_method_metadata( |
| 269 | prefix, layer |
| 270 | ) |
| 271 | if resolved_pattern is None: |
| 272 | return None |
| 273 | source, quant_key_str, target_pattern, quant_spec, table = resolved_pattern |
| 274 | else: |
| 275 | if isinstance(layer, LinearBase): |
| 276 | source = OnlineQuantizationSource.linear |
| 277 | quant_spec = self.args.linear |
| 278 | table = _ONLINE_LINEAR_METHODS |
| 279 | elif isinstance(layer, RoutedExperts): |
| 280 | source = OnlineQuantizationSource.moe |
| 281 | quant_spec = self.args.moe |
| 282 | table = _ONLINE_MOE_METHODS |
| 283 | else: |
| 284 | return None |
| 285 | |
| 286 | if should_ignore_layer( |
| 287 | prefix, |
| 288 | ignore=self.ignored_layers, |
| 289 | fused_mapping=self.packed_modules_mapping, |
| 290 | use_fnmatch=True, |
| 291 | ): |
| 292 | return None |
| 293 | quant_key_str = str(quant_spec) |
| 294 | target_pattern = None |
| 295 | |
| 296 | quant_method_cls = self._get_method_cls(quant_spec, table, layer) |
| 297 | if quant_method_cls is None: |
| 298 | return None |
| 299 | assert quant_spec is not None |
| 300 | return source, quant_key_str, target_pattern, quant_spec, quant_method_cls |
| 301 | |
| 302 | def _resolve_targets_quant_method_metadata( |
| 303 | self, prefix: str, layer: torch.nn.Module |
| 304 | ) -> ( |
| 305 | tuple[OnlineQuantizationSource, str, str, QuantSpec, dict[QuantKey, type]] |
| 306 | | None |
| 307 | ): |
| 308 | """Resolve target-pattern quantization metadata and dispatch table. |
| 309 | |
| 310 | Args: |
| 311 | prefix: Fully qualified layer name. |
| 312 | layer: Layer matched against configured target patterns. |
| 313 | |
| 314 | Returns: |
| 315 | A tuple of source, quantization key string, target pattern, spec, |
| 316 | and dispatch table. Returns None when no pattern applies or the |
| 317 | layer is ignored. |
| 318 | |
| 319 | """ |
| 320 | assert self.args.targets is not None |
| 321 | ignored = should_ignore_layer( |
| 322 | prefix, |
| 323 | ignore=self.ignored_layers, |
| 324 | fused_mapping=self.packed_modules_mapping, |
| 325 | use_fnmatch=True, |
| 326 | ) |
| 327 | matches = _find_matching_targets( |
| 328 | prefix, self.args.targets, fused_mapping=self.packed_modules_mapping |
| 329 | ) |
| 330 | if ignored and matches: |
| 331 | raise ValueError( |
| 332 | f"Layer {prefix} matches both quantization_config.ignore " |
| 333 | f"and quantization_config.targets ({matches}); a layer may " |
| 334 | f"not be referenced by both." |
| 335 | ) |
| 336 | if ignored or not matches: |
| 337 | return None |
| 338 | if len(matches) > 1: |
| 339 | raise ValueError( |
| 340 | f"Layer {prefix} matches multiple quantization_config." |
| 341 | f"targets patterns: {matches}. Each layer may match at most " |
| 342 | f"one target." |
| 343 | ) |
| 344 | target_pattern = matches[0] |
| 345 | quant_key_str = self.args.targets[target_pattern] |
| 346 | shorthand = _ONLINE_SHORTHANDS[quant_key_str] |
| 347 | if isinstance(layer, LinearBase): |
| 348 | quant_spec = shorthand.linear |
| 349 | table = _ONLINE_LINEAR_METHODS |
| 350 | elif isinstance(layer, RoutedExperts): |
| 351 | quant_spec = shorthand.moe |
| 352 | table = _ONLINE_MOE_METHODS |
| 353 | else: |
| 354 | raise ValueError( |
| 355 | f"Layer {prefix} was matched by quantization_config.targets " |
| 356 | f"({target_pattern}), but online quantization is not supported for " |
| 357 | f"{type(layer).__name__}." |
| 358 | ) |
| 359 | if quant_spec is None: |
| 360 | raise ValueError( |
| 361 | f"targets pattern {target_pattern} = {quant_key_str} does " |
| 362 | f"not define a QuantSpec for {type(layer).__name__} layers " |
| 363 | f"(matched at {prefix})." |
| 364 | ) |
| 365 | return ( |
| 366 | OnlineQuantizationSource.targets, |
| 367 | quant_key_str, |
| 368 | target_pattern, |
| 369 | quant_spec, |
| 370 | table, |
| 371 | ) |
| 372 | |
| 373 | def get_quant_method( |
| 374 | self, layer: torch.nn.Module, prefix: str |
| 375 | ) -> "QuantizeMethodBase | None": |
| 376 | # `targets` takes precedence over `moe` and `linear` and is exclusive. |
| 377 | resolved = self.resolve_quant_method_cls(layer, prefix) |
| 378 | if resolved is not None: |
| 379 | source, quant_key_str, target_pattern, _, quant_method_cls = resolved |
| 380 | self.quantized_layers[prefix] = ( |
| 381 | source.value, |
| 382 | quant_key_str, |
| 383 | target_pattern, |
| 384 | ) |
| 385 | if isinstance(layer, RoutedExperts): |
| 386 | assert issubclass(quant_method_cls, FusedMoEMethodBase) |
| 387 | return quant_method_cls(moe=layer.moe_config) |
| 388 | |
| 389 | assert issubclass(quant_method_cls, OnlineLinearBase) |
| 390 | return quant_method_cls() |
| 391 | |
| 392 | if isinstance(layer, LinearBase): |
| 393 | return UnquantizedLinearMethod() |
| 394 | if isinstance(layer, RoutedExperts): |
| 395 | return UnquantizedFusedMoEMethod(layer.moe_config) |
| 396 | return None |