Designing High-Performance GPU Kernels with TileLang: Tensor-Core GEMM, Fused Softmax, FlashAttention, and Autotuning

In this tutorial, we explore TileLang as a high-level Python domain-specific language for designing and compiling performance-oriented GPU kernels through TVM. We begin by validating the CUDA environment 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 directly with TileLang’s shared-memory tiles, register fragments, pipelined loops, parallel iteration primitives, reductions, and tensor-core GEMM operators while allowing the compiler to manage thread mapping, memory layouts, synchronization, vectorization, and low-level CUDA instruction generation. We also compare our kernels against PyTorch and cuBLAS baselines, inspect generated CUDA source, evaluate memory and compute throughput, and use autotuning to identify architecture-dependent kernel configurations.

Copy CodeCopiedUse a different Browser

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 missing. Stable wheel first, nightly as a fallback.""" try: import tilelang return except ImportError: pass print(">> installing tilelang (this pulls a bundled TVM, ~1-3 min)\n") _sh(f"{sys.executable} -m pip install -q tilelang") try: import tilelang return except ImportError: print(">> stable wheel unusable, trying nightly channel") _sh(f"{sys.executable} -m pip install -q tilelang " f"-f https://tile-ai.github.io/whl/nightly") import tilelang if os.path.isdir("/usr/local/cuda"): os.environ.setdefault("CUDA_HOME", "/usr/local/cuda") os.environ["PATH"] = os.environ.get("PATH", "") + ":/usr/local/cuda/bin" _bootstrap() import torch import torch.nn.functional 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 events.""" for _ in range(warmup): fn() torch.cuda.synchronize() start, end = torch.cuda.Event(True), torch.cuda.Event(True) start.record() for _ in range(rep): fn() end.record() torch.cuda.synchronize() return start.elapsed_time(end) / rep def check(actual: torch.Tensor, ref: torch.Tensor, name: str, tol: float = 2e-2) -> bool: """Relative-Frobenius-norm check. Far more meaningful than atol for fp16.""" a, r = actual.float(), ref.float() rel = (a - r).norm() / r.norm().clamp_min(1e-12) amax = (a - r).abs().max().item() ok = bool(rel Change runtime type -> GPU." DEV = torch.device("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.name} (sm_{SM}, {PROPS.multi_processor_count} SMs, " f"{PROPS.total_memory/2**30:.1f} GiB)") SMEM_CAP = 48 * 1024 if SM 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, K, bm, bn, bk, num_stages=st, threads=128) a = torch.randn(M, K, device=DEV, dtype=torch.float16) b = torch.randn(K, N, device=DEV, dtype=torch.float16) c = kernel(a, b) check(c, a @ b, "matmul 2048^3") flops = 2 * M * N * K 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 lines 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 = K = 2048 a = torch.randn(M, K, device=DEV, dtype=torch.float16) b = torch.randn(K, N, device=DEV, dtype=torch.float16) flops = 2 * M * N * K 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}") results = [] for bm, bn, bk, st, thr, swz in candidates: need = smem_bytes(bm, bn, bk, st) tag = f"{bm}x{bn}x{bk}" if need > SMEM_CAP: print(f" {tag:>16} {st:>4} {thr:>4} {str(swz):>4} {need//1024:>5}KB " f" SKIPPED (over smem budget)") continue try: k = make_matmul(M, N, K, bm, bn, bk, num_stages=st, threads=thr, use_swizzle=swz) c = k(a, b) ok = (c - (a @ b)).float().norm() / (a @ b).float().norm() 16} {st:>4} {thr:>4} {str(swz):>4} {need//1024:>5}KB " f"{ms:>8.3f} {flops/(ms*1e-3)/1e12:>9.2f}" f"{'' if ok else ' 16} {st:>4} {thr:>4} {str(swz):>4} - " f"failed: {type(e).name}") if results: best = min(results) print(f"\n winner: {best[1]}, stages={best[2]}, threads={best[3]}, " f"swizzle={best[4]} ({best[0]:.3f} ms)") print(" Takeaway: the best schedule is arch- and shape-dependent, which is") print(" exactly why section 7 exists.")

We implement a tiled tensor-core matrix multiplication kernel that moves input tiles through global memory, shared memory, and register fragments. We control tile dimensions, pipeline stages, thread counts, and L2 swizzling while allowing TileLang to generate tensor-core instructions, synchronization, and memory-transfer logic. We then benchmark several schedule configurations, verify their numerical accuracy, and identify the highest-performing architecture-dependent kernel configuration.

Copy CodeCopiedUse a different Browser

@tilelang.jit(out_idx=[-1]) def make_matmul_bias_gelu(M: int, N: int, K: 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 main(A: T.Tensor((M, K), dtype), B: T.Tensor((K, 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(K, 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 main def section_4(): banner("4. EPILOGUE FUSION — GEMM + bias + GELU in one kernel") M, N, K = 4096, 4096, 1024 st = DEFAULT_STAGES while smem_bytes(128, 128, 32, st) > SMEM_CAP and st > 1: st -= 1 fused = make_matmul_bias_gelu(M, N, K, 128, 128, 32, num_stages=st, threads=128) a = (torch.randn(M, K, device=DEV, dtype=torch.float16) / K 0.25) b = (torch.randn(K, N, device=DEV, dtype=torch.float16) / K 0.25) bias = torch.randn(N, device=DEV, dtype=torch.float16) out = fused(a, b, bias) ref = F.gelu(((a @ b).float() + bias.float()), approximate="tanh") check(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 traffic 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 main(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 main 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, device=DEV, dtype=torch.float16) * 3.0 y = sm(x) check(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 memory bound; the interesting bit is that the two-pass") print(" max/sum reduction never left registers.")

We extend the matrix multiplication kernel by fusing bias addition and the GELU activation directly into the register-resident accumulator. We reduce intermediate global-memory traffic by completing the epilogue before writing the final output tensor and compare the fused implementation with eager PyTorch execution. We also implement a row-wise softmax kernel using fragment-level maximum and sum reductions while keeping the normalization process largely within registers.

Copy CodeCopiedUse a different Browser

@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) shape = [batch, seq_len, heads, dim] NEG = -1.0e30 @T.prim_func def main(Q: T.Tensor(shape, dtype), K: T.Tensor(shape, dtype), V: T.Tensor(shape, dtype), O: T.Tensor(shape, 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 k in T.Pipelined(loop_range, num_stages=num_stages): T.copy(K[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 >= k * 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.e

[truncated for AI cost control]