# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math from functools import cache from typing import TYPE_CHECKING, Any import torch from vllm.platforms import current_platform from vllm.utils.import_utils import has_tilelang from vllm.utils.math_utils import cdiv # TileLang is used for MHC on CUDA and ROCm. Keep non-GPU imports cheap so # registering the Python wrapper modules does not require TileLang everywhere. if TYPE_CHECKING or current_platform.is_cuda_alike(): if not has_tilelang(): raise ImportError( "tilelang is required for mhc but is not installed. Install it with " "`pip install tilelang`." ) import tilelang import tilelang.language as T else: tilelang = None # type: ignore[assignment] T = None # type: ignore[assignment] ENABLE_PDL = current_platform.is_arch_support_pdl() and current_platform.is_cuda() @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 pass_configs: dict[tilelang.PassConfigKey, Any] = { tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, } if current_platform.is_cuda(): pass_configs[tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL] = 10 @tilelang.jit( pass_configs=pass_configs, ) def mhc_pre_big_fuse_tilelang( gemm_out_mul, gemm_out_sqrsum, hc_scale, hc_base, residual, post_mix, comb_mix, layer_input, 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, ): """Deeply fused kernels, everything other than gemm & sqrsum in mHC pre block.""" num_tokens = T.dynamic("num_tokens") hc_mult3 = hc_mult * (2 + hc_mult) hidden_block = math.gcd(512, 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] 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] / (hc_mult * 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 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): 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): 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) 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] 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( pass_configs=pass_configs, ) 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, 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, ): 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, 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[0] = T.rsqrt(rms[0] / (hc_mult * 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 T.get_thread_binding() < 32: 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 ) # 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) 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] 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( pass_configs=pass_configs, ) 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.float32) # 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] += weight_t[i_nt * tile_n + n, j, h_idx] * 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( pass_configs=pass_configs, ) 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( pass_configs=pass_configs, ) 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 = x[i_n, i_k] 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( pass_configs=pass_configs, ) 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 = x[token_idx, i_k] 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: 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: reduced_acc = T.alloc_var(T.float32, init=0.0) for i_w in T.unroll(num_warps): reduced_acc += 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 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, i_m] sqrsum[0, token_idx] = reduced_sqr if ENABLE_PDL: T.pdl_trigger() @tilelang.jit( pass_configs=pass_configs, ) 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()