Woosuk KwonandCodex
9b49f92344
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Codex <noreply@openai.com>
1857 lines
66 KiB
Python
1857 lines
66 KiB
Python
# 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()
|
||
|
||
|
||
@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 = 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 T.get_thread_binding() < 32:
|
||
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 T.get_thread_binding() >= 32:
|
||
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 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
|
||
)
|
||
|
||
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.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
|
||
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.float32, 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()
|