Designing High-Performance GPU Kernels with TileLang: Tensor-Core GEMM, Fused Softmax, FlashAttention, and Autotuning
Key Takeaways
- •TileLang enables developers to author optimized GPU kernels at the tile level in Python while the compiler automatically handles thread mapping, memory layouts, synchronization, and CUDA instruction generation.
- •The tutorial progressively implements and validates kernels for vector addition, tensor-core matrix multiplication, fused GEMM epilogues with bias and GELU, row-wise softmax, and FlashAttention against PyTorch baselines.
- •TileLang's autotuning decorator searches across tile sizes, pipeline depths, and thread counts to automatically identify optimal architecture-dependent configurations, addressing the challenge that ideal parameters vary across GPU generations like Ampere and Hopper.
- •Kernel fusion techniques demonstrated in the tutorial reduce intermediate global-memory traffic by completing operations such as bias addition and GELU activation within register-resident accumulators before writing final outputs.
- •The FlashAttention implementation processes query, key, and value tiles without materializing the full attention-score matrix, reducing memory footprint from quadratic to linear space using online softmax updates.
- •Generated-source inspection, device-side printing, and a built-in profiler together provide developers with visibility into compiler-emitted tensor-core operations, asynchronous copies, and synchronization barriers for debugging and optimization.

TileLang is a high-level Python domain-specific language (DSL) for designing and compiling performance-oriented GPU kernels through TVM. As large language models and other transformer architectures push computational demands higher, the ability to author custom GPU kernels that fully utilize tensor cores—NVIDIA's dedicated matrix-multiply accelerators—has become increasingly critical for training and inference efficiency. However, writing such kernels traditionally requires deep expertise in CUDA C++, warp-level programming, and hardware-specific memory hierarchies. TileLang addresses this gap by letting developers express computations at the tile level while the compiler handles the low-level details. This tutorial provides a comprehensive walkthrough of TileLang's capabilities, beginning with environment validation and the establishment of reusable benchmarking and numerical-verification utilities. From there, it progressively implements vector addition, tiled tensor-core matrix multiplication, schedule exploration, fused GEMM epilogues, row-wise softmax, and FlashAttention.
Throughout the tutorial, developers work directly with TileLang's shared-memory tiles, register fragments, pipelined loops, parallel iteration primitives, reductions, and tensor-core GEMM operators. The compiler handles thread mapping, memory layouts, synchronization, vectorization, and low-level CUDA instruction generation automatically. Kernels are benchmarked against PyTorch and cuBLAS baselines, with inspection of generated CUDA source, evaluation of memory and compute throughput, and autotuning to identify architecture-dependent kernel configurations. TileLang is available on GitHub.
The tutorial begins by configuring the Google Colab CUDA environment, installing TileLang with a nightly fallback, and importing the required PyTorch and TileLang modules. Reusable benchmarking, validation, and reporting utilities are defined to measure kernel latency and compare numerical outputs using relative error. A TileLang vector-add kernel is then implemented, executed on the GPU, compared against PyTorch's bandwidth, and inspected for the CUDA source generated by the compiler.
Next, a tiled tensor-core matrix multiplication kernel is implemented, moving input tiles through global memory, shared memory, and register fragments. Tile dimensions, pipeline stages, thread counts, and L2 swizzling are controlled manually, while TileLang generates tensor-core instructions, synchronization, and memory-transfer logic. Several schedule configurations are benchmarked, their numerical accuracy verified, and the highest-performing architecture-dependent kernel configuration identified. The full tutorial code is available here.
The matrix multiplication kernel is then extended by fusing bias addition and the GELU activation directly into the register-resident accumulator. This approach reduces intermediate global-memory traffic by completing the epilogue before writing the final output tensor. Kernel fusion like this is a well-established technique for reducing memory bandwidth bottlenecks, which often dominate latency in large-scale neural network workloads. The fused implementation is compared against eager PyTorch execution. Additionally, a row-wise softmax kernel is implemented using fragment-level maximum and sum reductions, keeping the normalization process largely within registers.
A fused FlashAttention forward kernel is implemented to process query, key, and value tiles without materializing the full attention-score matrix in global memory. Online softmax updates are applied using running maxima, normalization sums, rescaling factors, and tiled tensor-core matrix multiplications. FlashAttention, originally introduced by researchers at Stanford, has become a widely adopted technique for reducing the memory footprint of self-attention from quadratic space to linear space, and it is now integrated into major frameworks including PyTorch's scaled dot-product attention API. Both causal and non-causal attention are validated against PyTorch scaled dot-product attention, with latency and computational throughput compared.
An autotuning search space is defined across matrix tile sizes, K-block dimensions, pipeline depths, and thread counts, with configurations exceeding the shared-memory budget filtered out. TileLang's autotuning decorator is used to compile, benchmark, validate, and cache multiple kernel schedules for the same matrix-multiplication workload. The selected kernel is executed, verified against PyTorch, and evaluated for achieved latency and tensor-core throughput. This automated search addresses a practical challenge in GPU kernel development: optimal tile sizes and pipeline depths vary across GPU architectures such as Ampere and Hopper, making manual tuning brittle when targeting multiple hardware generations.
TileLang's debugging and introspection workflow is introduced through device-side printing, generated CUDA inspection, and the built-in kernel profiler. Compiler-emitted landmarks are examined, including tensor-core operations, asynchronous copies, synchronization barriers, and matrix-load instructions. All tutorial sections are organized into a fault-tolerant runner that records execution status, reports timing information, and prints a compact TileLang programming reference.
The tutorial demonstrates how TileLang translates tile-level Python programs into optimized GPU kernels without requiring manual management of thread indices, warp-level data layouts, tensor-core instructions, or asynchronous memory barriers. The implemented and validated kernels cover bandwidth-bound elementwise operations, compute-intensive GEMM workloads, fused neural-network epilogues, register-resident reductions, and online-softmax attention. The exploration also shows how block dimensions, shared-memory consumption, pipeline depth, thread count, tile shape, and L2 swizzling influence performance across different GPU architectures. Generated-source inspection, device-side debugging, profiling, and automated schedule search together establish a complete workflow for developing, verifying, benchmarking, and refining custom TileLang kernels.