| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | import torch.nn.functional as F |
| 6 | |
| 7 | from vllm import _custom_ops as ops |
| 8 | from vllm._aiter_ops import rocm_aiter_ops |
| 9 | from vllm.model_executor.custom_op import CustomOp |
| 10 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 11 | GroupShape, |
| 12 | get_fp8_min_max, |
| 13 | group_broadcast, |
| 14 | prep_scale_for_group_broadcast, |
| 15 | ) |
| 16 | from vllm.platforms import current_platform |
| 17 | from vllm.utils.deep_gemm import ( |
| 18 | DeepGemmQuantScaleFMT, |
| 19 | is_deep_gemm_e8m0_used, |
| 20 | is_deep_gemm_supported, |
| 21 | ) |
| 22 | |
| 23 | _FP8_DTYPE = current_platform.fp8_dtype() |
| 24 | _FP8_MIN, _FP8_MAX = get_fp8_min_max() |
| 25 | _FP8_MIN_SCALING_FACTOR = 1.0 / (_FP8_MAX * 512.0) |
| 26 | |
| 27 | |
| 28 | # --8<-- [start:quant_fp8] |
| 29 | @CustomOp.register("quant_fp8") |
| 30 | class QuantFP8(CustomOp): |
| 31 | """Quantize input tensor to FP8 (per-tensor, per-token, per-channel, or per-group). |
| 32 | This CustomOp supports both static and dynamic quantization. |
| 33 | """ |
| 34 | |
| 35 | # --8<-- [end:quant_fp8] |
| 36 | |
| 37 | def __init__( |
| 38 | self, |
| 39 | static: bool, |
| 40 | group_shape: GroupShape, |
| 41 | num_token_padding: int | None = None, |
| 42 | column_major_scales: bool = False, |
| 43 | tma_aligned_scales: bool = False, |
| 44 | use_ue8m0: bool | None = None, # for Torch compile |
| 45 | compile_native: bool = True, |
| 46 | ): |
| 47 | """Args: |
| 48 | static: static or dynamic quantization |
| 49 | group_shape: quantization group shape (PER_TOKEN, PER_TENSOR, |
| 50 | PER_CHANNEL, or arbitrary block size) |
| 51 | num_token_padding: Pad the token dimension of output to this |
| 52 | size |
| 53 | tma_aligned_scales: For group quantization, output scales in |
| 54 | TMA-aligned layout |
| 55 | column_major_scales: For group quantization, output scales in |
| 56 | column major format |
| 57 | compile_native: Manually compile forward_native if compile mode > None |
| 58 | |
| 59 | """ |
| 60 | super().__init__(compile_native=compile_native) |
| 61 | self.static = static |
| 62 | self.group_shape = group_shape |
| 63 | self.use_per_token_if_dynamic = group_shape == GroupShape.PER_TOKEN |
| 64 | self.num_token_padding = num_token_padding |
| 65 | self.column_major_scales = column_major_scales |
| 66 | self.tma_aligned_scales = tma_aligned_scales |
| 67 | self.use_ue8m0 = is_deep_gemm_e8m0_used() if use_ue8m0 is None else use_ue8m0 |
| 68 | self.use_deep_gemm_supported = is_deep_gemm_supported() |
| 69 | |
| 70 | self.use_aiter = rocm_aiter_ops.is_linear_fp8_enabled() |
| 71 | |
| 72 | self.is_group_quant = group_shape.is_per_group() |
| 73 | if self.is_group_quant: |
| 74 | self.group_size = group_shape.col |
| 75 | else: |
| 76 | self.use_per_token_if_dynamic = group_shape == GroupShape.PER_TOKEN |
| 77 | if not static: |
| 78 | assert group_shape in (GroupShape.PER_TOKEN, GroupShape.PER_TENSOR), ( |
| 79 | "Only per-token or per-tensor scales are supported for dynamic " |
| 80 | "non-group quantization." |
| 81 | ) |
| 82 | |
| 83 | def forward_cuda( |
| 84 | self, |
| 85 | x: torch.Tensor, |
| 86 | scale: torch.Tensor | None = None, |
| 87 | scale_ub: torch.Tensor | None = None, |
| 88 | use_triton: bool = False, |
| 89 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 90 | from vllm.model_executor.layers.quantization.utils import fp8_utils |
| 91 | |
| 92 | if ( |
| 93 | self.is_group_quant |
| 94 | and self.use_ue8m0 |
| 95 | and self.use_deep_gemm_supported |
| 96 | and not self.static |
| 97 | and self.group_size == 128 |
| 98 | and (DeepGemmQuantScaleFMT.from_oracle() == DeepGemmQuantScaleFMT.UE8M0) |
| 99 | ): |
| 100 | return fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm( |
| 101 | x, |
| 102 | group_size=self.group_size, |
| 103 | use_ue8m0=True, |
| 104 | ) |
| 105 | |
| 106 | if self.is_group_quant and not self.static: |
| 107 | assert scale is None, "Dynamic group quantization does not use scale" |
| 108 | |
| 109 | return fp8_utils.per_token_group_quant_fp8( |
| 110 | x, |
| 111 | group_size=self.group_size, |
| 112 | column_major_scales=self.column_major_scales, |
| 113 | tma_aligned_scales=self.tma_aligned_scales, |
| 114 | dtype=_FP8_DTYPE, |
| 115 | use_ue8m0=self.use_ue8m0, |
| 116 | ) |
| 117 | |
| 118 | assert (scale is not None) == self.static |
| 119 | assert scale_ub is None or ( |
| 120 | not self.static |
| 121 | and self.group_shape == GroupShape.PER_TOKEN |
| 122 | and scale_ub.numel() == 1 |
| 123 | ) |
| 124 | |
| 125 | return ops.scaled_fp8_quant( |
| 126 | x, |
| 127 | scale, |
| 128 | num_token_padding=self.num_token_padding, |
| 129 | scale_ub=scale_ub, |
| 130 | use_per_token_if_dynamic=self.use_per_token_if_dynamic, |
| 131 | group_shape=(self.group_shape.row, self.group_shape.col) |
| 132 | if self.static |
| 133 | else None, |
| 134 | ) |
| 135 | |
| 136 | def forward_hip( |
| 137 | self, |
| 138 | x: torch.Tensor, |
| 139 | scale: torch.Tensor | None = None, |
| 140 | scale_ub: torch.Tensor | None = None, |
| 141 | use_triton: bool = False, |
| 142 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 143 | use_aiter_quant = self.use_aiter and scale_ub is None and x.is_contiguous() |
| 144 | use_aiter_per_tensor_quant = ( |
| 145 | use_aiter_quant and self.group_shape.is_per_tensor() |
| 146 | ) |
| 147 | use_aiter_per_token_quant = use_aiter_quant and self.group_shape.is_per_token() |
| 148 | |
| 149 | use_aiter_per_group_quant = use_aiter_quant and self.group_shape.is_per_group() |
| 150 | |
| 151 | if use_aiter_per_group_quant: |
| 152 | return rocm_aiter_ops.group_fp8_quant(x, self.group_size) |
| 153 | if use_aiter_per_tensor_quant: |
| 154 | return rocm_aiter_ops.per_tensor_quant(x, _FP8_DTYPE, scale) |
| 155 | if use_aiter_per_token_quant: |
| 156 | return rocm_aiter_ops.per_token_quant(x, _FP8_DTYPE, scale) |
| 157 | |
| 158 | # Fall back to the CUDA implementation. Dispatch on the class, not |
| 159 | # through self: a subclass may override forward_cuda with a different |
| 160 | # signature (_DecodeConcatQuantFP8 does), and a virtual call would |
| 161 | # re-enter it with mismatched arguments. |
| 162 | return QuantFP8.forward_cuda(self, x, scale, scale_ub) |
| 163 | |
| 164 | def forward_xpu( |
| 165 | self, |
| 166 | x: torch.Tensor, |
| 167 | scale: torch.Tensor | None = None, |
| 168 | scale_ub: torch.Tensor | None = None, |
| 169 | use_triton: bool = False, |
| 170 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 171 | if self.is_group_quant and not self.static: |
| 172 | from vllm.model_executor.layers.quantization.utils import fp8_utils |
| 173 | |
| 174 | return fp8_utils.per_token_group_quant_fp8( |
| 175 | x, |
| 176 | group_size=self.group_size, |
| 177 | column_major_scales=self.column_major_scales, |
| 178 | dtype=_FP8_DTYPE, |
| 179 | use_ue8m0=self.use_ue8m0, |
| 180 | ) |
| 181 | # Dispatch on the class, for the same reason as in forward_hip. |
| 182 | return QuantFP8.forward_cuda(self, x, scale, scale_ub, use_triton) |
| 183 | |
| 184 | def forward_native( |
| 185 | self, |
| 186 | x: torch.Tensor, |
| 187 | scale: torch.Tensor | None = None, |
| 188 | scale_ub: torch.Tensor | None = None, |
| 189 | use_triton: bool = False, |
| 190 | ): |
| 191 | if self.is_group_quant and not self.static: |
| 192 | assert scale is None, "Dynamic group quantization does not use scale" |
| 193 | return self._quantize_group_native(x) |
| 194 | |
| 195 | assert (scale is not None) == self.static |
| 196 | assert scale_ub is None or ( |
| 197 | not self.static |
| 198 | and self.group_shape == GroupShape.PER_TOKEN |
| 199 | and scale_ub.numel() == 1 |
| 200 | ) |
| 201 | |
| 202 | if scale is None: |
| 203 | if self.group_shape == GroupShape.PER_TOKEN: |
| 204 | x_max, _ = x.abs().max(dim=-1) |
| 205 | x_max = x_max.unsqueeze(-1).to(torch.float32) |
| 206 | if scale_ub is not None: |
| 207 | x_max = x_max.clamp(max=scale_ub) |
| 208 | else: |
| 209 | x_max = x.abs().max().unsqueeze(-1).to(torch.float32) |
| 210 | |
| 211 | scale = (x_max / _FP8_MAX).clamp(min=_FP8_MIN_SCALING_FACTOR) |
| 212 | else: |
| 213 | scale = prep_scale_for_group_broadcast(scale, x, self.group_shape) |
| 214 | |
| 215 | # Even for dynamic per-token scales, |
| 216 | # reciprocal performs slightly better than division |
| 217 | out = ( |
| 218 | x.to(torch.float32) |
| 219 | * group_broadcast(scale.to(torch.float32), x.shape[-2:]).reciprocal() |
| 220 | ) |
| 221 | out = out.clamp(_FP8_MIN, _FP8_MAX).to(_FP8_DTYPE) |
| 222 | |
| 223 | # This currently generates an extra Triton kernel in compilation. |
| 224 | # Fortunately, we don't use padding if compiling. |
| 225 | # TODO(luka): benchmark torch._scaled_mm to hopefully remove padding |
| 226 | # in general. |
| 227 | if self.num_token_padding is not None: |
| 228 | padding = max(self.num_token_padding - out.size(0), 0) |
| 229 | out = F.pad(out, (0, 0, 0, padding), "constant", 0.0) |
| 230 | |
| 231 | return out, scale |
| 232 | |
| 233 | def _quantize_group_native( |
| 234 | self, x: torch.Tensor |
| 235 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 236 | orig_shape = x.shape |
| 237 | hidden_dim = x.shape[-1] |
| 238 | num_groups = (hidden_dim + self.group_size - 1) // self.group_size |
| 239 | padded_dim = num_groups * self.group_size |
| 240 | |
| 241 | if padded_dim != hidden_dim: |
| 242 | padding = padded_dim - hidden_dim |
| 243 | x = F.pad(x, (0, padding), mode="constant", value=0.0) |
| 244 | |
| 245 | x_grouped = x.view(-1, num_groups, self.group_size) |
| 246 | absmax = x_grouped.abs().max(dim=-1, keepdim=True)[0].float() |
| 247 | scales_raw = absmax / _FP8_MAX |
| 248 | if self.use_ue8m0: |
| 249 | scales_raw = torch.exp2(torch.ceil(torch.log2(scales_raw))) |
| 250 | scales = (scales_raw).clamp(min=_FP8_MIN_SCALING_FACTOR) |
| 251 | |
| 252 | x_scaled = x_grouped / scales |
| 253 | x_quant = x_scaled.clamp(_FP8_MIN, _FP8_MAX).to(_FP8_DTYPE) |
| 254 | |
| 255 | x_quant = x_quant.view(-1, padded_dim) |
| 256 | if padded_dim != hidden_dim: |
| 257 | x_quant = x_quant[..., :hidden_dim] |
| 258 | x_quant = x_quant.view(orig_shape) |
| 259 | |
| 260 | scales = scales.squeeze(-1) |
| 261 | scales = scales.reshape(orig_shape[:-1] + (num_groups,)) |
| 262 | |
| 263 | if self.column_major_scales: |
| 264 | scales = scales.transpose(-2, -1).contiguous().transpose(-1, -2) |
| 265 | |
| 266 | return x_quant, scales |