mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-11-01 08:23:24 +08:00
Cunxiao NiandLei Wang
3db18726c7
* [Example] Update kernel compilation in examples to use @tilelang.jit - Refactored multiple examples to eliminate the use of `tilelang.compile` for kernel creation, directly invoking the functions instead. - Added `@tilelang.jit` decorators with appropriate output indices to enhance performance and maintainability. - Improved code clarity by simplifying the kernel invocation process across various examples, ensuring consistency in how kernels are defined and executed. * format * Update example_tilelang_sparse_gqa_decode_varlen_indice.py * Update example_dequant_gemm_fine_grained.py * Update example_gemm_autotune.py --------- Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com>
70 lines
1.9 KiB
Python
70 lines
1.9 KiB
Python
import tilelang
|
|
import tilelang.language as T
|
|
|
|
|
|
@tilelang.jit(out_idx=[-1])
|
|
def matmul(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float"):
|
|
|
|
@T.prim_func
|
|
def gemm_schedule(
|
|
A: T.Tensor((M, K), dtype),
|
|
B: T.Tensor((K, N), dtype),
|
|
C: T.Tensor((M, N), dtype),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
|
|
A_shared = T.alloc_shared((block_M, block_K), dtype)
|
|
B_shared = T.alloc_shared((block_K, block_N), dtype)
|
|
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
|
|
|
# Enable rasterization for better L2 Cache Locality
|
|
T.use_swizzle(panel_size=10)
|
|
|
|
# Clear the local buffer
|
|
T.clear(C_local)
|
|
|
|
# Auto pipeline the computation
|
|
for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
|
|
T.copy(A[by * block_M, ko * block_K], A_shared)
|
|
|
|
# Instead of using
|
|
# T.copy(B[k * block_K, bx * block_N], B_shared)
|
|
# we can also use Parallel to auto map the thread
|
|
# bindings and vectorize the copy operation.
|
|
for k, j in T.Parallel(block_K, block_N):
|
|
B_shared[k, j] = B[ko * block_K + k, bx * block_N + j]
|
|
|
|
T.gemm(A_shared, B_shared, C_local)
|
|
|
|
T.copy(C_local, C[by * block_M, bx * block_N])
|
|
|
|
return gemm_schedule
|
|
|
|
|
|
def main():
|
|
kernel = matmul(1024, 1024, 1024, 128, 128, 32)
|
|
|
|
import torch
|
|
|
|
a = torch.randn(1024, 1024).cuda().half()
|
|
b = torch.randn(1024, 1024).cuda().half()
|
|
|
|
c = kernel(a, b)
|
|
|
|
ref_c = a @ b
|
|
|
|
print("c:")
|
|
print(c)
|
|
print("ref_c:")
|
|
print(ref_c)
|
|
|
|
torch.testing.assert_close(c, ref_c, rtol=1e-2, atol=1e-2)
|
|
print("All check passed.")
|
|
|
|
# Get CUDA Source
|
|
print("CUDA Source:")
|
|
print(kernel.get_kernel_source())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|