| 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.utils.import_utils import has_helion |
| 12 | |
| 13 | if not has_helion(): |
| 14 | raise ImportError( |
| 15 | "Helion kernel requires helion to be installed. " |
| 16 | "Install it with: pip install helion" |
| 17 | ) |
| 18 | |
| 19 | import helion |
| 20 | import helion.language as hl |
| 21 | |
| 22 | from vllm import ir |
| 23 | from vllm.kernels.helion.register import register_kernel |
| 24 | from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding |
| 25 | |
| 26 | logger = init_logger(__name__) |
| 27 | |
| 28 | |
| 29 | def _compute_cos_sin_cache( |
| 30 | max_position_embeddings, rotary_dim, device="cuda", dtype=torch.float |
| 31 | ): |
| 32 | inv_freq = 1.0 / ( |
| 33 | 10000 |
| 34 | ** (torch.arange(0, rotary_dim, 2, device=device, dtype=dtype) / rotary_dim) |
| 35 | ) |
| 36 | |
| 37 | t = torch.arange(max_position_embeddings, device=device, dtype=dtype) |
| 38 | |
| 39 | freqs = torch.einsum("i,j -> ij", t, inv_freq) |
| 40 | cos = freqs.cos() |
| 41 | sin = freqs.sin() |
| 42 | cache = torch.cat((cos, sin), dim=-1) |
| 43 | return cache |
| 44 | |
| 45 | |
| 46 | def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: |
| 47 | # TODO(xiaohongchen1991): it is difficult for kernel author to cover |
| 48 | # all input property combination. Currently, dtypes are fixed. We need |
| 49 | # optimization to bucket/skip some combinations |
| 50 | num_tokens_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] |
| 51 | num_heads_pair = [ |
| 52 | (16, 8), |
| 53 | (32, 8), |
| 54 | (64, 8), |
| 55 | ] |
| 56 | head_dim = 128 |
| 57 | in_dtype: torch.dtype = torch.bfloat16 |
| 58 | rotary_ratio = 1.0 |
| 59 | is_neox = True |
| 60 | eps = 1e-6 |
| 61 | device = "cuda" |
| 62 | inputs = {} |
| 63 | |
| 64 | for num_tokens, (num_q_heads, num_kv_heads) in product( |
| 65 | num_tokens_list, num_heads_pair |
| 66 | ): |
| 67 | total_dim = (num_q_heads + 2 * num_kv_heads) * head_dim |
| 68 | qkv = torch.empty( |
| 69 | num_tokens, total_dim, dtype=in_dtype, device=device |
| 70 | ).uniform_(-0.1, 0.1) |
| 71 | positions = torch.arange(num_tokens, dtype=torch.long, device=device) |
| 72 | q_weight = torch.empty(head_dim, dtype=in_dtype, device=device).uniform_( |
| 73 | 0.8, 1.2 |
| 74 | ) |
| 75 | k_weight = torch.empty(head_dim, dtype=in_dtype, device=device).uniform_( |
| 76 | 0.8, 1.2 |
| 77 | ) |
| 78 | rotary_dim = int(head_dim * rotary_ratio) |
| 79 | cos_sin_cache = _compute_cos_sin_cache(40960, rotary_dim) |
| 80 | cos_sin_cache = cos_sin_cache.to(in_dtype) |
| 81 | |
| 82 | config_key = CaseKey( |
| 83 | { |
| 84 | "q_heads": num_q_heads, |
| 85 | "kv_heads": num_kv_heads, |
| 86 | "num_tokens": num_tokens, |
| 87 | } |
| 88 | ) |
| 89 | inputs[config_key] = ( |
| 90 | qkv, |
| 91 | num_q_heads, |
| 92 | num_kv_heads, |
| 93 | num_kv_heads, |
| 94 | head_dim, |
| 95 | eps, |
| 96 | q_weight, |
| 97 | k_weight, |
| 98 | cos_sin_cache, |
| 99 | is_neox, |
| 100 | positions.view(-1), |
| 101 | ) |
| 102 | |
| 103 | return inputs |
| 104 | |
| 105 | |
| 106 | _pick_cache: dict[tuple[int, int, int], CaseKey | None] = {} |
| 107 | |
| 108 | |
| 109 | def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None: |
| 110 | """Pick the best pre-tuned config for the given input shape. |
| 111 | |
| 112 | Selection strategy: |
| 113 | 1. Find the closest q_heads among available configs |
| 114 | (exact match preferred). |
| 115 | 2. Find the closest kv_heads among available configs |
| 116 | (exact match preferred). |
| 117 | 3. Among the num_tokens values tuned for that q_heads and q_heads, pick |
| 118 | the smallest num_tokens >= the input's num_tokens. If the input is |
| 119 | larger than all available num_tokens, fall back to the largest. |
| 120 | """ |
| 121 | if not config_keys: |
| 122 | return None |
| 123 | |
| 124 | qkv, q_heads, kv_heads, *_ = args |
| 125 | num_tokens = qkv.shape[0] |
| 126 | |
| 127 | cache_key = (num_tokens, q_heads, kv_heads) |
| 128 | cached = _pick_cache.get(cache_key) |
| 129 | if cached is not None: |
| 130 | return cached |
| 131 | |
| 132 | configs: dict[int, dict[int, list[int]]] = {} |
| 133 | for key in config_keys: |
| 134 | if key.is_default(): |
| 135 | continue |
| 136 | configs.setdefault(key["q_heads"], {}).setdefault(key["kv_heads"], []).append( |
| 137 | key["num_tokens"] |
| 138 | ) |
| 139 | |
| 140 | if not configs: |
| 141 | return None |
| 142 | |
| 143 | best_q_heads = min(configs, key=lambda s: abs(s - q_heads)) |
| 144 | best_kv_heads = min(configs[best_q_heads], key=lambda s: abs(s - kv_heads)) |
| 145 | available_num_tokens = sorted(configs[best_q_heads][best_kv_heads]) |
| 146 | best_num_tokens = next( |
| 147 | (n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1] |
| 148 | ) |
| 149 | |
| 150 | result = CaseKey( |
| 151 | { |
| 152 | "q_heads": best_q_heads, |
| 153 | "kv_heads": best_kv_heads, |
| 154 | "num_tokens": best_num_tokens, |
| 155 | } |
| 156 | ) |
| 157 | _pick_cache[cache_key] = result |
| 158 | return result |
| 159 | |
| 160 | |
| 161 | def baseline( |
| 162 | qkv: torch.Tensor, # [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim] |
| 163 | num_heads_q: int, |
| 164 | num_heads_k: int, |
| 165 | num_heads_v: int, |
| 166 | head_dim: int, |
| 167 | eps: float, |
| 168 | q_weight: torch.Tensor, |
| 169 | k_weight: torch.Tensor, |
| 170 | cos_sin_cache: torch.Tensor, # [max_position, rotary_dim] |
| 171 | is_neox: bool, |
| 172 | position_ids: torch.Tensor, # [num_tokens], |
| 173 | forced_token_heads_per_warp: int = -1, # dummy |
| 174 | ) -> None: |
| 175 | q_size = num_heads_q * head_dim |
| 176 | kv_size = num_heads_k * head_dim |
| 177 | |
| 178 | q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1) |
| 179 | |
| 180 | q_by_head = q.view(*q.shape[:-1], q.shape[-1] // head_dim, head_dim) |
| 181 | q_by_head = ir.ops.rms_norm(q_by_head, q_weight, eps) |
| 182 | q = q_by_head.view(q.shape) |
| 183 | |
| 184 | k_by_head = k.view(*k.shape[:-1], k.shape[-1] // head_dim, head_dim) |
| 185 | k_by_head = ir.ops.rms_norm(k_by_head, k_weight, eps) |
| 186 | k = k_by_head.view(k.shape) |
| 187 | |
| 188 | q, k = RotaryEmbedding.forward_static( |
| 189 | position_ids, q, k, head_dim, cos_sin_cache.shape[1], cos_sin_cache, is_neox |
| 190 | ) |
| 191 | qkv[:, :q_size].copy_(q) |
| 192 | qkv[:, q_size : q_size + kv_size].copy_(k) |
| 193 | |
| 194 | |
| 195 | # Overwrite autotune_baseline_atol and autotune_baseline_rtol |
| 196 | # if too many configs failed due to baseline check during autotuning |
| 197 | @register_kernel( |
| 198 | mutates_args=["qkv"], |
| 199 | config_picker=pick_config, |
| 200 | input_generator=generate_inputs, |
| 201 | helion_settings=helion.Settings( |
| 202 | autotune_baseline_fn=baseline, |
| 203 | autotune_baseline_atol=5e-2, |
| 204 | autotune_baseline_rtol=5e-2, |
| 205 | ignore_warnings=[helion.exc.TensorOperationInWrapper], |
| 206 | ), |
| 207 | ) # type: ignore[misc] |
| 208 | def fused_qk_norm_rope( |
| 209 | qkv: torch.Tensor, # [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim] |
| 210 | num_heads_q: int, |
| 211 | num_heads_k: int, |
| 212 | num_heads_v: int, |
| 213 | head_dim: int, |
| 214 | eps: float, |
| 215 | q_weight: torch.Tensor, |
| 216 | k_weight: torch.Tensor, |
| 217 | cos_sin_cache: torch.Tensor, # [max_position, rotary_dim] |
| 218 | is_neox: bool, |
| 219 | position_ids: torch.Tensor, # [num_tokens], |
| 220 | forced_token_heads_per_warp: int = -1, # dummy |
| 221 | ) -> None: |
| 222 | assert qkv.ndim == 2 |
| 223 | num_tokens = qkv.shape[0] |
| 224 | total_heads = num_heads_q + num_heads_k + num_heads_v |
| 225 | assert qkv.shape[1] == total_heads * head_dim |
| 226 | hl.specialize(qkv.shape[1]) |
| 227 | |
| 228 | assert cos_sin_cache.ndim == 2 |
| 229 | max_position, rotary_dim = cos_sin_cache.shape |
| 230 | hl.specialize(max_position) |
| 231 | hl.specialize(rotary_dim) |
| 232 | assert rotary_dim % 2 == 0 |
| 233 | assert rotary_dim <= head_dim |
| 234 | embed_dim = rotary_dim // 2 |
| 235 | |
| 236 | hl.specialize(num_heads_q) |
| 237 | hl.specialize(num_heads_k) |
| 238 | hl.specialize(num_heads_v) |
| 239 | hl.specialize(head_dim) |
| 240 | |
| 241 | assert position_ids.ndim == 1 and position_ids.shape[0] == num_tokens |
| 242 | hl.specialize(position_ids.shape[0]) |
| 243 | |
| 244 | assert q_weight.ndim == 1 and q_weight.shape[0] == head_dim |
| 245 | hl.specialize(q_weight.shape[0]) |
| 246 | assert k_weight.ndim == 1 and k_weight.shape[0] == head_dim |
| 247 | hl.specialize(k_weight.shape[0]) |
| 248 | |
| 249 | assert qkv.dtype == q_weight.dtype and q_weight.dtype == k_weight.dtype |
| 250 | assert position_ids.dtype == torch.int64 |
| 251 | |
| 252 | assert qkv.is_contiguous() |
| 253 | assert position_ids.is_contiguous() |
| 254 | assert q_weight.is_contiguous() |
| 255 | assert k_weight.is_contiguous() |
| 256 | assert cos_sin_cache.is_contiguous() |
| 257 | |
| 258 | qk_heads = num_heads_q + num_heads_k |
| 259 | |
| 260 | qkv = qkv.view(num_tokens, -1, head_dim) |
| 261 | |
| 262 | for tile_m, tile_gn, tile_n in hl.tile( |
| 263 | [num_tokens, qk_heads, head_dim], block_size=[1, None, head_dim] |
| 264 | ): |
| 265 | x_blk = qkv[tile_m, tile_gn, tile_n].to(dtype=torch.float32) |
| 266 | |
| 267 | rms = x_blk.pow(2).sum(dim=-1) |
| 268 | rms = torch.rsqrt(rms * (1.0 / head_dim) + eps) |
| 269 | |
| 270 | use_q_weight = (tile_gn.index < num_heads_q)[None, :, None] |
| 271 | w_blk = torch.where( |
| 272 | use_q_weight, q_weight[None, None, tile_n], k_weight[None, None, tile_n] |
| 273 | ) |
| 274 | |
| 275 | x_blk = (x_blk * rms[:, :, None]).to(qkv.dtype) * w_blk |
| 276 | |
| 277 | qkv[tile_m, tile_gn, tile_n] = x_blk |
| 278 | |
| 279 | pos_id = position_ids[tile_m] |
| 280 | cos_blk = cos_sin_cache[pos_id, hl.arange(embed_dim)] |
| 281 | sin_blk = cos_sin_cache[pos_id, hl.arange(embed_dim) + embed_dim] |
| 282 | |
| 283 | if is_neox: |
| 284 | x1_offset = hl.arange(embed_dim) |
| 285 | x2_offset = x1_offset + embed_dim |
| 286 | else: |
| 287 | x1_offset = hl.arange(embed_dim) * 2 |
| 288 | x2_offset = x1_offset + 1 |
| 289 | |
| 290 | x1_blk = qkv[tile_m, tile_gn, x1_offset] |
| 291 | x2_blk = qkv[tile_m, tile_gn, x2_offset] |
| 292 | |
| 293 | o1_blk = x1_blk * cos_blk[:, None, :] - x2_blk * sin_blk[:, None, :] |
| 294 | o2_blk = x2_blk * cos_blk[:, None, :] + x1_blk * sin_blk[:, None, :] |
| 295 | |
| 296 | qkv[tile_m, tile_gn, x1_offset] = o1_blk |
| 297 | qkv[tile_m, tile_gn, x2_offset] = o2_blk |