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 < 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), 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 main 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, device=DEV) b = torch.randn(N, device=DEV) c = add(a, b) check(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 <- both are pure bandwidth, so a tie” f” is the correct outcome”) src = add.get_kernel_source() print(“n — generated device code (first 32 lines) —“) for line in src.splitlines()[:32]: print(” ” + line) print(” ———————————————“) We configure the Google Colab CUDA environment, install TileLang with a nightly fallback, and import the required PyTorch and TileLang modules. We define reusable benchmarking, validation, and reporting utilities that measure kernel latency and compare numerical outputs using relative error. We then implement a TileLang vector-add kernel, execute it on the GPU, compare its bandwidth with PyTorch, and inspect the CUDA source generated by the compiler. Copy CodeCopiedUse a different Browser @tilelang.jit(out_idx=[-1]) def make_matmul(M: int, N: int, K: 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 main(A: T.Tensor((M, K), dtype), B: T.Tensor((K, 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, enable=True) T.clear(C_local) 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) T.copy(C_local, C[by * block_M, bx * block_N]) return main def smem_bytes(block_M, block_N, block_K, stages, itemsize=2): return (block_M * block_K + block_K * block_N) * itemsize * stages def section_2(): banner(“2. MEMORY HIERARCHY — tiled tensor-core GEMM”) M = N = K = 2048 bm, bn, bk, st = 128, 128, 32, DEFAULT_STAGES while 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, 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,