| 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.logger import init_logger |
| 11 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 12 | get_fp8_min_max, |
| 13 | ) |
| 14 | from vllm.platforms import current_platform |
| 15 | from vllm.utils.import_utils import has_helion |
| 16 | |
| 17 | if not has_helion(): |
| 18 | raise ImportError( |
| 19 | "Helion kernel requires helion to be installed. " |
| 20 | "Install it with: pip install helion" |
| 21 | ) |
| 22 | |
| 23 | import helion |
| 24 | import helion.language as hl |
| 25 | |
| 26 | from vllm.kernels.helion.register import register_kernel |
| 27 | |
| 28 | logger = init_logger(__name__) |
| 29 | |
| 30 | |
| 31 | def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: |
| 32 | # TODO(xiaohongchen1991): it is difficult for kernel author to cover all |
| 33 | # input property combination. Currently, dtypes are fixed. We need |
| 34 | # optimization to bucket/skip some combinations |
| 35 | num_tokens_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] |
| 36 | hidden_size_list = [2048, 4096, 5120] |
| 37 | group_size_list = [128] |
| 38 | in_dtype: torch.dtype = torch.bfloat16 |
| 39 | out_dtype: torch.dtype = current_platform.fp8_dtype() |
| 40 | scale_dtype: torch.dtype = torch.float32 |
| 41 | |
| 42 | use_ue8m0 = False |
| 43 | column_major = False |
| 44 | fp8_min, fp8_max = get_fp8_min_max() |
| 45 | eps = 1e-10 |
| 46 | |
| 47 | inputs = {} |
| 48 | |
| 49 | for hidden_size, group_size, num_tokens in product( |
| 50 | hidden_size_list, group_size_list, num_tokens_list |
| 51 | ): |
| 52 | input = torch.randn(num_tokens, hidden_size, device="cuda", dtype=in_dtype) |
| 53 | output_q = torch.empty(input.shape, device=input.device, dtype=out_dtype) |
| 54 | output_s = torch.empty( |
| 55 | (num_tokens, hidden_size // group_size), |
| 56 | device=input.device, |
| 57 | dtype=scale_dtype, |
| 58 | ) |
| 59 | config_key = CaseKey( |
| 60 | { |
| 61 | "hidden_size": hidden_size, |
| 62 | "group_size": group_size, |
| 63 | "num_tokens": num_tokens, |
| 64 | } |
| 65 | ) |
| 66 | inputs[config_key] = ( |
| 67 | input, |
| 68 | output_q, |
| 69 | output_s, |
| 70 | group_size, |
| 71 | eps, |
| 72 | fp8_min, |
| 73 | fp8_max, |
| 74 | use_ue8m0, |
| 75 | column_major, |
| 76 | False, |
| 77 | ) |
| 78 | |
| 79 | return inputs |
| 80 | |
| 81 | |
| 82 | _pick_cache: dict[tuple[int, int, int], CaseKey | None] = {} |
| 83 | |
| 84 | |
| 85 | def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None: |
| 86 | """Pick the best pre-tuned config for the given input shape. |
| 87 | |
| 88 | Selection strategy: |
| 89 | 1. Find the closest hidden_size among available configs |
| 90 | (exact match preferred). |
| 91 | 2. Find the closest group_size among available configs |
| 92 | (exact match preferred). |
| 93 | 3. Among the num_tokens values tuned for that hidden_size and group_size, pick |
| 94 | the smallest num_tokens >= the input's num_tokens. If the input is |
| 95 | larger than all available num_tokens, fall back to the largest. |
| 96 | """ |
| 97 | if not config_keys: |
| 98 | return None |
| 99 | |
| 100 | input, _, _, group_size, *_ = args |
| 101 | num_tokens, hidden_size = input.shape |
| 102 | |
| 103 | cache_key = (num_tokens, group_size, hidden_size) |
| 104 | cached = _pick_cache.get(cache_key) |
| 105 | if cached is not None: |
| 106 | return cached |
| 107 | |
| 108 | configs: dict[int, dict[int, list[int]]] = {} |
| 109 | for key in config_keys: |
| 110 | if key.is_default(): |
| 111 | continue |
| 112 | configs.setdefault(key["hidden_size"], {}).setdefault( |
| 113 | key["group_size"], [] |
| 114 | ).append(key["num_tokens"]) |
| 115 | |
| 116 | if not configs: |
| 117 | return None |
| 118 | |
| 119 | best_hidden_size = min(configs, key=lambda s: abs(s - hidden_size)) |
| 120 | best_group_size = min(configs[best_hidden_size], key=lambda s: abs(s - group_size)) |
| 121 | available_num_tokens = sorted(configs[best_hidden_size][best_group_size]) |
| 122 | best_num_tokens = next( |
| 123 | (n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1] |
| 124 | ) |
| 125 | |
| 126 | result = CaseKey( |
| 127 | { |
| 128 | "hidden_size": best_hidden_size, |
| 129 | "group_size": best_group_size, |
| 130 | "num_tokens": best_num_tokens, |
| 131 | } |
| 132 | ) |
| 133 | _pick_cache[cache_key] = result |
| 134 | return result |
| 135 | |
| 136 | |
| 137 | def baseline( |
| 138 | input: torch.Tensor, # [num_tokens, hidden_size] |
| 139 | output_q: torch.Tensor, # [num_tokens, hidden_size] |
| 140 | output_s: torch.Tensor, # [num_tokens, groups_per_row] |
| 141 | group_size: int, |
| 142 | eps: float, |
| 143 | fp8_min: float, |
| 144 | fp8_max: float, |
| 145 | scale_ue8m0: bool, |
| 146 | dummy_is_scale_transposed: bool = False, |
| 147 | dummy_is_tma_aligned: bool = False, |
| 148 | ) -> None: |
| 149 | num_tokens, hidden_size = input.shape |
| 150 | groups_per_row = hidden_size // group_size |
| 151 | |
| 152 | x = input.view(num_tokens, groups_per_row, group_size).to(torch.float32) |
| 153 | s = torch.clamp(torch.amax(torch.abs(x), dim=-1), min=eps) / fp8_max |
| 154 | if scale_ue8m0: |
| 155 | s = torch.exp2(torch.ceil(torch.log2(s))) |
| 156 | y = torch.clamp(x / s[:, :, None], fp8_min, fp8_max) |
| 157 | |
| 158 | output_s.copy_(s) |
| 159 | output_q.copy_(y.view(num_tokens, hidden_size).to(output_q.dtype)) |
| 160 | |
| 161 | |
| 162 | @register_kernel( |
| 163 | mutates_args=["output_q", "output_s"], |
| 164 | config_picker=pick_config, |
| 165 | input_generator=generate_inputs, |
| 166 | helion_settings=helion.Settings( |
| 167 | autotune_baseline_fn=baseline, |
| 168 | ), |
| 169 | ) # type: ignore[misc] |
| 170 | def per_token_group_fp8_quant( |
| 171 | input: torch.Tensor, # [num_tokens, hidden_size] |
| 172 | output_q: torch.Tensor, # [num_tokens, hidden_size] |
| 173 | output_s: torch.Tensor, # [num_tokens, groups_per_row] |
| 174 | group_size: int, |
| 175 | eps: float, |
| 176 | fp8_min: float, |
| 177 | fp8_max: float, |
| 178 | scale_ue8m0: bool, |
| 179 | # Unused dummy args |
| 180 | # Kept for consistency with existing kernel interface |
| 181 | dummy_is_scale_transposed: bool = False, |
| 182 | dummy_is_tma_aligned: bool = False, |
| 183 | ) -> None: |
| 184 | # This code assumes batch_dim and num_tokens are flattened |
| 185 | assert input.ndim == 2 |
| 186 | num_tokens, hidden_size = input.shape |
| 187 | hl.specialize(hidden_size) |
| 188 | hl.specialize(group_size) |
| 189 | |
| 190 | groups_per_row = output_s.shape[1] |
| 191 | hl.specialize(groups_per_row) |
| 192 | assert hidden_size % group_size == 0 and hidden_size // group_size == groups_per_row |
| 193 | assert output_s.ndim == 2 and output_s.dtype == torch.float32 |
| 194 | |
| 195 | input = input.view(num_tokens, -1, group_size) |
| 196 | output_q = output_q.view(num_tokens, -1, group_size) |
| 197 | for tile_m, tile_gn, tile_n in hl.tile( |
| 198 | [num_tokens, groups_per_row, group_size], block_size=[1, None, group_size] |
| 199 | ): |
| 200 | x_blk = input[tile_m, tile_gn, tile_n] |
| 201 | y_s_blk = torch.clamp(torch.amax(torch.abs(x_blk), dim=-1), min=eps) |
| 202 | y_s_blk = y_s_blk / fp8_max |
| 203 | |
| 204 | if scale_ue8m0: |
| 205 | y_s_blk = torch.exp2(torch.ceil(torch.log2(y_s_blk))) |
| 206 | |
| 207 | y_q_blk = torch.clamp(x_blk / y_s_blk[:, :, None], fp8_min, fp8_max).to( |
| 208 | output_q.dtype |
| 209 | ) |
| 210 | |
| 211 | output_s[tile_m, tile_gn] = y_s_blk |
| 212 | output_q[tile_m, tile_gn, tile_n] = y_q_blk |