Designing High-Performance GPU Kernels with TileLang: Tensor-Core GEMM, Fused Softmax, FlashAttention, and Autotuning
In this tutorial, we discover TileLang as a high-level Python domain-specific language for designing and compiling performance-oriented GPU kernels by TVM. We start by validating the CUDA atmosphere and establishing reusable benchmarking and numerical-verification utilities, then progressively implement vector addition, tiled tensor-core matrix multiplication, schedule exploration, fused GEMM epilogues, row-wise softmax, and FlashAttention. Throughout the tutorial, we work immediately with TileLang’s shared-memory tiles, register fragments, pipelined loops, parallel iteration primitives, reductions, and tensor-core GEMM operators whereas permitting the compiler to handle thread mapping, reminiscence layouts, synchronization, vectorization, and low-level CUDA instruction technology. We additionally evaluate our kernels in opposition to PyTorch and cuBLAS baselines, examine generated CUDA supply, consider reminiscence and compute throughput, and use autotuning to establish architecture-dependent kernel configurations.
import os
import sys
import math
import time
import subprocess
import traceback
def _sh(cmd: str) -> int:
print(f"$ {cmd}", flush=True)
return subprocess.run(cmd, shell=True).returncode
def _bootstrap():
"""Install tilelang if lacking. Stable wheel first, nightly as a fallback."""
attempt:
import tilelang
return
besides ImportError:
go
print(">> putting in tilelang (this pulls a bundled TVM, ~1-3 min)n")
_sh(f"{sys.executable} -m pip set up -q tilelang")
attempt:
import tilelang
return
besides ImportError:
print(">> secure wheel unusable, attempting nightly channel")
_sh(f"{sys.executable} -m pip set up -q tilelang "
f"-f https://tile-ai.github.io/whl/nightly")
import tilelang
if os.path.isdir("/usr/native/cuda"):
os.environ.setdefault("CUDA_HOME", "/usr/native/cuda")
os.environ["PATH"] = os.environ.get("PATH", "") + ":/usr/native/cuda/bin"
_bootstrap()
import torch
import torch.nn.practical as F
import tilelang
import tilelang.language as T
def banner(title: str):
print("n" + "=" * 78)
print(f" {title}")
print("=" * 78, flush=True)
def bench(fn, warmup: int = 10, rep: int = 50) -> float:
"""Median-ish latency in milliseconds, measured with CUDA occasions."""
for _ in vary(warmup):
fn()
torch.cuda.synchronize()
begin, finish = torch.cuda.Event(True), torch.cuda.Event(True)
begin.file()
for _ in vary(rep):
fn()
finish.file()
torch.cuda.synchronize()
return begin.elapsed_time(finish) / rep
def examine(precise: torch.Tensor, ref: torch.Tensor, identify: str, tol: float = 2e-2) -> bool:
"""Relative-Frobenius-norm examine. Far extra significant than atol for fp16."""
a, r = precise.float(), ref.float()
rel = (a - r).norm() / r.norm().clamp_min(1e-12)
amax = (a - r).abs().max().merchandise()
okay = bool(rel < tol) and math.isfinite(amax)
print(f" [{'PASS' if ok else 'FAIL'}] {identify}: rel_err={rel:.3e} max_abs={amax:.3e}")
return okay
banner("0. ENVIRONMENT")
assert torch.cuda.is_available(), "No GPU. Runtime -> Change runtime sort -> GPU."
DEV = torch.system("cuda")
PROPS = torch.cuda.get_device_properties(0)
CC = torch.cuda.get_device_capability(0)
SM = CC[0] * 10 + CC[1]
print(f" tilelang : {getattr(tilelang, '__version__', 'unknown')}")
print(f" torch : {torch.__version__}")
print(f" GPU : {PROPS.identify} (sm_{SM}, {PROPS.multi_processor_count} SMs, "
f"{PROPS.total_memory/2**30:.1f} GiB)")
SMEM_CAP = 48 * 1024 if SM < 80 else 96 * 1024
DEFAULT_STAGES = 2 if SM < 80 else 3
print(f" smem funds: {SMEM_CAP//1024} KB/block, default num_stages={DEFAULT_STAGES}")
print(" observe: tilelang caches compiled kernels in ~/.tilelang/cache, so a")
print(" second run of this cell is dramatically quicker.")
@tilelang.jit(out_idx=[-1])
def make_vector_add(N: int, block_N: int = 256, dtype: str = "float32"):
@T.prim_func
def foremost(A: T.Tensor((N,), dtype),
B: T.Tensor((N,), dtype),
C: T.Tensor((N,), dtype)):
with T.Kernel(T.ceildiv(N, block_N), threads=256) as bx:
for i in T.Parallel(block_N):
C[bx * block_N + i] = A[bx * block_N + i] + B[bx * block_N + i]
return foremost
def section_1():
banner("1. HELLO, TILE — vector add + generated CUDA")
N = 1 << 22
add = make_vector_add(N, block_N=256)
a = torch.randn(N, system=DEV)
b = torch.randn(N, system=DEV)
c = add(a, b)
examine(c, a + b, "vector_add")
ms = bench(lambda: add(a, b))
gbs = 3 * N * 4 / (ms * 1e-3) / 1e9
ref_ms = bench(lambda: a + b)
print(f" tilelang: {ms*1e3:8.1f} us ({gbs:6.1f} GB/s)")
print(f" torch : {ref_ms*1e3:8.1f} us <- each are pure bandwidth, so a tie"
f" is the proper end result")
src = add.get_kernel_source()
print("n --- generated system code (first 32 traces) ---")
for line in src.splitlines()[:32]:
print(" " + line)
print(" ---------------------------------------------")
We configure the Google Colab CUDA atmosphere, set up TileLang with a nightly fallback, and import the required PyTorch and TileLang modules. We outline reusable benchmarking, validation, and reporting utilities that measure kernel latency and evaluate numerical outputs utilizing relative error. We then implement a TileLang vector-add kernel, execute it on the GPU, evaluate its bandwidth with PyTorch, and examine the CUDA supply generated by the compiler.
@tilelang.jit(out_idx=[-1])
def make_matmul(M: int, N: int, Ok: int,
block_M: int = 128, block_N: int = 128, block_K: int = 32,
num_stages: int = 3, threads: int = 128,
use_swizzle: bool = False,
dtype: str = "float16", accum_dtype: str = "float"):
@T.prim_func
def foremost(A: T.Tensor((M, Ok), dtype),
B: T.Tensor((Ok, N), dtype),
C: T.Tensor((M, N), dtype)):
with T.Kernel(T.ceildiv(N, block_N),
T.ceildiv(M, block_M),
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)
if use_swizzle:
T.use_swizzle(panel_size=10, allow=True)
T.clear(C_local)
for ko in T.Pipelined(T.ceildiv(Ok, block_K), num_stages=num_stages):
T.copy(A[by * block_M, ko * block_K], A_shared)
T.copy(B[ko * block_K, bx * block_N], B_shared)
T.gemm(A_shared, B_shared, C_local)
T.copy(C_local, C[by * block_M, bx * block_N])
return foremost
def smem_bytes(block_M, block_N, block_K, levels, itemsize=2):
return (block_M * block_K + block_K * block_N) * itemsize * levels
def section_2():
banner("2. MEMORY HIERARCHY — tiled tensor-core GEMM")
M = N = Ok = 2048
bm, bn, bk, st = 128, 128, 32, DEFAULT_STAGES
whereas smem_bytes(bm, bn, bk, st) > SMEM_CAP and st > 1:
st -= 1
print(f" config: {bm}x{bn}x{bk}, num_stages={st}, "
f"smem={smem_bytes(bm,bn,bk,st)/1024:.0f} KB")
kernel = make_matmul(M, N, Ok, bm, bn, bk, num_stages=st, threads=128)
a = torch.randn(M, Ok, system=DEV, dtype=torch.float16)
b = torch.randn(Ok, N, system=DEV, dtype=torch.float16)
c = kernel(a, b)
examine(c, a @ b, "matmul 2048^3")
flops = 2 * M * N * Ok
ms = bench(lambda: kernel(a, b))
ms_ref = bench(lambda: a @ b)
print(f" tilelang : {ms:7.3f} ms -> {flops/(ms*1e-3)/1e12:6.2f} TFLOP/s")
print(f" cuBLAS : {ms_ref:7.3f} ms -> {flops/(ms_ref*1e-3)/1e12:6.2f} TFLOP/s")
print(f" ratio : {ms_ref/ms*100:5.1f}% of cuBLAS from ~20 traces of Python")
src = kernel.get_kernel_source()
for needle in ("mma.sync", "wgmma", "ldmatrix", "cp.async", "tl::gemm"):
if needle in src:
print(f" emitted: {needle}")
return kernel
def section_3():
banner("3. KNOBS — sweeping the schedule by hand")
M = N = Ok = 2048
a = torch.randn(M, Ok, system=DEV, dtype=torch.float16)
b = torch.randn(Ok, N, system=DEV, dtype=torch.float16)
flops = 2 * M * N * Ok
candidates = [
(64, 64, 32, 2, 128, False),
(128, 128, 32, 2, 128, False),
(128, 128, 32, 2, 128, True),
(128, 128, 32, 3, 128, False),
(128, 128, 64, 2, 256, False),
(128, 256, 32, 2, 256, False),
]
print(f" {'tile':>16} {'stg':>4} {'thr':>4} {'swz':>4} {'smem':>7} "
f"{'ms':>8} {'TFLOP/s':>9}")
outcomes = []
for bm, bn, bk, st, thr, swz in candidates:
want = smem_bytes(bm, bn, bk, st)
tag = f"{bm}x{bn}x{bk}"
if want > SMEM_CAP:
print(f" {tag:>16} {st:>4} {thr:>4} {str(swz):>4} {want//1024:>5}KB "
f" SKIPPED (over smem funds)")
proceed
attempt:
okay = make_matmul(M, N, Ok, bm, bn, bk, num_stages=st,
threads=thr, use_swizzle=swz)
c = okay(a, b)
okay = (c - (a @ b)).float().norm() / (a @ b).float().norm() < 2e-2
ms = bench(lambda: okay(a, b), warmup=5, rep=20)
outcomes.append((ms, tag, st, thr, swz))
print(f" {tag:>16} {st:>4} {thr:>4} {str(swz):>4} {want//1024:>5}KB "
f"{ms:>8.3f} {flops/(ms*1e-3)/1e12:>9.2f}"
f"{'' if okay else ' <- NUMERICALLY WRONG'}")
besides Exception as e:
print(f" {tag:>16} {st:>4} {thr:>4} {str(swz):>4} - "
f"failed: {sort(e).__name__}")
if outcomes:
greatest = min(outcomes)
print(f"n winner: {greatest[1]}, levels={greatest[2]}, threads={greatest[3]}, "
f"swizzle={greatest[4]} ({greatest[0]:.3f} ms)")
print(" Takeaway: the very best schedule is arch- and shape-dependent, which is")
print(" precisely why part 7 exists.")
We implement a tiled tensor-core matrix multiplication kernel that strikes enter tiles by world reminiscence, shared reminiscence, and register fragments. We management tile dimensions, pipeline levels, thread counts, and L2 swizzling whereas permitting TileLang to generate tensor-core directions, synchronization, and memory-transfer logic. We then benchmark a number of schedule configurations, confirm their numerical accuracy, and establish the highest-performing architecture-dependent kernel configuration.
@tilelang.jit(out_idx=[-1])
def make_matmul_bias_gelu(M: int, N: int, Ok: int,
block_M: int = 128, block_N: int = 128, block_K: int = 32,
num_stages: int = 2, threads: int = 128,
dtype: str = "float16", accum_dtype: str = "float"):
@T.prim_func
def foremost(A: T.Tensor((M, Ok), dtype),
B: T.Tensor((Ok, N), dtype),
Bias: T.Tensor((N,), dtype),
C: T.Tensor((M, N), dtype)):
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M),
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)
Bias_shared = T.alloc_shared((block_N,), dtype)
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
T.clear(C_local)
T.copy(Bias[bx * block_N], Bias_shared)
for ko in T.Pipelined(T.ceildiv(Ok, block_K), num_stages=num_stages):
T.copy(A[by * block_M, ko * block_K], A_shared)
T.copy(B[ko * block_K, bx * block_N], B_shared)
T.gemm(A_shared, B_shared, C_local)
for i, j in T.Parallel(block_M, block_N):
C_local[i, j] = C_local[i, j] + Bias_shared[j]
for i, j in T.Parallel(block_M, block_N):
C_local[i, j] = C_local[i, j] / (
1.0 + T.exp(-1.5957691216 * (
C_local[i, j] + 0.044715 * C_local[i, j]
* C_local[i, j] * C_local[i, j])))
T.copy(C_local, C[by * block_M, bx * block_N])
return foremost
def section_4():
banner("4. EPILOGUE FUSION — GEMM + bias + GELU in a single kernel")
M, N, Ok = 4096, 4096, 1024
st = DEFAULT_STAGES
whereas smem_bytes(128, 128, 32, st) > SMEM_CAP and st > 1:
st -= 1
fused = make_matmul_bias_gelu(M, N, Ok, 128, 128, 32, num_stages=st, threads=128)
a = (torch.randn(M, Ok, system=DEV, dtype=torch.float16) / Ok ** 0.25)
b = (torch.randn(Ok, N, system=DEV, dtype=torch.float16) / Ok ** 0.25)
bias = torch.randn(N, system=DEV, dtype=torch.float16)
out = fused(a, b, bias)
ref = F.gelu(((a @ b).float() + bias.float()), approximate="tanh")
examine(out, ref, "matmul+bias+gelu", tol=3e-2)
ms_f = bench(lambda: fused(a, b, bias))
ms_e = bench(lambda: F.gelu(a @ b + bias, approximate="tanh"))
print(f" fused (1 kernel) : {ms_f:7.3f} ms")
print(f" torch (3 kernels) : {ms_e:7.3f} ms")
print(f" speedup : {ms_e/ms_f:5.2f}x")
print(f" HBM site visitors saved : ~{2*M*N*2/2**20:.0f} MiB of intermediate reads/writes")
@tilelang.jit(out_idx=[-1])
def make_softmax(M: int, N: int, block_M: int = 4, threads: int = 128,
dtype: str = "float16", accum_dtype: str = "float"):
@T.prim_func
def foremost(X: T.Tensor((M, N), dtype),
Y: T.Tensor((M, N), dtype)):
with T.Kernel(T.ceildiv(M, block_M), threads=threads) as bx:
X_shared = T.alloc_shared((block_M, N), dtype)
X_local = T.alloc_fragment((block_M, N), accum_dtype)
row_max = T.alloc_fragment((block_M,), accum_dtype)
row_sum = T.alloc_fragment((block_M,), accum_dtype)
T.copy(X[bx * block_M, 0], X_shared)
T.copy(X_shared, X_local)
T.reduce_max(X_local, row_max, dim=1, clear=True)
for i, j in T.Parallel(block_M, N):
X_local[i, j] = T.exp(X_local[i, j] - row_max[i])
T.reduce_sum(X_local, row_sum, dim=1)
for i, j in T.Parallel(block_M, N):
X_local[i, j] = X_local[i, j] / row_sum[i]
T.copy(X_local, Y[bx * block_M, 0])
return foremost
def section_5():
banner("5. REDUCTIONS — row-wise softmax")
M, N = 8192, 1024
sm = make_softmax(M, N, block_M=4, threads=128)
x = torch.randn(M, N, system=DEV, dtype=torch.float16) * 3.0
y = sm(x)
examine(y, torch.softmax(x.float(), dim=-1), "softmax")
ms = bench(lambda: sm(x))
ms_ref = bench(lambda: torch.softmax(x, dim=-1))
gbs = 2 * M * N * 2 / (ms * 1e-3) / 1e9
print(f" tilelang: {ms*1e3:7.1f} us ({gbs:6.1f} GB/s)")
print(f" torch : {ms_ref*1e3:7.1f} us")
print(" Both are reminiscence sure; the fascinating bit is that the two-pass")
print(" max/sum discount by no means left registers.")
We lengthen the matrix multiplication kernel by fusing bias addition and the GELU activation immediately into the register-resident accumulator. We scale back intermediate global-memory site visitors by finishing the epilogue earlier than writing the ultimate output tensor and evaluate the fused implementation with keen PyTorch execution. We additionally implement a row-wise softmax kernel utilizing fragment-level most and sum reductions whereas preserving the normalization course of largely inside registers.
@tilelang.jit(out_idx=[-1])
def make_flash_attn(batch: int, heads: int, seq_len: int, dim: int,
is_causal: bool = False,
block_M: int = 64, block_N: int = 64,
num_stages: int = 1, threads: int = 128,
dtype: str = "float16", accum_dtype: str = "float"):
scale = 1.0 / math.sqrt(dim)
form = [batch, seq_len, heads, dim]
NEG = -1.0e30
@T.prim_func
def foremost(Q: T.Tensor(form, dtype),
Ok: T.Tensor(form, dtype),
V: T.Tensor(form, dtype),
O: T.Tensor(form, dtype)):
with T.Kernel(T.ceildiv(seq_len, block_M), heads, batch,
threads=threads) as (bx, by, bz):
Q_shared = T.alloc_shared([block_M, dim], dtype)
K_shared = T.alloc_shared([block_N, dim], dtype)
V_shared = T.alloc_shared([block_N, dim], dtype)
acc_s = T.alloc_fragment([block_M, block_N], accum_dtype)
acc_s_cast = T.alloc_fragment([block_M, block_N], dtype)
acc_o = T.alloc_fragment([block_M, dim], accum_dtype)
m_prev = T.alloc_fragment([block_M], accum_dtype)
m_cur = T.alloc_fragment([block_M], accum_dtype)
alpha = T.alloc_fragment([block_M], accum_dtype)
p_sum = T.alloc_fragment([block_M], accum_dtype)
logsum = T.alloc_fragment([block_M], accum_dtype)
T.copy(Q[bz, bx * block_M:(bx + 1) * block_M, by, :], Q_shared)
T.fill(acc_o, 0)
T.fill(logsum, 0)
T.fill(m_cur, NEG)
loop_range = (T.ceildiv((bx + 1) * block_M, block_N)
if is_causal else T.ceildiv(seq_len, block_N))
for okay in T.Pipelined(loop_range, num_stages=num_stages):
T.copy(Ok[bz, k * block_N:(k + 1) * block_N, by, :], K_shared)
if is_causal:
for i, j in T.Parallel(block_M, block_N):
acc_s[i, j] = T.if_then_else(
bx * block_M + i >= okay * block_N + j, 0.0, NEG)
else:
T.clear(acc_s)
T.gemm(Q_shared, K_shared, acc_s, transpose_B=True)
T.copy(V[bz, k * block_N:(k + 1) * block_N, by, :], V_shared)
T.copy(m_cur, m_prev)
T.reduce_max(acc_s, m_cur, dim=1, clear=False)
for i in T.Parallel(block_M):
alpha[i] = T.exp((m_prev[i] - m_cur[i]) * scale)
for i, j in T.Parallel(block_M, block_N):
acc_s[i, j] = T.exp((acc_s[i, j] - m_cur[i]) * scale)
T.reduce_sum(acc_s, p_sum, dim=1)
for i in T.Parallel(block_M):
logsum[i] = logsum[i] * alpha[i] + p_sum[i]
for i, j in T.Parallel(block_M, dim):
acc_o[i, j] = acc_o[i, j] * alpha[i]
T.copy(acc_s, acc_s_cast)
T.gemm(acc_s_cast, V_shared, acc_o)
for i, j in T.Parallel(block_M, dim):
acc_o[i, j] = acc_o[i, j] / logsum[i]
T.copy(acc_o, O[bz, bx * block_M:(bx + 1) * block_M, by, :])
return foremost
def section_6():
banner("6. FLASHATTENTION — fused consideration ahead")
B, H, S, D = 4, 8, 1024, 64
q = torch.randn(B, S, H, D, system=DEV, dtype=torch.float16)
okay = torch.randn(B, S, H, D, system=DEV, dtype=torch.float16)
v = torch.randn(B, S, H, D, system=DEV, dtype=torch.float16)
def torch_ref(causal: bool):
qt, kt, vt = (t.transpose(1, 2) for t in (q, okay, v))
o = F.scaled_dot_product_attention(qt, kt, vt, is_causal=causal)
return o.transpose(1, 2).contiguous()
for causal in (False, True):
tag = "causal" if causal else "full"
attn = make_flash_attn(B, H, S, D, is_causal=causal,
block_M=64, block_N=64,
num_stages=1 if SM < 80 else 2, threads=128)
o = attn(q, okay, v)
ref = torch_ref(causal)
examine(o, ref, f"flash_attn ({tag})", tol=3e-2)
flops = 4 * B * H * S * S * D * (0.5 if causal else 1.0)
ms = bench(lambda: attn(q, okay, v), warmup=5, rep=20)
ms_ref = bench(lambda: torch_ref(causal), warmup=5, rep=20)
print(f" {tag:>6}: tilelang {ms:6.3f} ms "
f"({flops/(ms*1e-3)/1e12:5.2f} TFLOP/s) "
f"torch SDPA {ms_ref:6.3f} ms "
f"({flops/(ms_ref*1e-3)/1e12:5.2f} TFLOP/s)")
print(" ~70 traces of Python for a fused, causal, tensor-core consideration.")
print(" The upstream repo pushes the identical construction to FlashMLA-level perf")
print(" on H100 with warp specialisation and TMA.")
We implement a fused FlashAttention ahead kernel that processes question, key, and worth tiles with out materializing the complete attention-score matrix in world reminiscence. We apply on-line softmax updates utilizing operating maxima, normalization sums, rescaling components, and tiled tensor-core matrix multiplications. We validate each causal and non-causal consideration in opposition to PyTorch scaled dot-product consideration and evaluate their latency and computational throughput.
def section_7():
banner("7. AUTOTUNING — @tilelang.autotune")
M = N = Ok = 2048
def configs():
out = []
for bm in (64, 128):
for bn in (128, 256):
for bk in (32, 64):
for levels in (2, DEFAULT_STAGES):
for thr in (128, 256):
if smem_bytes(bm, bn, bk, levels) > SMEM_CAP:
proceed
cfg = dict(block_M=bm, block_N=bn, block_K=bk,
num_stages=levels, threads=thr)
if cfg not in out:
out.append(cfg)
return out[:8]
house = configs()
print(f" looking {len(house)} configurations "
f"(every one is an actual nvcc compile + benchmark)")
@tilelang.autotune(configs=house, warmup=10, rep=20)
@tilelang.jit(out_idx=[-1])
def tuned_matmul(M: int, N: int, Ok: int,
block_M: int = 128, block_N: int = 128, block_K: int = 32,
num_stages: int = 2, threads: int = 128,
dtype: str = "float16", accum_dtype: str = "float"):
@T.prim_func
def foremost(A: T.Tensor((M, Ok), dtype),
B: T.Tensor((Ok, N), dtype),
C: T.Tensor((M, N), dtype)):
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M),
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)
T.clear(C_local)
for ko in T.Pipelined(T.ceildiv(Ok, block_K), num_stages=num_stages):
T.copy(A[by * block_M, ko * block_K], A_shared)
T.copy(B[ko * block_K, bx * block_N], B_shared)
T.gemm(A_shared, B_shared, C_local)
T.copy(C_local, C[by * block_M, bx * block_N])
return foremost
a = torch.randn(M, Ok, system=DEV, dtype=torch.float16)
b = torch.randn(Ok, N, system=DEV, dtype=torch.float16)
t0 = time.time()
from tilelang.autotuner import set_autotune_inputs
with set_autotune_inputs(a, b):
greatest = tuned_matmul(M, N, Ok)
print(f" tuning took {time.time()-t0:.0f} s")
c = greatest(a, b)
examine(c, a @ b, "autotuned matmul")
ms = bench(lambda: greatest(a, b))
print(f" greatest kernel: {ms:.3f} ms -> {2*M*N*Ok/(ms*1e-3)/1e12:.2f} TFLOP/s")
print(" Re-running this cell hits the on-disk autotuner cache and is immediate.")
print(" Turn caching off with TILELANG_AUTO_TUNING_DISABLE_CACHE=1.")
We outline an autotuning search house throughout matrix tile sizes, Ok-block dimensions, pipeline depths, and thread counts whereas filtering configurations that exceed the shared-memory funds. We use TileLang’s autotuning decorator to compile, benchmark, validate, and cache a number of kernel schedules for a similar matrix-multiplication workload. We then execute the chosen kernel, confirm its output in opposition to PyTorch, and report the achieved latency and tensor-core throughput.
@tilelang.jit
def make_print_demo(N: int = 128, dtype: str = "float32"):
"""T.print emits a guarded device-side printf. Keep the grid tiny."""
@T.prim_func
def foremost(A: T.Tensor((N,), dtype), B: T.Tensor((N,), dtype)):
with T.Kernel(1, threads=32) as bx:
A_local = T.alloc_fragment((N,), dtype)
T.copy(A, A_local)
for i in T.Parallel(N):
A_local[i] = A_local[i] * 2.0
T.print(A_local[0], msg="A_local[0] after doubling")
T.copy(A_local, B)
return foremost
def section_8():
banner("8. INTROSPECTION — T.print, sources, profiler")
attempt:
dbg = make_print_demo(128)
a = torch.ones(128, system=DEV)
b = torch.empty(128, system=DEV)
dbg(a, b)
torch.cuda.synchronize()
print(f" T.print fired above; host examine: b[0] = {b[0].merchandise()} (anticipate 2.0)")
besides Exception as e:
print(f" T.print demo skipped: {sort(e).__name__}: {e}")
kern = make_matmul(1024, 1024, 1024, 128, 128, 32,
num_stages=min(2, DEFAULT_STAGES), threads=128)
src = kern.get_kernel_source()
print(f"n get_kernel_source(): {len(src.splitlines())} traces of CUDA")
print(" grep-worthy landmarks:")
for needle in ("__global__", "extern "C"", "mma", "cp.async",
"__syncthreads", "ldmatrix"):
n = src.rely(needle)
if n:
print(f" {needle:<16} x{n}")
attempt:
prof = kern.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
print(f"n profiler.do_bench(): {prof.do_bench():.3f} ms "
f"(it synthesises its personal inputs)")
besides Exception as e:
print(f"n profiler skipped: {sort(e).__name__}: {e}")
print("n Other instruments price figuring out:")
print(" kernel.get_host_source() - the launch wrapper")
print(" T.device_assert(cond, msg) - device-side assertions")
print(" TILELANG_DISABLE_CACHE=1 - drive recompilation")
print(" tilelang.instruments.plot_layout - visualise fragment/shared layouts")
print(" env TILELANG_CACHE_DIR - the place compiled kernels reside")
CHEATSHEET = """
SCOPES T.alloc_shared(form, dtype) __shared__ tile
T.alloc_fragment(form, dtype) registers, format inferred
T.alloc_var(dtype) one scalar per thread
T.alloc_barrier(n) mbarrier (Hopper pipelines)
MOVEMENT T.copy(src, dst) vectorized + casting transfer
T.async_copy(src, dst) express cp.async
T.tma_copy(...) Hopper bulk async
T.transpose(src, dst) shared-memory transpose
COMPUTE T.gemm(A_s, B_s, C_f,
transpose_A/B=...,
coverage=T.GemmWarpPolicy.*) tile matmul on tensor cores
T.gemm_sp(...) 2:4 structured sparsity
T.reduce_sum/max/min(buf, out, dim=, clear=)
T.cumsum / T.cummax scans
T.clear(buf) / T.fill(buf, v)
LOOPS T.Parallel(*extents) elementwise -> threads
T.Pipelined(n, num_stages=okay) software program pipeline
T.serial(n) plain sequential loop
ANNOTATIONS T.use_swizzle(panel_size=, allow=) L2 rasterization
T.annotate_layout({buf: format}) express layouts
T.annotate_l2_hit_ratio(buf, r)
ENTRY POINTS @tilelang.jit(out_idx=[-1]) final arg is the return worth
@tilelang.autotune(configs=...) stack above @jit
kernel.get_kernel_source()
kernel.get_profiler().do_bench()
ENV VARS TILELANG_CACHE_DIR, TILELANG_DISABLE_CACHE,
TILELANG_AUTO_TUNING_DISABLE_CACHE,
TILELANG_AUTO_TUNING_MAX_CPU_COUNT
WHERE TO GO NEXT examples/flash_attention fwd + bwd, autotuned
examples/dequantize_gemm W4A16 with LOP3 tips
examples/deepseek_mla MLA decode, ~80 traces
examples/linear_attention RetNet / Mamba
github.com/tile-ai/tilelang-puzzles 10 graded workouts
tilelang.com full docs + API reference
"""
SECTIONS = [
("1 hello / vector add", section_1),
("2 tiled GEMM", section_2),
("3 schedule sweep", section_3),
("4 epilogue fusion", section_4),
("5 softmax reduction", section_5),
("6 flash attention", section_6),
("7 autotuning", section_7),
("8 introspection", section_8),
]
def foremost(solely=None):
"""solely: non-compulsory listing of 1-based part numbers, e.g. foremost([2, 6])."""
t_start = time.time()
standing = []
for idx, (identify, fn) in enumerate(SECTIONS, begin=1):
if solely and idx not in solely:
proceed
t0 = time.time()
attempt:
fn()
standing.append((identify, "okay", time.time() - t0))
besides Exception:
print("n !! part failed — persevering with with the restn")
traceback.print_exc()
standing.append((identify, "FAILED", time.time() - t0))
lastly:
torch.cuda.synchronize()
banner("9. SUMMARY & CHEATSHEET")
for identify, st, dt in standing:
print(f" {st:>6} {identify:<24} {dt:6.1f} s")
print(f"n whole: {time.time()-t_start:.0f} s")
print(CHEATSHEET)
if __name__ == "__main__":
foremost()
We introduce TileLang’s debugging and introspection workflow by device-side printing, generated CUDA inspection, and the built-in kernel profiler. We study compiler-emitted landmarks similar to tensor-core operations, asynchronous copies, synchronization obstacles, and matrix-load directions. We lastly arrange all tutorial sections right into a fault-tolerant runner that data execution standing, studies timing data, and prints a compact TileLang programming reference.
In conclusion, we constructed a sensible understanding of how TileLang interprets tile-level Python applications into optimized GPU kernels with out requiring us to manually handle thread indices, warp-level information layouts, tensor-core directions, or asynchronous reminiscence obstacles. We applied and validated kernels that cowl bandwidth-bound elementwise operations, compute-intensive GEMM workloads, fused neural-network epilogues, register-resident reductions, and online-softmax consideration. We additionally examined how block dimensions, shared-memory consumption, pipeline depth, thread rely, tile form, and L2 swizzling affect efficiency throughout completely different GPU architectures. Finally, we used generated-source inspection, device-side debugging, profiling, and automated schedule search to ascertain a whole workflow for growing, verifying, benchmarking, and refining customized TileLang kernels.
Check out the Full Code here. Also, be happy to comply with us on Twitter and don’t neglect to hitch our 150k+ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.
Need to companion with us for selling your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar and many others.? Connect with us
The put up Designing High-Performance GPU Kernels with TileLang: Tensor-Core GEMM, Fused Softmax, FlashAttention, and Autotuning appeared first on MarkTechPost.
