forked from ccf-ai-infra/Intro-ops
96 lines
3.4 KiB
Python
96 lines
3.4 KiB
Python
from __future__ import annotations
|
|
from json.encoder import INFINITY
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
|
|
|
|
@tilelang.jit(
|
|
pass_configs={
|
|
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
|
},
|
|
)
|
|
def softmax_kernel(src, BLOCK_N: int, BLOCK_M: int):
|
|
log2_e = 1.44269504
|
|
N, M = T.const("N, M")
|
|
dtype = T.float32
|
|
src: T.Tensor((N, M), dtype)
|
|
out = T.empty((N, M), dtype)
|
|
|
|
# TODO: implement a tiled row-wise softmax kernel.
|
|
#
|
|
# Suggested steps:
|
|
# 1. Launch one TileLang kernel over row tiles.
|
|
# 2. Allocate fragments for src, out, temporary exp values, row max, row sum, and lse.
|
|
# 3. Initialize the running log-sum-exp state.
|
|
# 4. In a first T.Serial loop over column tiles:
|
|
# - copy the input tile into a fragment
|
|
# - reduce to get the tile max
|
|
# - compute exp2-based temporary values
|
|
# - reduce to get the tile sum
|
|
# - update the running lse
|
|
# 5. In a second T.Serial loop over column tiles:
|
|
# - copy the input tile again
|
|
# - normalize with the final lse
|
|
# - copy the result tile to global memory
|
|
with T.Kernel(T.ceildiv(N, BLOCK_N), threads=256) as bx:
|
|
src_frag = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
|
|
temp_exp_frag = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
|
|
out_frag = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
|
|
tile_max_frag = T.alloc_fragment((BLOCK_N,), dtype)
|
|
tile_sum_frag = T.alloc_fragment((BLOCK_N,), dtype)
|
|
tile_lse_frag = T.alloc_fragment((BLOCK_N,), dtype)
|
|
lse2_frag = T.alloc_fragment((BLOCK_N,), dtype)
|
|
|
|
for i in T.Parallel(BLOCK_N):
|
|
lse2_frag[i] = -T.infinity(dtype)
|
|
|
|
row_base = bx * BLOCK_N
|
|
# First pass: compute final lse
|
|
for k in T.Serial(T.ceildiv(M, BLOCK_M)):
|
|
col_base = k * BLOCK_M
|
|
T.copy(src[row_base : row_base + BLOCK_N,
|
|
col_base : col_base + BLOCK_M],
|
|
src_frag
|
|
)
|
|
|
|
T.reduce_max(src_frag, tile_max_frag, dim=1)
|
|
|
|
for i, j in T.Parallel(BLOCK_N, BLOCK_M):
|
|
temp_exp_frag[i, j] = T.exp2(
|
|
(src_frag[i, j] - tile_max_frag[i]) * log2_e
|
|
)
|
|
|
|
T.reduce_sum(temp_exp_frag, tile_sum_frag, dim=1)
|
|
|
|
for i in T.Parallel(BLOCK_N):
|
|
tile_lse_frag[i] = (
|
|
tile_max_frag[i] * log2_e
|
|
+ T.log2(tile_sum_frag[i])
|
|
)
|
|
for i in T.Parallel(BLOCK_N):
|
|
base = T.max(lse2_frag[i], tile_lse_frag[i])
|
|
lse2_frag[i] = base + T.log2(
|
|
T.exp2(lse2_frag[i] - base)
|
|
+ T.exp2(tile_lse_frag[i] - base)
|
|
)
|
|
|
|
# Second pass: normalization
|
|
for k in T.Serial(T.ceildiv(M, BLOCK_M)):
|
|
col_base = k * BLOCK_M
|
|
T.copy(src[row_base : row_base + BLOCK_N,
|
|
col_base : col_base + BLOCK_M],
|
|
src_frag
|
|
)
|
|
|
|
for i, j in T.Parallel(BLOCK_N, BLOCK_M):
|
|
out_frag[i, j] = T.exp2(
|
|
src_frag[i, j] * log2_e
|
|
- lse2_frag[i]
|
|
)
|
|
|
|
T.copy(out_frag,
|
|
out[row_base : row_base + BLOCK_N,
|
|
col_base : col_base + BLOCK_M])
|
|
return out
|