TileLang: Write, Autotune, and Ship Fused Tensor‑Core Kernels from Python

When a few dozen lines of Python talk to tensor cores

TileLang is a Python DSL (built on TVM) that lets you express kernels at the tile/fragment level and have the toolchain emit low‑level CUDA (tensor‑core intrinsics, cp.async, ldmatrix, etc.). The practical payoff is simple: keep reductions and accumulators in registers, fuse epilogues so you avoid M×N round trips to HBM (GPU high‑bandwidth memory), and autotune the tiling and pipeline parameters so kernels actually hit the hardware sweet spot.

Here’s a compact, pragmatic guide for engineering leads and performance engineers: what you need to run the examples, what the tutorial shows, how to validate correctness and performance, and the operational work to do before shipping tuned kernels.

What you’ll see and what it means

  • Tile primitives: alloc_shared (shared tiles), alloc_fragment (register fragments), T.Pipelined(…) for software pipelines, and T.gemm for tile‑level tensor‑core matmul.
  • Small, testable kernels: vector add (bandwidth bound), tiled tensor‑core GEMM, fused GEMM epilogue (bias + GELU), row‑wise softmax and reductions, and a fused FlashAttention forward.
  • Autotuning: compile and benchmark many schedules (tile sizes, pipeline depth, threads), filter by shared‑memory budget, and cache best kernels on disk (for example, ~/.tilelang/cache).
  • Introspection: device‑side printf (T.print), kernel.get_kernel_source() to inspect emitted CUDA, and kernel.get_profiler().do_bench() for microbenchmarks.

Practical defaults the tutorial uses (for quick reference)

  • bench helper: warmup=10, rep=50; timing is median-ish latency in milliseconds and measured with CUDA events.
  • Numerical checks use a relative Frobenius‑norm with tol = 2e-2 (the tutorial comment: “Relative-Frobenius-norm check. Far more meaningful than atol for fp16.”).
  • Device detection: the notebook queries torch.cuda.get_device_properties(0) and get_device_capability(0), and encodes compute capability as SM = CC[0]*10 + CC[1].
  • Shared‑memory heuristic: SMEM_CAP = 48 * 1024 bytes if SM < 80 else 96 * 1024 bytes; DEFAULT_STAGES = 2 if SM < 80 else 3, a tutorial heuristic, not a universal rule.
  • Typical matmul defaults: block_M=128, block_N=128, block_K=32, num_stages=3, threads=128, dtype=float16, accum_dtype=float.

Prerequisites, what you need to check before you run anything

  • GPU: an NVIDIA GPU with tensor cores (the examples target modern Ampere/Hopper‑class behaviors). Different GPU generations have different shared‑memory and register resources; the tutorial uses a small SM heuristic, verify it on your hardware.
  • CUDA / drivers: a CUDA toolkit and driver compatible with your GPU. The tutorial invokes the CUDA toolchain during compilation (nvcc or NVRTC depending on build). Check the TileLang repository README for exact supported versions before you run large autotunes.
  • Python & PyTorch: the notebook relies on torch for device querying, correctness baselines (a + b, softmax, scaled_dot_product_attention) and easy GPU transfers. Use a recent PyTorch that matches your CUDA runtime.
  • TileLang install: the tutorial suggests “Stable wheel first, nightly as a fallback.” If the stable pip wheel is not available, the notebook falls back to the nightly index (pip -f https://tile-ai.github.io/whl/nightly …). Always confirm the exact wheel or commit in the repo you run.
  • Build toolchain: compilation of kernels typically invokes nvcc; ensure nvcc is on PATH (or follow the notebook’s CUDA_HOME logic that adds /usr/local/cuda/bin when present).

Device capabilities: what to inspect and why

When you query the device, look at these concrete fields (via PyTorch):

  • torch.cuda.get_device_properties(0).major and .minor (compute capability)
  • torch.cuda.get_device_properties(0).total_memory (global memory)
  • For shared memory per SM use cudaDeviceGetAttribute or the equivalent driver API to read sharedMemPerMultiprocessor if you need absolute guarantees.

The tutorial uses SM = CC[0]*10 + CC[1] as a compact encoding and derives SMEM_CAP and DEFAULT_STAGES from it. Treat that as a convenience heuristic and validate it on your exact GPU (H100/A100/T4 have different SM and shared memory behaviors and partitioning options).

Why tile‑level control matters (and what to watch for)

Two performance rules dominate on modern NVIDIA GPUs:

  • Keep accumulators in registers. Using T.alloc_fragment to hold the per‑tile accumulator keeps the expensive reduction work register‑resident and avoids spills to shared or global memory.
  • Reduce HBM trips. Fuse the epilogue (bias add, activation such as GELU) so you perform those ops while the accumulator is in registers and then write final results once. That avoids writing and re‑reading intermediate M×N tiles from HBM.

That is why the tutorial implements a fused GEMM epilogue and reports HBM traffic saved for example shapes. The tutorial prints a framed result for the epilogue run (example output shown literally in the notebook):

” fused (1 kernel) : {ms_f:7.3f} ms”, ” torch (3 kernels) : {ms_e:7.3f} ms”, ” speedup : {ms_e/ms_f:5.2f}x”, ” HBM traffic saved : ~{2*M*N*2/2**20:.0f} MiB of intermediate reads/writes”

That output shows the effect: fusion can give multi‑x speedups on memory‑bound workloads. Caveats: fusion increases register pressure and can change rounding order, so numerical validation and register and shared‑memory monitoring are mandatory.

FlashAttention: the FLOP model and what the example measures

FlashAttention‑style kernels compute softmax online and avoid materializing the full S×S attention matrix. The tutorial’s forward attention FLOP model is stated as:

flops = 4 * B * H * S * S * D * (0.5 if causal else 1.0)

Here’s how that factor breaks down (why 4×):

  • Q⋅Kᵀ matmul: roughly 2 * B * H * S * S * D FLOPs (multiply and add)
  • softmax (scale, exp, div) and elementwise ops: a few ops per element, the tutorial aggregates these into a modest constant factor
  • softmax⋅V matmul: another ~2 * B * H * S * S * D FLOPs

The tutorial lumps the matmuls and minimal softmax cost into the 4× factor and halves the count for causal because causal attention skips half the pairs on average. When you report TFLOP/s, state the exact FLOP model you used, the tutorial prints the formula with the benchmark outputs so readers know what’s included.

A practical snippet from the tutorial notes: “~70 lines of Python for a fused, causal, tensor-core attention.” The notebook validates its fused forward against torch.nn.functional.scaled_dot_product_attention for both causal and non‑causal cases.

Autotuning: how it works and how much it costs

Autotuning in the tutorial enumerates candidate schedules (bm, bn, bk, stages, threads), filters configurations that exceed smem_bytes, and compiles and benchmarks each candidate. The smem_bytes formula used by the tutorial is:

smem_bytes(block_M, block_N, block_K, stages, itemsize=2) = (block_M * block_K + block_K * block_N) * itemsize * stages

The notebook prints the tuning action with a frank line:

” searching {len(space)} configurations (each one is a real nvcc compile + benchmark)”

Practical implications:

  • Each new candidate typically triggers a compile step and an execution measurement on the GPU. Depending on your machine, that can be tens of seconds to minutes per candidate, so wide grids can take minutes to hours on a single GPU.
  • TileLang caches compiled kernels to disk (for example, ~/.tilelang/cache). Re‑running a previously compiled candidate hits the cache and is fast; disable with TILELANG_AUTO_TUNING_DISABLE_CACHE=1 when you need fresh compiles.
  • To manage cost, prune the search with the shared‑memory filter, seed with reasonable defaults, run a narrow search first, and persist the cache in a shared volume or baked image so CI and cluster runs reuse artifacts.

Debugging and introspection you’ll actually use

  • kernel.get_kernel_source() returns the generated CUDA; the tutorial searches that source for tokens such as “mma.sync”, “ldmatrix”, “cp.async”, and “__syncthreads” as heuristics that specialized instructions and async copies were emitted.
  • T.print emits a guarded device‑side printf. A noted caveat in the notebook: “T.print emits a guarded device-side printf. Keep the grid tiny.”
  • Use kernel.get_profiler().do_bench() to microbenchmark the compiled kernel under the same runtime conditions the autotuner uses.

Important diagnostic checklist:

  • Inspect kernel.get_kernel_source() and confirm the compiled PTX targets your desired SM (check compute capability in the build notes).
  • Run the kernel on small shapes and verify numeric equivalence against PyTorch baselines (relative Frobenius and elementwise checks, see below).
  • Measure register usage and shared memory consumption in profiler outputs. If you see spills, reduce tile sizes or remove fusion.

Numerical validation, concrete steps

FP16 code gives speed but requires careful validation. The tutorial’s baseline checks are useful; expand them into a short validation suite:

  • Forward correctness: compare outputs to a PyTorch reference using relative Frobenius norm (tutorial default tol = 2e-2). For small matrices also check elementwise max absolute error.
  • Backward correctness for training: run gradient checks (finite differences) on small shapes; verify gradients match within tolerance and that training loss decreases on a small synthetic task.
  • Accumulator sanity: consider float32 accumulators while prototyping. The tutorial uses accum_dtype=”float” for many kernels.
  • Reproducibility: fix random seeds, run multiple times, and ensure numerical variation is acceptable for your application. Fusion changes operation order and can affect reproducibility.

Suggested validation tolerances (starting points): relative Frobenius ≤ 2e‑2 for large fp16 matrices; tighten for critical models or use float32 accumulators where necessary. Always fail the build or the integration test if a tuned kernel exceeds your project’s numeric thresholds.

Compact cheat sheet (what you’ll reach for)

  • Memory: T.alloc_shared(shape, dtype), T.alloc_fragment(shape, dtype), T.alloc_var(dtype), T.alloc_barrier(n)
  • Movement: T.copy (sync), T.async_copy (cp.async-like), T.tma_copy (large transfers)
  • Compute: T.gemm (tile matmul), T.reduce_sum/max, T.clear/T.fill
  • Loops & pipelines: T.Parallel(…), T.Pipelined(iter_count, num_stages=k), T.serial(n)
  • Annotations: T.use_swizzle(…), T.annotate_layout, T.annotate_l2_hit_ratio
  • Entry points & tuning: @tilelang.jit(out_idx=[-1]), @tilelang.autotune(configs=…), kernel.get_kernel_source(), kernel.get_profiler().do_bench()
  • Env control: TILELANG_CACHE_DIR, TILELANG_DISABLE_CACHE, TILELANG_AUTO_TUNING_DISABLE_CACHE

Key questions a curious engineering lead will ask

  • How do I know the compiler actually targeted tensor cores and async copies?

    Inspect kernel.get_kernel_source() and look for tokens such as “mma.sync”, “ldmatrix”, and “cp.async”. These are heuristics indicating the codegen emitted tensor‑core and async copy sequences. Also confirm the compiled binary targets the correct compute capability and validate performance against cuBLAS/PyTorch baselines.

  • How much does autotuning cost?

    Each new candidate in the tutorial typically incurs a compile and a benchmark (nvcc or NVRTC), so expect per‑candidate times from tens of seconds to minutes depending on hardware. Wide searches can therefore take minutes to hours on a single GPU. Prune the space, persist caches (for example, ~/.tilelang/cache), and bake tuned artifacts into images to amortize cost.

  • What speedup should I expect from epilogue fusion?

    The tutorial’s fused GEMM plus bias and GELU example prints a multi‑x speedup versus a three‑kernel eager PyTorch sequence and reports HBM traffic saved for the example shapes. Exact numbers depend on shape, dtype, and device. Treat the example as an existence proof and run the same microbenchmarks on your target hardware to quantify gains.

  • Are fused FP16 kernels safe for training?

    They can be, but you must validate. The tutorial uses a relative Frobenius check (tol ≈ 2e‑2). For production training add gradient finite‑difference checks, consider float32 accumulators during development, and keep unfused fallbacks for debugging.

  • Does the tutorial cover FlashAttention backward pass?

    The notebook demonstrates a fused forward FlashAttention and validates it against torch.nn.functional.scaled_dot_product_attention. Backward training implementations exist elsewhere in the examples repository; deploying training‑grade fused attention requires implementing and validating the backward pass as well.

Prioritized checklist before you ship tuned kernels

  1. Run a small representative shape on the target GPU, call kernel.get_kernel_source(), and verify emitted tokens (mma.sync, cp.async, ldmatrix) and correct outputs against PyTorch.
  2. Autotune a narrow grid first; record per‑candidate compile and bench times and budget the full tuning run. Persist the cache to a shared volume or bake compiled artifacts into your container image.
  3. Validate numerics: forward relative Frobenius, elementwise max for small dims, gradient finite‑differences for backward correctness. Fail CI on regressions.
  4. Monitor register and shared memory usage; if fusion causes spills, reduce tile sizes or peel back fusion.
  5. Package tuned kernels (cache and metadata) with versioned artifacts and a reproducible build manifest so production runs do not re‑tune on ephemeral nodes.

Operational caveats and final advice

TileLang gives you fine, explicit control over tiling, staging and register fragments without sinking into raw CUDA assembly. That puts high‑performance kernels within reach of teams that prefer Python and want compiler help. The hard work is not learning the DSL; it is validating numerics, managing autotune cost, and baking tuned artifacts into deployment pipelines so customers see reliable, reproducible speedups.

Keep these guardrails in mind:

  • Consider the tutorial’s SMEM_CAP and DEFAULT_STAGES as convenience heuristics, test and adapt them to your exact GPU.
  • Always include unfused fallbacks for correctness debugging, fusion changes operation ordering and can subtly alter results.
  • Persist compiled kernels across CI and production to avoid repeat compiles in ephemeral environments (TILELANG_CACHE_DIR). If you must disable caching for debugging, use TILELANG_AUTO_TUNING_DISABLE_CACHE=1.
  • When you report TFLOP/s, publish the full FLOP model and the hardware and software stack (GPU model, CUDA/toolkit, driver, TileLang commit). Performance numbers without that context are hard to act on.

TileLang sits in a sweet spot: enough explicit tile control to target tensor cores and async copies, but high‑level enough to be written in dozens of lines of Python. Use the primitives, autotune wisely, validate numerics, and operationalize the cache. Do that, and those few dozen lines will whisper high throughput to your GPUs.

0