mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-23 07:00:29 +08:00
e9a608e222
* fix CI bugs in hopper * lint fix * Update bulk_copy.cc * Refactor bulk copy logic in LowerBulkCopy function - Removed unnecessary blank lines for improved code readability. - Enhanced stride validation by checking for null pointers in global stride calculations, ensuring robustness against symbolic strides. - Updated pass configuration handling in dynamic tile language tests to streamline dynamic alignment and TMA lower pass settings. * test fix * ci fix * Update flash-attention dependencies and clean up example code - Downgraded `flash-attn` dependency version in `requirements-test.txt` to `<=2.2.0`. - Removed unused imports and commented-out code in various example files to enhance readability and maintainability. - Updated the `flashattn` function signature to include default parameters for `block_M`, `block_N`, `num_stages`, and `threads`. - Cleaned up the `example_mha_fwd_varlen.py` and `example_mha_bwd_wgmma_pipelined.py` files by removing unnecessary comments and improving code clarity. - Deleted the `example_mha_inference.py` file as it is no longer needed. * Update CI workflow to remove `--user` flag from pip install commands - Removed the `--user` flag from the pip install commands in both the development and testing sections of the CI workflow to ensure proper installation of dependencies in the virtual environment. * Update CI workflow to include `--no-user` flag in pip install commands - Added the `--no-user` flag to the pip install commands in both the development and testing sections of the CI workflow to ensure dependencies are installed correctly within the virtual environment. * Update CI workflow to include `--no-user` flag in pip install command for wheel mode - Added the `--no-user` flag to the pip install command in the wheel mode section of the CI workflow to ensure dependencies are installed correctly within the virtual environment. * test fix * avoid conflict with system environments * test fix * add commnets --------- Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com> Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>
318 lines
10 KiB
Python
318 lines
10 KiB
Python
import tilelang
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
import torch
|
|
|
|
|
|
def ref_program(A, B, BlockMask, block_M, block_N, block_K):
|
|
M, K = A.shape
|
|
N = B.shape[1]
|
|
ref_c = torch.zeros((M, N), dtype=torch.float16, device=A.device)
|
|
for i in range(M // block_M):
|
|
for j in range(N // block_N):
|
|
accu = torch.zeros((block_M, block_N), dtype=torch.float32, device=A.device)
|
|
for k in range(K // block_K):
|
|
if torch.any(BlockMask[i, j, k]):
|
|
accu += A[i * block_M:(i + 1) * block_M, k * block_K:(k + 1) * block_K].to(
|
|
torch.float32) @ B[k * block_K:(k + 1) * block_K,
|
|
j * block_N:(j + 1) * block_N].to(torch.float32)
|
|
ref_c[i * block_M:(i + 1) * block_M, j * block_N:(j + 1) * block_N] = (
|
|
accu.to(torch.float16))
|
|
return ref_c
|
|
|
|
|
|
def blocksparse_matmul_global(
|
|
M,
|
|
N,
|
|
K,
|
|
condition_dim,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
num_stages,
|
|
thread_num,
|
|
enable_rasteration,
|
|
dtype="float16",
|
|
accum_dtype="float",
|
|
):
|
|
|
|
block_mask_shape = (M // block_M, N // block_N, K // block_K, condition_dim)
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((M, K), dtype),
|
|
B: T.Tensor((K, N), dtype),
|
|
BlockMask: T.Tensor(block_mask_shape, "bool"),
|
|
C: T.Tensor((M, N), dtype),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=thread_num) 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)
|
|
C_shared = T.alloc_shared((block_M, block_N), dtype)
|
|
|
|
T.use_swizzle(panel_size=10, enable=enable_rasteration)
|
|
T.clear(C_local)
|
|
|
|
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
|
|
if T.any_of(BlockMask[by, bx, k, :]):
|
|
T.copy(A[by * block_M, k * block_K], A_shared)
|
|
T.copy(B[k * block_K, bx * block_N], B_shared)
|
|
T.gemm(A_shared, B_shared, C_local)
|
|
|
|
T.copy(C_local, C_shared)
|
|
T.copy(C_shared, C[by * block_M, bx * block_N])
|
|
|
|
return main
|
|
|
|
|
|
def blocksparse_matmul_shared(
|
|
M,
|
|
N,
|
|
K,
|
|
condition_dim,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
num_stages,
|
|
thread_num,
|
|
enable_rasteration,
|
|
dtype="float16",
|
|
accum_dtype="float",
|
|
):
|
|
|
|
block_mask_shape = (M // block_M, N // block_N, K // block_K, condition_dim)
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((M, K), dtype),
|
|
B: T.Tensor((K, N), dtype),
|
|
BlockMask: T.Tensor(block_mask_shape, "bool"),
|
|
C: T.Tensor((M, N), dtype),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=thread_num) 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)
|
|
block_mask_shared = T.alloc_shared(condition_dim, "bool")
|
|
C_shared = T.alloc_shared((block_M, block_N), dtype)
|
|
|
|
T.use_swizzle(panel_size=10, enable=enable_rasteration)
|
|
T.clear(C_local)
|
|
|
|
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=2):
|
|
for i in T.serial(condition_dim):
|
|
block_mask_shared[i] = BlockMask[by, bx, k, i]
|
|
# or T.any_of(block_mask_local[0:condition_dim])
|
|
# or T.any_of(block_mask_local[:])
|
|
if T.any_of(block_mask_shared):
|
|
T.copy(A[by * block_M, k * block_K], A_shared)
|
|
T.copy(B[k * block_K, bx * block_N], B_shared)
|
|
T.gemm(A_shared, B_shared, C_local)
|
|
|
|
T.copy(C_local, C_shared)
|
|
T.copy(C_shared, C[by * block_M, bx * block_N])
|
|
|
|
return main
|
|
|
|
|
|
def blocksparse_matmul_local(
|
|
M,
|
|
N,
|
|
K,
|
|
condition_dim,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
num_stages,
|
|
thread_num,
|
|
enable_rasteration,
|
|
dtype="float16",
|
|
accum_dtype="float",
|
|
):
|
|
|
|
block_mask_shape = (M // block_M, N // block_N, K // block_K, condition_dim)
|
|
|
|
@T.prim_func
|
|
def main(
|
|
A: T.Tensor((M, K), dtype),
|
|
B: T.Tensor((K, N), dtype),
|
|
BlockMask: T.Tensor(block_mask_shape, "bool"),
|
|
C: T.Tensor((M, N), dtype),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=thread_num) 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)
|
|
block_mask_local = T.alloc_local(condition_dim, "bool")
|
|
C_shared = T.alloc_shared((block_M, block_N), dtype)
|
|
|
|
T.use_swizzle(panel_size=10, enable=enable_rasteration)
|
|
T.clear(C_local)
|
|
|
|
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=2):
|
|
for i in T.serial(condition_dim):
|
|
block_mask_local[i] = BlockMask[by, bx, k, i]
|
|
# or T.any_of(block_mask_local[0:condition_dim])
|
|
# or T.any_of(block_mask_local[:])
|
|
if T.any_of(block_mask_local):
|
|
T.copy(A[by * block_M, k * block_K], A_shared)
|
|
T.copy(B[k * block_K, bx * block_N], B_shared)
|
|
T.gemm(A_shared, B_shared, C_local)
|
|
|
|
T.copy(C_local, C_shared)
|
|
T.copy(C_shared, C[by * block_M, bx * block_N])
|
|
|
|
return main
|
|
|
|
|
|
def run_block_sparse_matmul_global(M=1024, N=1024, K=1024, sparsity=0.5, condition_dim=2):
|
|
block_M = 128
|
|
block_N = 128
|
|
block_K = 32
|
|
num_stages = 2
|
|
thread_num = 128
|
|
enable_rasteration = True
|
|
|
|
# Initialize input matrices A and B on the GPU with half precision
|
|
a = torch.randn(M, K).cuda().half()
|
|
b = torch.randn(K, N).cuda().half()
|
|
|
|
func = blocksparse_matmul_global(
|
|
M,
|
|
N,
|
|
K,
|
|
condition_dim,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
num_stages,
|
|
thread_num,
|
|
enable_rasteration,
|
|
)
|
|
kernel = tilelang.compile(func, out_idx=-1)
|
|
# Create block mask with desired sparsity
|
|
mask_shape = (M // block_M, N // block_N, K // block_K)
|
|
block_mask = torch.rand(mask_shape).cuda() > sparsity
|
|
block_mask = block_mask.view(mask_shape + (1,)).repeat(1, 1, 1, condition_dim)
|
|
# random set the last dimension to be False
|
|
block_mask[:, :, :, 0] = False
|
|
|
|
# Run the compiled kernel (either tuned or default) with the inputs
|
|
c = kernel(a, b, block_mask)
|
|
|
|
# Compute the reference result using the naive PyTorch implementation
|
|
ref_c = ref_program(a, b, block_mask, block_M, block_N, block_K)
|
|
|
|
torch.testing.assert_close(c, ref_c, rtol=1e-2, atol=1e-2)
|
|
|
|
|
|
def run_block_sparse_matmul_shared(M=1024, N=1024, K=1024, sparsity=0.5, condition_dim=2):
|
|
block_M = 128
|
|
block_N = 128
|
|
block_K = 32
|
|
num_stages = 2
|
|
thread_num = 128
|
|
enable_rasteration = True
|
|
|
|
# Initialize input matrices A and B on the GPU with half precision
|
|
a = torch.randn(M, K).cuda().half()
|
|
b = torch.randn(K, N).cuda().half()
|
|
|
|
func = blocksparse_matmul_shared(
|
|
M,
|
|
N,
|
|
K,
|
|
condition_dim,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
num_stages,
|
|
thread_num,
|
|
enable_rasteration,
|
|
)
|
|
kernel = tilelang.compile(
|
|
func,
|
|
out_idx=-1,
|
|
pass_configs={
|
|
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
|
|
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
|
})
|
|
# Create block mask with desired sparsity
|
|
mask_shape = (M // block_M, N // block_N, K // block_K)
|
|
block_mask = torch.rand(mask_shape).cuda() > sparsity
|
|
block_mask = block_mask.view(mask_shape + (1,)).repeat(1, 1, 1, condition_dim)
|
|
# random set the last dimension to be False
|
|
block_mask[:, :, :, 0] = False
|
|
|
|
# Run the compiled kernel (either tuned or default) with the inputs
|
|
c = kernel(a, b, block_mask)
|
|
|
|
# Compute the reference result using the naive PyTorch implementation
|
|
ref_c = ref_program(a, b, block_mask, block_M, block_N, block_K)
|
|
|
|
torch.testing.assert_close(c, ref_c, rtol=1e-2, atol=1e-2)
|
|
|
|
|
|
def run_block_sparse_matmul_local(M=1024, N=1024, K=1024, sparsity=0.5, condition_dim=2):
|
|
block_M = 128
|
|
block_N = 128
|
|
block_K = 32
|
|
num_stages = 2
|
|
thread_num = 128
|
|
enable_rasteration = True
|
|
|
|
# Initialize input matrices A and B on the GPU with half precision
|
|
a = torch.randn(M, K).cuda().half()
|
|
b = torch.randn(K, N).cuda().half()
|
|
|
|
func = blocksparse_matmul_local(
|
|
M,
|
|
N,
|
|
K,
|
|
condition_dim,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
num_stages,
|
|
thread_num,
|
|
enable_rasteration,
|
|
)
|
|
kernel = tilelang.compile(
|
|
func,
|
|
out_idx=-1,
|
|
pass_configs={
|
|
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
|
|
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
|
})
|
|
# Create block mask with desired sparsity
|
|
mask_shape = (M // block_M, N // block_N, K // block_K)
|
|
block_mask = torch.rand(mask_shape).cuda() > sparsity
|
|
block_mask = block_mask.view(mask_shape + (1,)).repeat(1, 1, 1, condition_dim)
|
|
# random set the last dimension to be False
|
|
block_mask[:, :, :, 0] = False
|
|
|
|
# Run the compiled kernel (either tuned or default) with the inputs
|
|
c = kernel(a, b, block_mask)
|
|
|
|
# Compute the reference result using the naive PyTorch implementation
|
|
ref_c = ref_program(a, b, block_mask, block_M, block_N, block_K)
|
|
|
|
torch.testing.assert_close(c, ref_c, rtol=1e-2, atol=1e-2)
|
|
|
|
|
|
def test_block_sparse_matmul_global():
|
|
run_block_sparse_matmul_global(M=1024, N=1024, K=1024, sparsity=0.5, condition_dim=2)
|
|
|
|
|
|
def test_block_sparse_matmul_shared():
|
|
run_block_sparse_matmul_shared(M=1024, N=1024, K=1024, sparsity=0.5, condition_dim=2)
|
|
|
|
|
|
def test_block_sparse_matmul_local():
|
|
run_block_sparse_matmul_local(M=1024, N=1024, K=1024, sparsity=0.5, condition_dim=2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|