mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-20 06:30:07 +08:00
Lei Wang
bf8a6fc142
* [Refactor] Improve flash attention example and layout comparison logic - Removed unnecessary annotation for `lse_local_split` in the flash attention example to streamline the code. - Updated the handling of `lse_local_split` to utilize parallel processing for better performance. - Refactored kernel compilation and profiling logic to enhance clarity and maintainability in the flash attention example. - Added a condition in `FragmentNode::IsEqual` to handle broadcast cases, improving the robustness of layout comparisons. * lint fix * [Enhancement] Add support for shared memory scope in Fill operation - Introduced handling for `shared.dyn` and `shared` memory scopes in the Fill operation. - Implemented parallel operation and layout inference for improved performance in shared memory scenarios. - Updated thread loop partitioning and vectorization logic to accommodate new memory scope handling. * [Refactor] Remove deprecated decorator and enhance Cython kernel handling - Removed the deprecated decorator from the main module and added a new implementation in the utils module for better organization. - Introduced a pointer map in the Cython kernel adapter to manage pointer arguments, improving runtime shape resolution. - Updated the Cython kernel wrapper to utilize the new pointer map for handling kernel arguments. - Enhanced error checking in the tensor utility functions to ensure static shapes are enforced. - Added a new proxy module for buffer and tensor handling, streamlining the interface for TIR programs. * [Feature] Add matrix multiplication test and kernel implementation - Introduced a new test file `test_tilelang_language_ptr.py` that implements a matrix multiplication function using TileLang's primitives. - The `matmul_test` function defines a kernel for performing tile-level GEMM operations with customizable block sizes and data types. - Added a `run_matmul` function to compile and execute the kernel, along with a test function to validate the implementation. - Updated the `proxy.py` file to enhance type handling for buffer and tensor proxies, ensuring compatibility with TIR programs. - Minor formatting improvements in `deprecated.py` for better readability. * lint fix * [Refactor] Update tensor creation in matrix multiplication test - Replaced `T.Tensor.from_ptr` with `T.make_tensor` in `matmul_test` for improved clarity and consistency. - Updated imports in `__init__.py` to include `make_tensor`. - Added `make_tensor` function in `proxy.py` to streamline tensor creation from pointers. * [Refactor] Update tensor definitions across multiple files - Replaced instances of `T.Tensor` with updated tensor definitions in various benchmark and example files to enhance consistency and clarity. - Adjusted tensor shapes and types in functions related to matrix multiplication, attention mechanisms, and other operations. - Improved documentation in README and example files to reflect changes in tensor usage. * lint fix * [Refactor] Update tensor types in attention and matrix multiplication examples - Replaced instances of `T.Tensor` with `T.SharedTensor` and `T.FragmentTensor` in various attention and matrix multiplication functions to improve consistency and clarity. - Adjusted tensor definitions in benchmark and example files to align with the new tensor types. - Enhanced the overall structure and readability of the code by standardizing tensor usage across multiple files. * lint fix * [Refactor] Update tensor types in GEMM example and test files - Replaced instances of `T.Tensor` with `T.LocalTensor` and `T.Buffer` in the GEMM example and related test functions to improve consistency and clarity. - Enhanced the overall structure of the code by standardizing tensor usage across multiple files, aligning with recent updates in tensor definitions. * [Refactor] Update tensor usage in customize.py - Replaced instances of `T.Tensor` with `T.Buffer` in the `reshape` and `view` functions to enhance consistency with recent tensor definitions. - Improved code clarity by standardizing buffer usage across the file. * [Refactor] Update tensor types in test_tilelang_transform_annotate_device_regions.py - Replaced instances of `T.Tensor` with `T.Buffer` in the `before` and `expected` methods of the `TestAnnotateThreadExtent` and `TestAnnotateDeviceScope` classes to enhance consistency with recent tensor definitions. - Improved code clarity by standardizing buffer usage across the test file. * [Refactor] Update tensor types to SharedBuffer and FragmentBuffer - Replaced instances of `T.SharedTensor` and `T.FragmentTensor` with `T.SharedBuffer` and `T.FragmentBuffer` across multiple benchmark, example, and test files to enhance consistency with recent tensor definitions. - Improved code clarity and structure by standardizing buffer usage in attention and matrix multiplication functions. * [Refactor] Introduce Tensor alias for Buffer in proxy.py - Added a new alias `Tensor` for `Buffer` in `proxy.py` to facilitate JIT compilation, ensuring that inputs and outputs are mapped with `torch.Tensor`. - This change enhances clarity and consistency in tensor usage across the codebase.
350 lines
15 KiB
Python
350 lines
15 KiB
Python
from tilelang import tvm as tvm
|
|
import tilelang.testing
|
|
import tilelang.language as T
|
|
|
|
tilelang.testing.set_random_seed(0)
|
|
|
|
|
|
def chunk_scan_fwd(batch, seqlen, chunk_size, ngroups, nheads, headdim, dstate, block_M, block_N,
|
|
block_K, block_Dstate, num_stages, threads):
|
|
dtype = "float16"
|
|
accum_dtype = "float"
|
|
nchunks = T.ceildiv(seqlen, chunk_size)
|
|
p = 1.44269504
|
|
|
|
@T.prim_func
|
|
def main(cb: T.Tensor((batch, nchunks, ngroups, chunk_size, chunk_size), dtype), x: T.Tensor(
|
|
(batch, seqlen, nheads, headdim), dtype), dt: T.Tensor(
|
|
(batch, nheads, nchunks, chunk_size), dtype), dA_cumsum: T.Tensor(
|
|
(batch, nheads, nchunks, chunk_size), dtype),
|
|
C: T.Tensor((batch, seqlen, ngroups, dstate), dtype), prev_states: T.Tensor(
|
|
(batch, nchunks, nheads, headdim, dstate), dtype), D: T.Tensor(
|
|
(nheads), dtype), Output: T.Tensor((batch, seqlen, nheads, headdim), dtype)):
|
|
with T.Kernel(
|
|
nheads,
|
|
T.ceildiv(chunk_size, block_M) * T.ceildiv(headdim, block_N),
|
|
batch * nchunks,
|
|
threads=threads) as (bz, bx, by):
|
|
acc_o = T.alloc_fragment((block_M, block_N), accum_dtype)
|
|
cb_shared = T.alloc_shared((block_M, block_K), dtype, scope="shared.dyn")
|
|
cb_local = T.alloc_fragment((block_M, block_K), dtype)
|
|
dA_cs_k_shared = T.alloc_shared((block_K), dtype, scope="shared")
|
|
dA_cs_k_local = T.alloc_fragment((block_K), accum_dtype)
|
|
dA_cs_m_local = T.alloc_fragment((block_M), accum_dtype)
|
|
dt_shared = T.alloc_shared((block_K), dtype, scope="shared")
|
|
dt_local = T.alloc_fragment((block_K), accum_dtype)
|
|
x_shared = T.alloc_shared((block_K, block_N), dtype, scope="shared.dyn")
|
|
dA_cs_m_shared = T.alloc_shared((block_M), dtype, scope="shared")
|
|
scale_m_local = T.alloc_fragment((block_M), accum_dtype)
|
|
C_shared = T.alloc_shared((block_M, block_Dstate), dtype)
|
|
prev_state_shared = T.alloc_shared((block_N, block_Dstate), dtype)
|
|
D_local = T.alloc_fragment((1), accum_dtype)
|
|
x_residual_shared = T.alloc_shared((block_M, block_N), dtype, scope="shared.dyn")
|
|
x_residual_local = T.alloc_fragment((block_M, block_N), accum_dtype)
|
|
|
|
batch_idx = by % batch
|
|
chunk_idx = by // batch
|
|
# m: chunk_size
|
|
# n : headdim
|
|
m_idx = bx // T.ceildiv(headdim, block_N)
|
|
n_idx = bx % T.ceildiv(headdim, block_N)
|
|
|
|
T.copy(dA_cumsum[batch_idx, bz, chunk_idx, m_idx * block_M:(m_idx + 1) * block_M],
|
|
dA_cs_m_shared)
|
|
T.copy(dA_cs_m_shared, dA_cs_m_local)
|
|
T.clear(acc_o)
|
|
|
|
for i in T.Parallel(block_M):
|
|
scale_m_local[i] = T.exp2(dA_cs_m_local[i] * p)
|
|
T.copy(
|
|
C[batch_idx, chunk_idx * chunk_size + m_idx * block_M:chunk_idx * chunk_size +
|
|
(m_idx + 1) * block_M, bz // (nheads // ngroups), 0:block_Dstate], C_shared)
|
|
T.copy(
|
|
prev_states[batch_idx, chunk_idx, bz, n_idx * block_N:(n_idx + 1) * block_N,
|
|
0:block_Dstate], prev_state_shared)
|
|
T.gemm(C_shared, prev_state_shared, acc_o, transpose_B=True)
|
|
for i, j in T.Parallel(block_M, block_N):
|
|
acc_o[i, j] *= scale_m_local[i]
|
|
|
|
loop_range = T.ceildiv((m_idx + 1) * block_M, block_K)
|
|
|
|
for k in T.Pipelined(loop_range, num_stages=num_stages):
|
|
T.copy(
|
|
cb[batch_idx, chunk_idx, bz // (nheads // ngroups),
|
|
m_idx * block_M:(m_idx + 1) * block_M, k * block_K:(k + 1) * block_K],
|
|
cb_shared)
|
|
T.copy(cb_shared, cb_local)
|
|
T.copy(dA_cumsum[batch_idx, bz, chunk_idx, k * block_K:(k + 1) * block_K],
|
|
dA_cs_k_shared)
|
|
T.copy(dA_cs_k_shared, dA_cs_k_local)
|
|
for i, j in T.Parallel(block_M, block_K):
|
|
cb_local[i,
|
|
j] = cb_local[i,
|
|
j] * T.exp2(dA_cs_m_local[i] * p - dA_cs_k_local[j] * p)
|
|
T.copy(dt[batch_idx, bz, chunk_idx, k * block_K:(k + 1) * block_K], dt_shared)
|
|
T.copy(dt_shared, dt_local)
|
|
for i, j in T.Parallel(block_M, block_K):
|
|
cb_local[i, j] *= dt_local[j]
|
|
for i, j in T.Parallel(block_M, block_K):
|
|
cb_local[i, j] = T.if_then_else(m_idx * block_M + i >= k * block_K + j,
|
|
cb_local[i, j], 0)
|
|
T.copy(
|
|
x[batch_idx, chunk_idx * chunk_size + k * block_K:chunk_idx * chunk_size +
|
|
(k + 1) * block_K, bz, n_idx * block_N:(n_idx + 1) * block_N], x_shared)
|
|
T.gemm(cb_local, x_shared, acc_o)
|
|
|
|
D_local[0] = D[bz]
|
|
T.copy(
|
|
x[batch_idx, chunk_idx * chunk_size + m_idx * block_M:chunk_idx * chunk_size +
|
|
(m_idx + 1) * block_M, bz, n_idx * block_N:(n_idx + 1) * block_N],
|
|
x_residual_shared)
|
|
T.copy(x_residual_shared, x_residual_local)
|
|
for i, j in T.Parallel(block_M, block_N):
|
|
acc_o[i, j] += x_residual_local[i, j] * D_local[0]
|
|
|
|
T.copy(
|
|
acc_o,
|
|
Output[batch_idx, chunk_idx * chunk_size + m_idx * block_M:chunk_idx * chunk_size +
|
|
(m_idx + 1) * block_M, bz, n_idx * block_N:(n_idx + 1) * block_N])
|
|
|
|
return main
|
|
|
|
|
|
def run_chunk_scan(batch,
|
|
seqlen,
|
|
chunk_size,
|
|
ngroups,
|
|
nheads,
|
|
headdim,
|
|
dstate,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
block_Dstate,
|
|
num_stages=2,
|
|
threads=128):
|
|
program = chunk_scan_fwd(batch, seqlen, chunk_size, ngroups, nheads, headdim, dstate, block_M,
|
|
block_N, block_K, block_Dstate, num_stages, threads)
|
|
|
|
kernel = tilelang.compile(program, out_idx=[7])
|
|
profiler = kernel.get_profiler()
|
|
|
|
def ref_program(cb, x, dt, dA_cumsum, C, prev_states, D):
|
|
import torch
|
|
from einops import rearrange, repeat
|
|
"""
|
|
Argument:
|
|
cb: (batch, nchunks, ngroups, chunk_size, chunk_size)
|
|
x: (batch, seqlen, nheads, headdim)
|
|
dt: (batch, nheads, nchunks, chunk_size)
|
|
dA_cumsum: (batch, nheads, nchunks, chunk_size)
|
|
C: (batch, seqlen, ngroups, dstate)
|
|
prev_states: (batch, nchunks, nheads, headdim, dstate)
|
|
D: (nheads, headdim) or (nheads,)
|
|
z: (batch, seqlen, nheads, headdim)
|
|
Return:
|
|
out: (batch, seqlen, nheads, headdim)
|
|
"""
|
|
_, _, ngroups, _, _ = cb.shape
|
|
batch, seqlen, nheads, headdim = x.shape
|
|
# _, _, ngroups, dstate = B.shape
|
|
# assert B.shape == (batch, seqlen, ngroups, dstate)
|
|
_, _, nchunks, chunk_size = dt.shape
|
|
assert seqlen == nchunks * chunk_size
|
|
# assert C.shape == B.shape
|
|
# B = repeat(B, "b l g d -> b l (g h) d", h=nheads // ngroups)
|
|
C = repeat(C, "b l g d -> b l (g h) d", h=nheads // ngroups)
|
|
cb = repeat(cb, "b c g l s -> b c (g h) l s", h=nheads // ngroups)
|
|
# CB = torch.einsum("bclhn,bcshn->bchls", rearrange(C, "b (c l) h n -> b c l h n", c=nchunks),
|
|
# rearrange(B, "b (c s) h n -> b c s h n", c=nchunks))
|
|
# (batch, nheads, nchunks, chunksize, chunksize)
|
|
dt_segment_sum = dA_cumsum[:, :, :, :, None] - dA_cumsum[:, :, :, None, :]
|
|
decay = torch.exp(dt_segment_sum)
|
|
scores_decay = cb * rearrange(decay, "b h c l s -> b c h l s")
|
|
causal_mask = torch.tril(
|
|
torch.ones(chunk_size, chunk_size, device=x.device, dtype=bool), diagonal=0)
|
|
scores_decay = scores_decay.masked_fill(~causal_mask, 0)
|
|
out = torch.einsum('bchls,bhcs,bcshp->bclhp', scores_decay.to(x.dtype), dt.to(x.dtype),
|
|
rearrange(x, "b (c s) h p -> b c s h p", c=nchunks))
|
|
state_decay_out = torch.exp(rearrange(dA_cumsum, "b h c l -> b c l h 1"))
|
|
out_prev = torch.einsum('bclhn,bchpn->bclhp',
|
|
rearrange(C, "b (c l) h n -> b c l h n", c=nchunks),
|
|
prev_states.to(C.dtype)) * state_decay_out
|
|
out = out + out_prev
|
|
out = rearrange(out, "b c l h p -> b (c l) h p")
|
|
if D is not None:
|
|
if D.dim() == 1:
|
|
D = rearrange(D, "h -> h 1")
|
|
out = out + x * D
|
|
return out
|
|
|
|
profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2, max_mismatched_ratio=0.05)
|
|
|
|
|
|
def chunk_state_fwd(batch,
|
|
seqlen,
|
|
chunk_size,
|
|
ngroups,
|
|
nheads,
|
|
headdim,
|
|
dstate,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
num_stages=2,
|
|
threads=128):
|
|
dtype = "float16"
|
|
accum_dtype = "float"
|
|
nchunks = T.ceildiv(seqlen, chunk_size)
|
|
p = 1.44269504
|
|
|
|
@T.prim_func
|
|
def main(B: T.Tensor((batch, seqlen, ngroups, dstate), dtype), x: T.Tensor(
|
|
(batch, seqlen, nheads, headdim), dtype), dt: T.Tensor(
|
|
(batch, nheads, nchunks, chunk_size), dtype), dA_cumsum: T.Tensor(
|
|
(batch, nheads, nchunks, chunk_size), dtype), Output: T.Tensor(
|
|
(batch, nchunks, nheads, headdim, dstate), dtype)):
|
|
with T.Kernel(
|
|
nheads,
|
|
T.ceildiv(headdim, block_M) * T.ceildiv(dstate, block_N),
|
|
batch * nchunks,
|
|
threads=threads) as (bz, bx, by):
|
|
x_shared = T.alloc_shared((block_K, block_M), dtype)
|
|
x_local = T.alloc_fragment((block_K, block_M), dtype)
|
|
xt_local = T.alloc_fragment((block_M, block_K), dtype)
|
|
B_shared = T.alloc_shared((block_K, block_N), dtype)
|
|
dt_shared = T.alloc_shared((block_K), dtype)
|
|
dA_cumsum_shared = T.alloc_shared((block_K), dtype)
|
|
acc_o = T.alloc_fragment((block_M, block_N), accum_dtype)
|
|
scale = T.alloc_fragment((block_K), accum_dtype)
|
|
dA_cs_last = T.alloc_fragment((1), accum_dtype)
|
|
dA_cumsum_local = T.alloc_fragment((block_K), accum_dtype)
|
|
dt_local = T.alloc_fragment((block_K), accum_dtype)
|
|
|
|
loop_range = T.ceildiv(chunk_size, block_K)
|
|
|
|
batch_idx = by % batch
|
|
chunk_idx = by // batch
|
|
m_idx = bx // T.ceildiv(dstate, block_N)
|
|
n_idx = bx % T.ceildiv(dstate, block_N)
|
|
|
|
dA_cs_last[0] = dA_cumsum[batch_idx, bz, chunk_idx, chunk_size - 1]
|
|
T.clear(acc_o)
|
|
for k in T.Pipelined(loop_range, num_stages=num_stages):
|
|
T.copy(
|
|
x[batch_idx, chunk_idx * chunk_size + k * block_K:chunk_idx * chunk_size +
|
|
(k + 1) * block_K, bz, m_idx * block_M:(m_idx + 1) * block_M], x_shared)
|
|
T.copy(dA_cumsum[batch_idx, bz, chunk_idx, k * block_K:(k + 1) * block_K],
|
|
dA_cumsum_shared)
|
|
T.copy(dt[batch_idx, bz, chunk_idx, k * block_K:(k + 1) * block_K], dt_shared)
|
|
T.copy(dA_cumsum_shared, dA_cumsum_local)
|
|
T.copy(dt_shared, dt_local)
|
|
for i in T.Parallel(block_K):
|
|
scale[i] = T.exp2(dA_cs_last[0] * p - dA_cumsum_local[i] * p) * dt_local[i]
|
|
T.copy(x_shared, x_local)
|
|
for i, j in T.Parallel(block_M, block_K):
|
|
xt_local[i, j] = x_local[j, i] * scale[j]
|
|
T.copy(
|
|
B[batch_idx, chunk_idx * chunk_size + k * block_K:chunk_idx * chunk_size +
|
|
(k + 1) * block_K, bz // (nheads // ngroups),
|
|
n_idx * block_N:(n_idx + 1) * block_N], B_shared)
|
|
T.gemm(xt_local, B_shared, acc_o)
|
|
T.copy(
|
|
acc_o, Output[batch_idx, chunk_idx, bz, m_idx * block_M:(m_idx + 1) * block_M,
|
|
n_idx * block_N:(n_idx + 1) * block_N])
|
|
|
|
return main
|
|
|
|
|
|
def run_chunk_state(batch,
|
|
seqlen,
|
|
chunk_size,
|
|
ngroups,
|
|
nheads,
|
|
headdim,
|
|
dstate,
|
|
block_M,
|
|
block_N,
|
|
block_K,
|
|
num_stages=2,
|
|
threads=128):
|
|
program = chunk_state_fwd(batch, seqlen, chunk_size, ngroups, nheads, headdim, dstate, block_M,
|
|
block_N, block_K, num_stages, threads)
|
|
|
|
kernel = tilelang.compile(program, out_idx=[4])
|
|
profiler = kernel.get_profiler()
|
|
|
|
def ref_program(B, x, dt, dA_cumsum):
|
|
"""
|
|
Argument:
|
|
B: (batch, seqlen, ngroups, headdim)
|
|
x: (batch, seqlen, nheads, headdim)
|
|
dt: (batch, nheads, nchunks, chunk_size)
|
|
dA_cumsum: (batch, nheads, nchunks, chunk_size)
|
|
Return:
|
|
states: (batch, nchunks, nheads, headdim, dstate)
|
|
"""
|
|
# Check constraints.
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from einops import rearrange, repeat
|
|
|
|
batch, seqlen, nheads, headdim = x.shape
|
|
dstate = B.shape[-1]
|
|
_, _, nchunks, chunk_size = dt.shape
|
|
assert seqlen <= nchunks * chunk_size
|
|
assert x.shape == (batch, seqlen, nheads, headdim)
|
|
assert dt.shape == (batch, nheads, nchunks, chunk_size)
|
|
ngroups = B.shape[2]
|
|
assert nheads % ngroups == 0
|
|
assert B.shape == (batch, seqlen, ngroups, dstate)
|
|
B = repeat(B, "b l g d -> b l (g h) d", h=nheads // ngroups)
|
|
assert dA_cumsum.shape == (batch, nheads, nchunks, chunk_size)
|
|
if seqlen < nchunks * chunk_size:
|
|
x = F.pad(x, (0, 0, 0, 0, 0, nchunks * chunk_size - seqlen))
|
|
B = F.pad(B, (0, 0, 0, 0, 0, nchunks * chunk_size - seqlen))
|
|
x = rearrange(x, "b (c l) h p -> b c l h p", l=chunk_size)
|
|
B = rearrange(B, "b (c l) ... -> b c l ...", l=chunk_size)
|
|
decay_states = torch.exp((dA_cumsum[:, :, :, -1:] - dA_cumsum))
|
|
return torch.einsum("bclhn,bhcl,bhcl,bclhp->bchpn", B.to(x.dtype), decay_states.to(x.dtype),
|
|
dt.to(x.dtype), x)
|
|
|
|
profiler.assert_allclose(ref_program, atol=1e-2, rtol=1e-2, max_mismatched_ratio=0.05)
|
|
|
|
|
|
def test_chunk_scan():
|
|
run_chunk_scan(
|
|
batch=8,
|
|
seqlen=2048,
|
|
chunk_size=256,
|
|
ngroups=1,
|
|
nheads=8,
|
|
headdim=64,
|
|
dstate=128,
|
|
block_M=64,
|
|
block_N=64,
|
|
block_K=64,
|
|
block_Dstate=128,
|
|
num_stages=2,
|
|
threads=128)
|
|
|
|
|
|
def test_chunk_state():
|
|
run_chunk_state(
|
|
batch=8,
|
|
seqlen=2048,
|
|
chunk_size=256,
|
|
ngroups=1,
|
|
nheads=8,
|
|
headdim=64,
|
|
dstate=128,
|
|
block_M=64,
|
|
block_N=64,
|
|
block_K=64,
|
|
num_stages=2,
|
|
threads=128)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|