DEV Community

Shrijith Venkatramana
Shrijith Venkatramana

Posted on AI-assisted

TileLang: Writing LLM GPU Kernels by Thinking in Tiles

A modern LLM can spend most of its time doing something that looks almost embarrassingly simple:

C = A @ B
Enter fullscreen mode Exit fullscreen mode

The mathematics is simple.

Making an NVIDIA H100 execute that multiplication efficiently is a different problem.

You have to decide which pieces of A and B live in HBM, which are copied into shared memory, which values stay in registers, how threads cooperate, which Tensor Core instruction performs the multiply-accumulate, how memory transfers overlap with computation, and how the whole thing behaves when the matrix dimensions do not fit your preferred tile sizes.

That is why two implementations of the same equation can have very different runtimes.

This is the problem TileLang is designed to make easier.

TileLang is a Python-based domain-specific language for writing high-performance kernels around a tiled programming model. It sits on top of the TVM compiler infrastructure and gives developers explicit control over things such as tile sizes, memory placement, layouts, parallelism, tensorization, and software pipelining. At the same time, it tries to keep the programming model much simpler than writing every detail in CUDA C++.

The important idea is simple:

Instead of thinking about millions of individual GPU threads, think about a small number of tiles moving through a hierarchy of memories and compute units.

Once you see GPUs this way, TileLang starts to make sense.

1. Why LLM performance eventually becomes a kernel problem

Modern frameworks already do a great deal for you.

Suppose you write:

y = torch.matmul(x, weight)
Enter fullscreen mode Exit fullscreen mode

PyTorch does not execute that line directly. It eventually dispatches to highly optimized kernels, often from libraries such as cuBLAS or other specialized implementations.

For conventional operators, that is exactly what you want.

The trouble starts when the operation you need is unusual.

LLM inference contains many cases like this:

dequantize weights
       ↓
matrix multiplication
       ↓
scaling
       ↓
activation
       ↓
another transformation
Enter fullscreen mode Exit fullscreen mode

Or attention:

Q
 \
  QK^T
     ↓
   softmax
     ↓
    × V
Enter fullscreen mode Exit fullscreen mode

Or a newer architecture such as DeepSeek's Multi-Head Latent Attention, where the kernel has a particular combination of projections, KV-cache access patterns, reductions, and data movement.

The model gives you a mathematical computation graph.

The GPU cares about something closer to:

HBM
 ↓
L2
 ↓
shared memory
 ↓
registers
 ↓
Tensor Cores / vector ALUs
Enter fullscreen mode Exit fullscreen mode

The distance between those two descriptions is where kernel engineering lives.

This has been a recurring problem in AI systems for years.

In 2017 NVIDIA's Volta architecture introduced Tensor Cores, making specialized matrix operations a first-class hardware feature. In 2018 Tianqi Chen and collaborators introduced TVM as a compiler stack for mapping tensor computations onto diverse hardware. In 2019 Philippe Tillet, H. T. Kung, and David Cox published Triton, explicitly using tiles as the central abstraction for neural-network computation. In 2022 Tri Dao and collaborators showed with FlashAttention that even the algorithmic organization of memory movement could radically change attention performance.

TileLang belongs to this lineage.

Its authors include researchers from Peking University and Microsoft Research. The project was open-sourced in January 2025, and the work subsequently appeared as an ICLR 2026 oral paper under the title TileLang: Bridge Programmability and Performance in Modern Neural Kernels.

So TileLang is easier to understand as part of a long progression:

CUDA
  ↓
GPU libraries such as cuBLAS
  ↓
TVM and tensor compilers
  ↓
Triton
  ↓
TileLang
Enter fullscreen mode Exit fullscreen mode

Each step tries to let programmers express more useful structure without forcing them to manually specify every machine instruction.

2. The basic idea: a GPU should reuse data

Start with matrix multiplication.

Suppose:

A is M × K
B is K × N
C is M × N
Enter fullscreen mode Exit fullscreen mode

The mathematical definition is:

C[i,j] = sum_k A[i,k] * B[k,j]
Enter fullscreen mode Exit fullscreen mode

For a developer, this looks like a triple loop:

for i in range(M):
    for j in range(N):
        for k in range(K):
            C[i, j] += A[i, k] * B[k, j]
Enter fullscreen mode Exit fullscreen mode

But imagine implementing that literally on a GPU.

Every output element repeatedly needs values from A and B.

That creates enormous memory traffic.

Now divide the matrices into tiles.

For example:

A tile: 128 × 32
B tile:  32 × 128
C tile: 128 × 128
Enter fullscreen mode Exit fullscreen mode

A thread block can load the two input tiles into shared memory:

global memory
     ↓
+-------------+    +-------------+
| A 128 × 32  |    | B 32 × 128  |
+-------------+    +-------------+
       ↓                  ↓
    shared memory
       ↓                  ↓
        Tensor Cores
              ↓
       C 128 × 128
Enter fullscreen mode Exit fullscreen mode

Now the same values of A and B are reused many times.

This is the fundamental reason tiling works.

A back-of-the-envelope calculation

Consider:

A = 4096 × 4096
B = 4096 × 4096
Enter fullscreen mode Exit fullscreen mode

using FP16.

The mathematical work is approximately:

2 × 4096 × 4096 × 4096
≈ 137 billion FLOPs
Enter fullscreen mode Exit fullscreen mode

The three matrices occupy roughly:

A: 4096² × 2 bytes ≈ 32 MB
B: 4096² × 2 bytes ≈ 32 MB
C: 4096² × 2 bytes ≈ 32 MB

Total ≈ 96 MB
Enter fullscreen mode Exit fullscreen mode

That is the idealized amount of data you would need to read/write if the input matrices could be perfectly reused from on-chip storage.

Now consider a naive one-output-at-a-time implementation.

Each of the roughly 16.8 million output elements needs 4096 values from A and 4096 from B.

Very roughly:

16.8M × 8192 × 2 bytes
≈ 275 GB
Enter fullscreen mode Exit fullscreen mode

of input traffic.

The arithmetic is unchanged.

The memory traffic is radically different.

That is why GPU optimization is frequently about moving each value fewer times, rather than doing fewer mathematical operations.

On an H100 SXM, NVIDIA specifies roughly 3.35 TB/s of HBM bandwidth and about 1.98 PFLOPS of FP16 Tensor Core throughput. Those numbers illustrate the scale of the imbalance: the machine has enormous compute capability, but getting data to the compute units efficiently is part of the problem.

This is the intuition behind TileLang.

A tile is a unit of work and a unit of data movement.

3. What TileLang code actually looks like

TileLang deliberately resembles Python.

A simplified matrix-multiplication kernel looks roughly like this:

import tilelang
import tilelang.language as T

@tilelang.jit
def matmul(A, B, M, N, K,
           block_M=128,
           block_N=128,
           block_K=32):

    dtype = T.float16
    accum_dtype = T.float32

    @T.prim_func
    def kernel(
        A: T.Tensor((M, K), dtype),
        B: T.Tensor((K, N), dtype),
        C: T.Tensor((M, N), dtype),
    ):
        with T.Kernel(
            T.ceildiv(N, block_N),
            T.ceildiv(M, block_M),
            threads=128
        ) as (bx, by):

            A_shared = T.alloc_shared(
                (block_M, block_K), dtype
            )

            B_shared = T.alloc_shared(
                (block_K, block_N), dtype
            )

            C_local = T.alloc_fragment(
                (block_M, block_N), accum_dtype
            )

            T.clear(C_local)

            for k in T.Pipelined(
                T.ceildiv(K, block_K),
                num_stages=3
            ):
                T.copy(
                    A[by * block_M, k * block_K],
                    A_shared
                )

                T.copy(
                    B[k * block_K, bx * block_N],
                    B_shared
                )

                T.gemm(
                    A_shared,
                    B_shared,
                    C_local
                )

            T.copy(
                C_local,
                C[by * block_M, bx * block_N]
            )

    return kernel
Enter fullscreen mode Exit fullscreen mode

The important lines are these:

A_shared = T.alloc_shared(...)
B_shared = T.alloc_shared(...)
C_local = T.alloc_fragment(...)
Enter fullscreen mode Exit fullscreen mode

They explicitly describe the memory hierarchy.

And:

T.copy(...)
T.gemm(...)
Enter fullscreen mode Exit fullscreen mode

describe movement and computation at the tile level.

Finally:

T.Pipelined(...)
Enter fullscreen mode Exit fullscreen mode

describes how multiple iterations should overlap.

The current TileLang programming model treats these operations as first-class tile operations, while scheduling features such as parallelization, layout annotations, swizzling, and pipelining can be added separately.

That separation is one of the central ideas.

4. TileLang separates what you compute from how you schedule it

Consider two questions.

Question one:

What mathematical computation should happen?
Enter fullscreen mode Exit fullscreen mode

For matrix multiplication:

C = A @ B
Enter fullscreen mode Exit fullscreen mode

Question two:

How should the GPU execute it?
Enter fullscreen mode Exit fullscreen mode

Possible answers include:

128 × 128 output tiles
32-wide K tiles
128 threads
shared-memory staging
Tensor Core GEMM
3-stage pipeline
swizzled memory layout
Enter fullscreen mode Exit fullscreen mode

These are different concerns.

TileLang makes this separation explicit.

The paper describes four important scheduling dimensions:

thread binding
memory layout
tensorization
pipeline
Enter fullscreen mode Exit fullscreen mode

For example:

T.gemm(...)
Enter fullscreen mode Exit fullscreen mode

expresses the tile computation.

Meanwhile:

T.Pipelined(...)
Enter fullscreen mode Exit fullscreen mode

controls overlapping data movement and compute.

And:

T.annotate_layout(...)
Enter fullscreen mode Exit fullscreen mode

can influence how data is physically distributed.

This matters because the same mathematical operation may need a different schedule for different GPUs.

A schedule that works well on an A100 may behave differently on an H100.

A layout suitable for NVIDIA hardware may need a different implementation on AMD.

TileLang therefore keeps much of the programming model shared while allowing the backend to generate CUDA, HIP, LLVM, Metal, and other target-specific code paths. The current project documentation also lists targets such as CUDA, HIP, Metal, WebGPU and CPU execution.

This is an important distinction from writing CUDA directly.

CUDA gives you extremely detailed control.

TileLang tries to give you the level of control where the performance decisions are meaningful, while allowing the compiler to derive many of the mechanical details.

5. Why this matters specifically for LLMs

The strongest use cases for TileLang are not ordinary textbook matrix multiplications.

They are kernels where several operations need to be carefully fused and scheduled together.

FlashAttention

Normal attention is:

S = QK^T / sqrt(d)

P = softmax(S)

O = PV
Enter fullscreen mode Exit fullscreen mode

A straightforward implementation materializes S.

For sequence length n, that matrix is:

n × n
Enter fullscreen mode Exit fullscreen mode

So memory grows quadratically.

FlashAttention changed the implementation strategy.

Instead of producing the entire attention matrix in HBM, it processes blocks of Q and K, keeps intermediate values on chip, and performs the softmax incrementally.

The algorithm became famous because it showed that reducing memory traffic can matter as much as reducing arithmetic. The original FlashAttention paper described this as an IO-aware formulation and demonstrated substantial end-to-end speedups for Transformer workloads.

TileLang gives you the programming primitives needed to express this sort of computation.

Conceptually:

Q tile ─────────────┐
                    │
K tile ─────────────┼──> QKᵀ
                    │
                    └──> online softmax
                              │
V tile ───────────────────────┘
                              ↓
                         output tile
Enter fullscreen mode Exit fullscreen mode

The current TileLang paper reports that on H100 it can express pipeline schedules comparable in complexity to those used by FlashAttention-3. In the authors' evaluation, their FlashAttention implementation outperformed the compared FlashAttention-3, Triton, and PyTorch baselines for the tested workloads, with performance remaining close to FlashAttention-3 at longer sequence lengths.

Quantized LLM inference

Now consider a common LLM inference pattern:

weights stored as INT4 / FP4
             ↓
        dequantization
             ↓
      matrix multiplication
Enter fullscreen mode Exit fullscreen mode

You could perform these as separate kernels:

kernel 1: dequantize
kernel 2: GEMM
Enter fullscreen mode Exit fullscreen mode

That creates extra memory traffic.

A better implementation may keep the dequantized values in registers or another on-chip representation and feed them directly into the matrix operation.

TileLang's own paper includes an FP4/FP16 weight-only GEMM implementation. The example explicitly allocates packed weights, performs conversion into a local tile, and then feeds that tile into GEMM.

That is the kind of operation where the abstraction becomes useful:

compressed data
      ↓
 tile
      ↓
dequantize
      ↓
 tile
      ↓
 Tensor Core GEMM
Enter fullscreen mode Exit fullscreen mode

There is no need to interpret the operation as millions of independent scalar instructions.

Think in tiles.

DeepSeek-style attention

Another example is Multi-Head Latent Attention.

TileLang's paper includes a FlashMLA implementation with tiled Q, K, positional components, score accumulation, reductions and pipelining. The implementation is expressed in Python and uses TileLang primitives such as T.gemm, T.reduce_max, T.reduce_sum, T.Pipelined, shared-memory allocations and fragment allocations.

This is a useful example because it is much closer to the reality of modern LLM systems than a simple GEMM.

The kernel itself is an algorithmic object.

It is simultaneously:

computation
+
memory-management strategy
+
parallel schedule
+
hardware mapping
Enter fullscreen mode Exit fullscreen mode

TileLang is designed around that combined reality.

6. The part beginners usually miss: the compiler is doing a lot of work

It is easy to look at:

T.gemm(A_shared, B_shared, C_local)
Enter fullscreen mode Exit fullscreen mode

and assume that TileLang is merely a prettier syntax for CUDA.

That undersells what the compiler does.

TileLang programs are progressively lowered through a compiler pipeline involving TileLang's frontend and AST, TVM's intermediate representation, optimization passes and backend code generation. The compiler performs things such as layout inference, thread mapping, pipeline derivation and other transformations.

One particularly important idea is layout inference.

Suppose you write:

C_local = 128 × 128
Enter fullscreen mode Exit fullscreen mode

That does not mean every thread owns one simple 128 × 128 array.

The compiler needs to determine how that logical tile is distributed across the hardware.

Conceptually:

logical tile

128 × 128
    ↓
   split
    ↓
warps / threads
    ↓
register fragments
    ↓
Tensor Core instruction layout
Enter fullscreen mode Exit fullscreen mode

This is one of the reasons tile-level programming sits at an interesting point in the abstraction hierarchy.

At the Python level, you can reason about matrices.

At the hardware level, the machine is executing thread-level instructions.

TileLang lives in the middle.

The ICLR 2026 version of the work goes further with tile inference and tile recommendation. Tile inference uses the structure of a fused tile program to infer missing configuration information, while tile recommendation uses hardware information and heuristics to suggest configurations. The authors report fused attention implementations in under 80 lines of Python, with code-size reductions of up to 90% compared with manual implementations in their experiments.

This points toward an interesting division of labor:

Developer:
    define algorithmic dataflow

Compiler:
    derive much of the boring scheduling machinery

Developer:
    intervene when performance requires it
Enter fullscreen mode Exit fullscreen mode

That last line matters.

TileLang is not trying to remove the performance engineer.

It is trying to make the performance engineer work at a more useful level.

7. The economics of a faster kernel

Why spend engineering time on something that saves 10 microseconds?

Because LLM kernels often run an enormous number of times.

Imagine a serving system processing:

1,000,000 kernel invocations / second
Enter fullscreen mode Exit fullscreen mode

Suppose an optimization reduces one hot kernel from:

20 us → 15 us
Enter fullscreen mode Exit fullscreen mode

The saving is:

5 us × 1,000,000
= 5 seconds of GPU compute time per second
Enter fullscreen mode Exit fullscreen mode

That is approximately equivalent to freeing:

5 GPU-seconds / second
≈ 5 continuously occupied GPUs
Enter fullscreen mode Exit fullscreen mode

The exact operational value depends on whether you are throughput-bound, latency-bound, batch-constrained, or simply paying for reserved capacity.

But the basic equation is powerful:

value of optimization
≈ invocations × time saved × cost of compute
Enter fullscreen mode Exit fullscreen mode

This is why highly optimized kernels can justify weeks of engineering work.

And there is another effect.

Suppose a particular kernel consumes 40% of the runtime of an inference service.

A 2x improvement to that kernel does not make the entire model 2x faster.

If everything else stays unchanged:

old runtime:

40 units kernel
60 units everything else
-------------------------
100 total

new runtime:

20 units kernel
60 units everything else
-------------------------
80 total
Enter fullscreen mode Exit fullscreen mode

The whole application improves by:

100 / 80 = 1.25x
Enter fullscreen mode Exit fullscreen mode

This is simply Amdahl's law.

So the right question is not:

"Can TileLang make this kernel faster?"

It is:

"How much of my actual serving or training workload is this kernel responsible for?"

For production systems, that distinction matters more than a benchmark headline.

8. A useful mental model for developers

You can think of the ecosystem like this.

PyTorch

You express the model.

y = torch.matmul(x, w)
Enter fullscreen mode Exit fullscreen mode

The framework and libraries choose the implementation.

CUDA

You control almost everything.

threads
warps
shared memory
registers
synchronization
instructions
layouts
Enter fullscreen mode Exit fullscreen mode

You get enormous control and a corresponding engineering burden.

Triton

You work with blocks and tiles while the compiler handles much of the lower-level mapping.

This was already a major shift from CUDA. Philippe Tillet's 2019 paper explicitly argued that tiles could provide a more productive way to build custom neural-network kernels while approaching vendor-library performance.

TileLang

You make the tile and its dataflow even more explicit:

tile
 ↓
memory placement
 ↓
tile operation
 ↓
layout
 ↓
pipeline
 ↓
hardware instruction
Enter fullscreen mode Exit fullscreen mode

The interesting part is that TileLang does not force all of these decisions into one monolithic kernel description.

Its design explicitly separates dataflow-oriented tile operations from scheduling primitives. That makes it possible to begin with a relatively simple kernel and progressively introduce more hardware-aware control as performance requires it.

9. Where TileLang fits in your own work

You probably should not begin an LLM optimization project by writing everything in TileLang.

A practical progression is:

PyTorch operator
      ↓
profile
      ↓
identify a real bottleneck
      ↓
existing optimized kernel?
      ↓
yes → use it
no
      ↓
custom Triton / TileLang kernel
      ↓
profile again
      ↓
inspect generated code and memory behavior
      ↓
tune tile sizes, layout, pipeline
Enter fullscreen mode Exit fullscreen mode

TileLang becomes particularly interesting when the bottleneck has one or more of these properties:

• unusual fusion
• quantized computation
• custom attention
• unusual reductions
• expensive memory movement
• need for explicit shared-memory/register placement
• architecture-specific pipeline requirements
• desire to target multiple accelerator backends
Enter fullscreen mode Exit fullscreen mode

The deeper lesson is larger than TileLang itself.

LLM performance engineering is gradually becoming a discipline of data movement design.

The equation may say:

C = A @ B
Enter fullscreen mode Exit fullscreen mode

but the real implementation problem is:

Which A?
Which B?
Where do they live?
When are they loaded?
Who loads them?
Who reuses them?
Which threads own them?
Which instruction computes them?
Can the next tile load while this one computes?
Enter fullscreen mode Exit fullscreen mode

That is what the tile abstraction makes visible.

And that is why a language that looks like Python can still be talking directly about shared memory, register fragments, Tensor Cores, layouts and asynchronous pipelines.

The fundamental unit is no longer the scalar.

It is the tile.

Conclusion

The progression from CUDA to TVM, Triton and TileLang reflects a recurring tension in systems programming.

High-level abstractions give productivity.

Low-level control gives performance.

TileLang's approach is to move the abstraction boundary upward without pretending that hardware details have disappeared.

For LLM developers, that is useful because today's important kernels are increasingly fused, quantized and architecture-specific. The question is often less "how do I express this operation?" and more "how do I move this data through the machine while doing the operation?"

Once you start seeing an attention kernel as a stream of tiles moving between HBM, shared memory, registers and Tensor Cores, the code becomes much easier to reason about.

And that is the real idea behind TileLang.

When you next profile an LLM and discover that a tiny custom kernel is consuming a surprising amount of runtime, would you reach for CUDA, Triton, or a tile-oriented language like TileLang—and what would determine that choice?

For implementation details and runnable examples, the current TileLang documentation and kernel examples are the natural next step.

Top comments (0)