Source-linked AI summary
TileLang: A Composable Tiled Programming Model for AI Systems
Lei Wang, Yu Cheng, Yining Shi, Zhengju Tang, Zhiwen Mo, Wenhao Xie, Lingxiao Ma, Yuqing Xia, Jilong Xue, Fan Yang, Zhi Yang
TL;DR
Writing high-performance AI kernels remains difficult because hardware-specific scheduling and optimization details are complex, while existing DSLs leave usability and expressiveness gaps. TileLang separates dataflow from scheduling through composable tile operators, scheduling primitives, and compiler automation. Across real-world kernels on NVIDIA and AMD GPUs, it achieves performance comparable to or sometimes exceeding specialized vendor libraries and other DSL-based approaches.
Problem
High-performance AI kernels follow clear tile-based dataflow but remain difficult to write because developers must manually manage hardware-specific optimizations and existing DSLs can hide critical controls.
Method
TileLang uses a Python-like tiled programming model that separates dataflow from scheduling, combines composable tile operators with scheduling annotations, and automates layout, parallelization, and pipeline optimizations.
Results
TileLang achieves performance comparable to, and sometimes exceeding, specialized vendor libraries and other DSL-based approaches across NVIDIA and AMD GPUs.
Takeaways & Limitations
TileLang provides a unified block-and-thread programming approach with transparent scheduling capabilities for balancing compiler automation and expert performance control.
Takeaways & Limitations
TileLang currently depends on CUTLASS and manually wrapped CUDA/HIP code for built-in operators, with self-hosting planned as future work.
Abstract
from arXiv · showhide
Modern AI workloads rely heavily on optimized computing kernels for both training and inference. These AI kernels follow well-defined data-flow patterns, such as moving tiles between DRAM and SRAM and performing a sequence of computations on those tiles. However, writing high-performance kernels remains complex despite the clarity of these patterns. Achieving peak performance requires careful, hardware-centric optimizations to fully leverage modern accelerators. While domain-specific compilers attempt to reduce the burden of writing high-performance kernels, they often struggle with usability and expressiveness gaps. In this paper, we present TileLang, a generalized tiled programming model for more efficient AI Kernel programming. TileLang decouples scheduling space (thread binding, layout, tensorize and pipeline) from dataflow, and encapsulated them as a set of customization annotations and primitives. This approach allows users to focus on the kernel's data-flow itself, while leaving most other optimizations to compilers. We conduct comprehensive experiments on commonly-used devices, across numerous experiments, our evaluation shows that TileLang can achieve state-of-the-art performance in key kernels, demonstrating that its unified block-and-thread paradigm and transparent scheduling capabilities deliver both the power and flexibility demanded by modern AI system development.
1 INTRODUCTION
TileLang addresses the difficulty of writing high-performance AI kernels by separating kernel dataflow from scheduling and combining compiler automation with fine-grained user control. It reports performance comparable to or exceeding specialized libraries and other DSLs across NVIDIA and AMD GPUs.
- Motivation: AI kernels use tile-based dataflow, but high performance requires manually handling thread binding, memory layout, and hardware-specific instructions.These optimizations affect parallelism, memory access, tensorization, and accelerator utilization.
- Motivation: Existing DSLs simplify kernel programming but can hide thread behavior, memory layouts, and address-space details needed for performance-critical kernels.The paper identifies quantized-weight matrix multiplication as an example requiring vectorized datatype conversion and hardware-aligned layouts.
- TileLang Approach: TileLang decouples dataflow from scheduling, letting users express composable tile operators while compilers explore scheduling strategies.Scheduling primitives and annotations remain available for manually fine-tuning performance-critical aspects.
- TileLang Approach: TileLang provides a Python frontend and compiler passes for layout inference, automatic parallelization, pipeline derivation, parameter simplification, and loop-tail splitting.These capabilities target both ease of writing and efficient low-level code generation.
- Evaluation: TileLang achieves performance comparable to, and sometimes exceeding, specialized vendor libraries and DSL-based approaches across NVIDIA and AMD GPUs.The evaluation covers real-world AI kernels.
2 A TileLang Example
The TileLang example illustrates Python-based tile programming that exposes dataflow and memory placement while compiling programs through intermediate representation into executable optimized code.
- Example Context: TileLang is designed around Pythonic programming while making dataflow more transparent than approaches that obscure optimization details.The paper contrasts this design with systems whose dataflow or scheduling remains difficult to inspect and modify.
- GEMM Example: A GEMM kernel example defines inputs, outputs, shapes, data types, execution context, on-chip memory, and dataflow using Python constructs.TileLang supports imperative constructs such as if-else, for, and while, with explicit type annotations for device code generation.
- GEMM Example: TileLang exposes tiles, memory placement, pipelining, and operator calls so developers can manage data movement and computation with fine-grained control.The example uses global, shared, and register memory hierarchies to address bandwidth utilization and latency.
- Compilation: Automating low-level scheduling and synchronization lets developers focus on algorithm design while maintaining computational efficiency.The example is positioned as a productivity aid for tile-level matrix computation.
- Compilation: The tilelang.compile call lowers the program into an intermediate representation and then into executable optimized code.The paper presents the source program, lowered IR, and generated CUDA C code as successive stages.
3 The Tile Language
TileLang separates tile-level dataflow operators from scheduling primitives, combining a data-centric programming model with fine-grained performance control. Its compilation pipeline and hardware-aware abstractions support optimized data movement, computation, and execution.
- Compilation Pipeline: TileLang programs pass through parsing, IR building, optimization, and code generation stages.The pipeline transforms Python programs into TileLang AST, TVM IR, optimized IR, and backend code.
- Dataflow Operators: TileLang expresses computational semantics through tile-level operators such as T.copy, T.gemm, and T.reduce.These operators form the dataflow layer of the programming model.
- Scheduling Primitives: Scheduling primitives expose control over parallelism, pipelining, and memory layout independently of dataflow operators.This separation lets developers tune performance-critical execution aspects while retaining high-level computational semantics.
- Tile-based Programming Model: Tiles are shaped data regions owned and manipulated by warps, thread blocks, or equivalent parallel units.T.Kernel provides execution context such as thread-block indices and thread counts for automatic inference and manual thread control.
- Memory Allocation: TileLang explicitly places tile buffers in hardware memory spaces through user-facing allocation intrinsics.The model exposes physical memory placement rather than relying entirely on opaque compiler decisions.
- Tile Operators: TileLang operators abstract low-level tile implementation while exposing Lower and InferLayout interfaces for lowering and layout inference.Representative operators include GEMM, Copy, and Parallel.
4 Scheduling Design and Automation
TileLang separates dataflow from scheduling and automates memory-layout and thread-binding decisions through composable layout abstractions and priority-based inference. It also supports hardware-specific instruction paths while exposing trade-offs between control, portability, and compilation cost.
- TileLang organizes scheduling design around memory layout, thread binding, tensorization, and pipeline automation alongside dataflow.
- 4.1 Memory Layout Composition: Layout functions map iteration variables to physical memory indices, supporting composition, padding, and swizzling for shared-memory access.
- 4.1 Memory Layout Composition: Fragment layouts partition block-level register files across threads, representing each thread’s register position and local-register index.
- 4.1 Memory Layout Composition: TileLang derives block layouts compositionally: a m16k16 warp layout becomes m32k16, then a four-warp m128k16 block layout.
- 4.2 Thread Binding: Layout inference resolves buffer layouts and thread bindings top-down by operator priority, aligning flexible buffers with stricter operators such as GEMM.
- 4.3 Leveraging High-Performance Hardware Instructions: TileLang supports hardware instructions through library or direct implementations, but template-based libraries can limit layout control and increase compilation time.
5 Numerical Experiments
TileLang is evaluated across diverse AI operator workloads, GPU architectures, and state-of-the-art baselines. The results show competitive or superior performance across attention, GEMM, latent attention, and dequantized matmul kernels.
- Evaluation setup: TileLang is evaluated on NVIDIA H100, A100, and AMD MI300X GPUs across representative AI operator workloads.The workloads include attention variants, GEMM, and dequantized GEMM kernels.
- Attention performance: TileLang achieves speedups of 1.36×, 1.41×, and 1.70× over FlashAttention-3, Triton, and PyTorch, respectively.Its performance remains close to FlashAttention-3 for longer sequence lengths such as 8k.
- Attention performance: TileLang achieves average speedups of 1.77× and 2.10× over Triton in the chunk-scan and chunk-state Linear Attention experiments.
- MLA performance: On H100, TileLang reaches 98% of FlashMLA performance with around 70 lines of Python code, while on MI300X it reaches 95% of AITER performance.The H100 result includes a 1075.9× speedup over Torch; the MI300X result includes a 129.2× speedup over Torch.
- GEMM performance: Across RTX 4090, A100, H100, and MI300X, TileLang achieves vendor-library speedups of 1.10×, 0.97×, 1.00×, and 1.04×, respectively.Against Triton, the corresponding speedups are 1.08×, 1.03×, 1.13×, and 1.25×.
- Dequantized matmul: TileLang achieves a maximum 7.65× speedup over cuBLAS-WFP16AFP16 and average speedups of 1.04× over Marlin and 1.62× over BitsandBytes.The results cover WINT2AINT8, WINT4AFP16, and WNF4AFP16 configurations.
6 Conclusions and Discussions
TileLang combines Python-based tiled programming with automated optimization and optional fine-grained control over memory layouts, buffers, threads, and pipelines. The paper identifies self-hosting, distributed execution, cost modeling, dynamic-shape tuning, and broader backend support as future directions.
- Conclusions and Discussions: TileLang lets users program at tile granularity while explicitly declaring memory-hierarchy buffers and optionally controlling individual thread behavior.
- Conclusions and Discussions: TileLang abstracts optimization details such as pipelining by default while retaining frontend support for explicit pipeline implementation.It also supports dynamic parameters, dynamic shapes, and other advanced features for kernel libraries.
- Future work: Future work includes a self-hosting Tile Library to remove current dependencies on CUTLASS and manually wrapped CUDA/HIP code.
- Future work: Future directions include tile-level communication primitives, a TileLang cost model, and dynamic-shape tile-configuration tuning.
- Future work: The authors plan to extend TileLang across CPUs, NPUs, and other hardware platforms through generalized multibackend support.
- Conclusions and Discussions: TileLang is open-sourced to support future development and community contributions.
A Operator shapes in our benchmark
The benchmark appendix lists matrix, FlashAttention, and Linear Attention shapes used in the evaluation.
- Matrix shapes: Table 2 lists the matrix shapes used in the benchmark.
- FlashAttention shapes: Table 3 lists the FlashAttention shapes used in the benchmark.
- Linear Attention shapes: Table 4 lists the Linear Attention shapes used in the benchmark.
B.1 Matrix Multiplication (Matmul)
Figure 16 presents a kernel implementation of matrix multiplication.
- Matrix Multiplication (Matmul): Figure 16 shows the kernel implementation of matrix multiplication.
- Matrix Multiplication (Matmul): The figure concerns matrix multiplication kernel implementation.
- Matrix Multiplication (Matmul): The appendix includes a visual representation of a matrix multiplication kernel.
B.2 Dequantized Matrix Multiplication
TileLang is used to implement a weight-only quantized matrix multiplication with mixed-precision operands.
- B.2 Dequantized Matrix Multiplication: TileLang implements weight-only quantization for a matrix multiplication.The figure specifies the operand formats W_FP4_E2M1 and A_FP16.
- B.2 Dequantized Matrix Multiplication: The illustrated matrix multiplication combines FP4_E2M1 weights with FP16 activations.
- B.2 Dequantized Matrix Multiplication: The implementation showcases TileLang support for mixed-precision computations through a simple form.
B.3 FlashMLA Implementation
The section presents an implementation of FlashMLA with TileLang.
- B.3 FlashMLA Implementation: TileLang is used to implement FlashMLA.
- B.3 FlashMLA Implementation: The figure documents the implementation of FlashMLA using TileLang.
- B.3 FlashMLA Implementation: The supplied section materials include a figure describing FlashMLA implementation with TileLang.