Source: MarkTechPost
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.
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("https://www.marktechpost.com/usr/local/cuda"): os.environ.setdefault("CUDA_HOME", "https://www.marktechpost.com/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.
@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, 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() < 2e-2 ms = bench(lambda: k(a, b), warmup=5, rep=20) results.append((ms, tag, st, thr, swz)) print(f" {tag:>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 ' <- NUMERICALLY WRONG'}") except Exception as e: print(f" {tag:>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.
@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.
@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.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 main def section_6(): banner("6. FLASHATTENTION — fused attention forward") B, H, S, D = 4, 8, 1024, 64 q = torch.randn(B, S, H, D, device=DEV, dtype=torch.float16) k = torch.randn(B, S, H, D, device=DEV, dtype=torch.float16) v = torch.randn(B, S, H, D, device=DEV, dtype=torch.float16) def torch_ref(causal: bool): qt, kt, vt = (t.transpose(1, 2) for t in (q, k, 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, k, v) ref = torch_ref(causal) check(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, k, 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 lines of Python for a fused, causal, tensor-core attention.") print(" The upstream repo pushes the same structure to FlashMLA-level perf") print(" on H100 with warp specialisation and TMA.") We implement a fused FlashAttention forward kernel that processes query, key, and value tiles without materializing the full attention-score matrix in global memory. We apply online softmax updates using running maxima, normalization sums, rescaling factors, and tiled tensor-core matrix multiplications. We validate both causal and non-causal attention against PyTorch scaled dot-product attention and compare their latency and computational throughput.
def section_7(): banner("7. AUTOTUNING — @tilelang.autotune") M = N = K = 2048 def configs(): out = [] for bm in (64, 128): for bn in (128, 256): for bk in (32, 64): for stages in (2, DEFAULT_STAGES): for thr in (128, 256): if smem_bytes(bm, bn, bk, stages) > SMEM_CAP: continue cfg = dict(block_M=bm, block_N=bn, block_K=bk, num_stages=stages, threads=thr) if cfg not in out: out.append(cfg) return out[:8] space = configs() print(f" searching {len(space)} configurations " f"(each one is a real nvcc compile + benchmark)") @tilelang.autotune(configs=space, warmup=10, rep=20) @tilelang.jit(out_idx=[-1]) def tuned_matmul(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), 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(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 a = torch.randn(M, K, device=DEV, dtype=torch.float16) b = torch.randn(K, N, device=DEV, dtype=torch.float16) t0 = time.time() from tilelang.autotuner import set_autotune_inputs with set_autotune_inputs(a, b): best = tuned_matmul(M, N, K) print(f" tuning took {time.time()-t0:.0f} s") c = best(a, b) check(c, a @ b, "autotuned matmul") ms = bench(lambda: best(a, b)) print(f" best kernel: {ms:.3f} ms -> {2*M*N*K/(ms*1e-3)/1e12:.2f} TFLOP/s") print(" Re-running this cell hits the on-disk autotuner cache and is instant.") print(" Turn caching off with TILELANG_AUTO_TUNING_DISABLE_CACHE=1.") We define an autotuning search space across matrix tile sizes, K-block dimensions, pipeline depths, and thread counts while filtering configurations that exceed the shared-memory budget. We use TileLang’s autotuning decorator to compile, benchmark, validate, and cache multiple kernel schedules for the same matrix-multiplication workload. We then execute the selected kernel, verify its output against 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 main(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 main def section_8(): banner("8. INTROSPECTION — T.print, sources, profiler") try: dbg = make_print_demo(128) a = torch.ones(128, device=DEV) b = torch.empty(128, device=DEV) dbg(a, b) torch.cuda.synchronize() print(f" T.print fired above; host check: b[0] = {b[0].item()} (expect 2.0)") except Exception as e: print(f" T.print demo skipped: {type(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())} lines of CUDA") print(" grep-worthy landmarks:") for needle in ("__global__", "extern "C"", "mma", "cp.async", "__syncthreads", "ldmatrix"): n = src.count(needle) if n: print(f" {needle:<16} x{n}") try: 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 own inputs)") except Exception as e: print(f"n profiler skipped: {type(e).__name__}: {e}") print("n Other tools worth knowing:") print(" kernel.get_host_source() - the launch wrapper") print(" T.device_assert(cond, msg) - device-side assertions") print(" TILELANG_DISABLE_CACHE=1 - force recompilation") print(" tilelang.tools.plot_layout - visualise fragment/shared layouts") print(" env TILELANG_CACHE_DIR - where compiled kernels live") CHEATSHEET = """ SCOPES T.alloc_shared(shape, dtype) __shared__ tile T.alloc_fragment(shape, dtype) registers, layout inferred T.alloc_var(dtype) one scalar per thread T.alloc_barrier(n) mbarrier (Hopper pipelines) MOVEMENT T.copy(src, dst) vectorized + casting move T.async_copy(src, dst) explicit 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=..., policy=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=k) software pipeline T.serial(n) plain sequential loop ANNOTATIONS T.use_swizzle(panel_size=, enable=) L2 rasterization T.annotate_layout({buf: layout}) explicit layouts T.annotate_l2_hit_ratio(buf, r) ENTRY POINTS @tilelang.jit(out_idx=[-1]) last arg is the return value @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 tricks examples/deepseek_mla MLA decode, ~80 lines examples/linear_attention RetNet / Mamba github.com/tile-ai/tilelang-puzzles 10 graded exercises 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 main(only=None): """only: optional list of 1-based section numbers, e.g. main([2, 6]).""" t_start = time.time() status = [] for idx, (name, fn) in enumerate(SECTIONS, start=1): if only and idx not in only: continue t0 = time.time() try: fn() status.append((name, "ok", time.time() - t0)) except Exception: print("n !! section failed — continuing with the restn") traceback.print_exc() status.append((name, "FAILED", time.time() - t0)) finally: torch.cuda.synchronize() banner("9. SUMMARY & CHEATSHEET") for name, st, dt in status: print(f" {st:>6} {name:<24} {dt:6.1f} s") print(f"n total: {time.time()-t_start:.0f} s") print(CHEATSHEET) if __name__ == "__main__": main() We introduce TileLang’s debugging and introspection workflow through device-side printing, generated CUDA inspection, and the built-in kernel profiler. We examine compiler-emitted landmarks such as tensor-core operations, asynchronous copies, synchronization barriers, and matrix-load instructions. We finally organize all tutorial sections into a fault-tolerant runner that records execution status, reports timing information, and prints a compact TileLang programming reference.
In conclusion, we built a practical understanding of how TileLang translates tile-level Python programs into optimized GPU kernels without requiring us to manually manage thread indices, warp-level data layouts, tensor-core instructions, or asynchronous memory barriers. We implemented and validated kernels that cover bandwidth-bound elementwise operations, compute-intensive GEMM workloads, fused neural-network epilogues, register-resident reductions, and online-softmax attention. We also examined how block dimensions, shared-memory consumption, pipeline depth, thread count, tile shape, and L2 swizzling influence performance across different GPU architectures. Finally, we used generated-source inspection, device-side debugging, profiling, and automated schedule search to establish a complete workflow for developing, verifying, benchmarking, and refining custom TileLang kernels.
Check out the Full Code here. Also, feel free to follow us on Twitter and don’t forget to join 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 partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us
Sana Hassan
Sana Hassan, a consulting intern at Marktechpost and dual-degree student at IIT Madras, is passionate about applying technology and AI to address real-world challenges. With a keen interest in solving practical problems, he brings a fresh perspective to the intersection of AI and real-life solutions.