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

TileLang, a high-level Python domain-specific language for designing GPU kernels through TVM, enables developers to implement tensor-core GEMM, fused softmax, FlashAttention, and autotuning while managing shared-memory tiles, register fragments, and pipelined loops. The tutorial validates CUDA environments, benchmarks kernels against PyTorch and cuBLAS baselines, and uses autotuning to identify architecture-dependent configurations.

In this tutorial, we explore TileLang https://github.com/tile-ai/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. python 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. python @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. python @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. python @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. python 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. python @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 rest\n" 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 and don’t forget to join our Twitter https://x.com/intent/follow?screen name=marktechpost and Subscribe to 150k+ML SubReddit https://www.reddit.com/r/machinelearningnews/ . Wait are you on telegram? our Newsletter https://www.aidevsignals.com/ now you can join us on telegram as well. https://t.me/machinelearningresearchnews Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us https://forms.gle/wbash1wF6efRj8G58 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.