| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from vllm.logger import init_logger |
| 7 | from vllm.model_executor.layers.quantization.base_config import ( |
| 8 | QuantizationConfig, |
| 9 | QuantizeMethodBase, |
| 10 | ) |
| 11 | from vllm.model_executor.utils import is_weights_pre_processed |
| 12 | from vllm.platforms import current_platform |
| 13 | from vllm.utils.torch_utils import is_quantized_kv_cache |
| 14 | from vllm.v1.kv_cache_interface import kv_cache_uses_per_token_head_scales |
| 15 | |
| 16 | logger = init_logger(__name__) |
| 17 | |
| 18 | |
| 19 | class KVCacheScaleParameter(torch.nn.Parameter): |
| 20 | """Scalar parameter for KV-cache scales. |
| 21 | |
| 22 | Initialized to -1.0 (an invalid sentinel) so call sites just write |
| 23 | `KVCacheScaleParameter()`. The `weight_loader` accepts shape `()` or |
| 24 | `(1,)` and rejects anything else — per-head scales go through a separate |
| 25 | path (compressed-tensors' `_tp_aware_loader`), not this one. Per-instance |
| 26 | overrides still work because instance attribute assignment shadows this |
| 27 | class-level loader. |
| 28 | """ |
| 29 | |
| 30 | def __new__(cls) -> "KVCacheScaleParameter": |
| 31 | return super().__new__(cls, torch.tensor(-1.0), requires_grad=False) |
| 32 | |
| 33 | @staticmethod |
| 34 | def weight_loader(param: torch.nn.Parameter, loaded_weight: torch.Tensor) -> None: |
| 35 | if loaded_weight.numel() != 1: |
| 36 | raise ValueError( |
| 37 | f"KV-cache scale expects a scalar weight, got shape " |
| 38 | f"{tuple(loaded_weight.shape)}" |
| 39 | ) |
| 40 | param.data.copy_(loaded_weight.reshape(())) |
| 41 | |
| 42 | |
| 43 | class BaseKVCacheMethod(QuantizeMethodBase): |
| 44 | """Quant method that adds `_k_scale` and `_v_scale` attributes to the |
| 45 | Attention layer to support loading those scaling factors from checkpoints. |
| 46 | The k/v_scale will be used to: |
| 47 | - quantize k/v_cache entries before saving them to the cache |
| 48 | - dequantize k/v_cache entries before fetching them from the cache |
| 49 | |
| 50 | Args: |
| 51 | quant_config: the appropriate QuantizationConfig |
| 52 | |
| 53 | """ |
| 54 | |
| 55 | supports_pre_processed_weights = True |
| 56 | |
| 57 | def __init__(self, quant_config: QuantizationConfig): |
| 58 | self.quant_config = quant_config |
| 59 | |
| 60 | def create_weights(self, layer: torch.nn.Module): |
| 61 | """Create "weight" (aka q_scale, k_scale and v_scale) |
| 62 | for an attention layer. |
| 63 | """ |
| 64 | # Initialize the Q and KV cache scales to -1.0, an invalid value. |
| 65 | # If the q and k/v_scales appear in the checkpoint, it will be |
| 66 | # overwritten when loading weights. |
| 67 | layer.q_scale = KVCacheScaleParameter() |
| 68 | layer.k_scale = KVCacheScaleParameter() |
| 69 | layer.v_scale = KVCacheScaleParameter() |
| 70 | # Initialize P = softmax(QK^T) scales |
| 71 | layer.prob_scale = KVCacheScaleParameter() |
| 72 | |
| 73 | def apply(self, layer: torch.nn.Module) -> torch.Tensor: |
| 74 | raise RuntimeError(f"{self.__class__.__name__}.apply should not be called.") |
| 75 | |
| 76 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 77 | if is_weights_pre_processed(): |
| 78 | # Scale buffers were exported with post-processed values; drop the |
| 79 | # consumed scale placeholders and rebuild the host-side copies. |
| 80 | for name in ("k_scale", "v_scale", "q_scale", "prob_scale"): |
| 81 | if hasattr(layer, name): |
| 82 | delattr(layer, name) |
| 83 | if hasattr(layer, "_k_scale"): |
| 84 | layer._k_scale_float = layer._k_scale.item() |
| 85 | layer._v_scale_float = layer._v_scale.item() |
| 86 | layer._q_scale_float = layer._q_scale.item() |
| 87 | layer._k_scale_cpu = layer._k_scale.detach().to("cpu") |
| 88 | layer._v_scale_cpu = layer._v_scale.detach().to("cpu") |
| 89 | return |
| 90 | |
| 91 | # skip if there are no weights to process (for example, weight reloading) |
| 92 | if not hasattr(layer, "q_scale"): |
| 93 | assert not hasattr(layer, "k_scale") |
| 94 | assert not hasattr(layer, "v_scale") |
| 95 | assert not hasattr(layer, "prob_scale") |
| 96 | return |
| 97 | |
| 98 | # Per-token-head quantized KV cache: scales are computed dynamically |
| 99 | # per (token, head) in the kernel at cache-write time. |
| 100 | if kv_cache_uses_per_token_head_scales(layer.kv_cache_dtype): |
| 101 | layer._k_scale.copy_(1.0) |
| 102 | layer._v_scale.copy_(1.0) |
| 103 | layer._k_scale_float = 1.0 |
| 104 | layer._v_scale_float = 1.0 |
| 105 | del layer.k_scale |
| 106 | del layer.v_scale |
| 107 | del layer.q_scale |
| 108 | del layer.prob_scale |
| 109 | return |
| 110 | |
| 111 | # If the kv-cache is not quantized, we enforce the k/v_scale to be 1.0 |
| 112 | # regardless whether the kv-scale is available in the checkpoint. |
| 113 | # No need to process kv scales after loading if we are going to |
| 114 | # calculate them on the fly. |
| 115 | if is_quantized_kv_cache(layer.kv_cache_dtype): |
| 116 | if layer.k_scale > 0.0 and layer.v_scale > 0.0: |
| 117 | # We prefer to use separate k_scale and v_scale if present |
| 118 | k_scale = layer.k_scale.to("cpu").tolist() |
| 119 | v_scale = layer.v_scale.to("cpu").tolist() |
| 120 | if current_platform.is_fp8_fnuz(): |
| 121 | k_scale *= 2 |
| 122 | v_scale *= 2 |
| 123 | elif layer.k_scale < 0.0 and layer.v_scale < 0.0: |
| 124 | # If no scales were loaded (both scales are invalid negative |
| 125 | # values), use the default value of 1.0 |
| 126 | k_scale = 1.0 |
| 127 | v_scale = 1.0 |
| 128 | else: |
| 129 | # If we find a single kv_scale in the checkpoint, we remap |
| 130 | # kv_scale to k_scale during weight loading, and duplicate |
| 131 | # k_scale to v_scale here |
| 132 | assert layer.k_scale > 0.0 |
| 133 | scale_to_duplicate = max(layer.k_scale, layer.v_scale) |
| 134 | k_scale = scale_to_duplicate.to("cpu").tolist() |
| 135 | v_scale = scale_to_duplicate.to("cpu").tolist() |
| 136 | if current_platform.is_fp8_fnuz(): |
| 137 | k_scale *= 2 |
| 138 | v_scale *= 2 |
| 139 | |
| 140 | if not isinstance(k_scale, float) or not isinstance(v_scale, float): |
| 141 | raise ValueError( |
| 142 | "Only support per-tensor scaling factor for fp8 KV cache" |
| 143 | ) |
| 144 | |
| 145 | if layer.q_scale < 0.0: |
| 146 | logger.warning_once( |
| 147 | "Checkpoint does not provide a q scaling factor. " |
| 148 | "Setting it to k_scale. This only matters for " |
| 149 | "FP8 Attention backends (flash-attn or flashinfer)." |
| 150 | ) |
| 151 | layer._q_scale.copy_(k_scale) |
| 152 | layer._q_scale_float = k_scale |
| 153 | |
| 154 | # These are used in the final Attention.forward() |
| 155 | layer._k_scale.copy_(k_scale) |
| 156 | layer._v_scale.copy_(v_scale) |
| 157 | layer._k_scale_float = k_scale |
| 158 | layer._v_scale_float = v_scale |
| 159 | # Sync host (cpu) scale copies read by AITER fused kernels. |
| 160 | layer._k_scale_cpu.fill_(k_scale) |
| 161 | layer._v_scale_cpu.fill_(v_scale) |
| 162 | if k_scale == 1.0 and v_scale == 1.0 and "e5m2" not in layer.kv_cache_dtype: |
| 163 | logger.warning_once( |
| 164 | "Using KV cache scaling factor 1.0 for fp8_e4m3. " |
| 165 | "If this is unintended, verify that k/v_scale " |
| 166 | "scaling factors are properly set in the checkpoint." |
| 167 | ) |
| 168 | |
| 169 | if layer.q_scale > 0.0: |
| 170 | q_scale = layer.q_scale |
| 171 | if current_platform.is_fp8_fnuz(): |
| 172 | q_scale *= 2 |
| 173 | else: |
| 174 | q_scale = 1.0 |
| 175 | if layer.prob_scale > 0.0: |
| 176 | prob_scale = layer.prob_scale |
| 177 | if current_platform.is_fp8_fnuz(): |
| 178 | prob_scale *= 2 |
| 179 | else: |
| 180 | prob_scale = 1.0 |
| 181 | |
| 182 | is_singleton_float = ( |
| 183 | lambda x: isinstance(x, float) |
| 184 | or isinstance(x, torch.Tensor) |
| 185 | and x.numel() == 1 |
| 186 | and x.is_floating_point() |
| 187 | ) |
| 188 | if not is_singleton_float(q_scale) or not is_singleton_float(prob_scale): |
| 189 | raise ValueError( |
| 190 | "Only support per-tensor scaling factorfor fp8-quantized Q/prob" |
| 191 | ) |
| 192 | |
| 193 | # These are used in the final Attention.forward() |
| 194 | layer._q_scale.copy_(q_scale) |
| 195 | layer._q_scale_float = ( |
| 196 | q_scale.item() if isinstance(q_scale, torch.Tensor) else q_scale |
| 197 | ) |
| 198 | |
| 199 | layer._prob_scale.copy_(prob_scale) |
| 200 | if layer.kv_cache_dtype == "fp8" and (q_scale == 1.0 or prob_scale == 1.0): |
| 201 | logger.warning_once( |
| 202 | f"Using uncalibrated q_scale {q_scale} and/or prob_scale " |
| 203 | f"{prob_scale} with fp8 attention. This may cause accuracy " |
| 204 | "issues. Please make sure q/prob scaling factors are " |
| 205 | "available in the fp8 checkpoint." |
| 206 | ) |
| 207 | |
| 208 | del layer.k_scale |
| 209 | del layer.v_scale |
| 210 | del layer.q_scale |
| 211 | del layer.prob_scale |