# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from __future__ import annotations import math from dataclasses import dataclass from functools import cache from typing import Any import torch from vllm.model_executor.warmup.jit_warmup import WarmupIntRange, zip_inputs from vllm.model_executor.warmup.jit_warmup_tilelang_helper import ( TileLangLaunchSpec, VllmTileLangJitKernel, kernel_launcher, make_tilelang_warmup_tensor, ) from vllm.platforms import current_platform from vllm.tilelang_utils import T, tilelang, tilelang_jit from vllm.utils.math_utils import cdiv ENABLE_PDL = current_platform.is_arch_support_pdl() and current_platform.is_cuda() # TileLang's HIP AllReduce shuffle phase ignores a nonzero reduction-group # offset, so the 64-thread RMSNorm reduction must be wave64-aligned and use # threads 0..63. _MIX_ON_TRAILING_THREADS = current_platform.is_rocm() def _is_mix_thread(): if _MIX_ON_TRAILING_THREADS: return T.get_thread_binding() >= 64 return T.get_thread_binding() < 32 def _is_collapse_thread(): if _MIX_ON_TRAILING_THREADS: return T.get_thread_binding() < 64 return T.get_thread_binding() >= 32 @cache def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: device_props = torch.cuda.get_device_properties(0) n_sms = device_props.multi_processor_count split_k = n_sms // grid_size if k is not None: # avoid split_k for small k num_block_k = cdiv(k, block_k) split_k = min(split_k, num_block_k // 4) split_k = max(split_k, 1) return split_k # ``mhc_fused_tilelang`` splits the hidden size across a fixed thread block, so # every (tile_n, n_thr) pair is one more kernel to compile. Tuned on GB300 # (hc_mult 4, hidden_size 5120 and 7168) against the separate post + split-k # GEMM it replaces; the bands are coarse because the loss against a per-shape # optimum is ~1-2%, while one fixed config gives up to 18% at the top of the # range. _FUSED_POST_PRE_N_THR = 128 _FUSED_POST_PRE_MAX_TOKENS = int(__import__("os").environ.get("VLLM_MHC_FUSED_MAX_TOKENS", "32")) def mhc_fused_post_pre_split_config( num_tokens: int, hidden_size: int, hc_mult: int ) -> tuple[int, int, int] | None: """Pick ``(tile_n, n_splits, n_thr)`` for the fused post + pre-norm GEMM. Returns None when a separate post kernel followed by a split-k GEMM is the better choice, or when the hidden size does not divide evenly across the fused kernel's block. One source of truth for the dispatch decision, the compile key and the launch, which must agree. """ if num_tokens > _FUSED_POST_PRE_MAX_TOKENS: return None n_thr = _FUSED_POST_PRE_N_THR if hidden_size % n_thr: return None # The projection tiles supply the parallelism a small token count cannot, # so the tile grows with the batch: 1.3x at one token and 1.03x at 32 # against a split-k GEMM given the split its own estimator computes. if num_tokens < 16: tile_n, n_splits = 2, 8 else: tile_n, n_splits = 6, 8 # The kernel drops projection tiles and k-slices it cannot fill evenly. mix_size = hc_mult * (hc_mult + 2) while tile_n > 1 and mix_size % tile_n: tile_n -= 1 while n_splits > 1 and hidden_size % (n_splits * n_thr): n_splits //= 2 return tile_n, n_splits, n_thr def require_fused_post_pre_config( num_tokens: int, hidden_size: int, hc_mult: int ) -> tuple[int, int, int]: """The config for a shape the caller has already committed to fusing.""" config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) if config is None: raise ValueError( "the fused mHC post + pre-norm GEMM does not cover num_tokens=" f"{num_tokens} at hidden_size={hidden_size}" ) return config def mhc_fused_post_pre_splits(hidden_size: int, hc_mult: int) -> tuple[int, ...]: """Every split-k factor the fused post + pre-norm GEMM path can pick.""" configs = ( mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) for num_tokens in range(1, _FUSED_POST_PRE_MAX_TOKENS + 1) ) return tuple(sorted({config[1] for config in configs if config is not None})) @tilelang_jit def mhc_pre_big_fuse_tilelang( gemm_out_mul, gemm_out_sqrsum, hc_scale, hc_base, residual, post_mix, comb_mix, layer_input, pre_mix_in, pre_mix_out, aux_out, hidden_size: int, rms_eps: float, hc_pre_eps: float, hc_sinkhorn_eps: float, hc_post_mult_value: float, sinkhorn_repeat: int, n_splits: int = 16, hc_mult: int = 4, use_pre_mix_in: bool = False, save_pre_mix: bool = False, rms_numel: int = 0, write_aux: bool = False, ): """Fuse coefficient generation and residual collapse after the projection. With save_pre_mix, store the new pre-mix and collapse with pre_mix_in, or select stream zero when use_pre_mix_in is false. """ num_tokens = T.dynamic("num_tokens") hc_mult3 = hc_mult * (2 + hc_mult) hidden_block = math.gcd(512, hidden_size) if rms_numel == 0: rms_numel = hc_mult * hidden_size gemm_out_mul: T.Tensor[[n_splits, num_tokens, hc_mult3], T.float32] # type: ignore[no-redef, valid-type] gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32] # type: ignore[no-redef, valid-type] hc_scale: T.Tensor[[3], T.float32] # type: ignore[no-redef, valid-type] hc_base: T.Tensor[[hc_mult3], T.float32] # type: ignore[no-redef, valid-type] residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] # outputs post_mix: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32] # type: ignore[no-redef, valid-type] layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] pre_mix_in: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] pre_mix_out: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] aux_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] with T.Kernel(num_tokens, threads=96) as i: if ENABLE_PDL: T.pdl_sync() ################################################################## # _pre_norm_fn_fwd_norm rms = T.alloc_fragment(1, T.float32) mixes = T.alloc_fragment(hc_mult3, T.float32) T.clear(mixes) rms[0] = 0 for i_split in T.serial(n_splits): rms[0] += gemm_out_sqrsum[i_split, i] rms[0] = T.rsqrt(rms[0] / rms_numel + rms_eps) for j in T.Parallel(hc_mult3): mixes[j] = 0 for i_split in T.serial(n_splits): mixes[j] += gemm_out_mul[i_split, i, j] mixes[j] *= rms[0] mixes_shared = T.alloc_shared(hc_mult3, T.float32) T.copy(mixes, mixes_shared) if T.get_thread_binding() < 32: ################################################################## # _pre_split_mixes_fwd (post & comb) cm = T.alloc_fragment((hc_mult, hc_mult), T.float32) for j in T.Parallel(hc_mult): if save_pre_mix: pre_mix_out[i, j] = ( T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j]) + hc_pre_eps ) post_mix[i, j] = ( T.sigmoid( mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult] ) * hc_post_mult_value ) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = ( mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2] + hc_base[j * hc_mult + k + hc_mult * 2] ) ################################################################## # _sinkhorn_fwd row_sum = T.alloc_fragment(hc_mult, T.float32) col_sum = T.alloc_fragment(hc_mult, T.float32) # comb = comb.softmax(-1) + eps row_max = T.alloc_fragment(hc_mult, T.float32) T.reduce_max(cm, row_max, dim=1) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = T.exp(cm[j, k] - row_max[j]) T.reduce_sum(cm, row_sum, dim=1) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps # comb = comb / (comb.sum(-2) + eps) T.reduce_sum(cm, col_sum, dim=0) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps) for _ in T.serial(sinkhorn_repeat - 1): # comb = comb / (comb.sum(-1) + eps) T.reduce_sum(cm, row_sum, dim=1) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps) # comb = comb / (comb.sum(-2) + eps) T.reduce_sum(cm, col_sum, dim=0) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps) # save comb_mix to global memory for j, k in T.Parallel(hc_mult, hc_mult): comb_mix[i, j * hc_mult + k] = cm[j, k] else: ################################################################## # _pre_split_mixes_fwd (pre) pre_mix_shared = T.alloc_shared(hc_mult, T.float32) for j in T.Parallel(hc_mult): if use_pre_mix_in: pre_mix_shared[j] = pre_mix_in[i, j] elif save_pre_mix: pre_mix_shared[j] = T.if_then_else(j == 0, 1.0, 0.0) else: pre_mix_shared[j] = ( T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j]) + hc_pre_eps ) ################################################################### # _pre_apply_mix_fwd for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2): xs = T.alloc_shared((hc_mult, hidden_block), T.float32) xl = T.alloc_fragment((hc_mult, hidden_block), T.float32) T.copy(residual[i, 0, i0_h * hidden_block], xs) T.copy(xs, xl) ol = T.alloc_fragment(hidden_block, T.float32) T.clear(ol) if write_aux: # Aux consumers (draft models) take the plain stream mean # of the same residual this collapse reads. aux = T.alloc_fragment(hidden_block, T.float32) T.clear(aux) for i_hc in T.serial(hc_mult): pre = pre_mix_shared[i_hc] for i1_h in T.Parallel(hidden_block): ol[i1_h] += pre * xl[i_hc, i1_h] if write_aux: aux[i1_h] += xl[i_hc, i1_h] if write_aux: for i1_h in T.Parallel(hidden_block): aux_out[i, i0_h * hidden_block + i1_h] = T.bfloat16( aux[i1_h] / hc_mult ) T.copy(ol, layer_input[i, i0_h * hidden_block]) if ENABLE_PDL: T.pdl_trigger() # Copied from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/mhc.py#L478 @tilelang_jit def mhc_pre_big_fuse_with_norm_tilelang( gemm_out_mul, gemm_out_sqrsum, hc_scale, hc_base, residual, post_mix, comb_mix, layer_input, norm_weight, pre_mix_in, pre_mix_out, aux_out, hidden_size: int, rms_eps: float, hc_pre_eps: float, hc_sinkhorn_eps: float, hc_post_mult_value: float, sinkhorn_repeat: int, norm_eps: float, n_splits: int = 16, hc_mult: int = 4, gemm_last_dim: int = -1, use_pre_mix_in: bool = False, save_pre_mix: bool = False, rms_numel: int = 0, split_mode: str = "fused", write_aux: bool = False, ): # Split modes require shifted mHC: input collapse uses a carried pre-mix. assert split_mode in ("fused", "stats", "input") assert split_mode == "fused" or save_pre_mix num_tokens = T.dynamic("num_tokens") hc_mult3 = hc_mult * (2 + hc_mult) if gemm_last_dim < 0: gemm_last_dim = hc_mult3 if rms_numel == 0: rms_numel = hc_mult * hidden_size hidden_block = math.gcd(1024, hidden_size) gemm_out_mul: T.Tensor[[n_splits, num_tokens, gemm_last_dim], T.float32] # type: ignore[no-redef, valid-type] gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32] # type: ignore[no-redef, valid-type] hc_scale: T.Tensor[[3], T.float32] # type: ignore[no-redef, valid-type] hc_base: T.Tensor[[hc_mult3], T.float32] # type: ignore[no-redef, valid-type] residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] post_mix: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32] # type: ignore[no-redef, valid-type] layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] norm_weight: T.Tensor[[hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] pre_mix_in: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] pre_mix_out: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] aux_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] with T.Kernel(num_tokens, threads=96) as i: if ENABLE_PDL: T.pdl_sync() mixes_shared = T.alloc_shared(hc_mult3, T.float32) if split_mode != "input": rms = T.alloc_fragment(1, T.float32) mixes = T.alloc_fragment(hc_mult3, T.float32) T.clear(mixes) rms[0] = 0 for i_split in T.serial(n_splits): rms[0] += gemm_out_sqrsum[i_split, i] rms[0] = T.rsqrt(rms[0] / rms_numel + rms_eps) for j in T.Parallel(hc_mult3): mixes[j] = 0 for i_split in T.serial(n_splits): mixes[j] += gemm_out_mul[i_split, i, j] mixes[j] *= rms[0] T.copy(mixes, mixes_shared) if split_mode != "input" and _is_mix_thread(): cm = T.alloc_fragment((hc_mult, hc_mult), T.float32) for j in T.Parallel(hc_mult): if save_pre_mix: pre_mix_out[i, j] = ( T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j]) + hc_pre_eps ) post_mix[i, j] = ( T.sigmoid( mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult] ) * hc_post_mult_value ) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = ( mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2] + hc_base[j * hc_mult + k + hc_mult * 2] ) row_sum = T.alloc_fragment(hc_mult, T.float32) col_sum = T.alloc_fragment(hc_mult, T.float32) row_max = T.alloc_fragment(hc_mult, T.float32) T.reduce_max(cm, row_max, dim=1) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = T.exp(cm[j, k] - row_max[j]) T.reduce_sum(cm, row_sum, dim=1) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps T.reduce_sum(cm, col_sum, dim=0) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps) for _ in T.serial(sinkhorn_repeat - 1): T.reduce_sum(cm, row_sum, dim=1) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps) T.reduce_sum(cm, col_sum, dim=0) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps) for j, k in T.Parallel(hc_mult, hc_mult): comb_mix[i, j * hc_mult + k] = cm[j, k] elif split_mode != "stats" and _is_collapse_thread(): pre_mix_shared = T.alloc_shared(hc_mult, T.float32) for j in T.Parallel(hc_mult): if use_pre_mix_in: pre_mix_shared[j] = pre_mix_in[i, j] elif save_pre_mix: pre_mix_shared[j] = T.if_then_else(j == 0, 1.0, 0.0) else: pre_mix_shared[j] = ( T.sigmoid(mixes_shared[j] * hc_scale[0] + hc_base[j]) + hc_pre_eps ) # Pass 1: stash unnormalized weighted-sum output in shared memory # as bf16 (matches the rounding that RMSNorm would see) while # accumulating the per-position squared sum. output_shared = T.alloc_shared(hidden_size, T.bfloat16) sumsq_per_pos = T.alloc_fragment(hidden_block, T.float32) T.clear(sumsq_per_pos) for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2): xs = T.alloc_shared((hc_mult, hidden_block), T.bfloat16) xl = T.alloc_fragment((hc_mult, hidden_block), T.float32) T.copy(residual[i, 0, i0_h * hidden_block], xs) T.copy(xs, xl) ol = T.alloc_fragment(hidden_block, T.float32) T.clear(ol) if write_aux: # Aux consumers (draft models) take the plain stream mean # of the same residual this collapse reads. aux = T.alloc_fragment(hidden_block, T.float32) T.clear(aux) for i_hc in T.serial(hc_mult): pre = pre_mix_shared[i_hc] for i1_h in T.Parallel(hidden_block): ol[i1_h] += pre * xl[i_hc, i1_h] if write_aux: aux[i1_h] += xl[i_hc, i1_h] if write_aux: for i1_h in T.Parallel(hidden_block): aux_out[i, i0_h * hidden_block + i1_h] = T.bfloat16( aux[i1_h] / hc_mult ) if save_pre_mix: # Keep the BF16 boundary before the delayed input RMSNorm. rounded = T.alloc_fragment(hidden_block, T.bfloat16) T.copy(ol, rounded) for i1_h in T.Parallel(hidden_block): value = T.float32(rounded[i1_h]) sumsq_per_pos[i1_h] += value * value output_shared[i0_h * hidden_block + i1_h] = rounded[i1_h] else: for i1_h in T.Parallel(hidden_block): sumsq_per_pos[i1_h] += ol[i1_h] * ol[i1_h] output_shared[i0_h * hidden_block + i1_h] = T.bfloat16(ol[i1_h]) sumsq = T.alloc_fragment(1, T.float32) T.reduce_sum(sumsq_per_pos, sumsq, dim=0) rsqrt_norm = T.alloc_fragment(1, T.float32) rsqrt_norm[0] = T.rsqrt(sumsq[0] / hidden_size + norm_eps) # Pass 2: scale by rsqrt * norm_weight and write the result to HBM. for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2): w_shared = T.alloc_shared(hidden_block, T.bfloat16) w_local = T.alloc_fragment(hidden_block, T.float32) T.copy(norm_weight[i0_h * hidden_block], w_shared) T.copy(w_shared, w_local) ol = T.alloc_fragment(hidden_block, T.float32) for i1_h in T.Parallel(hidden_block): ol[i1_h] = ( output_shared[i0_h * hidden_block + i1_h] * rsqrt_norm[0] * w_local[i1_h] ) T.copy(ol, layer_input[i, i0_h * hidden_block]) if ENABLE_PDL: T.pdl_trigger() @tilelang_jit def mhc_pre_big_fuse_broadcast_with_norm_tilelang( gemm_out_mul, gemm_out_sqrsum, hc_scale, hc_base, residual, residual_out, post_mix, comb_mix, layer_input, norm_weight, hidden_size: int, rms_eps: float, hc_pre_eps: float, hc_sinkhorn_eps: float, hc_post_mult_value: float, sinkhorn_repeat: int, norm_eps: float, n_splits: int = 1, hc_mult: int = 4, gemm_last_dim: int = -1, ): num_tokens = T.dynamic("num_tokens") hc_mult3 = hc_mult * (2 + hc_mult) if gemm_last_dim < 0: gemm_last_dim = hc_mult3 hidden_block = math.gcd(1024, hidden_size) gemm_out_mul: T.Tensor[[n_splits, num_tokens, gemm_last_dim], T.float32] # type: ignore[no-redef, valid-type] gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32] # type: ignore[no-redef, valid-type] hc_scale: T.Tensor[[3], T.float32] # type: ignore[no-redef, valid-type] hc_base: T.Tensor[[hc_mult3], T.float32] # type: ignore[no-redef, valid-type] residual: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] residual_out: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] post_mix: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type] comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32] # type: ignore[no-redef, valid-type] layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] norm_weight: T.Tensor[[hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] with T.Kernel(num_tokens, threads=96) as i: rms = T.alloc_fragment(1, T.float32) mixes = T.alloc_fragment(hc_mult3, T.float32) T.clear(mixes) rms[0] = 0 if ENABLE_PDL: T.pdl_sync() for i_split in T.serial(n_splits): rms[0] += gemm_out_sqrsum[i_split, i] # rms scale rms[0] = T.rsqrt(rms[0] / hidden_size + rms_eps) for j in T.Parallel(hc_mult3): mixes[j] = 0 for i_split in T.serial(n_splits): mixes[j] += gemm_out_mul[i_split, i, j] mixes[j] *= rms[0] mixes_shared = T.alloc_shared(hc_mult3, T.float32) T.copy(mixes, mixes_shared) if _is_mix_thread(): cm = T.alloc_fragment((hc_mult, hc_mult), T.float32) for j in T.Parallel(hc_mult): post_mix[i, j] = ( T.sigmoid( mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult] ) * hc_post_mult_value ) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = ( mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2] + hc_base[j * hc_mult + k + hc_mult * 2] ) row_sum = T.alloc_fragment(hc_mult, T.float32) col_sum = T.alloc_fragment(hc_mult, T.float32) row_max = T.alloc_fragment(hc_mult, T.float32) T.reduce_max(cm, row_max, dim=1) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = T.exp(cm[j, k] - row_max[j]) T.reduce_sum(cm, row_sum, dim=1) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps T.reduce_sum(cm, col_sum, dim=0) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps) for _ in T.serial(sinkhorn_repeat - 1): T.reduce_sum(cm, row_sum, dim=1) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps) T.reduce_sum(cm, col_sum, dim=0) for j, k in T.Parallel(hc_mult, hc_mult): cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps) for j, k in T.Parallel(hc_mult, hc_mult): comb_mix[i, j * hc_mult + k] = cm[j, k] else: pre_mix_shared = T.alloc_shared(hc_mult, T.float32) for j in T.Parallel(hc_mult): pre_mix_shared[j] = ( T.sigmoid( mixes_shared[j] * hc_scale[0] + hc_base[j], ) + hc_pre_eps ) output_shared = T.alloc_shared(hidden_size, T.bfloat16) sumsq_per_pos = T.alloc_fragment(hidden_block, T.float32) T.clear(sumsq_per_pos) for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2): xs = T.alloc_shared(hidden_block, T.bfloat16) xl = T.alloc_fragment(hidden_block, T.float32) T.copy(residual[i, i0_h * hidden_block], xs) T.copy(xs, xl) for i_hc, i1_h in T.Parallel(hc_mult, hidden_block): residual_out[i, i_hc, i0_h * hidden_block + i1_h] = xs[i1_h] ol = T.alloc_fragment(hidden_block, T.float32) T.clear(ol) for i_hc in T.serial(hc_mult): pre = pre_mix_shared[i_hc] for i1_h in T.Parallel(hidden_block): ol[i1_h] += pre * xl[i1_h] for i1_h in T.Parallel(hidden_block): sumsq_per_pos[i1_h] += ol[i1_h] * ol[i1_h] output_shared[i0_h * hidden_block + i1_h] = T.bfloat16(ol[i1_h]) sumsq = T.alloc_fragment(1, T.float32) T.reduce_sum(sumsq_per_pos, sumsq, dim=0) rsqrt_norm = T.alloc_fragment(1, T.float32) rsqrt_norm[0] = T.rsqrt(sumsq[0] / hidden_size + norm_eps) for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2): w_shared = T.alloc_shared(hidden_block, T.bfloat16) w_local = T.alloc_fragment(hidden_block, T.float32) T.copy(norm_weight[i0_h * hidden_block], w_shared) T.copy(w_shared, w_local) ol = T.alloc_fragment(hidden_block, T.float32) for i1_h in T.Parallel(hidden_block): ol[i1_h] = ( output_shared[i0_h * hidden_block + i1_h] * rsqrt_norm[0] * w_local[i1_h] ) T.copy(ol, layer_input[i, i0_h * hidden_block]) if ENABLE_PDL: T.pdl_trigger() @tilelang_jit def mhc_fused_tilelang( comb_mix, residual_in, post_mix, x_in, weight_t, yp_out, rp_out, residual_out, hc: int, hidden: int, n_out: int, n_thr: int = 256, h_blk: int = 256, tile_n: int = 1, split_k: int = 1, ) -> tilelang.JITKernel: """Fused mhc post-mapping + pre-norm GEMM FMA.""" m = T.dynamic("num_tokens") split_k = T.dynamic("split_k") h = hidden h_blk = math.gcd(hidden, h_blk) h_per_split = h // split_k n_tiles = n_out // tile_n comb_mix: T.Tensor((m, hc, hc), T.float32) # type: ignore[no-redef, valid-type] residual_in: T.Tensor((m, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type] post_mix: T.Tensor((m, hc), T.float32) # type: ignore[no-redef, valid-type] x_in: T.Tensor((m, h), T.bfloat16) # type: ignore[no-redef, valid-type] weight_t: T.Tensor((n_out, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type] yp_out: T.Tensor((split_k, m, n_out), T.float32) # type: ignore[no-redef, valid-type] rp_out: T.Tensor((split_k, m), T.float32) # type: ignore[no-redef, valid-type] residual_out: T.Tensor((m, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type] h_iters = h_per_split // n_thr num_warps = n_thr // 32 with T.Kernel(m, n_tiles, split_k, threads=n_thr) as (i_n, i_nt, i_ks): tid = T.get_thread_binding() warp_id = tid // 32 lane = tid % 32 s_warp = T.alloc_shared((num_warps, tile_n + 1), T.float32) s_post = T.alloc_shared((hc,), T.float32) s_comb = T.alloc_shared((hc, hc), T.float32) pm = T.alloc_local((hc,), T.float32) cm = T.alloc_local((hc, hc), T.float32) acc = T.alloc_local((tile_n,), T.float32) sqr = T.alloc_local((1,), T.float32) new_r = T.alloc_local((hc,), T.float32) T.clear(acc) T.clear(sqr) h_split_start = i_ks * h_per_split if ENABLE_PDL: T.pdl_sync() T.copy(post_mix[i_n, 0], s_post) T.copy(comb_mix[i_n, 0, 0], s_comb) for j in T.unroll(hc): pm[j] = s_post[j] for j in T.unroll(hc): for k in T.unroll(hc): cm[k, j] = s_comb[k, j] # Each thread owns h_iters elements of the k-split's h slice. for it in T.serial(h_iters): h_idx = h_split_start + it * n_thr + tid # Compute new residual from layer output and past residual for j in T.unroll(hc): new_r[j] = pm[j] * x_in[i_n, h_idx] for k in T.unroll(hc): new_r[j] += cm[k, j] * residual_in[i_n, k, h_idx] # populate residual_out and compute sqr sum if i_nt == 0: for j in T.unroll(hc): residual_out[i_n, j, h_idx] = new_r[j] sqr[0] += new_r[j] * new_r[j] # Per-thread FMA into acc[n] for n in T.unroll(tile_n): for j in T.unroll(hc): acc[n] += T.cast(weight_t[i_nt * tile_n + n, j, h_idx], T.float32) * new_r[j] for n in T.unroll(tile_n): acc[n] = T.warp_reduce_sum(acc[n]) if i_nt == 0: sqr[0] = T.warp_reduce_sum(sqr[0]) # Cross-warp reduce via shared mem if lane == 0: for n in T.unroll(tile_n): s_warp[warp_id, n] = acc[n] if i_nt == 0: s_warp[warp_id, tile_n] = sqr[0] T.sync_threads() # Warp 0 does the final cross-warp sum and writes outputs if warp_id == 0: if lane < tile_n: v = T.alloc_var(T.float32, init=0.0) for w in T.unroll(num_warps): v += s_warp[w, lane] yp_out[i_ks, i_n, i_nt * tile_n + lane] = v if i_nt == 0 and lane == 0: v2 = T.alloc_var(T.float32, init=0.0) for w in T.unroll(num_warps): v2 += s_warp[w, tile_n] rp_out[i_ks, i_n] = v2 if ENABLE_PDL: T.pdl_trigger() @tilelang_jit def mhc_post_tilelang( a, b, c, d, x, hc: int, hidden: int, n_thr: int = 128, h_blk: int = 1024, ) -> tilelang.JITKernel: # rename for shorter code n = T.dynamic("num_tokens") h = hidden h_blk = math.gcd(hidden, h_blk) a: T.Tensor((n, hc, hc), T.float32) # type: ignore[no-redef, valid-type] b: T.Tensor((n, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type] c: T.Tensor((n, hc), T.float32) # type: ignore[no-redef, valid-type] d: T.Tensor((n, h), T.bfloat16) # type: ignore[no-redef, valid-type] x: T.Tensor((n, hc, h), T.bfloat16) # type: ignore[no-redef, valid-type] with T.Kernel(n, threads=n_thr) as i_n: b_shared = T.alloc_shared((hc, h_blk), T.bfloat16) d_shared = T.alloc_shared(h_blk, T.bfloat16) x_local = T.alloc_fragment((hc, h_blk), T.float32) b_local = T.alloc_fragment((hc, h_blk), T.float32) d_local = T.alloc_fragment(h_blk, T.float32) a_local = T.alloc_fragment((hc, hc), T.float32) c_local = T.alloc_fragment(hc, T.float32) if ENABLE_PDL: T.pdl_sync() T.copy(a[i_n, 0, 0], a_local) T.copy(c[i_n, 0], c_local) for i0_h in T.Serial(T.ceildiv(h, h_blk)): T.copy(b[i_n, 0, i0_h * h_blk], b_shared) T.copy(d[i_n, i0_h * h_blk], d_shared) T.copy(b_shared, b_local) T.copy(d_shared, d_local) for i_hco, i1_h in T.Parallel(hc, h_blk): x_local[i_hco, i1_h] = c_local[i_hco] * d_local[i1_h] for i_hci in T.vectorized(hc): x_local[i_hco, i1_h] += a_local[i_hci, i_hco] * b_local[i_hci, i1_h] T.copy(x_local, x[i_n, 0, i0_h * h_blk]) if ENABLE_PDL: T.pdl_trigger() @tilelang_jit def hc_prenorm_gemm_tilelang( x, fn, out, sqrsum, hidden_size: int, hc_mult: int = 4, n_out: int = 24, n_thr: int = 512, tile_n: int = 12, n_splits: int = 1, ) -> tilelang.JITKernel: num_tokens = T.dynamic("num_tokens") hc_hidden_size = hc_mult * hidden_size k_per_split = hc_hidden_size // n_splits k_iters = k_per_split // n_thr n_tiles = T.ceildiv(n_out, tile_n) x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16) # type: ignore[no-redef, valid-type] fn: T.Tensor((n_out, hc_hidden_size), T.float32) # type: ignore[no-redef, valid-type] out: T.Tensor((n_splits, num_tokens, n_out), T.float32) # type: ignore[no-redef, valid-type] sqrsum: T.Tensor((n_splits, num_tokens), T.float32) # type: ignore[no-redef, valid-type] with T.Kernel(num_tokens, n_tiles, n_splits, threads=n_thr) as ( i_n, i_t, i_s, ): tid = T.get_thread_binding() acc = T.alloc_local((tile_n,), T.float32) sqr = T.alloc_local((1,), T.float32) T.clear(acc) T.clear(sqr) if ENABLE_PDL: T.pdl_sync() for it in T.serial(k_iters): i_k = i_s * k_per_split + it * n_thr + tid x_val = T.cast(x[i_n, i_k], T.float32) for i_o in T.unroll(tile_n): out_idx = i_t * tile_n + i_o if out_idx < n_out: acc[i_o] += x_val * fn[out_idx, i_k] if i_t == 0: sqr[0] += x_val * x_val for i_o in T.unroll(tile_n): acc[i_o] = T.warp_reduce_sum(acc[i_o]) if i_t == 0: sqr[0] = T.warp_reduce_sum(sqr[0]) lane = tid % 32 warp_id = tid // 32 num_warps = n_thr // 32 warp_acc = T.alloc_shared((num_warps, tile_n), T.float32) warp_sqr = T.alloc_shared(num_warps, T.float32) if lane == 0: for i_o in T.unroll(tile_n): warp_acc[warp_id, i_o] = acc[i_o] if i_t == 0: warp_sqr[warp_id] = sqr[0] T.sync_threads() if warp_id == 0: if lane < tile_n: reduced_acc = T.alloc_var(T.float32, init=0.0) for i_w in T.unroll(num_warps): reduced_acc += warp_acc[i_w, lane] out_idx = i_t * tile_n + lane if out_idx < n_out: out[i_s, i_n, out_idx] = reduced_acc if lane == 0 and i_t == 0: reduced_sqr = T.alloc_var(T.float32, init=0.0) for i_w in T.unroll(num_warps): reduced_sqr += warp_sqr[i_w] sqrsum[i_s, i_n] = reduced_sqr if ENABLE_PDL: T.pdl_trigger() @tilelang_jit def hc_prenorm_gemm_block_m_tilelang( x, fn, out, sqrsum, hidden_size: int, hc_mult: int = 4, n_out: int = 24, n_thr: int = 512, tile_n: int = 12, block_m: int = 2, ) -> tilelang.JITKernel: num_tokens = T.dynamic("num_tokens") hc_hidden_size = hc_mult * hidden_size k_iters = hc_hidden_size // n_thr n_tiles = T.ceildiv(n_out, tile_n) m_tiles = T.ceildiv(num_tokens, block_m) x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16) # type: ignore[no-redef, valid-type] fn: T.Tensor((n_out, hc_hidden_size), T.float32) # type: ignore[no-redef, valid-type] out: T.Tensor((1, num_tokens, n_out), T.float32) # type: ignore[no-redef, valid-type] sqrsum: T.Tensor((1, num_tokens), T.float32) # type: ignore[no-redef, valid-type] with T.Kernel(m_tiles, n_tiles, threads=n_thr) as (i_mt, i_t): tid = T.get_thread_binding() acc = T.alloc_local((block_m, tile_n), T.float32) sqr = T.alloc_local((block_m,), T.float32) T.clear(acc) T.clear(sqr) if ENABLE_PDL: T.pdl_sync() for it in T.serial(k_iters): i_k = it * n_thr + tid fn_val = T.alloc_local((tile_n,), T.float32) for i_o in T.unroll(tile_n): out_idx = i_t * tile_n + i_o if out_idx < n_out: fn_val[i_o] = fn[out_idx, i_k] else: fn_val[i_o] = 0.0 for i_m in T.unroll(block_m): token_idx = i_mt * block_m + i_m if token_idx < num_tokens: x_val = T.cast(x[token_idx, i_k], T.float32) for i_o in T.unroll(tile_n): acc[i_m, i_o] += x_val * fn_val[i_o] if i_t == 0: sqr[i_m] += x_val * x_val for i_m in T.unroll(block_m): for i_o in T.unroll(tile_n): acc[i_m, i_o] = T.warp_reduce_sum(acc[i_m, i_o]) if i_t == 0: sqr[i_m] = T.warp_reduce_sum(sqr[i_m]) lane = tid % 32 warp_id = tid // 32 num_warps = n_thr // 32 warp_acc = T.alloc_shared((num_warps, block_m, tile_n), T.float32) warp_sqr = T.alloc_shared((num_warps, block_m), T.float32) if lane == 0: for i_m in T.unroll(block_m): for i_o in T.unroll(tile_n): warp_acc[warp_id, i_m, i_o] = acc[i_m, i_o] if i_t == 0: warp_sqr[warp_id, i_m] = sqr[i_m] T.sync_threads() if warp_id == 0: reduced_acc = T.alloc_local((block_m,), T.float32) reduced_sqr = T.alloc_local((block_m,), T.float32) T.clear(reduced_acc) T.clear(reduced_sqr) for i_m in T.unroll(block_m): token_idx = i_mt * block_m + i_m if token_idx < num_tokens: if lane < tile_n: for i_w in T.unroll(num_warps): reduced_acc[i_m] += warp_acc[i_w, i_m, lane] out_idx = i_t * tile_n + lane if out_idx < n_out: out[0, token_idx, out_idx] = reduced_acc[i_m] if lane == 0 and i_t == 0: for i_w in T.unroll(num_warps): reduced_sqr[i_m] += warp_sqr[i_w, i_m] sqrsum[0, token_idx] = reduced_sqr[i_m] if ENABLE_PDL: T.pdl_trigger() @tilelang_jit def hc_head_fuse_tilelang( residual, fn, hc_scale, hc_base, out, hidden_size: int, rms_eps: float, hc_eps: float, hc_mult: int = 4, n_thr: int = 128, h_blk: int = 1024, ): """Two-pass fused kernel for hc_head. Pass 1: accumulate per-token squared sum and hc_mult dot-products (projections onto fn rows) using cross-thread reducers. Pass 2: apply sigmoid-gated weighted sum of residual channels to output. Avoids materialising mixes / rsqrt / pre tensors to global memory. """ num_tokens = T.dynamic("num_tokens") hc_dim = hc_mult * hidden_size h_block = math.gcd(h_blk, hidden_size) n_h = hidden_size // h_block residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type] fn: T.Tensor[[hc_mult, hc_dim], T.float32] # type: ignore[no-redef,valid-type] hc_scale: T.Tensor[[1], T.float32] # type: ignore[no-redef,valid-type] hc_base: T.Tensor[[hc_mult], T.float32] # type: ignore[no-redef,valid-type] out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type] with T.Kernel(num_tokens, threads=n_thr) as i: if ENABLE_PDL: T.pdl_sync() # ------------------------------------------------------------------ # Pass 1 – for each residual channel m_c and h_block: # • accumulate squared sum (for RMS norm denominator) # • accumulate hc_mult dot-products with fn rows # ------------------------------------------------------------------ sqrsum_r = T.alloc_reducer((1,), T.float32, replication="all") mixes_r = T.alloc_reducer((hc_mult,), T.float32, replication="all") T.fill(sqrsum_r, 0.0) T.fill(mixes_r, 0.0) for m_c in T.serial(hc_mult): for i_h in T.serial(n_h): x_local = T.alloc_fragment(h_block, T.float32) T.copy(residual[i, m_c, i_h * h_block], x_local) for k in T.Parallel(h_block): sqrsum_r[0] += x_local[k] * x_local[k] for m_m in T.unroll(hc_mult): fn_local = T.alloc_fragment(h_block, T.float32) T.copy(fn[m_m, m_c * hidden_size + i_h * h_block], fn_local) for k in T.Parallel(h_block): mixes_r[m_m] += x_local[k] * fn_local[k] T.finalize_reducer(sqrsum_r) T.finalize_reducer(mixes_r) # ------------------------------------------------------------------ # Compute pre_mix = sigmoid(mix * rsqrt * scale + base) + eps # ------------------------------------------------------------------ pre_mix_shared = T.alloc_shared(hc_mult, T.float32) rsqrt_val = T.alloc_fragment(1, T.float32) rsqrt_val[0] = T.rsqrt(sqrsum_r[0] / hc_dim + rms_eps) for m in T.Parallel(hc_mult): pre_mix_shared[m] = ( T.sigmoid(mixes_r[m] * rsqrt_val[0] * hc_scale[0] + hc_base[m]) + hc_eps ) # ------------------------------------------------------------------ # Pass 2 – apply_mix: pipelined weighted sum over residual channels # ------------------------------------------------------------------ for i0_h in T.Pipelined(n_h, num_stages=2): xs = T.alloc_shared((hc_mult, h_block), T.bfloat16) xl = T.alloc_fragment((hc_mult, h_block), T.float32) T.copy(residual[i, 0, i0_h * h_block], xs, disable_tma=True) T.copy(xs, xl) ol = T.alloc_fragment(h_block, T.float32) T.clear(ol) for i_hc in T.serial(hc_mult): pre = pre_mix_shared[i_hc] for i1_h in T.Parallel(h_block): ol[i1_h] += pre * xl[i_hc, i1_h] T.copy(ol, out[i, i0_h * h_block], disable_tma=True) if ENABLE_PDL: T.pdl_trigger() class HcPrenormGemmTileLangKernel( VllmTileLangJitKernel["HcPrenormGemmTileLangKernel.CompileKey"] ): @dataclass(frozen=True) class CompileKey: hidden_size: int hc_mult: int n_out: int n_thr: int tile_n: int n_splits: int use_block_m: bool block_m: int @staticmethod def kernel(compile_key: CompileKey) -> Any: if compile_key.use_block_m: return hc_prenorm_gemm_block_m_tilelang return hc_prenorm_gemm_tilelang def dispatch( # type: ignore[override] self, *, num_tokens: int, hc_hidden_size: int, n_thr: int = 512, tile_n: int = 12, n_splits: int = 1, **compile_key_fields: int, ) -> CompileKey: use_default_config = n_splits == 1 and tile_n == 12 and n_thr == 512 use_block_m = use_default_config and num_tokens >= 1024 use_small_tile = ( use_default_config and not use_block_m and num_tokens < 128 and hc_hidden_size % 1024 == 0 ) effective_n_thr = 1024 if use_small_tile else n_thr effective_tile_n = 4 if use_small_tile else tile_n return self.CompileKey( **compile_key_fields, n_thr=effective_n_thr, tile_n=effective_tile_n, n_splits=n_splits, use_block_m=use_block_m, block_m=2 if use_block_m else 1, ) def get_warmup_keys( self, vllm_config: Any, *, hidden_size: int, hc_mult: int, n_out: int, ) -> list[CompileKey]: max_tokens = vllm_config.scheduler_config.max_num_batched_tokens if max_tokens <= 0: return [] return self._trace_dispatch(self.dispatch)( num_tokens=WarmupIntRange( 1, max_tokens + 1, advance=lambda value: ( 128 if value < 128 else 1024 if value < 1024 else max_tokens + 1 ), ), hc_hidden_size=hidden_size * hc_mult, hidden_size=hidden_size, hc_mult=hc_mult, n_out=n_out, ) def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: num_tokens = 128 if compile_key.use_block_m: num_tokens = 1024 elif compile_key.n_thr == 1024: num_tokens = 1 hc_hidden_size = compile_key.hidden_size * compile_key.hc_mult return dict( x=make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hc_hidden_size), fn=make_tilelang_warmup_tensor( torch.float32, compile_key.n_out, hc_hidden_size ), out=make_tilelang_warmup_tensor( torch.float32, compile_key.n_splits, num_tokens, compile_key.n_out, ), sqrsum=make_tilelang_warmup_tensor( torch.float32, compile_key.n_splits, num_tokens ), hidden_size=compile_key.hidden_size, hc_mult=compile_key.hc_mult, tile_n=compile_key.tile_n, n_thr=compile_key.n_thr, n_splits=compile_key.n_splits, ) @kernel_launcher def __call__( self, x: Any, fn: Any, out: Any, sqrsum: Any, hidden_size: int, hc_mult: int, tile_n: int = 12, n_thr: int = 512, n_splits: int = 1, ) -> TileLangLaunchSpec: assert out.shape[0] == n_splits assert sqrsum.shape[0] == n_splits assert x.shape[1] == hc_mult * hidden_size assert x.shape[1] % n_splits == 0 assert (x.shape[1] // n_splits) % n_thr == 0 compile_key = self.dispatch( num_tokens=x.shape[0], hc_hidden_size=x.shape[1], hidden_size=hidden_size, hc_mult=hc_mult, n_out=fn.shape[0], n_thr=n_thr, tile_n=tile_n, n_splits=n_splits, ) kernel_arg = ( compile_key.block_m if compile_key.use_block_m else compile_key.n_splits ) return (compile_key,), ( x, fn, out, sqrsum, hidden_size, hc_mult, fn.shape[0], compile_key.n_thr, compile_key.tile_n, kernel_arg, ) class MhcPreBigFuseTileLangKernel( VllmTileLangJitKernel["MhcPreBigFuseTileLangKernel.CompileKey"] ): @dataclass(frozen=True) class CompileKey: hidden_size: int hc_mult: int n_splits: int use_norm_weight: bool is_broadcast: bool rms_eps: float hc_pre_eps: float hc_sinkhorn_eps: float hc_post_mult_value: float sinkhorn_repeat: int norm_eps: float @staticmethod def kernel(compile_key: CompileKey) -> Any: if compile_key.is_broadcast: return mhc_pre_big_fuse_broadcast_with_norm_tilelang if compile_key.use_norm_weight: return mhc_pre_big_fuse_with_norm_tilelang return mhc_pre_big_fuse_tilelang @staticmethod def _kernel_args( compile_key: CompileKey, *, gemm_out_mul: Any, gemm_out_sqrsum: Any, hc_scale: Any, hc_base: Any, residual: Any, residual_out: Any | None, post_mix: Any, comb_mix: Any, layer_input: Any, norm_weight: Any | None, ) -> tuple[Any, ...]: common_args = ( gemm_out_mul, gemm_out_sqrsum, hc_scale, hc_base, residual, ) output_args = (post_mix, comb_mix, layer_input) eps_args = ( compile_key.rms_eps, compile_key.hc_pre_eps, compile_key.hc_sinkhorn_eps, compile_key.hc_post_mult_value, compile_key.sinkhorn_repeat, ) middle_args: tuple[Any, ...] if compile_key.is_broadcast: assert residual_out is not None assert norm_weight is not None middle_args = (residual_out, *output_args, norm_weight) elif compile_key.use_norm_weight: assert norm_weight is not None # layer_input stands in for the aux buffer on both epilogues: same # shape and dtype, and write_aux is off here so it is never written. middle_args = ( *output_args, norm_weight, post_mix, post_mix, layer_input, ) else: middle_args = (*output_args, post_mix, post_mix, layer_input) norm_eps_args = (compile_key.norm_eps,) if compile_key.use_norm_weight else () return ( *common_args, *middle_args, compile_key.hidden_size, *eps_args, *norm_eps_args, compile_key.n_splits, compile_key.hc_mult, ) def dispatch( # type: ignore[override] self, *, hidden_size: int, hc_mult: int, n_splits: int, is_broadcast: bool, use_norm_weight: bool, use_pre_gemm_splits: bool = False, pre_gemm_k: int = 0, sinkhorn_repeat: int, norm_eps: float, broadcast_norm_eps: float = 0.0, num_tokens: int = 1, use_fused_tilelang: bool = False, **compile_key_fields: float, ) -> CompileKey: pre_gemm_n_splits = ( compute_num_split(64, pre_gemm_k, (num_tokens + 63) // 64) if use_pre_gemm_splits else n_splits ) # The epilogue reduces over whatever the fused kernel split the GEMM # into, so both have to read the same config. fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) actual_n_splits = ( fused_config[1] if use_fused_tilelang and fused_config is not None else pre_gemm_n_splits ) actual_norm_eps = broadcast_norm_eps if is_broadcast else norm_eps actual_use_norm_weight = use_norm_weight or is_broadcast return self.CompileKey( **compile_key_fields, hidden_size=hidden_size, hc_mult=hc_mult, n_splits=actual_n_splits, use_norm_weight=actual_use_norm_weight, is_broadcast=is_broadcast, sinkhorn_repeat=sinkhorn_repeat, norm_eps=actual_norm_eps if actual_use_norm_weight else 0.0, ) def get_warmup_keys( self, vllm_config: Any, *, hidden_size: int, hc_mult: int, use_norm_weight: bool, rms_eps: float, hc_pre_eps: float, hc_sinkhorn_eps: float, hc_post_mult_value: float, sinkhorn_repeat: int, norm_eps: float | tuple[float, ...], broadcast_norm_eps: float | tuple[float, ...] | None = None, include_pre_gemm_splits: bool = False, include_broadcast_splits: bool = False, ) -> list[CompileKey]: max_tokens = vllm_config.scheduler_config.max_num_batched_tokens if max_tokens <= 0: return [] warmup_cases = zip_inputs( dict( is_broadcast=False, use_pre_gemm_splits=False, pre_gemm_k=0, use_fused_tilelang=False, ), dict( is_broadcast=False, use_pre_gemm_splits=include_pre_gemm_splits, pre_gemm_k=hc_mult * hidden_size, use_fused_tilelang=False, ), dict( is_broadcast=include_broadcast_splits, use_pre_gemm_splits=include_broadcast_splits, pre_gemm_k=hidden_size, use_fused_tilelang=False, ), dict( is_broadcast=False, use_pre_gemm_splits=False, pre_gemm_k=0, use_fused_tilelang=True, ), ) return self._trace_dispatch(self.dispatch)( warmup_cases, num_tokens=WarmupIntRange( 1, max_tokens + 1, advance=lambda value: 8 if value < 8 else cdiv(value, 64) * 64 + 1, ), n_splits=1, hidden_size=hidden_size, hc_mult=hc_mult, use_norm_weight=use_norm_weight, rms_eps=rms_eps, hc_pre_eps=hc_pre_eps, hc_sinkhorn_eps=hc_sinkhorn_eps, hc_post_mult_value=hc_post_mult_value, sinkhorn_repeat=sinkhorn_repeat, norm_eps=norm_eps, broadcast_norm_eps=broadcast_norm_eps if broadcast_norm_eps is not None else norm_eps, _when=lambda *, use_fused_tilelang, num_tokens: ( not use_fused_tilelang or num_tokens <= 16 ), ) def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: hidden_size = compile_key.hidden_size hc_mult = compile_key.hc_mult hc_mult3 = hc_mult * (2 + hc_mult) num_tokens = 1 gemm_out_mul = make_tilelang_warmup_tensor( torch.float32, compile_key.n_splits, num_tokens, hc_mult3 ) gemm_out_sqrsum = make_tilelang_warmup_tensor( torch.float32, compile_key.n_splits, num_tokens ) hc_scale = make_tilelang_warmup_tensor(torch.float32, 3) hc_base = make_tilelang_warmup_tensor(torch.float32, hc_mult3) post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult) comb_mix = make_tilelang_warmup_tensor( torch.float32, num_tokens, hc_mult * hc_mult ) layer_input = make_tilelang_warmup_tensor( torch.bfloat16, num_tokens, hidden_size ) residual_shape = ( (num_tokens, hidden_size) if compile_key.is_broadcast else (num_tokens, hc_mult, hidden_size) ) residual = make_tilelang_warmup_tensor(torch.bfloat16, *residual_shape) residual_out = ( make_tilelang_warmup_tensor( torch.bfloat16, num_tokens, hc_mult, hidden_size ) if compile_key.is_broadcast else None ) norm_weight = ( make_tilelang_warmup_tensor(torch.bfloat16, hidden_size) if compile_key.is_broadcast or compile_key.use_norm_weight else None ) return dict( gemm_out_mul=gemm_out_mul, gemm_out_sqrsum=gemm_out_sqrsum, hc_scale=hc_scale, hc_base=hc_base, residual=residual, post_mix=post_mix, comb_mix=comb_mix, layer_input=layer_input, rms_eps=compile_key.rms_eps, hc_pre_eps=compile_key.hc_pre_eps, hc_sinkhorn_eps=compile_key.hc_sinkhorn_eps, hc_post_mult_value=compile_key.hc_post_mult_value, sinkhorn_repeat=compile_key.sinkhorn_repeat, residual_out=residual_out, norm_weight=norm_weight, norm_eps=compile_key.norm_eps, ) @kernel_launcher def __call__( self, gemm_out_mul: Any, gemm_out_sqrsum: Any, hc_scale: Any, hc_base: Any, residual: Any, post_mix: Any, comb_mix: Any, layer_input: Any, rms_eps: float, hc_pre_eps: float, hc_sinkhorn_eps: float, hc_post_mult_value: float, sinkhorn_repeat: int, *, residual_out: Any | None = None, norm_weight: Any | None = None, norm_eps: float = 0.0, ) -> TileLangLaunchSpec: is_broadcast = residual_out is not None use_norm_weight = norm_weight is not None n_splits = gemm_out_mul.shape[0] if is_broadcast: assert residual_out is not None hidden_size = residual.shape[-1] hc_mult = residual_out.shape[-2] else: hc_mult = residual.shape[-2] hidden_size = residual.shape[-1] compile_key = self.dispatch( hidden_size=hidden_size, hc_mult=hc_mult, n_splits=n_splits, is_broadcast=is_broadcast, use_norm_weight=use_norm_weight, rms_eps=rms_eps, hc_pre_eps=hc_pre_eps, hc_sinkhorn_eps=hc_sinkhorn_eps, hc_post_mult_value=hc_post_mult_value, sinkhorn_repeat=sinkhorn_repeat, norm_eps=norm_eps, broadcast_norm_eps=norm_eps, ) return ( (compile_key,), self._kernel_args( compile_key, gemm_out_mul=gemm_out_mul, gemm_out_sqrsum=gemm_out_sqrsum, hc_scale=hc_scale, hc_base=hc_base, residual=residual, residual_out=residual_out, post_mix=post_mix, comb_mix=comb_mix, layer_input=layer_input, norm_weight=norm_weight, ), ) class MhcPostTileLangKernel(VllmTileLangJitKernel["MhcPostTileLangKernel.CompileKey"]): @dataclass(frozen=True) class CompileKey: hidden_size: int hc_mult: int @staticmethod def kernel() -> Any: return mhc_post_tilelang def dispatch( # type: ignore[override] self, **compile_key_fields: int, ) -> CompileKey: return self.CompileKey(**compile_key_fields) def get_warmup_keys( self, *, hidden_size: int, hc_mult: int, ) -> list[CompileKey]: return self._trace_dispatch(self.dispatch)( hidden_size=hidden_size, hc_mult=hc_mult, ) def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: num_tokens = 1 hidden_size = compile_key.hidden_size hc_mult = compile_key.hc_mult comb_mix = make_tilelang_warmup_tensor( torch.float32, num_tokens, hc_mult, hc_mult ) residual = make_tilelang_warmup_tensor( torch.bfloat16, num_tokens, hc_mult, hidden_size ) post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult) layer_input = make_tilelang_warmup_tensor( torch.bfloat16, num_tokens, hidden_size ) out = make_tilelang_warmup_tensor( torch.bfloat16, num_tokens, hc_mult, hidden_size ) return dict( comb_mix=comb_mix, residual=residual, post_mix=post_mix, layer_input=layer_input, out=out, hc_mult=hc_mult, hidden_size=hidden_size, ) @kernel_launcher def __call__( self, comb_mix: Any, residual: Any, post_mix: Any, layer_input: Any, out: Any, hc_mult: int, hidden_size: int, ) -> TileLangLaunchSpec: return (), ( comb_mix, residual, post_mix, layer_input, out, hc_mult, hidden_size, ) class MhcFusedTileLangKernel( VllmTileLangJitKernel["MhcFusedTileLangKernel.CompileKey"] ): @dataclass(frozen=True) class CompileKey: hidden_size: int hc_mult: int n_splits: int tile_n: int n_thr: int @staticmethod def kernel() -> Any: return mhc_fused_tilelang def dispatch( # type: ignore[override] self, *, num_tokens: int, hidden_size: int, hc_mult: int, ) -> CompileKey: # Subscripts rather than an unpack: the warmup tracer parses this body # and allows only plain assignments before the return. config = require_fused_post_pre_config(num_tokens, hidden_size, hc_mult) return self.CompileKey( hidden_size=hidden_size, hc_mult=hc_mult, n_splits=config[1], tile_n=config[0], n_thr=config[2], ) def get_warmup_keys( self, vllm_config: Any, *, hidden_size: int, hc_mult: int, ) -> list[CompileKey]: max_tokens = vllm_config.scheduler_config.max_num_batched_tokens return self._trace_dispatch(self.dispatch)( num_tokens=WarmupIntRange(1, max_tokens + 1), hidden_size=hidden_size, hc_mult=hc_mult, _when=lambda *, num_tokens: num_tokens <= _FUSED_POST_PRE_MAX_TOKENS, ) def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: num_tokens = 1 if compile_key.tile_n == 2 else 16 hidden_size = compile_key.hidden_size hc_mult = compile_key.hc_mult hc_mult3 = hc_mult * (2 + hc_mult) comb_mix = make_tilelang_warmup_tensor( torch.float32, num_tokens, hc_mult, hc_mult ) residual_in = make_tilelang_warmup_tensor( torch.bfloat16, num_tokens, hc_mult, hidden_size ) post_mix = make_tilelang_warmup_tensor(torch.float32, num_tokens, hc_mult) x_in = make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hidden_size) weight_t = make_tilelang_warmup_tensor( torch.bfloat16, hc_mult3, hc_mult, hidden_size ) return dict( comb_mix=comb_mix, residual_in=residual_in, post_mix=post_mix, x_in=x_in, weight_t=weight_t, hc_mult=hc_mult, hidden_size=hidden_size, hc_mult3=hc_mult3, ) @kernel_launcher def __call__( self, comb_mix: Any, residual_in: Any, post_mix: Any, x_in: Any, weight_t: Any, hc_mult: int, hidden_size: int, hc_mult3: int, ) -> TileLangLaunchSpec: num_tokens = residual_in.shape[0] tile_n, n_splits, n_thr = require_fused_post_pre_config( num_tokens, hidden_size, hc_mult ) yp_out = residual_in.new_empty( (n_splits, num_tokens, hc_mult3), dtype=torch.float32 ) rp_out = residual_in.new_empty((n_splits, num_tokens), dtype=torch.float32) residual_out = residual_in.new_empty(residual_in.shape) return ( (), ( comb_mix, residual_in, post_mix, x_in, weight_t, yp_out, rp_out, residual_out, hc_mult, hidden_size, hc_mult3, ), dict(n_thr=n_thr, tile_n=tile_n), (yp_out, rp_out, residual_out), ) class HcHeadFusedTileLangKernel( VllmTileLangJitKernel["HcHeadFusedTileLangKernel.CompileKey"] ): @dataclass(frozen=True) class CompileKey: hidden_size: int hc_mult: int rms_eps: float hc_eps: float @staticmethod def kernel() -> Any: return hc_head_fuse_tilelang def dispatch( # type: ignore[override] self, *, hidden_size: int, hc_mult: int, rms_eps: float, hc_eps: float, ) -> CompileKey: return self.CompileKey( hidden_size=hidden_size, hc_mult=hc_mult, rms_eps=rms_eps, hc_eps=hc_eps, ) def get_warmup_keys( self, *, hidden_size: int, hc_mult: int, rms_eps: float, hc_eps: float, ) -> list[CompileKey]: return self._trace_dispatch(self.dispatch)( hidden_size=hidden_size, hc_mult=hc_mult, rms_eps=rms_eps, hc_eps=hc_eps, ) def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: num_tokens = 1 hidden_size = compile_key.hidden_size hc_mult = compile_key.hc_mult residual = make_tilelang_warmup_tensor( torch.bfloat16, num_tokens, hc_mult, hidden_size ) fn = make_tilelang_warmup_tensor(torch.float32, hc_mult, hc_mult * hidden_size) hc_scale = make_tilelang_warmup_tensor(torch.float32, 1) hc_base = make_tilelang_warmup_tensor(torch.float32, hc_mult) out = make_tilelang_warmup_tensor(torch.bfloat16, num_tokens, hidden_size) return dict( residual=residual, fn=fn, hc_scale=hc_scale, hc_base=hc_base, out=out, hidden_size=hidden_size, rms_eps=compile_key.rms_eps, hc_eps=compile_key.hc_eps, hc_mult=hc_mult, ) @kernel_launcher def __call__( self, residual: Any, fn: Any, hc_scale: Any, hc_base: Any, out: Any, hidden_size: int, rms_eps: float, hc_eps: float, hc_mult: int, ) -> TileLangLaunchSpec: return (), ( residual, fn, hc_scale, hc_base, out, hidden_size, rms_eps, hc_eps, hc_mult, ) _HC_PRENORM_GEMM_TILELANG_KERNEL = HcPrenormGemmTileLangKernel() _MHC_PRE_BIG_FUSE_TILELANG_KERNEL = MhcPreBigFuseTileLangKernel() _MHC_POST_TILELANG_KERNEL = MhcPostTileLangKernel() _MHC_FUSED_TILELANG_KERNEL = MhcFusedTileLangKernel() _HC_HEAD_FUSED_TILELANG_KERNEL = HcHeadFusedTileLangKernel()