| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # Adapted from https://github.com/sgl-project/sglang/blob/4cb53ecd0cffceb6dee5c011a58f65997a86f151/python/sglang/srt/layers/quantization/int8_kernel.py |
| 5 | import logging |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.platforms import current_platform |
| 10 | from vllm.triton_utils import tl, triton |
| 11 | |
| 12 | logger = logging.getLogger(__name__) |
| 13 | |
| 14 | |
| 15 | def block_dequant( |
| 16 | x_q_block: torch.Tensor, |
| 17 | x_s: torch.Tensor, |
| 18 | block_size: list[int], |
| 19 | ) -> torch.Tensor: |
| 20 | """This function conducts block-wise dequantization. |
| 21 | The inputs are block-wise quantization tensor `x_q_block`, |
| 22 | block-wise quantization scale and the block size. |
| 23 | The outputs are dequantized tensor. |
| 24 | """ |
| 25 | block_n, block_k = block_size[0], block_size[1] |
| 26 | n, k = x_q_block.shape |
| 27 | n_tiles = (n + block_n - 1) // block_n |
| 28 | k_tiles = (k + block_k - 1) // block_k |
| 29 | assert n_tiles == x_s.shape[0] |
| 30 | assert k_tiles == x_s.shape[1] |
| 31 | |
| 32 | x_dq_block = x_q_block.to(torch.float32) |
| 33 | |
| 34 | for i in range(k_tiles): |
| 35 | for j in range(n_tiles): |
| 36 | x_dq_block[ |
| 37 | j * block_n : min((j + 1) * block_n, n), |
| 38 | i * block_k : min((i + 1) * block_k, k), |
| 39 | ] *= x_s[j][i] |
| 40 | |
| 41 | return x_dq_block |
| 42 | |
| 43 | |
| 44 | if current_platform.is_rocm(): |
| 45 | |
| 46 | @triton.jit |
| 47 | def round_int8(x): |
| 48 | return tl.extra.hip.libdevice.round(x).to(tl.int8) |
| 49 | |
| 50 | |
| 51 | elif current_platform.is_xpu(): |
| 52 | |
| 53 | @triton.jit |
| 54 | def round_int8(x): |
| 55 | return tl.extra.intel.libdevice.round(x).to(tl.int8) |
| 56 | |
| 57 | else: |
| 58 | |
| 59 | @triton.jit |
| 60 | def round_int8(x): |
| 61 | return tl.extra.cuda.libdevice.round(x).to(tl.int8) |
| 62 | |
| 63 | |
| 64 | @triton.jit |
| 65 | def _per_token_quant_int8( |
| 66 | x_ptr, |
| 67 | xq_ptr, |
| 68 | scale_ptr, |
| 69 | stride_x, |
| 70 | stride_xq, |
| 71 | N, |
| 72 | BLOCK: tl.constexpr, |
| 73 | ): |
| 74 | # Adapted from https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282 |
| 75 | row_id = tl.program_id(0) |
| 76 | |
| 77 | cols = tl.arange(0, BLOCK) |
| 78 | mask = cols < N |
| 79 | |
| 80 | x = tl.load(x_ptr + row_id * stride_x + cols, mask=mask, other=0.0).to(tl.float32) |
| 81 | absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10) |
| 82 | scale_x = absmax / 127 |
| 83 | x_q = x * (127 / absmax) |
| 84 | x_q = round_int8(x_q) |
| 85 | |
| 86 | tl.store(xq_ptr + row_id * stride_xq + cols, x_q, mask=mask) |
| 87 | tl.store(scale_ptr + row_id, scale_x) |
| 88 | |
| 89 | |
| 90 | def per_token_quant_int8(x): |
| 91 | original_shape = x.shape |
| 92 | if x.dim() > 2: |
| 93 | x = x.view(-1, original_shape[-1]) |
| 94 | M = x.numel() // x.shape[-1] |
| 95 | N = x.shape[-1] |
| 96 | x_q = torch.empty((M, N), device=x.device, dtype=torch.int8) |
| 97 | scales = torch.empty((M, 1), device=x.device, dtype=torch.float32) |
| 98 | BLOCK = triton.next_power_of_2(N) |
| 99 | # heuristics for number of warps |
| 100 | num_warps = min(max(BLOCK // 256, 1), 8) |
| 101 | x = x.contiguous() |
| 102 | _per_token_quant_int8[(M,)]( |
| 103 | x, |
| 104 | x_q, |
| 105 | scales, |
| 106 | stride_x=x.stride(-2), |
| 107 | stride_xq=x_q.stride(-2), |
| 108 | N=N, |
| 109 | BLOCK=BLOCK, |
| 110 | num_warps=num_warps, |
| 111 | num_stages=1, |
| 112 | ) |
| 113 | x_q = x_q.view(*original_shape) |
| 114 | scales = scales.view(*original_shape[:-1], 1) |
| 115 | return x_q, scales |
| 116 | |
| 117 | |
| 118 | @triton.jit |
| 119 | def _per_token_group_quant_int8( |
| 120 | # Pointers to inputs and output |
| 121 | y_ptr, |
| 122 | y_q_ptr, |
| 123 | y_s_ptr, |
| 124 | # Stride of input |
| 125 | y_stride, |
| 126 | # Columns of input |
| 127 | N, |
| 128 | # Avoid to divide zero |
| 129 | eps, |
| 130 | # Information for int8 |
| 131 | int8_min, |
| 132 | int8_max, |
| 133 | # Meta-parameters |
| 134 | BLOCK: tl.constexpr, |
| 135 | ): |
| 136 | """A Triton-accelerated function to perform per-token-group |
| 137 | quantization on a tensor. |
| 138 | |
| 139 | This function converts the tensor values into int8 values. |
| 140 | """ |
| 141 | # Map the program id to the row of X and Y it should compute. |
| 142 | g_id = tl.program_id(0) |
| 143 | y_ptr += g_id * y_stride |
| 144 | y_q_ptr += g_id * y_stride |
| 145 | y_s_ptr += g_id |
| 146 | |
| 147 | cols = tl.arange(0, BLOCK) # N <= BLOCK |
| 148 | mask = cols < N |
| 149 | |
| 150 | y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) |
| 151 | # Quant |
| 152 | _absmax = tl.maximum(tl.max(tl.abs(y)), eps) |
| 153 | y_s = _absmax / int8_max |
| 154 | y_q = tl.clamp(y / y_s, int8_min, int8_max).to(y_q_ptr.dtype.element_ty) |
| 155 | |
| 156 | tl.store(y_q_ptr + cols, y_q, mask=mask) |
| 157 | tl.store(y_s_ptr, y_s) |
| 158 | |
| 159 | |
| 160 | def per_token_group_quant_int8( |
| 161 | x: torch.Tensor, |
| 162 | group_size: int, |
| 163 | eps: float = 1e-10, |
| 164 | dtype: torch.dtype = torch.int8, |
| 165 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 166 | """Function to perform per-token-group quantization on an input tensor `x`. |
| 167 | |
| 168 | It converts the tensor values into signed int8 values and returns the |
| 169 | quantized tensor along with the scaling factor used for quantization. |
| 170 | |
| 171 | Args: |
| 172 | x: The input tensor with ndim >= 2. |
| 173 | group_size: The group size used for quantization. |
| 174 | eps: The minimum to avoid dividing zero. |
| 175 | dtype: The dype of output tensor. Note that only `torch.int8` |
| 176 | is supported for now. |
| 177 | |
| 178 | Returns: |
| 179 | tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the |
| 180 | scaling factor for quantization. |
| 181 | |
| 182 | """ |
| 183 | assert x.shape[-1] % group_size == 0, ( |
| 184 | "the last dimension of `x` cannot be divisible by `group_size`" |
| 185 | ) |
| 186 | assert x.is_contiguous(), "`x` is not contiguous" |
| 187 | |
| 188 | iinfo = torch.iinfo(dtype) |
| 189 | int8_max = iinfo.max |
| 190 | int8_min = iinfo.min |
| 191 | |
| 192 | x_q = torch.empty_like(x, device=x.device, dtype=dtype) |
| 193 | x_s = torch.empty( |
| 194 | x.shape[:-1] + (x.shape[-1] // group_size,), |
| 195 | device=x.device, |
| 196 | dtype=torch.float32, |
| 197 | ) |
| 198 | # Prefer native stable kernel on CUDA/ROCm when available. |
| 199 | if current_platform.is_cuda_alike(): |
| 200 | torch.ops._C.per_token_group_quant_int8( |
| 201 | x, x_q, x_s, group_size, eps, float(int8_min), float(int8_max) |
| 202 | ) |
| 203 | return x_q, x_s |
| 204 | |
| 205 | M = x.numel() // group_size |
| 206 | N = group_size |
| 207 | |
| 208 | BLOCK = triton.next_power_of_2(N) |
| 209 | # heuristics for number of warps |
| 210 | num_warps = min(max(BLOCK // 256, 1), 8) |
| 211 | num_stages = 1 |
| 212 | _per_token_group_quant_int8[(M,)]( |
| 213 | x, |
| 214 | x_q, |
| 215 | x_s, |
| 216 | group_size, |
| 217 | N, |
| 218 | eps, |
| 219 | int8_min=int8_min, |
| 220 | int8_max=int8_max, |
| 221 | BLOCK=BLOCK, |
| 222 | num_warps=num_warps, |
| 223 | num_stages=num_stages, |
| 224 | ) |
| 225 | |
| 226 | return x_q, x_s |