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 Browserimport 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 < tol) and math.isfinite(amax) print(f" [{'PASS' if ok else 'FAIL'}] {name}: rel_err={rel:.3e} max_abs={amax:.3e}") return ok banner("0. ENVIRONMENT") assert torch.cuda.is_available(), "No GPU. Runtime -> 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 < 80 else 96 * 1024 DEFAULT_STAGES = 2 if SM < 80 else 3 print(f" smem budget: {SMEM_CAP//1024} KB/block, default num_stages={DEFAULT_STAGES}") print(" note: tilelang caches compiled kernels in ~/.tilelang/cache, so a") print(" second run of this cell is dramatically faster.") @tilelang.jit(out_idx=[-1]) def make_vector_add(N: int, block_N: int = 256, dtype: str = "float32"): @T.prim_func def main(A: T.Tensor((N,), dtype),