| 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.register import register_kernel |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 13 | get_fp8_min_max, |
| 14 | ) |
| 15 | from vllm.platforms import current_platform |
| 16 | from vllm.utils.import_utils import has_helion |
| 17 | |
| 18 | if not has_helion(): |
| 19 | raise ImportError( |
| 20 | "Helion kernel requires helion to be installed. " |
| 21 | "Install it with: pip install helion" |
| 22 | ) |
| 23 | |
| 24 | import helion |
| 25 | import helion.language as hl |
| 26 | |
| 27 | logger = init_logger(__name__) |
| 28 | |
| 29 | |
| 30 | def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: |
| 31 | # TODO(xiaohongchen1991): it is difficult for kernel author to cover all input |
| 32 | # property combination. Currently, dtypes are fixed. We need optimization to |
| 33 | # bucket/skip some combinations |
| 34 | num_tokens_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] |
| 35 | hidden_size_list = [2048, 4096, 5120] |
| 36 | in_dtype: torch.dtype = torch.bfloat16 |
| 37 | out_dtype: torch.dtype = current_platform.fp8_dtype() |
| 38 | scale_dtype: torch.dtype = torch.float32 |
| 39 | inputs = {} |
| 40 | for num_tokens, hidden_size in product(num_tokens_list, hidden_size_list): |
| 41 | input = torch.randn(num_tokens, hidden_size, device="cuda", dtype=in_dtype) |
| 42 | result = torch.empty(input.shape, device=input.device, dtype=out_dtype) |
| 43 | scale = torch.empty((num_tokens, 1), device=input.device, dtype=scale_dtype) |
| 44 | # scale_ub clamps the per-token amax of |input|. Use a non-degenerate |
| 45 | # upper bound (midway between the mean and max of |input|) so clamping is |
| 46 | # partially active and the baseline comparison is meaningful. |
| 47 | # torch.mean(input) ~= 0 for the zero-mean input would collapse every |
| 48 | # scale to the floor and saturate the output. |
| 49 | input_abs = input.to(torch.float32).abs() |
| 50 | scale_ub = (0.5 * (input_abs.mean() + input_abs.amax())).to(scale_dtype) |
| 51 | |
| 52 | config_key = CaseKey({"hidden_size": hidden_size, "num_tokens": num_tokens}) |
| 53 | inputs[config_key] = (result, input, scale, scale_ub) |
| 54 | |
| 55 | return inputs |
| 56 | |
| 57 | |
| 58 | _pick_cache: dict[tuple[int, int], CaseKey | None] = {} |
| 59 | |
| 60 | |
| 61 | def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None: |
| 62 | """Pick the best pre-tuned config for the given input shape. |
| 63 | |
| 64 | Selection strategy: |
| 65 | 1. Find the closest hidden_size among available configs |
| 66 | (exact match preferred). |
| 67 | 2. Among the num_tokens values tuned for that hidden_size, pick |
| 68 | the smallest num_tokens >= the input's num_tokens. If the input is |
| 69 | larger than all available num_tokens, fall back to the largest. |
| 70 | """ |
| 71 | if not config_keys: |
| 72 | return None |
| 73 | |
| 74 | _, input, *_ = args |
| 75 | num_tokens, hidden_size = input.shape |
| 76 | |
| 77 | cache_key = (num_tokens, hidden_size) |
| 78 | cached = _pick_cache.get(cache_key) |
| 79 | if cached is not None: |
| 80 | return cached |
| 81 | |
| 82 | configs: dict[int, list[int]] = {} |
| 83 | for key in config_keys: |
| 84 | if key.is_default(): |
| 85 | continue |
| 86 | configs.setdefault(key["hidden_size"], []).append(key["num_tokens"]) |
| 87 | |
| 88 | if not configs: |
| 89 | return None |
| 90 | |
| 91 | best_hidden_size = min(configs, key=lambda s: abs(s - hidden_size)) |
| 92 | available_num_tokens = sorted(configs[best_hidden_size]) |
| 93 | best_num_tokens = next( |
| 94 | (n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1] |
| 95 | ) |
| 96 | |
| 97 | result = CaseKey({"hidden_size": best_hidden_size, "num_tokens": best_num_tokens}) |
| 98 | _pick_cache[cache_key] = result |
| 99 | return result |
| 100 | |
| 101 | |
| 102 | def baseline( |
| 103 | result: torch.Tensor, # [num_tokens, hidden_size] |
| 104 | input: torch.Tensor, # [num_tokens, hidden_size] |
| 105 | scale: torch.Tensor, # [num_tokens, 1] |
| 106 | scale_ub: torch.Tensor | None = None, # scalar tensor |
| 107 | ) -> None: |
| 108 | fp8_min, fp8_max = get_fp8_min_max() |
| 109 | min_scaling_factor = 1.0 / (fp8_max * 512.0) |
| 110 | |
| 111 | x = input.to(torch.float32) |
| 112 | s = torch.amax(torch.abs(x), dim=-1, keepdim=True) |
| 113 | if scale_ub is not None: |
| 114 | s = s.clamp(max=scale_ub) |
| 115 | s = (s * (1.0 / fp8_max)).clamp(min=min_scaling_factor) |
| 116 | y = (x / s).clamp(fp8_min, fp8_max) |
| 117 | |
| 118 | scale.copy_(s) |
| 119 | result.copy_(y.to(result.dtype)) |
| 120 | |
| 121 | |
| 122 | # Overwrite autotune_baseline_atol and autotune_baseline_rtol |
| 123 | # if too many configs failed due to baseline check during autotuning |
| 124 | @register_kernel( |
| 125 | mutates_args=["result", "scale"], |
| 126 | config_picker=pick_config, |
| 127 | input_generator=generate_inputs, |
| 128 | helion_settings=helion.Settings( |
| 129 | autotune_baseline_fn=baseline, |
| 130 | ignore_warnings=[helion.exc.TensorOperationInWrapper], |
| 131 | ), |
| 132 | ) |
| 133 | def dynamic_per_token_scaled_fp8_quant( |
| 134 | result: torch.Tensor, # [num_tokens, hidden_size] |
| 135 | input: torch.Tensor, # [num_tokens, hidden_size] |
| 136 | scale: torch.Tensor, # [num_tokens, 1] |
| 137 | scale_ub: torch.Tensor | None = None, # scalar tensor |
| 138 | ) -> None: |
| 139 | # This code assumes batch_dim and num_tokens are flattened |
| 140 | assert input.ndim == 2 |
| 141 | num_tokens, hidden_size = input.shape |
| 142 | hl.specialize(hidden_size) |
| 143 | |
| 144 | assert result.shape == input.shape |
| 145 | assert scale.shape[0] == num_tokens |
| 146 | assert scale.dtype == torch.float32 |
| 147 | assert input.stride()[-1] == 1 |
| 148 | assert result.stride()[-1] == 1 |
| 149 | |
| 150 | fp8_min, fp8_max = get_fp8_min_max() |
| 151 | min_scaling_factor = 1.0 / (fp8_max * 512.0) |
| 152 | |
| 153 | for tile_m in hl.tile(num_tokens, block_size=1): |
| 154 | s_blk = hl.zeros([tile_m], dtype=torch.float32) |
| 155 | for tile_n in hl.tile(hidden_size): |
| 156 | x_blk = input[tile_m, tile_n].to(dtype=torch.float32) |
| 157 | tmp_blk = torch.amax(torch.abs(x_blk), dim=-1) |
| 158 | s_blk = torch.maximum(s_blk, tmp_blk) |
| 159 | |
| 160 | if scale_ub is not None: |
| 161 | scale_ub_s = hl.load(scale_ub, []) |
| 162 | s_blk = s_blk.clamp(max=scale_ub_s) |
| 163 | s_blk = s_blk * (1.0 / fp8_max) |
| 164 | s_blk = s_blk.clamp(min=min_scaling_factor) |
| 165 | scale[tile_m, 0] = s_blk |
| 166 | |
| 167 | for tile_n in hl.tile(hidden_size): |
| 168 | x_blk = input[tile_m, tile_n].to(torch.float32) |
| 169 | y_blk = x_blk * (1.0 / s_blk[:, None]) |
| 170 | |
| 171 | result[tile_m, tile_n] = y_blk.clamp(fp8_min, fp8_max).to(result.dtype) |