Skip to content

GitLab

Unverified Commit 667632cc authored by guchaoyang's avatar guchaoyang Committed by GitHub
Browse files

Merge branch 'main' into dcu

parents d6dd2ddf a874e4e8
Hide whitespace changes
Inline Side-by-side
Showing with 2213 additions and 647 deletions
examples/grouped_gemm/example_grouped_gemm_fwd.py
View file @ 667632cc
......@@ -18,8 +18,7 @@ def torch_gmm(a, b, batch_sizes, batch_offsets_tensor, trans_b=False):
torch.Tensor: Resulting tensor after grouped matrix multiplication.
"""
assert a.shape[0] == sum(batch_sizes), "Sum of batch_sizes must equal the first dimension of a"
assert b.shape[0] == len(
batch_sizes), "The first dimension of b must match the length of batch_sizes"
assert b.shape[0] == len(batch_sizes), "The first dimension of b must match the length of batch_sizes"
# Initialize output tensor
output = torch.empty((sum(batch_sizes), b.shape[2]), device=a.device, dtype=a.dtype)
......@@ -38,15 +37,7 @@ def torch_gmm(a, b, batch_sizes, batch_offsets_tensor, trans_b=False):
@tilelang.jit(out_idx=[2])
def grouped_gemm(batch_sizes_list,
K,
N,
block_M,
block_N,
block_K,
num_stages=2,
threads=128,
dtype="float16"):
def grouped_gemm(batch_sizes_list, K, N, block_M, block_N, block_K, num_stages=2, threads=128, dtype=T.float16):
"""
args:
a (torch.Tensor): Input tensor of shape (M, K).
......@@ -54,46 +45,39 @@ def grouped_gemm(batch_sizes_list,
"""
batch_sum = sum(batch_sizes_list)
batch_count = len(batch_sizes_list)
accum_dtype = "float32"
accum_dtype = T.float32
total_m_blocks = sum((size + block_M - 1) // block_M for size in batch_sizes_list)
@T.prim_func
def kernel(
A: T.Tensor([batch_sum, K], dtype), # type: ignore
B: T.Tensor([batch_count, K, N], dtype), # type: ignore
C: T.Tensor([batch_sum, N], dtype), # type: ignore
batch_sizes: T.Tensor([batch_count], "int32"), # type: ignore
batch_offsets: T.Tensor([batch_count], "int32"), # type: ignore
batch_padded_offsets: T.Tensor([batch_count], "int32"), # type: ignore
A: T.Tensor([batch_sum, K], dtype), # type: ignore
B: T.Tensor([batch_count, K, N], dtype), # type: ignore
C: T.Tensor([batch_sum, N], dtype), # type: ignore
batch_sizes: T.Tensor([batch_count], T.int32), # type: ignore
batch_offsets: T.Tensor([batch_count], T.int32), # type: ignore
batch_padded_offsets: T.Tensor([batch_count], T.int32), # type: ignore
):
with T.Kernel(total_m_blocks, T.ceildiv(N, block_N), threads=threads) 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)
cur_batch_idx = T.alloc_local([1], "int32")
cur_batch_size = T.alloc_local([1], "int32")
cur_batch_idx = T.alloc_local([1], T.int32)
cur_batch_size = T.alloc_local([1], T.int32)
m_start_padded = bx * block_M
for i in range(batch_count):
in_cur_batch_idx = (m_start_padded >= batch_padded_offsets[i])
in_cur_batch_idx = m_start_padded >= batch_padded_offsets[i]
cur_batch_idx[0] = T.if_then_else(in_cur_batch_idx, i, cur_batch_idx[0])
cur_batch_size[0] = batch_sizes[cur_batch_idx[0]]
m_start = m_start_padded - batch_padded_offsets[cur_batch_idx[0]] + batch_offsets[
cur_batch_idx[0]]
actual_rows = T.max(
0,
T.min(block_M,
cur_batch_size[0] + batch_padded_offsets[cur_batch_idx[0]] - m_start_padded))
m_start = m_start_padded - batch_padded_offsets[cur_batch_idx[0]] + batch_offsets[cur_batch_idx[0]]
actual_rows = T.max(0, T.min(block_M, cur_batch_size[0] + batch_padded_offsets[cur_batch_idx[0]] - m_start_padded))
T.clear(C_local)
for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=num_stages):
T.copy(A[m_start:m_start + block_M, k * block_K:(k + 1) * block_K], A_shared)
T.copy(
B[cur_batch_idx[0], k * block_K:(k + 1) * block_K,
by * block_N:(by + 1) * block_N], B_shared)
T.copy(A[m_start : m_start + block_M, k * block_K : (k + 1) * block_K], A_shared)
T.copy(B[cur_batch_idx[0], k * block_K : (k + 1) * block_K, by * block_N : (by + 1) * block_N], B_shared)
T.gemm(A_shared, B_shared, C_local)
for i, j in T.Parallel(block_M, block_N):
......@@ -111,8 +95,7 @@ def construct_inputs(batch_sizes_list, K, M, trans_b, padding_M, device, dtype):
for i in range(batch_count - 1):
batch_offsets_list.append(batch_offsets_list[-1] + batch_sizes_list[i])
for i in range(batch_count - 1):
batch_padded_offsets_list.append(batch_padded_offsets_list[-1] +
math.ceil((batch_sizes_list[i]) / padding_M) * padding_M)
batch_padded_offsets_list.append(batch_padded_offsets_list[-1] + math.ceil((batch_sizes_list[i]) / padding_M) * padding_M)
A = torch.randn(batch_sum, K, device=device, dtype=dtype)
B = torch.randn(batch_count, K, M, device=device, dtype=dtype)
C = torch.empty(batch_sum, M, device=device, dtype=dtype)
......@@ -125,27 +108,16 @@ def construct_inputs(batch_sizes_list, K, M, trans_b, padding_M, device, dtype):
return A, B, C, batch_sizes, batch_offsets, batch_padded_offsets
def run_tilelang_grouped_gemm(batch_sizes_list,
K,
M,
block_M,
block_N,
block_K,
trans_b,
num_stages=2,
threads=128,
profile=False):
def run_tilelang_grouped_gemm(batch_sizes_list, K, M, block_M, block_N, block_K, trans_b, num_stages=2, threads=128, profile=False):
padding_M = block_M
batch_sum = sum(batch_sizes_list)
kernel = grouped_gemm(
tuple(batch_sizes_list), K, M, block_M, block_N, block_K, num_stages, threads)
kernel = grouped_gemm(tuple(batch_sizes_list), K, M, block_M, block_N, block_K, num_stages, threads)
# print(kernel.get_kernel_source())
device = torch.device("cuda")
dtype = torch.float16
A, B, C, batch_sizes, batch_offsets, batch_padded_offsets = construct_inputs(
batch_sizes_list, K, M, trans_b, padding_M, device, dtype)
A, B, C, batch_sizes, batch_offsets, batch_padded_offsets = construct_inputs(batch_sizes_list, K, M, trans_b, padding_M, device, dtype)
out = kernel(A, B, batch_sizes, batch_offsets, batch_padded_offsets)
ref_output = torch_gmm(A, B, batch_sizes, batch_offsets, trans_b)
# print(out)
......@@ -157,8 +129,7 @@ def run_tilelang_grouped_gemm(batch_sizes_list,
if profile:
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Auto)
latency = profiler.do_bench(
warmup=500, input_tensors=[A, B, batch_sizes, batch_offsets, batch_padded_offsets])
latency = profiler.do_bench(warmup=500, input_tensors=[A, B, batch_sizes, batch_offsets, batch_padded_offsets])
print(f"Latency: {latency} ms")
print(f"TFlops: {batch_sum * K * M * 2 / latency * 1e-9} TFlops")
......@@ -173,12 +144,11 @@ def test_grouped_gemm():
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
'--batch_sizes', type=str, default="64, 128", help='comma-separated batch sizes')
parser.add_argument('--K', type=int, default=8192, help='reduce dim')
parser.add_argument('--M', type=int, default=8192, help='output dim')
parser.add_argument('--trans_b', action="store_true", help="transpose B")
parser.add_argument('--profile', action="store_true", help="profile")
parser.add_argument("--batch_sizes", type=str, default="64, 128", help="comma-separated batch sizes")
parser.add_argument("--K", type=int, default=8192, help="reduce dim")
parser.add_argument("--M", type=int, default=8192, help="output dim")
parser.add_argument("--trans_b", action="store_true", help="transpose B")
parser.add_argument("--profile", action="store_true", help="profile")
args = parser.parse_args()
batch_sizes_list = [int(x) for x in args.batch_sizes.split(",")]
......@@ -190,14 +160,4 @@ if __name__ == "__main__":
num_stages = 2
threads = 256
run_tilelang_grouped_gemm(
batch_sizes_list,
K,
M,
block_M,
block_N,
block_K,
trans_b,
num_stages,
threads,
profile=args.profile)
run_tilelang_grouped_gemm(batch_sizes_list, K, M, block_M, block_N, block_K, trans_b, num_stages, threads, profile=args.profile)
examples/hadamard_transform/example_hadamard.py
View file @ 667632cc
......@@ -17,7 +17,7 @@ def is_pow_of_2(n):
def hadamard(b, n, dtype):
assert is_pow_of_2(n), "n must be a power of 2"
assert 2 <= n <= 32768, "n must be in [2, 32768]"
elem_size = {'float32': 4, 'float16': 2, 'bfloat16': 2}[dtype]
elem_size = {T.float32: 4, T.float16: 2, T.bfloat16: 2}[dtype]
logN = int(math.log2(n))
threads = [0, 1, 1, 1, 2, 4, 8, 16, 32, 32, 128, 256, 256, 256, 256, 256][logN]
......@@ -40,23 +40,21 @@ def hadamard(b, n, dtype):
# print(f'{exchange_round=}')
@T.macro
def warp_shfl(local: T.Tensor((thread_elem,), dtype), buf: T.Tensor((thread_elem,), dtype),
round: int):
def warp_shfl(local: T.Tensor((thread_elem,), dtype), buf: T.Tensor((thread_elem,), dtype), round: int):
tx = T.get_thread_binding(0)
for i in T.serial(round):
tx_stride = 1 << i
another_tx = tx ^ tx_stride
sign = (
tx >> i
) & 1 # get i-th lowest bit of tx, which determines the operation type for shared[tx, :]
sign = (tx >> i) & 1 # get i-th lowest bit of tx, which determines the operation type for shared[tx, :]
for j in T.Pipelined(thread_elem, num_stages=1):
buf[j] = T.tvm_warp_shuffle(
0xffffffff, # mask of all threads
0xFFFFFFFF, # mask of all threads
local[j],
another_tx % warp_size,
warp_size,
warp_size)
warp_size,
)
local[j] = T.if_then_else(sign == 0, local[j] + buf[j], buf[j] - local[j])
@T.prim_func
......@@ -78,10 +76,8 @@ def hadamard(b, n, dtype):
for j in T.serial(chunknum):
chunkbase = j * chunksize
for k in T.serial(chunksize // 2):
local[chunkbase +
k] = local[chunkbase + k] + local[chunkbase + k + chunksize // 2]
local[chunkbase + k + chunksize //
2] = local[chunkbase + k] - 2 * local[chunkbase + k + chunksize // 2]
local[chunkbase + k] = local[chunkbase + k] + local[chunkbase + k + chunksize // 2]
local[chunkbase + k + chunksize // 2] = local[chunkbase + k] - 2 * local[chunkbase + k + chunksize // 2]
# 3. Hadamard inside warp, n<=512
# In warp level, we rely on warp shuffle to exchange data inside each warp, without using shared memory
......@@ -131,28 +127,27 @@ def ref_program(x: torch.Tensor):
assert x.ndim == 2
dim = x.shape[-1]
assert is_pow_of_2(dim)
return F.linear(
x, torch.tensor(scipy.linalg.hadamard(dim, dtype=float), dtype=x.dtype, device=x.device))
return F.linear(x, torch.tensor(scipy.linalg.hadamard(dim, dtype=float), dtype=x.dtype, device=x.device))
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--batch', type=int, default=64, help='Batch size')
parser.add_argument('--dim', type=int, default=32768, help='Dimension')
parser.add_argument("--batch", type=int, default=64, help="Batch size")
parser.add_argument("--dim", type=int, default=32768, help="Dimension")
args = parser.parse_args()
B, D = args.batch, args.dim
x = torch.randn((B, D), device='cuda')
kernel = hadamard(B, D, 'float32')
x = torch.randn((B, D), device="cuda")
kernel = hadamard(B, D, T.float32)
y = kernel(x)
y_ref = ref_program(x)
torch.testing.assert_close(y, y_ref, atol=1e-2, rtol=1e-2)
print('All tests passed.')
print("All tests passed.")
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Auto)
latency = profiler.do_bench(warmup=100)
print("Tile-lang: {:.2f} ms".format(latency))
if __name__ == '__main__':
if __name__ == "__main__":
main()
examples/lazy_jit/lazyjit.en.ipynb 0 → 100644
View file @ 667632cc
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "5e0deecc",
"metadata": {},
"outputs": [],
"source": [
"import sys\n",
"from pathlib import Path\n",
"\n",
"sys.path.insert(0, str(Path.cwd().parent.parent.absolute()))\n",
"import tilelang\n",
"import torch\n",
"import tilelang.language as T"
]
},
{
"cell_type": "markdown",
"id": "1ca2c56d",
"metadata": {},
"source": [
"# Tilelang Lazy JIT"
]
},
{
"cell_type": "markdown",
"id": "156e7370",
"metadata": {},
"source": [
"## Tensor Annotation"
]
},
{
"cell_type": "markdown",
"id": "b070c109",
"metadata": {},
"source": [
"Tilelang Lazy JIT combines the jit generation and invocation logic.\n",
"\n",
"The function signature syntax is similar to triton but with significant enhancements, most notably allowing Tensor annotations:\n",
"\n",
"For example, the code below annotates a 2D Tensor with T.Tensor[[int, int], T.float16]\n",
"1. Each dimension is a compile-time constant; changing it triggers recompilation\n",
"2. Its dtype must be T.float16\n",
"\n",
"DType can also be Any or None in addition to a concrete type\n"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "60bf8954",
"metadata": {},
"outputs": [],
"source": [
"@tilelang.lazy_jit\n",
"def gemm(\n",
" A: T.Tensor[[int, int], T.float16],\n",
" B: T.Tensor[[int, int], T.float16],\n",
" out_dtype: T.dtype = T.float32,\n",
" block_M: int = 128,\n",
" block_N: int = 128,\n",
" block_K: int = 32,\n",
"):\n",
" M, K = A.shape\n",
" K, N = B.shape\n",
" C = T.empty((M, N), out_dtype)\n",
" with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=128) as (bx, by):\n",
" A_shared = T.alloc_shared((block_M, block_K), A.dtype)\n",
" B_shared = T.alloc_shared((block_K, block_N), B.dtype)\n",
" C_local = T.alloc_fragment((block_M, block_N), out_dtype)\n",
" T.clear(C_local)\n",
" for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):\n",
" T.copy(A[bx * block_M, k * block_K], A_shared)\n",
" T.copy(B[k * block_K, by * block_N], B_shared)\n",
" T.gemm(A_shared, B_shared, C_local)\n",
" T.copy(C_local, C[bx * block_M, by * block_N])\n",
" return C"
]
},
{
"cell_type": "markdown",
"id": "28f868fe",
"metadata": {},
"source": [
"Call the Tensor directly as an argument to trigger the full jit compile-and-run workflow:"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "ee13394a",
"metadata": {},
"outputs": [],
"source": [
"A = torch.randn(1024, 512, dtype=torch.float16, device=\"cuda\")\n",
"B = torch.randn(512, 256, dtype=torch.float16, device=\"cuda\")\n",
"C = gemm(A, B)\n",
"\n",
"# check output is correct\n",
"C_ref = (A @ B).float()\n",
"torch.testing.assert_close(C, C_ref, rtol=1e-2, atol=1e-2)"
]
},
{
"cell_type": "markdown",
"id": "c6705091",
"metadata": {},
"source": [
"Change the call-site arguments; if the compiler parameters differ, it recompiles:"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "d8aab5b7",
"metadata": {},
"outputs": [],
"source": [
"A = torch.randn(1024, 512, dtype=torch.float16, device=\"cuda\")\n",
"B = torch.randn(512, 1024, dtype=torch.float16, device=\"cuda\")\n",
"C = gemm(A, B, block_M=64, block_N=64)"
]
},
{
"cell_type": "markdown",
"id": "ce6b7391",
"metadata": {},
"source": [
"You can also manually call compile helpers to build a kernel\n",
"\n",
"1. `ker.compile` compiles the kernel\n",
"2. `ker.get_tir` retrieves the tir\n",
"3. `ker.par_compile` compiles in parallel"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "f3cf3a2d",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"2025-11-25 17:29:46 [TileLang:tilelang.cache.kernel_cache:WARNING]: Found kernel in memory cache. For better performance, consider using `@tilelang.jit` instead of direct kernel caching.\n"
]
}
],
"source": [
"kernel = gemm.compile(A, B, block_M=64, block_N=64)\n",
"C = kernel(A, B)"
]
},
{
"cell_type": "markdown",
"id": "921761b5",
"metadata": {},
"source": [
"## More Tensor Annotation"
]
},
{
"cell_type": "markdown",
"id": "4539e54e",
"metadata": {},
"source": [
"### Separate the implementation with macros"
]
},
{
"cell_type": "markdown",
"id": "ad96ba65",
"metadata": {},
"source": [
"Next we'll implement a simple gemm in several ways. For convenience, first write a macro that captures the main gemm logic:"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "171d4fe6",
"metadata": {},
"outputs": [],
"source": [
"@T.macro\n",
"def gemm_impl(A, B, C, M, N, K, block_M, block_N, block_K):\n",
" with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=128) as (bx, by):\n",
" A_shared = T.alloc_shared((block_M, block_K), A.dtype)\n",
" B_shared = T.alloc_shared((block_K, block_N), B.dtype)\n",
" C_local = T.alloc_fragment((block_M, block_N), C.dtype)\n",
" T.clear(C_local)\n",
" for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):\n",
" T.copy(A[bx * block_M, k * block_K], A_shared)\n",
" T.copy(B[k * block_K, by * block_N], B_shared)\n",
" T.gemm(A_shared, B_shared, C_local)\n",
" T.copy(C_local, C[bx * block_M, by * block_N])"
]
},
{
"cell_type": "markdown",
"id": "446a1acd",
"metadata": {},
"source": [
"### Mark dynamic shapes with T.dyn\n",
"\n",
"When some dimensions are dynamic, mark them with T.dyn. T.dyn can take a string argument to name the variable"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6a38aa95",
"metadata": {},
"outputs": [],
"source": [
"@tilelang.lazy_jit\n",
"def gemm_dyn_K(\n",
" A: T.Tensor[[int, T.dyn[\"K\"]], T.float16], # noqa: F821\n",
" B: T.Tensor[[T.dyn[\"K\"], int], T.float16], # noqa: F821\n",
"):\n",
" M, K = A.shape\n",
" K, N = B.shape\n",
" C = T.empty((M, N), T.float32)\n",
" gemm_impl(A, B, C, M, N, K, 128, 128, 32)\n",
" return C"
]
},
{
"cell_type": "markdown",
"id": "c60fd346",
"metadata": {},
"source": [
"Inspect the lazy_jit function signature: parameters with a `$` suffix are compile-time constants that may vary, and those with `$dyn` are runtime variables"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "c6992eb4",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'A': TensorAnnot(shape=[A_shape_0$, K$dyn], strides=None, dtype=dtype('float16')),\n",
" 'B': TensorAnnot(shape=[K$dyn, B_shape_1$], strides=None, dtype=dtype('float16'))}"
]
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"gemm_dyn_K.func.annot"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "fe6cfdc8",
"metadata": {},
"outputs": [],
"source": [
"A = torch.randn(1024, 512, dtype=torch.float16, device=\"cuda\")\n",
"B = torch.randn(512, 256, dtype=torch.float16, device=\"cuda\")\n",
"C = gemm_dyn_K(A, B)\n",
"C_ref = (A @ B).float()\n",
"torch.testing.assert_close(C, C_ref, rtol=1e-2, atol=1e-2)"
]
},
{
"cell_type": "markdown",
"id": "2ee97bf7",
"metadata": {},
"source": [
"### Use T.StridedTensor to annotate tensors with strides\n",
"\n",
"Annotation format: T.StridedTensor[Shape, Stride, DType]. Each Shape or Stride entry can be\n",
"* int: compile-time constant\n",
"* T.dyn: runtime value\n",
"\n",
"DType can be None or Any"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "9dde1dae",
"metadata": {},
"outputs": [],
"source": [
"from typing import Any\n",
"\n",
"\n",
"@tilelang.lazy_jit\n",
"def as_contingious(A: T.StridedTensor[[T.dyn, T.dyn], [T.dyn, T.dyn], Any]):\n",
" M, N = A.shape\n",
" B = T.empty((M, N), A.dtype)\n",
" block_M = 128\n",
" block_N = 128\n",
" with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=128) as (bx, by):\n",
" T.copy(\n",
" A[bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N],\n",
" B[bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N],\n",
" )\n",
" return B"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "dec2c0a7",
"metadata": {},
"outputs": [],
"source": [
"A = torch.randn(1024, 1024, device=\"cuda\")\n",
"B = as_contingious(A[::2, ::2])\n",
"B_ref = A[::2, ::2].contiguous()\n",
"torch.testing.assert_close(B, B_ref)"
]
},
{
"cell_type": "markdown",
"id": "f5fb20d6",
"metadata": {},
"source": [
"## More Annotation"
]
},
{
"cell_type": "markdown",
"id": "890df0a2",
"metadata": {},
"source": [
"### Annotate tensors with T.ptr\n",
"lazy_jit lets you declare a handle with T.ptr, but you must define its shape inside the function via T.match_buffer"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "0fc17af6",
"metadata": {},
"outputs": [],
"source": [
"@tilelang.lazy_jit\n",
"def gemm_ptr(\n",
" A: T.ptr,\n",
" B: T.ptr,\n",
" M: int,\n",
" N: int,\n",
" K: int,\n",
"):\n",
" A = T.match_buffer(A, (M, K), T.float16)\n",
" B = T.match_buffer(B, (K, N), T.float16)\n",
" C = T.empty((M, N), T.float32)\n",
" gemm_impl(A, B, C, M, N, K, block_M=128, block_N=128, block_K=32)\n",
" return C"
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "8e52a554",
"metadata": {},
"outputs": [],
"source": [
"A = torch.randn(1024, 512, dtype=torch.float16, device=\"cuda\")\n",
"B = torch.randn(512, 256, dtype=torch.float16, device=\"cuda\")\n",
"C = gemm_ptr(A, B, 1024, 256, 512)\n",
"C_ref = (A @ B).float()\n",
"torch.testing.assert_close(C, C_ref, rtol=1e-2, atol=1e-2)"
]
},
{
"cell_type": "markdown",
"id": "6b19ef90",
"metadata": {},
"source": [
"### Use T.int32 to annotate runtime variables\n",
"\n",
"lazy_jit lets you define runtime variables with T.int32 or other types, enabling a fully dynamic gemm similar to triton"
]
},
{
"cell_type": "code",
"execution_count": 14,
"id": "c1e7598a",
"metadata": {},
"outputs": [],
"source": [
"@tilelang.lazy_jit\n",
"def gemm_ptr_dyn(\n",
" A: T.ptr,\n",
" B: T.ptr,\n",
" M: T.int32,\n",
" N: T.int32,\n",
" K: T.int32,\n",
"):\n",
" A = T.match_buffer(A, (M, K), T.float16, strides=(K, 1))\n",
" B = T.match_buffer(B, (K, N), T.float16, strides=(N, 1))\n",
" C = T.empty((M, N), T.float32)\n",
" gemm_impl(A, B, C, M, N, K, block_M=128, block_N=128, block_K=32)\n",
" return C"
]
},
{
"cell_type": "code",
"execution_count": 15,
"id": "9e9a4c88",
"metadata": {},
"outputs": [],
"source": [
"A = torch.randn(1024, 512, dtype=torch.float16, device=\"cuda\")\n",
"B = torch.randn(512, 256, dtype=torch.float16, device=\"cuda\")\n",
"C = gemm_ptr_dyn(A, B, 1024, 256, 512)\n",
"C_ref = (A @ B).float()\n",
"torch.testing.assert_close(C, C_ref, rtol=1e-2, atol=1e-2)"
]
},
{
"cell_type": "markdown",
"id": "39166cb4",
"metadata": {},
"source": [
"## Compilation and parallel compilation"
]
},
{
"cell_type": "markdown",
"id": "8c6fbe08",
"metadata": {},
"source": [
"lazyjit and the original jit both support parallel compilation\n",
"\n",
"To avoid wasting memory with torch.tensor placeholders, use T.Tensor to create placeholders"
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "7222e57b",
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "c6d7f05cdfff412e9a527332438f7aa2",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Elaborating: 0%| | 0/8 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "14836065a21b41ae8fc34e8763ae49fc",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Parallel Compiling: 0%| | 0/8 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/plain": [
"[<tilelang.jit.kernel.JITKernel at 0x7f29c0072ed0>,\n",
" <tilelang.jit.kernel.JITKernel at 0x7f29c00882f0>,\n",
" <tilelang.jit.kernel.JITKernel at 0x7f29c00735f0>,\n",
" <tilelang.jit.kernel.JITKernel at 0x7f29c0088890>,\n",
" <tilelang.jit.kernel.JITKernel at 0x7f29c01f94c0>,\n",
" <tilelang.jit.kernel.JITKernel at 0x7f29c0073fe0>,\n",
" <tilelang.jit.kernel.JITKernel at 0x7f29c0070ce0>,\n",
" <tilelang.jit.kernel.JITKernel at 0x7f29c00732f0>]"
]
},
"execution_count": 16,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from itertools import product\n",
"\n",
"\n",
"def get_configs():\n",
" return [\n",
" {\n",
" \"A\": T.Tensor((1024, 1024), T.float32),\n",
" \"B\": T.Tensor((1024, 1024), T.float32),\n",
" \"block_M\": block_M,\n",
" \"block_N\": block_N,\n",
" \"block_K\": block_K,\n",
" }\n",
" for block_M, block_N, block_K in product([32, 64], repeat=3)\n",
" ]\n",
"\n",
"\n",
"gemm.par_compile(get_configs())"
]
},
{
"cell_type": "markdown",
"id": "5160d2cc",
"metadata": {},
"source": [
"## More convenient macros"
]
},
{
"cell_type": "markdown",
"id": "be44afc4",
"metadata": {},
"source": [
"tilelang macros are now upgraded:\n",
"\n",
"1. Allow `T.Ref` as an annotation, similar to C++ pass-by-reference\n",
"2. Allow returning multiple values\n",
"3. Allow nesting and recursion"
]
},
{
"cell_type": "markdown",
"id": "79575972",
"metadata": {},
"source": [
"### Passing references with T.Ref\n",
"\n",
"The reference via T.Ref can target a var or a buffer element"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "90eaa6e5",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"# import tilelang.language as T\n",
"\n",
"@T.prim_func\n",
"def foo(x_handle: T.handle):\n",
" x = T.match_buffer(x_handle, (2,), strides=(1,))\n",
" # with T.block(\"root\"):\n",
" bx = T.launch_thread(\"blockIdx.x\", 1)\n",
" tx = T.launch_thread(\"threadIdx.x\", 128)\n",
" ty = T.launch_thread(\"threadIdx.y\", 1)\n",
" tz = T.launch_thread(\"threadIdx.z\", 1)\n",
" with T.block(\"tilelang_root\"):\n",
" T.reads()\n",
" idx = T.Buffer((1,), \"int32\", scope=\"local.var\")\n",
" T.writes(x[T.min(1, idx[0]):T.min(1, idx[0]) + (T.max(1, idx[0]) + 1 - T.min(1, idx[0]))])\n",
" T.block_attr({\"tl.local_var_init\": {idx.data: 0}})\n",
" idx = T.alloc_buffer((1,), \"int32\", data=idx.data, scope=\"local.var\")\n",
" x[1] = T.float32(1.0)\n",
" _tmp: T.int32 = idx[0]\n",
" x[_tmp] = T.float32(1.0)"
]
},
"execution_count": 17,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"@T.macro\n",
"def macro_with_ref(x: T.Ref):\n",
" x = 1 # noqa: F841\n",
"\n",
"\n",
"@T.prim_func\n",
"def foo(x: T.Tensor((2,))):\n",
" with T.Kernel(1) as _:\n",
" # Supports constant indices\n",
" macro_with_ref(x[1])\n",
"\n",
" # Also supports variable indices\n",
" idx = T.alloc_var(T.int32, 0)\n",
" macro_with_ref(x[idx])\n",
"\n",
"\n",
"foo"
]
},
{
"cell_type": "markdown",
"id": "7bb447a2",
"metadata": {},
"source": [
"### Pass as arguments\n",
"\n",
"You can pass macros as parameters"
]
},
{
"cell_type": "code",
"execution_count": 18,
"id": "dc7bb779",
"metadata": {},
"outputs": [],
"source": [
"@tilelang.lazy_jit\n",
"def element_wise(\n",
" A: T.Tensor[[T.dyn], Any],\n",
" fn,\n",
"):\n",
" (N,) = A.shape\n",
" B = T.empty((N,), dtype=A.dtype)\n",
" block_N = 128\n",
" with T.Kernel(T.ceildiv(N, block_N), threads=128) as bx:\n",
" for i in T.Parallel(block_N):\n",
" idx = bx * block_N + i\n",
" B[idx] = fn(A[idx])\n",
" return B\n",
"\n",
"\n",
"@T.macro\n",
"def add_one(x):\n",
" return x + 1"
]
},
{
"cell_type": "code",
"execution_count": 19,
"id": "a89fdb44",
"metadata": {},
"outputs": [],
"source": [
"A = torch.randn(1024, device=\"cuda\")\n",
"B = element_wise(A, add_one)\n",
"B_ref = A + 1\n",
"torch.testing.assert_close(B, B_ref)"
]
},
{
"cell_type": "markdown",
"id": "ef6e403a",
"metadata": {},
"source": [
"### Macro recursion\n",
"\n",
"Macro can be recursive, even if it's rarely needed, as long as the termination condition is known at compile time"
]
},
{
"cell_type": "code",
"execution_count": 20,
"id": "7703cab5",
"metadata": {},
"outputs": [],
"source": [
"@T.macro\n",
"def n31(x, var: T.Ref):\n",
" if x == 1:\n",
" pass\n",
" elif x % 2 == 0:\n",
" var = var // 2\n",
" n31(x // 2, var)\n",
" else:\n",
" var = var * 3 + 1\n",
" n31(x * 3 + 1, var)\n",
"\n",
"\n",
"@tilelang.lazy_jit\n",
"def foo(A: T.Tensor[[1], T.int32], n: int):\n",
" with T.Kernel(1) as _:\n",
" n31(n, A[0])"
]
},
{
"cell_type": "code",
"execution_count": 21,
"id": "542ddd4e",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"tensor([18], device='cuda:0', dtype=torch.int32)"
]
},
"execution_count": 21,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"A = torch.tensor([100], dtype=torch.int32, device=\"cuda\")\n",
"foo(A, 5)\n",
"A"
]
},
{
"cell_type": "markdown",
"id": "dc30c2d2",
"metadata": {},
"source": [
"### Macro returning multiple values"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d5a2388f",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"# import tilelang.language as T\n",
"\n",
"@T.prim_func\n",
"def foo():\n",
" # with T.block(\"root\"):\n",
" x = T.launch_thread(\"blockIdx.x\", 32)\n",
" tx = T.launch_thread(\"threadIdx.x\", 128)\n",
" ty = T.launch_thread(\"threadIdx.y\", 1)\n",
" tz = T.launch_thread(\"threadIdx.z\", 1)\n",
" with T.block(\"tilelang_root\"):\n",
" T.reads()\n",
" T.writes()\n",
" s: T.int32 = T.sin(x)\n",
" c: T.int32 = T.cos(x)\n",
" a: T.int32 = s + c\n",
" b: T.int32 = s - c\n",
" T.evaluate(0)"
]
},
"execution_count": 22,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"@T.macro\n",
"def sincos(x):\n",
" return T.sin(x), T.cos(x)\n",
"\n",
"\n",
"@T.prim_func\n",
"def foo():\n",
" with T.Kernel(32) as x:\n",
" s, c = sincos(x)\n",
" a = s + c # noqa: F841\n",
" b = s - c # noqa: F841\n",
"\n",
"\n",
"foo"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "tilelang-dev_0",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.9"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
\ No newline at end of file
examples/lazy_jit/lazyjit.zh.ipynb 0 → 100644
View file @ 667632cc
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "5e0deecc",
"metadata": {},
"outputs": [],
"source": [
"import sys\n",
"from pathlib import Path\n",
"\n",
"sys.path.insert(0, str(Path.cwd().parent.parent.absolute()))\n",
"import tilelang\n",
"import torch\n",
"import tilelang.language as T"
]
},
{
"cell_type": "markdown",
"id": "1ca2c56d",
"metadata": {},
"source": [
"# Tilelang Lazy JIT"
]
},
{
"cell_type": "markdown",
"id": "156e7370",
"metadata": {},
"source": [
"## Tensor Annotation"
]
},
{
"cell_type": "markdown",
"id": "b070c109",
"metadata": {},
"source": [
"Tilelang Lazy JIT 将 jit 生成和调用的逻辑合并到一起\n",
"\n",
"函数签名的写法与 triton 相似，但做了大量增强，最主要的增强是允许对 Tensor 的标注：\n",
"\n",
"例如，下面的代码用 T.Tensor[[int, int], T.float16] 来标注了一个二维 Tensor\n",
"1. 它的每个维度都是编译期常量，如果改变，会触发重新编译\n",
"2. 它的类型必须是 T.float16\n",
"\n",
"DType 除了写确定的外，还可以写 Any 或者 None"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "60bf8954",
"metadata": {},
"outputs": [],
"source": [
"@tilelang.lazy_jit\n",
"def gemm(\n",
" A: T.Tensor[[int, int], T.float16],\n",
" B: T.Tensor[[int, int], T.float16],\n",
" out_dtype: T.dtype = T.float32,\n",
" block_M: int = 128,\n",
" block_N: int = 128,\n",
" block_K: int = 32,\n",
"):\n",
" M, K = A.shape\n",
" K, N = B.shape\n",
" C = T.empty((M, N), out_dtype)\n",
" with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=128) as (bx, by):\n",
" A_shared = T.alloc_shared((block_M, block_K), A.dtype)\n",
" B_shared = T.alloc_shared((block_K, block_N), B.dtype)\n",
" C_local = T.alloc_fragment((block_M, block_N), out_dtype)\n",
" T.clear(C_local)\n",
" for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):\n",
" T.copy(A[bx * block_M, k * block_K], A_shared)\n",
" T.copy(B[k * block_K, by * block_N], B_shared)\n",
" T.gemm(A_shared, B_shared, C_local)\n",
" T.copy(C_local, C[bx * block_M, by * block_N])\n",
" return C"
]
},
{
"cell_type": "markdown",
"id": "28f868fe",
"metadata": {},
"source": [
"直接将 Tensor 作为参数调用，即可触发完整的 jit 编译运行流程："
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "ee13394a",
"metadata": {},
"outputs": [],
"source": [
"A = torch.randn(1024, 512, dtype=torch.float16, device=\"cuda\")\n",
"B = torch.randn(512, 256, dtype=torch.float16, device=\"cuda\")\n",
"C = gemm(A, B)\n",
"\n",
"# check output is correct\n",
"C_ref = (A @ B).float()\n",
"torch.testing.assert_close(C, C_ref, rtol=1e-2, atol=1e-2)"
]
},
{
"cell_type": "markdown",
"id": "c6705091",
"metadata": {},
"source": [
"更改调用的参数，如果编译器参数发生了变化，会触发重新编译："
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "d8aab5b7",
"metadata": {},
"outputs": [],
"source": [
"A = torch.randn(1024, 512, dtype=torch.float16, device=\"cuda\")\n",
"B = torch.randn(512, 1024, dtype=torch.float16, device=\"cuda\")\n",
"C = gemm(A, B, block_M=64, block_N=64)"
]
},
{
"cell_type": "markdown",
"id": "ce6b7391",
"metadata": {},
"source": [
"你也可以手动调用 compile 函数编译 kernel\n",
"\n",
"1. `ker.compile` 编译 kernel\n",
"2. `ker.get_tir` 获取 tir\n",
"3. `ker.par_compile` 并行编译"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "f3cf3a2d",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"2025-11-25 17:29:46 [TileLang:tilelang.cache.kernel_cache:WARNING]: Found kernel in memory cache. For better performance, consider using `@tilelang.jit` instead of direct kernel caching.\n"
]
}
],
"source": [
"kernel = gemm.compile(A, B, block_M=64, block_N=64)\n",
"C = kernel(A, B)"
]
},
{
"cell_type": "markdown",
"id": "921761b5",
"metadata": {},
"source": [
"## More Tensor Annotation"
]
},
{
"cell_type": "markdown",
"id": "4539e54e",
"metadata": {},
"source": [
"### 用 macro 来分离实现"
]
},
{
"cell_type": "markdown",
"id": "ad96ba65",
"metadata": {},
"source": [
"接下来，我们会用各种方式来实现一个简单的 gemm，为了方便，我们先写一个 macro 把 gemm 的主要逻辑写出来："
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "171d4fe6",
"metadata": {},
"outputs": [],
"source": [
"@T.macro\n",
"def gemm_impl(A, B, C, M, N, K, block_M, block_N, block_K):\n",
" with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=128) as (bx, by):\n",
" A_shared = T.alloc_shared((block_M, block_K), A.dtype)\n",
" B_shared = T.alloc_shared((block_K, block_N), B.dtype)\n",
" C_local = T.alloc_fragment((block_M, block_N), C.dtype)\n",
" T.clear(C_local)\n",
" for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):\n",
" T.copy(A[bx * block_M, k * block_K], A_shared)\n",
" T.copy(B[k * block_K, by * block_N], B_shared)\n",
" T.gemm(A_shared, B_shared, C_local)\n",
" T.copy(C_local, C[bx * block_M, by * block_N])"
]
},
{
"cell_type": "markdown",
"id": "446a1acd",
"metadata": {},
"source": [
"### 用 T.dyn 标记动态 Shape\n",
"\n",
"当某些维度是动态的的时候，可以用 T.dyn 来标记。T.dyn 可以接受一个字符串参数，表示变量的名字"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6a38aa95",
"metadata": {},
"outputs": [],
"source": [
"@tilelang.lazy_jit\n",
"def gemm_dyn_K(\n",
" A: T.Tensor[[int, T.dyn[\"K\"]], T.float16], # noqa: F821\n",
" B: T.Tensor[[T.dyn[\"K\"], int], T.float16], # noqa: F821\n",
"):\n",
" M, K = A.shape\n",
" K, N = B.shape\n",
" C = T.empty((M, N), T.float32)\n",
" gemm_impl(A, B, C, M, N, K, 128, 128, 32)\n",
" return C"
]
},
{
"cell_type": "markdown",
"id": "c60fd346",
"metadata": {},
"source": [
"查看 lazy_jit 的函数签名，其中带有后缀`$` 的是不确定的编译期常量，带有 `$dyn` 的是运行时的变量"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "c6992eb4",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'A': TensorAnnot(shape=[A_shape_0$, K$dyn], strides=None, dtype=dtype('float16')),\n",
" 'B': TensorAnnot(shape=[K$dyn, B_shape_1$], strides=None, dtype=dtype('float16'))}"
]
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"gemm_dyn_K.func.annot"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "fe6cfdc8",
"metadata": {},
"outputs": [],
"source": [
"A = torch.randn(1024, 512, dtype=torch.float16, device=\"cuda\")\n",
"B = torch.randn(512, 256, dtype=torch.float16, device=\"cuda\")\n",
"C = gemm_dyn_K(A, B)\n",
"C_ref = (A @ B).float()\n",
"torch.testing.assert_close(C, C_ref, rtol=1e-2, atol=1e-2)"
]
},
{
"cell_type": "markdown",
"id": "2ee97bf7",
"metadata": {},
"source": [
"### 用 T.StridedTensor 标记带 stride 的 Tensor\n",
"\n",
"标记方法：T.StridedTensor[Shape, Stride, DType]，每个 Shape 或 Stride 可以写\n",
"* int: 表示编译期常量\n",
"* T.dyn：表示运行时常量\n",
"\n",
"DType 可以写 None 或 Any"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "9dde1dae",
"metadata": {},
"outputs": [],
"source": [
"from typing import Any\n",
"\n",
"\n",
"@tilelang.lazy_jit\n",
"def as_contingious(A: T.StridedTensor[[T.dyn, T.dyn], [T.dyn, T.dyn], Any]):\n",
" M, N = A.shape\n",
" B = T.empty((M, N), A.dtype)\n",
" block_M = 128\n",
" block_N = 128\n",
" with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(N, block_N), threads=128) as (bx, by):\n",
" T.copy(\n",
" A[bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N],\n",
" B[bx * block_M : (bx + 1) * block_M, by * block_N : (by + 1) * block_N],\n",
" )\n",
" return B"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "dec2c0a7",
"metadata": {},
"outputs": [],
"source": [
"A = torch.randn(1024, 1024, device=\"cuda\")\n",
"B = as_contingious(A[::2, ::2])\n",
"B_ref = A[::2, ::2].contiguous()\n",
"torch.testing.assert_close(B, B_ref)"
]
},
{
"cell_type": "markdown",
"id": "f5fb20d6",
"metadata": {},
"source": [
"## More Annotation"
]
},
{
"cell_type": "markdown",
"id": "890df0a2",
"metadata": {},
"source": [
"### 用 T.ptr 标注 Tensor\n",
"lazy_jit 允许你用 T.ptr 来声明一个 handle，但必须在函数内用 T.match_buffer 给它定义 shape"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "0fc17af6",
"metadata": {},
"outputs": [],
"source": [
"@tilelang.lazy_jit\n",
"def gemm_ptr(\n",
" A: T.ptr,\n",
" B: T.ptr,\n",
" M: int,\n",
" N: int,\n",
" K: int,\n",
"):\n",
" A = T.match_buffer(A, (M, K), T.float16)\n",
" B = T.match_buffer(B, (K, N), T.float16)\n",
" C = T.empty((M, N), T.float32)\n",
" gemm_impl(A, B, C, M, N, K, block_M=128, block_N=128, block_K=32)\n",
" return C"
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "8e52a554",
"metadata": {},
"outputs": [],
"source": [
"A = torch.randn(1024, 512, dtype=torch.float16, device=\"cuda\")\n",
"B = torch.randn(512, 256, dtype=torch.float16, device=\"cuda\")\n",
"C = gemm_ptr(A, B, 1024, 256, 512)\n",
"C_ref = (A @ B).float()\n",
"torch.testing.assert_close(C, C_ref, rtol=1e-2, atol=1e-2)"
]
},
{
"cell_type": "markdown",
"id": "6b19ef90",
"metadata": {},
"source": [
"### 用 T.int32 标注运行时变量\n",
"\n",
"lazy_jit 允许你用 T.int32 或其他类型来定义运行时变量，这样，你可以写一个完全动态的 gemm，这和 triton 非常相似"
]
},
{
"cell_type": "code",
"execution_count": 14,
"id": "c1e7598a",
"metadata": {},
"outputs": [],
"source": [
"@tilelang.lazy_jit\n",
"def gemm_ptr_dyn(\n",
" A: T.ptr,\n",
" B: T.ptr,\n",
" M: T.int32,\n",
" N: T.int32,\n",
" K: T.int32,\n",
"):\n",
" A = T.match_buffer(A, (M, K), T.float16, strides=(K, 1))\n",
" B = T.match_buffer(B, (K, N), T.float16, strides=(N, 1))\n",
" C = T.empty((M, N), T.float32)\n",
" gemm_impl(A, B, C, M, N, K, block_M=128, block_N=128, block_K=32)\n",
" return C"
]
},
{
"cell_type": "code",
"execution_count": 15,
"id": "9e9a4c88",
"metadata": {},
"outputs": [],
"source": [
"A = torch.randn(1024, 512, dtype=torch.float16, device=\"cuda\")\n",
"B = torch.randn(512, 256, dtype=torch.float16, device=\"cuda\")\n",
"C = gemm_ptr_dyn(A, B, 1024, 256, 512)\n",
"C_ref = (A @ B).float()\n",
"torch.testing.assert_close(C, C_ref, rtol=1e-2, atol=1e-2)"
]
},
{
"cell_type": "markdown",
"id": "39166cb4",
"metadata": {},
"source": [
"## 编译与并行编译"
]
},
{
"cell_type": "markdown",
"id": "8c6fbe08",
"metadata": {},
"source": [
"lazyjit 和原来的 jit 都支持并行编译\n",
"\n",
"为了防止 torch.tensor 白白浪费内存，可以使用 T.Tensor 来创建 placeholder"
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "7222e57b",
"metadata": {},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "c6d7f05cdfff412e9a527332438f7aa2",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Elaborating: 0%| | 0/8 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "14836065a21b41ae8fc34e8763ae49fc",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Parallel Compiling: 0%| | 0/8 [00:00<?, ?it/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/plain": [
"[<tilelang.jit.kernel.JITKernel at 0x7f29c0072ed0>,\n",
" <tilelang.jit.kernel.JITKernel at 0x7f29c00882f0>,\n",
" <tilelang.jit.kernel.JITKernel at 0x7f29c00735f0>,\n",
" <tilelang.jit.kernel.JITKernel at 0x7f29c0088890>,\n",
" <tilelang.jit.kernel.JITKernel at 0x7f29c01f94c0>,\n",
" <tilelang.jit.kernel.JITKernel at 0x7f29c0073fe0>,\n",
" <tilelang.jit.kernel.JITKernel at 0x7f29c0070ce0>,\n",
" <tilelang.jit.kernel.JITKernel at 0x7f29c00732f0>]"
]
},
"execution_count": 16,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"from itertools import product\n",
"\n",
"\n",
"def get_configs():\n",
" return [\n",
" {\n",
" \"A\": T.Tensor((1024, 1024), T.float32),\n",
" \"B\": T.Tensor((1024, 1024), T.float32),\n",
" \"block_M\": block_M,\n",
" \"block_N\": block_N,\n",
" \"block_K\": block_K,\n",
" }\n",
" for block_M, block_N, block_K in product([32, 64], repeat=3)\n",
" ]\n",
"\n",
"\n",
"gemm.par_compile(get_configs())"
]
},
{
"cell_type": "markdown",
"id": "5160d2cc",
"metadata": {},
"source": [
"## 更便利的 Macro"
]
},
{
"cell_type": "markdown",
"id": "be44afc4",
"metadata": {},
"source": [
"tilelang 的 macro 现在已经升级：\n",
"\n",
"1. 允许用 `T.Ref` 作为 annotation，这类似与 C++ 的引用传递\n",
"2. 允许返回多个值\n",
"3. 允许嵌套，递归"
]
},
{
"cell_type": "markdown",
"id": "79575972",
"metadata": {},
"source": [
"### T.Ref 传递引用\n",
"\n",
"T.Ref 传递的引用可以 var 也可以是 Buffer 的索引"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "90eaa6e5",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"# import tilelang.language as T\n",
"\n",
"@T.prim_func\n",
"def foo(x_handle: T.handle):\n",
" x = T.match_buffer(x_handle, (2,), strides=(1,))\n",
" # with T.block(\"root\"):\n",
" bx = T.launch_thread(\"blockIdx.x\", 1)\n",
" tx = T.launch_thread(\"threadIdx.x\", 128)\n",
" ty = T.launch_thread(\"threadIdx.y\", 1)\n",
" tz = T.launch_thread(\"threadIdx.z\", 1)\n",
" with T.block(\"tilelang_root\"):\n",
" T.reads()\n",
" idx = T.Buffer((1,), \"int32\", scope=\"local.var\")\n",
" T.writes(x[T.min(1, idx[0]):T.min(1, idx[0]) + (T.max(1, idx[0]) + 1 - T.min(1, idx[0]))])\n",
" T.block_attr({\"tl.local_var_init\": {idx.data: 0}})\n",
" idx = T.alloc_buffer((1,), \"int32\", data=idx.data, scope=\"local.var\")\n",
" x[1] = T.float32(1.0)\n",
" _tmp: T.int32 = idx[0]\n",
" x[_tmp] = T.float32(1.0)"
]
},
"execution_count": 17,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"@T.macro\n",
"def macro_with_ref(x: T.Ref):\n",
" x = 1 # noqa: F841\n",
"\n",
"\n",
"@T.prim_func\n",
"def foo(x: T.Tensor((2,))):\n",
" with T.Kernel(1) as _:\n",
" # 支持常量 index\n",
" macro_with_ref(x[1])\n",
"\n",
" # 也支持变量 index\n",
" idx = T.alloc_var(T.int32, 0)\n",
" macro_with_ref(x[idx])\n",
"\n",
"\n",
"foo"
]
},
{
"cell_type": "markdown",
"id": "7bb447a2",
"metadata": {},
"source": [
"### 当作参数传递\n",
"\n",
"你可以把 macro 当做参数传递"
]
},
{
"cell_type": "code",
"execution_count": 18,
"id": "dc7bb779",
"metadata": {},
"outputs": [],
"source": [
"@tilelang.lazy_jit\n",
"def element_wise(\n",
" A: T.Tensor[[T.dyn], Any],\n",
" fn,\n",
"):\n",
" (N,) = A.shape\n",
" B = T.empty((N,), dtype=A.dtype)\n",
" block_N = 128\n",
" with T.Kernel(T.ceildiv(N, block_N), threads=128) as bx:\n",
" for i in T.Parallel(block_N):\n",
" idx = bx * block_N + i\n",
" B[idx] = fn(A[idx])\n",
" return B\n",
"\n",
"\n",
"@T.macro\n",
"def add_one(x):\n",
" return x + 1"
]
},
{
"cell_type": "code",
"execution_count": 19,
"id": "a89fdb44",
"metadata": {},
"outputs": [],
"source": [
"A = torch.randn(1024, device=\"cuda\")\n",
"B = element_wise(A, add_one)\n",
"B_ref = A + 1\n",
"torch.testing.assert_close(B, B_ref)"
]
},
{
"cell_type": "markdown",
"id": "ef6e403a",
"metadata": {},
"source": [
"### Macro 递归\n",
"\n",
"虽然不知道有没有这种需求，但 macro 是可以递归的，但要求终止条件编译期间确定"
]
},
{
"cell_type": "code",
"execution_count": 20,
"id": "7703cab5",
"metadata": {},
"outputs": [],
"source": [
"@T.macro\n",
"def n31(x, var: T.Ref):\n",
" if x == 1:\n",
" pass\n",
" elif x % 2 == 0:\n",
" var = var // 2\n",
" n31(x // 2, var)\n",
" else:\n",
" var = var * 3 + 1\n",
" n31(x * 3 + 1, var)\n",
"\n",
"\n",
"@tilelang.lazy_jit\n",
"def foo(A: T.Tensor[[1], T.int32], n: int):\n",
" with T.Kernel(1) as _:\n",
" n31(n, A[0])"
]
},
{
"cell_type": "code",
"execution_count": 21,
"id": "542ddd4e",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"tensor([18], device='cuda:0', dtype=torch.int32)"
]
},
"execution_count": 21,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"A = torch.tensor([100], dtype=torch.int32, device=\"cuda\")\n",
"foo(A, 5)\n",
"A"
]
},
{
"cell_type": "markdown",
"id": "dc30c2d2",
"metadata": {},
"source": [
"### Macro 返回多个值"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d5a2388f",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"# import tilelang.language as T\n",
"\n",
"@T.prim_func\n",
"def foo():\n",
" # with T.block(\"root\"):\n",
" x = T.launch_thread(\"blockIdx.x\", 32)\n",
" tx = T.launch_thread(\"threadIdx.x\", 128)\n",
" ty = T.launch_thread(\"threadIdx.y\", 1)\n",
" tz = T.launch_thread(\"threadIdx.z\", 1)\n",
" with T.block(\"tilelang_root\"):\n",
" T.reads()\n",
" T.writes()\n",
" s: T.int32 = T.sin(x)\n",
" c: T.int32 = T.cos(x)\n",
" a: T.int32 = s + c\n",
" b: T.int32 = s - c\n",
" T.evaluate(0)"
]
},
"execution_count": 22,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"@T.macro\n",
"def sincos(x):\n",
" return T.sin(x), T.cos(x)\n",
"\n",
"\n",
"@T.prim_func\n",
"def foo():\n",
" with T.Kernel(32) as x:\n",
" s, c = sincos(x)\n",
" a = s + c # noqa: F841\n",
" b = s - c # noqa: F841\n",
"\n",
"\n",
"foo"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "tilelang-dev_0",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.9"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
\ No newline at end of file
examples/linear_attention/example_linear_attn_bwd.py
View file @ 667632cc
......@@ -13,20 +13,20 @@ from typing import Optional, Tuple
pass_configs={
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
})
}
)
def tl_fused_chunk_bwd_kernel(
B,
S,
H,
DK,
DV,
dtype: str = 'float16',
dtype: T.dtype = T.float16,
scale: float = None,
) -> torch.Tensor:
if scale is None:
scale = DK**-0.5
accum_dtype = 'float'
accum_dtype = T.float32
chunk_size = 64
BK = BV = 64 # Set to 128 can be faster, but has some numerical differences with FLA
......@@ -37,13 +37,13 @@ def tl_fused_chunk_bwd_kernel(
@T.prim_func
def fused_chunk_linear_attn_bwd(
Q: T.Tensor([B, S, H, DK], dtype), # type: ignore
K: T.Tensor([B, S, H, DK], dtype), # type: ignore
V: T.Tensor([B, S, H, DV], dtype), # type: ignore
dO: T.Tensor([B, S, H, DV], dtype), # type: ignore
dQ: T.Tensor([B, S, H, DK], accum_dtype), # type: ignore
dK: T.Tensor([B, S, H, DK], accum_dtype), # type: ignore
dV: T.Tensor([B, S, H, DV], accum_dtype), # type: ignore
Q: T.Tensor([B, S, H, DK], dtype), # type: ignore
K: T.Tensor([B, S, H, DK], dtype), # type: ignore
V: T.Tensor([B, S, H, DV], dtype), # type: ignore
dO: T.Tensor([B, S, H, DV], dtype), # type: ignore
dQ: T.Tensor([B, S, H, DK], accum_dtype), # type: ignore
dK: T.Tensor([B, S, H, DK], accum_dtype), # type: ignore
dV: T.Tensor([B, S, H, DV], accum_dtype), # type: ignore
):
with T.Kernel(NV, NK, B * H) as (i_v, i_k, i_bh):
i_b = i_bh // H
......@@ -66,11 +66,13 @@ def tl_fused_chunk_bwd_kernel(
dh = T.alloc_fragment([BK, BV], accum_dtype)
dh_shared = T.alloc_shared([BK, BV], dtype)
T.annotate_layout({
dq_shared: tilelang.layout.make_swizzled_layout(dq_shared),
dk_shared: tilelang.layout.make_swizzled_layout(dk_shared),
dv_shared: tilelang.layout.make_swizzled_layout(dv_shared)
})
T.annotate_layout(
{
dq_shared: tilelang.layout.make_swizzled_layout(dq_shared),
dk_shared: tilelang.layout.make_swizzled_layout(dk_shared),
dv_shared: tilelang.layout.make_swizzled_layout(dv_shared),
}
)
T.use_swizzle(10)
T.clear(h)
......@@ -78,10 +80,9 @@ def tl_fused_chunk_bwd_kernel(
# Calculate dQ
for i in T.Pipelined(0, NT):
T.copy(K[i_b, i * chunk_size:(i + 1) * chunk_size, i_h, i_k * BK:(i_k + 1) * BK], k)
T.copy(V[i_b, i * chunk_size:(i + 1) * chunk_size, i_h, i_v * BV:(i_v + 1) * BV], v)
T.copy(dO[i_b, i * chunk_size:(i + 1) * chunk_size, i_h, i_v * BV:(i_v + 1) * BV],
do)
T.copy(K[i_b, i * chunk_size : (i + 1) * chunk_size, i_h, i_k * BK : (i_k + 1) * BK], k)
T.copy(V[i_b, i * chunk_size : (i + 1) * chunk_size, i_h, i_v * BV : (i_v + 1) * BV], v)
T.copy(dO[i_b, i * chunk_size : (i + 1) * chunk_size, i_h, i_v * BV : (i_v + 1) * BV], do)
T.gemm(do, v, ds, transpose_B=True, clear_accum=True)
for row, col in T.Parallel(chunk_size, chunk_size):
......@@ -94,29 +95,19 @@ def tl_fused_chunk_bwd_kernel(
for row, col in T.Parallel(chunk_size, BK):
dq[row, col] *= scale
T.copy(dq, dq_shared)
T.atomic_add(
dQ[i_b, i * chunk_size:(i + 1) * chunk_size, i_h, i_k * BK:(i_k + 1) * BK],
dq_shared)
T.atomic_add(dQ[i_b, i * chunk_size : (i + 1) * chunk_size, i_h, i_k * BK : (i_k + 1) * BK], dq_shared)
# Calculate dK, dV (reversely)
for i in T.Pipelined(1, NT + 1):
start = NT - i
for row, col in T.Parallel(chunk_size, BK):
q[row, col] = Q[i_b, start * chunk_size + row, i_h, i_k * BK + col] * scale
T.copy(
K[i_b, start * chunk_size:(start + 1) * chunk_size, i_h,
i_k * BK:(i_k + 1) * BK], k)
T.copy(
V[i_b, start * chunk_size:(start + 1) * chunk_size, i_h,
i_v * BV:(i_v + 1) * BV], v)
T.copy(
dO[i_b, start * chunk_size:(start + 1) * chunk_size, i_h,
i_v * BV:(i_v + 1) * BV], do)
T.copy(K[i_b, start * chunk_size : (start + 1) * chunk_size, i_h, i_k * BK : (i_k + 1) * BK], k)
T.copy(V[i_b, start * chunk_size : (start + 1) * chunk_size, i_h, i_v * BV : (i_v + 1) * BV], v)
T.copy(dO[i_b, start * chunk_size : (start + 1) * chunk_size, i_h, i_v * BV : (i_v + 1) * BV], do)
# Calculate dk
T.gemm(
v, do, ds, transpose_B=True, clear_accum=True
) # ds here actually means `s`, but we simply reuse the buffer `ds`
T.gemm(v, do, ds, transpose_B=True, clear_accum=True) # ds here actually means `s`, but we simply reuse the buffer `ds`
for row, col in T.Parallel(chunk_size, chunk_size):
ds_shared[row, col] = T.if_then_else(row <= col, ds[row, col], 0)
T.gemm(ds_shared, q, dk, clear_accum=True)
......@@ -134,13 +125,9 @@ def tl_fused_chunk_bwd_kernel(
T.gemm(q, do, dh, transpose_A=True)
T.copy(dk, dk_shared)
T.atomic_add(
dK[i_b, start * chunk_size:(start + 1) * chunk_size, i_h,
i_k * BK:(i_k + 1) * BK], dk_shared)
T.atomic_add(dK[i_b, start * chunk_size : (start + 1) * chunk_size, i_h, i_k * BK : (i_k + 1) * BK], dk_shared)
T.copy(dv, dv_shared)
T.atomic_add(
dV[i_b, start * chunk_size:(start + 1) * chunk_size, i_h,
i_v * BV:(i_v + 1) * BV], dv_shared)
T.atomic_add(dV[i_b, start * chunk_size : (start + 1) * chunk_size, i_h, i_v * BV : (i_v + 1) * BV], dv_shared)
return fused_chunk_linear_attn_bwd
......@@ -155,34 +142,31 @@ def tl_fused_chunk_bwd(Q, K, V, dO):
return dQ.to(torch.float16), dK.to(torch.float16), dV.to(torch.float16)
def ref_program(q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
scale: Optional[float] = None) -> Tuple[torch.Tensor, torch.Tensor]:
def ref_program(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, scale: Optional[float] = None) -> Tuple[torch.Tensor, torch.Tensor]:
q, k, v = q.float(), k.float(), v.float()
if scale is None:
scale = q.shape[-1]**-0.5
scale = q.shape[-1] ** -0.5
chunk_size = 64
q = rearrange(q, 'b (n c) h d -> b h n c d', c=chunk_size) * scale
k = rearrange(k, 'b (n c) h d -> b h n c d', c=chunk_size)
v = rearrange(v, 'b (n c) h d -> b h n c d', c=chunk_size)
q = rearrange(q, "b (n c) h d -> b h n c d", c=chunk_size) * scale
k = rearrange(k, "b (n c) h d -> b h n c d", c=chunk_size)
v = rearrange(v, "b (n c) h d -> b h n c d", c=chunk_size)
kv = k.transpose(-1, -2) @ v
kv = kv.cumsum(2)
h = kv[:, :, -1, :, :]
kv = torch.cat([torch.zeros_like(kv[:, :, :1]), kv[:, :, :-1]], dim=2)
inter = q @ kv
intra = ((q @ k.transpose(-1, -2)).masked_fill_(
torch.triu(torch.ones(chunk_size, chunk_size, dtype=bool, device=q.device), diagonal=1),
0)) @ v
intra = (
(q @ k.transpose(-1, -2)).masked_fill_(torch.triu(torch.ones(chunk_size, chunk_size, dtype=bool, device=q.device), diagonal=1), 0)
) @ v
o = inter + intra
return rearrange(o, 'b h n c d -> b (n c) h d'), h
return rearrange(o, "b h n c d -> b (n c) h d"), h
def main(B=1, S=1024, H=16, D=128):
q = torch.randn((B, S, H, D), device='cuda', dtype=torch.float16, requires_grad=True)
k = torch.randn((B, S, H, D), device='cuda', dtype=torch.float16, requires_grad=True)
v = torch.randn((B, S, H, D), device='cuda', dtype=torch.float16, requires_grad=True)
do = torch.randn((B, S, H, D), device='cuda', dtype=torch.float16)
q = torch.randn((B, S, H, D), device="cuda", dtype=torch.float16, requires_grad=True)
k = torch.randn((B, S, H, D), device="cuda", dtype=torch.float16, requires_grad=True)
v = torch.randn((B, S, H, D), device="cuda", dtype=torch.float16, requires_grad=True)
do = torch.randn((B, S, H, D), device="cuda", dtype=torch.float16)
# qk norm is necessary for linear attn
q = l2norm_fwd(q)[0].requires_grad_(True)
......@@ -193,30 +177,27 @@ def main(B=1, S=1024, H=16, D=128):
o_ref, _ = ref_program(q, k, v)
o_ref.backward(do, retain_graph=True)
assert torch.allclose(
dq, q.grad, atol=1e-2, rtol=1e-2), f'dq max err: {(dq - q.grad).abs().max()}'
assert torch.allclose(
dk, k.grad, atol=1e-2, rtol=1e-2), f'dk max err: {(dk - k.grad).abs().max()}'
assert torch.allclose(
dv, v.grad, atol=1e-2, rtol=1e-2), f'dv max err: {(dv - v.grad).abs().max()}'
print('Passed all tests!✅')
assert torch.allclose(dq, q.grad, atol=1e-2, rtol=1e-2), f"dq max err: {(dq - q.grad).abs().max()}"
assert torch.allclose(dk, k.grad, atol=1e-2, rtol=1e-2), f"dk max err: {(dk - k.grad).abs().max()}"
assert torch.allclose(dv, v.grad, atol=1e-2, rtol=1e-2), f"dv max err: {(dv - v.grad).abs().max()}"
print("Passed all tests!✅")
# Benchmark
q.grad = k.grad = v.grad = None
o_ref, _ = fused_chunk_linear_attn(q, k, v, output_final_state=True, normalize=False)
t1 = do_bench(lambda: o_ref.backward(do, retain_graph=True), backend='cupti')
t2 = do_bench(lambda: tl_fused_chunk_bwd(q, k, v, do), backend='cupti')
print(f'Triton latency: {t1:.3f} ms')
print(f'TileLang latency: {t2:.3f} ms')
print(f'Speedup: {t1/t2:.3f}x')
t1 = do_bench(lambda: o_ref.backward(do, retain_graph=True), backend="cupti")
t2 = do_bench(lambda: tl_fused_chunk_bwd(q, k, v, do), backend="cupti")
print(f"Triton latency: {t1:.3f} ms")
print(f"TileLang latency: {t2:.3f} ms")
print(f"Speedup: {t1 / t2:.3f}x")
if __name__ == '__main__':
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--B', type=int, default=8, help='Batch size')
parser.add_argument('--S', type=int, default=1024, help='Seq len')
parser.add_argument('--H', type=int, default=32, help='Num heads')
parser.add_argument('--D', type=int, default=128, help='Head dim')
parser.add_argument("--B", type=int, default=8, help="Batch size")
parser.add_argument("--S", type=int, default=1024, help="Seq len")
parser.add_argument("--H", type=int, default=32, help="Num heads")
parser.add_argument("--D", type=int, default=128, help="Head dim")
args = parser.parse_args()
main(args.B, args.S, args.H, args.D)
examples/linear_attention/example_linear_attn_fwd.py
View file @ 667632cc
......@@ -14,20 +14,20 @@ from typing import Optional, Tuple
pass_configs={
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
})
},
)
def tl_fused_chunk_fwd_kernel(
B,
S,
H,
DK,
DV,
dtype: str = 'float16',
dtype: T.dtype = T.float16,
scale: float = None,
) -> torch.Tensor:
if scale is None:
scale = DK**-0.5
accum_dtype = 'float'
accum_dtype = T.float32
chunk_size = 64
BK = BV = 64 # Set to 128 can be faster, but has some numerical differences with FLA
......@@ -38,11 +38,12 @@ def tl_fused_chunk_fwd_kernel(
@T.prim_func
def fused_chunk_linear_attn_fwd(
Q: T.Tensor([B, S, H, DK], dtype), # type: ignore
K: T.Tensor([B, S, H, DK], dtype), # type: ignore
V: T.Tensor([B, S, H, DV], dtype), # type: ignore
O: T.Tensor([B, S, H, DV], accum_dtype), # type: ignore
final_state: T.Tensor([B, H, DK, DV], accum_dtype)): # type: ignore
Q: T.Tensor([B, S, H, DK], dtype), # type: ignore
K: T.Tensor([B, S, H, DK], dtype), # type: ignore
V: T.Tensor([B, S, H, DV], dtype), # type: ignore
O: T.Tensor([B, S, H, DV], accum_dtype), # type: ignore
final_state: T.Tensor([B, H, DK, DV], accum_dtype),
): # type: ignore
with T.Kernel(NV, NK, B * H) as (i_v, i_k, i_bh):
i_b = i_bh // H
i_h = i_bh % H
......@@ -65,8 +66,8 @@ def tl_fused_chunk_fwd_kernel(
for i in T.Pipelined(0, NT):
for row, col in T.Parallel(chunk_size, BK):
q[row, col] = Q[i_b, i * chunk_size + row, i_h, i_k * BK + col] * scale
T.copy(K[i_b, i * chunk_size:(i + 1) * chunk_size, i_h, i_k * BK:(i_k + 1) * BK], k)
T.copy(V[i_b, i * chunk_size:(i + 1) * chunk_size, i_h, i_v * BV:(i_v + 1) * BV], v)
T.copy(K[i_b, i * chunk_size : (i + 1) * chunk_size, i_h, i_k * BK : (i_k + 1) * BK], k)
T.copy(V[i_b, i * chunk_size : (i + 1) * chunk_size, i_h, i_v * BV : (i_v + 1) * BV], v)
T.gemm(q, k, s, clear_accum=True, transpose_B=True)
for row, col in T.Parallel(chunk_size, chunk_size):
......@@ -77,12 +78,10 @@ def tl_fused_chunk_fwd_kernel(
T.gemm(k, v, h, transpose_A=True)
T.gemm(q, h_shared, o)
T.copy(o, o_shared)
T.atomic_add(
O[i_b, i * chunk_size:(i + 1) * chunk_size, i_h, i_v * BV:(i_v + 1) * BV],
o_shared)
T.atomic_add(O[i_b, i * chunk_size : (i + 1) * chunk_size, i_h, i_v * BV : (i_v + 1) * BV], o_shared)
# Output final state
T.copy(h, final_state[i_b, i_h, i_k * BK:(i_k + 1) * BK, i_v * BV:(i_v + 1) * BV])
T.copy(h, final_state[i_b, i_h, i_k * BK : (i_k + 1) * BK, i_v * BV : (i_v + 1) * BV])
return fused_chunk_linear_attn_fwd
......@@ -91,38 +90,35 @@ def tl_fused_chunk_fwd(q, k, v):
B, S, H, D = q.shape
kernel = tl_fused_chunk_fwd_kernel(B, S, H, D, D)
print(kernel.get_kernel_source())
o = torch.zeros((B, S, H, D), device='cuda', dtype=torch.float32)
o = torch.zeros((B, S, H, D), device="cuda", dtype=torch.float32)
h = kernel(q, k, v, o)
return o, h
def ref_program(q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
scale: Optional[float] = None) -> Tuple[torch.Tensor, torch.Tensor]:
def ref_program(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, scale: Optional[float] = None) -> Tuple[torch.Tensor, torch.Tensor]:
q, k, v = q.float(), k.float(), v.float()
if scale is None:
scale = q.shape[-1]**-0.5
scale = q.shape[-1] ** -0.5
chunk_size = 64
q = rearrange(q, 'b (n c) h d -> b h n c d', c=chunk_size) * scale
k = rearrange(k, 'b (n c) h d -> b h n c d', c=chunk_size)
v = rearrange(v, 'b (n c) h d -> b h n c d', c=chunk_size)
q = rearrange(q, "b (n c) h d -> b h n c d", c=chunk_size) * scale
k = rearrange(k, "b (n c) h d -> b h n c d", c=chunk_size)
v = rearrange(v, "b (n c) h d -> b h n c d", c=chunk_size)
kv = k.transpose(-1, -2) @ v
kv = kv.cumsum(2)
h = kv[:, :, -1, :, :]
kv = torch.cat([torch.zeros_like(kv[:, :, :1]), kv[:, :, :-1]], dim=2)
inter = q @ kv
intra = ((q @ k.transpose(-1, -2)).masked_fill_(
torch.triu(torch.ones(chunk_size, chunk_size, dtype=bool, device=q.device), diagonal=1),
0)) @ v
intra = (
(q @ k.transpose(-1, -2)).masked_fill_(torch.triu(torch.ones(chunk_size, chunk_size, dtype=bool, device=q.device), diagonal=1), 0)
) @ v
o = inter + intra
return rearrange(o, 'b h n c d -> b (n c) h d'), h
return rearrange(o, "b h n c d -> b (n c) h d"), h
def main(B=1, S=512, H=16, D=128):
q = torch.randn((B, S, H, D), device='cuda', dtype=torch.float16)
k = torch.randn((B, S, H, D), device='cuda', dtype=torch.float16)
v = torch.randn((B, S, H, D), device='cuda', dtype=torch.float16)
q = torch.randn((B, S, H, D), device="cuda", dtype=torch.float16)
k = torch.randn((B, S, H, D), device="cuda", dtype=torch.float16)
v = torch.randn((B, S, H, D), device="cuda", dtype=torch.float16)
# qk norm is necessary for linear attn
q, _ = l2norm_fwd(q)
......@@ -131,25 +127,23 @@ def main(B=1, S=512, H=16, D=128):
o, h = tl_fused_chunk_fwd(q, k, v)
o_ref, h_ref = ref_program(q, k, v)
assert torch.allclose(o, o_ref, atol=1e-2, rtol=1e-2), f'o max err: {(o - o_ref).abs().max()}'
assert torch.allclose(h, h_ref, atol=1e-2, rtol=1e-2), f'h max err: {(h - h_ref).abs().max()}'
print('Passed all tests!✅')
assert torch.allclose(o, o_ref, atol=1e-2, rtol=1e-2), f"o max err: {(o - o_ref).abs().max()}"
assert torch.allclose(h, h_ref, atol=1e-2, rtol=1e-2), f"h max err: {(h - h_ref).abs().max()}"
print("Passed all tests!✅")
t1 = do_bench(
lambda: fused_chunk_linear_attn(q, k, v, output_final_state=True, normalize=False),
backend='cupti')
t2 = do_bench(lambda: tl_fused_chunk_fwd(q, k, v), backend='cupti')
print(f'Triton latency: {t1:.3f} ms')
print(f'TileLang latency: {t2:.3f} ms')
print(f'Speedup: {t1/t2:.3f}x')
t1 = do_bench(lambda: fused_chunk_linear_attn(q, k, v, output_final_state=True, normalize=False), backend="cupti")
t2 = do_bench(lambda: tl_fused_chunk_fwd(q, k, v), backend="cupti")
print(f"Triton latency: {t1:.3f} ms")
print(f"TileLang latency: {t2:.3f} ms")
print(f"Speedup: {t1 / t2:.3f}x")
if __name__ == '__main__':
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--B', type=int, default=8, help='Batch size')
parser.add_argument('--S', type=int, default=1024, help='Seq len')
parser.add_argument('--H', type=int, default=32, help='Num heads')
parser.add_argument('--D', type=int, default=128, help='Head dim')
parser.add_argument("--B", type=int, default=8, help="Batch size")
parser.add_argument("--S", type=int, default=1024, help="Seq len")
parser.add_argument("--H", type=int, default=32, help="Num heads")
parser.add_argument("--D", type=int, default=128, help="Head dim")
args = parser.parse_args()
main(args.B, args.S, args.H, args.D)
examples/linear_attention/example_mamba_chunk_scan.py
View file @ 667632cc
......@@ -9,6 +9,7 @@ import itertools
def chunk_scan_triton(cb, x, dt, dA_cumsum, C, states, D):
from mamba_ssm.ops.triton.ssd_chunk_scan import _chunk_scan_fwd
out, _ = _chunk_scan_fwd(cb, x, dt, dA_cumsum, C, states, D)
return out
......@@ -43,14 +44,15 @@ def ref_program(cb, x, dt, dA_cumsum, C, prev_states, D):
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)
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))
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_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:
......@@ -61,12 +63,7 @@ def ref_program(cb, x, dt, dA_cumsum, C, prev_states, D):
def get_configs():
iter_params = dict(
block_M=[64, 128, 256],
block_N=[32, 64],
block_K=[64, 128, 256],
block_Dstate=[128],
num_stages=[1, 2, 3, 4, 5])
iter_params = dict(block_M=[64, 128, 256], block_N=[32, 64], block_K=[64, 128, 256], block_Dstate=[128], num_stages=[1, 2, 3, 4, 5])
return [dict(zip(iter_params, values)) for values in itertools.product(*iter_params.values())]
......@@ -77,40 +74,42 @@ def get_configs():
tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
},
)
def chunk_scan_fwd(batch,
seqlen,
chunk_size,
ngroups,
nheads,
headdim,
dstate,
block_M=64,
block_N=64,
block_K=64,
block_Dstate=128,
num_stages=2,
threads=128):
dtype = "float16"
accum_dtype = "float"
def chunk_scan_fwd(
batch,
seqlen,
chunk_size,
ngroups,
nheads,
headdim,
dstate,
block_M=64,
block_N=64,
block_K=64,
block_Dstate=128,
num_stages=2,
threads=128,
):
dtype = T.float16
accum_dtype = T.float32
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), # type: ignore
x: T.Tensor((batch, seqlen, nheads, headdim), dtype), # type: ignore
dt: T.Tensor((batch, nheads, nchunks, chunk_size), dtype), # type: ignore
dA_cumsum: T.Tensor((batch, nheads, nchunks, chunk_size), dtype), # type: ignore
C: T.Tensor((batch, seqlen, ngroups, dstate), dtype), # type: ignore
prev_states: T.Tensor((batch, nchunks, nheads, headdim, dstate), dtype), # type: ignore
D: T.Tensor((nheads), dtype), # type: ignore
Output: T.Tensor((batch, seqlen, nheads, headdim), dtype) # type: ignore
cb: T.Tensor((batch, nchunks, ngroups, chunk_size, chunk_size), dtype), # type: ignore
x: T.Tensor((batch, seqlen, nheads, headdim), dtype), # type: ignore
dt: T.Tensor((batch, nheads, nchunks, chunk_size), dtype), # type: ignore
dA_cumsum: T.Tensor((batch, nheads, nchunks, chunk_size), dtype), # type: ignore
C: T.Tensor((batch, seqlen, ngroups, dstate), dtype), # type: ignore
prev_states: T.Tensor((batch, nchunks, nheads, headdim, dstate), dtype), # type: ignore
D: T.Tensor((nheads), dtype), # type: ignore
Output: T.Tensor((batch, seqlen, nheads, headdim), dtype), # type: ignore
):
with T.Kernel(
nheads,
T.ceildiv(chunk_size, block_M) * T.ceildiv(headdim, block_N),
batch * nchunks,
threads=threads) as (bz, bx, by):
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)
acc_o_shared = T.alloc_shared((block_M, block_N), dtype)
cb_shared = T.alloc_shared((block_M, block_K), dtype, scope="shared.dyn")
......@@ -136,27 +135,32 @@ def chunk_scan_fwd(batch,
m_idx = bx // T.ceildiv(headdim, block_N)
n_idx = bx % T.ceildiv(headdim, block_N)
T.annotate_layout({
acc_o_shared: tilelang.layout.make_swizzled_layout(acc_o_shared),
cb_shared: tilelang.layout.make_swizzled_layout(cb_shared),
x_residual_shared: tilelang.layout.make_swizzled_layout(x_residual_shared)
})
T.annotate_layout(
{
acc_o_shared: tilelang.layout.make_swizzled_layout(acc_o_shared),
cb_shared: tilelang.layout.make_swizzled_layout(cb_shared),
x_residual_shared: tilelang.layout.make_swizzled_layout(x_residual_shared),
}
)
T.no_set_max_nreg()
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_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)
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]
......@@ -165,34 +169,47 @@ def chunk_scan_fwd(batch,
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)
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_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)
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)
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)
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)
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]
......@@ -200,27 +217,40 @@ def chunk_scan_fwd(batch,
T.copy(acc_o, acc_o_shared)
T.copy(
acc_o_shared,
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])
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
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--batch', type=int, default=8, help='batch size')
parser.add_argument('--heads', type=int, default=80, help='heads')
parser.add_argument('--groups', type=int, default=1, help='groups')
parser.add_argument('--seq_len', type=int, default=4096, help='sequence length')
parser.add_argument('--chunk_size', type=int, default=256, help='chunk size')
parser.add_argument('--dim', type=int, default=64, help='dim')
parser.add_argument('--dstate', type=int, default=128, help='dstate')
parser.add_argument('--tune', action='store_true', help='tune configs')
parser.add_argument("--batch", type=int, default=8, help="batch size")
parser.add_argument("--heads", type=int, default=80, help="heads")
parser.add_argument("--groups", type=int, default=1, help="groups")
parser.add_argument("--seq_len", type=int, default=4096, help="sequence length")
parser.add_argument("--chunk_size", type=int, default=256, help="chunk size")
parser.add_argument("--dim", type=int, default=64, help="dim")
parser.add_argument("--dstate", type=int, default=128, help="dstate")
parser.add_argument("--tune", action="store_true", help="tune configs")
args = parser.parse_args()
batch, heads, groups, seq_len, chunk_size, dim, dstate = args.batch, args.heads, args.groups, args.seq_len, args.chunk_size, args.dim, args.dstate
batch, heads, groups, seq_len, chunk_size, dim, dstate = (
args.batch,
args.heads,
args.groups,
args.seq_len,
args.chunk_size,
args.dim,
args.dstate,
)
total_flops = 2 * batch * seq_len * chunk_size * heads * dim * 0.5 + 2 * batch * seq_len * heads * dim * dstate
if (not args.tune):
if not args.tune:
kernel = chunk_scan_fwd(
batch,
seq_len,
......@@ -234,7 +264,8 @@ if __name__ == "__main__":
block_K=64,
block_Dstate=128,
num_stages=2,
threads=128)
threads=128,
)
profiler = kernel.get_profiler(tilelang.TensorSupplyType.Normal)
profiler.assert_allclose(ref_program, rtol=0.01, atol=0.01)
print("All checks pass.")
......
examples/linear_attention/example_mamba_chunk_state.py
View file @ 667632cc
......@@ -10,6 +10,7 @@ import itertools
def chunk_state_triton(B, x, dt, dA_cumsum):
from mamba_ssm.ops.triton.ssd_chunk_state import _chunk_state_fwd
return _chunk_state_fwd(B, x, dt, dA_cumsum, states_in_fp32=False)
......@@ -41,46 +42,33 @@ def ref_program(B, x, dt, dA_cumsum):
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)
return torch.einsum("bclhn,bhcl,bhcl,bclhp->bchpn", B.to(x.dtype), decay_states.to(x.dtype), dt.to(x.dtype), x)
def get_configs():
iter_params = dict(
block_M=[64, 128], block_N=[32, 64, 128], block_K=[32, 64], num_stages=[1, 2, 3, 4, 5])
iter_params = dict(block_M=[64, 128], block_N=[32, 64, 128], block_K=[32, 64], num_stages=[1, 2, 3, 4, 5])
return [dict(zip(iter_params, values)) for values in itertools.product(*iter_params.values())]
@autotune(configs=get_configs(), warmup=10, rep=10)
@tilelang.jit(out_idx=[4])
def chunk_state_fwd(batch,
seqlen,
chunk_size,
ngroups,
nheads,
headdim,
dstate,
block_M=64,
block_N=64,
block_K=64,
num_stages=2,
threads=128):
dtype = "float16"
accum_dtype = "float"
def chunk_state_fwd(
batch, seqlen, chunk_size, ngroups, nheads, headdim, dstate, block_M=64, block_N=64, block_K=64, num_stages=2, threads=128
):
dtype = T.float16
accum_dtype = T.float32
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):
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)
......@@ -101,20 +89,24 @@ def chunk_state_fwd(batch,
m_idx = bx // T.ceildiv(dstate, block_N)
n_idx = bx % T.ceildiv(dstate, block_N)
T.annotate_layout({
x_shared: tilelang.layout.make_swizzled_layout(x_shared),
acc_o_shared: tilelang.layout.make_swizzled_layout(acc_o_shared)
})
T.annotate_layout(
{x_shared: tilelang.layout.make_swizzled_layout(x_shared), acc_o_shared: tilelang.layout.make_swizzled_layout(acc_o_shared)}
)
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)
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):
......@@ -123,47 +115,50 @@ def chunk_state_fwd(batch,
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)
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, acc_o_shared)
T.copy(
acc_o_shared,
Output[batch_idx, chunk_idx, bz, m_idx * block_M:(m_idx + 1) * block_M,
n_idx * block_N:(n_idx + 1) * block_N])
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
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--batch', type=int, default=8, help='batch size')
parser.add_argument('--heads', type=int, default=80, help='heads')
parser.add_argument('--groups', type=int, default=1, help='groups')
parser.add_argument('--seq_len', type=int, default=4096, help='sequence length')
parser.add_argument('--chunk_size', type=int, default=256, help='chunk size')
parser.add_argument('--dim', type=int, default=64, help='dim')
parser.add_argument('--dstate', type=int, default=128, help='dstate')
parser.add_argument('--tune', action='store_true', help='tune configs')
parser.add_argument("--batch", type=int, default=8, help="batch size")
parser.add_argument("--heads", type=int, default=80, help="heads")
parser.add_argument("--groups", type=int, default=1, help="groups")
parser.add_argument("--seq_len", type=int, default=4096, help="sequence length")
parser.add_argument("--chunk_size", type=int, default=256, help="chunk size")
parser.add_argument("--dim", type=int, default=64, help="dim")
parser.add_argument("--dstate", type=int, default=128, help="dstate")
parser.add_argument("--tune", action="store_true", help="tune configs")
args = parser.parse_args()
batch, heads, groups, seq_len, chunk_size, dim, dstate = args.batch, args.heads, args.groups, args.seq_len, args.chunk_size, args.dim, args.dstate
batch, heads, groups, seq_len, chunk_size, dim, dstate = (
args.batch,
args.heads,
args.groups,
args.seq_len,
args.chunk_size,
args.dim,
args.dstate,
)
total_flops = 2 * batch * seq_len * heads * dim * dstate
if (not args.tune):
if not args.tune:
kernel = chunk_state_fwd(
batch,
seq_len,
chunk_size,
groups,
heads,
dim,
dstate,
block_M=64,
block_N=128,
block_K=64,
num_stages=4,
threads=128)
batch, seq_len, chunk_size, groups, heads, dim, dstate, block_M=64, block_N=128, block_K=64, num_stages=4, threads=128
)
profiler = kernel.get_profiler(tilelang.TensorSupplyType.Normal)
profiler.assert_allclose(ref_program, rtol=0.01, atol=0.01)
print("All checks pass.")
......
examples/linear_attention/example_retention_fwd.py
View file @ 667632cc
......@@ -13,13 +13,12 @@ def chunk_retention_fwd_kernel(
H,
DK,
DV,
dtype: str = 'float16',
dtype: T.dtype = T.float16,
scale: float = None,
) -> torch.Tensor:
if scale is None:
scale = DK**-0.5
accum_dtype = 'float'
accum_dtype = T.float32
chunk_size = 64
BK = BV = 64 # Set to 128 can be faster, but has some numerical differences with FLA
......@@ -30,16 +29,16 @@ def chunk_retention_fwd_kernel(
@T.prim_func
def chunk_retention_fwd(
Q: T.Tensor([B, S, H, DK], dtype), # type: ignore
K: T.Tensor([B, S, H, DK], dtype), # type: ignore
V: T.Tensor([B, S, H, DV], dtype), # type: ignore
O: T.Tensor([NK, B, S, H, DV], dtype), # type: ignore
Q: T.Tensor([B, S, H, DK], dtype), # type: ignore
K: T.Tensor([B, S, H, DK], dtype), # type: ignore
V: T.Tensor([B, S, H, DV], dtype), # type: ignore
O: T.Tensor([NK, B, S, H, DV], dtype), # type: ignore
):
with T.Kernel(NV, NK, B * H) as (i_v, i_k, i_bh):
i_b = i_bh // H
i_h = i_bh % H
log_decay = T.alloc_var('float32')
log_decay = T.log2(1 - T.exp2(-5. - 1. * i_h)) # Head-specific log decay
log_decay = T.alloc_var(T.float32)
log_decay = T.log2(1 - T.exp2(-5.0 - 1.0 * i_h)) # Head-specific log decay
q = T.alloc_shared([chunk_size, BK], dtype)
k = T.alloc_shared([chunk_size, BK], dtype)
......@@ -56,14 +55,12 @@ def chunk_retention_fwd_kernel(
for i in T.Pipelined(0, NT):
for row, col in T.Parallel(chunk_size, BK):
q[row, col] = Q[i_b, i * chunk_size + row, i_h, i_k * BK + col] * scale
T.copy(K[i_b, i * chunk_size:(i + 1) * chunk_size, i_h, i_k * BK:(i_k + 1) * BK], k)
T.copy(V[i_b, i * chunk_size:(i + 1) * chunk_size, i_h, i_v * BV:(i_v + 1) * BV], v)
T.copy(K[i_b, i * chunk_size : (i + 1) * chunk_size, i_h, i_k * BK : (i_k + 1) * BK], k)
T.copy(V[i_b, i * chunk_size : (i + 1) * chunk_size, i_h, i_v * BV : (i_v + 1) * BV], v)
T.gemm(q, k, s, clear_accum=True, transpose_B=True)
for row, col in T.Parallel(chunk_size, chunk_size):
s_shared[row,
col] = T.if_then_else(row >= col, s[row, col] * T.exp2(
(row - col) * log_decay), 0)
s_shared[row, col] = T.if_then_else(row >= col, s[row, col] * T.exp2((row - col) * log_decay), 0)
T.copy(h, h_shared)
T.gemm(q, h_shared, o, clear_accum=True)
......@@ -75,9 +72,7 @@ def chunk_retention_fwd_kernel(
v[row, col] = v[row, col] * T.exp2((chunk_size - row - 1) * log_decay)
for row, col in T.Parallel(BK, BV):
h[row, col] = T.exp2(chunk_size * log_decay) * h[row, col]
T.copy(
o, O[i_k, i_b, i * chunk_size:(i + 1) * chunk_size, i_h,
i_v * BV:(i_v + 1) * BV])
T.copy(o, O[i_k, i_b, i * chunk_size : (i + 1) * chunk_size, i_h, i_v * BV : (i_v + 1) * BV])
T.gemm(k, v, h, transpose_A=True)
return chunk_retention_fwd
......@@ -89,24 +84,24 @@ def postprocess(o):
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--B', type=int, default=8, help='Batch size')
parser.add_argument('--S', type=int, default=4096, help='Seq len')
parser.add_argument('--H', type=int, default=32, help='Num heads')
parser.add_argument('--D', type=int, default=128, help='Head dim')
parser.add_argument("--B", type=int, default=8, help="Batch size")
parser.add_argument("--S", type=int, default=4096, help="Seq len")
parser.add_argument("--H", type=int, default=32, help="Num heads")
parser.add_argument("--D", type=int, default=128, help="Head dim")
args = parser.parse_args()
B, S, H, D = args.B, args.S, args.H, args.D
total_flops = 2.0 * B * S * S * H * D # causal
q = torch.randn((B, S, H, D), device='cuda', dtype=torch.float16)
k = torch.randn((B, S, H, D), device='cuda', dtype=torch.float16)
v = torch.randn((B, S, H, D), device='cuda', dtype=torch.float16)
q = torch.randn((B, S, H, D), device="cuda", dtype=torch.float16)
k = torch.randn((B, S, H, D), device="cuda", dtype=torch.float16)
v = torch.randn((B, S, H, D), device="cuda", dtype=torch.float16)
kernel = chunk_retention_fwd_kernel(B, S, H, D, D)
t = do_bench(lambda: postprocess(kernel(q, k, v)), warmup=25, rep=100)
print(f'Tilelang latency: {t:.3f} ms')
print(f'Tilelang TFLOPs: {total_flops/t * 1e-9}')
print(f"Tilelang latency: {t:.3f} ms")
print(f"Tilelang TFLOPs: {total_flops / t * 1e-9}")
if __name__ == '__main__':
if __name__ == "__main__":
main()
examples/minference/example_vertical_slash_sparse_attn.py
View file @ 667632cc
......@@ -15,12 +15,11 @@ from tilelang.profiler import do_bench
@tilelang.jit(out_idx=[3])
def _tl_vs_sparse_flashattn(batch, heads, seq_len, dim, vertical_size, slash_size):
block_M = 64
block_N = 64
num_stages = 2
threads = 128
scale = (1.0 / dim)**0.5 * 1.44269504
scale = (1.0 / dim) ** 0.5 * 1.44269504
shape = [batch, heads, seq_len, dim]
seq_blocks = (seq_len + block_M - 1) // block_M
......@@ -30,15 +29,13 @@ def _tl_vs_sparse_flashattn(batch, heads, seq_len, dim, vertical_size, slash_siz
offset_shape = count_shape + [slash_size]
index_shape = count_shape + [vertical_size]
vertical_size_round, slash_size_round = tilelang.next_power_of_2(
vertical_size), tilelang.next_power_of_2(slash_size)
vertical_size_round, slash_size_round = tilelang.next_power_of_2(vertical_size), tilelang.next_power_of_2(slash_size)
dtype = "float16"
accum_dtype = "float"
int_dtype = "int32"
dtype = T.float16
accum_dtype = T.float32
int_dtype = T.int32
def kernel_func(block_M, block_N, num_stages, threads):
@T.macro
def Prefetch(
K: T.Tensor(shape, dtype),
......@@ -53,32 +50,30 @@ def _tl_vs_sparse_flashattn(batch, heads, seq_len, dim, vertical_size, slash_siz
):
with T.attr("default", "async_scope", 1):
for i, j in T.Parallel(block_N, dim):
K_shared[i, j] = T.if_then_else(k + i < column_count,
K[bz, by, column_index[k + i], j], 0)
K_shared[i, j] = T.if_then_else(k + i < column_count, K[bz, by, column_index[k + i], j], 0)
with T.attr("default", "async_scope", 1):
for i, j in T.Parallel(block_N, dim):
V_shared[i, j] = T.if_then_else(k + i < column_count,
V[bz, by, column_index[k + i], j], 0)
V_shared[i, j] = T.if_then_else(k + i < column_count, V[bz, by, column_index[k + i], j], 0)
T.ptx_commit_group()
@T.macro
def Compute(
acc_s: T.FragmentBuffer([block_M, block_N], accum_dtype),
acc_s_cast: T.FragmentBuffer([block_M, block_N], dtype),
acc_o: T.FragmentBuffer([block_M, dim], accum_dtype),
scores_max: T.FragmentBuffer([block_M], accum_dtype),
scores_max_prev: T.FragmentBuffer([block_M], accum_dtype),
k: T.int32,
column_count: T.int32,
Q_shared: T.SharedBuffer([block_M, dim], dtype),
K_shared: T.SharedBuffer([block_N, dim], dtype),
V_shared: T.SharedBuffer([block_N, dim], dtype),
scores_scale: T.FragmentBuffer([block_M], accum_dtype),
scores_sum: T.FragmentBuffer([block_M], accum_dtype),
logsum: T.FragmentBuffer([block_M], accum_dtype),
count: T.int32,
acc_s: T.FragmentBuffer([block_M, block_N], accum_dtype),
acc_s_cast: T.FragmentBuffer([block_M, block_N], dtype),
acc_o: T.FragmentBuffer([block_M, dim], accum_dtype),
scores_max: T.FragmentBuffer([block_M], accum_dtype),
scores_max_prev: T.FragmentBuffer([block_M], accum_dtype),
k: T.int32,
column_count: T.int32,
Q_shared: T.SharedBuffer([block_M, dim], dtype),
K_shared: T.SharedBuffer([block_N, dim], dtype),
V_shared: T.SharedBuffer([block_N, dim], dtype),
scores_scale: T.FragmentBuffer([block_M], accum_dtype),
scores_sum: T.FragmentBuffer([block_M], accum_dtype),
logsum: T.FragmentBuffer([block_M], accum_dtype),
count: T.int32,
):
T.ptx_wait_group(count)
for i, j in T.Parallel(block_M, block_N):
......@@ -87,6 +82,8 @@ def _tl_vs_sparse_flashattn(batch, heads, seq_len, dim, vertical_size, slash_siz
T.copy(scores_max, scores_max_prev)
T.reduce_max(acc_s, scores_max, dim=1, clear=False)
for i in T.Parallel(block_M):
scores_max[i] = T.max(scores_max[i], scores_max_prev[i])
for i in T.Parallel(block_M):
scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale)
......@@ -106,17 +103,16 @@ def _tl_vs_sparse_flashattn(batch, heads, seq_len, dim, vertical_size, slash_siz
@T.prim_func
def vs_sparse_flashattn_ws(
Q: T.Tensor(shape, dtype),
K: T.Tensor(shape, dtype),
V: T.Tensor(shape, dtype),
Output: T.Tensor(shape, dtype),
BlockCount: T.Tensor(count_shape, int_dtype),
BlockOffset: T.Tensor(offset_shape, int_dtype),
ColumnCount: T.Tensor(count_shape, int_dtype),
ColumnIndex: T.Tensor(index_shape, int_dtype),
Q: T.Tensor(shape, dtype),
K: T.Tensor(shape, dtype),
V: T.Tensor(shape, dtype),
Output: T.Tensor(shape, dtype),
BlockCount: T.Tensor(count_shape, int_dtype),
BlockOffset: T.Tensor(offset_shape, int_dtype),
ColumnCount: T.Tensor(count_shape, int_dtype),
ColumnIndex: T.Tensor(index_shape, int_dtype),
):
with T.Kernel(T.ceildiv(seq_len, block_M), heads, batch, threads=256) as (bc, by, bz):
bx = T.ceildiv(seq_len, block_M) - 1 - bc
Q_shared = T.alloc_shared([block_M, dim], dtype)
K_shared = T.alloc_shared([2, block_N, dim], dtype)
......@@ -141,9 +137,11 @@ def _tl_vs_sparse_flashattn(batch, heads, seq_len, dim, vertical_size, slash_siz
T.create_list_of_mbarrier([128] * 9)
T.annotate_layout({
O_shared: tilelang.layout.make_swizzled_layout(O_shared),
})
T.annotate_layout(
{
O_shared: tilelang.layout.make_swizzled_layout(O_shared),
}
)
block_count[0] = BlockCount[bz, by, bx]
column_count[0] = ColumnCount[bz, by, bx]
......@@ -160,15 +158,15 @@ def _tl_vs_sparse_flashattn(batch, heads, seq_len, dim, vertical_size, slash_siz
if tid >= 128:
T.annotate_producer_reg_dealloc()
T.copy(Q[bz, by, bx * block_M:(bx + 1) * block_M, :], Q_shared)
T.copy(Q[bz, by, bx * block_M : (bx + 1) * block_M, :], Q_shared)
T.mbarrier_arrive(mbarrier=8)
for bi in T.serial(block_count[0]):
k = block_offset[bi]
T.mbarrier_wait_parity(mbarrier=bi % 2 + 4, parity=(((bi & 3) >> 1) ^ 1))
T.copy(K[bz, by, k:k + block_N, :], K_shared[bi % 2, :, :])
T.copy(K[bz, by, k : k + block_N, :], K_shared[bi % 2, :, :])
T.mbarrier_arrive(mbarrier=bi % 2)
T.mbarrier_wait_parity(mbarrier=bi % 2 + 6, parity=(((bi & 3) >> 1) ^ 1))
T.copy(V[bz, by, k:k + block_N, :], V_shared[bi % 2, :, :])
T.copy(V[bz, by, k : k + block_N, :], V_shared[bi % 2, :, :])
T.mbarrier_arrive(mbarrier=bi % 2 + 2)
else:
T.annotate_consumer_reg_alloc()
......@@ -179,37 +177,28 @@ def _tl_vs_sparse_flashattn(batch, heads, seq_len, dim, vertical_size, slash_siz
for bi in T.serial(block_count[0]):
k = block_offset[bi]
for i, j in T.Parallel(block_M, block_N):
acc_s[i, j] = T.if_then_else(bx * block_M + i >= k + j, 0,
-T.infinity(acc_s.dtype))
acc_s[i, j] = T.if_then_else(bx * block_M + i >= k + j, 0, -T.infinity(acc_s.dtype))
T.mbarrier_wait_parity(mbarrier=bi % 2, parity=((bi & 3) >> 1))
T.gemm(
Q_shared,
K_shared[bi % 2, :, :],
acc_s,
transpose_B=True,
policy=T.GemmWarpPolicy.FullRow)
T.gemm(Q_shared, K_shared[bi % 2, :, :], acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
T.mbarrier_arrive(mbarrier=bi % 2 + 4)
T.copy(scores_max, scores_max_prev)
T.reduce_max(acc_s, scores_max, dim=1, clear=False)
for i in T.Parallel(block_M):
scores_max[i] = T.max(scores_max[i], scores_max_prev[i])
for i in T.Parallel(block_M):
scores_scale[i] = T.exp2(scores_max_prev[i] * scale -
scores_max[i] * scale)
scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale)
for i, j in T.Parallel(block_M, block_N):
acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale)
for i, j in T.Parallel(block_M, dim):
acc_o[i, j] = acc_o[i, j] * scores_scale[i]
T.copy(acc_s, acc_s_cast)
T.mbarrier_wait_parity(mbarrier=bi % 2 + 2, parity=(((bi & 3) >> 1)))
T.gemm(
acc_s_cast,
V_shared[bi % 2, :, :],
acc_o,
policy=T.GemmWarpPolicy.FullRow)
T.mbarrier_wait_parity(mbarrier=bi % 2 + 2, parity=((bi & 3) >> 1))
T.gemm(acc_s_cast, V_shared[bi % 2, :, :], acc_o, policy=T.GemmWarpPolicy.FullRow)
T.mbarrier_arrive(mbarrier=bi % 2 + 6)
......@@ -219,38 +208,85 @@ def _tl_vs_sparse_flashattn(batch, heads, seq_len, dim, vertical_size, slash_siz
logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i]
if column_count[0] != 0:
Prefetch(K, V, K_shared_1, V_shared_1, column_index, column_count[0], 0, bz,
by)
Prefetch(K, V, K_shared_1, V_shared_1, column_index, column_count[0], 0, bz, by)
for bi in T.serial(T.ceildiv(column_count[0], block_N) - 1):
k = bi * block_N
if bi % 2 == 0:
Prefetch(K, V, K_shared_2, V_shared_2, column_index,
column_count[0], k + block_N, bz, by)
Compute(acc_s, acc_s_cast, acc_o, scores_max, scores_max_prev, k,
column_count[0], Q_shared, K_shared_1, V_shared_1,
scores_scale, scores_sum, logsum, 1)
Prefetch(K, V, K_shared_2, V_shared_2, column_index, column_count[0], k + block_N, bz, by)
Compute(
acc_s,
acc_s_cast,
acc_o,
scores_max,
scores_max_prev,
k,
column_count[0],
Q_shared,
K_shared_1,
V_shared_1,
scores_scale,
scores_sum,
logsum,
1,
)
else:
Prefetch(K, V, K_shared_1, V_shared_1, column_index,
column_count[0], k + block_N, bz, by)
Compute(acc_s, acc_s_cast, acc_o, scores_max, scores_max_prev, k,
column_count[0], Q_shared, K_shared_2, V_shared_2,
scores_scale, scores_sum, logsum, 1)
Prefetch(K, V, K_shared_1, V_shared_1, column_index, column_count[0], k + block_N, bz, by)
Compute(
acc_s,
acc_s_cast,
acc_o,
scores_max,
scores_max_prev,
k,
column_count[0],
Q_shared,
K_shared_2,
V_shared_2,
scores_scale,
scores_sum,
logsum,
1,
)
if T.ceildiv(column_count[0], block_N) % 2 == 0:
Compute(acc_s, acc_s_cast, acc_o, scores_max, scores_max_prev,
T.ceildiv(column_count[0], block_N) * block_N - block_N,
column_count[0], Q_shared, K_shared_2, V_shared_2, scores_scale,
scores_sum, logsum, 0)
Compute(
acc_s,
acc_s_cast,
acc_o,
scores_max,
scores_max_prev,
T.ceildiv(column_count[0], block_N) * block_N - block_N,
column_count[0],
Q_shared,
K_shared_2,
V_shared_2,
scores_scale,
scores_sum,
logsum,
0,
)
else:
Compute(acc_s, acc_s_cast, acc_o, scores_max, scores_max_prev,
T.ceildiv(column_count[0], block_N) * block_N - block_N,
column_count[0], Q_shared, K_shared_1, V_shared_1, scores_scale,
scores_sum, logsum, 0)
Compute(
acc_s,
acc_s_cast,
acc_o,
scores_max,
scores_max_prev,
T.ceildiv(column_count[0], block_N) * block_N - block_N,
column_count[0],
Q_shared,
K_shared_1,
V_shared_1,
scores_scale,
scores_sum,
logsum,
0,
)
for i, j in T.Parallel(block_M, dim):
acc_o[i, j] /= logsum[i]
T.copy(acc_o, O_shared)
T.copy(O_shared, Output[bz, by, bx * block_M:(bx + 1) * block_M, :])
T.copy(O_shared, Output[bz, by, bx * block_M : (bx + 1) * block_M, :])
return vs_sparse_flashattn_ws
......@@ -466,11 +502,8 @@ def vertical_slash_sparse_attention(
import os
current_dir = os.path.dirname(os.path.abspath(__file__))
sources = [
os.path.join(current_dir, 'ops', 'kernels.cpp'),
os.path.join(current_dir, 'ops', 'vertical_slash_index.cu')
]
ops = load(name='convert', sources=sources, verbose=False)
sources = [os.path.join(current_dir, "ops", "kernels.cpp"), os.path.join(current_dir, "ops", "vertical_slash_index.cu")]
ops = load(name="convert", sources=sources, verbose=False)
convert_vertical_slash_indexes = ops.convert_vertical_slash_indexes
batch_size, num_heads, context_size, head_dim = query.shape
pad = (block_size_M - context_size) & (block_size_M - 1)
......@@ -481,15 +514,13 @@ def vertical_slash_sparse_attention(
value = torch.nn.functional.pad(value, [0, 0, 0, pad, 0, 0, 0, 0])
if head_dim not in [16, 32, 64, 128, 256, 512]:
target_dim = 2**math.ceil(math.log2(head_dim)) - head_dim
target_dim = 2 ** math.ceil(math.log2(head_dim)) - head_dim
query = torch.nn.functional.pad(query, [0, target_dim, 0, 0, 0, 0, 0, 0])
key = torch.nn.functional.pad(key, [0, target_dim, 0, 0, 0, 0, 0, 0])
value = torch.nn.functional.pad(value, [0, target_dim, 0, 0, 0, 0, 0, 0])
v_idx = v_idx.to(torch.int32).reshape((batch_size, num_heads, -1)).sort(
dim=-1, descending=False)[0]
s_idx = s_idx.to(torch.int32).reshape((batch_size, num_heads, -1)).sort(
dim=-1, descending=True)[0]
v_idx = v_idx.to(torch.int32).reshape((batch_size, num_heads, -1)).sort(dim=-1, descending=False)[0]
s_idx = s_idx.to(torch.int32).reshape((batch_size, num_heads, -1)).sort(dim=-1, descending=True)[0]
seqlens = torch.tensor([context_size] * query.shape[0], dtype=torch.int32, device=query.device)
sm_scale = head_dim**-0.5
......@@ -502,8 +533,7 @@ def vertical_slash_sparse_attention(
block_size_N,
)
tl_kernel = _tl_vs_sparse_flashattn(batch_size, num_heads, context_size, head_dim,
v_idx.shape[2], s_idx.shape[2])
tl_kernel = _tl_vs_sparse_flashattn(batch_size, num_heads, context_size, head_dim, v_idx.shape[2], s_idx.shape[2])
def run(is_triton: bool = True):
if is_triton:
......@@ -521,8 +551,7 @@ def vertical_slash_sparse_attention(
block_size_N,
)
else:
out = tl_kernel(query, key, value, block_count, block_offset, column_count,
column_index)
out = tl_kernel(query, key, value, block_count, block_offset, column_count, column_index)
return out[..., :context_size, :head_dim]
return run
......@@ -532,8 +561,7 @@ def sum_all_diagonal_matrix(mat: torch.tensor):
b, h, n, m = mat.shape
zero_mat = torch.zeros((b, h, n, n)).to(mat.device) # Zero matrix used for padding
mat_padded = torch.cat((zero_mat, mat, zero_mat), -1) # pads the matrix on left and right
mat_strided = mat_padded.as_strided(
(1, 1, n, n + m), (1, n * (2 * n + m), 2 * n + m + 1, 1)) # Change the strides
mat_strided = mat_padded.as_strided((1, 1, n, n + m), (1, n * (2 * n + m), 2 * n + m + 1, 1)) # Change the strides
sum_diags = torch.sum(mat_strided, 2) # Sums the resulting matrix's columns
return sum_diags[:, :, 1:]
......@@ -555,24 +583,23 @@ def main(argv=None):
vertical_size, slash_size = args.vertical_size, args.slash_size
torch.manual_seed(0)
q = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device='cuda', dtype=torch.float16)
k = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device='cuda', dtype=torch.float16)
v = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device='cuda', dtype=torch.float16)
q = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device="cuda", dtype=torch.float16)
k = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device="cuda", dtype=torch.float16)
v = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device="cuda", dtype=torch.float16)
q_len = SEQ_LEN
vertical_size, slash_size = min(q_len, vertical_size), min(q_len, slash_size)
last_q = 64
qk = torch.einsum('bhmk, bhnk -> bhmn', q[:, :, -last_q:, :], k)
qk = torch.einsum("bhmk, bhnk -> bhmn", q[:, :, -last_q:, :], k)
arange = torch.arange(last_q, device="cuda")
qk[:, :, :, -last_q:] = torch.where(arange[None, None, :, None] >= arange[None, None, None, :],
qk[:, :, :, -last_q:], -torch.inf)
qk[:, :, :, -last_q:] = torch.where(arange[None, None, :, None] >= arange[None, None, None, :], qk[:, :, :, -last_q:], -torch.inf)
qk = torch.nn.functional.softmax(qk, dim=-1, dtype=torch.float32)
vertical = qk.sum(-2, keepdim=True)
vertical[..., :30] = torch.inf
vertical_topk = torch.topk(vertical, vertical_size, -1).indices
slash = sum_all_diagonal_matrix(qk)[..., :-last_q + 1]
slash = sum_all_diagonal_matrix(qk)[..., : -last_q + 1]
slash[..., -30:] = torch.inf
slash = (q_len - 1) - torch.topk(slash, slash_size, -1).indices
......
examples/norm/rms_norm.py
View file @ 667632cc
......@@ -4,7 +4,7 @@ import tilelang.language as T
def rms_norm_splitk(M, N, blk_m, blk_k):
dtype = "float"
dtype = T.float
@T.prim_func
def main(A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype)):
......@@ -35,7 +35,7 @@ def rms_norm_splitk(M, N, blk_m, blk_k):
@tilelang.jit(out_idx=[-1], pass_configs={"tl.disable_tma_lower": True})
def rms_norm(M, N, blk_m):
dtype = "float"
dtype = T.float
@T.prim_func
def main(A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype)):
......@@ -45,7 +45,7 @@ def rms_norm(M, N, blk_m):
A_local = T.alloc_fragment((blk_m, N), dtype)
A_powsum = T.alloc_fragment((blk_m,), dtype)
T.copy(A[bx * blk_m:(bx + 1) * blk_m, :], A_shared)
T.copy(A[bx * blk_m : (bx + 1) * blk_m, :], A_shared)
T.copy(A_shared, A_local)
for i, j in T.Parallel(blk_m, N):
A_pow_local[i, j] = A_local[i, j] * A_local[i, j]
......@@ -54,7 +54,7 @@ def rms_norm(M, N, blk_m):
A_powsum[i] = T.rsqrt(A_powsum[i] / N + 1e-12)
for i, j in T.Parallel(blk_m, N):
A_local[i, j] *= A_powsum[i]
T.copy(A_local, B[bx * blk_m:(bx + 1) * blk_m, :])
T.copy(A_local, B[bx * blk_m : (bx + 1) * blk_m, :])
return main
......
examples/norm/test_rms_norm.py
View file @ 667632cc
......@@ -5,7 +5,7 @@ import tilelang.language as T
def rms_norm_splitk(M, N, blk_m, blk_k):
dtype = "float"
dtype = T.float
@T.prim_func
def main(A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype)):
......@@ -35,7 +35,7 @@ def rms_norm_splitk(M, N, blk_m, blk_k):
def rms_norm(M, N, blk_m):
dtype = "float"
dtype = T.float
@T.prim_func
def main(A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype)):
......@@ -45,7 +45,7 @@ def rms_norm(M, N, blk_m):
A_local = T.alloc_fragment((blk_m, N), dtype)
A_powsum = T.alloc_fragment((blk_m,), dtype)
T.copy(A[bx * blk_m:(bx + 1) * blk_m, :], A_shared)
T.copy(A[bx * blk_m : (bx + 1) * blk_m, :], A_shared)
T.copy(A_shared, A_local)
for i, j in T.Parallel(blk_m, N):
A_pow_local[i, j] = A_local[i, j] * A_local[i, j]
......@@ -54,7 +54,7 @@ def rms_norm(M, N, blk_m):
A_powsum[i] = T.rsqrt(A_powsum[i] / N + 1e-12)
for i, j in T.Parallel(blk_m, N):
A_local[i, j] *= A_powsum[i]
T.copy(A_local, B[bx * blk_m:(bx + 1) * blk_m, :])
T.copy(A_local, B[bx * blk_m : (bx + 1) * blk_m, :])
return main
......
examples/online_softmax/online_softmax.py
View file @ 667632cc
......@@ -9,19 +9,19 @@ from typing import Callable
def softmax_kernel(
M,
N,
dtype: str = "float16",
dtype: T.dtype = T.float16,
) -> "Callable":
BN = min(tl.next_power_of_2(N), 8192)
NN = tl.cdiv(N, BN)
accum_dtype = "float"
accum_dtype = T.float32
scale = 1.44269504 # log2(e)
@T.prim_func
def main(
X: T.Tensor([M, N], dtype),
Y: T.Tensor([M, N], dtype),
X: T.Tensor([M, N], dtype),
Y: T.Tensor([M, N], dtype),
):
with T.Kernel(M, threads=128) as (i_m):
x = T.alloc_fragment([BN], dtype)
......@@ -33,7 +33,7 @@ def softmax_kernel(
T.fill(lse, -T.infinity(accum_dtype))
for i_n in T.Pipelined(0, NN):
T.copy(X[i_m, i_n * BN:(i_n + 1) * BN], x)
T.copy(X[i_m, i_n * BN : (i_n + 1) * BN], x)
T.reduce_max(x, max_x, dim=0, clear=True)
......@@ -45,12 +45,12 @@ def softmax_kernel(
lse[0] = max_x[0] * scale + T.log2(T.exp2(lse[0] - max_x[0] * scale) + sum_exp_x[0])
for i_n in T.Pipelined(0, NN):
T.copy(X[i_m, i_n * BN:(i_n + 1) * BN], x)
T.copy(X[i_m, i_n * BN : (i_n + 1) * BN], x)
for j in T.Parallel(BN):
y[j] = T.exp2(x[j] * scale - lse[0])
T.copy(y, Y[i_m, i_n * BN:(i_n + 1) * BN])
T.copy(y, Y[i_m, i_n * BN : (i_n + 1) * BN])
return main
......@@ -69,4 +69,4 @@ t1 = do_bench(lambda: X.softmax(dim=1), warmup=25, rep=100)
t2 = do_bench(lambda: kernel(X), warmup=25, rep=100)
print(f"torch latency: {t1:.3f} ms")
print(f"TileLang latency: {t2:.3f} ms")
print(f"Speedup: {t1/t2:.3f}x")
print(f"Speedup: {t1 / t2:.3f}x")
examples/plot_layout/README.md
View file @ 667632cc
......@@ -10,7 +10,7 @@ from typing import Literal, Callable
from tilelang.intrinsics.utils import get_mma_micro_size
from tilelang.tools import plot_layout
def make_mma_load_base_layout(dtype: str = "float16",
def make_mma_load_base_layout(dtype: str = T.float16,
matrix: Literal["A", "B"] = "A",
transposed: bool = False) -> T.Fragment:
"""
......@@ -69,7 +69,7 @@ def make_mma_load_base_layout(dtype: str = "float16",
micro_size_s, _, micro_size_r = get_mma_micro_size(dtype)
transform_func = transform_func
inverse_mma_load_layout = IndexMap.from_func(transform_func, index_dtype="int32")
inverse_mma_load_layout = IndexMap.from_func(transform_func, index_dtype=T.int32)
def forward_thread(i: int, j: int) -> int:
"""
......@@ -94,7 +94,7 @@ def make_mma_load_base_layout(dtype: str = "float16",
# Create a 16×16 matrix layout for ldmatrix operations
base_layout = make_mma_load_base_layout(dtype="float16", matrix="A", transposed=False)
base_layout = make_mma_load_base_layout(dtype=T.float16, matrix="A", transposed=False)
# Print the layout structure (optional for debugging)
print(base_layout)
......
examples/plot_layout/fragment_mfma_load_a.py
View file @ 667632cc
......@@ -11,10 +11,9 @@ from tilelang.intrinsics.mfma_layout import (
)
def make_mfma_load_base_layout(dtype: str = "float16",
matrix: Literal["A", "B"] = "A",
k_dim: int = 16,
transposed: bool = False) -> T.Fragment:
def make_mfma_load_base_layout(
dtype: T.dtype = T.float16, matrix: Literal["A", "B"] = "A", k_dim: int = 16, transposed: bool = False
) -> T.Fragment:
"""
Create a layout function for storing MFMA results into a fragment buffer.
This layout is used in conjunction with `inverse_mfma_store_layout` to
......@@ -72,17 +71,15 @@ def make_mfma_load_base_layout(dtype: str = "float16",
# so the b matrix expected a transposed basic layout
transform_func: Callable = None
if matrix == "A":
transform_func = transform_func_sr_a if is_sr_axis_order else lambda i, j: transform_func_sr_a(
j, i)
transform_func = transform_func_sr_a if is_sr_axis_order else lambda i, j: transform_func_sr_a(j, i)
micro_size_s, micro_size_r = micro_size_x, micro_size_k
elif matrix == "B":
transform_func = transform_func_sr_b if is_sr_axis_order else lambda i, j: transform_func_sr_b(
j, i)
transform_func = transform_func_sr_b if is_sr_axis_order else lambda i, j: transform_func_sr_b(j, i)
micro_size_s, micro_size_r = micro_size_k, micro_size_y
else:
raise ValueError(f"Unsupported matrix {matrix}")
inverse_mma_load_layout = IndexMap.from_func(transform_func, index_dtype="int32")
inverse_mma_load_layout = IndexMap.from_func(transform_func, index_dtype=T.int32)
def forward_thread(i: int, j: int) -> int:
"""
......@@ -115,19 +112,16 @@ chunk = 2
from tilelang.tools import plot_layout
# ldmatrix layout 16x16
base_layout = make_mfma_load_base_layout(dtype="float16", matrix="A", transposed=False)
base_layout = make_mfma_load_base_layout(dtype=T.float16, matrix="A", transposed=False)
print(base_layout)
plot_layout(base_layout, name="base_layout")
# warp layout 32x32
warp_layout = base_layout.repeat([warp_rows, warp_cols],
repeat_on_thread=False,
lower_dim_first=False)
warp_layout = base_layout.repeat([warp_rows, warp_cols], repeat_on_thread=False, lower_dim_first=False)
print(warp_layout)
plot_layout(warp_layout, name="warp_layout")
# block layout 64x32
block_layout = warp_layout.repeat([block_rows, 1], repeat_on_thread=True,
lower_dim_first=True).replicate(block_cols)
block_layout = warp_layout.repeat([block_rows, 1], repeat_on_thread=True, lower_dim_first=True).replicate(block_cols)
print(block_layout)
plot_layout(block_layout, name="block_layout")
examples/plot_layout/fragment_mma_load_a.py
View file @ 667632cc
......@@ -5,9 +5,7 @@ from tvm.tir import IndexMap
from tilelang.intrinsics.utils import get_mma_micro_size
def make_mma_load_base_layout(dtype: str = "float16",
matrix: Literal["A", "B"] = "A",
transposed: bool = False) -> T.Fragment:
def make_mma_load_base_layout(dtype: T.dtype = T.float16, matrix: Literal["A", "B"] = "A", transposed: bool = False) -> T.Fragment:
"""
Create a layout function for storing MMA results into a fragment buffer.
This layout is used in conjunction with `inverse_mma_store_layout` to
......@@ -36,6 +34,7 @@ def make_mma_load_base_layout(dtype: str = "float16",
shared_16x16_to_mma_32x8_layout_sr_b,
shared_16x32_to_mma_32x16_layout_sr_b,
)
assert matrix in ["A", "B"], "matrix should be either A or B"
dtype_bits = DataType(dtype).bits
# s represents spatial axis
......@@ -67,17 +66,15 @@ def make_mma_load_base_layout(dtype: str = "float16",
# so the b matrix expected a transposed basic layout
transform_func: Callable = None
if matrix == "A":
transform_func = transform_func_sr_a if is_sr_axis_order else lambda i, j: transform_func_sr_a(
j, i)
transform_func = transform_func_sr_a if is_sr_axis_order else lambda i, j: transform_func_sr_a(j, i)
micro_size_s, micro_size_r = micro_size_x, micro_size_k
elif matrix == "B":
transform_func = transform_func_sr_b if is_sr_axis_order else lambda i, j: transform_func_sr_b(
j, i)
transform_func = transform_func_sr_b if is_sr_axis_order else lambda i, j: transform_func_sr_b(j, i)
micro_size_s, micro_size_r = micro_size_k, micro_size_y
else:
raise ValueError(f"Unsupported matrix {matrix}")
inverse_mma_load_layout = IndexMap.from_func(transform_func, index_dtype="int32")
inverse_mma_load_layout = IndexMap.from_func(transform_func, index_dtype=T.int32)
def forward_thread(i: int, j: int) -> int:
"""
......@@ -110,7 +107,7 @@ chunk = 2
from tilelang.tools import plot_layout
# ldmatrix layout 16x16
base_layout = make_mma_load_base_layout(dtype="float16", matrix="A", transposed=False)
base_layout = make_mma_load_base_layout(dtype=T.float16, matrix="A", transposed=False)
print(base_layout)
plot_layout(base_layout, name="base_layout")
......
examples/quickstart.py
View file @ 667632cc
......@@ -6,13 +6,12 @@ import tilelang.language as T
# target currently can be "cuda" or "hip" or "cpu".
# if not specified, it will be inferred from the input tensors during compile time
@tilelang.jit
def matmul(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float"):
def matmul(M, N, K, block_M, block_N, block_K, dtype=T.float16, accum_dtype=T.float32):
@T.prim_func
def matmul_relu_kernel(
A: T.Tensor((M, K), dtype),
B: T.Tensor((K, N), dtype),
C: T.Tensor((M, N), dtype),
A: T.Tensor((M, K), dtype),
B: T.Tensor((K, N), dtype),
C: T.Tensor((M, N), dtype),
):
# Initialize Kernel Context
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
......@@ -55,10 +54,9 @@ block_M = 128
block_N = 128
block_K = 32
# 1. Define the kernel (matmul) and compile/lower it into an executable module
# Define the kernel (matmul) and compile/lower it into an executable module
matmul_relu_kernel = matmul(M, N, K, block_M, block_N, block_K)
# 3. Test the kernel in Python with PyTorch data
# Test the kernel in Python with PyTorch data
import torch
# Create random input tensors on the GPU
......@@ -78,7 +76,7 @@ torch.testing.assert_close(c, ref_c, rtol=1e-2, atol=1e-2)
print("Kernel output matches PyTorch reference.")
# 4. Retrieve and inspect the generated CUDA source (optional)
# cuda_source = jit_kernel.get_kernel_source()
# cuda_source = matmul_relu_kernel.get_kernel_source()
# print("Generated CUDA kernel:\n", cuda_source)
# 5.Profile latency with kernel
......
examples/rand/rand_uint.py 0 → 100644
View file @ 667632cc
import tilelang
import tilelang.language as T
import torch
import triton
import triton.language as tl
@tilelang.jit
def tilelang_rand_1d(M=1024, seed=42):
num_per_thread = 128
threads = 1
blk_M = num_per_thread * threads
@T.prim_func
def rand_kernel(A: T.Tensor((M,), "uint32")):
with T.Kernel(T.ceildiv(M, threads * num_per_thread), threads=threads) as bx:
tx = T.get_thread_binding()
T.rng_init(seed, 0, bx * blk_M + tx * num_per_thread)
for i, j in T.Parallel(threads, num_per_thread):
offsets = (bx * threads + i) * num_per_thread
idx = offsets + j
if idx < M:
A[idx] = T.rng_rand()
return rand_kernel
@triton.jit
def triton_rand_1d(X, M, elements_per_thread, seed):
pid = tl.program_id(0)
offset = pid * elements_per_thread + tl.arange(0, elements_per_thread)
r0, r1, r2, r3 = tl.randint4x(seed, offset)
base_idx = offset * 4
tl.store(X + base_idx, r0, mask=base_idx < M)
tl.store(X + base_idx + 1, r1, mask=(base_idx + 1) < M)
tl.store(X + base_idx + 2, r2, mask=(base_idx + 2) < M)
tl.store(X + base_idx + 3, r3, mask=(base_idx + 3) < M)
def test_rand_1d(M, seed):
kernel = tilelang_rand_1d(M, seed)
tilelang_result = torch.empty(M, dtype=torch.uint32, device="cuda")
kernel(tilelang_result)
triton_result = torch.empty(M, dtype=torch.uint32, device="cuda")
grid = (triton.cdiv(M, 128),)
triton_rand_1d[grid](triton_result, tl.constexpr(M), tl.constexpr(128 // 4), seed)
torch.testing.assert_close(tilelang_result, triton_result)
if __name__ == "__main__":
test_rand_1d(1024, 42)
test_rand_1d(512, 123)
test_rand_1d(128, 0)
examples/seer_attention/block_sparse_attn_tilelang.py
View file @ 667632cc
......@@ -10,10 +10,7 @@ def get_sparse_attn_mask_from_topk(x, topk, use_dense_for_last_block=False):
bsz, num_head, downsample_len, _ = x.shape
# N_CTX = downsample_len * BLOCK
sparse_index = torch.topk(x, topk, dim=-1).indices
dense_mask = torch.full([bsz, num_head, downsample_len, downsample_len],
False,
dtype=torch.bool,
device=x.device)
dense_mask = torch.full([bsz, num_head, downsample_len, downsample_len], False, dtype=torch.bool, device=x.device)
dense_mask.scatter_(-1, sparse_index, True)
if use_dense_for_last_block:
dense_mask[:, :, -2:, :] = True
......@@ -30,38 +27,41 @@ def get_sparse_attn_mask_from_threshold(x, threshold, use_dense_for_last_block=F
@tilelang.jit(
out_idx=[4], pass_configs={
out_idx=[4],
pass_configs={
tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
})
},
)
def blocksparse_flashattn(batch, heads, seq_q, seq_kv, dim, downsample_len, is_causal):
block_M = 64
block_N = 64
num_stages = 0
threads = 128
scale = (1.0 / dim)**0.5 * 1.44269504 # log2(e)
scale = (1.0 / dim) ** 0.5 * 1.44269504 # log2(e)
q_shape = [batch, heads, seq_q, dim]
kv_shape = [batch, heads, seq_kv, dim]
block_mask_shape = [batch, heads, downsample_len, downsample_len]
dtype = "float16"
accum_dtype = "float"
block_mask_dtype = "int8"
dtype = T.float16
accum_dtype = T.float32
block_mask_dtype = T.int8
def kernel_func(block_M, block_N, num_stages, threads):
@T.macro
def Softmax(
acc_s: T.FragmentBuffer([block_M, block_N], accum_dtype),
acc_s_cast: T.FragmentBuffer([block_M, block_N], dtype),
scores_max: T.FragmentBuffer([block_M], accum_dtype),
scores_max_prev: T.FragmentBuffer([block_M], accum_dtype),
scores_scale: T.FragmentBuffer([block_M], accum_dtype),
scores_sum: T.FragmentBuffer([block_M], accum_dtype),
logsum: T.FragmentBuffer([block_M], accum_dtype),
acc_s: T.FragmentBuffer([block_M, block_N], accum_dtype),
acc_s_cast: T.FragmentBuffer([block_M, block_N], dtype),
scores_max: T.FragmentBuffer([block_M], accum_dtype),
scores_max_prev: T.FragmentBuffer([block_M], accum_dtype),
scores_scale: T.FragmentBuffer([block_M], accum_dtype),
scores_sum: T.FragmentBuffer([block_M], accum_dtype),
logsum: T.FragmentBuffer([block_M], accum_dtype),
):
T.copy(scores_max, scores_max_prev)
T.fill(scores_max, -T.infinity(accum_dtype))
T.reduce_max(acc_s, scores_max, dim=1, clear=False)
for i in T.Parallel(block_M):
scores_max[i] = T.max(scores_max[i], scores_max_prev[i])
# To do causal softmax, we need to set the scores_max to 0 if it is -inf
# This process is called Check_inf in FlashAttention3 code, and it only need to be done
# in the first ceil_div(kBlockM, kBlockN) steps.
......@@ -81,19 +81,19 @@ def blocksparse_flashattn(batch, heads, seq_q, seq_kv, dim, downsample_len, is_c
@T.macro
def Rescale(
acc_o: T.FragmentBuffer([block_M, dim], accum_dtype),
scores_scale: T.FragmentBuffer([block_M], accum_dtype),
acc_o: T.FragmentBuffer([block_M, dim], accum_dtype),
scores_scale: T.FragmentBuffer([block_M], accum_dtype),
):
for i, j in T.Parallel(block_M, dim):
acc_o[i, j] *= scores_scale[i]
@T.prim_func
def main(
Q: T.Tensor(q_shape, dtype),
K: T.Tensor(kv_shape, dtype),
V: T.Tensor(kv_shape, dtype),
BlockSparseMask: T.Tensor(block_mask_shape, block_mask_dtype),
Output: T.Tensor(q_shape, dtype),
Q: T.Tensor(q_shape, dtype),
K: T.Tensor(kv_shape, dtype),
V: T.Tensor(kv_shape, dtype),
BlockSparseMask: T.Tensor(block_mask_shape, block_mask_dtype),
Output: T.Tensor(q_shape, dtype),
):
with T.Kernel(T.ceildiv(seq_q, block_M), heads, batch, threads=threads) as (bx, by, bz):
Q_shared = T.alloc_shared([block_M, dim], dtype)
......@@ -110,7 +110,7 @@ def blocksparse_flashattn(batch, heads, seq_q, seq_kv, dim, downsample_len, is_c
logsum = T.alloc_fragment([block_M], accum_dtype)
block_mask = T.alloc_local([downsample_len], block_mask_dtype)
T.copy(Q[bz, by, bx * block_M:(bx + 1) * block_M, :], Q_shared)
T.copy(Q[bz, by, bx * block_M : (bx + 1) * block_M, :], Q_shared)
T.fill(acc_o, 0)
T.fill(logsum, 0)
T.fill(scores_max, -T.infinity(accum_dtype))
......@@ -122,33 +122,25 @@ def blocksparse_flashattn(batch, heads, seq_q, seq_kv, dim, downsample_len, is_c
for k in T.Pipelined(loop_range, num_stages=num_stages):
if block_mask[k] != 0:
T.copy(K[bz, by, k * block_N:(k + 1) * block_N, :], K_shared)
T.copy(K[bz, by, k * block_N : (k + 1) * block_N, :], K_shared)
if is_causal:
past_len = seq_kv - seq_q
for i, j in T.Parallel(block_M, block_N):
acc_s[i, j] = T.if_then_else(
bx * block_M + i + past_len >= k * block_N + j, 0,
-T.infinity(acc_s.dtype))
acc_s[i, j] = T.if_then_else(bx * block_M + i + past_len >= k * block_N + j, 0, -T.infinity(acc_s.dtype))
else:
T.clear(acc_s)
T.gemm(
Q_shared,
K_shared,
acc_s,
transpose_B=True,
policy=T.GemmWarpPolicy.FullRow)
Softmax(acc_s, acc_s_cast, scores_max, scores_max_prev, scores_scale,
scores_sum, logsum)
T.gemm(Q_shared, K_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow)
Softmax(acc_s, acc_s_cast, scores_max, scores_max_prev, scores_scale, scores_sum, logsum)
Rescale(acc_o, scores_scale)
T.copy(V[bz, by, k * block_N:(k + 1) * block_N, :], V_shared)
T.copy(V[bz, by, k * block_N : (k + 1) * block_N, :], V_shared)
T.gemm(acc_s_cast, V_shared, acc_o, policy=T.GemmWarpPolicy.FullRow)
for i, j in T.Parallel(block_M, dim):
acc_o[i, j] /= logsum[i]
T.copy(acc_o, O_shared)
T.copy(O_shared, Output[bz, by, bx * block_M:(bx + 1) * block_M, :])
T.copy(O_shared, Output[bz, by, bx * block_M : (bx + 1) * block_M, :])
return main
......@@ -163,44 +155,40 @@ def test_topk_sparse_attention():
torch.manual_seed(0)
# Create inputs
q = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device='cuda', dtype=torch.float16)
k = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device='cuda', dtype=torch.float16)
v = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device='cuda', dtype=torch.float16)
q = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device="cuda", dtype=torch.float16)
k = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device="cuda", dtype=torch.float16)
v = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device="cuda", dtype=torch.float16)
sm_scale = 1.0 / (D_HEAD**0.5)
# Create sparse mask (downsampled to block level)
downsample_factor = BLOCK
downsample_len = math.ceil(SEQ_LEN / downsample_factor)
x_ds = torch.randn([BATCH, N_HEADS, downsample_len, downsample_len],
device='cuda',
dtype=torch.float16)
x_ds = torch.randn([BATCH, N_HEADS, downsample_len, downsample_len], device="cuda", dtype=torch.float16)
x_ds[:, :, :, 0] = 100
block_mask = get_sparse_attn_mask_from_topk(x_ds, topk=TOPK)
# Run tilelang kernel
kernel = blocksparse_flashattn(
BATCH, N_HEADS, SEQ_LEN, SEQ_LEN, D_HEAD, downsample_len, is_causal=True)
kernel = blocksparse_flashattn(BATCH, N_HEADS, SEQ_LEN, SEQ_LEN, D_HEAD, downsample_len, is_causal=True)
tilelang_output = kernel(q, k, v, block_mask.to(torch.int8))
# Compute reference
# Expand block mask to full attention matrix
full_mask = torch.kron(block_mask.float(), torch.ones(BLOCK, BLOCK, device='cuda'))
full_mask = torch.kron(block_mask.float(), torch.ones(BLOCK, BLOCK, device="cuda"))
full_mask = full_mask[..., :SEQ_LEN, :SEQ_LEN].bool()
full_mask = full_mask & torch.tril(torch.ones_like(full_mask)) # Apply causal
# PyTorch reference implementation
attn = torch.einsum('bhsd,bhtd->bhst', q, k) * sm_scale
attn = attn.masked_fill(~full_mask, float('-inf'))
attn = torch.einsum("bhsd,bhtd->bhst", q, k) * sm_scale
attn = attn.masked_fill(~full_mask, float("-inf"))
attn = F.softmax(attn, dim=-1)
ref_output = torch.einsum('bhst,bhtd->bhsd', attn, v)
ref_output = torch.einsum("bhst,bhtd->bhsd", attn, v)
print("ref_output", ref_output)
print("tilelang_output", tilelang_output)
# Verify accuracy
assert torch.allclose(tilelang_output, ref_output, atol=1e-2, rtol=1e-2), \
"TileLang output doesn't match reference"
assert torch.allclose(tilelang_output, ref_output, atol=1e-2, rtol=1e-2), "TileLang output doesn't match reference"
print("Pass topk sparse attention test with qlen == klen")
......@@ -213,42 +201,40 @@ def test_topk_sparse_attention_qlen_lt_klen():
torch.manual_seed(0)
# Create inputs.
q = torch.randn(BATCH, N_HEADS, Q_LEN, D_HEAD, device='cuda', dtype=torch.float16)
k = torch.randn(BATCH, N_HEADS, K_LEN, D_HEAD, device='cuda', dtype=torch.float16)
v = torch.randn(BATCH, N_HEADS, K_LEN, D_HEAD, device='cuda', dtype=torch.float16)
q = torch.randn(BATCH, N_HEADS, Q_LEN, D_HEAD, device="cuda", dtype=torch.float16)
k = torch.randn(BATCH, N_HEADS, K_LEN, D_HEAD, device="cuda", dtype=torch.float16)
v = torch.randn(BATCH, N_HEADS, K_LEN, D_HEAD, device="cuda", dtype=torch.float16)
sm_scale = 1.0 / (D_HEAD**0.5)
downsample_factor = BLOCK
downsample_len = math.ceil(K_LEN / downsample_factor) # number of blocks along one dimension
x_ds = torch.randn(
BATCH, N_HEADS, downsample_len, downsample_len, device='cuda', dtype=torch.float16)
x_ds = torch.randn(BATCH, N_HEADS, downsample_len, downsample_len, device="cuda", dtype=torch.float16)
# Force the first column to be high so that the first block is always selected.
x_ds[:, :, :, 0] = 100
block_mask = get_sparse_attn_mask_from_topk(x_ds, topk=TOPK)
kernel = blocksparse_flashattn(
BATCH, N_HEADS, Q_LEN, K_LEN, D_HEAD, downsample_len, is_causal=True)
kernel = blocksparse_flashattn(BATCH, N_HEADS, Q_LEN, K_LEN, D_HEAD, downsample_len, is_causal=True)
print(kernel.get_kernel_source())
tilelang_output = kernel(q, k, v, block_mask.to(torch.int8))
past_len = K_LEN - Q_LEN
attn = torch.einsum('bhsd,bhtd->bhst', q, k) * sm_scale
attn = torch.einsum("bhsd,bhtd->bhst", q, k) * sm_scale
full_mask_full = torch.kron(block_mask.float(), torch.ones(BLOCK, BLOCK, device='cuda')).bool()
full_mask_full = torch.kron(block_mask.float(), torch.ones(BLOCK, BLOCK, device="cuda")).bool()
full_mask_full = full_mask_full[..., :K_LEN, :K_LEN]
effective_mask = full_mask_full[..., past_len:K_LEN, :] # shape: (B, H, Q_LEN, K_LEN)
i_global = torch.arange(past_len, K_LEN, device=k.device).unsqueeze(1) # shape: (Q_LEN, 1)
j_global = torch.arange(K_LEN, device=k.device).unsqueeze(0) # shape: (1, K_LEN)
causal_mask = (j_global <= i_global) # shape: (Q_LEN, K_LEN)
causal_mask = j_global <= i_global # shape: (Q_LEN, K_LEN)
final_mask = effective_mask & causal_mask # shape: (B, H, Q_LEN, K_LEN)
attn = attn.masked_fill(~final_mask, float('-inf'))
attn = attn.masked_fill(~final_mask, float("-inf"))
attn = F.softmax(attn, dim=-1)
ref_output = torch.einsum('bhst,bhtd->bhsd', attn, v)
ref_output = torch.einsum("bhst,bhtd->bhsd", attn, v)
print("ref_output", ref_output)
print("tilelang_output", tilelang_output)
......
examples/seer_attention/block_sparse_attn_triton.py
View file @ 667632cc
......@@ -15,10 +15,7 @@ def get_sparse_attn_mask_from_topk(x, topk, use_dense_for_last_block=False):
bsz, num_head, downsample_len, _ = x.shape
# N_CTX = downsample_len * BLOCK
sparse_index = torch.topk(x, topk, dim=-1).indices
dense_mask = torch.full([bsz, num_head, downsample_len, downsample_len],
False,
dtype=torch.bool,
device=x.device)
dense_mask = torch.full([bsz, num_head, downsample_len, downsample_len], False, dtype=torch.bool, device=x.device)
dense_mask.scatter_(-1, sparse_index, True)
if use_dense_for_last_block:
dense_mask[:, :, -2:, :] = True
......@@ -54,7 +51,6 @@ def _fwd_kernel_inner(
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
):
mask_val = tl.load(block_mask_ptr + k_block_col_idx * stride_bmask_n)
if mask_val == True:
......@@ -69,7 +65,7 @@ def _fwd_kernel_inner(
qk *= sm_scale
# the following is needed only when LAST_K_BLOCK or BLOCK_M < BLOCK_N
qk += tl.where(offs_m[:, None] + past_len >= (start_n + offs_n[None, :]), 0, float('-inf'))
qk += tl.where(offs_m[:, None] + past_len >= (start_n + offs_n[None, :]), 0, float("-inf"))
m_ij = tl.maximum(m_i, tl.max(qk, 1))
qk -= m_ij[:, None]
......@@ -149,7 +145,7 @@ def _fwd_kernel(
v_ptrs = V + off_v
mask_ptrs = block_mask_ptr + start_m * stride_bmm
m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
......@@ -185,24 +181,12 @@ def _fwd_kernel(
acc = acc * l_recip
acc = acc.to(Out.dtype.element_ty)
off_o = off_z * stride_oz + off_h * stride_oh + offs_m[:, None] * stride_om + offs_d[
None, :] * stride_od
off_o = off_z * stride_oz + off_h * stride_oh + offs_m[:, None] * stride_om + offs_d[None, :] * stride_od
out_ptrs = Out + off_o
tl.store(out_ptrs, acc, mask=offs_m[:, None] < N_CTX)
def _forward(ctx,
q,
k,
v,
block_sparse_mask,
sm_scale,
BLOCK_M=64,
BLOCK_N=64,
num_warps=None,
num_stages=1,
out=None):
def _forward(ctx, q, k, v, block_sparse_mask, sm_scale, BLOCK_M=64, BLOCK_N=64, num_warps=None, num_stages=1, out=None):
assert q.shape[-1] == k.shape[-1] == v.shape[-1]
assert k.shape[2] == v.shape[2]
o = out if out is not None else torch.empty_like(q).contiguous()
......@@ -247,7 +231,6 @@ def _forward(ctx,
class _sparse_attention(torch.autograd.Function):
@staticmethod
def forward(ctx, q, k, v, block_sparse_dense, sm_scale):
# shape constraints
......@@ -271,9 +254,9 @@ def test_topk_sparse_attention():
torch.manual_seed(0)
# Create inputs
q = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device='cuda', dtype=torch.bfloat16)
k = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device='cuda', dtype=torch.bfloat16)
v = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device='cuda', dtype=torch.bfloat16)
q = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device="cuda", dtype=torch.bfloat16)
k = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device="cuda", dtype=torch.bfloat16)
v = torch.randn(BATCH, N_HEADS, SEQ_LEN, D_HEAD, device="cuda", dtype=torch.bfloat16)
sm_scale = 1.0 / (D_HEAD**0.5)
# Create sparse mask (downsampled to block level)
......@@ -281,9 +264,7 @@ def test_topk_sparse_attention():
downsample_len = math.ceil(SEQ_LEN / downsample_factor)
print("downsample_len", downsample_len)
x_ds = torch.randn([BATCH, N_HEADS, downsample_len, downsample_len],
device='cuda',
dtype=torch.bfloat16)
x_ds = torch.randn([BATCH, N_HEADS, downsample_len, downsample_len], device="cuda", dtype=torch.bfloat16)
x_ds[:, :, :, 0] = 100
print("x_ds.shape", x_ds.shape)
block_mask = get_sparse_attn_mask_from_topk(x_ds, topk=TOPK)
......@@ -295,22 +276,21 @@ def test_topk_sparse_attention():
# Compute reference
# Expand block mask to full attention matrix
full_mask = torch.kron(block_mask.float(), torch.ones(BLOCK, BLOCK, device='cuda'))
full_mask = torch.kron(block_mask.float(), torch.ones(BLOCK, BLOCK, device="cuda"))
full_mask = full_mask[..., :SEQ_LEN, :SEQ_LEN].bool()
full_mask = full_mask & torch.tril(torch.ones_like(full_mask)) # Apply causal
# PyTorch reference implementation
attn = torch.einsum('bhsd,bhtd->bhst', q, k) * sm_scale
attn = attn.masked_fill(~full_mask, float('-inf'))
attn = torch.einsum("bhsd,bhtd->bhst", q, k) * sm_scale
attn = attn.masked_fill(~full_mask, float("-inf"))
attn = F.softmax(attn, dim=-1)
ref_output = torch.einsum('bhst,bhtd->bhsd', attn, v)
ref_output = torch.einsum("bhst,bhtd->bhsd", attn, v)
# print("ref_output", ref_output)
# print("triton_output", triton_output)
# Verify accuracy
assert torch.allclose(triton_output, ref_output, atol=1e-2, rtol=1e-2), \
"Triton output doesn't match reference"
assert torch.allclose(triton_output, ref_output, atol=1e-2, rtol=1e-2), "Triton output doesn't match reference"
print("Pass topk sparse attention test with qlen == klen")
......@@ -322,16 +302,15 @@ def test_topk_sparse_attention_qlt_kl():
torch.manual_seed(0)
# Create inputs.
q = torch.randn(BATCH, N_HEADS, Q_LEN, D_HEAD, device='cuda', dtype=torch.bfloat16)
k = torch.randn(BATCH, N_HEADS, K_LEN, D_HEAD, device='cuda', dtype=torch.bfloat16)
v = torch.randn(BATCH, N_HEADS, K_LEN, D_HEAD, device='cuda', dtype=torch.bfloat16)
q = torch.randn(BATCH, N_HEADS, Q_LEN, D_HEAD, device="cuda", dtype=torch.bfloat16)
k = torch.randn(BATCH, N_HEADS, K_LEN, D_HEAD, device="cuda", dtype=torch.bfloat16)
v = torch.randn(BATCH, N_HEADS, K_LEN, D_HEAD, device="cuda", dtype=torch.bfloat16)
# softmax scale
sm_scale = 1.0 / (D_HEAD**0.5)
downsample_factor = BLOCK
downsample_len = math.ceil(K_LEN / downsample_factor) # number of blocks along one dimension
x_ds = torch.randn(
BATCH, N_HEADS, downsample_len, downsample_len, device='cuda', dtype=torch.bfloat16)
x_ds = torch.randn(BATCH, N_HEADS, downsample_len, downsample_len, device="cuda", dtype=torch.bfloat16)
# Force the first column to be high so that the first block is always selected.
x_ds[:, :, :, 0] = 100
block_mask = get_sparse_attn_mask_from_topk(x_ds, topk=TOPK)
......@@ -340,26 +319,25 @@ def test_topk_sparse_attention_qlt_kl():
past_len = K_LEN - Q_LEN
attn = torch.einsum('bhsd,bhtd->bhst', q, k) * sm_scale
attn = torch.einsum("bhsd,bhtd->bhst", q, k) * sm_scale
full_mask_full = torch.kron(block_mask.float(), torch.ones(BLOCK, BLOCK, device='cuda')).bool()
full_mask_full = torch.kron(block_mask.float(), torch.ones(BLOCK, BLOCK, device="cuda")).bool()
full_mask_full = full_mask_full[..., :K_LEN, :K_LEN]
effective_mask = full_mask_full[..., past_len:K_LEN, :] # shape: (B, H, Q_LEN, K_LEN)
i_global = torch.arange(past_len, K_LEN, device=k.device).unsqueeze(1) # shape: (Q_LEN, 1)
j_global = torch.arange(K_LEN, device=k.device).unsqueeze(0) # shape: (1, K_LEN)
causal_mask = (j_global <= i_global) # shape: (Q_LEN, K_LEN)
causal_mask = j_global <= i_global # shape: (Q_LEN, K_LEN)
final_mask = effective_mask & causal_mask # shape: (B, H, Q_LEN, K_LEN)
attn = attn.masked_fill(~final_mask, float('-inf'))
attn = attn.masked_fill(~final_mask, float("-inf"))
attn = F.softmax(attn, dim=-1)
ref_output = torch.einsum('bhst,bhtd->bhsd', attn, v)
ref_output = torch.einsum("bhst,bhtd->bhsd", attn, v)
# Verify accuracy.
assert torch.allclose(triton_output, ref_output, atol=1e-2, rtol=1e-2), \
"Triton output doesn't match reference when qlen < klen"
assert torch.allclose(triton_output, ref_output, atol=1e-2, rtol=1e-2), "Triton output doesn't match reference when qlen < klen"
print("Pass topk sparse attention test with qlen < klen")
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment