ccf-ai-infra
/
Intro-ops
0
0
Fork 93
Code Issues Pull Requests 48 Actions Packages Projects Releases Wiki Activity

refactor: split tilelang student kernels 

Co-authored-by: wawahejun <hejunlbbc@gmail.com>
This commit is contained in:
yutianyu 2026-05-02 02:00:05 +08:00
parent 9db9af67b8
commit e4ce213d65
8 changed files with 128 additions and 114 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
ops/copy/tilelang/copy_tl.py
View File

@ -3,10 +3,11 @@ from __future__ import annotations
from functools import lru_cache
from dataclasses import dataclass
import tilelang
import tilelang.language as T
import torch
from ops.copy.tilelang.kernel import copy_kernel
def _tl_dtype(dtype: torch.dtype):
if dtype is torch.float16:
@ -15,29 +16,13 @@ def _tl_dtype(dtype: torch.dtype):
return T.float32
raise TypeError(f"unsupported TileLang dtype: {dtype}")
@tilelang.jit
def _copy_kernel(src, BLOCK_N: int, dtype):
N = T.const("N")
src: T.Tensor((N,), dtype)
out = T.empty((N,), dtype)
with T.Kernel(N // BLOCK_N, threads=256) as pid_n:
T.copy(
src[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N],
out[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N],
)
return out
def _block_n(n: int) -> int:
return 1024 if n % 1024 == 0 else n
@lru_cache(maxsize=32)
def _compiled_copy(n: int, block_n: int, dtype: torch.dtype):
return _copy_kernel.compile(N=n, BLOCK_N=block_n, dtype=_tl_dtype(dtype))
return copy_kernel.compile(N=n, BLOCK_N=block_n, dtype=_tl_dtype(dtype))
@dataclass

19
ops/copy/tilelang/kernel.py Normal file
View File

@ -0,0 +1,19 @@
from __future__ import annotations
import tilelang
import tilelang.language as T
@tilelang.jit
def copy_kernel(src, BLOCK_N: int, dtype):
N = T.const("N")
src: T.Tensor((N,), dtype)
out = T.empty((N,), dtype)
with T.Kernel(N // BLOCK_N, threads=256) as pid_n:
T.copy(
src[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N],
out[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N],
)
return out

30
ops/reduce_sum/tilelang/kernel.py Normal file
View File

@ -0,0 +1,30 @@
from __future__ import annotations
import tilelang
import tilelang.language as T
@tilelang.jit(
pass_configs={
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
},
)
def reduce_sum_kernel(src, BLOCK_N: int, BLOCK_M: int):
N, M = T.const("N, M")
dtype = T.float32
src: T.Tensor((N, M), dtype)
out = T.empty((N,), dtype)
with T.Kernel(N // BLOCK_N, threads=256) as pid_n:
src_local = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
out_local = T.alloc_fragment((BLOCK_N,), dtype)
T.clear(out_local)
for m_blk in T.Serial(M // BLOCK_M):
T.copy(src[pid_n * BLOCK_N, m_blk * BLOCK_M], src_local)
T.reduce_sum(src_local, out_local, dim=1, clear=False)
T.copy(out_local, out[pid_n * BLOCK_N])
return out

31
ops/reduce_sum/tilelang/reduce_sum_tl.py
View File

@ -3,40 +3,13 @@ from __future__ import annotations
from dataclasses import dataclass
from functools import lru_cache
import tilelang
import tilelang.language as T
import torch
@tilelang.jit(
pass_configs={
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
},
)
def _reduce_sum_kernel(src, BLOCK_N: int, BLOCK_M: int):
N, M = T.const("N, M")
dtype = T.float32
src: T.Tensor((N, M), dtype)
out = T.empty((N,), dtype)
with T.Kernel(N // BLOCK_N, threads=256) as pid_n:
src_local = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
out_local = T.alloc_fragment((BLOCK_N,), dtype)
T.clear(out_local)
for m_blk in T.Serial(M // BLOCK_M):
T.copy(src[pid_n * BLOCK_N, m_blk * BLOCK_M], src_local)
T.reduce_sum(src_local, out_local, dim=1, clear=False)
T.copy(out_local, out[pid_n * BLOCK_N])
return out
from ops.reduce_sum.tilelang.kernel import reduce_sum_kernel
@lru_cache(maxsize=32)
def _compiled_reduce_sum(n: int, m: int, block_n: int, block_m: int):
return _reduce_sum_kernel.compile(N=n, M=m, BLOCK_N=block_n, BLOCK_M=block_m)
return reduce_sum_kernel.compile(N=n, M=m, BLOCK_N=block_n, BLOCK_M=block_m)
def _blocks(n: int, m: int) -> tuple[int, int]:

50
ops/softmax/tilelang/kernel.py Normal file
View File

@ -0,0 +1,50 @@
from __future__ import annotations
import tilelang
import tilelang.language as T
@tilelang.jit(
pass_configs={
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: 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)
with T.Kernel(N // BLOCK_N, threads=256) as pid_n:
src_local = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
out_local = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
cur_exp = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
cur_max = T.alloc_fragment((BLOCK_N,), dtype)
cur_sum = T.alloc_fragment((BLOCK_N,), dtype)
lse = T.alloc_fragment((BLOCK_N,), dtype)
T.fill(lse, -T.infinity(dtype))
for m_blk in T.Serial(M // BLOCK_M):
T.copy(src[pid_n * BLOCK_N, m_blk * BLOCK_M], src_local)
T.reduce_max(src_local, cur_max, dim=1, clear=True)
for i, j in T.Parallel(BLOCK_N, BLOCK_M):
cur_exp[i, j] = T.exp2(src_local[i, j] * log2_e - cur_max[i] * log2_e)
T.reduce_sum(cur_exp, cur_sum, dim=1, clear=True)
for i in T.Parallel(BLOCK_N):
lse[i] = cur_max[i] * log2_e + T.log2(
T.exp2(lse[i] - cur_max[i] * log2_e) + cur_sum[i]
)
for m_blk in T.Serial(M // BLOCK_M):
T.copy(src[pid_n * BLOCK_N, m_blk * BLOCK_M], src_local)
for i, j in T.Parallel(BLOCK_N, BLOCK_M):
out_local[i, j] = T.exp2(src_local[i, j] * log2_e - lse[i])
T.copy(out_local, out[pid_n * BLOCK_N, m_blk * BLOCK_M])
return out

51
ops/softmax/tilelang/softmax_tl.py
View File

@ -3,60 +3,13 @@ from __future__ import annotations
from dataclasses import dataclass
from functools import lru_cache
import tilelang
import tilelang.language as T
import torch
@tilelang.jit(
pass_configs={
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: 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)
with T.Kernel(N // BLOCK_N, threads=256) as pid_n:
src_local = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
out_local = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
cur_exp = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype)
cur_max = T.alloc_fragment((BLOCK_N,), dtype)
cur_sum = T.alloc_fragment((BLOCK_N,), dtype)
lse = T.alloc_fragment((BLOCK_N,), dtype)
T.fill(lse, -T.infinity(dtype))
for m_blk in T.Serial(M // BLOCK_M):
T.copy(src[pid_n * BLOCK_N, m_blk * BLOCK_M], src_local)
T.reduce_max(src_local, cur_max, dim=1, clear=True)
for i, j in T.Parallel(BLOCK_N, BLOCK_M):
cur_exp[i, j] = T.exp2(src_local[i, j] * log2_e - cur_max[i] * log2_e)
T.reduce_sum(cur_exp, cur_sum, dim=1, clear=True)
for i in T.Parallel(BLOCK_N):
lse[i] = cur_max[i] * log2_e + T.log2(
T.exp2(lse[i] - cur_max[i] * log2_e) + cur_sum[i]
)
for m_blk in T.Serial(M // BLOCK_M):
T.copy(src[pid_n * BLOCK_N, m_blk * BLOCK_M], src_local)
for i, j in T.Parallel(BLOCK_N, BLOCK_M):
out_local[i, j] = T.exp2(src_local[i, j] * log2_e - lse[i])
T.copy(out_local, out[pid_n * BLOCK_N, m_blk * BLOCK_M])
return out
from ops.softmax.tilelang.kernel import softmax_kernel
@lru_cache(maxsize=32)
def _compiled_softmax(n: int, m: int, block_n: int, block_m: int):
return _softmax_kernel.compile(N=n, M=m, BLOCK_N=block_n, BLOCK_M=block_m)
return softmax_kernel.compile(N=n, M=m, BLOCK_N=block_n, BLOCK_M=block_m)
def _blocks(n: int, m: int) -> tuple[int, int]:

19
ops/vector_add/tilelang/kernel.py Normal file
View File

@ -0,0 +1,19 @@
from __future__ import annotations
import tilelang
import tilelang.language as T
@tilelang.jit
def vector_add_kernel(a, b, BLOCK_N: int, dtype):
N = T.const("N")
a: T.Tensor((N,), dtype)
b: T.Tensor((N,), dtype)
out = T.empty((N,), dtype)
with T.Kernel(N // BLOCK_N, threads=256) as pid_n:
base = pid_n * BLOCK_N
for i in T.Parallel(BLOCK_N):
out[base + i] = a[base + i] + b[base + i]
return out

21
ops/vector_add/tilelang/vector_add_tl.py
View File

@ -3,10 +3,11 @@ from __future__ import annotations
from dataclasses import dataclass
from functools import lru_cache
import tilelang
import tilelang.language as T
import torch
from ops.vector_add.tilelang.kernel import vector_add_kernel
def _tl_dtype(dtype: torch.dtype):
if dtype is torch.float16:
@ -15,29 +16,13 @@ def _tl_dtype(dtype: torch.dtype):
return T.float32
raise TypeError(f"unsupported TileLang dtype: {dtype}")
@tilelang.jit
def _vector_add_kernel(a, b, BLOCK_N: int, dtype):
N = T.const("N")
a: T.Tensor((N,), dtype)
b: T.Tensor((N,), dtype)
out = T.empty((N,), dtype)
with T.Kernel(N // BLOCK_N, threads=256) as pid_n:
base = pid_n * BLOCK_N
for i in T.Parallel(BLOCK_N):
out[base + i] = a[base + i] + b[base + i]
return out
def _block_n(n: int) -> int:
return 1024 if n % 1024 == 0 else n
@lru_cache(maxsize=32)
def _compiled_vector_add(n: int, block_n: int, dtype: torch.dtype):
return _vector_add_kernel.compile(N=n, BLOCK_N=block_n, dtype=_tl_dtype(dtype))
return vector_add_kernel.compile(N=n, BLOCK_N=block_n, dtype=_tl_dtype(dtype))
@dataclass