| 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.activation import SiluAndMul |
| 17 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 18 | get_fp8_min_max, |
| 19 | ) |
| 20 | from vllm.platforms import current_platform |
| 21 | from vllm.utils.import_utils import has_helion |
| 22 | |
| 23 | if not has_helion(): |
| 24 | raise ImportError( |
| 25 | "Helion kernel requires helion to be installed. " |
| 26 | "Install it with: pip install helion" |
| 27 | ) |
| 28 | |
| 29 | import helion |
| 30 | import helion.language as hl |
| 31 | |
| 32 | from vllm.kernels.helion.register import register_kernel |
| 33 | |
| 34 | logger = init_logger(__name__) |
| 35 | |
| 36 | |
| 37 | def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: |
| 38 | # TODO(xiaohongchen1991): it is difficult for kernel author to cover all input |
| 39 | # property combination. Currently, dtypes are fixed. We need optimization to |
| 40 | # bucket/skip some combinations |
| 41 | num_tokens_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] |
| 42 | intermediate_size_list = [6144, 12288, 25600] |
| 43 | |
| 44 | in_dtype: torch.dtype = torch.bfloat16 |
| 45 | out_dtype: torch.dtype = current_platform.fp8_dtype() |
| 46 | scale_dtype: torch.dtype = torch.float32 |
| 47 | group_size_list = [128] |
| 48 | inputs = {} |
| 49 | for intermediate_size, group_size, num_tokens in product( |
| 50 | intermediate_size_list, group_size_list, num_tokens_list |
| 51 | ): |
| 52 | input = torch.randn( |
| 53 | num_tokens, 2 * intermediate_size, device="cuda", dtype=in_dtype |
| 54 | ) |
| 55 | result = torch.empty( |
| 56 | num_tokens, intermediate_size, device=input.device, dtype=out_dtype |
| 57 | ) |
| 58 | scale = torch.empty( |
| 59 | (num_tokens, intermediate_size // group_size), |
| 60 | device=input.device, |
| 61 | dtype=scale_dtype, |
| 62 | ) |
| 63 | # scale_ub clamps the per-group amax of the SiLU-and-mul activation. Use |
| 64 | # a non-degenerate upper bound (midway between the mean and max of the |
| 65 | # activation magnitude) so clamping is partially active and the baseline |
| 66 | # comparison is meaningful. torch.mean(input) ~= 0 for the zero-mean |
| 67 | # input would collapse every scale to the floor and saturate the output. |
| 68 | # Mirrors tests/kernels/helion/test_silu_and_mul_per_block_quant.py. |
| 69 | act_abs = SiluAndMul.forward_native(input.to(torch.float32)).abs() |
| 70 | scale_ub = (0.5 * (act_abs.mean() + act_abs.amax())).to(scale_dtype) |
| 71 | |
| 72 | config_key = CaseKey( |
| 73 | { |
| 74 | "intermediate_size": intermediate_size, |
| 75 | "group_size": group_size, |
| 76 | "num_tokens": num_tokens, |
| 77 | } |
| 78 | ) |
| 79 | inputs[config_key] = (result, input, scale, group_size, scale_ub, False) |
| 80 | |
| 81 | return inputs |
| 82 | |
| 83 | |
| 84 | _pick_cache: dict[tuple[int, int, int], CaseKey | None] = {} |
| 85 | |
| 86 | |
| 87 | def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None: |
| 88 | """Pick the best pre-tuned config for the given input shape. |
| 89 | |
| 90 | Selection strategy: |
| 91 | 1. Find the closest intermediate_size among available configs |
| 92 | (exact match preferred). |
| 93 | 2. Find the closest group_size among available configs |
| 94 | (exact match preferred). |
| 95 | 3. Among the num_tokens values tuned for that intermediate_size and group_size, |
| 96 | pick the smallest num_tokens >= the input's num_tokens. If the input is |
| 97 | larger than all available num_tokens, fall back to the largest. |
| 98 | """ |
| 99 | if not config_keys: |
| 100 | return None |
| 101 | |
| 102 | result, _, _, group_size, *_ = args |
| 103 | num_tokens, intermediate_size = result.shape |
| 104 | |
| 105 | cache_key = (num_tokens, group_size, intermediate_size) |
| 106 | cached = _pick_cache.get(cache_key) |
| 107 | if cached is not None: |
| 108 | return cached |
| 109 | |
| 110 | configs: dict[int, dict[int, list[int]]] = {} |
| 111 | for key in config_keys: |
| 112 | if key.is_default(): |
| 113 | continue |
| 114 | configs.setdefault(key["intermediate_size"], {}).setdefault( |
| 115 | key["group_size"], [] |
| 116 | ).append(key["num_tokens"]) |
| 117 | |
| 118 | if not configs: |
| 119 | return None |
| 120 | |
| 121 | best_intermediate_size = min(configs, key=lambda s: abs(s - intermediate_size)) |
| 122 | best_group_size = min( |
| 123 | configs[best_intermediate_size], key=lambda s: abs(s - group_size) |
| 124 | ) |
| 125 | available_num_tokens = sorted(configs[best_intermediate_size][best_group_size]) |
| 126 | best_num_tokens = next( |
| 127 | (n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1] |
| 128 | ) |
| 129 | |
| 130 | result = CaseKey( |
| 131 | { |
| 132 | "intermediate_size": best_intermediate_size, |
| 133 | "group_size": best_group_size, |
| 134 | "num_tokens": best_num_tokens, |
| 135 | } |
| 136 | ) |
| 137 | _pick_cache[cache_key] = result |
| 138 | return result |
| 139 | |
| 140 | |
| 141 | def baseline( |
| 142 | out: torch.Tensor, # [num_tokens, intermediate_size] |
| 143 | input: torch.Tensor, # [num_tokens, 2 * intermediate_size] |
| 144 | scales: torch.Tensor, # [num_tokens, groups_per_row] |
| 145 | group_size: int, |
| 146 | scale_ub: torch.Tensor | None = None, # scalar tensor |
| 147 | is_scale_transposed: bool = False, |
| 148 | ) -> None: |
| 149 | num_tokens, intermediate_size = out.shape |
| 150 | groups_per_row = intermediate_size // group_size |
| 151 | quant_dtype = out.dtype |
| 152 | qtype_min: int | float |
| 153 | qtype_max: int | float |
| 154 | |
| 155 | if quant_dtype == torch.int8: |
| 156 | qtype_min, qtype_max = get_int8_min_max() |
| 157 | min_scaling_factor = get_int8_min_scaling_factor() |
| 158 | else: |
| 159 | qtype_min, qtype_max = get_fp8_min_max() |
| 160 | min_scaling_factor = 1.0 / (qtype_max * 512.0) |
| 161 | |
| 162 | act = SiluAndMul.forward_native(input.to(torch.float32)) |
| 163 | x_grouped = act.view(num_tokens, groups_per_row, group_size) |
| 164 | |
| 165 | s = torch.amax(torch.abs(x_grouped), dim=-1) |
| 166 | if scale_ub is not None: |
| 167 | s = s.clamp(max=scale_ub) |
| 168 | s = (s * (1.0 / qtype_max)).clamp(min=min_scaling_factor) |
| 169 | |
| 170 | y = x_grouped / s[:, :, None] |
| 171 | if quant_dtype == torch.int8: |
| 172 | y = y.round() |
| 173 | |
| 174 | scales.copy_(s) |
| 175 | out.copy_( |
| 176 | y.clamp(qtype_min, qtype_max).view(num_tokens, intermediate_size).to(out.dtype) |
| 177 | ) |
| 178 | |
| 179 | |
| 180 | @register_kernel( |
| 181 | mutates_args=["out", "scales"], |
| 182 | config_picker=pick_config, |
| 183 | input_generator=generate_inputs, |
| 184 | helion_settings=helion.Settings( |
| 185 | autotune_baseline_fn=baseline, |
| 186 | ignore_warnings=[helion.exc.TensorOperationInWrapper], |
| 187 | ), |
| 188 | ) # type: ignore[misc] |
| 189 | def silu_and_mul_per_block_quant( |
| 190 | out: torch.Tensor, # [num_tokens, intermediate_size] |
| 191 | input: torch.Tensor, # [num_tokens, 2 * intermediate_size] |
| 192 | scales: torch.Tensor, # [num_tokens, groups_per_row] |
| 193 | group_size: int, |
| 194 | scale_ub: torch.Tensor | None = None, # scalar tensor |
| 195 | is_scale_transposed: bool = False, # dummy |
| 196 | ) -> None: |
| 197 | # This code assumes batch_dim and num_tokens are flattened |
| 198 | assert input.ndim == 2 |
| 199 | num_tokens, two_intermediate_size = input.shape |
| 200 | hl.specialize(two_intermediate_size) |
| 201 | |
| 202 | assert two_intermediate_size % 2 == 0 |
| 203 | intermediate_size = two_intermediate_size // 2 |
| 204 | |
| 205 | assert out.shape[0] == num_tokens |
| 206 | assert out.shape[1] == intermediate_size |
| 207 | fp8_dtype = get_fp8_dtype() |
| 208 | assert out.dtype in [fp8_dtype, torch.int8] |
| 209 | |
| 210 | if scale_ub is not None: |
| 211 | assert out.dtype == fp8_dtype |
| 212 | assert scale_ub.dtype == torch.float32 |
| 213 | |
| 214 | assert scales.ndim == 2 and scales.dtype == torch.float32 |
| 215 | |
| 216 | assert scales.shape[0] == num_tokens |
| 217 | groups_per_row = scales.shape[1] |
| 218 | hl.specialize(groups_per_row) |
| 219 | assert ( |
| 220 | intermediate_size % group_size == 0 |
| 221 | and intermediate_size // group_size == groups_per_row |
| 222 | ) |
| 223 | |
| 224 | assert group_size in [64, 128] |
| 225 | hl.specialize(group_size) |
| 226 | |
| 227 | assert input.stride()[-1] == 1 |
| 228 | assert out.stride()[-1] == 1 |
| 229 | |
| 230 | quant_dtype = out.dtype |
| 231 | qtype_traits_min: int | float |
| 232 | qtype_traits_max: int | float |
| 233 | if quant_dtype == torch.int8: |
| 234 | qtype_traits_min, qtype_traits_max = get_int8_min_max() |
| 235 | min_scaling_factor = get_int8_min_scaling_factor() |
| 236 | else: |
| 237 | qtype_traits_min, qtype_traits_max = get_fp8_min_max() |
| 238 | min_scaling_factor = 1.0 / (qtype_traits_max * 512.0) |
| 239 | |
| 240 | qtype_max = float(qtype_traits_max) |
| 241 | |
| 242 | input = input.view(num_tokens, -1, group_size) |
| 243 | out = out.view(num_tokens, -1, group_size) |
| 244 | |
| 245 | for tile_m, tile_gn, tile_n in hl.tile( |
| 246 | [num_tokens, groups_per_row, group_size], block_size=[1, None, group_size] |
| 247 | ): |
| 248 | x_a_blk = input[tile_m, tile_gn, tile_n].to(torch.float32) |
| 249 | x_b_blk = hl.load( |
| 250 | input, |
| 251 | [tile_m, tile_gn.index + groups_per_row, tile_n], |
| 252 | extra_mask=(tile_gn.index + groups_per_row < 2 * groups_per_row)[ |
| 253 | None, :, None |
| 254 | ], |
| 255 | ).to(torch.float32) |
| 256 | x_blk = x_a_blk * torch.sigmoid(x_a_blk) * x_b_blk |
| 257 | s_blk = torch.amax(torch.abs(x_blk), dim=-1).to(torch.float32) |
| 258 | |
| 259 | if scale_ub is not None: |
| 260 | scale_ub_s = hl.load(scale_ub, []) |
| 261 | s_blk = s_blk.clamp(max=scale_ub_s) |
| 262 | s_blk = s_blk * (1.0 / qtype_max) |
| 263 | s_blk = s_blk.clamp(min=min_scaling_factor) |
| 264 | |
| 265 | scales[tile_m, tile_gn] = s_blk |
| 266 | if quant_dtype == torch.int8: |
| 267 | y_blk = (x_blk * (1.0 / s_blk[:, :, None])).round() |
| 268 | else: |
| 269 | y_blk = x_blk / s_blk[:, :, None] |
| 270 | |
| 271 | out[tile_m, tile_gn, tile_n] = y_blk.clamp( |
| 272 | qtype_traits_min, qtype_traits_max |
| 273 | ).to(out.dtype) |