wehub-resource-sync
239 lines
10 KiB
Python
239 lines
10 KiB
Python
"""A compiler pass that fuses add + rms_norm."""
|
|
|
|
from typing import Optional
|
|
|
|
import tvm
|
|
from tvm import relax
|
|
from tvm.relax.analysis import remove_all_unused
|
|
from tvm.relax.expr_functor import PyExprMutator, mutator
|
|
from tvm.script import tirx as T
|
|
|
|
from ..support.max_thread_check import get_max_num_threads_per_block
|
|
|
|
|
|
def _get_add_rms_norm_decode(hidden_size: int, eps: float, TX: int, in_dtype: str):
|
|
if in_dtype not in ("float16", "bfloat16"):
|
|
raise ValueError(f"Unsupported data type: {in_dtype}")
|
|
inv_hidden_size = T.float32(1.0 / float(hidden_size))
|
|
eps = T.float32(eps)
|
|
add_local_size = hidden_size // TX
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def decode_add_rms(pA: T.handle, pB: T.handle, pC: T.handle, pO: T.handle, pAdd: T.handle):
|
|
T.func_attr({"tirx.noalias": True, "tirx.is_scheduled": 1})
|
|
batch_size = T.int32()
|
|
A = T.match_buffer(pA, (batch_size, 1, hidden_size), in_dtype)
|
|
B = T.match_buffer(pB, (batch_size, 1, hidden_size), in_dtype)
|
|
C = T.match_buffer(pC, (hidden_size,), in_dtype)
|
|
out = T.match_buffer(pO, (batch_size, 1, hidden_size), in_dtype)
|
|
add = T.match_buffer(pAdd, (batch_size, 1, hidden_size), in_dtype)
|
|
add_local = T.sblock_alloc_buffer((hidden_size // TX,), in_dtype, scope="local")
|
|
sum_shared = T.sblock_alloc_buffer((batch_size, 1), scope="shared")
|
|
sum_local = T.sblock_alloc_buffer((TX, batch_size, 1), scope="local")
|
|
for v_bx in T.thread_binding(batch_size, thread="blockIdx.x"):
|
|
for v_tx in T.thread_binding(
|
|
TX,
|
|
thread="threadIdx.x",
|
|
annotations={
|
|
"pragma_auto_unroll_max_step": 256,
|
|
"pragma_unroll_explicit": 1,
|
|
},
|
|
):
|
|
for i in range(add_local_size):
|
|
with T.sblock("T_add"):
|
|
bx = T.axis.spatial(batch_size, v_bx)
|
|
h = T.axis.spatial(hidden_size, i * TX + v_tx)
|
|
add_local[h // TX] = A[bx, 0, h] + B[bx, 0, h]
|
|
with T.sblock("T_write_back"):
|
|
bx = T.axis.spatial(batch_size, v_bx)
|
|
v_ax1 = T.axis.spatial(1, 0)
|
|
h = T.axis.spatial(hidden_size, i * TX + v_tx)
|
|
add[bx, v_ax1, h] = add_local[h // TX]
|
|
with T.sblock("T_multiply_red_rf_init"):
|
|
tx, bx = T.axis.remap("SS", [v_tx, v_bx])
|
|
sum_local[tx, bx, 0] = T.float32(0)
|
|
for v_i, _j in T.grid(add_local_size, 1):
|
|
with T.sblock("T_multiply_red_rf_update"):
|
|
tx, bx, i = T.axis.remap("SSR", [v_tx, v_bx, v_i])
|
|
sum_local[tx, bx, 0] += T.float32(add_local[i]) * T.float32(add_local[i])
|
|
for _j in range(1):
|
|
for v_tx_2 in T.thread_binding(TX, thread="threadIdx.x"):
|
|
with T.sblock("T_multiply_red"):
|
|
tx, bx = T.axis.remap("RS", [v_tx_2, v_bx])
|
|
T.reads(sum_local[tx, bx, 0])
|
|
T.writes(sum_shared[bx, 0])
|
|
with T.init():
|
|
sum_shared[bx, 0] = T.float32(0)
|
|
sum_shared[bx, 0] += sum_local[tx, bx, 0]
|
|
for i in range(add_local_size):
|
|
for v_tx_2 in T.thread_binding(TX, thread="threadIdx.x"):
|
|
with T.sblock("T_cast_2"):
|
|
bx = T.axis.spatial(batch_size, v_bx)
|
|
h = T.axis.spatial(hidden_size, i * TX + v_tx_2)
|
|
out[bx, 0, h] = T.cast(
|
|
T.rsqrt(sum_shared[bx, 0] * inv_hidden_size + eps)
|
|
* T.float32(add_local[h // TX])
|
|
* T.float32(C[h]),
|
|
dtype=in_dtype,
|
|
)
|
|
|
|
return decode_add_rms
|
|
|
|
|
|
def _get_add_rms_norm_prefill(hidden_size: int, eps: float, TX: int, in_dtype: str):
|
|
if in_dtype not in ("float16", "bfloat16"):
|
|
raise ValueError(f"Unsupported data type: {in_dtype}")
|
|
inv_hidden_size = T.float32(1.0 / float(hidden_size))
|
|
eps = T.float32(eps)
|
|
add_local_size = hidden_size // TX
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def prefill_add_rms(pA: T.handle, pB: T.handle, pC: T.handle, pO: T.handle, pAdd: T.handle):
|
|
T.func_attr({"tirx.noalias": True, "tirx.is_scheduled": 1})
|
|
seq_len = T.int32()
|
|
A = T.match_buffer(pA, (1, seq_len, hidden_size), in_dtype)
|
|
B = T.match_buffer(pB, (1, seq_len, hidden_size), in_dtype)
|
|
C = T.match_buffer(pC, (hidden_size,), in_dtype)
|
|
out = T.match_buffer(pO, (1, seq_len, hidden_size), in_dtype)
|
|
add = T.match_buffer(pAdd, (1, seq_len, hidden_size), in_dtype)
|
|
add_local = T.sblock_alloc_buffer((hidden_size // TX,), in_dtype, scope="local")
|
|
sum_shared = T.sblock_alloc_buffer((1, seq_len), scope="shared")
|
|
sum_local = T.sblock_alloc_buffer((TX, 1, seq_len), scope="local")
|
|
for v_bx in T.thread_binding(seq_len, thread="blockIdx.x"):
|
|
for v_tx in T.thread_binding(
|
|
TX,
|
|
thread="threadIdx.x",
|
|
annotations={
|
|
"pragma_auto_unroll_max_step": 256,
|
|
"pragma_unroll_explicit": 1,
|
|
},
|
|
):
|
|
for v_i in range(add_local_size):
|
|
with T.sblock("T_add"):
|
|
bx = T.axis.spatial(seq_len, v_bx)
|
|
h = T.axis.spatial(hidden_size, v_i * TX + v_tx)
|
|
add_local[h // TX] = A[0, bx, h] + B[0, bx, h]
|
|
with T.sblock("T_write_back"):
|
|
bx = T.axis.spatial(seq_len, v_bx)
|
|
h = T.axis.spatial(hidden_size, v_i * TX + v_tx)
|
|
add[0, bx, h] = add_local[h // TX]
|
|
with T.sblock("T_multiply_red_rf_init"):
|
|
tx, bx = T.axis.remap("SS", [v_tx, v_bx])
|
|
sum_local[tx, 0, bx] = T.float32(0)
|
|
for v_i, _j in T.grid(add_local_size, 1):
|
|
with T.sblock("T_multiply_red_rf_update"):
|
|
tx, bx, i = T.axis.remap("SSR", [v_tx, v_bx, v_i])
|
|
sum_local[tx, 0, bx] += T.float32(add_local[i]) * T.float32(add_local[i])
|
|
for _j in range(1):
|
|
for v_tx_2 in T.thread_binding(TX, thread="threadIdx.x"):
|
|
with T.sblock("T_multiply_red"):
|
|
tx, bx = T.axis.remap("RS", [v_tx_2, v_bx])
|
|
with T.init():
|
|
sum_shared[0, bx] = T.float32(0)
|
|
sum_shared[0, bx] = sum_shared[0, bx] + sum_local[tx, 0, bx]
|
|
for v_i in range(add_local_size):
|
|
for v_tx_2 in T.thread_binding(TX, thread="threadIdx.x"):
|
|
with T.sblock("T_cast_2"):
|
|
bx = T.axis.spatial(seq_len, v_bx)
|
|
v1 = T.axis.spatial(hidden_size, v_i * TX + v_tx_2)
|
|
out[0, bx, v1] = T.cast(
|
|
T.rsqrt(sum_shared[0, bx] * inv_hidden_size + eps)
|
|
* T.float32(add_local[v1 // TX])
|
|
* T.float32(C[v1]),
|
|
dtype=in_dtype,
|
|
)
|
|
|
|
return prefill_add_rms
|
|
|
|
|
|
@tvm.transform.module_pass(opt_level=0, name="FuseAddRMSNorm")
|
|
class FuseAddRMSNorm:
|
|
"""A compiler pass that fuses add + rms_norm."""
|
|
|
|
def __init__(self, target: tvm.target.Target) -> None:
|
|
"""Initializer.
|
|
|
|
Parameters
|
|
----------
|
|
target : tvm.target.Target
|
|
Target device.
|
|
"""
|
|
self.target = target
|
|
|
|
def transform_module(self, mod: tvm.IRModule, _ctx: tvm.transform.PassContext) -> tvm.IRModule:
|
|
"""IRModule-level transformation."""
|
|
return _FuseAddRMSNormRewriter(mod.clone(), self.target).transform()
|
|
|
|
|
|
@mutator
|
|
class _FuseAddRMSNormRewriter(PyExprMutator):
|
|
def __init__(self, mod: tvm.IRModule, target: tvm.target.Target):
|
|
super().__init__(mod)
|
|
self.mod = mod
|
|
self.prefill_norm_gv: Optional[tvm.ir.GlobalVar] = None
|
|
self.decode_norm_gv: Optional[tvm.ir.GlobalVar] = None
|
|
self.TX = min(1024, get_max_num_threads_per_block(target))
|
|
|
|
def transform(self) -> tvm.IRModule:
|
|
"""Entry point of the transformation"""
|
|
for g_var, func in self.mod.functions_items():
|
|
if not isinstance(func, relax.Function):
|
|
continue
|
|
new_func = self.visit_expr(func)
|
|
new_func = remove_all_unused(new_func)
|
|
self.builder_.update_func(g_var, new_func)
|
|
return self.builder_.finalize()
|
|
|
|
def visit_call_(self, call: relax.Call) -> relax.Expr:
|
|
call = super().visit_call_(call)
|
|
|
|
# Match the "rms_norm(add(x1, x2), w)" pattern
|
|
if call.op != tvm.ir.Op.get("relax.nn.rms_norm") or call.ty.dtype not in [
|
|
"bfloat16",
|
|
"float16",
|
|
]:
|
|
return call
|
|
assert len(call.args) == 2
|
|
weight = call.args[1]
|
|
eps = call.attrs.epsilon
|
|
assert isinstance(call.args[0], relax.Var)
|
|
y = self.lookup_binding(call.args[0])
|
|
if not isinstance(y, relax.Call) or y.op != tvm.ir.Op.get("relax.add"):
|
|
return call
|
|
assert len(y.args) == 2
|
|
x1 = y.args[0]
|
|
x2 = y.args[1]
|
|
# Extra check
|
|
n, _, h = x1.ty.shape
|
|
h = int(h)
|
|
if h % self.TX != 0:
|
|
return call
|
|
|
|
is_prefill = n == 1
|
|
func_gv = self.prefill_norm_gv if is_prefill else self.decode_norm_gv
|
|
if func_gv is None:
|
|
if is_prefill:
|
|
func_gv = self.builder_.add_func(
|
|
_get_add_rms_norm_prefill(h, eps, self.TX, call.ty.dtype),
|
|
"fuse_add_norm_prefill",
|
|
)
|
|
self.prefill_norm_gv = func_gv
|
|
else:
|
|
func_gv = self.builder_.add_func(
|
|
_get_add_rms_norm_decode(h, eps, self.TX, call.ty.dtype),
|
|
"fuse_add_norm_decode",
|
|
)
|
|
self.decode_norm_gv = func_gv
|
|
|
|
tuple_output = self.builder_.emit(
|
|
relax.call_tir(
|
|
func_gv,
|
|
[x1, x2, weight],
|
|
out_ty=[x1.ty, x2.ty],
|
|
)
|
|
)
|
|
new_o = relax.TupleGetItem(tuple_output, 0)
|
|
new_y = self.builder_.emit(relax.TupleGetItem(tuple_output, 1))
|
|
self.set_var_remap(call.args[0], new_y)
|
|
return new_o
|