TileLang: Bridge Programmability and Performance in Modern Neural Kernels
Authors: Lei Wang
Yu Cheng
Yining Shi
Zhiwen Mo
Zhengju Tang
Wenhao Xie
Tong Wu
Lingxiao Ma
Yuqing Xia
Jilong Xue
Fan Yang
Zhi Yang
Venue: ICLR 2026 Oral (2026)
Keywords: compiler; AI; programming model
Primary area: infrastructure, software libraries, hardware, systems, etc.
TL;DR: We introduce TileLang, a controllable programming system for fused neural kernels.

Abstract: Modern AI algorithms increasingly adopt fused kernels for performance, but implementing them remains complex due to the lack of fine-grained control in existing compilers like Triton. We introduce TileLang, a controllable programming system for fused neural kernels. TileLang provides explicit tile-level primitives for memory placement, data movement, and parallel scheduling. To guide developers in hardware-aware programming, the TileLang introduces two key techniques: tile inference which models tile programs as fused graphs and automatically deduces tile configuration from partial annotations; and tile recommendation that suggests efficient tile configurations based on hardware profiles and heuristics. TileLang makes it easy to express a wide range of fused attention kernels in under 80 lines of Python code, reducing code size by up to 90% compared to manual implementations. Evaluations show that TileLang achieves up to 5x speedup over Triton on NVIDIA H100 and up to 6 on AMD GPUs, demonstrating its ability to bridge programmability and performance.

Published as a conference paper at ICLR 2026
TILELANG: BRIDGE PROGRAMMABILITY AND PERFORMANCE IN MODERN NEURAL KERNELS
Lei Wang', Yu Cheng', Yining Shi', Zhiwen Mo², Zhengju Tang', Wenhao Xie Tong Wu', Lingxiao Ma', Yuqing Jilong Xue', Fan Yang", Zhi Yang".
Peking University Imperial College London TileAl *Microsoft Research
Sangzhi@pku.edie en, tile-ai@outlcok.com
ABSTRACT
Achieving high performance in modern Al increasingly requires kernels co- designed with underlying hardware, but writing efficient kernels remains challeng- ing due hardware-level complexity and limited fine- mne-grained control in compilers like Triton. In this paper, introduce TILELANG, programmable tile-level sys- tem that provides explicit primitives for memory placement, data movement, and parallel scheduling. Using unified fused tile-level dataflow graph (FTG), TILE- LANG streamlines kernel development by unifying tile recommendation, which guides developers with hardware-aware defaults, and tile inference, which auto- mates completion through constraint propagation. TILELANN enables concise expression of wide range of Al algorithms in fewer than 70 lines of Python, reducing code size by up to 85, 5,5% compared with manual implementations. Our evaluation shows that TILELANG delivers 1.08x-100 108x-10.58x speedups over Triton on NVIDIA H100 x average) and 1.01 x- 1.56x on AMD GPUs (2.65 on average), effectively bridging programmability and performance. Our code available at phtttp:: http:///ggggoooommeeeaaaaaaaaaaa ait.ileelllltennnnnggnn
INTRODUCTION
The rapid progress of modern neural networks has driven growing demand for highly optimized compute kernels, particularly for memory-bound operations such as attention.
In recent years, modern attention algorithms
such Multi-Head Attenttin(MHHA) (aawwani
MLA Performance and Code Complexity H100
(ct al.) 2017). Multi-Head Latent Attention (MLA) (Liu et al.) 2024). Gated Query Atten-
Speedup Nzrralizedd to PyTorch
10a
841.26x
1040.56x
tion (GQA) (Ainslie et al. 2023), and Linear
Attention (Gu & Dao 2023 Dao & Gul 2024
151.19x-
Sun et al. 2023) Yang et al. 2024). increas- ingly demand fine-grained control over memory
10a
TileLang
hierarchy, scheduling, and data movement to fully utilize hardware capabilities. However, ex- isting systems like Triton (Tillet et al.) 2019)
10
PyTorch
Triton
FlashMLA
1.00x
lack programmable abstractions to support this level of control. For instance, FlashMLA relies on carefully pipelined computations and shared memory reuse, but Triton gives programmers no direct control over tile reuse or pipeline schedul- ing, restricting performance optimization.
10a
100
200
300
400
500
Code Lines (LOC)
Performance code size trade-off for MLA kernels on NVIDIA H100. Points closer to the top-left indicate better balance between perfor- mance and implementation simplicity. The anno- tated speedup values indicate performance gains over the Py implementation. manually write complex CUDA kernels sacrifice significant performance due to abstraction
Figure 1:
As result, developers often face a steep trade- off between achieving peak performance and maintaining programmability: they must either
"Corresponding author.

---

Published as a conference paper at ICLR 2026
mismatches. As illustrated in Figure the Triton implementation of MLA requires only 130 lines of code. whose convenience comes at a steep cost itts performance reaches only 14.2% of the wardtewittenn CUDA version (DeepSeek] 2025) (~500 lines) on NVIDIA H100 GPUs. Bridging the gap between programmability and performance requires addressing two key challenges. First, programming model must give developers precise control over data movement and computation, enabling direct interaction with hardware resources. Second, compiler must efficiently lower these high-level programs to GPU code, mapping abstractions onto hardware resources without adding programming complexity. Solving both challenges is essential to balance developer productivity with near-peak hardware performance on modern accelerators.
We introduce TILELANG, a controllable programming system for modern neural workloads. TILE- LANG provides programmable tile abstractions that developers express and optimize low-level kernel behaviors in a high-level, composable way. Unlike existing compilers such as Triton, which rely on opaque optimization passes, TILELANG gives developers explicit control over memory, data movement, layout, and parallel execution. Specifically, developers can allocate buffers in differ- ent hardware memory levels (alloc_shared, alloc_ fragment), orchestrate data transfers (copy), define custom memory layouts (annotate layout, _se_ swizzle), and fine-tune parallelism and pipelining strategies (Parallel, Pipelined).
alloc_fragment
Under programmable tile abstractions, TILELANG programs can be represented as a unified fused tile-level dataflow graph (FTG). By operating on this FTG, TILELANG enables fine-grained reasoning and optimization of Al kernels, guiding developers from high-level design choices to fully specified, hardware-efficient kernel configurations. It introduces two complementary techniques. First, tile recommendation analyzes the FTG along with partially specified configurations provide hardware- aware defaults for tile shapes, memory placement, and warp partitions, offering developers high- quality starting points that can be accepted, adjusted, or further tuned. Second, tile inference propagates shape and layout constraints across the FTG to complete the remaining configurations based on the partially annotated operators. It also automatically aligns buffer shapes, layouts, and memory allocations both downstream and upstream. This design blends flexible user control with automated optimization, yielding efficient kernels with far less manual effort.
As shown in Figure 1 TILELANG achieves on average 5.56x the performance of Triton and ap- proaches the hand-written CUDA version in performance, while requiring less than 16% the code size of the manual kernel and even fewer LOCs than Triton. This highlights TILELANG's G ability to attain a more favorable balance between programmability and performance, offering both high efficiency and low development effort. We also implement other modern Al Clernels-includee Dequantize Matmul (Wang et al.) 2024), Multi-Head Attention (MHA) (Vaswani et al. (017)) and Block-Sparse Attention (BSA) (Guo et al., 2024). Despite its deliberately streamlined interface, TILELANG achieves state-of-the-an throughput across heterogeneous GPUs, delivering speed-ups of up to 10.59x over Triton on an NVIDIA H100 and 11.56x on an AMD MI300X (AMD) 2024).
Our contributions are twofold: (1) programmable tile abstractions that let developers directly control and interact with hardware; and (2) tile recommendation and inference that guide developers with hardware- aware defaults and automatically complete configurations over unified FTG graph. We be- lieve TILELANG improves both the productivity and performance of modern Al kernel development.
2 RELATED WORK
AI kernel programming and optimizations. To simplify the development of AI kernels, libraries like FlashAttention-3 (Shah et al. 2024), CUTLASS (NVIDIA 2019), and ThunderKittens (Spector et al. 2025) rely on manual or template-driven designs. Triton (Tillet et al., 2019) provides a high-level Python DSL but restricts control over critical performance paths. Gluon (OpenAl (025) is built on Triton DSL and exposes lower memory hierarchies like shared memory and registers.
Helion (PyTorch) (025) works as a higher-level DSL and is designed to compile down to Triton. based programming model with sequential semantics. Tilus (Ding et al., (025) is another Python DSL for GPU programming, designed with thread- block-level granularity and tensors the core data type. Mojo (Godoy et al. 2025) combines Python's 's interoperability and CUDA-like syntax to build performance portable HPC science kernels. portmmaae--potabll
Cypress (Yadav et 2025) introduces
Frameworks such as Py PyTorch (Paszke et al. 2019). Graphene (Hagedom et al.; (023). MLIR (Lattner et al. 202T). and Welder (Shi et al.) 2023) take a compiler-centri approach. Unlike these works, TILELANG is tile- programmable language that automates layout and low-level configuration
2

---

Published as a conference paper at ICLR 2026
while giving users fine lie-ggrined control. Its flexible tile programming abstraction can help researchers obtain kernels for a broad range of AI operations, and enable advanced optimizations like software pipelining (Cheng et al. 2025) and warp specialization (Huang al. 2023).
Cost modeling. TANGRAM (Gao et al. 2019) optimizes dataflow across scheduling layers, along with performance modeling tool extended by SET (Cai et al. 2023) with Resource Allocation Trees. KPerfIR (Guan et al., 2025) adds instrumentation for profiling and pipeline reordering in Triton. ML-based predictors like Path Forward (Li et al. 2023) and NEUSIGHT (Lee et al.) (0255 also exist,
In contrast, TILELL An's tile-level analytical cost model uniquely captures both computation and data movement at tile granularity, supporting fusion-aware scheduling with high accuracy and usability.
An extended discussion of related work is provided in Appendix [A] covering classic tensor-level IRs (e.g.. XLA (Google) 2019)), polyhedral compilation (Griebl et al. 1998 Zhao et al. 202 I). loop-scheduling systems such as Halide (Ragan-Kelley et al. 2013), 3), the TVM stack (Chen et al.
2018 CUTLASS INVIDIA 2019), and TaichiLang (Hu et al.
2019 2020 2021)
3
PROGRAMMING MODEL
3.1 TILE LANGUAGE
Tile declarations. TILELANG elevates tile-a hyper-rectangular slice of a tensor-to first-class citizen. A tile may be owned by warp, thread block, or any programmer-define, parallel unit, and can be reshaped or re-partitioned at compile time. the FlashMLA kernel, the global matrices are consumed in tiles whose extents are parameterized by block_ H. block_N, __,, and related symbolic sizes. The T. Kernel structure establishes the kernel's 's launch configuration (e.g. bx, by, and the thread count), enabling both index derivation for each thread block and subsequent compiler analyses such as memory-access coalescing and loop tiling.
Tile placement. A distinguishing feature of TILELANG is the ability to map every tile buffer concrete level of the target accelerator's memory hierarchy via user-visible intrinsics, rather than relying on opaque compiler heuristics. T alloo shared reserves storage in low-latency, software-managed shared memory on NVIDIA GPUs (or an architecturally analogous space on other devices). T. alloc_ frrooffrrametttee places accumulator tiles in the register file. Although registers are scarcer than shared memory, their single-cycle latency is indispensable for performance-crition reductions. During compilation, a layout-inference pass distributes these register tiles across threads while respecting register-pressure constraints and bank conflicts.
Tile operators and schedulable primitives. Table Tin Appendix Cshowcases the representative sub- set of core building blocks that orchestrate computation and movement among tiles. Fundamental op- erators (T. copy, T. gemm, T reduce) act on tile operands directly, allowing the programmer to ex- press dense linear algebra, pointwise transforms, and reductions without resorting to scalarized loops. Orthogonal scheduling primitives expose fine-grained control over parallelism (T Parallel), pipelining (T. Pipelined,, and memory layout (T annotate layout, T. use swizzle).
3.2 TILELANG PHILOSOPHY
Tile-level tradeoff. The system adopts tiles as the central abstraction because this granularity provides practical balance between portability and performance. TILELANG models the GPU memory hierarchy and the major compute and data-movement units, exposing tile size, memory placement, warp partitioning, memory layout, and software pipelining as tunable dimensions. This design enables hardware-aware specialization on both NVIDIA and AMD while preserving a unified programming model. Remaining tradeoffs lie below the tile level, where extremely fine-grained hardware behavior cannot be captured through stable and portable API.
Novel tile abstraction. Unlike prior systems where a "tile" is essentially manually managed shared-ment memoryy buffer, TILELANG treats tiles as first-class IR constructs with explicit semantics for indexing, data movement, reuse, and pipelining. This makes tile behavior compiler-visible and supports systematic analysis and transformation. Consequently, TILELANG differs not only in surface syntax but also in the underlying IR, which enables principled optimization at tile granularity.
3

---

Published as a conference paper at ICLR 2026
3.3 A FLASH MULTI-HEAN LATENT ATTENTION EXAMPLE
By fusing high-level expressiveness with architecture-aware orchestration, TILELANG succinctly captures sophisticated AI algorithms such FlashMLA (Liu al.] 2024) while fully harnessing the performance envelope of modern GPU architectures. Figure 2 illustrates TILELANG'S developer-compile co-optimization model: the developer specifies key seci as tile config- uration, launch grid (block H. block N), buffer placement
Cecide tile corfigmacier
Recomeend
sonee (aad) bat sicCliak (egresse))) (eeeadorta
Giveion ledley drstive avrry Taye.
Recomsend
0 0.3tared Taalisz anacvillbsece it, tinl, stape)
Sataad Tallas, durvi((heecu Baai,) drype)
0.5e sueed Tallac shand[ttaak J. (e,ds)), dtgps)
503.3 (tg.lleeesyyrriiiie Mises.0) scoe.dtoe)
Leftre desinve data
www.eleeeeeedd Lyget Inference
T.csyyQlis, block t:dyy 13 HadkL 1. (.Paae)) L.usy/l.pe by Blaouitty 3) biccku, Q. (.
Gut Digitine sctensting
Inference
Pind cewteiieoooeeee....... Hanuul amstappec Tecpptiiibe, BRe
Dieck i, Tead, ataredt
saallo__shaaeeee a.alloc P_alloc_fragmeen fragment), swizzled layouts (T. annotate layout), and warp-level collaboration (T Pararllell) The compiler then infers the remaining low-level details, including latency-hiding pipelines, conflict-free memory layouts, and instruction selection for peak hardware performance. To balance flexibility with automation, TILELANG offers two developer- facing facilities. First, tile recom- mendation (Sec. (.2) supplies hardware-aware defaults that serve as high-quality starting points. Second, tile inference (Sec. (43) analytically propagates user-pros ser-iprodded or recommended hints to complete the schedule and guarantee consistency, Working concert, these facilities deliver near-optimal performance with limited manual tuning.
Eorqueurrtt, T.clest(ber_))
anod Kir_
Mere persitsiniigg
Recomend
T.geen() shared, saane
polisy!t dermuppelloy Fallco1)
T.prar())
policect Gerneupde FatiOn1)
Nato atiliie High perfirnacce iratractien
Inference
(L-urr, M.Ooore and polieatt feadioopal killnn)
Tee 4, 1 T.Peratteiitaas. dial:
acc,aEs, lepeai[i]
E.aaylasz, a, 0.stare))
T.cayyDsshaeee Oetpur by Miockutidy stock,s, (3)
Figure 2: FlashMLA TILELANG kernel example
SCHEDULING GUIDANCE AND AUTOMATION
4.1 TWO-STAGE FRAMEWORK
Optimization space. nigh-perffrmance kernel design in TILELANG begins with a tile-level program, represented as a fused tile-level graph (FTG) capturing dataflow and tiling structure- each node represents tile operator and each edge encodes data dependency. By operating on this unified graph, TILELANG exposes and reasons about hardware-awar optimizations across six key dimensions: tile size (affecting shared memory and register usage), memory placement (selecting appropriate memory scope), warp partitioning (how threads collaborate and bind within block), memory layour (how tile data is organized across memory levels), software pipelining (overlapping compute and data transfer, e.g., via TMA), and tensorization (mapping operations to CUDA or Tensor Cores).
Tile recommendation and inference. To efficiently explore the optimization space, TILELANG adopts a unified two-stage workflow over the FTG. In the first stage, tile recommendation analyzes the FTG provide hardware-aware defaults for partially annotated operators, covering dimensions such as initial tile shapes, memory placement, and warp partitioning (Section These recommendations shape the memory footprint, compute partitioning, and thread collaboration, providing high-quality starting points. In the second stage. leveraging the context from recommendation, tile inference prop- agates constraints through the FTG, automatically inferring the remaining configuration, including tile size, memory layout, software pipelining, and tensorization. It ensures consistency, compatibility, and hardware efficiency (Section 4.3) Together, these stages unify developer guidance and automated completion: recommendation narrows the design space with informed hints, while inference finalizes kernels with minimal manual effort.
fully specified,
In TileLang, the FTG defines the division of labor between developer control and system automation. Developers specify the FTG by composing tile operators; tile-level annotations such as tile sizes, memory placement, warp partitioning, and tensorization are optional. Given FTG, TileLAng's optimization pipeline performs tile recommendation and inference, propagating shape, layout, and memory constraints to complete missing details. This design lets developers concentrate on describ-
4

---

Published as a conference paper at ICLR 2026
ing computation while the system automatically finalizes and optimizes low-level configuration, supporting both fully automated and hint-guided usage.
Running example. Taking MLA as an example (Figure 2). TILELANG first performs tile recom- mendation as illustrated in Figure B Tile operators in the FTG expose tunable parameters--such as tile size, memory placement, and warp-partitioning strategies- as the user interface for these optimization knobs. For instance, in the first T. gemm operator (Figure annotations specify Q and KV tiles in shared memory, while S resides in registers. The S tile is further partitioned across columns using the "po 1 cy=Ful 1 CCC partitaaiti strategy. These decisions directly shape the memory footprint and influence data access patterns across the FTG. The cost model analyzes the FTG to estimate memory traffic, guiding the search toward configurations that minimize data movement.
memory placement
Stage 1: Recommend
Stage 2: Inference
Mocish KV-r
cim
Mociihi KV. shaned
dim
shaned
dmm
biocht shaned
Rucnet
CO0 Stamed
act & ragiteer
rogasor
T.gemmm
Q:ahared,
T.copy( Swregister, S_esst=7, polley=?)
T.gemnet S_caste7,
Lggemm[ G-shared, KV-shared,
Tcopyli saaattsha polley-allgathhe)
T.gemml schaats KV. KV-ahhaeed acc, roc_oaregister policysFultco)
KV-shared,
Seregister,
Scregister, polloy=FullCo)
transpose_B-Toom Trun,
Pollcy=FullCo)
policynFuliCat
Figure 3: Two-stage workflow of optimizing MLA example.
Tile inference completes the configuration by operating over the FTG. For example, once (output of the first T. gemm in Figure 3 and S_cast (input of the second T. gemm in Figure 3 are fixed in location, shape, and partitioning in the first step, inference automatically determines the tile placement and partitioning (e.g., all-gatte -alghrther or scatter) of copy (Figure 3 2 in the second step, ensuring consistency without manual effort. Beyond copy decisions, inference also derives memory layouts by mapping multi-diimessonal indices to physical addresses, explicitly considering vector- ization, coalescing, and bank conflicts. Finally, automates software pipelining and tensorization, ensuring that the resulting kernel configuration is efficient on the underlying hardware. TILELANG also provides ceatform-specine recommendations and inference (see Appendix D)
4.2 TILE RECOMMENDATION
(a) Hardware-Aligned Search:
Cost Model Guided Filtering
(c) Interactive Schedule
blockM: 16, ..)
Accest?
Final:
polley: ( (FullRow, [FullRow,Fulica)
blockM - 64
num_stages:{1 (1,
Candidates
Policy FullCol
Q.llloo(): ( Reg. Smem]
num #2
Q.alloo)) Smam
TheDosicc
(00.
Figure 4: Tile recommendation with cost model. In (b), we show a seatter plot of candidate schedules. The x-axis orders candidates by their predicted latency" and the y-axis shows their normalized scores defined as Intency of best candidate lasency of current candidate which lies in the range (0.1).
Roofline-based cost model. As outlined in Figure TILELANG uses static roofline- based cost model to evaluate candidate configurations, including tile shapes, memory placement strategies, and warp partitioning. Each FTG under given configuration is lowered into an IR that explicitly encodes compute and memory access patterns per tile. From this IR, the model statically estimates execution time as:
MemoryTraffic Computation,
Time max 1.
Lintrinsic
Bandwidth; Performance,
where i indexes memory hierarchy levels (e.g.. HBM, L2, L1), j indexes compute unit types (e.g., tensor cores, CUDA cores, SFUs), and imiinii accounts for inherent overheads such as kernel launch latency and loop prologue/epilogue costs. The roofline formulation assumes perfect overlap between computation and memory transfers, providing a tight upper bound for rapid evaluation without runtime profiling.
5

---

Published as a conference paper at ICLR 2026
Based the cost model, TILELANG generates actionable recommendations for kernel tuning, including tile shapes, memory placement, and warp partitioning. These recommendations form an interactive baseline: developers can accept, adjust, iteratively refine them across multiple rounds. This human-in-the-lowe workflow balances automation with expert insight, slashing tuning effort while preserving full design control.
Tile size. TILELANG presents a ranked shortlist of tile shapes that are multiples of the device's native tensor-core fragments and respect register and shared-memory limits. Each candidate shows predicted arithmetic intensity, memory traffic, and roofline utilisation. Developers can accept the top choice, pin alternatives for later benchmarking, or adjust dimensions manually.
Memory placement. Given chosen tile shape, TILELANG enumerates legal bindings of operands and temporaries to registers or shared memory. flagging options that exceed capacity. Each binding includes estimated pipeline stalls and effective bandwidth, letting developers quickly explore trade- offs and commit or refine placements.
Warp partition. To ensure sufficient thread-level parallelism, TILELANG proposes warp partitions that evenly cover the output tile and match the SM topology. With predicted occupancy and com- putc-memory overlap, developers can select, benchmark, or override, retaining full control while benefiting from data-driven guidance.
4.3 TILE INFERENCE
Layout inference. While memory placement and computation partitioning in Section 4.2 decide where tensors reside and how computation is split, layout inference determines how mineiionaa indices are converted into physical memory addresses- taking into account vectorization, memory coalescing, and bank conflict avoidance. In other words, layout is not about which memory scope is used, but how data is accessed within that scope. Once placement and partitioning are fixed, the system can then infer an appropriate layout to ensure efficient low-level memory access.
TILELANG supports high-level indexing into multi-dimensional arrays (e.g., A [i, k)), which is eventually lowered to physical memory addresses through a hierarchy of abstractions. At the physical level, layouts modeled as linear address expressions of the form E; 34 where y; is the index along dimension i, and s; is its stride. To capture such mappings, TILELANG introduces a composable Layout algebra based IterVar-al a loop iterator that carries range and stride information. This allows layout transformations (e.g.. transposes) to be expressed as algebraic mappings, such as lambda i, 5: (i, i). Formally, a layout becomes function f K K" converting high-level indices into memory addresses. Additionally, TILELANG defines Fragment layouts-a specialized extension where ʃ IK" - K2, mapping each index thread's register ID and its local offset. This enables precise modeling of intra-thread register allocation. Although a buffer of size N theoretically allows (NV) memory layouts, the set of feasible layouts is significantly constrained by hardware. Global memory prefers coalesced access, shared memory requires bank conflict avoidance, and Tensor Core instructions impose strict layout requirements. To explore these constraints, TILELANG employs a greedy strategy that derives valid layouts by enforcing layout rules on selected tile operators.
Layout inference
Free Layout Inference
@Shared
e
@Shared
0
Rhared
0
Switried
Swissted
Shared
T.gemm(
QShared,
Mathuu
MatMui
MuiMul
KVShared,
acc_s,
Meuridti Leyout rebrence
transpose B-True
Bag.
Hardware Efficient Loop Loyout inforence Shared
Ordit Sonthar
Reduce Pios
Common Layout Inference
scom @
T.reduce_ran rax(3cc_5, scores_rox, dim-1)
@
O Flgg
Hepo
10
10
Curpurt
Figure 5: Layout Inference mechanism in TILELANG

---

Published as a conference paper at ICLR 2026
We propose hierarchical layout inference algorithm that operates an FTG. As illustrated in Figure(5] FlashMLA can be represented as a FTG, where nodes are tile-level operators (c.g. matmul, softmax) and edges encode data dependencies. The graph captures how Q, K, and V tiles are loaded into shared memory, attention scores computed and normalized in registers, and final outputs written back. This structure makes memory movement and parallelism explicit, enabling layout inference and efficient scheduling.
Our goal is to synthesize memory layouts that optimize low-level execution efficiency while preserving high-level tensor semantics. The inference process is modeled as a constraint propagation algorithm (Algorithm[]) i Appendix E that iteratively traverses the FTG and incrementally refines the layout mapping L until convergence. As illustrated in Figure the algorithm integrates three complementary inference strategies: (1) Striet Layout Inference (Figure
enforces operator-speciff constraints for hardware-sensittve primitives such as tensor core GEMM, including swizzled shared memory layouts and MMA-aligned register allocations; (2) Common Layout Inference (Figure
propagates layout decisions through structurally aligned operators (e.g., reductions), ensuring consistent thread bindings and register reuse; and (3) Free Layout Inference (Figure
handles the remaining unconstrained layouts by partitioning them into subgraphs via connected component analysis. For each subgraph, the partitioning scheme with the lowest register usage is selected. This step also determines the loop layout using the hardware cost model, which specifies thread binding and vectorization length to maximize memory coalescing and minimize bank conflicts. This unified inference pipeline supports composable, performacc-portable portable layout generation and seamlessly bridges high-level loop indexing with low-level memory organization.
Pipeline inference. TILELANG automatically infers pipelined schedule from sequential program. As shown Figure (a), operations like copy and gemm are overlapped to increase parallelism. The system analyzes dependencies in the FTG and generates a structured pipeline that preserves exccution correctness, exposing only a single num num_stages parameter to users. Additionally, TILELANG applies Warp Specialization to fully exploit asynchronous copy instructions on Hopper GPUs, inserting synchronization barriers where necessary to maintain correct data dependencies. The detailed inference procedure is described in Appendix
Original Serialized Schedule
Iteration
Copy
MatMul(Q, KV)
Copy MatMul(O_Pe K_PE)
Pipelineednum, Pipplineddnumsstgeee=/
Iteration
MatMul(Q,
Copy
MatMullQ K_PE)
Pipaline
Inference
Producer
Copy
Copy K_PE
Copy
Copy
Pipelined Schedule
Consumer
MatMul(O,
MatMul(O_ K_PE)
MatMul(O, KV)
MatMui(Q_PE,, K.PE)
Figure Pipeline Inference mechanism in TILELLNG
Instruction inference. In TILELANG, while low-level hardware instructions such as 4p4a or mma can be manually invoked via source injection or inline PTX (NVIDIA) 2021), choosing the most appropriate instruction based on input shapes and data types can be challenging. To address this, TileLAnG integrates with high-level Tile Libraries like NVIDIA's cute (NVIDIA) 2019) and AMD's ck (AMD) (025), which abstract hardware-speciffc details and automatically choose efficient instructions based on input configurations. These libraries expose standardized tile-based APIs and TILELANG supports their invocation via a unified T. ccall_extern interface, simplifying development while ensuring performance portability.
(e.g., t1:
EVALUATION
TILELANG realized as a Pythonic DSL whose compiler lowers high-level tile programs to hardware- specialized kernels through modular IR and oode-generation pipeline. TileLang is implemented on top of the TVM backend, but our main contributions sit above TVM. TILELANG provides the tile abstraction, the FTG IR, and its own optimization passes, while TVM supplies the low-level code generation backend. As illustrated in Fig. 7 TILELANG adopts a five-stage compilation workflow: (1) tile-level code is written a Python- based DSL: (2) the compiler translates the AST into the TileLang AST; (3) a FTG is constructed from the TensorIR: (4) series of optimization passes in TileLang and TVM is applied; and (5) the optimized IR is finally lowered to CUDA, or other backends.
7

---

Published as a conference paper at ICLR 2026
Optimizations
LLVM IR CUDA C/CH+
Tilelang Python
Parser
TileLang
Builder
FTG
TVM
Codegen
Program AST
AST
HIP C/C++
FIG Pass
TVM
Figure 7: Overall TILELANG workflow.
5.1 EXPERIMENTAL SETUP
Hardware platforms. We assess the performance of TILELANG two leading GPU architectures: NVIDIA and AMD, which dominate contemporary accelerator ecosystems. Our evaluation employs state-of-the-art hardware, including the NVIDIA H100 (80GB) (NVIDIA 2023) and the AMD Instinct MI300X (192GB) (AMD 2023). The NVIDIA H100 leverages CUDA 12.8, while the
AMD MI300X utilizes ROCm 6.2.0. Both GPUs are benchmarked under the Ubuntu 20.04 operating system to ensure consistency in environmental configurations.
AI kernels. To evaluate system performance, we analyze nine representative operators: (1) GEMM, (2) fused dequantized GEMM AFPIS,
(3) Attention, (4) Multi-Head Latent Attention, (5) Block Sparse Attention, (6) 2D Convolution, (7) Chunk Gated Delta Net, (8) Vertical Slash Sparse
Attention, and (9) Attention Sink. Shape configurations are provided in Appendix N
Baselines, Our comparative analysis considers the following baselines: (1) PyTorch Induc-
tor-torch.. - ma tmul for GEMM, SDPA (PyTorch 2023) for attention, and other operators com-
piled via Inductor: (2) Triton implementations, including GemLite (Mobius ML (2024) and MLA from SGLang (Zheng et al., 2024]; (3) ThunderKittens (TK) (Spector et al. [2025) a template-
based framework for high- hegh-peeformance AI kernels on NVIDIA GPUs; and (4) Highly optimized
libraries, including CUTLASS (NVIDIA 2019) and Composable Kernel (AMD) (025) for GEMM,
Marlin (Frantar et al. 2025) for dequantized GEMM, FlashAttention- V3 (Dao 2023) for MHA,
AITER (AMD) 2025) for MLA, and Block Sparse Attention (Guo et al. 2024 for sparse attention.
We evaluate kernel performance versus code complexity (Section 5.2) and present ablation results (Section 5.3
with cost model and tuning time analyses in Appendices and H
5.2 KERNEL PERFORMANCE
Matrix Multiplication. TILELANG achieves high performance with low code complexity across diverse GEMM configurations, demonstrating 1.18-1. 1 speedup over PyTorch on NVIDIA H100, while maintaining competitive performance (0.94-1.05x) on AMD MI300X. It also de- livers 1.08-1.43x speedup over Triton with minimal kernel code, enabled by automated inference that abstracts low-level hardware details such as TMA and pipeline scheduling. Compared with TK, TILELANG achieves 0.99-1.11x speedups while reducing code complexity by 77%. Its cost-model guidance and automated tile inference eliminate manual tuning. TK depends on curated CUDA templates, limiting to NVIDIA GPUs, whereas TILELANG supports multiple hardware backends.
Low-Bit Matmul. For WINT4AAPMI GEMM, TILELANG achieves 1.35-3.81 speedups over PyTorch and up 1.55x over Triton on H100, while outperforming the specialized Marlin kernel with far simpler code. On MI300X, it delivers on average 0.96 x over Triton. These gains arise because TILELAnG exposes low-level memory, dequantization, and layout controls that are abstracted away in Triton.
Convolution. On H100, TILELANG achieves 1.24-1.79x and 1.10-1.97x x speedups over Py Torch and Triton, respectively, with reduced code complexity. These gains come from its instruction inference mechanism, which maps data movement efficiently to TMA im2 col. On MI300X, the improvements are even larger, reaching 1.29-6 6.80x over PyTorch and 1.02-3.10x over Triton.
Flash Attention. TILELANG achieves efficient attention computation with concise code across sequence lengths. On H100 and MI300X. it delivers 1.08-1.58x and 1.22-1.37x speedups over Triton, while matching the performance of FlashAttention- V3 (0.98x and 0.96x on average). These results stem from TILELANG's ability infer and apply platform-specifi partitioning and pipelining strategies that exploit specialized compute units. TILELANG achieves up to 1.10x speedup over TK while significantly reducing code complexity (from 185 lines to 66), highlighting its programma-
8

---

Published as a conference paper at ICLR 2026
TileLang Py Taach
Thunderkittess Marlin
FlashAtbention BSA
CUTS
Spoosuu Normalzzed to
GEMM-FP16XFP16
GEMM-FP GEMM-F116XINT4 16XINT4
Conv2d
MHA
15
12
1.22x
1.5
10-
1.14g
1.96x
2.91x
1.0
1.0x
1,00kg
1.00x
1.00k
50
100
150
200
40c
20
40
200
400
6on
800
Code Lines (LOC)
Code Lines
Code (LOC)
Code Lines (LOC)
Block-Sparse-MHH.
Chunk- Dennk-Gaadd-Delt-Neee Vertical-Slach Sparse
Attention-Sine
Spocdup Noma
50-
20
20. 30.38k
42 85X
109.13x
25-
200
114
10
100g
1,00x
100
200
200
50
100
150
50
75
100
125
60
80
Code Lines (LOC)
Code Lines (LOC)
Code Lines (LOC)
Code Lines (LOC)
(a) Performance vs. code complexity across operator configurations NVIDIA H100 GPU.
Tillang
I
Pytorrhh
Triton
+ Composable Kemei
GEMM-FP16XFPL
GEMM FP16XINT4
Conv2d
Spocdup Nemmalzea to Pyturch
1.0
0.74k
0.57x
F 84k
5.0
2.98k
100
1.11x
1.00x
2.5
200
400
600
100
150
20
40
60
80
500
1000
1500
Code Lines (LDC)
Code Lines (LOC)
Code Lines (LOC)
Code Lines (LOC)
Spenodup Normalted
Biock-Sparse-MHA
MLA
Vertical Sisttttarrse-tttteeennn
Attention-Sinks
PyTora PyTorch
38
400
28 1.25x
300
239.70x
1.77
200
200
16X
100
1.00%
48. 71x
60x
100
125
40
en
80
Code Lines (LOC)
Code Lnnss (LOC)
Code Lines (LOC)
Code Lines (LOC)
(b) Performance vs. code complexity across operator configurations on AMD MI300X GPU.
Figure 8: Performance and code complexity on an NVIDIA H100 GPU and AMD MI300X GPU.
The y-axis denotes the speedup relative to PyTorch, while the x-axis indicates lines of code (LOC).
Ideal solutions appear toward the top-left corner.
bility. By combining tile-level guidance with automated inference, TILELANG streamlines kernel development. This is particularly valuable for complex attention operators, where TK often requires extensive manual tuning of tiling, warp partitioning, layout, and pipelines.
Flash MLA. As shown in Figure I TILELANG achieves 4.06-10.59 10.99x speedups over Triton on H100, with substantially reduced code complexity. It matches the latency of the specialized FlashMLA kernel while reducing code complexity by 6.86x. On MI300X, TILELANG delivers 5.64-12.97x gains over Triton and slightly outperforms the hand-tuned ROCm library AITER (1.05x). These improvements arise from warp specialization and automated TMA mapping.
Block Sparse Attention. TILELANG achieves acceleration of 3.42-7.87x and 1.22-1.37 over Triton with less code on H100 and MI300X, respectively. On H100, it matches BlockSparse (BSA) latency (0.91-1.82x) while greatly reducing complexity. Implementing block-sparse MHA requires only adding two lines to the standard MHA code (Appendix 0.5)
Chunk Gated Delta Net. On H100, TILELANG achieves 15.88-77033xx speedups over Py by fusing complex operations into a single kernel. Compared to Triton, attains 1.10-1.45x speedups with 39% fewer lines of code. These gains come from automated tile recommendation and inference, which optimize memory placement and partitioning for efficient hardware utilization.
Vertical Slash Sparse Attention. TILELANG delivers 108.55-280. $45-280.411 and 105.01-363.52 speedups over PyTorch H100 and MI300X, largely by fusing the sparse attention operation into a single efficient kernel. Compared to Triton, achieves 1.16-1.97x and 1.19-1.60x speedups on H100 and MI300X, respectively, while cutting the code size by roughly half.

---

Published as a conference paper at ICLR 2026
Attention Sinks. For attention with the sinks mechanism, TILELANG achieves 14.21 -2x.577 and 5.35-9.11x speedups over Py TyTcchh on H100 and MI300X, respectively, enabled by TILELANG's FTG-based fusion into single optimized kemel. Against Triton, it reaches 1.13-1.30x H100 and 2.x2-2.69x on MI300X. The attention-sink variant differs only slightly from standard MHA, showing that TILELLNGG readily supports diverse attention patterns with minimal effort.
5.3 ABLATION STUDIES
To help clarify what contributes to the speedups over the baseline, we perform an ablation study on FLashMLA as a representative example. Starting from TANELLANG version that uses manually crafted scheduling heuristics (TL-Heuristic), we progressively enable three components: (i) cost- model-guided tiling (+Tile), which improves the compute-memory ratio and cache use; (ii) cost- mooellggided memory placement (+Alloc), which chooses efficient buffer locations and reduces register spilling; and (iii) warp partitioning (+Partition), which improves intra-warp load balance. Performance is measured at each stage relative to the Triton baseline.
H100
MI300x
Spandup vs thton
12.5
12x
Speeeda vs Triton
15
13
10.0
10x
TL-eeristic
5.0
+oorition
2.5
Baseiin -
JK
2K
4K
8K
IK
2x 4x
Sequence Length
KV Sequence Length
Figure 9: Ablation study for FlashMLA on both H100 and MI300X GPUs. As shown in Figure each of the evaluated optimizations provides measurable performance gains over the Triton baseline, validating their effectiveness. Our analysis also highlights architecture- specific behaviors (illustrated in Appendix D). On H100, tiling (+Tile) yields modest speedup of 1.31 x. Building on this, warp partitioning (+Partition) provides the dominant contribution, delivering an additional 4.34x improvement. On MI300X, allocation placement (+Alloc) serves the primary optimization, achieving a 6.56 x speedup. When further combined with tiling (+Tile), the overall gain increases by another 1.75 improvement.
5.4 COMPARE WITH MORE RECENT DSLS
We further compare TILELANG with recent DSLs including Helion, Gluon, and Tilus. Our evalua- tion covers the latest Hopper-supported implemen- tations across GEMM, MHA, and Mamba-chunk- scan. Figure I0 reports the GEMM results, with additional results in Appendix L On GEMM, TILE- LANG achieves 1.15 ~ 1.62x, 1.52 ~ 1.83x, and 1.87 ~ 2.12x speedups over Helion, Gluon, and Tilus, respectively, while using fewer lines of code. These improvements largely stem from TILELANG's compiler-visible tile abstraction, which enables more structured optimization than existing DSLs. Note that Tilus is not yet fully optimized for Hopper features such as WGMMA and TMA, and on Ampere/Ada TILELANG is still slightly better (see Appendix D).
GEMM Performance and Code Complexity on H100
TileLang
1.4
Speeduu Normalized
PyTorch
to Py TTrch
1.2
1.25x
Gluan
Helion
1.00x
0.99x
Tilus
0.8
0.77x
0.63g
20
10
80
100
120
Code Lines (LOC)
Figure 10: Performance and code-size parison of GEMM kernels on H100 across recent systems.
CONCLUSION
TILELANG offers a controllable tile-level programming model with graph-based optimizations via tile recommendation and inference. By combining automated configuration with fine-grained developer control, it streamlines kernel development and delivers significant speedups. It enables rapid experimentation with emerging AI algorithms, such custom attention, sparsity, and quantization. TILELANG also lowers barriers for systems-aware research across diverse hardware platforms.
10

---

Published as a conference paper at ICLR 2026
7 REPRODUCIBILITY STATEMENT
Operator shapes used in our benchmarks are drawn from widely adopted, real-world AI models (c.g.. GPT-OSS, DeepSeek V3, Qwen3-Next). A list of these operator configurations is included in Appendix N and the corresponding TILELANG code of kernels used in evaluation provided in Appendix The system implementation and scripts for reproducing our experiments are publicly available, ensuring full reproducibility.
We provide a detailed description of our experimental setup in Section 5.1
ACKNOWLEDGEMENT
We would like to thank all users and developers who have utilized and contributed Tilel.ang since its open-source release. We express special gratitude to DeepSeek for their pioneering adoption of TileLang in the development of new model algorithms. We also thank them for providing valuable feedback and significant code contributions to the community. This work was supported by the National Natural Science Foundation China under Grant No. 62572009.
REFERENCES
Joshua Ainslie, James Lee- Lee-Thorp, Michiel De Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.1 13245, 2023.
AMD,
Amd TM 3 architecture. Technical report, AMD, 2023.
URL https
AMD. White paper: Introducing amd cd 3 architecture, 2024. URL
https
tecchtthhhhho docs/
AMD. ROCm aiter, 2025. httpss//g...hbb.co////01
AMD. AMD Composable Kernel, 2025.
https://github. wwt.ecommmooo/.ee hata/ROCm/coooe ROCm/cooposalllee
kernel
Jingwei Cai, Yuchen Wei, Zuotong Wu, Sen Peng, and Kaisheng Ma. Inter-layer scheduling space definition and exploration for tiled accelerators. Proceedings of the 50th Annual International Symposium Computer Architecture, pp. 1-17, 2023.
Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, et al. [TVM}: An automated (End-to-End) optimizing compiler for deep learning. In 13th USENIX Symposium Operating Systems Design and Implementation (OSDI 18), pp. 578,594, 2018
Yu Cheng, Lei Wang, Yining Shi, Yuqing Xia, Lingxiao Ma, Jilong Xue, Yang Wang, Zhiwen Mo, Feiyang Chen, Fan Yang, et al. Pipethreader: Software-defined pipelining for efficient dnn execution. In 19th USENIX Symposium Operating Systems Design and Implementation (OSDI 25), 2025.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arX1v:2307.086// 2023.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, arXiv preprint arXivv2405.2 2024.
DeepSeek.
2025. http://////tttbb.oom///eeeeeeeea a
Yaoyao Ding, Bohan Hou, Xiao Zhang, Allan Lin, Tianqi Chen, Cody Yu Hao, Yida Wang, and Gen- nady Pekhimenko. Tilus: tile-level gpgpu programming language for low-precision computation. arXiv preprint arXiv:2504. /2984, 2025.
11

---

Published as a conference paper at ICLR 2026
Siyuan Feng, Bohan Hou, Hongyi Jin, Wuwei Lin, Junru Shao, Ruihang Lai, Zihao Ye, Lianmin Zheng, Cody Hao Yu, Yong Yu, et al. Tensorir: An abstraction for automatic tensorized program optimization. In Proceedings of the 28th ACM International Conference Architectural Support for Programming Languages and Operating Systems, Volume 2, pp. 804-817, 2023.
Elias Frantar, Roberto L Castro, Jiale Chen, Torsten Hoefler, and Dan Alistarh. Marlin: Mixed- precision auto-regressive parallel inference on large language models. In Proceedings of the 30th ACM SIGPLAN Annual Symposium on Principles and Practice of Purallel Programming, pp. 239-251, 2025.
Mingyu Gao, Xuan Yang, Jing Pu, Mark Horowitz, and Christos Kozyrakis. Tangram: Optimized coarse-grained dataflow for scalable nn accelerators. In Proceedings of the twenty-Fourth Interna- tional Conference on Architectural Support for Programming Languages and Operating Systems, pp. 807-820. 2019.
William F Godoy, Tatiana Melnichenko, Pedro Valero-Lara, Wael Elwasif, Philip Fackler, Rafael Fer- reira Da Silva, Keita Teranishi, and Jeffrey S Vetter. Mojo: Mlir-based peefformance-porttlee hpc science kernels on gpus for the python ecosystem. arXiv preprint arXiv:2509.21039 2025.
Google. XLA, 2019. https: : http://wwwtttseee http///wtetsrrooooorrrrr/o//////
Martin Griebl, Christian Lengauer, and Sabine Wetzel. Code generation in the polytope model. In Proceedings. 1998 International Conference on Parallel Architectures and Compilation Techniques (Cat. No. 98EX192), pp. 106-111. IEEE, 1998.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv: 02.00252, 2023.
Yue Guan, Yuanwei Fang, Keren Zhou, Corbin Robeck, Manman Ren, Zhongkai Yu, Yufei Ding, and Adnan Aziz. Kperfir: Towards a open and compiler-centric ecosystem for gpu kernel performance tooling on modern ai workloads. 2025.
Junxian Guo, Haotian Tang, Shang Yang. Zhekai Zhang, Zhijian Liu, and Song Han. Block Sparse At- tention. https: /tggt/ttttbbb../ com mi -an hatptttthaaaa-aaaakc--------e-ee -ettparrre--ttt.....e on 2024.
Bastian Hagedorn, Bin Fan, Hanfeng Chen, Cris Cecka, Michael Garland, and Vinod Grover. Graphene: An for optimized tensor computations on gpus. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3, pp. 302-313, 2023.
Yuanming Hu, Tzu-Mao Li, Luke Anderson, Jonathan Ragan- and Frédo Durand. Taichi: language for high-performance computation on spatially sparse data structures. ACM Transactions on Graphics (TOG), 38(6):201, 2019.
Yuanming Hu, Luke Anderson, Tzu-Mao Li, Qi Sun, Nathan Carr, Jonathan Kelle,,,lley, and Frédo Durand. Difftaichi: Differentiable programming for physical simulation. ICLR, 2020.
Yuanming Hu, Jiafeng Liu, Xuanda Yang, Mingkuan Xu, Ye Kuang, Weiwei Xu, Qiang Dai,
William Freeman, and Frédo Durand. Quantaichi: A compiler for quantized simulations. ACM Transactions on Graphics (TOG), 40(4), 2021.
Guyue Huang, Yang Bai, Liu Liu, Yuke Wang, Bei Yu, Yufei Ding, and Yuan Xie. Alcop: Automatic load-compute pipelining in deep learning compiler for ai-gpus. Proceedings of Machine Learning and Systems, 5:680-694 2023.
Hyoukjun Kwon, Prasanth Chatarasi, Vivek Sarkar, Tushar Krishna. Michael Pellauer, and Angshu- man Parashar. Macstro: A data-centric approach to understand reuse, performance, and hardware cost of dnn mappings. IEEE micro, 40(3):20-29 2020.
Chris Lattner, Mehdi Amini, Uday Bondhugula, Albert Cohen, Andy Davis, Jacques Pienaar, River Riddle, Tatiana Shpeisman, Nicolas Vasilache, and Oleksandr Zinenko. Mlir: Scaling compiler infrastructure for domain specific computation. In 2021 IEEE/ACM International Symposium on Code Generation and Optimization (CGO), pp. 24. IEEE, 2021. httpt::////// :tta:://mmir.. iilm..
lorg/
12

---

Published as a conference paper at ICLR 2026
Sconho Lee. Amar Phanishayee, and Divya Mahajan. Forecasting gpu performance for deep learning training and inference. In Proceedings of the 30th ACM Interationall Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1. pp. 493-508, 2025.
Ying Li, Yifan Sun, and Adwait Jog. Path forward beyond simulators: Fast and accurate gpu execution time prediction for dnn workloads. In Proceedings of the 56th Annual IEEE/ACM International Symposium on Microarchtteture, pp. 380-394, 2023.
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng. Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437, 2024.
Mobius ML. Gemlite: Lightweight generative machine learning engine. https://github.cr com/ mooiiusm//gmmm11 2024.
NVIDIA. NVIDIA cutlass, 2019.
NVIDIA.
PTX ISA. 2021.
com/cuda/
NVIDIA, Nvidia h100 tensor core gpu architecture. Technical report,, NVIDIA Corporation, 2023.
OpenAl. Gluon, 2025. python/examp les/ gluon
Angshuman Parashar, Priyanka Raina, Yakun Sophia Shao, Yu-Hsin Chen, Victor A Ying, Anurag Mukkara, Rangharajan Venkatesan, Brucek Khailany, Stephen Keekler, and Joel Emer. Timeloop: A systematic approach to dnn accelerator evaluation. In 2019 IEEE international symposium on performance analysis of systems and software (ISPASS). pp. 304-315. IEEE, 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019. httppss//pptorrhh.orgg
Py PyToch. torch..n.functional..aalel dot prodtppooottttieee attention,
2023.
URL https
//000//wwaccra..cc////////////21 sca beaeled__dt_prrccct__tteeetooood
nn. functional.
PyTorch. Helion: A High-Level DSL for Performant and Portable ML Kernels, 2025.
https
Jonathan Ragan-Kelley, Connelly Barnes, Andrew Adams, Sylvain Paris, Frédo Durand, and Saman Amarasinghe. Halide: A language and compiler for optimizing parallelism, locality, and recom- putation image processing pipelines. In Proceedings of the 34th ACM SIGPLAN Conference on Programming Language Design and Implementation, PLDI '13, pp. 519-530, New York, NY, USA, 2013. ACM. ISBN 978- 777-1-45033-2014-6 doi: 10.1 00.145//491996624227766 URL
http://doi. acm. 1Ttg
Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. arXiv preprint arXiv:2407.08608, 2024.
Yining Shi, Zhi Yang, Jilong Xue, Lingxiao Ma, Yuqing Xia, Ziming Miao, Yuxiao Guo, Fan Yang, and Lidong Zhou. Welder: Scheduling deep learning memory access via tile-graph. In 17th USENIX Symposium on Operating Systems Design and Implementation (OSDI 23), pp. 701-718, 2023.
Benjamin F Spector, Simran Arora, Aaryan Singhal, Daniel Fu, and Christopher Ré. Thunderkittens: Simple, fast, and adorable ai kernels. In Proceedings the 13th International Conference on Learn- ing Representations, 2025, URL http:://pppppeeeeoeoeeeee/oaoan/ http://ggtttbb..com/aaaeeeeeeee http:://tttt...oo//aa//eee////nr
13

---

Published as a conference paper at ICLR 2026
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprins arXiv:2307.08621, 2023.
Philippe Tillet, H. T. Kung, and David Cox. Triton: An Intermediale Language and Compiler for Tiled Neural Network Computations, pp. 10-19. Association for Computing Machinery, New York, NY. USA, 2019. ISBN 9781450367196. URL /ttp:::///ol.oo0/1..1115/////// 3329973
Nicolas Vasilache, Oleksandr Zinenko, Theodoros Theodoridis, Priya Goyal, Zachary DeVito, William S. Moses, Sven Verdoolaege, Andrew Adams, and Albert Cohen. Tensor comprehensions: Framework-agnostir high-performance machine learning abstractions. CoRR, abs/1802.04730, 2018. URL BtLLtttpp/axxx http:///xww..ooo/aa/////// 1802. 04730
org
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Haojie Wang, Jidong Zhai, Mingyu Gao. Zixuan Ma, Shizhi Tang, Liyan Zheng, Yuanzhi Li, Kaiyuan Rong, Yuanyong Chen, and Zhihao Jia. PET: Optimizing tensor programs with par- tially equivalent transformations and automated corrections. In 15th USENIX Symposium Operating Systems Design and Implementation (OSDI 21), pp. 37-54. USENIX Association,
July 2021. ISBN 19991-999933-22299 URL https: /tt::://wwww.. www usenix. org/conference/ wnnpettttttttttinnoooooooo wang
Lei Wang, Lingxiao Ma, Shijie Cao, Quanlu Zhang, Jilong Xue, Yining Shi, Ningxin Zheng, Ziming Miao, Fan Yang. Ting Cao, et al. Ladder: Enabling efficient (lo-Prrecssoo] deep learning computing through hardware-aware tensor transformation. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), pp. 307-323, 2024.
Mengdi Wu, Xinhao Cheng, Shengyu Liu, Chunan Shi, Jianan Ji, Kit Ao, Praveen Velliengiri, Xupeng Miao, Oded Padon, and Zhihao Jia. Mirage: A multi-level superoptimizer for tensor programs. In 19th USENIX Symposium on Operating Systems Design and Implementation (OSDI 25), Boston, MA, July 2025. USENIX Association.
Rohan Yadav, Michael Garland, Alex Aiken, and Michael Bauer. Task-based tensor computations on modern gpus. Proceedings of the ACM on Programming Languages, 9PPLDD):3964420, 2025.
Songlin Yang. Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. Advances in neural information processing systems, 37: 115491-115522 2024.
Jie Zhao, Bojie Li, Wang Nie, Zhen Geng, Renwei Zhang, Xiong Gao, Bin Cheng, Chen Wu, Yun Cheng, Zheng Li, Peng Di, Kun Zhang, and Xuefeng Jin. Akg: Automatic kernel generation for neural processing units using polyhedral transformations. In Proceedings of the 42nd ACM SIGPLAN International Conference on Programming Language Design and Implementation, PLDI 2021, pp. 1248. New York, NY, USA, 2021. Association for Computing Machinery. doi: (11145/3453443.445444 3453483. 3454106
ISBN 9781450383912
URL /ttLhttps////o...ooo http://o...r.....111 1145/
Lianmin Zheng, Chengfan Jia, Minmin Sun, Zhao Wu, Cody Hao Yu, Ameer Haj-Ali, Yida Wang, Jun Yang, Danyang Zhuo, Koushik Sen, et al. Ansor: Generating Perfo--erformancc) tensor programs for deep learning. In 14th USENIX symposium on operating systems design and implementation (OSDI 20), pp. 863-879, 2020.
Lianmin Zheng, Liangsheng Yin. Zhiqiang Xie, Chuyue Livia Sun, Jeff Huang, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, lon Stoica, Joseph E Gonzalez, et al. Sglang: Efficient execution of structured language model programs. Advances in Neural Information P Processing Systems, 37: 62557-62583, 2024.
14

---

Published as a conference paper at ICLR 2026
Our appendix is organized follows:
Appendix
Extended discussion of related work.
Appendix
B
Details of the MLA algorithm.
Appendix
Semantics of a partial list of TILELANG primitives.
Appendix
D
Platform-specifie scheduling
Appendix
Layout inference algorithm of TILELLAGG
Appendix
Pipeline inference algorithm of TILELANG.
Appendix
Evaluation of the cost model.
Appendix
Tuning time measurements.
Appendix
Matmul implementation difference between TVM, Triton, and TILELANG.
Appendix
Comparison of FTG IR and TVM IR.
Appendix
K
Comparison of FLashMLA implementations different architectures.
Appendix
Comparison with recent D
Appendix
M
Commit hashes of each baseline.
Appendix
Operator shapes used our benchmark.
Appendix
TILELANG code of kernels used in the evaluation.
A EXTENDED DISCUSSION OF RELATED WORK
AI kernel generators For kernel generating and optimization, Ansor Zheng et 2020 builds many kernel combinations by sampling programs from hierarchical search space. PET (Wang et al. 2021) moves forward to partially equivalent transformations and automated correction for newly discovered kernels. TensorIR (Feng et al., 2023) generates new kernels by generalizing the loop nest representation used in existing machine learning compilers. Mirage (Wu et al. 2025) proposes a uniform representation of a tensor program at each level of the latest GPU compute hierarchy to find custom kernels. AKG (Zhao et al. 2021) leverages polyhedral schedulers to perform much wider class of transformations automatically generate kernels NPUs. TileLang also utilizes multi-level programming interfaces and autotuning techniques based on tile-level cost model for kernel generation and optimization.
Tensor-level IRs, The polyhedral model, and Loop synthesizers. Traditional approaches ad- dress program optimization at different abstraction levels: Tensor-level IRs (c.g., XLA (Google 2019h) lower tensor programs via pattern-matched templates (e.g., LLVM, CUDA). Polyhedral mod- els (Griebl et al. 1998) (e.g., TC (Vasilache et al.] (018)) automate affine loop transforms, mainly for al., (013)) generate loop nests guided by user-defined schedules. TILELANG targets distinct programming model and control granularity. differs fundamentally by introducing tiles as first-class programming units. It offers programmable control over fusion strategies, memory hierarchy, and parallelism. This enables developers to design fused kernels with both high performance and portability across hardware.
DNN layers. Loop synthesizers (e.g., Halide (Ragan-Kelle
TVM (Chen et al. 2018). TILELANG builds upon TVM's IR and arithmetic passes. However, unlike TVM's schedule-driven loop generation from high-level compute definitions, TILELANG offers explicit, tile-level programmability and control over memory, fusion, and parallelism. This enables much finer-grained kernel customization beyond what TVM can achieve. For instance, TVM cannot fully express advanced algorithms like Attentiten (FA) or Multi-Head Latent Attention (MLA), which demand precise management of memory hierarchy and execution order cader--aaalllities that TileLang supports.
Warp Partition. Warp Partition (WP) is key component of TILELANG's exccution model, building directly on the tile abstraction. Given a specified tile size, WP allows further partitioning of the tile along each dimension across multiple warps. For example, consider GEMM operation C A@B,
15

---

Published as a conference paper at ICLR 2026
where A E IRRM E BERN and C RMMN, The output tile C can be partitioned along either the M or axis, corresponding to full-row or full-column warp-pann suup-artitioninng strategies, respectively. By giving users explicit control over warp partitioning, TILELANG enables fine- iine-griieed management of resources such as register usage within each warp. This, in turn, allows users to better control the performance of operations. Such flexibility is crucial for mapping computations efficiently to hardware, especially when optimizing diverse and performance-sensitiy kernel workloads.
CuTe library. While both TILELANG and CuTe (NVIDIA) 2019) share this high-level goal, their underlying mechanisms differ: Cu'Te relies on shape/stride pairs, whereas TILEELANG encodes the mappings using explicit arithmetic expressions. This arithmetic formulation offers advantages by more directly capturing index transformations and enabling more flexible, composable manipulations, allowing for clear definition and description in the DSL frontend.
The rooffine-guided cost model. as the sested-loop-basee modeling in Timeloop (Parashar et al. (2019). the data-centric representation in Maestro (Kwon al. 2020). contrast these methods, our work leverages a tile-level program- ming abstraction, which naturally lends itself to tile-centric cost model. This enables us to accurately capture both computation and data movement at the tile granularity, while maintaining simplicity and enhanced support for modeling operator fusion. This design strikes a balance between accuracy and usability, making it effective for guiding schedule selection without introducing excessive complexity.
Several analytical modeling approaches have been proposed, such
TaichiLang. Although both Taichi (Hu et al. 2019 2020 202 I) and TileLang target GPU workloads, they differ fundamentally in abstraction and intended use. Taichi provides high-level, scalar-loop DSL with automatic parallelization, SNode-based data layouts, and strong autodiff support, making it well-suited for scientific computing and simulation. TileLang, by contrast, is designed for deep- learning kernels such as attention and GEMM, where peak performance requires explicit control of tile shapes, memory placement across shared/I.DS and registers, multi-stage pipelining, warp specialization, and instruction selection (Tensor Core MMA or AMD MFMA). These capabilities are not directly expressible Taichi's fully automatic model.
B MLA ALGORITHM
Instead of storing full-sized key and value matrices, MLA projects input token embeddings into lower-dimensional latent space using down-projection matrix:
Ze = X, Wdowa-
The latent vector 2, is then used to reconstruct the key and value representations:
k;- = 2, W
up: V, = 2, W
up
To incorporate positional information, Rotary Positional Embedding (RoPE) is applied to the recon- structed keys and queries:
k oo RoPE(k;).
Queries are also compressed using similar process reduce activation memory:
q: z, @
MLA further enhances computational efficiency through a technique known as matrix absorption. which reorders matrix multiplications to optimize performance. This approach enables the key and value inputs to share the same latent representation 2,, thereby reducing redundancy and memory usage. In the adopted configuration, MLA employs single shared key-value (KV) head, with a head dimension of 512.
C PARTIAL LIST OF TILELANG PRIMITIVES
Table illustrates the expressiveness of the TILELANG intermediate representation. To support efficient code generation across diverse hardware backends, TILELANG decouples the definition of algorithms from their optimization. The Dataflow Centric Tile Operators define the functional seman- ties of the workload (e.g., matrix multiplication, atomic updates), while the Scheduling Primitives expose critical optimization handles such as loop pipelining ("Pipelined") and layout transformation ('annotate layottlaaaotaaaaa to maximize hardware utilization and locality without altering the algorithmic correctness.
16

---

Published as a conference paper at ICLR 2026
Table 1: partial list of primitives supported by TILELANG.
Dataflow Centric Tile Operators
Scheduling Primitives
copy
data movement among hicrarchy memory.
Parsllel
Parallolization of loop iterations over threads.
genm
matrix multiplication different GPUs,
Pipelined
Enables pipelining overlap data transfers with
computation.
reduce
reduction operator (e.g. sum, min, max) ex-
annotate_ layout
Definition of custom memory layouts to minimize
ploiting warpyblock-lowe parallelism.
bank conflits optimize thread binding.
atomic
atomic operations to ensure thread-safe up- shared global memory.
use_awizale
Improves L2 cache locality via swizzled access pat- terns.
Warp Specialization
barrier_arrive
Signals the arrival at synchronzatiion point (mburrirr) for producer!
consumer coordination.
barriex_ech
execution until specific conditions (e.g. transaction counts) met.
D PLATFORM-SPECIFIC SCHEDULING
TILELANG also provides platform-specific recommendations and inference. Taking MLA as an example, we illustrate how TILELANG performs tile recommendation and inference based on the code shown Figure 2
On the H100, each SM features 228 KiB of shared memory and 256 KiB register file, whereas the MI300X provides 64 KiB of Local Data Share (LDS) and a total of 512 KiB in registers. Given these architectural differences, TILELANG first recommends different tile configurations.
Table 2: Comparison of specifications between NVIDIA H100 SXM and AMD MI300X.
Specification
NVIDIA H100 SXM
AMD MI300X
Clock Frequency
1.83 GHz
2.10 GHz
DDR Memory Bandwidth
3.35 TB/s
5.30 TB/s
L2 Bandwidth
9.45 TB/s
16.63 TB/s
L1/Shared Memory
30.92 TB/s
81.72 TB/s
Compute Units (SMs/CUs)
132 SMs
304 CUs
Shared Memory per SM/CU
228 KiB
64 KiB
Register File per SM/CU
256 KiB
512 KiB
Peak FP16 Performance
989 TFLOPs
1307 TFLOPs
As shown in Figure II for memory placement, users may initially allocate the Q tile to shared memory the MI300X. However, this approach fails due to the limited capacity of shared memory. TILELANG detects this constraint and instead recommends placing both Q and acc_s in registers. In contrast, on the H100, both tiles fit comfortably in shared memory and are placed there accordingly. For Software Pipelining. TILELANG disables pipelining on the MI300X to support larger tile sizes and reduce register pressure, whereas on the H100, pipelining overlap. Tile sizes are also adjusted accordingly to fit each platform's resource constraints. For Warp Partitioning, users may initially adopt a default policy for the two gemm operators, which often leads to sub-optimal performance. TILELANG addresses this by analyzing the underlying hardware and recommending platform-spede gemm operators use the FullCol scheme, partitioning acc_ and acc_ vertically to match the Tensor Core shape. In contrast, TILELANG applies a Fu FullRt FullRow policy on the MI300X, partitioning tiles horizontally.
enabled maximize pipeline
partitioning strategies, as illustrated in Figure B On the H100, both
17

---

Published as a conference paper at ICLR 2026
KV shared,
000
Sharee,
MlockN
KV ienewrr
Warp
KV Ailandd
Partition
Warp
RV
transpose. B-Trus pollyyFFllica)
acc Irep
Teapy
blockN
ROC_O pegutao
Warpgroup
Toogy!
Elocke
Werparoup
acc_D pegitrac
progestett
KV Doladd
Tgemmm
T.garmm
T.gemm
5 chared,
Papatedd
KV W.,hhared,
acc
Werpghoup
SCG
Warpsroup
PolloyeFulCCe)
(a) The Warp Partition schedule recommended by TILELANG for MLA on H100.
blockN
Weep Partition
cim
Portition
Tgeemm(
Q shared, KV_ shared,
blookt
KV
KKt jahaned memmrr!
(aherod! mameeyy
0,1
blockN
acc Papolee
acc_ (rgiend
transpose B=True, polley=FullRom)
blockM
Warp
a00,
acc_o (regitter)
Prppoted
ja
(rughtter)
T.gemm( acc S, sisahaee,, acc policy=
T.gemm
pagtener
T.gemm
Warp 2,3
RCC. raginted
ACC_ . Drepteteer
Warp 2,3
(b) The Warp Partition schedule recommended by TILELANG for MLA on MI300X.
Figure 11: Cooperative workflow between tile-recommendattin and inference stages NVIDIA H100 and AMD MI300X GPUs
E LAYOUT INFERENCE ALGORITHM
Goal and Scope. Algorithms and compute hardware-aware, globally consistent layout map- ping over a fused tile graph (FTG). Starting from partially annotated buffers and operator semanties, the pass infers concrete memory layouts hhobal//hared/registter fragments), resolves aliasing, and re- turns both the final layout map L and the loop-level binding/predication maps (ForMap, P PredMap) used by the loop-lowering routine in Algorithm 3 The inference explicitly respects coalescing, bank-conffict avoidance, and Tensor-Core- foienffrindll register tiling, while minimizing register pressure for the remaining degrees of freedom.
Three-Phase Inference.
The core procedure IaayutIInereecce executes in three stages:
Phase I: Strict constraints ("STRICT"). Each operator visited once with RunInferStep at level "STRICT". This enforces hard constraints dictated by hardware- sensitive primitives (e.g., swizzled shared-memory layouts and MMA-aligned register tiles for GEMM). All layouts fixed here recorded into Lstrict and treated as immutable thereafter.
Phase II: Common propagation ("COMMON"). A worklist over all operators drives fixed-point propagation. When an update materializes for buffer, its users are enqueued. This phase spreads compatible, non-rigid constraints across the FTG until convergence, ensuring consistent thread bindings and compatible address formulas across producers and consumers.
Phase III: Free choices with register-cost minimization ("FREE"), Remaining uncon- strained buffers are partitioned by connected components (w.r.t. uses and aliasing). For each component, the algorithm enumerates candidate "roots": it snapshots state, seeds inference from root with level "FREE", and greedily extends to other members. Can- didates that trigger a conflict (layout mismatch or iterator normalization errors) are dis- carded. Among feasible candidates, it selects the minimizing total fragment registers via SumF zagment lightweight, local backtracking that controls search while favoring low register pressure.
then restores and commits the best snapshot. This realizes a
Key Subroutines. RunInferStep constructs the per-operator context (target, thread bounds, current T., analyzer, and any out-of-hound info) and calls op InferTayout (args, level)
18

---

Published as a conference paper at ICLR 2026
Algorithm Hardware-Aware Layout Inference over Fused Tile Graph
I: procedure Layoutlnference(FGG, AnnotatedLayouts, Target)
2: (Ops, UseList, AliasGroups, Var, ThreadBounds, BufferOOB) CollectFrom/R(FTG Target)
3: I. = l.=coy(AAnotttedl.. .ayouus
4: =_strict==mpty_mapp
Phase Striet constraints
6: for = to 10psl do
7: RunlnferStep(i. "STRICT", false, L, L strict)
8: for all (huf, lay) in I. do
9: I._strict[buf] lay
10:
Phase II: Common propagation
11: q=queue_of of all op indices()
12: FinishlnferQueeee"" "COMMOO" L.,ttrcc
13:
Phase III: Free choices, minimize register cost 14: for all comp in ComnetediCompooents(PPs UseList, AliasGroups) do
15: best = (infinity. null)
(reg_cost, payload)
16: for all root in comp members do
17:
ops_bak- Snapshot(Ops)
18:
I_ copy(L)
19:
ok = Try/lnforFromRoootrooo, "FREE", L_tmp, L_strict)
20:
if ok then
21:
for all other in comp.members. other root do ok ThylnfeerFromRoootahhee,, "FREE", L_pp,, L strict)
22:
23:
if ok then
24:
cost = Sww.urragmennngiiststnlnmmmm candidate = (cost, (Snapshot(Ops).. L. L_tm))
25:
26:
best = Regeeeeesseeeeeeeet,,,,, candidate)
27:
Restore(Ops, ops_bak)
28:
assert(best I= null)
29:
lop__sapp L_best) best,payload
30:
ApplySnapshot(Ops, ops_snap)
31:
I best
32:
D Alias completion
33: for all (var, buffers) in AliasGroups do
34:
if exists in buffers such that L[b] defined then
35: ref= L.[b]
36: for all buf in buffers do if [[buf] is undefined then L[buf] ReshapeliNeeededree,, shape(bul))
37:
38:
39: for all (buffer, dummy) in Usel do
40:
if scope(buffcr) "local.fre fragal..f then
[sseri[[[buffer] is defined)
42: (ForMap, PredMap) ColllccLLoopaaauussssaPPPaaaaaas
return (L. ForMap. PredMap)
19

---

Published as a conference paper at ICLR 2026
Algorithm 2 Core Subroutines for Layout Inference
I: procedure RunInferSteplop p_d, level, update_ queie, L, strict)
2: op Ops[op_id
args = (target, thread_bounds bounds ThreadBoundsjop.idd layout_map L, analyzer, buffer_oob Buffer-
OOB[op_id|)
4: updates = anp.Inferlayout(arggs level)
b list of (buffer, layout)
5: for all (huf, lay_ new) in updates do
6:
if buf in L then
if scope(buu)) = "local "ragee...agmmen"" and level "STRICT" buf not strict then if TragneennContaiiss([[bu]] lay_new) then L/buf] lay_new
8:
9:
10:
PropagateAllas(but, lay lay_new, L, update_qqeue)
continue
12:
assert[sEquallLayout(fffou lay_ (ay_)ew)
13:
PropagateAllassbuf,, lay_new, new, L, update_ queue)
14:
else
15:
L[buf]] lay new
PropagateAlias(but,, new, update queue) if qudate__qeuee queue then
17:
18:
enqueue all users(bul)
19: procedure Finishlnferuucuc(lvee,,, L, strict, q)
20: while q is empty do
21:
id q.popO
22:
RunlnferStep(id. level, true, L, strict)
23: procedure PropagateAlias(src buf, sre_ layout, L, update_queue)
24:
for all sib in AliasGroupsssre_buf..ooo bufraeee var] where sib != buf do if shape(src layouuy == shape(sib) then
25:
26:
1gt sre_soyluouu
27:
else
28:
tgt = Reshapo6sre_layout, shape(sib))
29: if sib in then
30:
assent(IsFquall aaoou((_[ssssby tgt))
31: else
32:
[[[ibb] if update queue then
33:
34:
enqueue all users(sib)
35: procedure TryInfcrrroRRoorrroot, level, I__mp, l_strict)
36: success false
37: try
RunlnferStep(root, level. true, tmp, strict)
39: FinishlnferQueueelleelll L tmp, L strict, q)
40: success = true
41: catch LayoutConflict NormalizclterError
42: success false
43: return success
44: function SumFFragmennRgisteess((]]
45: total 0
46: (buf, layout) in do
47: if layoutkkndd = "Fragment" then
48: total = total + product(ayouutoourpu_gggg
49:
return total
20

---

Given a loop and Published as a conference paper at ICLR 2026
Algorithm 3 Loop Lowering:
Binding, Vectorization, Predication
1:
hayp:paoopplaaaaaaaaaoooomoooooo ForMap, PredMap, thread_var var)
2: loop_laye hap_layutt=-- F
3: parallel_low
(not skip_thread p and
(not local registerr regisserr_onll(oorrrr only(ForLoop)) and
(not store_into sos_into_ooccaaFong Joorl(Forl .oop))
4: if parallei_looy then
5:
Forl.oop Partitionl wwrttttooooooooooooolooooo .oop(Forl.oop. Loop, thread var, analyzer, layolty
6: mos_non__ocae
touches_noot Bacresssoo__oooooeeeeeee _or_ooocaaooorooopppp....
7: has has_reducer cetoontiinsseduuereeeteee htt./ooroooo.oo...........
has mas_cast__oe
cast_ops- contains containnnnooneedocciiinn _l____atsst_____ store(Forl SarreFool..op.boor .oop.body)
9: if thas_non_loceh or has_cast_ops) and (not has reducer) then
10:
ForLoop= = VectorizeLLoop(ForL.ooo
11: if ForLoop in PredMap and parallel_loog loop then
12:
return (PTTeenFFsseePettMaaFFooooo Loop], ForL.oop)
13:
return ForLoop
to obtain layout updates. Updates are merged into L with two safeguards: (i) for pre-existing buffers, strict equality is required unless the buffer a register Fragment and the new layout contains the old one (per FragmentContains when not in "STRICT" and not locked by Larieti (ii) every update triggers PropagateAlio which reshapes
in which case safe refinement is allowed
sibling shapes as needed and enforces alias-wise equality, enqueueing users when in worklist mode. TryInferF cownnerrFrooooottooooooooooooooo runs a guarded, queue-based inference seeded chosen root and catches wwt.yyuuttooolc////oooooooo
ze zetrerrrro mark a candidate infeasible. SumF zagent Regist accumulates the product of each fragment layout's output shape, serving as a proxy for total register footprint.
Alias Completion and Validity Checks. After the three phases, each alias group is revis- ited: if any sibling has a layout, the rest are filled by reshaping that layout to their shapes. The pass asserts that all local. fragment buffers used in the IR are defined. Finally, LollectLoopLayout sAndP redicate summarizes loop-level binding decisions and out- of-bounds predicates into (ForMap, PredMap).
Loop Lowering: Binding, Vectorization, Predication (Alg.
(ForMap, PredMap), the lowering proceeds as follows:
Thread binding. If the loop is parallelizable (not skipped, not purely local-register, and it touches non-local memory), Partit iertitiins iitiiiioooooppp binds iterations to the hardware thread vari- able using the loop layout from ForMap, aligning with the previously inferred thread/block organization.
Vectorization. If the loop body either touches non-local memory performs non-reduction cast stores, and there is no reducer present, fectorin zeooop is applied. This realizes the vector length implied by the chosen layout, improving coalescing and matching hardware vector widths while mitigating bank conflicts.
Predication. If the loop may encounter boundary conditions and parallel, the loop body is guarded with I TfhenFll 1.l.l....... se using the predicate from P PredMap, ensuring safe accesses without sacrificing parallel throughput.
Discussion. The division into STRICTICOMMONFREE keeps the search tractable: rigid, hardware-mandate forms are locked first; compatible information is then propagated to convergence; and only the remaining degrees of freedom are explored via component-local, anapshot-and-choose search guided by a register-cost objective. Alias propagation guarantees storage-consister ad- dress mappings. while fragment-aware refinement enables safe specialization of register tiling. The produced (c, ForMp, PredMap) bridge high-level tile indices and low-level memory/thread orga- nization, enabling performant, portable lowering across backends.
21

---

Published as a conference paper at ICLR 2026
Algorithm 4 Pipeline Inference
1: procedure PipelineInferennee// NumStages, InitOrderMap, InitStageMap)
2: OrderMap copy(Ini(OrdeeeMa)))
3: StageMap copy(IniiStageeMpp)
for all serial loop f.body do
5: if L E OrderMap and L E StageMap then
6:
continue
7: if T. & NumStages then continue
8:
9:
n NumStages L.]
10:
root L.body
11:
seq Flatten (oSeq
12:
Infos []
for 0 |seq] -
14:
(R. W. = RWCollec((seq(]))
15:
Infos.push closspuus_bbcc((SS50 (SfagelnfoRR W, i, C))
16:
$ = CollectCooyyReadsstte (Infos)
17:
PropagaieProduucers(tiioo,, S)
18:
ComputeLLas(USe((n((ss))
19:
order _d__iixx
20:
for all p in Infos do
21:
if FirstStage(p) and lastt_use * -1 then continue
22:
23:
p.order order_idxx order idx +- order_ids p.stage for all q Infos do FirstStage(()) and g.last_use _se = p.original_id then
25:
26:
g.order = order idx; order idx <- order idx + 1; g.stage = 0
27:
assert(order idx = IInfos)
28: & TailCopyCount(Ifoss
29:
if > 0 and 2 then
30:
for all p Infos
31:
p.order = (p.order + k) mod [Infos]
32:
if not p.copy and not p.producer then
33:
p.stage p.stage
34:
orders [p.order for Infos]
35:
stages jp.stage for Infos]
36:
ApplySSotwaarPipeeinee... orders, stages, OrderMap. StageMap)
37:
return (OrderMap, StageMap)
F PIPELINE INFERENCE ALGORITHM
Goal and Scope. Algorithm computes a software-pipelined schedule for serial loops that expose staged data movement and computation. Given input function f, an priori upper bound on the number of pipeline stages NumStages, and optional initial order/stage annotations (InitOrde rMap, InitSta ageMap), the pass produces pair of maps (OrderMap, StageMap). For each eligible serial loop in f.body, the algorithm assigns (i) total order index to every statement in the flattened loop body and (ii) stage id in (O...., 1}. where NumSt a [.), thereby enabling backend-specific software pipelining and overlapped execution copies and compute.
Loop Selection and Linearization. The outer procedure PipelineInffrrncc first filters serial loops: if a loop L already has entries in both Orde MMap and StageMap, or if it lacks stage budget in NumStages, it skipped. For each remaining loop, the body is linearized into a sequence seq via PlattenToSeg, which yields stable, single-pass order of statements. Each sequence element is then summarized into a StageInfo Info record containing its read set, write set, original index, and Boolean flag indicating whether the statement performs global-to-shared copy.
Read/Write Classification and Copy Detection. statement and classifies its memory behavior into three components: set of read regions R,
The helper RNCollect (Algorithm traverses
22

---

Published as a conference paper at ICLR 2026
Algorithm 5 Core Subroutines for Pipeline Inference
1: function RWCollec((stmt)
2: R. []: W = []: C. = false; within fälse; isg false
Visit(stmt):
4: on BufferStore(b,, idxs. v):
5: W += Region(h, idxs); isg false; Visit(v);
6:
if isg and scope(h) ("shared", "shared.dyn") then C' true;
BufferLoadd(,, idxs):
+- Region(b, idxs):
9:
if scope(b) = "global" and not within then isg = true;
10:
on IrThenElse(c, a. b):
11:
within truc; Visit(c); within false; Visit(a);
if b.defined)) then Visit(b);
13: return (R, W.C)
14:
15: function Collec(CopyReadd(nneo)
16: set()
in Infos where p.copy do 18: for all r in p.reads do
19: S.add(r:buffer)
20: return S
21:
22: procedure Propagateerroduccessnnfoo S)
23: all p in Infos where p.copy do
24: upd = true
25: while upd do
26: upd false
27:
for all q Infos where g. g.copy and q.original_ids p.original_idh idx do exists w writes with bubuffer then
28:
29:
9. prooducer true; upd true
30:
for all in q.reads do S.add((:buffer)
31:
32:
33: procedure ComputeLaastsse((nnoss))
34: for all p in Infos where FirstStage(p) do
35:
for i = p.original_ie _ix + to [Infos] - do
36:
if exists in Infos, in p. writes with r.buffer w.buffer and MayConflict(--rciion, w.region) then
37:
p.last use = max(pplast__ueee _uess
38:
39: function TailCopyCount((nfo))
40: c= O, mn Infos]; mer 0
41: for all in Infos do
42:
if FirstStage(p) then
43:
c c 1; mm min(mnn,
44: else
45:
p.oder)
46: if mn
> mz then
47: return
48: else
49: return
50:
51: procedure ApplySootwarrePipelne(..,, orders, stages, OrderMap. StageMap)
52: OrderMap[L] []] f- orders
53: StageMap[ +- stages
54:
55: function FirstStage(p)
56: return p.copy or p. p.producer
58: function MayConflict(a, b)
59: return Intersect(IntSeer (IntSe((()n,,t))) IntSet(b)) * Nothing
23

---

Published as a conference paper at ICLR 2026
set of write regions W, and a Boolean flag C' for copy-like behavior. The visitor marks loads and stores according to buffer scope (e.g,, global, shared, shared. dyn), and uses a simple state machine over the control-flow context (within, isg) to detect global loads that feed subsequent shared-memory stores. Whenever Bu ferfferttrre into a shared buffer is preceded by such a global read, the statement is classified as copy (C = true), allowing the later stages to identify candidate prologuc/epilogue moves for software pipelining.
Producer Propagation and Lifetime Analysis. Given the per-statement summaries, CollectCopyReads aggregates the set S of buffers read by copy statements, which serves as the seed for producer discovery. PropagateProduce Prducerss then iteratively walks backwards over the sequence mark statements that produce any of the buffers in S before a copy: whenever non-copy statement writes to a buffer in S,, it is labeled producer, and its own read buffers are added S. This fixed-point propagation captures multi-hop producer chains that eventually feed global-to-shared copies.
Next, ComputeLastUse computes a conservative last-use index for each first-stage statement (copy or producer). For a given p. the pass scans later infos and checks whether any read region of a later statement may conflict with any write region of p, using MayConf 1. ict and an interval-set intersection test. The largest index where such conflict occurs is recorded p.last use, providing lifetime window that guides stage assignment and the positioning of prefetch-like operations.
Stage Assignment and Tail Rotation. With producer/copy labels and last-use information in place, P Pipel inn In assigns initial order and stage for each info in a single forward pass. The core policy is that first-stage statement p that is dead within the steady-state (i.e., Fi FirstStage(p) and p.last_use = 1) participates directly in the final pipelined schedule; otherwise, such statements are skipped in this phase until their consumers are placed. For each selected p, the algorithm emits the current order index with stage n (typically the last stage), then searches for matching first-stage statements & whose last use equals p.original idx and places those q immediately after at stage O. This yields an interleaving of first-stage and steady-state work that respects data dependencies and the inferred lifetimes. The invariant or order_ide [Infos] is asserted at the end to guarantee that all statements receive a unique order.
To further improve the pipeline structure, Ta MeilCooyCoout detects whether first-stage statements form a contiguous tail segment in the assigned order. counts the number of first-stage infos c, and tracks the minimum order index among them (mn) and the maximum order index of non-first-stage statements (mz). mn > mæ, first-stage statements appear strictly after all other work, and the function returns c; otherwise, it returns -.. When a positive tail count k is found and at least two stages are available (n 2 2), the algorithm rotates the schedule by l positions in a modular fashion and decrements the stage of non-copy, non- non-producer statements. Intuitively, this rotation shifts tail copies into the prologue while pulling steady-state computation earlier, yielding more balanced pipeline across the n stages.
Map Materialization and Backend Interface. After rotation, the final per-statement orders and stages are collected into arrays orders and stages, which are committed to the global maps via ApplySoft
ipeline, For each loop OrderMap [L] records a permutation of the flattened body, and StageMap [L] records a stage id for each element in that permutation. These maps serve as the contract between the high-level pipeline inference and downstream backends: code generators can exploit the stage structure to schedule prefetches, overlaps of global global-to-sharee to-shared copies with compute, and explicit prologue/epilogue code, without re-running dependence analysis on the original IR.
Discussion. The pipeline inference algorithm deliberately decouples (i) classification of copy and producer statements, (ii) lifetime and conflict analysis, and (ii) stage-aware ordering and optional schedule rotation. The use of region-based read/write summaries and conservative MayConflicn ict checks ensures correctness under aliasing and partially overlapping accesses. At the same time, the simple rotation heuristic (Tail (ailCoopyCount captures common pattern in GPU kernels where global-to-she th-rre transfers form a logical prologue or epilogue. By emitting (OrderMap, StageMap) instead of directly rewriting the IR, the pass remains backend-agnostic
24

---

Published as a conference paper at ICLR 2026
while still exposing enough structure for aggressive software pipelining and latency hiding across diverse hardware targets.
G EFFECTIVENESS OF THE COST MODEL
TILELANG employs an analytical cost model to prune suboptimal candidates and prioritize high- potential ones. This approach yields schedules that match or closely approach the performance of the best results of brute-force search exhaustive autotuning, while requiring orders of magnitude less tuning effort. For example, on GEMM-FPI6X FP16 workloads derived from models such as LLaMA-70B, TILELANG prunes 95% of candidate schedules, retaining only the top 5%. Despite this aggressive pruning, it achieves on average 98.47% of the performance (in TFLOPS) of the best configurations found by exhaustive search, substantially reducing compilation time with only negligible performance loss. This also serves as an evaluation of the effectiveness of our roofline- based analytical cost model.
M
N
K
Predicted-Topx Best (TFLOPS)
512
1024
8192
100.0%
512
12288
12288
99.9%
512
28672
8192
98.7%
2048
12288
49152
100.0%
4096
1024
7168
100.0%
4096
14336
14336
100.0%
4096
28672
8192
99.6%
8192
8192
28672
100.0%
8192
28672
8192
100.0%
16384
1024
7168
98.4%
Table Accuracy of our analytical cost model: predicted top-5% schedules retain over 98% the best performance while pruning 95% of candidate schedules.
H TUNING TIME
As demonstrated Table TILELANG leverages hardware-awarc recommendation mechanism to efficiently automate the design of high-performance computational kernels. The system achieves average tuning durations of approximately 10 seconds across both NVIDIA H100 and AMD MI300X accelerators. For the most complex operations, extended tuning times average 13.69 seconds H100 and 15.08 seconds on MI300X, reflecting the scalability of our approach under computationally intensive workloads.
We also conducted direct tuning-time comparison with Ansor/AutoTVM and Triton for the GEMM and 2D convolution kernel on H100. TileLang and Triton tuned with 20 configs, and the number of trials is set to 100 in Ansor. The results in Table Bp show that TileLang tunes markedly faster than both frameworks elppeill vs. TVM Ansor, which requires much longer empirical search. This improvement comes from TileLang's 's first-class tile IR, which defines a far more structured optimization space, and from our cost-model-quided inference, which avoids large brute-force searches.
Table 4: Average Tuning Times for Different Operators
Operation
GEMM
DequantGEMM
FlashMHA
FlashMLA
FlashBSA
H100 Time (s)
9.05
9.15
13.48
13.69
13.46
M1300 Time (s)
10.99
11.10
14.67
15.03
15.08
25

---

Published as a conference paper at ICLR 2026
Table Comparison of Average Tuning Times for GEMM
Operation
GEMMI
GEMM2
GEMM3
GEMM4
TileLang Time (s)
11.81
11.78
14.59
14.31
Triton Time (s)
18.43
18.24
20.15
20.07
Ansor Time (s)
518.52
455.51
3007.00
4142.05
Table 6: Comparison of Average Tuning Times for Conv2D
Operation
Conv2D1
Conv2D2
Conv2D3
Conv2D4
Conv2DS
Conv2D6
Conv2D7
Cony2D8
TilcLang Time (s)
11.56
17.49
17.65
17.41
19.18
18.87
17.50
17.21
Triton Time (8)
17.56
18.76
37,59
38.57
37.61
18.82
19.13
18.97
MATMUL IMPLEMENTATION DIFFS: TVM vs. TILELANG VS. TRITON
TILELANG achieves significant code-size reduction through its fundamentally different tile abstrac- tion. Instead of manipulating raw pointers, TILELANG represents tiles as first-class IR constructs, endowed with explicit semantics for indexing, data movement, and pipelining. High-level primitives such copy, gemm, and pipe innd enable the compiler automatically perform address com- putation and pipeline orchestration. In contrast, Triton requires programmers to manually manage memory access via pointer and offset arithmetic, resulting in lower level of abstraction.
Moreover, TILELANG supports maintaining execution context via T. Kernel, which obviates the need for developers explicitly compute grid dimensions and launch kernels. This design choice further reduces code size by eliminating boilerplate associated with kernel invocation.
In TVM, matrix multiplication kernels are typically expressed through simple tensor-compute defini- tions. Performance optimization is achieved by applying hand-written schedules, which can transform the computation into nigh-performance form. While this manual scheduling process may require tens to hundreds of lines of Python code, TVM also provides automatic scheduling mechanisms such as Ansor to explore schedule configurations. After scheduling, the tensor expressions are compiled and built into executable code.
However, TVM's scheduling inderinngnddeedeeent compute-expressine abstraction has limited expressiveness for certain operators, such as FlashAttention and irregular sparse kernels. Automatic scheduling in TVM also faces challenges when dealing with very large search space and when targeting new hardware backends with insufficient operator expressiveness. TILELANG addresses these limitations by employing a human-ing in- -iat-iiiittt--oppp tho-loop methodology to enhance expressiveness for complex and irregular workloads, and by leveraging cost-model-dmie driven scheduling approach to mitigate the search-space explosion issue.
J COMPARISON OF FTG IR AND TVM IR
Figure 13 (a) shows FTG IR produced after the Pipeline Inference Pass, which explicitly ma- terializes a software pipeline with a prologue (prefetch the first A/B tiles and clear the accumula- tor), steady-state loop that overlaps compute and prefetch, and an epilogue (final compute and waike-back). This IR emphasizes schedule semantics and readability: block tiling, shared- staging, and double buffering are expressed with concise, high-level primitives such as T. Kernel, silloc_sharede
loc_ fr aragtett T. clear, T. copy, and gemm. The same semantics (b), which preserves the pipeline structure while making GPU execution details explicit: b dx/ wwwdddddhheedx bindings, scoped buffers in shared/local memory, unrolled initialization, asynchronous copies (e.g., cp async) with commit/wait groups, and spreic- GEMM call. In short, (a) captures tile-wise schedule for readability and porta- bility, whereas (b) exposes fine-graa fnn-grrined GPU mechanisms for maximum control and performance, structure.
are then lowered to TVM IR Figure (b),
without changing the pologue-steady--tttee
26

---

Published as a conference paper at ICLR 2026
TileLang
3 Triton
8,2t,, bupte, CLDTT,
strideaan an stcidiaak,
445 Jecrett
3: T.tossccoor, ky, Care)) T.tesectss 80, cise).
ACTISATIIN:
1
vit C. osttesess Sieck,0, T.oeticcves, staduit is Cit, bi::
LLodio BL0X
A-hhece Trallic-theeetthhee Slock.t) depss
man 000 LL.oii BLOCK SIZE
12
C.lecat T.allor /twp://pttitiitttt......mm block.5), TechesetC leca))
2ld
25
firxt.ple, pit,n
200P 5325
22
yor to sa wew.ldddtt.-.tiiieee stack
E150 2i2.. vincran p1..
S14.E Sixet_aleee (cold
T.crrgtilie blacc, stasck, (.sterd)) T.pancc a,caame (.locat)
3id (pis
35
36
T.crpue Lacal, 1034 Stock block_a)
offici- + GHiLi HLCC
ofFLA SLAACAACE FLOOL
11 Mo0A Lenrecc,
aCOUMLAAEE CO ST2E 5.00 1000..02
digeetl fooa523
TOr In cangREO, 11 CESHOC, SCEE
TVM
BLOCK. ether-d
PaSCN accoralatoc?
eccurslatzz
piza 2.00K 5225 stoade
te-pl Sw..aaaaensiiee O "A", stgpe)
paza BL00K 522E DOC.
(e. 10, name- 13", degee In, /nggos
te-cenare axtaCC (), name (1")
affICCE sig BLOCK SEDC
BLOOL SIZE 30
45 affecn pid,n SC2,, sranDOn BLBOK 30
c,am ELETE srride_on offeLLEE Moo)) creideee =ff :]
laatea to oo,es, Mlarrpptacccu..oom coppp)
CAAA CoZES.On)s Neas) CoFFL_nnieee
Rik, D. paatppececcnn segeo,, aalsen).
17 445 natmoice, 10:
11. Punc cwwatouee fanttaangi
angs-
HETAI
furcto
14
ccscurable bwn.complaammmtte
15 A, D, C: teech. Teross
buasesdaro), b.STESONEL), c.asosaac)), C.NcECdRCL), ACTIVN activNtion
12
Figure 12: Side-by-side diff of minimal GEMM implementations in TVM, TILELANG, and Triton (first page).
S7 pela
def formit
ST.peta.frre def Hetnoll A: T.TesseecON, 10, *140018'), B: T.TenoorDD,, 1), "fostti"),
A: T.Tosear((M, 10, "Mastt))))
B: T.Toxeset(K, 10, "Maatto0")
C: T.Toxssecch, 16, "Maaat""").
C) TeTseeeeo 10 *Cootte
):
with T.terrel
cdlli Mldul),
colldiie ModuO) theadstirrals) as Dx, 19):
3:
ALears + T.decl anrfeee x (ledi,.)) stype, cossstehared"))
be - Titheved coataitt.caalditt Heeck H), "pleck2ak.x")
C. - T.dect baffarcc12s)) 2, accam ctps, scspe="locat")
RLmares - surfeccarre x B.ed,..), stape, scoss""tteeed")
A Tallaa shacve((biacc 1, Blek 10, Micesto")
Buffee Allocation
22
28
23
by t.teod bisding culldii Hleck "sleck2dk.y)
for in lu
Lid I. E.theoal bi "Uhvoollkk."))
B.atared T.allas atavee/Cacct i, bleck Mices2er)
24
13
C.accal23] F.
ia T.ueol1(0)
Claaa T.allec.eepenntCoao Heduu)) "Flcalte")
36
35
T sp. "urrtt 0.a cace, 9, 3isate, 0, Lptt..cmmii_gpwwa
T.sts sp. Aanett, A.maared caca, 9, A,sata, 6,
27
T.ctasr(C_eeca))
T.seraaleeiieeee black.0 10:
Teyyyykiit
blacka)),
*et::gemmser33 128,
blacks
T.cgy(Kee MasLL MaccA)], ksharel)
123,
Rats Leep with Pipeldee Amostatiee for T.metallcesteeed K03: shared, R,shazed C,teral)
T.westi())
ppt.rpppssncc "itoaass""""iitt"" A.staved data,
26
jit.sppacyyc"siinnttt R.staved data, f.stu_comit_grrusp
ceppybbby Strck., Kl, a
24
26
Tatt croppt)
CeDp a stackk, by strckul), 8.svare)
22
T.calt, estuant
34
31
Cassate the last stage
22
22
T. L.ten.acoooo atr date, bladktt bledull 1), L. ptr (t.aheee data, bladka kieduk 1), t.teraccssstti (tt.ccostttttt.ooaaa... (e_oooa 128,
"tl:igumm sar128 128,
35
T.aeenck, shazed, a,staced, c.tecat)
34
37
36
Casy the seout to the octout buffee Leea, clay stoduut, be btodk()))
37
15
16
3
(a) Fused Tile Graph IR
(b) Lowered Tensor IR
Figure 13: A byde-by-side comparison showing how tile-level FTG-IR is lowered into Tensor IR.
27

---

Published as a conference paper at ICLR 2026
K COMPARISON OF FLASHMLA IMPLEMENTATIONS ON DIFFERENT ARCHITECTURES: NVIDIA vs. AMD
and Stare
(hi6,
14
14
13
000, (utwo))
/mell 1815 stre Wee. mn wec.hceeed
Far 1. (a0
For in Toliiile (hha
1 atl, 2i Candlie
Bor k. la "Nonstas Sdeck
all cares
Figure 14: Comparison of FlashMLA implementations targeting NVIDIA (left) and AMD (right)
architectures.
Figure 14 Iillustrates the code-level divergences between the FlashMLA implementations for NVIDIA and AMD architectures. While the high-level algorithmic structure remains unified, the implemen- tation diverges to exploit distinct architectural strengths. First, Region highlights the memory scope allocation for Query (Q) tiles, where the NVIDIA backend utilizes shared memory while the AMD backend prioritizes register memory. Similarly, Regions 2 and 5 collectively illustrate the divergence in intermediate accumulator storage: NVIDIA buffers these results in shared memory, whereas AMD maintains them directly in registers. Regarding output strategy, Regions 3 and 7 show that NVIDIA stages the output tile in shared memory to ensure coalesced transactions, in contrast to AMD, which performs a direct copy from registers to global memory. Finally, Regions and 6 depict the adaptation of GEMM policies, employing a FullCol strategy for NVIDIA and FullRow for AMD) to ensure optimal instruction performance.
L COMPARISON WITH RECENT SYSTEMS
Table 7 Tevallates Causal MHA on an H100 (B- 64, H 64, D 128). TileLang more concise than Tilus (Ding et al. 2025) (66 vs. 83 LOC) while achieving a 59%-75% performance gain. Table 8 evaaluates the same MHA on an RTX 4090 (B- 16, H 32, D 128). TileLang achieves the best performance among baselines and outperforms the latest Tilus by 3%-4% with significantly
fewer lines of code.
seq seq
Tilus (ms)
TileLang (ms)
1024
6.85
4.29
2048
25.65
15.69
4096
97.14
55.44
Table 7: Performance of Causal MHA on H100 (B 64, H = 64, D 128).
28

---

Published as a conference paper at ICLR 2026
seq
len
FA2 (iflops)
Triton (Ulops)
Tilus (tflops)
TileLang (tflops)
MIHA
1024
137.58
106.61
133.82
143.40
2048
151.46
121.44
154.36
159.03
4096
162.70
129.75
159.26
162.79
8192
164.10
134.40
161.60
165.56
LoC
389
197
393
138
Table 8: Performance and Lines Code(LoC) of Causal MHA on RTX 4090 (B 16, H 32,
128).
(M, K)
TileLang (ms)
Gluon (ms)
Helion (ms)
Tilus (ms)
8192, 1024, 8192
0.16
0.30
0.26
0.31
8192, 8192, 8192
2.22
1.64
2.97
8192, 28672, 8192
5.09
7.82
5.87
10.19
8192, 8192, 28672
5.09
7.77
5.89
10.02
Table Comparison of TileLang, Gluon (LoC 68), Helion (LoC 24), and Tilus (LoC 110) on GEMM workloads.
seq_len len
TileLang (ms)
Helion (ms)
1024
0.17
0.19
2048
0.33
0.36
4096
0.65
0.71
8192
1.28
1.41
16384
2.53
3.01
32768
5.08
5.56
Table 10: Performance of Mamba-chunk-sen on H100, with batch size 80 attention heads, model dimension 64, dstate 128, and sequence lengths ranging from 1024 to 32768. Helion LOC=116, TileLang LOC=114.
29

---

Published as a conference paper at ICLR 2026
examples that support execution on Hopper, covering both GEMM and Mamba-chunk scan workloads. As shown in Table 9 TileLang achieves 1.15-1.62x, 1.52-1.83 and 1.87-2.12x speedups over Helion, Gluon, and Tilus, respectively.
We evaluate the latest official Helion PPyTorch 2025) and Gluon (OpenAl 2025
On Mamba-2-chunk-scaan TileLang further provides 1.10-1.19x speedups Helion, reported Table [0]
The differences attributable to design limitations in existing DSLs. Helion lacks effective tilerrcoommendation system, making optimization difficult and causing long tuning times (over 20 minutes). Gluon lacks appropriate abstractions and interfaces for pipeline scheduling, which makes it difficult to achieve effective overlap of computation and memory operations. Second, Gluon operates lower programming abstraction level, requiring users manually make larger number of design and optimization decisions. As a result, writing high-perte high-performance kernels in Gluon is considerably more challenging. This also reflected by the fact that the Gluon implementation of GEMM requires substantially more lines of code than the corresponding TileLang implementation. Tilus exposes brread-lock-level control over shared memory and registers, but without tile abstraction and critical tile-related optimization. In contrast, TileLang's tile-recommendation mechanism efficiently identifies good tile configurations, and its overlaps computation and memory I/O. These additions provide clearer qualitative and quantitative comparison among recent DSLS and further highlight TILELANG's performance and usability advantages.
inference strategy generates an effective schedule that
M BASELINE COMMIT HASHES
To ensure reproducibility, we provide the specific commit hashes for the baselines used in experiments that were not specified in the main text (see Table II
Table 11: Commit hashes for baseline frameworks.
System
Commit Hash
Gluon
Helion
Tilus
Ansor
ThunderKittens
Marlin
Block-Sparse-Atteen Attention
Kerpe
Tilus (Variant?)
FlashAttention-2 (FA2)
TritonBench
N OPERATOR SHAPES IN OUR BENCHMARK
30

---

Published as a conference paper at ICLR 2026
Table 12: Matrix shapes in FP16 Matmul evaluation
DO
D1
D2
D)
M
8192
8192
8192
8192
N
1024
8192
28672
8192
K
8192
8192
8192
28672
Table 13: Matrix shapes in our Fused Dequantize-Matmi evaluation
MO
M1
M2
M3
M
N
1024
8192
28672
8192
K
8192
8192
8192
28672
Table 14: FlashAttention and Block Sparse Attention(with 50%, 90% sparsity) shapes our evalua-
tion
FAO
FA1
FA2
FA3
FA4
FA5
FA6
FA7
batch
64
64
64
64
64
64
64
64
nheads
64
64
64
64
64
64
64
64
seq_len len
1024
2048
4096
8192
1024
2048
4096
8192
head dim
128
128
128
128
128
128
128
128
causal
false
false
false
false
true
true
Table 15: FlashMLA shapes in our evaluation
FMLA0
FMLAI
FMLA2
FMLA3
batch
64
64
64
64
nheads
128
128
128
128
seq_len
1024
2048
4096
8192
head_ dim
512
512
512
512
pe_ e__dim
64
64
64
64
causal
false
false
false
false
Table 16: Convolution-2D shapes in our evaluation
ConvO
Conv1
Conv2
Conv3
Conv4
Conv5
Conv6
Conv7
N
128
128
128
128
128
128
128
128
C
2048
512
512
512
256
1024
512
64
H
7
7
14
7
14
14
28
56
W
7
14
7
14
28
F
512
2048
512
512
256
256
128
64
K
3
S
I
2
D
1
P
O
O
O
G
I
I
1
1
HO
7
7
7
14
14
28
56
WO
7
7
7
14
14
28
56
Count
3
31

---

Published as a conference paper at ICLR 2026
Table 17: Chunk-Gated-Dellae. kernel shapes in evaluation
CGDNO
CGDN1
CGDN2
CGDN3
CGDN4
CGDNS
batch
1
1
64
64
64
nheads
32
32
32
32
32
32
seq len
16384
32768
65536
1024
2048
4096
head_din dim
128
128
128
128
128
128
Table 18: Vertical Slash Sparse Attention shapes in our evaluation
VSSAO
VSSAI
VSSA2
VSSA3
batch
nheads
seq_len
8192
16384
32768
65536
head dim
64
64
64
64
vertical size
1000
1000
800
1000
slash size
600
200
600
600
Table 19: Attention Sink shapes in our evaluation
SinkO
Sink1
Sink2
Sink3
batch
1
nheads
64
64
64
64
kv kv_heads
8
8
8
seq_len
1024
2048
4096
8192
@v_seq_len
1024
2048
4096
8192
head dim
64
64
64
casual
true
true
true
true
32

---

Published as a conference paper at ICLR 2026
O KERNEL IMPLEMENTATIONS
O.1 MATRIX MULTIPLICATION (MATMUL)
BLilelang, lit
(A leneso,, B: I.lensor, :.Tensoz): with T. Kernel (N // block_N, м // block_M, chreadswthreads)
byl:
A_shared c.alloo__haaredbbll
block_K)
B_shared r.alloc_shared Kk block_N)
c_lccal frallossfragmentt (black_M, block_N)
9
T.clear(C_local)
T.Pipelined(K c.cypy (A]by bleck_X, & [], shaaardd r.copy block_,, block_N], B_shared: T.genm (Aaaha, B_shared, C_local)
sun_stagesm2):
11
13
15
T. I.copy (C_local, Ciby bleck_X, bx block_N))
Figure 15: Kemee Implementation of Matrix Multiplication.
O.2 DEQUANTIZED MATRIX MULTIPLICATION
BLilelang. fit
2
def dequant gemg (A: T.Tensor, B: Tenena C: Te
with Kernol (T coildiv p_partitionn, threads= (raduca_thread, n_portittonn) (bx, fmmierossiz_____ R_quant_local T.alloo_lomal (
A_local T.allce_locs
in_dtypet
storaga_dtype)
B_dequantias leepuaantie_llcall ftalloo_lcoall l (mmiro_mmmeeeeaaeeeeee in_dtypet
F.allos_looall ((1,), accum_dte sccum_dtype)
reducced_acoumrres Io accun_dtype)
10
T.clear (accun_res] for [T for v in (T.oillvv.., black_K11 :
12
(micro_sise__k)
1
Alby, block_K kr micro_size_k
for vectorizzed (mierro__a compprssloopss
8 quant leclloollll __qanttlocall [v] b[
n_partition
ko (reduce_threat + bearo__siieekcoooreeeeeee
micro_ ctt_r_sttee___ereaaeeeeee v,
T.call_extern
T.address_co qua__aaaa_eaallll
T.address_of (B_éequatiire
dtype=in_dtype,
for ki in T.serial (micro_size_k):
accum_res [0] +s Aocool [ki] B s_dequattize_lcaalpe
with T.altr(
recuureduces (lambda
y, (2. (accum_dtypee 0111,
preduce_scope"" r.reintespret
uin:: (Esiie
htnde="""
11
T.evaluste:
T. Lewa__theee_allreeeee
2.uint32
accum_res[0]
True, reducee__cccunrres([0 kr, dLype="handle",
1)
kz O= Cloy, bx n_partitiion + mi] reduoed__ccum_
Figure 16: Implementation of Weight-Only Quantization (WFF4 ApPi6) Matmul using TILE- LANG, showcasing support for mixed-precision computations via simple form.
33

---

Published as a conference paper at ICLR 2026
O.3 FLASH ATTENTION IMPLEMENTATION
Btilelang lit
def saattattentiion T.Tensor, T.nensor, Tenessrr, Output: T.Tanso")
3
with [eRernll
(div black heads, batch, threadsathreads)
by, 0z)
Q_shared 1.allca_shared [block_N, cim], dtype)
K_shared T.alloo_sharedi [blook_M, 61m], drype)
Y_shared T.allee_sharedi [block_N, cim], dtype)
0_shared T.allloeshaaree (blookkM,, 61m], drype)
9
acc_s = s.alloe_frogmant (block_M, block_8], accun_dtype)
acq_s_oast fragmmment
block_8l, dtypet
acc_o = t.alloc_frogmans ([block_M, dim), accun_dtype)
7.alloo fragment (bbookkMM) socun_dtype)
scores_max_prase T.alloc_ fraggmennt [blook_M), stype)
scores_scale soares Toa ([block_M, dppe)
scorcs_sum = F. logsum a.aloe_fragment ([block_],, scsum_dtype)
(block_M), accum_dtype)
T.copy (Q,b,, bx block (bx 1) block_N, (_share)) T.f:ll T.fill(acc__ (acc_e, 0)
19
T. (logsum,
T.fill (scoress_c
Loop_range = cesiddireiiii bleek_N), T.ceildiv(
(bx 1) bleck_X, block(N)) if is_causal else (.ceildiv block_NI)
for in P.Pipelineddlooppc range, nun mun_sttgge-nuunLtagee T.copyiK[bz, k bl
block_N, by, K_shared)
if is_causel: waarallllibbocke...... block_8): for acc_sti, jl then (bx bleck_X >- block je C,
else:
taclear
T.geam|Q_sheres K_shared, acc_s, Cranspose_D-True, T.copy (acores_oaeeaee acorea_nak__rev)
FullRRwi
[.fiillssoore_eal reduooo__a (aco__ scores_maX, cim=i, clear-False)
for in P.Parallel scores_scalin for i, j in I.Parallei ace_sli, T. redd:co___u (a6c_5, socres_sum, dim=') for in T.Parallel : logaum(1) logsum[1]
beeppp/ccoree__ax_reee scale scores_max scale)
block_N);
T.exp2 (acc_alá, 1] scale scores_max[1] scalel
scores_scale(t] acores_sum[1
T.copy laco_s, aco_s_caat]
] T.Parallel (block_M, dim): ace_oli, v- acorea_scale_i]
T.copy (bb
block (i
b'ock_M, by, s), V_shared)
T.geamlaco_s_c c_oast, V_shared, ace_o, Gem.yy-T. PullRow) 1, T.Parallet (block_M, dim!:
acc_oli, j] /- logsum[i] T.copy (ace_e, 0_ahared)
T. T.copy (O_shared, Output [bz, bx black (bx
block_M, by, :])
Figure 17: Implementation of Flash Attention with TILELANG.
34

---

Published as a conference paper at ICLR 2026
0.4 FLASHMLA IMPLEMENTATION
3t BLilelang jit
def Flash_nn ((ooo(((aaiiii (lbatch, heads, dimi, dlype),
3
0: T.Tensor
T.nenaor ([batch, beads, se_din], dbype),
KV: Iensor (botch, seqlen_kv, kv wy_hhed_niie num, dim), dtype), Tenaor seqlen kv_head_num num, pe_dim), dbype), Output: r.rensor I. heads, dim,, dtype),
9
with K.Kernei (batch, heads // min (block_H, kv_group_numl, threada-256j as (bx, byl: 0_shared T.alloe_sharedi (btrodddiilooo___ 61m], drype)
S_shared T.allct_shared [block_H, block_N), ctype)
Toal (b pe_dim), dtype)
KV wV_ahered w.w.llo__hhareeeee (black__ dim, dtype)
"_pe_shared T.al (ha
pe_dim), dtype)
0 o_shared = = (block_ cim], dtype)
I.e (block__,, block_Nl, accun_dtype)
scc_o = r.clloe_ ( (block_ dim), aecun_dtype)
I.elloc_fragman: (bboock_H) accun_dtype)
19
scores_max_prav = I.alloo_ (.blosk_)) accum dtype)
scores_scale I. ([block_H,, ascus_dtype)
scorcs_sum T.alloo_ (block_H), [), aecum_dtype)) logsum w.alloc_fragmee Hblock_]),, sccum_dtype)
cur_kv_head T.use_swieelee10)
Iky_group_num block_H)
T.copy (bx, by VALID 3LOCK_I (by 1) VALID BLOCKJI, :h, (_shared)
T.copy (2_pe [bx, by BLOK (by 1) HALKK_ (_pe_shared) E.E1ll(ace, 0)
T. (logsum, T. (scorrs__e_ T.infinicy (egeeniiiyyecummmtyeeeee.
loop_range T.ceildiv block
for in
nun_stages-21:
T.copy 1&V (bx, & b:ock_N: 1k 2) block our sur_kv_head, :), XV_shared)
T.copy [bx, k bi T.clear (acc_s)
block_N, cur_xv_head, t], K_pe_ahazed]
T.genm!
Q_shared, KVLshared, acc_s, Keenspose_B-True hewpr:/eeeeeeeeeyyy..o....
T.genm!
Q_pe_shared, K_pe_ahared,
ace_s,
transpose BBrrue,
pollcy=T. GemmKarpPol
T.oopy [ayyreerr___,, acorea_nak__ree
T. [scores_mak, #T.infintty(accmm_dyya
T. reduce__aae (aco_á, scores_max, cim=i, clear-False)
T.Parallel (block_H)
scores_scalel1
scale scores_max scale)
in TeParallel (block_R, (lock_N):
ace_sli,
T.exp2 (acc_alı, il scale scores_max[1] scalel
redssss predsos_aum aastion_aummmmo (3cc_5, socres_sum, dim=')
54
T. cacpyy 8_shared)
in T.Parallel logsum[i] = lcgsum[i] scores_scale(s] + scores_sum(i)
(blok_H)) ock_H):
J. T.Parallel (block_H, dim):
acc_oli,
scores_acale'i]
T.gemm(S_aharren
WV_shared, acs_e, pal.ey-TT.ommoorrgeeeeeeeeeee hetpyrTT...mmooollll.l.........e
for i, in 1.n.T.T.. [.l.rrllllllblcck. diml:
aec_a(1, 3] /= logaum(1] T.copy (acc_e, o_shared)
T.copy (0_shared, Output [bx, by BLOCK bby 1) BLOCK_,, :1)
Figure 18: Implementation of FlashMLA with TILELANG.
35

---

Published as a conference paper at ICLR 2026
0.5 BLOCK SPARSE ATTENTION IMPLEMENTATION
3t BLllelang. lit
def blockaparse_atttn (O: T.Tensor, T.Tensor, T.rarsor, RlockMask: T.Tensor, Output: T.Tensor)
3
with KeKernll
Idiv black heads, batch, threads-threads)
by,
s_shared 1.allca_shared (block_M, cim], dtype)
K_shared T.alloo_sharedi [blook_M, 61m], drype)
V_shared T.allee_shared [block_N, cim], dtype)
0_shared w...lllllshaarreee (blookk_M 61m], drype)
9
acc_s h.slloo_frogmant (block_M, block_8], accum_dtype)
acs_s_oast frallloonraggnttcc
block_Nl, dtypet
acc_o = E.elloc_frogmans (bboock_M, dim), accun_dtype)
7.alloo
foaglooeeeeeeet fragment (bbookkMM] sccum_dtype)
scores_max_prast T.alloc_ fraggment [blook_M), stype)
scores_scaue scale Toa (blosk_M), accum_dtype) scorcs_sum = F. logsum haaloo_fragment ([block_],, scsum_dtype)
([block_M,
_ppee_dtype)
T.copy (Q,b,, bx block (bx 1) block_N, (_share)) r.fill (acc_e, 0)
19
T. (logsum, T.fill (1corlesccrs_axx -I.infinity (nceinityycccummttype.))
Loop_range =
T.min cet. bloek_N), T.ceildiv(
(bx 1) bleck_X, block(N)) if is_causal else (.ceildiv block_NI)
for in Peppeeined
nur_st.pps-nuunstapee)
if BlockMask [bz, Ex, by, kl: 7. P.copy [bz, k block (k block by, ih Ksshared) if is_ceusal: for i, J in (elockk block_N):
acc_sli, 51 thif block
Lanenity block_N 3,
else:
T.clear (ace_s)
T. ge (Q_shared, K_shared, ace_s, tranapose_Bi BBTrue,
PuLLRoww)
D.ccpy (scores_max, scores_max_pzee
7. (scoressccee_aa,
(acuu_tyeacuu
w.reduce_maxlaacee. scores_nax, dim-l, clea:-False)
for in (blockk for scores_scale([] T.Parallel Iblook_M, block_N);
scale scores_mex[1] scelet
a0c_a(1, ]] T.exp2 (ace_s(t, 1) scale scorea_rax(1] scale) T. Ereduce_sumtacee scorea_sus dim-1!
T.Porallel (black_M:
logsum[i] logsum[i] scores_scale(i] scores_sum[i]
Tooppy @apo a6c_s_cast)
for in
dim):
a6c_e(1, ]] ^= [cores_scale["]]
E.copy (Vibs, T.gems (acs_a_cast, V_shared, sea_e, CammWarpPa Fo..Fallllylllllleeoooo
block_N:
block_N, th, V_shared)
for 1, in T.Parallel (block_M dimt: ags_a(1, 3] /= logaum(1) T.copy (acc_s, o_shared)
T.copy (0_shared, Output [bx, black_Me (bx 1) block_M, by, :])
Figure 19: Implementation of Block Sparse Flash Attention with TILELANG.
36


Reviews:
[{"review_id": "jkoMiNfVsD", "reviewer": "Reviewer_Bjv9", "summary": "This paper introduces TILELANG, a novel programming system that bridges the programmability-performance gap in AI accelerator programming. It treats \"tiles\" (hyper-rectangular tensor slices) as first-class entities in a Pythonic DSL, enabling explicit control over memory, data movement, and scheduling. Core innovations include:\n\nTile Recommendation: Uses a roofline model to suggest hardware-aware defaults (tile sizes, memory placement, warp partitioning).\nTile Inference: Autocompletes low-level details (memory layouts, schedules) via constraint propagation from partial specifications or recommendations.\nOn NVIDIA H100 and AMD MI300X GPUs, TILELANG matches or exceeds hand-tuned libraries (FlashAttention, CUTLASS) and outperforms Triton, with code size closer to manual implementations but simpler.", "questions": "1. What performance trade-offs were made to create a unified API for both NVIDIA and AMD? Does this abstraction hide hardware-specific features that might be key for hitting peak performance on a single platform?\n\n2. How would TILELANG's tile-based abstractions extend to highly irregular computations, such as sparse operations or Graph Neural Networks, where data access patterns are not uniform?\n\n3. The appendix reports impressive tuning times of around 10-15 seconds. However, this is not directly compared to the tuning time of other auto-tuning frameworks (e.g., TVM's Ansor/AutoTVM, Triton's autotuner) for the same tasks. A direct comparison would better highlight the efficiency advantage of the recommendation/inference approach over empirical search.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "Solves a Critical Problem: Precisely targets the core trade-off between programmability and peak performance in AI kernel development, a highly relevant issue.\n\nInnovative Human-in-the-Loop Framework: The novel \"Tile Recommendation + Tile Inference\" model is highly effective, lowering the barrier for entry while retaining expert control and drastically reducing manual tuning effort.\n\nElegant Programming Model: The tile-centric DSL is clean and expressive, allowing developers to control low-level hardware behavior (memory, parallelism) in a high-level language with concise code.\n\nHighly Convincing Evaluation: Achieves state-of-the-art (SOTA) or near-SOTA performance on both NVIDIA and AMD GPUs with significantly less code than hand-tuned libraries, demonstrating its effectiveness and cross-platform capability.", "weaknesses": "1. The roofline-based analytical model ignores complex microarchitectural details, which could lead to suboptimal suggestions in certain scenarios.\n\n2. There are several works  about kernel generators that are not cited properly.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tckQI9N6gU", "reviewer": "Reviewer_CYSa", "summary": "This work proposes TileLang, a framework to help simplify GPU kernel programming. The framework automates decisions around memory movement and scheduling, while exposing tile primitives and Pythonic operators. The work evaluates on NVIDIA and AMD GPUs and across a breadth of workloads.", "questions": "See questions above.", "limitations": "", "rating": 4, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "The results are strong and the kernel examples / level of simplicity is compelling. This seems like a useful system for the AI community. The baselines for evaluation are well-chosen and the chosen level of abstraction to expose to the developer (between programming primitives, compiler automations) seems effective.", "weaknesses": "It is not very precise how the automation works within TileLang / what the underlying algorithms are. The automation is very nice, but it's not clear if it breaks down in any settings since that's not fully analyzed/explained to the reader. The paper makes clear *what* is done, but not *how* it's accomplished. \n- It is not clear how TileLang handles the memory access patterns for NVIDIA vs. AMD.\n- It’s is not clear where the line is that divides the developer’s tasks from what TileLang automates. \n- It’s clear that the code is compact from the plots, but it’s not explained what deltas from existing frameworks contribute to the code size delta. \n- It is not clear what contributes to the speed ups over the baselines\n\nThe paper does not clearly compare to prior work, which makes it difficult to assess novelty. For instance, ThunderKittens (over a year old) is not mentioned in the related works or intro despite also using automatically optimized tile primitives and Pythonic functions. Mojo is another compiled language that works on both AMD and NVIDIA with tiles. In the intro Contribution #1 is tile abstractions, which has been shown in extensive prior work and it’s not clear how TileLang differs. Contribution #2 is tile configuration. Again prior frameworks provide default swizzle patterns and register layouts for matrix instruction alignment; how does the proposed method differ? \n\nThe writing could be clearer in a few places: \n- What is the backend of TileLang? The intro doesn’t explain what language / frameworks it’s built on. For instance, the first time it’s mentioned that TileLang is a compiler/uses compiler passes is on L122. \n- L122: Why are compiler analyses needed versus making the load/store functions use coalesced load patterns and loop tiling to begin with? What additional value is the compiler providing? \n\nResults:\n- Does the MI300X kernel match the performance of FlashAttention-3 on the H100? It would be useful to know the absolute TFLOP numbers everywhere, since PyTorch is a constantly changing baseline. For instance, some analysis has shown that PyTorch SDPA achieves 20-100 TFLOPs in some AMD Docker containers, so it is difficult to understand the results.\n- Is there an analysis that shows that the “Pipelined” abstraction is sufficient for all warp specialization a user might want to do? I notice that synchronization is abstracted away from the user. Is there ever a case where the user would want to manage synchronization points more carefully and how does TileLang support this if so? \n- Are the appendix kernels the same on NVIDIA and AMD? \n\nOverall this is a very solid and useful system, but the research aspects -- how the philosophy and contributions are novel, and how the algorithms work -- have not been sufficiently explained.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zholP9w7RO", "reviewer": "Reviewer_CKBd", "summary": "The authors propose a new domain-specific language (DSL) for developing efficient kernels on modern GPUs, called TileLang. Similar to the widely used tile-level kernel programming language Triton, TileLang also treats tiles as first-class citizens. However, unlike Triton, TileLang exposes more control over underlying hardware components—such as memory scope, tile layout, and thread mapping—to the programmer. This additional control allows developers to implement more efficient kernels and achieve better performance in cases where the Triton compiler cannot automatically apply these optimizations. Extensive experiments demonstrate the advantages of TileLang over existing languages.", "questions": "Thank you to the authors for submitting to ICLR 2026!\n\nI have been following TileLang for quite some time. It is a solid piece of work that makes meaningful academic contributions and provides a robust implementation suitable for real-world model deployment. In my view, TileLang may be an even better fit for top systems conferences like ASPLOS or OSDI, where the strong implementation and in-depth design details could be more fully appreciated.\n\nI have the following suggestions for improving the paper:\n\n**1. Be more conservative with claims**\n\nIn the introduction, the authors use the MLA workload on H100 to highlight TileLang’s advantages over Triton, showing TileLang as over 5x faster. Other workloads reportedly achieve up to 10x speedups. While Triton can indeed be suboptimal in many cases, performance differences of several times are uncommon. The current phrasing may give readers the impression that TileLang consistently outperforms Triton by a large margin. I suggest reporting average speedups and explaining the reasons behind the 5x improvement on Hopper’s MLA workload—such as potential issues in the Triton program implementation or current compiler limitations.\n\n**2. Add experiments on more hardware**\n\nThe evaluation primarily focuses on H100 and MI300. To demonstrate the generality of TileLang, it would be valuable to include benchmarks on other widely used architectures like Ampere and Blackwell.\nFor example, I previously - two or three months ago - benchmarked MHA implementations of TileLang and Tilus [1] on an RTX 4090 (q=k=v=4096, qkv heads=32, dim=128, using the example code from both GitHub projects). TileLang achieved 128 TFLOPs, while Tilus reached 148 TFLOPs. My motivation was to test whether TileLang’s pipelined loop effectively handles the prefetching optimization used in FlashAttention, which differs from that in matrix multiplication.\n\n**3. Compare with more recent related works**\n\nTriton’s programming model is elegant and easy to use, but its abstraction level is higher than that of the underlying hardware. As hardware evolves, achieving peak performance increasingly requires complex optimizations (e.g., software pipelining, warp specialization), which are difficult for the Triton compiler to apply automatically. Consequently, several tile-based DSLs have been proposed that sit between CUDA C and Triton in abstraction level.\nTo my knowledge, these include Triton Gluon [2], Tilus [1], and TileLang (the subject of this paper). A qualitative comparison among these DSLs would strengthen the paper. If possible, adding quantitative results -- even preliminary ones -- would make the comparison more compelling.\n\n**Minor Suggestion**\n\nConsider adding a short section acknowledging TileLang’s implementation foundation in TVM, to give appropriate credit.\n\n**Summary**\n\nOverall, TileLang is a novel and well-executed work. I strongly support acceptance and would champion this paper. Happy to increase the score if above concerns are addressed.\n\n**References**\n\n- [1] Tilus: A Tile-Level GPGPU Programming Language for Low-Precision Computation (https://arxiv.org/abs/2504.12984, https://github.com/NVIDIA/tilus)\n- [2] Triton Gluon: https://github.com/triton-lang/triton/tree/main/python/tutorials/gluon", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. The problem addressed by this work is highly important. We need effective DSLs to simplify writing peak-performance kernels for modern GPUs, which are increasingly complex to program.\n2. Both the implementation and experimental results are solid.\n3. The paper is well-written and easy to follow.", "weaknesses": "1. Some claims could be more conservative.\n2. Experiments are limited to a few hardware platforms.\n3. It would strengthen the paper to compare with more recent related works.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vYAbSmj1WN", "reviewer": "Reviewer_7JBZ", "summary": "TileLang is a tile-level domain-specific language (DSL) for developing high-performance AI kernels. It is built on top of TVM, and provides {unified fused tile-level dataflow graph (FTG), tile recommendation, tile inference}, and explicit control over {memory placement, data movement, parallel scheduling} to bridge the gap between programmability and performance.\n\nOverall, I like the paper.", "questions": "**Q1.** I was wondering about the benchmarking details - do your performance numbers include the Python wrapper overhead, or are they just measuring the kernel execution time? Also, what warmup strategy did you use and how many repetitions for the statistics?\n\n**Q2.** It would be helpful to see some measurements of compilation and code generation time in your benchmarks. This would help understand the overhead for first-time graph execution, especially for JIT compilation scenarios.\n\n**Q3.** Similar to W3 asking for direct code comparisons, could you provide a direct IR-level comparison between TileLang and TVM for a simple matmul example? It would be very helpful to see the actual IR representations side-by-side to understand how different TileLang's tile-level abstractions are from TVM's TensorIR at the IR level.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 4, "strengths": "**S1.** Excellent contribution from both research and engineering perspectives. TileLang addresses a real pain point in ML kernel development by bridging the programmability-performance gap, which will greatly help in migrating ML kernels to new hardware platforms, as well as developing new ML kernels on existing hardware.\n\n**S2.** Strong empirical results across diverse kernels and hardware platforms. The benchmarks demonstrate substantial speedups (over Triton, PyTorch) and sufficiently close to the hand-crafted kernels (e.g. FlashMLA).\n\n**S3.** The unified FTG abstraction combined with tile recommendation and tile inference provides a principled approach to kernel optimization. The two-stage workflow effectively balances automation with fine-grained developer control.", "weaknesses": "(Here, the weaknesses are more about how the paper writing can be improved to provide more informative and intuitive insights to the readers, but less about the technical contributions.)\n\n**W1.** Missing comparison with TVM. Since TileLang is built directly on top of TVM, the paper lacks benchmarks comparing against TVM itself. While comparisons with Triton/Torch/hand-crafted kernels are valuable, without TVM we're missing a key component to understand TileLang's contribution over its foundation.\n\n**W2.** Lack of software architecture overview. The relationship between TileLang and TVM is unclear. I would like to see an illustration showing the compilation pipeline: Python TileLang DSL → FTG → TVM IR (TensorIR/Relay IR?) → CUDA. What is the relationship between FTG and TVM's IRs? Is FTG a standalone IR or must it combine with other IRs to fully represent the computation? This is not stated clearly in the paper and would help distinguish TileLang's real contribution from TVM's existing capabilities.\n\n**W3.** Missing direct code comparisons with TVM/Triton/etc. Instead of just describing syntactic differences and abstraction levels, showing actual side-by-side Python code comparing TileLang vs TVM vs Triton would be much more intuitive. (\"Shut up and show me the code\").\n\n**W4.** No comparison with TaichiLang, which sits at a higher abstraction level and supports autodiff. A performance/LoC/syntax comparison would help position TileLang in the landscape of kernel programming frameworks. Or, at least, a mention of the differences.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)