| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Iterable, Mapping |
| 5 | from types import MappingProxyType |
| 6 | |
| 7 | from compressed_tensors import CompressionFormat |
| 8 | from compressed_tensors.quantization import QuantizationStrategy |
| 9 | from torch.nn import Module |
| 10 | |
| 11 | from vllm.model_executor.layers.quantization.utils.config_utils import ( |
| 12 | find_matching_patterns, |
| 13 | is_equal_or_regex_match, |
| 14 | ) |
| 15 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 16 | kFp8Static128BlockSym, |
| 17 | kFp8StaticChannelSym, |
| 18 | kFp8StaticTensorSym, |
| 19 | ) |
| 20 | from vllm.model_executor.parameter import ( |
| 21 | BlockQuantScaleParameter, |
| 22 | ChannelQuantScaleParameter, |
| 23 | PerTensorScaleParameter, |
| 24 | ) |
| 25 | |
| 26 | # Maps quantization strategy to the corresponding scale parameter type. |
| 27 | # Shared across compressed-tensor scheme classes (w8a16_fp8, w8a8_fp8, …). |
| 28 | STRATEGY_TO_PARAMETER_TYPE = { |
| 29 | QuantizationStrategy.BLOCK: BlockQuantScaleParameter, |
| 30 | QuantizationStrategy.CHANNEL: ChannelQuantScaleParameter, |
| 31 | QuantizationStrategy.TENSOR: PerTensorScaleParameter, |
| 32 | } |
| 33 | |
| 34 | # Maps quantization strategy to the vLLM weight-quant key used for |
| 35 | # kernel selection. Shared across compressed-tensor scheme classes. |
| 36 | STRATEGY_TO_WEIGHT_QUANT_KEY = { |
| 37 | QuantizationStrategy.BLOCK: kFp8Static128BlockSym, |
| 38 | QuantizationStrategy.CHANNEL: kFp8StaticChannelSym, |
| 39 | QuantizationStrategy.TENSOR: kFp8StaticTensorSym, |
| 40 | } |
| 41 | |
| 42 | |
| 43 | def is_activation_quantization_format(format: str) -> bool: |
| 44 | _ACTIVATION_QUANTIZATION_FORMATS = [ |
| 45 | CompressionFormat.naive_quantized.value, |
| 46 | CompressionFormat.int_quantized.value, |
| 47 | CompressionFormat.float_quantized.value, |
| 48 | CompressionFormat.nvfp4_pack_quantized.value, |
| 49 | ] |
| 50 | return format in _ACTIVATION_QUANTIZATION_FORMATS |
| 51 | |
| 52 | |
| 53 | def should_ignore_layer( |
| 54 | layer_name: str | None, |
| 55 | ignore: Iterable[str] = tuple(), |
| 56 | fused_mapping: Mapping[str, list[str]] = MappingProxyType({}), |
| 57 | use_fnmatch: bool = False, |
| 58 | ) -> bool: |
| 59 | if layer_name is None: |
| 60 | return False |
| 61 | per_shard_matches = find_matching_patterns( |
| 62 | layer_name, ignore, fused_mapping, use_fnmatch=use_fnmatch |
| 63 | ) |
| 64 | shards_ignored = [len(matches) > 0 for matches in per_shard_matches] |
| 65 | if any(shards_ignored) and not all(shards_ignored): |
| 66 | raise ValueError( |
| 67 | f"Found different quantization schemes for the shards of " |
| 68 | f"{layer_name}. vLLM requires all to use the same scheme." |
| 69 | ) |
| 70 | return all(shards_ignored) |
| 71 | |
| 72 | |
| 73 | def find_matched_target( |
| 74 | layer_name: str | None, |
| 75 | module: Module, |
| 76 | targets: Iterable[str], |
| 77 | fused_mapping: Mapping[str, list[str]] = MappingProxyType({}), |
| 78 | ) -> str | None: |
| 79 | """Helper function to look up which "target" in the compressed-tensors |
| 80 | config that a layer corresponds to. |
| 81 | |
| 82 | Recall that a compressed-tensors configs has a concept of |
| 83 | config_groups, where each layer can be quantized with a different |
| 84 | scheme. |
| 85 | |
| 86 | targets in each config_group will be a list of either layer names |
| 87 | (or regexes corresponding to layer names) or names of torch Modules. |
| 88 | |
| 89 | First, we try to match the layer_name with a target |
| 90 | Second, we try to match the module's name with a target |
| 91 | Third, we try to map the layer_name to a list of fused module names. |
| 92 | *All* component module names must match in order for a match to be |
| 93 | successful. A successful match returns the first component target |
| 94 | |
| 95 | Args: |
| 96 | layer_name: layer name |
| 97 | module: torch.nn.Module |
| 98 | targets: list of targets to match the layer against |
| 99 | fused_mapping: map from fused layer names to its components |
| 100 | |
| 101 | """ |
| 102 | if layer_name is None: |
| 103 | layer_name = "" |
| 104 | |
| 105 | matched_target = ( |
| 106 | _find_first_match(layer_name, targets) |
| 107 | or _match_fused_layer(layer_name, targets, fused_mapping) |
| 108 | or _find_first_match(module.__class__.__name__, targets, True) |
| 109 | ) |
| 110 | |
| 111 | return matched_target |
| 112 | |
| 113 | |
| 114 | def _find_first_match( |
| 115 | value: str, targets: Iterable[str], check_contains: bool = False |
| 116 | ) -> str | None: |
| 117 | """Returns first element of target that matches value either |
| 118 | exactly or as a regex after 're:'. If check_contains is set to True, |
| 119 | additionally checks if the target string is contained within the value. |
| 120 | |
| 121 | Args: |
| 122 | value: string to compare the list of targets against |
| 123 | targets: list of targets to match the layer against |
| 124 | check_contains: whether or not to do a substring match |
| 125 | |
| 126 | """ |
| 127 | for target in targets: |
| 128 | if is_equal_or_regex_match(value, target, check_contains=check_contains): |
| 129 | return target |
| 130 | return None |
| 131 | |
| 132 | |
| 133 | def _match_fused_layer( |
| 134 | layer_name: str, |
| 135 | target_layers: Iterable[str], |
| 136 | fused_mapping: Mapping[str, list[str]], |
| 137 | ) -> str | None: |
| 138 | """Match a fused layer name to its corresponding individual layer in |
| 139 | target_layers. Returns first value in fused_mapping which matches targets |
| 140 | |
| 141 | Implements an "all" matching strategy where a fused layer matches iff |
| 142 | "all" of its components match |
| 143 | |
| 144 | Args: |
| 145 | layer_name: layer name |
| 146 | target_layers: list of targets to match the layer against |
| 147 | fused_mapping: map from fused layer names to its components |
| 148 | |
| 149 | Examples: |
| 150 | layer_name = "model.layers.0.self_attn.qkv_proj" |
| 151 | target_layers = ["model.layers.0.self_attn.q_proj", |
| 152 | "model.layers.0.self_attn.k_proj", |
| 153 | "model.layers.0.self_attn.v_proj"] |
| 154 | |
| 155 | """ |
| 156 | # find layer_name in mapping |
| 157 | fused = next((key for key in fused_mapping if layer_name.endswith(key)), None) |
| 158 | if fused is None: |
| 159 | return None |
| 160 | |
| 161 | # expand path of unfused components |
| 162 | unfused_paths = [ |
| 163 | layer_name.replace(fused, unfused) for unfused in fused_mapping[fused] |
| 164 | ] |
| 165 | |
| 166 | # for each unfused component, find a match in targets |
| 167 | unfused_matches: list[str | None] = [] |
| 168 | for unfused in unfused_paths: |
| 169 | for target in target_layers: |
| 170 | if is_equal_or_regex_match(unfused, target): |
| 171 | unfused_matches.append(target) |
| 172 | break |
| 173 | else: |
| 174 | unfused_matches.append(None) |
| 175 | |
| 176 | return unfused_matches[0] if all(unfused_matches) else None |