| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from itertools import product |
| 5 | from typing import Any |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.kernels.helion.case_key import CaseKey |
| 10 | from vllm.kernels.helion.utils import ( |
| 11 | get_fp8_dtype, |
| 12 | get_int8_min_max, |
| 13 | get_int8_min_scaling_factor, |
| 14 | ) |
| 15 | from vllm.logger import init_logger |
| 16 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 17 | get_fp8_min_max, |
| 18 | ) |
| 19 | from vllm.platforms import current_platform |
| 20 | from vllm.utils.import_utils import has_helion |
| 21 | |
| 22 | if not has_helion(): |
| 23 | raise ImportError( |
| 24 | "Helion kernel requires helion to be installed. " |
| 25 | "Install it with: pip install helion" |
| 26 | ) |
| 27 | |
| 28 | import helion |
| 29 | import helion.language as hl |
| 30 | |
| 31 | from vllm.kernels.helion.register import register_kernel |
| 32 | |
| 33 | logger = init_logger(__name__) |
| 34 | |
| 35 | |
| 36 | def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: |
| 37 | # TODO(xiaohongchen1991): it is difficult for kernel author to cover all |
| 38 | # input property combination. Currently, dtypes are fixed. We need |
| 39 | # optimization to bucket/skip some combinations |
| 40 | num_tokens_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] |
| 41 | hidden_size_list = [2048, 4096, 5120] |
| 42 | group_size_list = [128] |
| 43 | in_dtype: torch.dtype = torch.bfloat16 |
| 44 | out_dtype: torch.dtype = current_platform.fp8_dtype() |
| 45 | scale_dtype: torch.dtype = torch.float32 |
| 46 | inputs = {} |
| 47 | |
| 48 | for hidden_size, group_size, num_tokens in product( |
| 49 | hidden_size_list, group_size_list, num_tokens_list |
| 50 | ): |
| 51 | input = torch.randn(num_tokens, hidden_size, device="cuda", dtype=in_dtype) |
| 52 | result = torch.empty(input.shape, device=input.device, dtype=out_dtype) |
| 53 | scale = torch.empty( |
| 54 | (num_tokens, hidden_size // group_size), |
| 55 | device=input.device, |
| 56 | dtype=scale_dtype, |
| 57 | ) |
| 58 | residual = torch.randn_like(input) |
| 59 | weight = torch.normal( |
| 60 | mean=1.0, |
| 61 | std=1.0, |
| 62 | size=(hidden_size,), |
| 63 | dtype=input.dtype, |
| 64 | device=input.device, |
| 65 | ) |
| 66 | epsilon = 1e-6 |
| 67 | # scale_ub clamps the per-group amax of the RMS-normed, weighted output. |
| 68 | # Use a non-degenerate upper bound (midway between the mean and max of |
| 69 | # that magnitude) so clamping is partially active and the baseline |
| 70 | # comparison is meaningful. torch.mean(input) ~= 0 for the zero-mean |
| 71 | # input would collapse every scale to the floor and saturate the output. |
| 72 | # Mirrors the reference normalization in baseline() below. |
| 73 | x = input.to(torch.float32) + residual.to(torch.float32) |
| 74 | rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + epsilon) |
| 75 | x_norm_abs = ((x * rms).to(input.dtype) * weight).abs().to(torch.float32) |
| 76 | scale_ub = (0.5 * (x_norm_abs.mean() + x_norm_abs.amax())).to(scale_dtype) |
| 77 | |
| 78 | config_key = CaseKey( |
| 79 | { |
| 80 | "hidden_size": hidden_size, |
| 81 | "group_size": group_size, |
| 82 | "num_tokens": num_tokens, |
| 83 | } |
| 84 | ) |
| 85 | inputs[config_key] = ( |
| 86 | result, |
| 87 | input, |
| 88 | weight, |
| 89 | scale, |
| 90 | epsilon, |
| 91 | scale_ub, |
| 92 | residual, |
| 93 | group_size, |
| 94 | False, |
| 95 | ) |
| 96 | |
| 97 | return inputs |
| 98 | |
| 99 | |
| 100 | _pick_cache: dict[tuple[int, int, int], CaseKey | None] = {} |
| 101 | |
| 102 | |
| 103 | def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None: |
| 104 | """Pick the best pre-tuned config for the given input shape. |
| 105 | |
| 106 | Selection strategy: |
| 107 | 1. Find the closest hidden_size among available configs |
| 108 | (exact match preferred). |
| 109 | 2. Find the closest group_size among available configs |
| 110 | (exact match preferred). |
| 111 | 2. Among the num_tokens values tuned for that hidden_size and group_size, pick |
| 112 | the smallest num_tokens >= the input's num_tokens. If the input is |
| 113 | larger than all available num_tokens, fall back to the largest. |
| 114 | """ |
| 115 | if not config_keys: |
| 116 | return None |
| 117 | |
| 118 | _, input, _, _, _, _, _, group_size, *_ = args |
| 119 | num_tokens, hidden_size = input.shape |
| 120 | |
| 121 | cache_key = (num_tokens, group_size, hidden_size) |
| 122 | cached = _pick_cache.get(cache_key) |
| 123 | if cached is not None: |
| 124 | return cached |
| 125 | |
| 126 | configs: dict[int, dict[int, list[int]]] = {} |
| 127 | for key in config_keys: |
| 128 | if key.is_default(): |
| 129 | continue |
| 130 | configs.setdefault(key["hidden_size"], {}).setdefault( |
| 131 | key["group_size"], [] |
| 132 | ).append(key["num_tokens"]) |
| 133 | |
| 134 | if not configs: |
| 135 | return None |
| 136 | |
| 137 | best_hidden_size = min(configs, key=lambda s: abs(s - hidden_size)) |
| 138 | best_group_size = min(configs[best_hidden_size], key=lambda s: abs(s - group_size)) |
| 139 | available_num_tokens = sorted(configs[best_hidden_size][best_group_size]) |
| 140 | best_num_tokens = next( |
| 141 | (n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1] |
| 142 | ) |
| 143 | |
| 144 | result = CaseKey( |
| 145 | { |
| 146 | "hidden_size": best_hidden_size, |
| 147 | "group_size": best_group_size, |
| 148 | "num_tokens": best_num_tokens, |
| 149 | } |
| 150 | ) |
| 151 | _pick_cache[cache_key] = result |
| 152 | return result |
| 153 | |
| 154 | |
| 155 | def baseline( |
| 156 | result: torch.Tensor, # [num_tokens, hidden_size] |
| 157 | input: torch.Tensor, # [num_tokens, hidden_size] |
| 158 | weight: torch.Tensor, # [hidden_size] |
| 159 | scale: torch.Tensor, # [num_tokens, groups_per_row] |
| 160 | epsilon: float, |
| 161 | scale_ub: torch.Tensor | None, # [] |
| 162 | residual: torch.Tensor | None, # [num_tokens, hidden_size] |
| 163 | group_size: int, |
| 164 | is_scale_transposed: bool, |
| 165 | ) -> None: |
| 166 | num_tokens, hidden_size = input.shape |
| 167 | groups_per_row = hidden_size // group_size |
| 168 | quant_dtype = result.dtype |
| 169 | qtype_min: int | float |
| 170 | qtype_max: int | float |
| 171 | |
| 172 | if quant_dtype == torch.int8: |
| 173 | qtype_min, qtype_max = get_int8_min_max() |
| 174 | min_scaling_factor = get_int8_min_scaling_factor() |
| 175 | else: |
| 176 | qtype_min, qtype_max = get_fp8_min_max() |
| 177 | min_scaling_factor = 1.0 / (qtype_max * 512.0) |
| 178 | |
| 179 | x = input.to(torch.float32) |
| 180 | if residual is not None: |
| 181 | x = x + residual |
| 182 | residual.copy_(x.to(residual.dtype)) |
| 183 | |
| 184 | rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + epsilon) |
| 185 | x_norm = (x * rms).to(input.dtype) * weight |
| 186 | x_grouped = x_norm.view(num_tokens, groups_per_row, group_size).to(torch.float32) |
| 187 | |
| 188 | s = torch.amax(torch.abs(x_grouped), dim=-1).to(torch.float32) |
| 189 | if scale_ub is not None: |
| 190 | s = s.clamp(max=scale_ub) |
| 191 | s = (s * (1.0 / qtype_max)).clamp(min=min_scaling_factor) |
| 192 | |
| 193 | y = x_grouped / s[:, :, None] |
| 194 | if quant_dtype == torch.int8: |
| 195 | y = y.round() |
| 196 | |
| 197 | scale.copy_(s) |
| 198 | result.copy_( |
| 199 | y.clamp(qtype_min, qtype_max).view(num_tokens, hidden_size).to(result.dtype) |
| 200 | ) |
| 201 | |
| 202 | |
| 203 | @register_kernel( |
| 204 | mutates_args=["result", "scale", "residual"], |
| 205 | config_picker=pick_config, |
| 206 | input_generator=generate_inputs, |
| 207 | helion_settings=helion.Settings( |
| 208 | autotune_baseline_fn=baseline, |
| 209 | ignore_warnings=[helion.exc.TensorOperationInWrapper], |
| 210 | ), |
| 211 | ) # type: ignore[misc] |
| 212 | def rms_norm_per_block_quant( |
| 213 | result: torch.Tensor, # [num_tokens, hidden_size] |
| 214 | input: torch.Tensor, # [num_tokens, hidden_size] |
| 215 | weight: torch.Tensor, # [hidden_size] |
| 216 | scale: torch.Tensor, # [num_tokens, groups_per_row] |
| 217 | epsilon: float, |
| 218 | scale_ub: torch.Tensor | None, # [] |
| 219 | residual: torch.Tensor | None, # [num_tokens, hidden_size] |
| 220 | group_size: int, |
| 221 | is_scale_transposed: bool, # dummy |
| 222 | ) -> None: |
| 223 | # This code assumes batch_dim and num_tokens are flattened |
| 224 | assert input.ndim == 2 |
| 225 | num_tokens, hidden_size = input.shape |
| 226 | hl.specialize(hidden_size) |
| 227 | hl.specialize(group_size) |
| 228 | |
| 229 | groups_per_row = scale.shape[1] |
| 230 | hl.specialize(groups_per_row) |
| 231 | assert hidden_size % group_size == 0 and hidden_size // group_size == groups_per_row |
| 232 | assert scale.shape[0] == num_tokens |
| 233 | assert scale.dtype == torch.float32 |
| 234 | if scale.stride(1) > 1: |
| 235 | assert is_scale_transposed |
| 236 | |
| 237 | fp8_dtype = get_fp8_dtype() |
| 238 | assert result.dtype in [fp8_dtype, torch.int8] |
| 239 | assert result.is_contiguous() and input.is_contiguous() |
| 240 | |
| 241 | if scale_ub is not None: |
| 242 | assert result.dtype == fp8_dtype |
| 243 | assert scale_ub.dtype == torch.float32 |
| 244 | |
| 245 | assert input.dtype == weight.dtype |
| 246 | |
| 247 | if residual is not None: |
| 248 | assert residual.dtype == input.dtype |
| 249 | |
| 250 | assert group_size in [64, 128] |
| 251 | |
| 252 | quant_dtype = result.dtype |
| 253 | qtype_traits_min: int | float |
| 254 | qtype_traits_max: int | float |
| 255 | if quant_dtype == torch.int8: |
| 256 | qtype_traits_min, qtype_traits_max = get_int8_min_max() |
| 257 | min_scaling_factor = get_int8_min_scaling_factor() |
| 258 | else: |
| 259 | qtype_traits_min, qtype_traits_max = get_fp8_min_max() |
| 260 | min_scaling_factor = 1.0 / (qtype_traits_max * 512.0) |
| 261 | |
| 262 | qtype_max = float(qtype_traits_max) |
| 263 | |
| 264 | for tile_m in hl.tile(num_tokens, block_size=1): |
| 265 | rms = hl.zeros([tile_m], dtype=torch.float32) |
| 266 | for tile_n in hl.tile(hidden_size): |
| 267 | x_blk = input[tile_m, tile_n].to(torch.float32) |
| 268 | if residual is not None: |
| 269 | x_blk = x_blk + residual[tile_m, tile_n] |
| 270 | rms = rms + x_blk.pow(2).sum(dim=-1) |
| 271 | |
| 272 | rms = torch.rsqrt(rms * (1.0 / hidden_size) + epsilon) |
| 273 | |
| 274 | m_idx = tile_m.begin + hl.arange(tile_m.block_size) |
| 275 | m_blk = m_idx[:, None, None] |
| 276 | for tile_gn, tile_n in hl.tile( |
| 277 | [groups_per_row, group_size], block_size=[None, group_size] |
| 278 | ): |
| 279 | gn_idx = tile_gn.index |
| 280 | n_offset = tile_n.index |
| 281 | n_idx = gn_idx[:, None] * group_size + n_offset[None, :] |
| 282 | n_blk = n_idx[None, :, :] |
| 283 | mask = (gn_idx < groups_per_row)[None, :, None] |
| 284 | |
| 285 | x_blk = hl.load(input, [m_blk, n_blk], extra_mask=mask).to( |
| 286 | dtype=torch.float32 |
| 287 | ) |
| 288 | if residual is not None: |
| 289 | r_blk = hl.load(residual, [m_blk, n_blk], extra_mask=mask) |
| 290 | x_blk = x_blk + r_blk |
| 291 | |
| 292 | w_blk = hl.load(weight, [n_blk], extra_mask=mask) |
| 293 | x_norm_blk = (x_blk * rms[:, None, None]).to(input.dtype) * w_blk |
| 294 | s_blk = torch.amax(torch.abs(x_norm_blk), dim=-1).to(torch.float32) |
| 295 | |
| 296 | if scale_ub is not None: |
| 297 | scale_ub_s = hl.load(scale_ub, []) |
| 298 | s_blk = s_blk.clamp(max=scale_ub_s) |
| 299 | |
| 300 | s_blk = s_blk * (1.0 / qtype_max) |
| 301 | s_blk = s_blk.clamp(min=min_scaling_factor) |
| 302 | |
| 303 | scale[tile_m, tile_gn] = s_blk |
| 304 | |
| 305 | if quant_dtype == torch.int8: |
| 306 | y_blk = (x_norm_blk * (1.0 / s_blk[:, :, None])).round() |
| 307 | else: |
| 308 | y_blk = x_norm_blk / s_blk[:, :, None] |
| 309 | |
| 310 | y_blk = y_blk.clamp(qtype_traits_min, qtype_traits_max).to(result.dtype) |
| 311 | hl.store(result, [m_blk, n_blk], y_blk, extra_mask=mask) |
| 312 | |
| 313 | if residual is not None: |
| 314 | hl.store( |
| 315 | residual, [m_blk, n_blk], x_blk.to(residual.dtype), extra_mask=mask |
| 316 | ) |