| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import dataclass |
| 5 | from typing import TYPE_CHECKING |
| 6 | |
| 7 | import regex as re |
| 8 | |
| 9 | from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead |
| 10 | |
| 11 | if TYPE_CHECKING: |
| 12 | import torch |
| 13 | |
| 14 | from .inc import INCConfig |
| 15 | |
| 16 | |
| 17 | @dataclass(frozen=True) |
| 18 | class INCLayerConfig: |
| 19 | bits: int |
| 20 | group_size: int | tuple[int, int] |
| 21 | sym: bool |
| 22 | packing_format: str |
| 23 | backend: str |
| 24 | data_type: str |
| 25 | quantized: bool |
| 26 | |
| 27 | @property |
| 28 | def is_gptq(self) -> bool: |
| 29 | return "gptq" in self.packing_format or "gptq" in self.backend |
| 30 | |
| 31 | @property |
| 32 | def is_awq(self) -> bool: |
| 33 | return "awq" in self.packing_format or "awq" in self.backend |
| 34 | |
| 35 | @property |
| 36 | def is_wna16_int(self) -> bool: |
| 37 | return self.data_type == "int" and self.quantized |
| 38 | |
| 39 | @property |
| 40 | def is_mxfp4(self) -> bool: |
| 41 | return "mx_fp" in self.data_type and self.bits == 4 |
| 42 | |
| 43 | @property |
| 44 | def is_mxfp8(self) -> bool: |
| 45 | return "mx_fp" in self.data_type and self.bits == 8 |
| 46 | |
| 47 | @property |
| 48 | def is_fp8_block(self) -> bool: |
| 49 | return ( |
| 50 | self.data_type == "fp" |
| 51 | and self.bits == 8 |
| 52 | and self.sym |
| 53 | and self.quantized |
| 54 | and self.packing_format == "auto_round:fp8" |
| 55 | and isinstance(self.group_size, tuple) |
| 56 | and len(self.group_size) == 2 |
| 57 | ) |
| 58 | |
| 59 | |
| 60 | class INCConfigParser: |
| 61 | def __init__(self, config: "INCConfig") -> None: |
| 62 | self._config = config |
| 63 | |
| 64 | def resolve(self, layer: "torch.nn.Module", layer_name: str) -> INCLayerConfig: |
| 65 | bits, group_size, sym = self._resolve_raw(layer, layer_name) |
| 66 | return INCLayerConfig( |
| 67 | bits=bits, |
| 68 | group_size=group_size, |
| 69 | sym=sym, |
| 70 | packing_format=self._config.packing_format, |
| 71 | backend=self._config.backend, |
| 72 | data_type=self._config.data_type, |
| 73 | quantized=bits < 16, |
| 74 | ) |
| 75 | |
| 76 | def get_layer_config( |
| 77 | self, layer: "torch.nn.Module", layer_name: str |
| 78 | ) -> tuple[int, int | tuple[int, int], bool]: |
| 79 | layer_config = self.resolve(layer, layer_name) |
| 80 | return layer_config.bits, layer_config.group_size, layer_config.sym |
| 81 | |
| 82 | def _resolve_raw( |
| 83 | self, layer: "torch.nn.Module", layer_name: str |
| 84 | ) -> tuple[int, int | tuple[int, int], bool]: |
| 85 | REGEX_SPECIAL_CHARS = set(r"*+?^$()[]{}|\\") |
| 86 | |
| 87 | def is_explicitly_configured(name: str) -> bool: |
| 88 | """Return True if *name* has an explicit entry in extra_config, |
| 89 | either via exact key match or via a regex pattern key.""" |
| 90 | if not self._config.extra_config: |
| 91 | return False |
| 92 | if name in self._config.extra_config: |
| 93 | return True |
| 94 | for pattern in self._config.extra_config: |
| 95 | if not isinstance(pattern, str) or not any( |
| 96 | c in REGEX_SPECIAL_CHARS for c in pattern |
| 97 | ): |
| 98 | continue |
| 99 | try: |
| 100 | if re.search(re.compile(pattern), name) is not None: |
| 101 | return True |
| 102 | except re.error: |
| 103 | continue |
| 104 | return False |
| 105 | |
| 106 | def get_config( |
| 107 | name: str, quantized: bool = True |
| 108 | ) -> tuple[int, int | tuple[int, int], bool]: |
| 109 | if not self._config.extra_config: |
| 110 | return ( |
| 111 | self._config.weight_bits if quantized else 16, |
| 112 | self._config.group_size if quantized else -1, |
| 113 | self._config.sym if quantized else True, |
| 114 | ) |
| 115 | |
| 116 | if name in self._config.extra_config: |
| 117 | cfg = self._config.extra_config[name] |
| 118 | return ( |
| 119 | cfg.get("bits", self._config.weight_bits if quantized else 16), |
| 120 | self._normalize_group_size( |
| 121 | cfg.get( |
| 122 | "group_size", |
| 123 | self._config.group_size if quantized else -1, |
| 124 | ) |
| 125 | ), |
| 126 | cfg.get("sym", self._config.sym if quantized else True), |
| 127 | ) |
| 128 | |
| 129 | regex_special_chars = set(r"*+?^$()[]{}|\\") |
| 130 | for pattern, cfg in self._config.extra_config.items(): |
| 131 | if not isinstance(pattern, str) or not any( |
| 132 | c in regex_special_chars for c in pattern |
| 133 | ): |
| 134 | continue |
| 135 | |
| 136 | try: |
| 137 | if re.search(re.compile(pattern), name) is not None: |
| 138 | return ( |
| 139 | cfg.get( |
| 140 | "bits", |
| 141 | self._config.weight_bits if quantized else 16, |
| 142 | ), |
| 143 | self._normalize_group_size( |
| 144 | cfg.get( |
| 145 | "group_size", |
| 146 | self._config.group_size if quantized else -1, |
| 147 | ) |
| 148 | ), |
| 149 | cfg.get("sym", self._config.sym if quantized else True), |
| 150 | ) |
| 151 | except re.error: |
| 152 | continue |
| 153 | |
| 154 | return ( |
| 155 | self._config.weight_bits if quantized else 16, |
| 156 | self._config.group_size if quantized else -1, |
| 157 | self._config.sym if quantized else True, |
| 158 | ) |
| 159 | |
| 160 | if self._config.extra_config and layer_name in self._config.extra_config: |
| 161 | return get_config(layer_name) |
| 162 | |
| 163 | # Suffix match: handle cases where extra_config keys use short names |
| 164 | # (e.g. "lm_head") but the layer_name is fully qualified |
| 165 | # (e.g. "model.language_model.lm_head") due to model nesting. |
| 166 | if self._config.extra_config: |
| 167 | for cfg_key in self._config.extra_config: |
| 168 | if layer_name.endswith(f".{cfg_key}"): |
| 169 | return get_config(cfg_key) |
| 170 | |
| 171 | quantized = not isinstance(layer, ParallelLMHead) |
| 172 | if self._config.block_name_to_quantize: |
| 173 | quantized = any( |
| 174 | layer_name.startswith(name) |
| 175 | for name in self._config.block_name_to_quantize |
| 176 | ) |
| 177 | |
| 178 | if self._config.extra_config and "fusedmoe" in layer.__class__.__name__.lower(): |
| 179 | moe_configs = [ |
| 180 | get_config(name, quantized) |
| 181 | for name in self._config.extra_config |
| 182 | if name.startswith(layer_name) |
| 183 | ] |
| 184 | if moe_configs: |
| 185 | if len(set(moe_configs)) == 1: |
| 186 | return moe_configs[0] |
| 187 | raise ValueError( |
| 188 | f"Fused MoE layer '{layer_name}' requires " |
| 189 | f"consistent quant config for all sub-layers" |
| 190 | ) |
| 191 | |
| 192 | if self._config.extra_config: |
| 193 | for fusion_key, sub_keys in self._config.packed_modules_mapping.items(): |
| 194 | if fusion_key in layer_name and layer_name.count(fusion_key) == 1: |
| 195 | sub_names = [ |
| 196 | layer_name.replace(fusion_key, sub_key) for sub_key in sub_keys |
| 197 | ] |
| 198 | # Only trigger if at least one sub_name is explicitly |
| 199 | # configured in extra_config (via exact match or regex). |
| 200 | # This prevents false matches when a short fusion_key |
| 201 | # (e.g. "qkv") is merely a substring of a longer layer |
| 202 | # name (e.g. "in_proj_qkvz") and none of the generated |
| 203 | # sub_names are actually configured. |
| 204 | if not any(is_explicitly_configured(n) for n in sub_names): |
| 205 | continue |
| 206 | sub_configs = [get_config(name, quantized) for name in sub_names] |
| 207 | if len(set(sub_configs)) == 1: |
| 208 | return sub_configs[0] |
| 209 | raise ValueError( |
| 210 | f"Fused module '{layer_name}' requires " |
| 211 | f"consistent quant config for {sub_names}" |
| 212 | ) |
| 213 | |
| 214 | return get_config(layer_name, quantized) |
| 215 | |
| 216 | @staticmethod |
| 217 | def _normalize_group_size( |
| 218 | group_size: int | list[int] | tuple[int, int], |
| 219 | ) -> int | tuple[int, int]: |
| 220 | """Normalize INC group_size into either an int or a 2-D int tuple.""" |
| 221 | if isinstance(group_size, (list, tuple)): |
| 222 | if len(group_size) != 2 or not all( |
| 223 | isinstance(value, int) for value in group_size |
| 224 | ): |
| 225 | raise ValueError( |
| 226 | "INC block-wise FP8 requires group_size to be a 2-D " |
| 227 | f"integer sequence, but found {group_size!r}." |
| 228 | ) |
| 229 | return (group_size[0], group_size[1]) |
| 230 | |
| 231 | if not isinstance(group_size, int): |
| 232 | raise ValueError( |
| 233 | "INC group_size must be an int or a 2-D integer sequence, " |
| 234 | f"but found {type(group_size).__name__}." |
| 235 | ) |
| 236 | |
| 237 | return group_size |