| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from __future__ import annotations |
| 5 | |
| 6 | import math |
| 7 | from dataclasses import dataclass |
| 8 | from functools import cache |
| 9 | from typing import Any |
| 10 | |
| 11 | import torch |
| 12 | |
| 13 | from vllm.model_executor.warmup.jit_warmup import WarmupIntRange, zip_inputs |
| 14 | from vllm.model_executor.warmup.jit_warmup_tilelang_helper import ( |
| 15 | TileLangLaunchSpec, |
| 16 | VllmTileLangJitKernel, |
| 17 | kernel_launcher, |
| 18 | make_tilelang_warmup_tensor, |
| 19 | ) |
| 20 | from vllm.platforms import current_platform |
| 21 | from vllm.tilelang_utils import T, tilelang, tilelang_jit |
| 22 | from vllm.utils.math_utils import cdiv |
| 23 | |
| 24 | ENABLE_PDL = current_platform.is_arch_support_pdl() and current_platform.is_cuda() |
| 25 | |
| 26 | # TileLang's HIP AllReduce shuffle phase ignores a nonzero reduction-group |
| 27 | # offset, so the 64-thread RMSNorm reduction must be wave64-aligned and use |
| 28 | # threads 0..63. |
| 29 | _MIX_ON_TRAILING_THREADS = current_platform.is_rocm() |
| 30 | |
| 31 | |
| 32 | def _is_mix_thread(): |
| 33 | if _MIX_ON_TRAILING_THREADS: |
| 34 | return T.get_thread_binding() >= 64 |
| 35 | return T.get_thread_binding() < 32 |
| 36 | |
| 37 | |
| 38 | def _is_collapse_thread(): |
| 39 | if _MIX_ON_TRAILING_THREADS: |
| 40 | return T.get_thread_binding() < 64 |
| 41 | return T.get_thread_binding() >= 32 |
| 42 | |
| 43 | |
| 44 | @cache |
| 45 | def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: |
| 46 | device_props = torch.cuda.get_device_properties(0) |
| 47 | n_sms = device_props.multi_processor_count |
| 48 | split_k = n_sms // grid_size |
| 49 | if k is not None: |
| 50 | # avoid split_k for small k |
| 51 | num_block_k = cdiv(k, block_k) |
| 52 | split_k = min(split_k, num_block_k // 4) |
| 53 | split_k = max(split_k, 1) |
| 54 | return split_k |
| 55 | |
| 56 | |
| 57 | # ``mhc_fused_tilelang`` splits the hidden size across a fixed thread block, so |
| 58 | # every (tile_n, n_thr) pair is one more kernel to compile. Tuned on GB300 |
| 59 | # (hc_mult 4, hidden_size 5120 and 7168) against the separate post + split-k |
| 60 | # GEMM it replaces; the bands are coarse because the loss against a per-shape |
| 61 | # optimum is ~1-2%, while one fixed config gives up to 18% at the top of the |
| 62 | # range. |
| 63 | _FUSED_POST_PRE_N_THR = 128 |
| 64 | _FUSED_POST_PRE_MAX_TOKENS = 32 |
| 65 | |
| 66 | |
| 67 | def mhc_fused_post_pre_split_config( |
| 68 | num_tokens: int, hidden_size: int, hc_mult: int |
| 69 | ) -> tuple[int, int, int] | None: |
| 70 | """Pick ``(tile_n, n_splits, n_thr)`` for the fused post + pre-norm GEMM. |
| 71 | |
| 72 | Returns None when a separate post kernel followed by a split-k GEMM is the |
| 73 | better choice, or when the hidden size does not divide evenly across the |
| 74 | fused kernel's block. One source of truth for the dispatch decision, the |
| 75 | compile key and the launch, which must agree. |
| 76 | """ |
| 77 | if num_tokens > _FUSED_POST_PRE_MAX_TOKENS: |
| 78 | return None |
| 79 | n_thr = _FUSED_POST_PRE_N_THR |
| 80 | if hidden_size % n_thr: |
| 81 | return None |
| 82 | # The projection tiles supply the parallelism a small token count cannot, |
| 83 | # so the tile grows with the batch: 1.3x at one token and 1.03x at 32 |
| 84 | # against a split-k GEMM given the split its own estimator computes. |
| 85 | if num_tokens < 16: |
| 86 | tile_n, n_splits = 2, 8 |
| 87 | else: |
| 88 | tile_n, n_splits = 6, 8 |
| 89 | # The kernel drops projection tiles and k-slices it cannot fill evenly. |
| 90 | mix_size = hc_mult * (hc_mult + 2) |
| 91 | while tile_n > 1 and mix_size % tile_n: |
| 92 | tile_n -= 1 |
| 93 | while n_splits > 1 and hidden_size % (n_splits * n_thr): |
| 94 | n_splits //= 2 |
| 95 | return tile_n, n_splits, n_thr |
| 96 | |
| 97 | |
| 98 | def require_fused_post_pre_config( |
| 99 | num_tokens: int, hidden_size: int, hc_mult: int |
| 100 | ) -> tuple[int, int, int]: |
| 101 | """The config for a shape the caller has already committed to fusing.""" |
| 102 | config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) |
| 103 | if config is None: |
| 104 | raise ValueError( |
| 105 | "the fused mHC post + pre-norm GEMM does not cover num_tokens=" |
| 106 | f"{num_tokens} at hidden_size={hidden_size}" |
| 107 | ) |
| 108 | return config |
| 109 | |
| 110 | |
| 111 | def mhc_fused_post_pre_splits(hidden_size: int, hc_mult: int) -> tuple[int, ...]: |
| 112 | """Every split-k factor the fused post + pre-norm GEMM path can pick.""" |
| 113 | configs = ( |
| 114 | mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) |
| 115 | for num_tokens in range(1, _FUSED_POST_PRE_MAX_TOKENS + 1) |
| 116 | ) |
| 117 | return tuple(sorted({config[1] for config in configs if config is not None})) |
| 118 | |
| 119 | |
| 120 | @tilelang_jit |
| 121 | def mhc_pre_big_fuse_tilelang( |
| 122 | gemm_out_mul, |
| 123 | gemm_out_sqrsum, |
| 124 | hc_scale, |
| 125 | hc_base, |
| 126 | residual, |
| 127 | post_mix, |
| 128 | comb_mix, |
| 129 | layer_input, |
| 130 | pre_mix_in, |
| 131 | pre_mix_out, |
| 132 | aux_out, |
| 133 | hidden_size: int, |
| 134 | rms_eps: float, |
| 135 | hc_pre_eps: float, |
| 136 | hc_sinkhorn_eps: float, |
| 137 | hc_post_mult_value: float, |
| 138 | sinkhorn_repeat: int, |
| 139 | n_splits: int = 16, |
| 140 | hc_mult: int = 4, |
| 141 | use_pre_mix_in: bool = False, |
| 142 | save_pre_mix: bool = False, |
| 143 | rms_numel: int = 0, |
| 144 | write_aux: bool = False, |
| 145 | ): |
| 146 | """Fuse coefficient generation and residual collapse after the projection. |
| 147 | |
| 148 | With save_pre_mix, store the new pre-mix and collapse with pre_mix_in, |
| 149 | or select stream zero when use_pre_mix_in is false. |
| 150 | """ |
| 151 | num_tokens = T.dynamic("num_tokens") |
| 152 | hc_mult3 = hc_mult * (2 + hc_mult) |
| 153 | hidden_block = math.gcd(512, hidden_size) |
| 154 | if rms_numel == 0: |
| 155 | rms_numel = hc_mult * hidden_size |
| 156 | |
| 157 | gemm_out_mul: T.Tensor[[n_splits, num_tokens, hc_mult3], T.float32] # type: ignore[no-redef, valid-type] |
| 158 | gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32] # type: ignore[no-redef, valid-type] |
| 159 | hc_scale: T.Tensor[[3], T.float32] # type: ignore[no-redef, valid-type] |
| 160 | hc_base: T.Tensor[[hc_mult3], T.float32] # type: ignore[no-redef, valid-type] |
| 161 | residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] |
| 162 | # outputs |
| 163 | post_mix: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] |
| 164 | comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32] # type: ignore[no-redef, valid-type] |
| 165 | layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] |
| 166 | |
| 167 | pre_mix_in: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] |
| 168 | pre_mix_out: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] |
| 169 | aux_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] |
| 170 | |
| 171 | with T.Kernel(num_tokens, threads=96) as i: |
| 172 | if ENABLE_PDL: |
| 173 | T.pdl_sync() |
| 174 | ################################################################## |
| 175 | # _pre_norm_fn_fwd_norm |
| 176 | rms = T.alloc_fragment(1, T.float32) |
| 177 | mixes = T.alloc_fragment(hc_mult3, T.float32) |
| 178 | T.clear(mixes) |
| 179 | rms[0] = 0 |
| 180 | for i_split in T.serial(n_splits): |
| 181 | rms[0] += gemm_out_sqrsum[i_split, i] |
| 182 | rms[0] = T.rsqrt(rms[0] / rms_numel + rms_eps) |
| 183 | for j in T.Parallel(hc_mult3): |
| 184 | mixes[j] = 0 |
| 185 | for i_split in T.serial(n_splits): |
| 186 | mixes[j] += gemm_out_mul[i_split, i, j] |
| 187 | mixes[j] *= rms[0] |
| 188 | mixes_shared = T.alloc_shared(hc_mult3, T.float32) |
| 189 | T.copy(mixes, mixes_shared) |
| 190 | |
| 191 | if T.get_thread_binding() < 32: |
| 192 | ################################################################## |
| 193 | # _pre_split_mixes_fwd (post & comb) |
| 194 | cm = T.alloc_fragment((hc_mult, hc_mult), T.float32) |
| 195 | for j in T.Parallel(hc_mult): |
| 196 | if save_pre_mix: |
| 197 | pre_mix_out[i, j] = ( |
| 198 | T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j]) |
| 199 | + hc_pre_eps |
| 200 | ) |
| 201 | post_mix[i, j] = ( |
| 202 | T.sigmoid( |
| 203 | mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult] |
| 204 | ) |
| 205 | * hc_post_mult_value |
| 206 | ) |
| 207 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 208 | cm[j, k] = ( |
| 209 | mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2] |
| 210 | + hc_base[j * hc_mult + k + hc_mult * 2] |
| 211 | ) |
| 212 | |
| 213 | ################################################################## |
| 214 | # _sinkhorn_fwd |
| 215 | row_sum = T.alloc_fragment(hc_mult, T.float32) |
| 216 | col_sum = T.alloc_fragment(hc_mult, T.float32) |
| 217 | |
| 218 | # comb = comb.softmax(-1) + eps |
| 219 | row_max = T.alloc_fragment(hc_mult, T.float32) |
| 220 | T.reduce_max(cm, row_max, dim=1) |
| 221 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 222 | cm[j, k] = T.exp(cm[j, k] - row_max[j]) |
| 223 | T.reduce_sum(cm, row_sum, dim=1) |
| 224 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 225 | cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps |
| 226 | |
| 227 | # comb = comb / (comb.sum(-2) + eps) |
| 228 | T.reduce_sum(cm, col_sum, dim=0) |
| 229 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 230 | cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps) |
| 231 | |
| 232 | for _ in T.serial(sinkhorn_repeat - 1): |
| 233 | # comb = comb / (comb.sum(-1) + eps) |
| 234 | T.reduce_sum(cm, row_sum, dim=1) |
| 235 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 236 | cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps) |
| 237 | |
| 238 | # comb = comb / (comb.sum(-2) + eps) |
| 239 | T.reduce_sum(cm, col_sum, dim=0) |
| 240 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 241 | cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps) |
| 242 | |
| 243 | # save comb_mix to global memory |
| 244 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 245 | comb_mix[i, j * hc_mult + k] = cm[j, k] |
| 246 | else: |
| 247 | ################################################################## |
| 248 | # _pre_split_mixes_fwd (pre) |
| 249 | pre_mix_shared = T.alloc_shared(hc_mult, T.float32) |
| 250 | for j in T.Parallel(hc_mult): |
| 251 | if use_pre_mix_in: |
| 252 | pre_mix_shared[j] = pre_mix_in[i, j] |
| 253 | elif save_pre_mix: |
| 254 | pre_mix_shared[j] = T.if_then_else(j == 0, 1.0, 0.0) |
| 255 | else: |
| 256 | pre_mix_shared[j] = ( |
| 257 | T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j]) |
| 258 | + hc_pre_eps |
| 259 | ) |
| 260 | ################################################################### |
| 261 | # _pre_apply_mix_fwd |
| 262 | for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2): |
| 263 | xs = T.alloc_shared((hc_mult, hidden_block), T.float32) |
| 264 | xl = T.alloc_fragment((hc_mult, hidden_block), T.float32) |
| 265 | T.copy(residual[i, 0, i0_h * hidden_block], xs) |
| 266 | T.copy(xs, xl) |
| 267 | |
| 268 | ol = T.alloc_fragment(hidden_block, T.float32) |
| 269 | T.clear(ol) |
| 270 | if write_aux: |
| 271 | # Aux consumers (draft models) take the plain stream mean |
| 272 | # of the same residual this collapse reads. |
| 273 | aux = T.alloc_fragment(hidden_block, T.float32) |
| 274 | T.clear(aux) |
| 275 | |
| 276 | for i_hc in T.serial(hc_mult): |
| 277 | pre = pre_mix_shared[i_hc] |
| 278 | for i1_h in T.Parallel(hidden_block): |
| 279 | ol[i1_h] += pre * xl[i_hc, i1_h] |
| 280 | if write_aux: |
| 281 | aux[i1_h] += xl[i_hc, i1_h] |
| 282 | |
| 283 | if write_aux: |
| 284 | for i1_h in T.Parallel(hidden_block): |
| 285 | aux_out[i, i0_h * hidden_block + i1_h] = T.bfloat16( |
| 286 | aux[i1_h] / hc_mult |
| 287 | ) |
| 288 | |
| 289 | T.copy(ol, layer_input[i, i0_h * hidden_block]) |
| 290 | |
| 291 | if ENABLE_PDL: |
| 292 | T.pdl_trigger() |
| 293 | |
| 294 | |
| 295 | # Copied from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/mhc.py#L478 |
| 296 | |
| 297 | |
| 298 | @tilelang_jit |
| 299 | def mhc_pre_big_fuse_with_norm_tilelang( |
| 300 | gemm_out_mul, |
| 301 | gemm_out_sqrsum, |
| 302 | hc_scale, |
| 303 | hc_base, |
| 304 | residual, |
| 305 | post_mix, |
| 306 | comb_mix, |
| 307 | layer_input, |
| 308 | norm_weight, |
| 309 | pre_mix_in, |
| 310 | pre_mix_out, |
| 311 | aux_out, |
| 312 | hidden_size: int, |
| 313 | rms_eps: float, |
| 314 | hc_pre_eps: float, |
| 315 | hc_sinkhorn_eps: float, |
| 316 | hc_post_mult_value: float, |
| 317 | sinkhorn_repeat: int, |
| 318 | norm_eps: float, |
| 319 | n_splits: int = 16, |
| 320 | hc_mult: int = 4, |
| 321 | gemm_last_dim: int = -1, |
| 322 | use_pre_mix_in: bool = False, |
| 323 | save_pre_mix: bool = False, |
| 324 | rms_numel: int = 0, |
| 325 | split_mode: str = "fused", |
| 326 | write_aux: bool = False, |
| 327 | ): |
| 328 | # Split modes require shifted mHC: input collapse uses a carried pre-mix. |
| 329 | assert split_mode in ("fused", "stats", "input") |
| 330 | assert split_mode == "fused" or save_pre_mix |
| 331 | num_tokens = T.dynamic("num_tokens") |
| 332 | hc_mult3 = hc_mult * (2 + hc_mult) |
| 333 | if gemm_last_dim < 0: |
| 334 | gemm_last_dim = hc_mult3 |
| 335 | if rms_numel == 0: |
| 336 | rms_numel = hc_mult * hidden_size |
| 337 | hidden_block = math.gcd(1024, hidden_size) |
| 338 | |
| 339 | gemm_out_mul: T.Tensor[[n_splits, num_tokens, gemm_last_dim], T.float32] # type: ignore[no-redef, valid-type] |
| 340 | gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32] # type: ignore[no-redef, valid-type] |
| 341 | hc_scale: T.Tensor[[3], T.float32] # type: ignore[no-redef, valid-type] |
| 342 | hc_base: T.Tensor[[hc_mult3], T.float32] # type: ignore[no-redef, valid-type] |
| 343 | residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] |
| 344 | post_mix: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] |
| 345 | comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32] # type: ignore[no-redef, valid-type] |
| 346 | layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] |
| 347 | norm_weight: T.Tensor[[hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] |
| 348 | |
| 349 | pre_mix_in: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] |
| 350 | pre_mix_out: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] |
| 351 | aux_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] |
| 352 | |
| 353 | with T.Kernel(num_tokens, threads=96) as i: |
| 354 | if ENABLE_PDL: |
| 355 | T.pdl_sync() |
| 356 | mixes_shared = T.alloc_shared(hc_mult3, T.float32) |
| 357 | if split_mode != "input": |
| 358 | rms = T.alloc_fragment(1, T.float32) |
| 359 | mixes = T.alloc_fragment(hc_mult3, T.float32) |
| 360 | T.clear(mixes) |
| 361 | rms[0] = 0 |
| 362 | |
| 363 | for i_split in T.serial(n_splits): |
| 364 | rms[0] += gemm_out_sqrsum[i_split, i] |
| 365 | rms[0] = T.rsqrt(rms[0] / rms_numel + rms_eps) |
| 366 | for j in T.Parallel(hc_mult3): |
| 367 | mixes[j] = 0 |
| 368 | for i_split in T.serial(n_splits): |
| 369 | mixes[j] += gemm_out_mul[i_split, i, j] |
| 370 | mixes[j] *= rms[0] |
| 371 | T.copy(mixes, mixes_shared) |
| 372 | |
| 373 | if split_mode != "input" and _is_mix_thread(): |
| 374 | cm = T.alloc_fragment((hc_mult, hc_mult), T.float32) |
| 375 | for j in T.Parallel(hc_mult): |
| 376 | if save_pre_mix: |
| 377 | pre_mix_out[i, j] = ( |
| 378 | T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j]) |
| 379 | + hc_pre_eps |
| 380 | ) |
| 381 | post_mix[i, j] = ( |
| 382 | T.sigmoid( |
| 383 | mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult] |
| 384 | ) |
| 385 | * hc_post_mult_value |
| 386 | ) |
| 387 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 388 | cm[j, k] = ( |
| 389 | mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2] |
| 390 | + hc_base[j * hc_mult + k + hc_mult * 2] |
| 391 | ) |
| 392 | |
| 393 | row_sum = T.alloc_fragment(hc_mult, T.float32) |
| 394 | col_sum = T.alloc_fragment(hc_mult, T.float32) |
| 395 | |
| 396 | row_max = T.alloc_fragment(hc_mult, T.float32) |
| 397 | T.reduce_max(cm, row_max, dim=1) |
| 398 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 399 | cm[j, k] = T.exp(cm[j, k] - row_max[j]) |
| 400 | T.reduce_sum(cm, row_sum, dim=1) |
| 401 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 402 | cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps |
| 403 | |
| 404 | T.reduce_sum(cm, col_sum, dim=0) |
| 405 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 406 | cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps) |
| 407 | |
| 408 | for _ in T.serial(sinkhorn_repeat - 1): |
| 409 | T.reduce_sum(cm, row_sum, dim=1) |
| 410 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 411 | cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps) |
| 412 | |
| 413 | T.reduce_sum(cm, col_sum, dim=0) |
| 414 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 415 | cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps) |
| 416 | |
| 417 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 418 | comb_mix[i, j * hc_mult + k] = cm[j, k] |
| 419 | elif split_mode != "stats" and _is_collapse_thread(): |
| 420 | pre_mix_shared = T.alloc_shared(hc_mult, T.float32) |
| 421 | for j in T.Parallel(hc_mult): |
| 422 | if use_pre_mix_in: |
| 423 | pre_mix_shared[j] = pre_mix_in[i, j] |
| 424 | elif save_pre_mix: |
| 425 | pre_mix_shared[j] = T.if_then_else(j == 0, 1.0, 0.0) |
| 426 | else: |
| 427 | pre_mix_shared[j] = ( |
| 428 | T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j]) |
| 429 | + hc_pre_eps |
| 430 | ) |
| 431 | |
| 432 | # Pass 1: stash unnormalized weighted-sum output in shared memory |
| 433 | # as bf16 (matches the rounding that RMSNorm would see) while |
| 434 | # accumulating the per-position squared sum. |
| 435 | output_shared = T.alloc_shared(hidden_size, T.bfloat16) |
| 436 | sumsq_per_pos = T.alloc_fragment(hidden_block, T.float32) |
| 437 | T.clear(sumsq_per_pos) |
| 438 | |
| 439 | for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2): |
| 440 | xs = T.alloc_shared((hc_mult, hidden_block), T.bfloat16) |
| 441 | xl = T.alloc_fragment((hc_mult, hidden_block), T.float32) |
| 442 | T.copy(residual[i, 0, i0_h * hidden_block], xs) |
| 443 | T.copy(xs, xl) |
| 444 | |
| 445 | ol = T.alloc_fragment(hidden_block, T.float32) |
| 446 | T.clear(ol) |
| 447 | if write_aux: |
| 448 | # Aux consumers (draft models) take the plain stream mean |
| 449 | # of the same residual this collapse reads. |
| 450 | aux = T.alloc_fragment(hidden_block, T.float32) |
| 451 | T.clear(aux) |
| 452 | |
| 453 | for i_hc in T.serial(hc_mult): |
| 454 | pre = pre_mix_shared[i_hc] |
| 455 | for i1_h in T.Parallel(hidden_block): |
| 456 | ol[i1_h] += pre * xl[i_hc, i1_h] |
| 457 | if write_aux: |
| 458 | aux[i1_h] += xl[i_hc, i1_h] |
| 459 | |
| 460 | if write_aux: |
| 461 | for i1_h in T.Parallel(hidden_block): |
| 462 | aux_out[i, i0_h * hidden_block + i1_h] = T.bfloat16( |
| 463 | aux[i1_h] / hc_mult |
| 464 | ) |
| 465 | |
| 466 | if save_pre_mix: |
| 467 | # Keep the BF16 boundary before the delayed input RMSNorm. |
| 468 | rounded = T.alloc_fragment(hidden_block, T.bfloat16) |
| 469 | T.copy(ol, rounded) |
| 470 | for i1_h in T.Parallel(hidden_block): |
| 471 | value = T.float32(rounded[i1_h]) |
| 472 | sumsq_per_pos[i1_h] += value * value |
| 473 | output_shared[i0_h * hidden_block + i1_h] = rounded[i1_h] |
| 474 | else: |
| 475 | for i1_h in T.Parallel(hidden_block): |
| 476 | sumsq_per_pos[i1_h] += ol[i1_h] * ol[i1_h] |
| 477 | output_shared[i0_h * hidden_block + i1_h] = T.bfloat16(ol[i1_h]) |
| 478 | |
| 479 | sumsq = T.alloc_fragment(1, T.float32) |
| 480 | T.reduce_sum(sumsq_per_pos, sumsq, dim=0) |
| 481 | rsqrt_norm = T.alloc_fragment(1, T.float32) |
| 482 | rsqrt_norm[0] = T.rsqrt(sumsq[0] / hidden_size + norm_eps) |
| 483 | |
| 484 | # Pass 2: scale by rsqrt * norm_weight and write the result to HBM. |
| 485 | for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2): |
| 486 | w_shared = T.alloc_shared(hidden_block, T.bfloat16) |
| 487 | w_local = T.alloc_fragment(hidden_block, T.float32) |
| 488 | T.copy(norm_weight[i0_h * hidden_block], w_shared) |
| 489 | T.copy(w_shared, w_local) |
| 490 | |
| 491 | ol = T.alloc_fragment(hidden_block, T.float32) |
| 492 | for i1_h in T.Parallel(hidden_block): |
| 493 | ol[i1_h] = ( |
| 494 | output_shared[i0_h * hidden_block + i1_h] |
| 495 | * rsqrt_norm[0] |
| 496 | * w_local[i1_h] |
| 497 | ) |
| 498 | |
| 499 | T.copy(ol, layer_input[i, i0_h * hidden_block]) |
| 500 | |
| 501 | if ENABLE_PDL: |
| 502 | T.pdl_trigger() |
| 503 | |
| 504 | |
| 505 | @tilelang_jit |
| 506 | def mhc_pre_big_fuse_broadcast_with_norm_tilelang( |
| 507 | gemm_out_mul, |
| 508 | gemm_out_sqrsum, |
| 509 | hc_scale, |
| 510 | hc_base, |
| 511 | residual, |
| 512 | residual_out, |
| 513 | post_mix, |
| 514 | comb_mix, |
| 515 | layer_input, |
| 516 | norm_weight, |
| 517 | hidden_size: int, |
| 518 | rms_eps: float, |
| 519 | hc_pre_eps: float, |
| 520 | hc_sinkhorn_eps: float, |
| 521 | hc_post_mult_value: float, |
| 522 | sinkhorn_repeat: int, |
| 523 | norm_eps: float, |
| 524 | n_splits: int = 1, |
| 525 | hc_mult: int = 4, |
| 526 | gemm_last_dim: int = -1, |
| 527 | ): |
| 528 | num_tokens = T.dynamic("num_tokens") |
| 529 | hc_mult3 = hc_mult * (2 + hc_mult) |
| 530 | if gemm_last_dim < 0: |
| 531 | gemm_last_dim = hc_mult3 |
| 532 | hidden_block = math.gcd(1024, hidden_size) |
| 533 | |
| 534 | gemm_out_mul: T.Tensor[[n_splits, num_tokens, gemm_last_dim], T.float32] # type: ignore[no-redef, valid-type] |
| 535 | gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32] # type: ignore[no-redef, valid-type] |
| 536 | hc_scale: T.Tensor[[3], T.float32] # type: ignore[no-redef, valid-type] |
| 537 | hc_base: T.Tensor[[hc_mult3], T.float32] # type: ignore[no-redef, valid-type] |
| 538 | residual: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] |
| 539 | residual_out: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] |
| 540 | post_mix: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] |
| 541 | comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32] # type: ignore[no-redef, valid-type] |
| 542 | layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] |
| 543 | norm_weight: T.Tensor[[hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] |
| 544 | |
| 545 | with T.Kernel(num_tokens, threads=96) as i: |
| 546 | rms = T.alloc_fragment(1, T.float32) |
| 547 | mixes = T.alloc_fragment(hc_mult3, T.float32) |
| 548 | T.clear(mixes) |
| 549 | rms[0] = 0 |
| 550 | |
| 551 | if ENABLE_PDL: |
| 552 | T.pdl_sync() |
| 553 | |
| 554 | for i_split in T.serial(n_splits): |
| 555 | rms[0] += gemm_out_sqrsum[i_split, i] |
| 556 | # rms scale |
| 557 | rms[0] = T.rsqrt(rms[0] / hidden_size + rms_eps) |
| 558 | for j in T.Parallel(hc_mult3): |
| 559 | mixes[j] = 0 |
| 560 | for i_split in T.serial(n_splits): |
| 561 | mixes[j] += gemm_out_mul[i_split, i, j] |
| 562 | mixes[j] *= rms[0] |
| 563 | mixes_shared = T.alloc_shared(hc_mult3, T.float32) |
| 564 | T.copy(mixes, mixes_shared) |
| 565 | |
| 566 | if _is_mix_thread(): |
| 567 | cm = T.alloc_fragment((hc_mult, hc_mult), T.float32) |
| 568 | for j in T.Parallel(hc_mult): |
| 569 | post_mix[i, j] = ( |
| 570 | T.sigmoid( |
| 571 | mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult] |
| 572 | ) |
| 573 | * hc_post_mult_value |
| 574 | ) |
| 575 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 576 | cm[j, k] = ( |
| 577 | mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2] |
| 578 | + hc_base[j * hc_mult + k + hc_mult * 2] |
| 579 | ) |
| 580 | |
| 581 | row_sum = T.alloc_fragment(hc_mult, T.float32) |
| 582 | col_sum = T.alloc_fragment(hc_mult, T.float32) |
| 583 | |
| 584 | row_max = T.alloc_fragment(hc_mult, T.float32) |
| 585 | T.reduce_max(cm, row_max, dim=1) |
| 586 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 587 | cm[j, k] = T.exp(cm[j, k] - row_max[j]) |
| 588 | T.reduce_sum(cm, row_sum, dim=1) |
| 589 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 590 | cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps |
| 591 | |
| 592 | T.reduce_sum(cm, col_sum, dim=0) |
| 593 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 594 | cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps) |
| 595 | |
| 596 | for _ in T.serial(sinkhorn_repeat - 1): |
| 597 | T.reduce_sum(cm, row_sum, dim=1) |
| 598 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 599 | cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps) |
| 600 | |
| 601 | T.reduce_sum(cm, col_sum, dim=0) |
| 602 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 603 | cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps) |
| 604 | |
| 605 | for j, k in T.Parallel(hc_mult, hc_mult): |
| 606 | comb_mix[i, j * hc_mult + k] = cm[j, k] |
| 607 | else: |
| 608 | pre_mix_shared = T.alloc_shared(hc_mult, T.float32) |
| 609 | for j in T.Parallel(hc_mult): |
| 610 | pre_mix_shared[j] = ( |
| 611 | T.sigmoid( |
| 612 | mixes_shared[j] * hc_scale[0] + hc_base[j], |
| 613 | ) |
| 614 | + hc_pre_eps |
| 615 | ) |
| 616 | |
| 617 | output_shared = T.alloc_shared(hidden_size, T.bfloat16) |
| 618 | sumsq_per_pos = T.alloc_fragment(hidden_block, T.float32) |
| 619 | T.clear(sumsq_per_pos) |
| 620 | |
| 621 | for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2): |
| 622 | xs = T.alloc_shared(hidden_block, T.bfloat16) |
| 623 | xl = T.alloc_fragment(hidden_block, T.float32) |
| 624 | T.copy(residual[i, i0_h * hidden_block], xs) |
| 625 | T.copy(xs, xl) |
| 626 | |
| 627 | for i_hc, i1_h in T.Parallel(hc_mult, hidden_block): |
| 628 | residual_out[i, i_hc, i0_h * hidden_block + i1_h] = xs[i1_h] |
| 629 | |
| 630 | ol = T.alloc_fragment(hidden_block, T.float32) |
| 631 | T.clear(ol) |
| 632 | |
| 633 | for i_hc in T.serial(hc_mult): |
| 634 | pre = pre_mix_shared[i_hc] |
| 635 | for i1_h in T.Parallel(hidden_block): |
| 636 | ol[i1_h] += pre * xl[i1_h] |
| 637 | |
| 638 | for i1_h in T.Parallel(hidden_block): |
| 639 | sumsq_per_pos[i1_h] += ol[i1_h] * ol[i1_h] |
| 640 | output_shared[i0_h * hidden_block + i1_h] = T.bfloat16(ol[i1_h]) |
| 641 | |
| 642 | sumsq = T.alloc_fragment(1, T.float32) |
| 643 | T.reduce_sum(sumsq_per_pos, sumsq, dim=0) |
| 644 | rsqrt_norm = T.alloc_fragment(1, T.float32) |
| 645 | rsqrt_norm[0] = T.rsqrt(sumsq[0] / hidden_size + norm_eps) |
| 646 | |
| 647 | for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2): |
| 648 | w_shared = T.alloc_shared(hidden_block, T.bfloat16) |
| 649 | w_local = T.alloc_fragment(hidden_block, T.float32) |
| 650 | T.copy(norm_weight[i0_h * hidden_block], w_shared) |
| 651 | T.copy(w_shared, w_local) |
| 652 | |
| 653 | ol = T.alloc_fragment(hidden_block, T.float32) |
| 654 | for i1_h in T.Parallel(hidden_block): |
| 655 | ol[i1_h] = ( |
| 656 | output_shared[i0_h * hidden_block + i1_h] |
| 657 | * rsqrt_norm[0] |
| 658 | * w_local[i1_h] |
| 659 | ) |
| 660 | |
| 661 | T.copy(ol, layer_input[i, i0_h * hidden_block]) |
| 662 | |
| 663 | if ENABLE_PDL: |
| 664 | T.pdl_trigger() |
| 665 | |
| 666 | |
| 667 | @tilelang_jit |
| 668 | def mhc_fused_tilelang( |
| 669 | comb_mix, |
| 670 | residual_in, |
| 671 | post_mix, |
| 672 | x_in, |
| 673 | weight_t, |
| 674 | yp_out, |
| 675 | rp_out, |
| 676 | residual_out, |
| 677 | hc: int, |
| 678 | hidden: int, |
| 679 | n_out: int, |
| 680 | n_thr: int = 256, |
| 681 | h_blk: int = 256, |
| 682 | tile_n: int = 1, |
| 683 | split_k: int = 1, |
| 684 | ) -> tilelang.JITKernel: |
| 685 | """Fused mhc post-mapping + pre-norm GEMM FMA.""" |
| 686 | m = T.dynamic("num_tokens") |
| 687 | split_k = T.dynamic("split_k") |
| 688 | h = hidden |
| 689 | h_blk = math.gcd(hidden, h_blk) |
| 690 | h_per_split = h // split_k |
| 691 | n_tiles = n_out // tile_n |
| 692 | |
| 693 | comb_mix: T.Tensor((m, hc, hc), T.float32) # type: ignore[no-redef, valid-type] |
| 694 | residual_in: T.Tensor((m, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type] |
| 695 | post_mix: T.Tensor((m, hc), T.float32) # type: ignore[no-redef, valid-type] |
| 696 | x_in: T.Tensor((m, h), T.bfloat16) # type: ignore[no-redef, valid-type] |
| 697 | weight_t: T.Tensor((n_out, hc, h), T.float32) # type: ignore[no-redef, valid-type] |
| 698 | yp_out: T.Tensor((split_k, m, n_out), T.float32) # type: ignore[no-redef, valid-type] |
| 699 | rp_out: T.Tensor((split_k, m), T.float32) # type: ignore[no-redef, valid-type] |
| 700 | residual_out: T.Tensor((m, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type] |
| 701 | |
| 702 | h_iters = h_per_split // n_thr |
| 703 | num_warps = n_thr // 32 |
| 704 | |
| 705 | with T.Kernel(m, n_tiles, split_k, threads=n_thr) as (i_n, i_nt, i_ks): |
| 706 | tid = T.get_thread_binding() |
| 707 | warp_id = tid // 32 |
| 708 | lane = tid % 32 |
| 709 | |
| 710 | s_warp = T.alloc_shared((num_warps, tile_n + 1), T.float32) |
| 711 | s_post = T.alloc_shared((hc,), T.float32) |
| 712 | s_comb = T.alloc_shared((hc, hc), T.float32) |
| 713 | |
| 714 | pm = T.alloc_local((hc,), T.float32) |
| 715 | cm = T.alloc_local((hc, hc), T.float32) |
| 716 | acc = T.alloc_local((tile_n,), T.float32) |
| 717 | sqr = T.alloc_local((1,), T.float32) |
| 718 | new_r = T.alloc_local((hc,), T.float32) |
| 719 | |
| 720 | T.clear(acc) |
| 721 | T.clear(sqr) |
| 722 | h_split_start = i_ks * h_per_split |
| 723 | |
| 724 | if ENABLE_PDL: |
| 725 | T.pdl_sync() |
| 726 | |
| 727 | T.copy(post_mix[i_n, 0], s_post) |
| 728 | T.copy(comb_mix[i_n, 0, 0], s_comb) |
| 729 | |
| 730 | for j in T.unroll(hc): |
| 731 | pm[j] = s_post[j] |
| 732 | for j in T.unroll(hc): |
| 733 | for k in T.unroll(hc): |
| 734 | cm[k, j] = s_comb[k, j] |
| 735 | |
| 736 | # Each thread owns h_iters elements of the k-split's h slice. |
| 737 | for it in T.serial(h_iters): |
| 738 | h_idx = h_split_start + it * n_thr + tid |
| 739 | |
| 740 | # Compute new residual from layer output and past residual |
| 741 | for j in T.unroll(hc): |
| 742 | new_r[j] = pm[j] * x_in[i_n, h_idx] |
| 743 | for k in T.unroll(hc): |
| 744 | new_r[j] += cm[k, j] * residual_in[i_n, k, h_idx] |
| 745 | |
| 746 | # populate residual_out and compute sqr sum |
| 747 | if i_nt == 0: |
| 748 | for j in T.unroll(hc): |
| 749 | residual_out[i_n, j, h_idx] = new_r[j] |
| 750 | sqr[0] += new_r[j] * new_r[j] |
| 751 | |
| 752 | # Per-thread FMA into acc[n] |
| 753 | for n in T.unroll(tile_n): |
| 754 | for j in T.unroll(hc): |
| 755 | acc[n] += weight_t[i_nt * tile_n + n, j, h_idx] * new_r[j] |
| 756 | |
| 757 | for n in T.unroll(tile_n): |
| 758 | acc[n] = T.warp_reduce_sum(acc[n]) |
| 759 | if i_nt == 0: |
| 760 | sqr[0] = T.warp_reduce_sum(sqr[0]) |
| 761 | |
| 762 | # Cross-warp reduce via shared mem |
| 763 | if lane == 0: |
| 764 | for n in T.unroll(tile_n): |
| 765 | s_warp[warp_id, n] = acc[n] |
| 766 | if i_nt == 0: |
| 767 | s_warp[warp_id, tile_n] = sqr[0] |
| 768 | T.sync_threads() |
| 769 | |
| 770 | # Warp 0 does the final cross-warp sum and writes outputs |
| 771 | if warp_id == 0: |
| 772 | if lane < tile_n: |
| 773 | v = T.alloc_var(T.float32, init=0.0) |
| 774 | for w in T.unroll(num_warps): |
| 775 | v += s_warp[w, lane] |
| 776 | yp_out[i_ks, i_n, i_nt * tile_n + lane] = v |
| 777 | |
| 778 | if i_nt == 0 and lane == 0: |
| 779 | v2 = T.alloc_var(T.float32, init=0.0) |
| 780 | for w in T.unroll(num_warps): |
| 781 | v2 += s_warp[w, tile_n] |
| 782 | rp_out[i_ks, i_n] = v2 |
| 783 | |
| 784 | if ENABLE_PDL: |
| 785 | T.pdl_trigger() |
| 786 | |
| 787 | |
| 788 | @tilelang_jit |
| 789 | def mhc_post_tilelang( |
| 790 | a, |
| 791 | b, |
| 792 | c, |
| 793 | d, |
| 794 | x, |
| 795 | hc: int, |
| 796 | hidden: int, |
| 797 | n_thr: int = 128, |
| 798 | h_blk: int = 1024, |
| 799 | ) -> tilelang.JITKernel: |
| 800 | # rename for shorter code |
| 801 | n = T.dynamic("num_tokens") |
| 802 | h = hidden |
| 803 | |
| 804 | h_blk = math.gcd(hidden, h_blk) |
| 805 | a: T.Tensor((n, hc, hc), T.float32) # type: ignore[no-redef, valid-type] |
| 806 | b: T.Tensor((n, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type] |
| 807 | c: T.Tensor((n, hc), T.float32) # type: ignore[no-redef, valid-type] |
| 808 | d: T.Tensor((n, h), T.bfloat16) # type: ignore[no-redef, valid-type] |
| 809 | x: T.Tensor((n, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type] |
| 810 | with T.Kernel(n, threads=n_thr) as i_n: |
| 811 | b_shared = T.alloc_shared((hc, h_blk), T.bfloat16) |
| 812 | d_shared = T.alloc_shared(h_blk, T.bfloat16) |
| 813 | |
| 814 | x_local = T.alloc_fragment((hc, h_blk), T.float32) |
| 815 | b_local = T.alloc_fragment((hc, h_blk), T.float32) |
| 816 | d_local = T.alloc_fragment(h_blk, T.float32) |
| 817 | |
| 818 | a_local = T.alloc_fragment((hc, hc), T.float32) |
| 819 | c_local = T.alloc_fragment(hc, T.float32) |
| 820 | if ENABLE_PDL: |
| 821 | T.pdl_sync() |
| 822 | T.copy(a[i_n, 0, 0], a_local) |
| 823 | T.copy(c[i_n, 0], c_local) |
| 824 | |
| 825 | for i0_h in T.Serial(T.ceildiv(h, h_blk)): |
| 826 | T.copy(b[i_n, 0, i0_h * h_blk], b_shared) |
| 827 | T.copy(d[i_n, i0_h * h_blk], d_shared) |
| 828 | |
| 829 | T.copy(b_shared, b_local) |
| 830 | T.copy(d_shared, d_local) |
| 831 | for i_hco, i1_h in T.Parallel(hc, h_blk): |
| 832 | x_local[i_hco, i1_h] = c_local[i_hco] * d_local[i1_h] |
| 833 | for i_hci in T.vectorized(hc): |
| 834 | x_local[i_hco, i1_h] += a_local[i_hci, i_hco] * b_local[i_hci, i1_h] |
| 835 | |
| 836 | T.copy(x_local, x[i_n, 0, i0_h * h_blk]) |
| 837 | if ENABLE_PDL: |
| 838 | T.pdl_trigger() |
| 839 | |
| 840 | |
| 841 | @tilelang_jit |
| 842 | def hc_prenorm_gemm_tilelang( |
| 843 | x, |
| 844 | fn, |
| 845 | out, |
| 846 | sqrsum, |
| 847 | hidden_size: int, |
| 848 | hc_mult: int = 4, |
| 849 | n_out: int = 24, |
| 850 | n_thr: int = 512, |
| 851 | tile_n: int = 12, |
| 852 | n_splits: int = 1, |
| 853 | ) -> tilelang.JITKernel: |
| 854 | num_tokens = T.dynamic("num_tokens") |
| 855 | hc_hidden_size = hc_mult * hidden_size |
| 856 | k_per_split = hc_hidden_size // n_splits |
| 857 | k_iters = k_per_split // n_thr |
| 858 | n_tiles = T.ceildiv(n_out, tile_n) |
| 859 | |
| 860 | x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16) # type: ignore[no-redef, valid-type] |
| 861 | fn: T.Tensor((n_out, hc_hidden_size), T.float32) # type: ignore[no-redef, valid-type] |
| 862 | out: T.Tensor((n_splits, num_tokens, n_out), T.float32) # type: ignore[no-redef, valid-type] |
| 863 | sqrsum: T.Tensor((n_splits, num_tokens), T.float32) # type: ignore[no-redef, valid-type] |
| 864 | |
| 865 | with T.Kernel(num_tokens, n_tiles, n_splits, threads=n_thr) as ( |
| 866 | i_n, |
| 867 | i_t, |
| 868 | i_s, |
| 869 | ): |
| 870 | tid = T.get_thread_binding() |
| 871 | acc = T.alloc_local((tile_n,), T.float32) |
| 872 | sqr = T.alloc_local((1,), T.float32) |
| 873 | T.clear(acc) |
| 874 | T.clear(sqr) |
| 875 | |
| 876 | if ENABLE_PDL: |
| 877 | T.pdl_sync() |
| 878 | |
| 879 | for it in T.serial(k_iters): |
| 880 | i_k = i_s * k_per_split + it * n_thr + tid |
| 881 | x_val = T.cast(x[i_n, i_k], T.float32) |
| 882 | for i_o in T.unroll(tile_n): |
| 883 | out_idx = i_t * tile_n + i_o |
| 884 | if out_idx < n_out: |
| 885 | acc[i_o] += x_val * fn[out_idx, i_k] |
| 886 | if i_t == 0: |
| 887 | sqr[0] += x_val * x_val |
| 888 | |
| 889 | for i_o in T.unroll(tile_n): |
| 890 | acc[i_o] = T.warp_reduce_sum(acc[i_o]) |
| 891 | if i_t == 0: |
| 892 | sqr[0] = T.warp_reduce_sum(sqr[0]) |
| 893 | |
| 894 | lane = tid % 32 |
| 895 | warp_id = tid // 32 |
| 896 | num_warps = n_thr // 32 |
| 897 | warp_acc = T.alloc_shared((num_warps, tile_n), T.float32) |
| 898 | warp_sqr = T.alloc_shared(num_warps, T.float32) |
| 899 | |
| 900 | if lane == 0: |
| 901 | for i_o in T.unroll(tile_n): |
| 902 | warp_acc[warp_id, i_o] = acc[i_o] |
| 903 | if i_t == 0: |
| 904 | warp_sqr[warp_id] = sqr[0] |
| 905 | T.sync_threads() |
| 906 | |
| 907 | if warp_id == 0: |
| 908 | if lane < tile_n: |
| 909 | reduced_acc = T.alloc_var(T.float32, init=0.0) |
| 910 | for i_w in T.unroll(num_warps): |
| 911 | reduced_acc += warp_acc[i_w, lane] |
| 912 | out_idx = i_t * tile_n + lane |
| 913 | if out_idx < n_out: |
| 914 | out[i_s, i_n, out_idx] = reduced_acc |
| 915 | if lane == 0 and i_t == 0: |
| 916 | reduced_sqr = T.alloc_var(T.float32, init=0.0) |
| 917 | for i_w in T.unroll(num_warps): |
| 918 | reduced_sqr += warp_sqr[i_w] |
| 919 | sqrsum[i_s, i_n] = reduced_sqr |
| 920 | |
| 921 | if ENABLE_PDL: |
| 922 | T.pdl_trigger() |
| 923 | |
| 924 | |
| 925 | @tilelang_jit |
| 926 | def hc_prenorm_gemm_block_m_tilelang( |
| 927 | x, |
| 928 | fn, |
| 929 | out, |
| 930 | sqrsum, |
| 931 | hidden_size: int, |
| 932 | hc_mult: int = 4, |
| 933 | n_out: int = 24, |
| 934 | n_thr: int = 512, |
| 935 | tile_n: int = 12, |
| 936 | block_m: int = 2, |
| 937 | ) -> tilelang.JITKernel: |
| 938 | num_tokens = T.dynamic("num_tokens") |
| 939 | hc_hidden_size = hc_mult * hidden_size |
| 940 | k_iters = hc_hidden_size // n_thr |
| 941 | n_tiles = T.ceildiv(n_out, tile_n) |
| 942 | m_tiles = T.ceildiv(num_tokens, block_m) |
| 943 | |
| 944 | x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16) # type: ignore[no-redef, valid-type] |
| 945 | fn: T.Tensor((n_out, hc_hidden_size), T.float32) # type: ignore[no-redef, valid-type] |
| 946 | out: T.Tensor((1, num_tokens, n_out), T.float32) # type: ignore[no-redef, valid-type] |
| 947 | sqrsum: T.Tensor((1, num_tokens), T.float32) # type: ignore[no-redef, valid-type] |
| 948 | |
| 949 | with T.Kernel(m_tiles, n_tiles, threads=n_thr) as (i_mt, i_t): |
| 950 | tid = T.get_thread_binding() |
| 951 | acc = T.alloc_local((block_m, tile_n), T.float32) |
| 952 | sqr = T.alloc_local((block_m,), T.float32) |
| 953 | T.clear(acc) |
| 954 | T.clear(sqr) |
| 955 | |
| 956 | if ENABLE_PDL: |
| 957 | T.pdl_sync() |
| 958 | |
| 959 | for it in T.serial(k_iters): |
| 960 | i_k = it * n_thr + tid |
| 961 | fn_val = T.alloc_local((tile_n,), T.float32) |
| 962 | for i_o in T.unroll(tile_n): |
| 963 | out_idx = i_t * tile_n + i_o |
| 964 | if out_idx < n_out: |
| 965 | fn_val[i_o] = fn[out_idx, i_k] |
| 966 | else: |
| 967 | fn_val[i_o] = 0.0 |
| 968 | for i_m in T.unroll(block_m): |
| 969 | token_idx = i_mt * block_m + i_m |
| 970 | if token_idx < num_tokens: |
| 971 | x_val = T.cast(x[token_idx, i_k], T.float32) |
| 972 | for i_o in T.unroll(tile_n): |
| 973 | acc[i_m, i_o] += x_val * fn_val[i_o] |
| 974 | if i_t == 0: |
| 975 | sqr[i_m] += x_val * x_val |
| 976 | |
| 977 | for i_m in T.unroll(block_m): |
| 978 | for i_o in T.unroll(tile_n): |
| 979 | acc[i_m, i_o] = T.warp_reduce_sum(acc[i_m, i_o]) |
| 980 | if i_t == 0: |
| 981 | sqr[i_m] = T.warp_reduce_sum(sqr[i_m]) |
| 982 | |
| 983 | lane = tid % 32 |
| 984 | warp_id = tid // 32 |
| 985 | num_warps = n_thr // 32 |
| 986 | warp_acc = T.alloc_shared((num_warps, block_m, tile_n), T.float32) |
| 987 | warp_sqr = T.alloc_shared((num_warps, block_m), T.float32) |
| 988 | |
| 989 | if lane == 0: |
| 990 | for i_m in T.unroll(block_m): |
| 991 | for i_o in T.unroll(tile_n): |
| 992 | warp_acc[warp_id, i_m, i_o] = acc[i_m, i_o] |
| 993 | if i_t == 0: |
| 994 | warp_sqr[warp_id, i_m] = sqr[i_m] |
| 995 | T.sync_threads() |
| 996 | |
| 997 | if warp_id == 0: |
| 998 | reduced_acc = T.alloc_local((block_m,), T.float32) |
| 999 | reduced_sqr = T.alloc_local((block_m,), T.float32) |
| 1000 | T.clear(reduced_acc) |
| 1001 | T.clear(reduced_sqr) |
| 1002 | for i_m in T.unroll(block_m): |
| 1003 | token_idx = i_mt * block_m + i_m |
| 1004 | if token_idx < num_tokens: |
| 1005 | if lane < tile_n: |
| 1006 | for i_w in T.unroll(num_warps): |
| 1007 | reduced_acc[i_m] += warp_acc[i_w, i_m, lane] |
| 1008 | out_idx = i_t * tile_n + lane |
| 1009 | if out_idx < n_out: |
| 1010 | out[0, token_idx, out_idx] = reduced_acc[i_m] |
| 1011 | if lane == 0 and i_t == 0: |
| 1012 | for i_w in T.unroll(num_warps): |
| 1013 | reduced_sqr[i_m] += warp_sqr[i_w, i_m] |
| 1014 | sqrsum[0, token_idx] = reduced_sqr[i_m] |
| 1015 | |
| 1016 | if ENABLE_PDL: |
| 1017 | T.pdl_trigger() |
| 1018 | |
| 1019 | |
| 1020 | @tilelang_jit |
| 1021 | def hc_head_fuse_tilelang( |
| 1022 | residual, |
| 1023 | fn, |
| 1024 | hc_scale, |
| 1025 | hc_base, |
| 1026 | out, |
| 1027 | hidden_size: int, |
| 1028 | rms_eps: float, |
| 1029 | hc_eps: float, |
| 1030 | hc_mult: int = 4, |
| 1031 | n_thr: int = 128, |
| 1032 | h_blk: int = 1024, |
| 1033 | ): |
| 1034 | """Two-pass fused kernel for hc_head. |
| 1035 | |
| 1036 | Pass 1: accumulate per-token squared sum and hc_mult dot-products |
| 1037 | (projections onto fn rows) using cross-thread reducers. |
| 1038 | Pass 2: apply sigmoid-gated weighted sum of residual channels to output. |
| 1039 | |
| 1040 | Avoids materialising mixes / rsqrt / pre tensors to global memory. |
| 1041 | """ |
| 1042 | num_tokens = T.dynamic("num_tokens") |
| 1043 | hc_dim = hc_mult * hidden_size |
| 1044 | h_block = math.gcd(h_blk, hidden_size) |
| 1045 | n_h = hidden_size // h_block |
| 1046 | |
| 1047 | residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type] |
| 1048 | fn: T.Tensor[[hc_mult, hc_dim], T.float32] # type: ignore[no-redef,valid-type] |
| 1049 | hc_scale: T.Tensor[[1], T.float32] # type: ignore[no-redef,valid-type] |
| 1050 | hc_base: T.Tensor[[hc_mult], T.float32] # type: ignore[no-redef,valid-type] |
| 1051 | out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type] |
| 1052 | |
| 1053 | with T.Kernel(num_tokens, threads=n_thr) as i: |
| 1054 | if ENABLE_PDL: |
| 1055 | T.pdl_sync() |
| 1056 | |
| 1057 | # ------------------------------------------------------------------ |
| 1058 | # Pass 1 – for each residual channel m_c and h_block: |
| 1059 | # • accumulate squared sum (for RMS norm denominator) |
| 1060 | # • accumulate hc_mult dot-products with fn rows |
| 1061 | # ------------------------------------------------------------------ |
| 1062 | sqrsum_r = T.alloc_reducer((1,), T.float32, replication="all") |
| 1063 | mixes_r = T.alloc_reducer((hc_mult,), T.float32, replication="all") |
| 1064 | T.fill(sqrsum_r, 0.0) |
| 1065 | T.fill(mixes_r, 0.0) |
| 1066 | |
| 1067 | for m_c in T.serial(hc_mult): |
| 1068 | for i_h in T.serial(n_h): |
| 1069 | x_local = T.alloc_fragment(h_block, T.float32) |
| 1070 | T.copy(residual[i, m_c, i_h * h_block], x_local) |
| 1071 | |
| 1072 | for k in T.Parallel(h_block): |
| 1073 | sqrsum_r[0] += x_local[k] * x_local[k] |
| 1074 | |
| 1075 | for m_m in T.unroll(hc_mult): |
| 1076 | fn_local = T.alloc_fragment(h_block, T.float32) |
| 1077 | T.copy(fn[m_m, m_c * hidden_size + i_h * h_block], fn_local) |
| 1078 | for k in T.Parallel(h_block): |
| 1079 | mixes_r[m_m] += x_local[k] * fn_local[k] |
| 1080 | |
| 1081 | T.finalize_reducer(sqrsum_r) |
| 1082 | T.finalize_reducer(mixes_r) |
| 1083 | |
| 1084 | # ------------------------------------------------------------------ |
| 1085 | # Compute pre_mix = sigmoid(mix * rsqrt * scale + base) + eps |
| 1086 | # ------------------------------------------------------------------ |
| 1087 | pre_mix_shared = T.alloc_shared(hc_mult, T.float32) |
| 1088 | rsqrt_val = T.alloc_fragment(1, T.float32) |
| 1089 | rsqrt_val[0] = T.rsqrt(sqrsum_r[0] / hc_dim + rms_eps) |
| 1090 | for m in T.Parallel(hc_mult): |
| 1091 | pre_mix_shared[m] = ( |
| 1092 | T.sigmoid(mixes_r[m] * rsqrt_val[0] * hc_scale[0] + hc_base[m]) + hc_eps |
| 1093 | ) |
| 1094 | |
| 1095 | # ------------------------------------------------------------------ |
| 1096 | # Pass 2 – apply_mix: pipelined weighted sum over residual channels |
| 1097 | # ------------------------------------------------------------------ |
| 1098 | for i0_h in T.Pipelined(n_h, num_stages=2): |
| 1099 | xs = T.alloc_shared((hc_mult, h_block), T.bfloat16) |
| 1100 | xl = T.alloc_fragment((hc_mult, h_block), T.float32) |
| 1101 | T.copy(residual[i, 0, i0_h * h_block], xs, disable_tma=True) |
| 1102 | T.copy(xs, xl) |
| 1103 | |
| 1104 | ol = T.alloc_fragment(h_block, T.float32) |
| 1105 | T.clear(ol) |
| 1106 | for i_hc in T.serial(hc_mult): |
| 1107 | pre = pre_mix_shared[i_hc] |
| 1108 | for i1_h in T.Parallel(h_block): |
| 1109 | ol[i1_h] += pre * xl[i_hc, i1_h] |
| 1110 | |
| 1111 | T.copy(ol, out[i, i0_h * h_block], disable_tma=True) |
| 1112 | |
| 1113 | if ENABLE_PDL: |
| 1114 | T.pdl_trigger() |
| 1115 | |
| 1116 | |
| 1117 | class HcPrenormGemmTileLangKernel( |
| 1118 | VllmTileLangJitKernel["HcPrenormGemmTileLangKernel.CompileKey"] |
| 1119 | ): |
| 1120 | @dataclass(frozen=True) |
| 1121 | class CompileKey: |
| 1122 | hidden_size: int |
| 1123 | hc_mult: int |
| 1124 | n_out: int |
| 1125 | n_thr: int |
| 1126 | tile_n: int |
| 1127 | n_splits: int |
| 1128 | use_block_m: bool |
| 1129 | block_m: int |
| 1130 | |
| 1131 | @staticmethod |
| 1132 | def kernel(compile_key: CompileKey) -> Any: |
| 1133 | if compile_key.use_block_m: |
| 1134 | return hc_prenorm_gemm_block_m_tilelang |
| 1135 | return hc_prenorm_gemm_tilelang |
| 1136 | |
| 1137 | def dispatch( # type: ignore[override] |
| 1138 | self, |
| 1139 | *, |
| 1140 | num_tokens: int, |
| 1141 | hc_hidden_size: int, |
| 1142 | n_thr: int = 512, |
| 1143 | tile_n: int = 12, |
| 1144 | n_splits: int = 1, |
| 1145 | **compile_key_fields: int, |
| 1146 | ) -> CompileKey: |
| 1147 | use_default_config = n_splits == 1 and tile_n == 12 and n_thr == 512 |
| 1148 | use_block_m = use_default_config and num_tokens >= 1024 |
| 1149 | use_small_tile = ( |
| 1150 | use_default_config |
| 1151 | and not use_block_m |
| 1152 | and num_tokens < 128 |
| 1153 | and hc_hidden_size % 1024 == 0 |
| 1154 | ) |
| 1155 | effective_n_thr = 1024 if use_small_tile else n_thr |
| 1156 | effective_tile_n = 4 if use_small_tile else tile_n |
| 1157 | return self.CompileKey( |
| 1158 | **compile_key_fields, |
| 1159 | n_thr=effective_n_thr, |
| 1160 | tile_n=effective_tile_n, |
| 1161 | n_splits=n_splits, |
| 1162 | use_block_m=use_block_m, |
| 1163 | block_m=2 if use_block_m else 1, |
| 1164 | ) |
| 1165 | |
| 1166 | def get_warmup_keys( |
| 1167 | self, |
| 1168 | vllm_config: Any, |
| 1169 | *, |
| 1170 | hidden_size: int, |
| 1171 | hc_mult: int, |
| 1172 | n_out: int, |
| 1173 | ) -> list[CompileKey]: |
| 1174 | max_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 1175 | if max_tokens <= 0: |
| 1176 | return [] |
| 1177 | return self._trace_dispatch(self.dispatch)( |
| 1178 | num_tokens=WarmupIntRange( |
| 1179 | 1, |
| 1180 | max_tokens + 1, |
| 1181 | advance=lambda value: ( |
| 1182 | 128 if value < 128 else 1024 if value < 1024 else max_tokens + 1 |
| 1183 | ), |
| 1184 | ), |
| 1185 | hc_hidden_size=hidden_size * hc_mult, |
| 1186 | hidden_size=hidden_size, |
| 1187 | hc_mult=hc_mult, |
| 1188 | n_out=n_out, |
| 1189 | ) |
| 1190 | |
| 1191 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 1192 | num_tokens = 128 |
| 1193 | if compile_key.use_block_m: |
| 1194 | num_tokens = 1024 |
| 1195 | elif compile_key.n_thr == 1024: |
| 1196 | num_tokens = 1 |
| 1197 | hc_hidden_size = compile_key.hidden_size * compile_key.hc_mult |
| 1198 | return dict( |
| 1199 | x=make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hc_hidden_size), |
| 1200 | fn=make_tilelang_warmup_tensor( |
| 1201 | torch.float32, compile_key.n_out, hc_hidden_size |
| 1202 | ), |
| 1203 | out=make_tilelang_warmup_tensor( |
| 1204 | torch.float32, |
| 1205 | compile_key.n_splits, |
| 1206 | num_tokens, |
| 1207 | compile_key.n_out, |
| 1208 | ), |
| 1209 | sqrsum=make_tilelang_warmup_tensor( |
| 1210 | torch.float32, compile_key.n_splits, num_tokens |
| 1211 | ), |
| 1212 | hidden_size=compile_key.hidden_size, |
| 1213 | hc_mult=compile_key.hc_mult, |
| 1214 | tile_n=compile_key.tile_n, |
| 1215 | n_thr=compile_key.n_thr, |
| 1216 | n_splits=compile_key.n_splits, |
| 1217 | ) |
| 1218 | |
| 1219 | @kernel_launcher |
| 1220 | def __call__( |
| 1221 | self, |
| 1222 | x: Any, |
| 1223 | fn: Any, |
| 1224 | out: Any, |
| 1225 | sqrsum: Any, |
| 1226 | hidden_size: int, |
| 1227 | hc_mult: int, |
| 1228 | tile_n: int = 12, |
| 1229 | n_thr: int = 512, |
| 1230 | n_splits: int = 1, |
| 1231 | ) -> TileLangLaunchSpec: |
| 1232 | assert out.shape[0] == n_splits |
| 1233 | assert sqrsum.shape[0] == n_splits |
| 1234 | assert x.shape[1] == hc_mult * hidden_size |
| 1235 | assert x.shape[1] % n_splits == 0 |
| 1236 | assert (x.shape[1] // n_splits) % n_thr == 0 |
| 1237 | compile_key = self.dispatch( |
| 1238 | num_tokens=x.shape[0], |
| 1239 | hc_hidden_size=x.shape[1], |
| 1240 | hidden_size=hidden_size, |
| 1241 | hc_mult=hc_mult, |
| 1242 | n_out=fn.shape[0], |
| 1243 | n_thr=n_thr, |
| 1244 | tile_n=tile_n, |
| 1245 | n_splits=n_splits, |
| 1246 | ) |
| 1247 | kernel_arg = ( |
| 1248 | compile_key.block_m if compile_key.use_block_m else compile_key.n_splits |
| 1249 | ) |
| 1250 | return (compile_key,), ( |
| 1251 | x, |
| 1252 | fn, |
| 1253 | out, |
| 1254 | sqrsum, |
| 1255 | hidden_size, |
| 1256 | hc_mult, |
| 1257 | fn.shape[0], |
| 1258 | compile_key.n_thr, |
| 1259 | compile_key.tile_n, |
| 1260 | kernel_arg, |
| 1261 | ) |
| 1262 | |
| 1263 | |
| 1264 | class MhcPreBigFuseTileLangKernel( |
| 1265 | VllmTileLangJitKernel["MhcPreBigFuseTileLangKernel.CompileKey"] |
| 1266 | ): |
| 1267 | @dataclass(frozen=True) |
| 1268 | class CompileKey: |
| 1269 | hidden_size: int |
| 1270 | hc_mult: int |
| 1271 | n_splits: int |
| 1272 | use_norm_weight: bool |
| 1273 | is_broadcast: bool |
| 1274 | rms_eps: float |
| 1275 | hc_pre_eps: float |
| 1276 | hc_sinkhorn_eps: float |
| 1277 | hc_post_mult_value: float |
| 1278 | sinkhorn_repeat: int |
| 1279 | norm_eps: float |
| 1280 | |
| 1281 | @staticmethod |
| 1282 | def kernel(compile_key: CompileKey) -> Any: |
| 1283 | if compile_key.is_broadcast: |
| 1284 | return mhc_pre_big_fuse_broadcast_with_norm_tilelang |
| 1285 | if compile_key.use_norm_weight: |
| 1286 | return mhc_pre_big_fuse_with_norm_tilelang |
| 1287 | return mhc_pre_big_fuse_tilelang |
| 1288 | |
| 1289 | @staticmethod |
| 1290 | def _kernel_args( |
| 1291 | compile_key: CompileKey, |
| 1292 | *, |
| 1293 | gemm_out_mul: Any, |
| 1294 | gemm_out_sqrsum: Any, |
| 1295 | hc_scale: Any, |
| 1296 | hc_base: Any, |
| 1297 | residual: Any, |
| 1298 | residual_out: Any | None, |
| 1299 | post_mix: Any, |
| 1300 | comb_mix: Any, |
| 1301 | layer_input: Any, |
| 1302 | norm_weight: Any | None, |
| 1303 | ) -> tuple[Any, ...]: |
| 1304 | common_args = ( |
| 1305 | gemm_out_mul, |
| 1306 | gemm_out_sqrsum, |
| 1307 | hc_scale, |
| 1308 | hc_base, |
| 1309 | residual, |
| 1310 | ) |
| 1311 | output_args = (post_mix, comb_mix, layer_input) |
| 1312 | eps_args = ( |
| 1313 | compile_key.rms_eps, |
| 1314 | compile_key.hc_pre_eps, |
| 1315 | compile_key.hc_sinkhorn_eps, |
| 1316 | compile_key.hc_post_mult_value, |
| 1317 | compile_key.sinkhorn_repeat, |
| 1318 | ) |
| 1319 | middle_args: tuple[Any, ...] |
| 1320 | if compile_key.is_broadcast: |
| 1321 | assert residual_out is not None |
| 1322 | assert norm_weight is not None |
| 1323 | middle_args = (residual_out, *output_args, norm_weight) |
| 1324 | elif compile_key.use_norm_weight: |
| 1325 | assert norm_weight is not None |
| 1326 | # layer_input stands in for the aux buffer on both epilogues: same |
| 1327 | # shape and dtype, and write_aux is off here so it is never written. |
| 1328 | middle_args = ( |
| 1329 | *output_args, |
| 1330 | norm_weight, |
| 1331 | post_mix, |
| 1332 | post_mix, |
| 1333 | layer_input, |
| 1334 | ) |
| 1335 | else: |
| 1336 | middle_args = (*output_args, post_mix, post_mix, layer_input) |
| 1337 | norm_eps_args = (compile_key.norm_eps,) if compile_key.use_norm_weight else () |
| 1338 | return ( |
| 1339 | *common_args, |
| 1340 | *middle_args, |
| 1341 | compile_key.hidden_size, |
| 1342 | *eps_args, |
| 1343 | *norm_eps_args, |
| 1344 | compile_key.n_splits, |
| 1345 | compile_key.hc_mult, |
| 1346 | ) |
| 1347 | |
| 1348 | def dispatch( # type: ignore[override] |
| 1349 | self, |
| 1350 | *, |
| 1351 | hidden_size: int, |
| 1352 | hc_mult: int, |
| 1353 | n_splits: int, |
| 1354 | is_broadcast: bool, |
| 1355 | use_norm_weight: bool, |
| 1356 | use_pre_gemm_splits: bool = False, |
| 1357 | pre_gemm_k: int = 0, |
| 1358 | sinkhorn_repeat: int, |
| 1359 | norm_eps: float, |
| 1360 | broadcast_norm_eps: float = 0.0, |
| 1361 | num_tokens: int = 1, |
| 1362 | use_fused_tilelang: bool = False, |
| 1363 | **compile_key_fields: float, |
| 1364 | ) -> CompileKey: |
| 1365 | pre_gemm_n_splits = ( |
| 1366 | compute_num_split(64, pre_gemm_k, (num_tokens + 63) // 64) |
| 1367 | if use_pre_gemm_splits |
| 1368 | else n_splits |
| 1369 | ) |
| 1370 | # The epilogue reduces over whatever the fused kernel split the GEMM |
| 1371 | # into, so both have to read the same config. |
| 1372 | fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) |
| 1373 | actual_n_splits = ( |
| 1374 | fused_config[1] |
| 1375 | if use_fused_tilelang and fused_config is not None |
| 1376 | else pre_gemm_n_splits |
| 1377 | ) |
| 1378 | actual_norm_eps = broadcast_norm_eps if is_broadcast else norm_eps |
| 1379 | actual_use_norm_weight = use_norm_weight or is_broadcast |
| 1380 | return self.CompileKey( |
| 1381 | **compile_key_fields, |
| 1382 | hidden_size=hidden_size, |
| 1383 | hc_mult=hc_mult, |
| 1384 | n_splits=actual_n_splits, |
| 1385 | use_norm_weight=actual_use_norm_weight, |
| 1386 | is_broadcast=is_broadcast, |
| 1387 | sinkhorn_repeat=sinkhorn_repeat, |
| 1388 | norm_eps=actual_norm_eps if actual_use_norm_weight else 0.0, |
| 1389 | ) |
| 1390 | |
| 1391 | def get_warmup_keys( |
| 1392 | self, |
| 1393 | vllm_config: Any, |
| 1394 | *, |
| 1395 | hidden_size: int, |
| 1396 | hc_mult: int, |
| 1397 | use_norm_weight: bool, |
| 1398 | rms_eps: float, |
| 1399 | hc_pre_eps: float, |
| 1400 | hc_sinkhorn_eps: float, |
| 1401 | hc_post_mult_value: float, |
| 1402 | sinkhorn_repeat: int, |
| 1403 | norm_eps: float | tuple[float, ...], |
| 1404 | broadcast_norm_eps: float | tuple[float, ...] | None = None, |
| 1405 | include_pre_gemm_splits: bool = False, |
| 1406 | include_broadcast_splits: bool = False, |
| 1407 | ) -> list[CompileKey]: |
| 1408 | max_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 1409 | if max_tokens <= 0: |
| 1410 | return [] |
| 1411 | |
| 1412 | warmup_cases = zip_inputs( |
| 1413 | dict( |
| 1414 | is_broadcast=False, |
| 1415 | use_pre_gemm_splits=False, |
| 1416 | pre_gemm_k=0, |
| 1417 | use_fused_tilelang=False, |
| 1418 | ), |
| 1419 | dict( |
| 1420 | is_broadcast=False, |
| 1421 | use_pre_gemm_splits=include_pre_gemm_splits, |
| 1422 | pre_gemm_k=hc_mult * hidden_size, |
| 1423 | use_fused_tilelang=False, |
| 1424 | ), |
| 1425 | dict( |
| 1426 | is_broadcast=include_broadcast_splits, |
| 1427 | use_pre_gemm_splits=include_broadcast_splits, |
| 1428 | pre_gemm_k=hidden_size, |
| 1429 | use_fused_tilelang=False, |
| 1430 | ), |
| 1431 | dict( |
| 1432 | is_broadcast=False, |
| 1433 | use_pre_gemm_splits=False, |
| 1434 | pre_gemm_k=0, |
| 1435 | use_fused_tilelang=True, |
| 1436 | ), |
| 1437 | ) |
| 1438 | return self._trace_dispatch(self.dispatch)( |
| 1439 | warmup_cases, |
| 1440 | num_tokens=WarmupIntRange( |
| 1441 | 1, |
| 1442 | max_tokens + 1, |
| 1443 | advance=lambda value: 8 if value < 8 else cdiv(value, 64) * 64 + 1, |
| 1444 | ), |
| 1445 | n_splits=1, |
| 1446 | hidden_size=hidden_size, |
| 1447 | hc_mult=hc_mult, |
| 1448 | use_norm_weight=use_norm_weight, |
| 1449 | rms_eps=rms_eps, |
| 1450 | hc_pre_eps=hc_pre_eps, |
| 1451 | hc_sinkhorn_eps=hc_sinkhorn_eps, |
| 1452 | hc_post_mult_value=hc_post_mult_value, |
| 1453 | sinkhorn_repeat=sinkhorn_repeat, |
| 1454 | norm_eps=norm_eps, |
| 1455 | broadcast_norm_eps=broadcast_norm_eps |
| 1456 | if broadcast_norm_eps is not None |
| 1457 | else norm_eps, |
| 1458 | _when=lambda *, use_fused_tilelang, num_tokens: ( |
| 1459 | not use_fused_tilelang or num_tokens <= 16 |
| 1460 | ), |
| 1461 | ) |
| 1462 | |
| 1463 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 1464 | hidden_size = compile_key.hidden_size |
| 1465 | hc_mult = compile_key.hc_mult |
| 1466 | hc_mult3 = hc_mult * (2 + hc_mult) |
| 1467 | num_tokens = 1 |
| 1468 | |
| 1469 | gemm_out_mul = make_tilelang_warmup_tensor( |
| 1470 | torch.float32, compile_key.n_splits, num_tokens, hc_mult3 |
| 1471 | ) |
| 1472 | gemm_out_sqrsum = make_tilelang_warmup_tensor( |
| 1473 | torch.float32, compile_key.n_splits, num_tokens |
| 1474 | ) |
| 1475 | hc_scale = make_tilelang_warmup_tensor(torch.float32, 3) |
| 1476 | hc_base = make_tilelang_warmup_tensor(torch.float32, hc_mult3) |
| 1477 | post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult) |
| 1478 | comb_mix = make_tilelang_warmup_tensor( |
| 1479 | torch.float32, num_tokens, hc_mult * hc_mult |
| 1480 | ) |
| 1481 | layer_input = make_tilelang_warmup_tensor( |
| 1482 | torch.bfloat16, num_tokens, hidden_size |
| 1483 | ) |
| 1484 | residual_shape = ( |
| 1485 | (num_tokens, hidden_size) |
| 1486 | if compile_key.is_broadcast |
| 1487 | else (num_tokens, hc_mult, hidden_size) |
| 1488 | ) |
| 1489 | residual = make_tilelang_warmup_tensor(torch.bfloat16, *residual_shape) |
| 1490 | residual_out = ( |
| 1491 | make_tilelang_warmup_tensor( |
| 1492 | torch.bfloat16, num_tokens, hc_mult, hidden_size |
| 1493 | ) |
| 1494 | if compile_key.is_broadcast |
| 1495 | else None |
| 1496 | ) |
| 1497 | norm_weight = ( |
| 1498 | make_tilelang_warmup_tensor(torch.bfloat16, hidden_size) |
| 1499 | if compile_key.is_broadcast or compile_key.use_norm_weight |
| 1500 | else None |
| 1501 | ) |
| 1502 | return dict( |
| 1503 | gemm_out_mul=gemm_out_mul, |
| 1504 | gemm_out_sqrsum=gemm_out_sqrsum, |
| 1505 | hc_scale=hc_scale, |
| 1506 | hc_base=hc_base, |
| 1507 | residual=residual, |
| 1508 | post_mix=post_mix, |
| 1509 | comb_mix=comb_mix, |
| 1510 | layer_input=layer_input, |
| 1511 | rms_eps=compile_key.rms_eps, |
| 1512 | hc_pre_eps=compile_key.hc_pre_eps, |
| 1513 | hc_sinkhorn_eps=compile_key.hc_sinkhorn_eps, |
| 1514 | hc_post_mult_value=compile_key.hc_post_mult_value, |
| 1515 | sinkhorn_repeat=compile_key.sinkhorn_repeat, |
| 1516 | residual_out=residual_out, |
| 1517 | norm_weight=norm_weight, |
| 1518 | norm_eps=compile_key.norm_eps, |
| 1519 | ) |
| 1520 | |
| 1521 | @kernel_launcher |
| 1522 | def __call__( |
| 1523 | self, |
| 1524 | gemm_out_mul: Any, |
| 1525 | gemm_out_sqrsum: Any, |
| 1526 | hc_scale: Any, |
| 1527 | hc_base: Any, |
| 1528 | residual: Any, |
| 1529 | post_mix: Any, |
| 1530 | comb_mix: Any, |
| 1531 | layer_input: Any, |
| 1532 | rms_eps: float, |
| 1533 | hc_pre_eps: float, |
| 1534 | hc_sinkhorn_eps: float, |
| 1535 | hc_post_mult_value: float, |
| 1536 | sinkhorn_repeat: int, |
| 1537 | *, |
| 1538 | residual_out: Any | None = None, |
| 1539 | norm_weight: Any | None = None, |
| 1540 | norm_eps: float = 0.0, |
| 1541 | ) -> TileLangLaunchSpec: |
| 1542 | is_broadcast = residual_out is not None |
| 1543 | use_norm_weight = norm_weight is not None |
| 1544 | n_splits = gemm_out_mul.shape[0] |
| 1545 | if is_broadcast: |
| 1546 | assert residual_out is not None |
| 1547 | hidden_size = residual.shape[-1] |
| 1548 | hc_mult = residual_out.shape[-2] |
| 1549 | else: |
| 1550 | hc_mult = residual.shape[-2] |
| 1551 | hidden_size = residual.shape[-1] |
| 1552 | compile_key = self.dispatch( |
| 1553 | hidden_size=hidden_size, |
| 1554 | hc_mult=hc_mult, |
| 1555 | n_splits=n_splits, |
| 1556 | is_broadcast=is_broadcast, |
| 1557 | use_norm_weight=use_norm_weight, |
| 1558 | rms_eps=rms_eps, |
| 1559 | hc_pre_eps=hc_pre_eps, |
| 1560 | hc_sinkhorn_eps=hc_sinkhorn_eps, |
| 1561 | hc_post_mult_value=hc_post_mult_value, |
| 1562 | sinkhorn_repeat=sinkhorn_repeat, |
| 1563 | norm_eps=norm_eps, |
| 1564 | broadcast_norm_eps=norm_eps, |
| 1565 | ) |
| 1566 | return ( |
| 1567 | (compile_key,), |
| 1568 | self._kernel_args( |
| 1569 | compile_key, |
| 1570 | gemm_out_mul=gemm_out_mul, |
| 1571 | gemm_out_sqrsum=gemm_out_sqrsum, |
| 1572 | hc_scale=hc_scale, |
| 1573 | hc_base=hc_base, |
| 1574 | residual=residual, |
| 1575 | residual_out=residual_out, |
| 1576 | post_mix=post_mix, |
| 1577 | comb_mix=comb_mix, |
| 1578 | layer_input=layer_input, |
| 1579 | norm_weight=norm_weight, |
| 1580 | ), |
| 1581 | ) |
| 1582 | |
| 1583 | |
| 1584 | class MhcPostTileLangKernel(VllmTileLangJitKernel["MhcPostTileLangKernel.CompileKey"]): |
| 1585 | @dataclass(frozen=True) |
| 1586 | class CompileKey: |
| 1587 | hidden_size: int |
| 1588 | hc_mult: int |
| 1589 | |
| 1590 | @staticmethod |
| 1591 | def kernel() -> Any: |
| 1592 | return mhc_post_tilelang |
| 1593 | |
| 1594 | def dispatch( # type: ignore[override] |
| 1595 | self, |
| 1596 | **compile_key_fields: int, |
| 1597 | ) -> CompileKey: |
| 1598 | return self.CompileKey(**compile_key_fields) |
| 1599 | |
| 1600 | def get_warmup_keys( |
| 1601 | self, |
| 1602 | *, |
| 1603 | hidden_size: int, |
| 1604 | hc_mult: int, |
| 1605 | ) -> list[CompileKey]: |
| 1606 | return self._trace_dispatch(self.dispatch)( |
| 1607 | hidden_size=hidden_size, |
| 1608 | hc_mult=hc_mult, |
| 1609 | ) |
| 1610 | |
| 1611 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 1612 | num_tokens = 1 |
| 1613 | hidden_size = compile_key.hidden_size |
| 1614 | hc_mult = compile_key.hc_mult |
| 1615 | comb_mix = make_tilelang_warmup_tensor( |
| 1616 | torch.float32, num_tokens, hc_mult, hc_mult |
| 1617 | ) |
| 1618 | residual = make_tilelang_warmup_tensor( |
| 1619 | torch.bfloat16, num_tokens, hc_mult, hidden_size |
| 1620 | ) |
| 1621 | post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult) |
| 1622 | layer_input = make_tilelang_warmup_tensor( |
| 1623 | torch.bfloat16, num_tokens, hidden_size |
| 1624 | ) |
| 1625 | out = make_tilelang_warmup_tensor( |
| 1626 | torch.bfloat16, num_tokens, hc_mult, hidden_size |
| 1627 | ) |
| 1628 | return dict( |
| 1629 | comb_mix=comb_mix, |
| 1630 | residual=residual, |
| 1631 | post_mix=post_mix, |
| 1632 | layer_input=layer_input, |
| 1633 | out=out, |
| 1634 | hc_mult=hc_mult, |
| 1635 | hidden_size=hidden_size, |
| 1636 | ) |
| 1637 | |
| 1638 | @kernel_launcher |
| 1639 | def __call__( |
| 1640 | self, |
| 1641 | comb_mix: Any, |
| 1642 | residual: Any, |
| 1643 | post_mix: Any, |
| 1644 | layer_input: Any, |
| 1645 | out: Any, |
| 1646 | hc_mult: int, |
| 1647 | hidden_size: int, |
| 1648 | ) -> TileLangLaunchSpec: |
| 1649 | return (), ( |
| 1650 | comb_mix, |
| 1651 | residual, |
| 1652 | post_mix, |
| 1653 | layer_input, |
| 1654 | out, |
| 1655 | hc_mult, |
| 1656 | hidden_size, |
| 1657 | ) |
| 1658 | |
| 1659 | |
| 1660 | class MhcFusedTileLangKernel( |
| 1661 | VllmTileLangJitKernel["MhcFusedTileLangKernel.CompileKey"] |
| 1662 | ): |
| 1663 | @dataclass(frozen=True) |
| 1664 | class CompileKey: |
| 1665 | hidden_size: int |
| 1666 | hc_mult: int |
| 1667 | n_splits: int |
| 1668 | tile_n: int |
| 1669 | n_thr: int |
| 1670 | |
| 1671 | @staticmethod |
| 1672 | def kernel() -> Any: |
| 1673 | return mhc_fused_tilelang |
| 1674 | |
| 1675 | def dispatch( # type: ignore[override] |
| 1676 | self, |
| 1677 | *, |
| 1678 | num_tokens: int, |
| 1679 | hidden_size: int, |
| 1680 | hc_mult: int, |
| 1681 | ) -> CompileKey: |
| 1682 | # Subscripts rather than an unpack: the warmup tracer parses this body |
| 1683 | # and allows only plain assignments before the return. |
| 1684 | config = require_fused_post_pre_config(num_tokens, hidden_size, hc_mult) |
| 1685 | return self.CompileKey( |
| 1686 | hidden_size=hidden_size, |
| 1687 | hc_mult=hc_mult, |
| 1688 | n_splits=config[1], |
| 1689 | tile_n=config[0], |
| 1690 | n_thr=config[2], |
| 1691 | ) |
| 1692 | |
| 1693 | def get_warmup_keys( |
| 1694 | self, |
| 1695 | vllm_config: Any, |
| 1696 | *, |
| 1697 | hidden_size: int, |
| 1698 | hc_mult: int, |
| 1699 | ) -> list[CompileKey]: |
| 1700 | max_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 1701 | return self._trace_dispatch(self.dispatch)( |
| 1702 | num_tokens=WarmupIntRange(1, max_tokens + 1), |
| 1703 | hidden_size=hidden_size, |
| 1704 | hc_mult=hc_mult, |
| 1705 | _when=lambda *, num_tokens: num_tokens <= _FUSED_POST_PRE_MAX_TOKENS, |
| 1706 | ) |
| 1707 | |
| 1708 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 1709 | num_tokens = 1 if compile_key.tile_n == 2 else 16 |
| 1710 | hidden_size = compile_key.hidden_size |
| 1711 | hc_mult = compile_key.hc_mult |
| 1712 | hc_mult3 = hc_mult * (2 + hc_mult) |
| 1713 | comb_mix = make_tilelang_warmup_tensor( |
| 1714 | torch.float32, num_tokens, hc_mult, hc_mult |
| 1715 | ) |
| 1716 | residual_in = make_tilelang_warmup_tensor( |
| 1717 | torch.bfloat16, num_tokens, hc_mult, hidden_size |
| 1718 | ) |
| 1719 | post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult) |
| 1720 | x_in = make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hidden_size) |
| 1721 | weight_t = make_tilelang_warmup_tensor( |
| 1722 | torch.float32, hc_mult3, hc_mult, hidden_size |
| 1723 | ) |
| 1724 | return dict( |
| 1725 | comb_mix=comb_mix, |
| 1726 | residual_in=residual_in, |
| 1727 | post_mix=post_mix, |
| 1728 | x_in=x_in, |
| 1729 | weight_t=weight_t, |
| 1730 | hc_mult=hc_mult, |
| 1731 | hidden_size=hidden_size, |
| 1732 | hc_mult3=hc_mult3, |
| 1733 | ) |
| 1734 | |
| 1735 | @kernel_launcher |
| 1736 | def __call__( |
| 1737 | self, |
| 1738 | comb_mix: Any, |
| 1739 | residual_in: Any, |
| 1740 | post_mix: Any, |
| 1741 | x_in: Any, |
| 1742 | weight_t: Any, |
| 1743 | hc_mult: int, |
| 1744 | hidden_size: int, |
| 1745 | hc_mult3: int, |
| 1746 | ) -> TileLangLaunchSpec: |
| 1747 | num_tokens = residual_in.shape[0] |
| 1748 | tile_n, n_splits, n_thr = require_fused_post_pre_config( |
| 1749 | num_tokens, hidden_size, hc_mult |
| 1750 | ) |
| 1751 | yp_out = residual_in.new_empty( |
| 1752 | (n_splits, num_tokens, hc_mult3), dtype=torch.float32 |
| 1753 | ) |
| 1754 | rp_out = residual_in.new_empty((n_splits, num_tokens), dtype=torch.float32) |
| 1755 | residual_out = residual_in.new_empty(residual_in.shape) |
| 1756 | return ( |
| 1757 | (), |
| 1758 | ( |
| 1759 | comb_mix, |
| 1760 | residual_in, |
| 1761 | post_mix, |
| 1762 | x_in, |
| 1763 | weight_t, |
| 1764 | yp_out, |
| 1765 | rp_out, |
| 1766 | residual_out, |
| 1767 | hc_mult, |
| 1768 | hidden_size, |
| 1769 | hc_mult3, |
| 1770 | ), |
| 1771 | dict(n_thr=n_thr, tile_n=tile_n), |
| 1772 | (yp_out, rp_out, residual_out), |
| 1773 | ) |
| 1774 | |
| 1775 | |
| 1776 | class HcHeadFusedTileLangKernel( |
| 1777 | VllmTileLangJitKernel["HcHeadFusedTileLangKernel.CompileKey"] |
| 1778 | ): |
| 1779 | @dataclass(frozen=True) |
| 1780 | class CompileKey: |
| 1781 | hidden_size: int |
| 1782 | hc_mult: int |
| 1783 | rms_eps: float |
| 1784 | hc_eps: float |
| 1785 | |
| 1786 | @staticmethod |
| 1787 | def kernel() -> Any: |
| 1788 | return hc_head_fuse_tilelang |
| 1789 | |
| 1790 | def dispatch( # type: ignore[override] |
| 1791 | self, |
| 1792 | *, |
| 1793 | hidden_size: int, |
| 1794 | hc_mult: int, |
| 1795 | rms_eps: float, |
| 1796 | hc_eps: float, |
| 1797 | ) -> CompileKey: |
| 1798 | return self.CompileKey( |
| 1799 | hidden_size=hidden_size, |
| 1800 | hc_mult=hc_mult, |
| 1801 | rms_eps=rms_eps, |
| 1802 | hc_eps=hc_eps, |
| 1803 | ) |
| 1804 | |
| 1805 | def get_warmup_keys( |
| 1806 | self, |
| 1807 | *, |
| 1808 | hidden_size: int, |
| 1809 | hc_mult: int, |
| 1810 | rms_eps: float, |
| 1811 | hc_eps: float, |
| 1812 | ) -> list[CompileKey]: |
| 1813 | return self._trace_dispatch(self.dispatch)( |
| 1814 | hidden_size=hidden_size, |
| 1815 | hc_mult=hc_mult, |
| 1816 | rms_eps=rms_eps, |
| 1817 | hc_eps=hc_eps, |
| 1818 | ) |
| 1819 | |
| 1820 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 1821 | num_tokens = 1 |
| 1822 | hidden_size = compile_key.hidden_size |
| 1823 | hc_mult = compile_key.hc_mult |
| 1824 | residual = make_tilelang_warmup_tensor( |
| 1825 | torch.bfloat16, num_tokens, hc_mult, hidden_size |
| 1826 | ) |
| 1827 | fn = make_tilelang_warmup_tensor(torch.float32, hc_mult, hc_mult * hidden_size) |
| 1828 | hc_scale = make_tilelang_warmup_tensor(torch.float32, 1) |
| 1829 | hc_base = make_tilelang_warmup_tensor(torch.float32, hc_mult) |
| 1830 | out = make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hidden_size) |
| 1831 | return dict( |
| 1832 | residual=residual, |
| 1833 | fn=fn, |
| 1834 | hc_scale=hc_scale, |
| 1835 | hc_base=hc_base, |
| 1836 | out=out, |
| 1837 | hidden_size=hidden_size, |
| 1838 | rms_eps=compile_key.rms_eps, |
| 1839 | hc_eps=compile_key.hc_eps, |
| 1840 | hc_mult=hc_mult, |
| 1841 | ) |
| 1842 | |
| 1843 | @kernel_launcher |
| 1844 | def __call__( |
| 1845 | self, |
| 1846 | residual: Any, |
| 1847 | fn: Any, |
| 1848 | hc_scale: Any, |
| 1849 | hc_base: Any, |
| 1850 | out: Any, |
| 1851 | hidden_size: int, |
| 1852 | rms_eps: float, |
| 1853 | hc_eps: float, |
| 1854 | hc_mult: int, |
| 1855 | ) -> TileLangLaunchSpec: |
| 1856 | return (), ( |
| 1857 | residual, |
| 1858 | fn, |
| 1859 | hc_scale, |
| 1860 | hc_base, |
| 1861 | out, |
| 1862 | hidden_size, |
| 1863 | rms_eps, |
| 1864 | hc_eps, |
| 1865 | hc_mult, |
| 1866 | ) |
| 1867 | |
| 1868 | |
| 1869 | _HC_PRENORM_GEMM_TILELANG_KERNEL = HcPrenormGemmTileLangKernel() |
| 1870 | _MHC_PRE_BIG_FUSE_TILELANG_KERNEL = MhcPreBigFuseTileLangKernel() |
| 1871 | _MHC_POST_TILELANG_KERNEL = MhcPostTileLangKernel() |
| 1872 | _MHC_FUSED_TILELANG_KERNEL = MhcFusedTileLangKernel() |
| 1873 | _HC_HEAD_FUSED_TILELANG_KERNEL = HcHeadFusedTileLangKernel() |