| 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 import _custom_ops as ops |
| 8 | from vllm.model_executor.layers.quantization.compressed_tensors.triton_scaled_mm import ( # noqa: E501 |
| 9 | triton_scaled_mm, |
| 10 | ) |
| 11 | from vllm.model_executor.layers.quantization.utils import replace_parameter |
| 12 | from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( |
| 13 | convert_to_channelwise, |
| 14 | ) |
| 15 | from vllm.platforms import current_platform |
| 16 | from vllm.utils.torch_utils import direct_register_custom_op |
| 17 | |
| 18 | from .BlockScaledMMLinearKernel import ( |
| 19 | Fp8BlockScaledMMLinearKernel, |
| 20 | ) |
| 21 | from .cutlass import CutlassInt8ScaledMMLinearKernel |
| 22 | from .ScaledMMLinearKernel import ( |
| 23 | Int8ScaledMMLinearLayerConfig, |
| 24 | ) |
| 25 | |
| 26 | |
| 27 | class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel): |
| 28 | @classmethod |
| 29 | def is_supported( |
| 30 | cls, compute_capability: int | None = None |
| 31 | ) -> tuple[bool, str | None]: |
| 32 | if current_platform.is_cuda_alike() or current_platform.is_xpu(): |
| 33 | return True, None |
| 34 | return False, "requires ROCm, CUDA or XPU." |
| 35 | |
| 36 | @classmethod |
| 37 | def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 38 | return True, None |
| 39 | |
| 40 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 41 | w_q, _, i_s, _, _ = self._get_layer_params(layer) |
| 42 | w_q_name, w_s_name, i_s_name, i_zp_name, azp_adj_name = self.layer_param_names |
| 43 | |
| 44 | replace_parameter( |
| 45 | layer, |
| 46 | w_q_name, |
| 47 | torch.nn.Parameter(w_q.t().data, requires_grad=False), |
| 48 | ) |
| 49 | |
| 50 | # WEIGHT SCALE |
| 51 | # Triton kernel supports only per-tensor and per-channel. |
| 52 | # If we have a fused module (QKV, MLP) with per tensor scales (thus N |
| 53 | # scales being passed to the kernel), convert to the per-channel case. |
| 54 | is_fused_module = len(layer.logical_widths) > 1 |
| 55 | weight_scale = getattr(layer, w_s_name) |
| 56 | if is_fused_module and not self.config.is_channelwise: |
| 57 | weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths) |
| 58 | replace_parameter( |
| 59 | layer, |
| 60 | w_s_name, |
| 61 | torch.nn.Parameter(weight_scale.data, requires_grad=False), |
| 62 | ) |
| 63 | |
| 64 | # INPUT SCALE |
| 65 | if self.config.is_static_input_scheme: |
| 66 | assert i_s is not None |
| 67 | |
| 68 | if self.config.input_symmetric: |
| 69 | replace_parameter( |
| 70 | layer, |
| 71 | i_s_name, |
| 72 | torch.nn.Parameter(i_s.max(), requires_grad=False), |
| 73 | ) |
| 74 | setattr(layer, i_zp_name, None) |
| 75 | else: |
| 76 | input_zero_point = getattr(layer, i_zp_name) |
| 77 | |
| 78 | # Reconstruct the ranges to find a single scale and azp |
| 79 | int8_traits = torch.iinfo(torch.int8) |
| 80 | azps = input_zero_point.to(dtype=torch.int32) |
| 81 | range_max = (i_s * (int8_traits.max - azps)).max() |
| 82 | range_min = (i_s * (int8_traits.min - azps)).min() |
| 83 | |
| 84 | scale = (range_max - range_min) / (int8_traits.max - int8_traits.min) |
| 85 | replace_parameter( |
| 86 | layer, |
| 87 | i_s_name, |
| 88 | torch.nn.Parameter(scale, requires_grad=False), |
| 89 | ) |
| 90 | |
| 91 | # AZP loaded as int8 but used as int32 |
| 92 | azp = (int8_traits.min - range_min / scale).to(dtype=torch.int32) |
| 93 | replace_parameter( |
| 94 | layer, |
| 95 | i_zp_name, |
| 96 | torch.nn.Parameter(azp, requires_grad=False), |
| 97 | ) |
| 98 | else: |
| 99 | setattr(layer, i_s_name, None) |
| 100 | setattr(layer, i_zp_name, None) |
| 101 | |
| 102 | # azp_adj is the AZP adjustment term, used to account for weights. |
| 103 | # It does not depend on scales or azp, so it is the same for |
| 104 | # static and dynamic quantization. |
| 105 | # See csrc/quantization/w8a8/cutlass/Epilogues.md for the math. |
| 106 | if not self.config.input_symmetric: |
| 107 | weight = getattr(layer, w_q_name) |
| 108 | # weight is already transposed to [K, N], sum over K (dim=0) |
| 109 | azp_adj = weight.sum(dim=0, keepdim=True, dtype=torch.int32) |
| 110 | if self.config.is_static_input_scheme: |
| 111 | # Fold azp into azp_adj for the per-tensor case |
| 112 | azp_adj = getattr(layer, i_zp_name) * azp_adj |
| 113 | setattr( |
| 114 | layer, |
| 115 | azp_adj_name, |
| 116 | torch.nn.Parameter(azp_adj, requires_grad=False), |
| 117 | ) |
| 118 | else: |
| 119 | setattr(layer, azp_adj_name, None) |
| 120 | |
| 121 | def apply_weights( |
| 122 | self, |
| 123 | layer: torch.nn.Module, |
| 124 | x: torch.Tensor, |
| 125 | bias: torch.Tensor | None = None, |
| 126 | ) -> torch.Tensor: |
| 127 | w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer) |
| 128 | |
| 129 | symmetric = azp_adj is None |
| 130 | x_q, x_s, x_zp = ops.scaled_int8_quant( |
| 131 | x.contiguous(), i_s, i_zp, symmetric=symmetric |
| 132 | ) |
| 133 | |
| 134 | out = triton_scaled_mm( |
| 135 | x_q, w_q, scale_a=x_s, scale_b=w_s, out_dtype=x.dtype, bias=bias |
| 136 | ) |
| 137 | |
| 138 | if azp_adj is not None: |
| 139 | # Asymmetric quantization: subtract the zero-point correction. |
| 140 | # D = scale_a * scale_b * (A_q @ B_q - azp * azp_adj) + bias |
| 141 | # triton_scaled_mm already computed scale_a * scale_b * (A_q @ B_q) + bias |
| 142 | # so we subtract scale_a * scale_b * azp * azp_adj |
| 143 | # |
| 144 | # x_s: [M, 1] or scalar, w_s: [N, 1] or scalar, azp_adj: [1, N] |
| 145 | # Reshape w_s from [N, 1] to [1, N] for proper broadcasting. |
| 146 | w_s_row = w_s.view(1, -1) if w_s.dim() > 0 else w_s |
| 147 | static = i_zp is not None |
| 148 | if not static and x_zp is not None: |
| 149 | # Dynamic per-token: azp is per-token, azp_adj is per-channel |
| 150 | # x_zp: [M, 1], azp_adj: [1, N] |
| 151 | out -= x_s * w_s_row * (x_zp * azp_adj).to(x.dtype) |
| 152 | else: |
| 153 | # Static per-tensor: azp already folded into azp_adj |
| 154 | out -= (x_s * w_s_row * azp_adj).to(x.dtype) |
| 155 | |
| 156 | return out |
| 157 | |
| 158 | |
| 159 | class TritonFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): |
| 160 | @classmethod |
| 161 | def is_supported(cls, compute_capability=None): |
| 162 | if not (current_platform.is_cuda_alike() or current_platform.is_xpu()): |
| 163 | return False, "only CUDA-alike and XPU devices are supported." |
| 164 | return True, None |
| 165 | |
| 166 | def apply_block_scaled_mm( |
| 167 | self, |
| 168 | A: torch.Tensor, |
| 169 | B: torch.Tensor, |
| 170 | As: torch.Tensor, |
| 171 | Bs: torch.Tensor, |
| 172 | ) -> torch.Tensor: |
| 173 | return torch.ops.vllm.w8a8_triton_block_scaled_mm_func( |
| 174 | A, |
| 175 | B, |
| 176 | As, |
| 177 | Bs, |
| 178 | list(self.weight_group_shape), |
| 179 | self.config.out_dtype, |
| 180 | ) |
| 181 | |
| 182 | |
| 183 | # TODO we should be able to change the type of block_size to GroupShape |
| 184 | # after we resolve GroupShape compilation issue |
| 185 | # https://github.com/vllm-project/vllm/issues/25270 |
| 186 | def _w8a8_triton_block_scaled_mm_func( |
| 187 | qx: torch.Tensor, |
| 188 | weight: torch.Tensor, |
| 189 | x_scale: torch.Tensor, |
| 190 | weight_scale: torch.Tensor, |
| 191 | block_size: list[int], |
| 192 | output_dtype: torch.dtype, |
| 193 | ) -> torch.Tensor: |
| 194 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 195 | w8a8_triton_block_scaled_mm, |
| 196 | ) |
| 197 | |
| 198 | return w8a8_triton_block_scaled_mm( |
| 199 | qx, weight, x_scale, weight_scale, block_size, output_dtype |
| 200 | ) |
| 201 | |
| 202 | |
| 203 | def _w8a8_triton_block_scaled_mm_fake( |
| 204 | qx: torch.Tensor, |
| 205 | weight: torch.Tensor, |
| 206 | x_scale: torch.Tensor, |
| 207 | weight_scale: torch.Tensor, |
| 208 | block_size: list[int], |
| 209 | output_dtype: torch.dtype, |
| 210 | ) -> torch.Tensor: |
| 211 | return torch.empty( |
| 212 | (qx.size(0), weight.size(0)), dtype=output_dtype, device=qx.device |
| 213 | ) |
| 214 | |
| 215 | |
| 216 | direct_register_custom_op( |
| 217 | "w8a8_triton_block_scaled_mm_func", |
| 218 | _w8a8_triton_block_scaled_mm_func, |
| 219 | fake_impl=_w8a8_triton_block_scaled_mm_fake, |
| 220 | ) |