| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | from collections.abc import Sequence |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm import _custom_ops as ops |
| 10 | from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 |
| 11 | from vllm.model_executor.layers.quantization.utils import replace_parameter |
| 12 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 13 | GroupShape, |
| 14 | QuantKey, |
| 15 | kFp8StaticTensorSym, |
| 16 | ) |
| 17 | from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( |
| 18 | CUTLASS_BLOCK_FP8_SUPPORTED, |
| 19 | convert_to_channelwise, |
| 20 | cutlass_fp8_supported, |
| 21 | ) |
| 22 | from vllm.model_executor.utils import set_weight_attrs |
| 23 | from vllm.platforms import current_platform |
| 24 | |
| 25 | from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel |
| 26 | from .ScaledMMLinearKernel import ( |
| 27 | FP8ScaledMMLinearKernel, |
| 28 | FP8ScaledMMLinearLayerConfig, |
| 29 | Int8ScaledMMLinearKernel, |
| 30 | Int8ScaledMMLinearLayerConfig, |
| 31 | ) |
| 32 | |
| 33 | |
| 34 | class CutlassInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel): |
| 35 | @classmethod |
| 36 | def is_supported( |
| 37 | cls, compute_capability: int | None = None |
| 38 | ) -> tuple[bool, str | None]: |
| 39 | if not current_platform.is_cuda(): |
| 40 | return False, "requires CUDA." |
| 41 | return True, None |
| 42 | |
| 43 | @classmethod |
| 44 | def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 45 | return True, None |
| 46 | |
| 47 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 48 | w_q_name, w_s_name, i_s_name, i_zp_name, azp_adj_name = self.layer_param_names |
| 49 | config = self.config |
| 50 | # WEIGHT |
| 51 | # Cutlass kernels need transposed weight. |
| 52 | weight = getattr(layer, w_q_name) |
| 53 | replace_parameter( |
| 54 | layer, |
| 55 | w_q_name, |
| 56 | torch.nn.Parameter(weight.t().data, requires_grad=False), |
| 57 | ) |
| 58 | |
| 59 | # WEIGHT SCALE |
| 60 | # Cutlass kernels support only per-tensor and per-channel. |
| 61 | # If we have a fused module (QKV, MLP) with per tensor scales (thus N |
| 62 | # scales being passed to the kernel), convert to the per-channel case. |
| 63 | is_fused_module = len(layer.logical_widths) > 1 |
| 64 | weight_scale = getattr(layer, w_s_name) |
| 65 | if is_fused_module and not config.is_channelwise: |
| 66 | weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths) |
| 67 | replace_parameter( |
| 68 | layer, |
| 69 | w_s_name, |
| 70 | torch.nn.Parameter(weight_scale.data, requires_grad=False), |
| 71 | ) |
| 72 | |
| 73 | # INPUT SCALE |
| 74 | if config.is_static_input_scheme: |
| 75 | input_scale = getattr(layer, i_s_name) |
| 76 | |
| 77 | if config.input_symmetric: |
| 78 | replace_parameter( |
| 79 | layer, |
| 80 | i_s_name, |
| 81 | torch.nn.Parameter(input_scale.max(), requires_grad=False), |
| 82 | ) |
| 83 | setattr(layer, i_zp_name, None) |
| 84 | else: |
| 85 | input_zero_point = getattr(layer, i_zp_name) |
| 86 | |
| 87 | # reconstruct the ranges |
| 88 | int8_traits = torch.iinfo(torch.int8) |
| 89 | azps = input_zero_point.to(dtype=torch.int32) |
| 90 | range_max = (input_scale * (int8_traits.max - azps)).max() |
| 91 | range_min = (input_scale * (int8_traits.min - azps)).min() |
| 92 | |
| 93 | scale = (range_max - range_min) / (int8_traits.max - int8_traits.min) |
| 94 | replace_parameter( |
| 95 | layer, i_s_name, torch.nn.Parameter(scale, requires_grad=False) |
| 96 | ) |
| 97 | |
| 98 | # AZP loaded as int8 but used as int32 |
| 99 | azp = (int8_traits.min - range_min / scale).to(dtype=torch.int32) |
| 100 | replace_parameter( |
| 101 | layer, i_zp_name, torch.nn.Parameter(azp, requires_grad=False) |
| 102 | ) |
| 103 | |
| 104 | # azp_adj is the AZP adjustment term, used to account for weights. |
| 105 | # It does not depend on scales or azp, so it is the same for |
| 106 | # static and dynamic quantization. |
| 107 | # For more details, see csrc/quantization/w8a8/cutlass/Epilogues.md |
| 108 | # https://github.com/vllm-project/vllm/blob/main/csrc/quantization/w8a8/cutlass/Epilogues.md |
| 109 | if not config.input_symmetric: |
| 110 | weight = getattr(layer, w_q_name) |
| 111 | azp_adj = weight.sum(dim=0, keepdim=True, dtype=torch.int32) |
| 112 | if config.is_static_input_scheme: |
| 113 | # cutlass_w8a8 requires azp to be folded into azp_adj |
| 114 | # in the per-tensor case |
| 115 | azp_adj = getattr(layer, i_zp_name) * azp_adj |
| 116 | setattr( |
| 117 | layer, |
| 118 | azp_adj_name, |
| 119 | torch.nn.Parameter(azp_adj, requires_grad=False), |
| 120 | ) |
| 121 | |
| 122 | def apply_weights( |
| 123 | self, |
| 124 | layer: torch.nn.Module, |
| 125 | x: torch.Tensor, |
| 126 | bias: torch.Tensor | None = None, |
| 127 | ) -> torch.Tensor: |
| 128 | w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer) |
| 129 | |
| 130 | # ops.scaled_int8_quant supports both dynamic and static quant: |
| 131 | # * dynamic, i_s is None and x_s computed from x. |
| 132 | # * static, i_s is scalar and x_s is i_s. |
| 133 | symmetric = azp_adj is None |
| 134 | x_q, x_s, x_zp = ops.scaled_int8_quant( |
| 135 | x.contiguous(), i_s, i_zp, symmetric=symmetric |
| 136 | ) |
| 137 | |
| 138 | if x_zp is not None: |
| 139 | # Currently, static is always per-tensor and dynamic is per-token |
| 140 | static = i_zp is not None |
| 141 | azp = None if static else x_zp |
| 142 | return ops.cutlass_scaled_mm_azp( |
| 143 | x_q, |
| 144 | w_q, |
| 145 | scale_a=x_s, |
| 146 | scale_b=w_s, |
| 147 | out_dtype=x.dtype, |
| 148 | azp_adj=azp_adj, |
| 149 | azp=azp, |
| 150 | bias=bias, |
| 151 | ) |
| 152 | return ops.cutlass_scaled_mm( |
| 153 | x_q, w_q, scale_a=x_s, scale_b=w_s, out_dtype=x.dtype, bias=bias |
| 154 | ) |
| 155 | |
| 156 | |
| 157 | class CutlassFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): |
| 158 | def __init__( |
| 159 | self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str] |
| 160 | ) -> None: |
| 161 | self.logical_output_size: int | None = None |
| 162 | super().__init__(c, layer_param_names) |
| 163 | |
| 164 | @classmethod |
| 165 | def is_supported( |
| 166 | cls, compute_capability: int | None = None |
| 167 | ) -> tuple[bool, str | None]: |
| 168 | if not cutlass_fp8_supported(): |
| 169 | return False, "CUTLASS FP8 kernels not available" |
| 170 | return True, None |
| 171 | |
| 172 | @classmethod |
| 173 | def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 174 | return True, None |
| 175 | |
| 176 | def input_quant_key(self) -> QuantKey | None: |
| 177 | """Only static per-tensor activation quantization is supported for external |
| 178 | quantization.""" |
| 179 | if self.config.activation_quant_key == kFp8StaticTensorSym: |
| 180 | return kFp8StaticTensorSym |
| 181 | return None |
| 182 | |
| 183 | @staticmethod |
| 184 | def _pad_to_alignment( |
| 185 | x: torch.Tensor, dim: int, alignment: int, value: float = 0.0 |
| 186 | ) -> torch.Tensor: |
| 187 | """Pad tensor ``x`` along ``dim`` to the next multiple of |
| 188 | ``alignment``.""" |
| 189 | remainder = x.shape[dim] % alignment |
| 190 | if remainder == 0: |
| 191 | return x |
| 192 | pad_size = alignment - remainder |
| 193 | pad_spec = [0] * (2 * x.dim()) |
| 194 | pad_spec[-(2 * dim + 1)] = pad_size |
| 195 | return torch.nn.functional.pad(x, pad_spec, value=value) |
| 196 | |
| 197 | @staticmethod |
| 198 | def padded_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: |
| 199 | if loaded_weight.shape != param.shape: |
| 200 | slices = tuple(slice(0, s) for s in loaded_weight.shape) |
| 201 | param.data[slices].copy_(loaded_weight) |
| 202 | else: |
| 203 | param.data.copy_(loaded_weight.view(param.shape)) |
| 204 | |
| 205 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 206 | weight_name, weight_scale_name, _, _ = self.layer_param_names |
| 207 | weight = getattr(layer, weight_name) |
| 208 | |
| 209 | # keep the logical output width so runtime can slice away static padding. |
| 210 | self.logical_output_size = weight.shape[1] |
| 211 | |
| 212 | pad_k = (16 - weight.shape[0] % 16) % 16 |
| 213 | pad_n = (16 - weight.shape[1] % 16) % 16 |
| 214 | if pad_k == 0 and pad_n == 0: |
| 215 | return |
| 216 | |
| 217 | # B is column-major [K, N] |
| 218 | padded_weight = torch.nn.functional.pad( |
| 219 | weight.t().contiguous(), |
| 220 | (0, pad_k, 0, pad_n), |
| 221 | ).t() |
| 222 | replace_parameter(layer, weight_name, padded_weight.data) |
| 223 | set_weight_attrs( |
| 224 | getattr(layer, weight_name), |
| 225 | { |
| 226 | "weight_loader": self.padded_weight_loader, |
| 227 | }, |
| 228 | ) |
| 229 | |
| 230 | weight_scale = getattr(layer, weight_scale_name, None) |
| 231 | if weight_scale is not None and pad_n > 0 and weight_scale.numel() > 1: |
| 232 | flat_scale = weight_scale.reshape(-1) |
| 233 | padded_scale = self._pad_to_alignment( |
| 234 | flat_scale, dim=0, alignment=16, value=1.0 |
| 235 | ).view(-1, *weight_scale.shape[1:]) |
| 236 | replace_parameter(layer, weight_scale_name, padded_scale.data) |
| 237 | set_weight_attrs( |
| 238 | getattr(layer, weight_name), |
| 239 | { |
| 240 | "weight_loader": self.padded_weight_loader, |
| 241 | }, |
| 242 | ) |
| 243 | |
| 244 | def apply_scaled_mm( |
| 245 | self, |
| 246 | *, |
| 247 | A: torch.Tensor, |
| 248 | B: torch.Tensor, |
| 249 | out_dtype: torch.dtype, |
| 250 | As: torch.Tensor, |
| 251 | Bs: torch.Tensor, |
| 252 | bias: torch.Tensor | None, |
| 253 | output_shape: list, |
| 254 | ) -> torch.Tensor: |
| 255 | padded_k, padded_n = B.shape |
| 256 | output_size = self.logical_output_size |
| 257 | assert output_size is not None |
| 258 | pad_k = padded_k - A.shape[1] |
| 259 | pad_n = padded_n - output_size |
| 260 | |
| 261 | if pad_k > 0: |
| 262 | A = self._pad_to_alignment(A, dim=1, alignment=16) |
| 263 | if pad_n > 0 and bias is not None: |
| 264 | bias = self._pad_to_alignment(bias, dim=0, alignment=16) |
| 265 | |
| 266 | output = ops.cutlass_scaled_mm( |
| 267 | A, B, out_dtype=out_dtype, scale_a=As, scale_b=Bs, bias=bias |
| 268 | ) |
| 269 | |
| 270 | if pad_n > 0: |
| 271 | output = output[..., :output_size].contiguous() |
| 272 | |
| 273 | return output.view(*output_shape[:-1], output_size) |
| 274 | |
| 275 | |
| 276 | class CutlassFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): |
| 277 | def __init__(self, config: FP8ScaledMMLinearLayerConfig) -> None: |
| 278 | super().__init__(config) |
| 279 | act_scale_descriptor = config.activation_quant_key.scale |
| 280 | self.quant_fp8 = QuantFP8( |
| 281 | static=act_scale_descriptor.static, |
| 282 | group_shape=act_scale_descriptor.group_shape, |
| 283 | num_token_padding=self.get_output_padding(), |
| 284 | use_ue8m0=False, |
| 285 | column_major_scales=True, |
| 286 | ) |
| 287 | |
| 288 | @classmethod |
| 289 | def is_supported(cls, compute_capability=None): |
| 290 | if not CUTLASS_BLOCK_FP8_SUPPORTED: |
| 291 | return ( |
| 292 | False, |
| 293 | "The device compute capability of" |
| 294 | f"{compute_capability} is not supported.", |
| 295 | ) |
| 296 | return True, None |
| 297 | |
| 298 | @classmethod |
| 299 | def can_implement(cls, config: FP8ScaledMMLinearLayerConfig): |
| 300 | can_implement_base, reason = super().can_implement(config) |
| 301 | if not can_implement_base: |
| 302 | return can_implement_base, reason |
| 303 | |
| 304 | act_quant_desc = config.activation_quant_key.scale |
| 305 | if act_quant_desc.group_shape != GroupShape(1, 128): |
| 306 | return ( |
| 307 | False, |
| 308 | "Supports only dynamic per token group activation " |
| 309 | "quantization with group_shape=(1,128).", |
| 310 | ) |
| 311 | return True, None |
| 312 | |
| 313 | def apply_block_scaled_mm( |
| 314 | self, |
| 315 | A: torch.Tensor, |
| 316 | B: torch.Tensor, |
| 317 | As: torch.Tensor, |
| 318 | Bs: torch.Tensor, |
| 319 | ) -> torch.Tensor: |
| 320 | out_dtype = self.config.out_dtype |
| 321 | return ops.cutlass_scaled_mm( |
| 322 | A, |
| 323 | B.T, |
| 324 | out_dtype=out_dtype, |
| 325 | scale_a=As, |
| 326 | scale_b=Bs.T, |
| 327 | ) |
| 328 | |
| 329 | |
| 330 | def cutlass_scaled_mm( |
| 331 | A: torch.Tensor, |
| 332 | B: torch.Tensor, |
| 333 | As: torch.Tensor, |
| 334 | Bs: torch.Tensor, |
| 335 | block_size: list[int], |
| 336 | output_dtype: torch.dtype = torch.float16, |
| 337 | ) -> torch.Tensor: |
| 338 | return ops.cutlass_scaled_mm( |
| 339 | A, |
| 340 | B.T, |
| 341 | out_dtype=output_dtype, |
| 342 | scale_a=As, |
| 343 | scale_b=Bs.T, |
| 344 | ) |