| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 8 | get_fp8_min_max, |
| 9 | group_broadcast, |
| 10 | ) |
| 11 | from vllm.platforms import current_platform |
| 12 | from vllm.utils.deep_gemm import _ceil_to_ue8m0, is_deep_gemm_e8m0_used |
| 13 | from vllm.utils.math_utils import round_up |
| 14 | |
| 15 | FP8_DTYPE = current_platform.fp8_dtype() |
| 16 | DEVICE = current_platform.device_type |
| 17 | |
| 18 | |
| 19 | def as_float32_tensor(x: float | torch.Tensor) -> torch.Tensor: |
| 20 | return torch.as_tensor(x, dtype=torch.float32, device=DEVICE) |
| 21 | |
| 22 | |
| 23 | def ref_dynamic_per_token_quant( |
| 24 | x: torch.Tensor, quant_dtype: torch.dtype, scale_ub: torch.Tensor | None = None |
| 25 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 26 | assert quant_dtype in [torch.int8, FP8_DTYPE] |
| 27 | if scale_ub is not None: |
| 28 | assert quant_dtype == FP8_DTYPE |
| 29 | |
| 30 | if quant_dtype == torch.int8: |
| 31 | qtype_traits = torch.iinfo(quant_dtype) |
| 32 | qtype_traits_min = qtype_traits.min |
| 33 | qtype_traits_max = qtype_traits.max |
| 34 | else: |
| 35 | qtype_traits_min, qtype_traits_max = get_fp8_min_max() |
| 36 | qtype_max = as_float32_tensor(qtype_traits_max) |
| 37 | s_1 = as_float32_tensor(1.0) |
| 38 | s_512 = as_float32_tensor(512.0) |
| 39 | |
| 40 | # For fp8, in order to match the cuda kernel output, we have to do exactly |
| 41 | # the same operations as in the corresponding fp8 kernel to prevent |
| 42 | # rounding errors. |
| 43 | |
| 44 | # Compute scales |
| 45 | x_token_max, _ = x.abs().max(dim=-1) |
| 46 | x_token_max = as_float32_tensor(x_token_max) |
| 47 | if scale_ub is not None: |
| 48 | x_token_max = x_token_max.clamp(max=scale_ub) |
| 49 | scales = (x_token_max / qtype_max)[:, None] |
| 50 | |
| 51 | # Quant |
| 52 | if quant_dtype == torch.int8: |
| 53 | iscales = as_float32_tensor(s_1 / scales) |
| 54 | torch_out = as_float32_tensor(x) * iscales |
| 55 | torch_out = torch_out.round() |
| 56 | torch_out = torch_out.clamp(qtype_traits_min, qtype_traits_max).to(quant_dtype) |
| 57 | else: |
| 58 | assert quant_dtype == FP8_DTYPE |
| 59 | min_scaling_factor = s_1 / (qtype_max * s_512) |
| 60 | scales = scales.clamp(min=min_scaling_factor) |
| 61 | torch_out = as_float32_tensor(x) / scales |
| 62 | torch_out = torch_out.clamp(qtype_traits_min, qtype_traits_max).to(quant_dtype) |
| 63 | |
| 64 | return torch_out, scales |
| 65 | |
| 66 | |
| 67 | # The int8 version is very similar. Incorporate the int8 version, like in |
| 68 | # ref_dynamic_per_token_quant, when we have a dynamic_per_tensor int8 quant |
| 69 | # kernel |
| 70 | def ref_dynamic_per_tensor_fp8_quant( |
| 71 | x: torch.Tensor, |
| 72 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 73 | fp8_traits_min, fp8_traits_max = get_fp8_min_max() |
| 74 | fp8_max = as_float32_tensor(fp8_traits_max) |
| 75 | one = as_float32_tensor(1.0) |
| 76 | |
| 77 | # For fp8, in order to match the cuda kernel output, we have to do exactly |
| 78 | # the same operations as in the corresponding fp8 kernel to prevent |
| 79 | # rounding errors. |
| 80 | |
| 81 | x_max = as_float32_tensor(x.abs().max()) |
| 82 | ref_scale = x_max / fp8_max |
| 83 | ref_iscale = one / ref_scale |
| 84 | ref_out = ( |
| 85 | (as_float32_tensor(x) * ref_iscale) |
| 86 | .clamp(fp8_traits_min, fp8_traits_max) |
| 87 | .to(FP8_DTYPE) |
| 88 | ) |
| 89 | return ref_out, ref_scale.view(1) |
| 90 | |
| 91 | |
| 92 | def native_w8a8_block_matmul( |
| 93 | A: torch.Tensor, |
| 94 | B: torch.Tensor, |
| 95 | As: torch.Tensor, |
| 96 | Bs: torch.Tensor, |
| 97 | block_size: list[int], |
| 98 | output_dtype: torch.dtype, |
| 99 | compute_type: torch.dtype = torch.float32, |
| 100 | ) -> torch.Tensor: |
| 101 | """This function performs matrix multiplication with block-wise |
| 102 | quantization using native torch. |
| 103 | It is agnostic to the input data type and can be used for both int8 and |
| 104 | fp8 data types. |
| 105 | |
| 106 | It takes two input tensors `A` and `B` (int8) with scales `As` and |
| 107 | `Bs` (float32). |
| 108 | The output is returned in the specified `output_dtype`. |
| 109 | """ |
| 110 | A = A.to(compute_type) |
| 111 | B = B.to(compute_type) |
| 112 | assert A.shape[-1] == B.shape[-1] |
| 113 | assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2 |
| 114 | assert len(block_size) == 2 |
| 115 | block_n, block_k = block_size[0], block_size[1] |
| 116 | assert (A.shape[-1] + block_k - 1) // block_k == As.shape[-1] |
| 117 | assert A.shape[:-1] == As.shape[:-1] |
| 118 | |
| 119 | M = A.numel() // A.shape[-1] |
| 120 | N, K = B.shape |
| 121 | origin_C_shape = A.shape[:-1] + (N,) |
| 122 | A = A.reshape(M, A.shape[-1]) |
| 123 | As = As.reshape(M, As.shape[-1]) |
| 124 | n_tiles = (N + block_n - 1) // block_n |
| 125 | k_tiles = (K + block_k - 1) // block_k |
| 126 | assert n_tiles == Bs.shape[0], f"{n_tiles} == {Bs.shape[0]}" |
| 127 | assert k_tiles == Bs.shape[1], f"{k_tiles} == {Bs.shape[1]}" |
| 128 | |
| 129 | C_shape = (M, N) |
| 130 | C = torch.zeros(C_shape, dtype=compute_type, device=A.device) |
| 131 | |
| 132 | A_tiles = [A[:, i * block_k : min((i + 1) * block_k, K)] for i in range(k_tiles)] |
| 133 | B_tiles = [ |
| 134 | [ |
| 135 | B[ |
| 136 | j * block_n : min((j + 1) * block_n, N), |
| 137 | i * block_k : min((i + 1) * block_k, K), |
| 138 | ] |
| 139 | for i in range(k_tiles) |
| 140 | ] |
| 141 | for j in range(n_tiles) |
| 142 | ] |
| 143 | C_tiles = [C[:, j * block_n : min((j + 1) * block_n, N)] for j in range(n_tiles)] |
| 144 | As_tiles = [As[:, i : i + 1] for i in range(k_tiles)] |
| 145 | |
| 146 | for i in range(k_tiles): |
| 147 | for j in range(n_tiles): |
| 148 | a = A_tiles[i] |
| 149 | b = B_tiles[j][i] |
| 150 | c = C_tiles[j] |
| 151 | s = As_tiles[i] * Bs[j][i] |
| 152 | c[:, :] += torch.matmul(a, b.t()) * s |
| 153 | |
| 154 | C = C.reshape(origin_C_shape).to(output_dtype) |
| 155 | return C |
| 156 | |
| 157 | |
| 158 | def native_per_token_group_quant_fp8( |
| 159 | x, group_size, eps=1e-10, dtype=torch.float8_e4m3fn |
| 160 | ): |
| 161 | """Function to perform per-token-group quantization on an input tensor |
| 162 | `x` using native torch.""" |
| 163 | assert x.shape[-1] % group_size == 0, ( |
| 164 | "the last dimension of `x` must be divisible by `group_size`" |
| 165 | ) |
| 166 | assert x.is_contiguous(), "`x` is not contiguous" |
| 167 | |
| 168 | finfo = torch.finfo(dtype) |
| 169 | fp8_min = finfo.min |
| 170 | fp8_max = finfo.max |
| 171 | |
| 172 | x_ = x.reshape(x.numel() // group_size, group_size) |
| 173 | amax = x_.abs().max(dim=-1, keepdim=True)[0].clamp(min=eps).to(torch.float32) |
| 174 | x_s = amax / fp8_max |
| 175 | if is_deep_gemm_e8m0_used(): |
| 176 | x_s = _ceil_to_ue8m0(x_s) |
| 177 | x_q = (x_ / x_s).clamp(min=fp8_min, max=fp8_max).to(dtype) |
| 178 | x_q = x_q.reshape(x.shape) |
| 179 | x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size,)) |
| 180 | |
| 181 | return x_q, x_s |
| 182 | |
| 183 | |
| 184 | def native_per_token_group_quant_int8( |
| 185 | x, group_size, eps=1e-10, dtype=torch.int8, *, round_to_nearest=True |
| 186 | ): |
| 187 | """Function to perform per-token-group quantization on an input tensor |
| 188 | `x` using native torch. |
| 189 | |
| 190 | It converts the tensor values into int8 values and returns the |
| 191 | quantized tensor along with the scaling factor used for quantization. |
| 192 | Per-token quantization rounds to nearest; block quantization truncates. |
| 193 | """ |
| 194 | assert x.shape[-1] % group_size == 0, ( |
| 195 | "the last dimension of `x` must be divisible by `group_size`" |
| 196 | ) |
| 197 | assert x.is_contiguous(), "`x` is not contiguous" |
| 198 | |
| 199 | iinfo = torch.iinfo(dtype) |
| 200 | int8_min = iinfo.min |
| 201 | int8_max = iinfo.max |
| 202 | |
| 203 | x_ = x.reshape(x.numel() // group_size, group_size) |
| 204 | # Use float32 for scale calculation for stability |
| 205 | amax = x_.abs().max(dim=-1, keepdim=True)[0].clamp(min=eps).to(torch.float32) |
| 206 | x_s = amax / int8_max |
| 207 | x_q = x_.to(torch.float32) / x_s |
| 208 | if round_to_nearest: |
| 209 | x_q = x_q.round() |
| 210 | x_q = x_q.clamp(min=int8_min, max=int8_max).to(dtype) |
| 211 | x_q = x_q.reshape(x.shape) |
| 212 | x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size,)) |
| 213 | |
| 214 | return x_q, x_s |
| 215 | |
| 216 | |
| 217 | DEFAULT_BLOCK_SHAPE = [128, 128] |
| 218 | |
| 219 | |
| 220 | def per_block_cast_to_int8( |
| 221 | x: torch.Tensor, |
| 222 | block_shape: list[int] = DEFAULT_BLOCK_SHAPE, |
| 223 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 224 | block_m, block_n = block_shape |
| 225 | assert x.dim() == 2 |
| 226 | m, n = x.shape |
| 227 | x_padded = torch.zeros( |
| 228 | (round_up(m, block_m), round_up(n, block_n)), dtype=x.dtype, device=x.device |
| 229 | ) |
| 230 | x_padded[:m, :n] = x |
| 231 | x_view = x_padded.view(-1, block_m, x_padded.size(1) // block_n, block_n) |
| 232 | x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4) |
| 233 | x_scaled = (x_view * (256.0 / x_amax)).to(torch.int8) |
| 234 | x_scaled_sub = x_scaled.view_as(x_padded)[:m, :n].contiguous() |
| 235 | scales = (x_amax / 256.0).view(x_view.size(0), x_view.size(2)) |
| 236 | return x_scaled_sub, scales |
| 237 | |
| 238 | |
| 239 | def dequant( |
| 240 | t: torch.Tensor, |
| 241 | scale: torch.Tensor | None, |
| 242 | block_shape: list[int] | None, |
| 243 | per_act_token_quant: bool, |
| 244 | out_dtype: torch.dtype | None = torch.float32, |
| 245 | ) -> torch.Tensor: |
| 246 | if scale is not None: |
| 247 | f32 = torch.float32 |
| 248 | if per_act_token_quant or block_shape is None: |
| 249 | return (t.to(f32) * scale).to(out_dtype) |
| 250 | else: |
| 251 | return (t.to(f32) * group_broadcast(scale, t.shape)).to(out_dtype) |
| 252 | else: |
| 253 | return t.to(out_dtype) |
| 254 | |
| 255 | |
| 256 | def native_batched_masked_quant_matmul( |
| 257 | A: torch.Tensor, |
| 258 | B: torch.Tensor, |
| 259 | C: torch.Tensor, |
| 260 | num_expert_tokens: torch.Tensor, |
| 261 | A_scale: torch.Tensor | None = None, |
| 262 | B_scale: torch.Tensor | None = None, |
| 263 | block_shape: list[int] | None = None, |
| 264 | per_act_token_quant: bool = False, |
| 265 | ) -> torch.Tensor: |
| 266 | num_expert_tokens_cpu = num_expert_tokens.clone() |
| 267 | num_expert_tokens_cpu = num_expert_tokens_cpu.to(device="cpu") |
| 268 | num_experts = num_expert_tokens.size(0) |
| 269 | |
| 270 | for e in range(num_experts): |
| 271 | num_tokens = num_expert_tokens_cpu[e] |
| 272 | if A.dtype.itemsize == 1 and block_shape is not None: |
| 273 | assert A_scale is not None and B_scale is not None |
| 274 | tmp = native_w8a8_block_matmul( |
| 275 | A[e], B[e], A_scale[e], B_scale[e], block_shape, C.dtype |
| 276 | ) |
| 277 | C[e, :num_tokens, :] = tmp[:num_tokens, :] |
| 278 | elif A.dtype.itemsize == 1 and block_shape is None: |
| 279 | assert A_scale is not None and B_scale is not None |
| 280 | A_dq = dequant(A[e], A_scale[e], block_shape, per_act_token_quant) |
| 281 | B_dq = dequant(B[e], B_scale[e], block_shape, per_act_token_quant) |
| 282 | C[e, :num_tokens, :] = (A_dq[:num_tokens] @ B_dq.transpose(0, 1)).to( |
| 283 | C.dtype |
| 284 | ) |
| 285 | else: |
| 286 | assert A_scale is None |
| 287 | assert B_scale is None |
| 288 | C[e, :num_tokens, :] = A[e, :num_tokens, :] @ B[e].transpose(0, 1) |
| 289 | |
| 290 | return C |