Part I: The mental model
CUDA (SIMT): you write per-thread code threadIdx, blockIdx, blockDim you index individual threads __shared__ float tile[...]; you declare scratchpad memory __syncthreads(); you place the barriers you also own: coalescing, vectorization, bank conflicts, register use Triton (per-block SPMD): you write per-tile code pid = tl.program_id(0) which block of work am I offs = pid*BLOCK + tl.arange(0,BLOCK) a whole tile of indices at once x = tl.load(ptr + offs, mask=offs<n) load the tile, masked at the edge tl.store(out + offs, x + y, mask=...) store the tile the compiler owns: thread mapping, coalescing, shared memory, sync
The one-sentence identity. Triton is a block-level
programming model for GPUs, where you write a single program that
operates on tiles and the compiler lowers those tile operations
into correct, coalesced, synchronized thread-level code, so you
keep control of the parallel decomposition and hand away the
per-thread bookkeeping. CUDA exposes the GPU as thousands
of scalar threads (the SIMT model) and asks you to manage the
memory hierarchy yourself. Triton raises the unit of programming
from the thread to the tile, a small multi-dimensional block of
values that the compiler treats as a first-class object. You
still decide how the problem is cut into blocks and how big each
tile is, because that is where the algorithmic insight lives, but
you never write threadIdx, never allocate shared
memory, and never place a __syncthreads.
Two consequences follow. First, the code you write is close to
the math. A softmax kernel loads a row, subtracts its max,
exponentiates, sums, and divides, in about ten lines that read
like NumPy over a block, and the compiler is responsible for
making that block-level program run well on real hardware.
Second, one kernel is portable across GPUs. Because you never
hard-coded a warp count or a shared-memory layout, the same
@triton.jit function compiles for different NVIDIA
generations and for AMD GPUs, with the tile-to-hardware mapping
decided per target by the compiler. The productivity jump is
large enough that Triton became the code-generation backend for
PyTorch 2's torch.compile, which means an enormous
fraction of the Triton kernels running in the world today were
written by a compiler, not a human.
@triton.jit add_kernel Python source captured as an AST
|
v
add_kernel[grid](...) JITFunction: specialize on arg dtypes + alignment,
| build a cache key, compile once per signature
v
code generator walk the AST, emit Triton IR (tt dialect, MLIR)
|
v
Triton IR (TTIR) hardware-agnostic tile program
| passes: coalesce, layout, pipeline, accel-matmul
v
Triton GPU IR (TTGIR) adds #blocked / #mma / #shared layouts, warps per CTA
|
v
LLVM IR (llir) tt / ttg ops lowered to the LLVM dialect, then LLVM IR
|
v
PTX -> cubin NVPTX backend emits PTX, ptxas assembles it (NVIDIA)
|
v
driver.launch(grid) load cubin, launch a grid of programs on the GPU
Everything in this chapter is verified against the project as of
July 2026. Triton moves fast and its compiler was rebuilt on
MLIR
a few years ago, so where an exact pass name or file path is
likely to have shifted I say so and stay at the level of the
concept. The lineage worth knowing, Triton began as Philippe
Tillet's research language (the
2019 MAPL paper
on tiled computations), matured under OpenAI, and now lives as the
community project triton-lang/triton with first-class
NVIDIA and AMD backends.
Part II: Using it
Triton is a Linux-and-GPU project. On a machine with a recent NVIDIA (Volta or newer) or AMD ROCm GPU the fastest way in is the wheel.
pip install triton
# triton also ships bundled inside recent PyTorch builds, so if you already
# have a CUDA-enabled torch, `import triton` may already work.
python -c "import triton, triton.language as tl; print(triton.__version__)"Building from source is only needed if you are working on the compiler itself. It is heavy, because the build compiles a pinned LLVM and MLIR, so most kernel authors should stay on the wheel.
git clone https://github.com/triton-lang/triton
cd triton
pip install -e . # compiles bundled LLVM/MLIR, takes a while the first timeOn macOS Triton installs and is pleasant to read and step through, but there is no local GPU to run on. The honest macOS workflow is reading the code and the tutorials locally and running kernels on a Linux GPU box.
Worked example: vector-add
The canonical first kernel adds two vectors. It shows the whole
programming model in one screen, program_id, a tile
of offsets, a boundary mask, tl.load and
tl.store, and a launch grid.
import torch
import triton
import triton.language as tl
@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
pid = tl.program_id(axis=0) # which block of BLOCK_SIZE elements
block_start = pid * BLOCK_SIZE
offsets = block_start + tl.arange(0, BLOCK_SIZE)
mask = offsets < n_elements # guard the ragged last block
x = tl.load(x_ptr + offsets, mask=mask) # load a whole tile from DRAM
y = tl.load(y_ptr + offsets, mask=mask)
tl.store(out_ptr + offsets, x + y, mask=mask)
def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
out = torch.empty_like(x)
n_elements = out.numel()
# grid is a function of the meta-parameters; ceil-div covers the tail
grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
add_kernel[grid](x, y, out, n_elements, BLOCK_SIZE=1024)
return out
x = torch.rand(98_432, device="cuda")
y = torch.rand(98_432, device="cuda")
torch.testing.assert_close(add(x, y), x + y)
Read it as follows. Torch tensors are passed straight in and
Triton treats each as a pointer to its first element plus the
metadata it needs. tl.program_id(axis=0) returns this
program's index along the first grid axis, the analog of CUDA's
blockIdx.x. tl.arange(0, BLOCK_SIZE)
materializes a compile-time-sized tile of indices, so
offsets is a whole block of positions rather than one
scalar. The mask makes the last, partial block safe,
masked-out lanes are simply not loaded or stored. Note that
BLOCK_SIZE is marked tl.constexpr, which
means it is a compile-time constant baked into the generated
kernel, and tl.arange requires its length to be a
power of two. The launch add_kernel[grid](...) uses
the square-bracket grid syntax, where grid is a
callable that receives the meta-parameters and returns the number
of program instances to spawn. Here that is
ceil(n_elements / BLOCK_SIZE), one program per tile.
Worked example: fused softmax
The second tutorial kernel is a fused softmax over the rows of a 2D tensor, and it is the example that first shows Triton earning its keep. A naive PyTorch softmax reads the input, writes the max, reads again for the exponentials, writes them, reads to sum, and writes the normalized result, several passes over DRAM. A fused Triton kernel loads each row once into fast on-chip memory, does all the arithmetic there, and writes the answer once.
@triton.jit
def softmax_kernel(out_ptr, in_ptr, in_row_stride, out_row_stride,
n_cols, BLOCK_SIZE: tl.constexpr):
row = tl.program_id(0) # one program per row
row_start = in_ptr + row * in_row_stride
col_offsets = tl.arange(0, BLOCK_SIZE) # BLOCK_SIZE >= n_cols
ptrs = row_start + col_offsets
mask = col_offsets < n_cols
# load the row; padded lanes get -inf so they never win the max or add mass
x = tl.load(ptrs, mask=mask, other=-float("inf"))
x = x - tl.max(x, axis=0) # numerically stable shift
num = tl.exp(x)
denom = tl.sum(num, axis=0) # reduction stays on-chip
y = num / denom
out_row = out_ptr + row * out_row_stride
tl.store(out_row + col_offsets, y, mask=mask)
def softmax(x: torch.Tensor) -> torch.Tensor:
n_rows, n_cols = x.shape
BLOCK_SIZE = triton.next_power_of_2(n_cols) # one tile spans the whole row
out = torch.empty_like(x)
softmax_kernel[(n_rows,)](
out, x, x.stride(0), out.stride(0), n_cols, BLOCK_SIZE=BLOCK_SIZE,
)
return out
The load-once, compute-on-chip, store-once shape is the whole
point, and it is the same memory-traffic argument that makes
fused attention fast, told at row scale here and at tile scale in
the FlashAttention chapter.
The tl.max and tl.sum calls are tile
reductions along axis=0, lowered by the compiler into
an efficient in-register and shared-memory reduction tree that you
never have to write. The stable shift by the row max is the same
trick derived on the softmax page, and
padding the tail with -inf keeps it correct when
BLOCK_SIZE is rounded up past n_cols.
The tutorial version goes one step further into a persistent
kernel, where the grid is capped at the number of resident
programs the GPU can hold and each program strides over multiple
rows in a loop, which amortizes launch overhead. It also picks
num_warps from the tile size. The idea to keep is the
fusion, not the exact persistent bookkeeping.
Common beginner mistakes
First, forgetting the mask. Without mask= on a load
or store, a partial final block reads and writes out of bounds,
which corrupts memory or faults. Every ragged dimension needs a
mask. Second, non-power-of-two tiles. tl.arange and
most tile shapes must be powers of two, so you size the block up
with triton.next_power_of_2 and rely on masks for the
remainder. Third, treating a tile like a Python list. A Triton
tile is a compiler value, not a runtime array, so you cannot index
it with a data-dependent scalar or take its Python
len. You express selection with tl.where
and masks. Fourth, expecting print debugging. The kernel body runs
on the GPU after compilation, so ordinary print does
nothing useful, and you reach instead for tl.device_print,
or you shrink to one block and compare against a torch reference,
or you dump the generated IR (Part VII).
Part III: When it is the right tool
Triton is the right tool when a standard library kernel does not exist for what you need and you want good performance quickly and portably. Custom fused elementwise-plus-reduction kernels, novel attention variants, quantized or mixed-dtype matmuls, and fused epilogues are its sweet spot, and it is also the shortest path for a researcher who wants a kernel that is fast enough without a month of CUDA. It is the code-generation target of PyTorch's TorchInductor, so many people ship Triton without writing it, and it underpins hand-written kernels in modern inference and training stacks. When you do want to read a serious hand-written Triton kernel in the wild, the attention kernels inside serving systems such as vLLM and SGLang are good specimens.
The honest cases for alternatives. NVIDIA's cuBLAS and cuDNN, and
the CUTLASS
template library, still win the last few percent of
peak on a fixed shape and a single GPU generation, which is why
production GEMM and convolution often call the vendor library
rather than a Triton kernel. Hand-written CUDA C++ gives total
control over every instruction and remains the answer when you
need warp-level primitives or hardware features that the DSL has
not yet surfaced. Numba's CUDA support puts you in Python but
keeps the per-thread SIMT model, so you still think in
threadIdx rather than tiles. JAX's
Pallas
is a kernel
DSL that actually lowers to Triton on GPU (and to Mosaic on TPU),
so it is a sibling and sometimes a front-end rather than a
competitor. Graph-level tensor compilers such as
TVM and
Halide
solve a different, whole-model problem. The tradeoff to state
plainly is that Triton trades a slice of peak performance
for a large gain in productivity and portability, and its job is
to get you most of the way to hand-tuned speed with a fraction of
the effort, on more than one vendor's hardware.
The architecture-shaped warning is about block size. The tile dimensions and the resulting register and shared-memory footprint decide how many program instances fit on each streaming multiprocessor at once, which is the occupancy that hides memory latency. Tiles too small waste memory bandwidth and launch overhead, tiles too large spill registers or exhaust shared memory and quietly serialize, and the best choice depends on the shape, the dtype, and the GPU. This is exactly the parameter Triton does not want you to guess, which is why autotuning exists.
Part IV: The full life of one kernel launch
The specimen is the vector-add above. Follow one call to
add_kernel[grid](x, y, out, n_elements, BLOCK_SIZE=1024)
from Python all the way to threads running on the device. The
same path runs for softmax, matmul, or any other
@triton.jit function.
Stage 1: the decorator and the JITFunction
@triton.jit does not compile anything. It wraps the
Python function in a JITFunction object (in the
runtime, historically triton/runtime/jit.py) and
captures the function's source and its abstract syntax tree.
Nothing about the GPU is decided yet, because the same source can
compile to many different kernels depending on the argument types
and the constexpr values it is eventually called with.
Stage 2: launch, specialization, and the cache key
The [grid] subscript plus the call is where work
begins. The runtime inspects the actual arguments and builds a
specialization signature. It records the element dtype of
each pointer argument, the types of the scalars, and the value of
every tl.constexpr argument, which is baked into the
kernel so a different BLOCK_SIZE is a different
kernel. It also records alignment facts, most importantly whether
each pointer is divisible by 16 bytes, which lets the compiler
emit wider vectorized loads, and whether certain integer arguments
are divisible by 16 or equal to 1. Those facts, plus the target
description (compute capability, warp size) and the Triton
version, form a cache key. If a compiled kernel already exists for
that key, in the in-process cache or on disk under
~/.triton/cache, the runtime skips straight to
launch. Otherwise it compiles once and stores the result. This is
why the first call to a kernel with a new signature is slow and
every later call is nearly free.
Stage 3: code generation to Triton IR
On a cache miss the code generator (in the compiler package)
walks the captured AST and emits Triton IR, the
tt dialect of MLIR, sometimes written TTIR. This is a
hardware-agnostic representation of the tile program. Loads and
stores become tt.load and tt.store over
tensor-of-pointer values, tl.arange becomes a range
op, the addition becomes an arith op over a tensor
value, reductions become tt.reduce, and a matmul
would become tt.dot. There is still no notion of
warps, threads, or shared memory. The IR describes tiles and the
operations over them, and its types carry tile shapes.
Stage 4: lowering to Triton GPU IR and choosing layouts
The pivotal stage converts Triton IR into Triton GPU
IR, the ttg dialect (TTGIR). This is where
the abstract tiles acquire layouts, attributes that say
exactly how the elements of a tile are distributed across the
warps and threads of a program instance and through shared memory.
The main ones are a blocked layout (#blocked, giving
each thread a small contiguous run of elements so that loads
coalesce), a shared-memory layout (#shared, often
swizzled to avoid bank conflicts), and, for tensor-core matmuls,
an MMA layout (#mma) for the accumulator plus
operand layouts for the inputs of tt.dot. A sequence
of MLIR passes runs over TTGIR, and while exact names drift, their
roles are stable, a coalescing pass picks blocked layouts that
make global memory access contiguous, a matmul-acceleration pass
rewrites tt.dot onto tensor-core MMA layouts, a
software-pipelining pass restructures loops so that loads for
future iterations are issued ahead of the compute that consumes
them (using asynchronous copies on hardware that supports them), a
layout-conversion cleanup removes redundant re-layouts, and a
shared-memory allocation pass plus a barrier-insertion pass place
the buffers and the synchronization you never wrote. For the
humble vector-add there is no dot and no loop, so this mostly
settles on a coalesced blocked layout and how many warps to use.
Stage 5: LLVM, PTX, and the binary
TTGIR is lowered to the LLVM dialect and then to ordinary LLVM IR
(the llir stage), at which point tile operations have
become per-thread instructions, shared-memory accesses, and
barriers. The NVIDIA backend runs LLVM's NVPTX target to produce
PTX, then invokes ptxas to assemble
PTX into a cubin, the actual machine code for the
target SM. The AMD backend follows the parallel road to AMDGCN
through the ROCm toolchain. The backends live under
third_party/nvidia and third_party/amd,
and the architecture is deliberately pluggable so a new accelerator
can register a backend that consumes TTGIR. The compiled result
keeps every intermediate around, reachable from Python through the
compiled kernel's asm dictionary under keys such as
ttir, ttgir, llir, and
ptx, which is exactly what you inspect in the labs.
Stage 6: the grid, the launcher, and execution
Back on the host, the grid callable is evaluated with the
meta-parameters to get the number of programs, here
ceil(98432 / 1024) = 97. A generated C launcher
stub unpacks the Python arguments into the kernel's ABI, and the
driver (the CUDA or HIP driver behind
triton.runtime.driver) loads the cubin and launches a
grid of 97 program instances, each a thread block of
num_warps * 32 threads. Every program computes its
own program_id, its tile of offsets, its mask, and
does its two loads, one add, and one store. The masked-off lanes
in program 96 quietly do nothing. That closes the loop, Python
call in, one lazily compiled and forever cached cubin, a grid of
tiles out, and the sum sitting in out.
Part V: Internals deep dives
Deep dive: the language, tiles, and block pointers
The user-facing language is triton.language,
conventionally imported as tl. Its core is a small,
deliberate vocabulary over tiles. Index construction with
tl.program_id, tl.num_programs, and
tl.arange. Memory with tl.load and
tl.store, each taking a tile of pointers plus an
optional mask and, for loads, an other
fill value. Construction with tl.zeros,
tl.full, and broadcasting. Elementwise math and
tl.where for data-dependent selection. Reductions
such as tl.sum, tl.max,
tl.min, and scans such as tl.cumsum, all
along a named axis. Atomics such as
tl.atomic_add for cross-program accumulation. And the
matmul primitive tl.dot(a, b), which multiplies two
tiles on the tensor cores and accumulates in fp32, subject to
minimum tile sizes and a precision knob (tf32 versus full
precision). The mental discipline is that everything is a tile
whose shape is known at compile time, so control flow that depends
on tile data is expressed with masks and
tl.where, not with Python branching over values.
Raw pointer arithmetic like x_ptr + offsets is fine
for one-dimensional access, but for tiled traversal of a matrix it
gets error-prone, so Triton offers block pointers.
tl.make_block_ptr(base, shape, strides, offsets,
block_shape, order) builds a structured pointer into a
tensor that knows the full tensor shape, the per-dimension strides,
the current block offsets, the tile shape, and the memory order of
the dimensions. You then tl.load(bptr,
boundary_check=(0, 1), padding_option="zero") to load a tile
with automatic edge handling, and tl.advance(bptr, (0,
BLOCK_K)) to slide the window along a dimension. Block
pointers move the index bookkeeping into the compiler, and on
newer hardware they map onto bulk tensor-memory transfers, so a
matmul main loop reads as advance, load A tile, load B tile,
accumulate with tl.dot, repeat. The illustrative
skeleton, with the boilerplate elided.
a_bp = tl.make_block_ptr(a_ptr, (M, K), (stride_am, stride_ak),
(pid_m * BLOCK_M, 0), (BLOCK_M, BLOCK_K), order=(1, 0))
b_bp = tl.make_block_ptr(b_ptr, (K, N), (stride_bk, stride_bn),
(0, pid_n * BLOCK_N), (BLOCK_K, BLOCK_N), order=(1, 0))
acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
for k in range(0, K, BLOCK_K):
a = tl.load(a_bp, boundary_check=(0, 1), padding_option="zero")
b = tl.load(b_bp, boundary_check=(0, 1), padding_option="zero")
acc += tl.dot(a, b) # tensor-core matmul, fp32 accumulate
a_bp = tl.advance(a_bp, (0, BLOCK_K))
b_bp = tl.advance(b_bp, (BLOCK_K, 0))
c = acc.to(tl.float16)
The order argument names which dimension is
contiguous in memory, which the compiler needs to choose a
coalesced layout, and it is a frequent source of confusion because
getting it wrong is correct but slow.
Deep dive: autotuning
The best tile sizes and the best num_warps and
num_stages depend on the problem shape and the GPU,
and Triton's answer is to search. The @triton.autotune
decorator takes a list of candidate triton.Config
objects and a key list of argument names. On the first
call for a given set of key values, the autotuner benchmarks every
config and caches the winner, and it re-searches only when a key
value changes.
@triton.autotune(
configs=[
triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 64, "GROUP_M": 8},
num_warps=8, num_stages=3),
triton.Config({"BLOCK_M": 64, "BLOCK_N": 64, "BLOCK_K": 32, "GROUP_M": 8},
num_warps=4, num_stages=4),
# ... more candidates ...
],
key=["M", "N", "K"], # re-tune only when these change
)
@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K, ...): ...
Two knobs deserve names. num_warps is how many warps
of 32 threads execute one program instance, so it sets the thread
block size and trades parallelism for register pressure, with 4 a
common default. num_stages is the depth of software
pipelining for loops that feed tl.dot, that is how
many iterations ahead the compiler prefetches operand tiles into
shared memory to overlap memory with compute, at the cost of more
shared memory. A companion decorator, @triton.heuristics,
computes a meta-parameter from the arguments at launch time, for
instance setting an EVEN_K flag when
K % BLOCK_K == 0 so the compiler can drop the masked
tail of the K loop. The trap worth stating, autotuning benchmarks
real launches, so the very first call to a freshly keyed kernel
pays for the whole search, which is why libraries pre-warm their
kernels and why an autotuned kernel feels slow exactly once.
Deep dive: the MLIR compiler stack and layouts
The reason Triton can be both simple to write and fast is that the
hard decisions live in the compiler, and since the MLIR rebuild
those decisions are explicit passes over typed IR rather than
hidden inside a monolith. The two dialects that matter most are
tt (Triton IR, hardware-agnostic tiles) and
ttg (Triton GPU IR, tiles plus layouts), with an
additional NVIDIA-specific dialect for features such as the
Hopper tensor-memory-accelerator path and warp specialization.
The load-bearing concept is the layout attribute. A
blocked layout is defined by how many elements each thread holds,
how threads fill a warp, and how warps fill the program instance,
plus an order, and choosing it well is what makes a global load
coalesce into a few wide transactions instead of many scattered
ones. An MMA layout describes the peculiar element distribution
that the tensor-core instructions produce, and much of the
compiler's cleverness is inserting the minimum number of
layout conversions (which cost shared-memory round trips) to move
data between the layout a load prefers and the layout a
tl.dot requires.
Once data layout is an explicit, typed attribute on every tile value, the compiler can reason about coalescing, tensor-core operand formats, and shared-memory staging as ordinary rewrites, and the programmer is freed from all of it while keeping control of the one thing that needs human judgment, the block decomposition. This is the same move that torchtitan makes at the cluster scale with DTensor placements, layout is metadata the system reasons over, told here at the granularity of threads within one GPU rather than GPUs within a cluster.
Deep dive: specialization, caching, and the driver
The runtime is what turns a Python function into a fast, reusable
binary. The JITFunction owns the specialization logic
from Stage 2, and its correctness rests on a subtle point, the
alignment and divisibility facts are assumptions the compiler is
allowed to bake in, so a kernel compiled under the assumption that
a pointer is 16-byte aligned must never be launched on an
unaligned pointer. The specializer therefore encodes those facts
into the cache key, and you can opt a given argument out with
do_not_specialize when the assumption would be unsafe
or would cause needless recompilation. Compiled kernels are cached
both in memory on the JITFunction and on disk under
~/.triton/cache (relocatable with
TRITON_CACHE_DIR), keyed by the signature and target,
so the compile cost is paid once per machine, not once per
process. The driver layer abstracts CUDA and HIP behind a common
interface for allocation, module loading, and launch, which is how
one launch path serves both vendors. The practical upshot is that
a Triton kernel behaves like a normal Python function after its
first call, and the cost model in your head should be, compile
once per new signature, then a cheap cache lookup and a driver
launch every time after.
Part VI: Reading the repository
The project splits cleanly into a Python front-end and runtime, a C++ MLIR compiler core, and per-vendor backends. Read it in that order. All locations described by role, since exact paths shift between releases.
Stage 0, the tutorials. Start in the
tutorials directory (under the Python package). Read
01-vector-add, then 02-fused-softmax,
then 03-matrix-multiplication, then
06-fused-attention. These are the intended
curriculum, each is runnable and benchmarked with
triton.testing, and they introduce program ids and
masks, then reductions and fusion, then tl.dot,
autotuning, and the L2-cache grouping trick, then a full
FlashAttention-style kernel. Question, at each step, which line
decides the block decomposition and which lines are pure tile
arithmetic the compiler will lower.
Stage 1, the language surface. Read the
triton/language package, especially the core module
that defines tiles, load, store,
dot, and the reductions, and the standard-library
module built on top of them. Question, which operations are
primitive (they lower directly to tt ops) versus
composed in Python from other tile ops?
Stage 2, the runtime. Read the JIT and autotuner
modules in triton/runtime. The JITFunction
is where specialization, the cache key, and the launch path live,
and the autotuner is short and readable, it is little more than
benchmark every config, remember the best per key. Question, what
exactly goes into the specialization signature, and where does the
divisibility-by-16 assumption enter?
Stage 3, the compiler driver. Read the compiler
package (triton/compiler), the code generator that
walks the AST into Triton IR and the top-level
compile function that sequences the pass pipeline down
to a binary. Question, in what order do the TTIR and TTGIR passes
run, and where is the target backend selected?
Stage 4, the MLIR core. Drop into C++. The
dialect definitions live under the include tree
(include/triton) and the passes under
lib. Read one pass end to end, the coalescing pass or
the matmul-acceleration pass are the most instructive, against the
dialect ops they rewrite. Question, how does a pass decide a
blocked layout, and what makes a layout conversion necessary?
Stage 5, the backends. The vendor backends live
under third_party/nvidia and
third_party/amd, each lowering TTGIR to LLVM and then
to PTX or AMDGCN and driving the assembler. Skim these to see the
pluggable seam, then stop, they are the deepest and most
hardware-specific part of the tree.
Where not to start. The Hopper-specific machinery (tensor memory accelerator paths, warp specialization, cluster launch) is fascinating but a distraction on a first read, and the C++ pass internals only make sense once you have watched the IR change in the labs. Read the IR before you read the pass that produced it.
Part VII: Hands-on labs
Labs 1 through 4 need one GPU. Lab 5 needs none beyond reading. Log and IR formats vary with the fast pace of the project.
Lab 1: vector-add and verify. Concept: the launch grid and masking.
# run the Part II add() and check it
import torch, triton
x = torch.rand(100_000, device="cuda"); y = torch.rand(100_000, device="cuda")
torch.testing.assert_close(add(x, y), x + y)
# now break it on purpose: drop the mask in add_kernel and rerun.
# on a size that is not a multiple of BLOCK_SIZE, watch it fault or corrupt.
Observe that the correct version passes and the mask-free version
fails only for sizes that are not a multiple of
BLOCK_SIZE, which is exactly the ragged final block.
Lab 2: see the compiler's output. Concept: the TTIR to PTX pipeline of Part IV.
# dump every IR stage to disk for any kernel you launch
TRITON_KERNEL_DUMP=1 TRITON_DUMP_DIR=./dump python your_kernel.py
# or dump the MLIR after each pass to stderr
MLIR_ENABLE_DUMP=1 python your_kernel.py 2> passes.log# or reach the strings directly from a compiled kernel
k = add_kernel.warmup(x, y, out, x.numel(), BLOCK_SIZE=1024, grid=(1,))
print(k.asm["ttir"][:800]) # hardware-agnostic tile IR
print(k.asm["ttgir"][:800]) # note the #blocked layout attributes
print(k.asm["ptx"][:800]) # the assembly ptxas turns into a cubin
Read the TTIR and find tt.load, tt.store,
and the range. Then read the TTGIR and find the
#blocked layout the coalescing pass chose. Then skim
the PTX and see the tile become per-thread loads. This single lab
makes Part IV concrete.
Lab 3: fuse a softmax and benchmark it. Concept: memory traffic and fusion.
import torch, triton
x = torch.randn(4096, 4096, device="cuda")
ref = torch.softmax(x, dim=1)
torch.testing.assert_close(softmax(x), ref, atol=1e-3, rtol=1e-3)
print(triton.testing.do_bench(lambda: softmax(x))) # your kernel, ms
print(triton.testing.do_bench(lambda: torch.softmax(x, 1))) # torch reference
Observe that the fused kernel moves the input across DRAM far
fewer times than an unfused sequence would, and compare the
timings. Then shrink the number of columns and watch the chosen
BLOCK_SIZE follow next_power_of_2.
Lab 4: autotune a matmul. Concept: the tile-size search of Part V.
# start from tutorial 03, then vary the config list and the key.
# 1) time a fixed (M, N, K) with the default configs.
# 2) remove all but one config and compare; feel the search cost on call one.
# 3) add num_stages=2 vs 4 variants and see which the tuner keeps.
# TRITON_PRINT_AUTOTUNING=1 makes the autotuner announce its winner.
Observe that the first call for a new (M, N, K) is
slow because it benchmarks every config, and later calls with the
same key are fast. This is the pre-warm behavior real libraries
rely on.
Lab 5: read a real kernel. Concept: block pointers and tl.dot at production scale.
# open the fused-attention tutorial (06) and read it against the
# FlashAttention writeup on this site. map every tl.load/tl.dot to a tile
# in the flash-attention algorithm, and find where the online-softmax
# running max and sum are carried across the K/V loop.Observe that the same tiling and online-softmax structure derived in the FlashAttention chapter and the attention page appears here as a loop over key/value blocks that never materializes the full score matrix. Triton is the language that makes that algorithm a few dozen readable lines.
Part VIII: Questions and model answers
Understanding checks. Answer aloud before reading.
1. What is Triton, in one sentence?
A Python-embedded DSL and MLIR-based compiler for GPU kernels, in which you write a single program over tiles and the compiler lowers the tile operations into coalesced, synchronized thread-level code, so you own the block decomposition and hand the per-thread details to the compiler.
2. How does the Triton programming model differ from CUDA's?
CUDA is SIMT, you write per-thread code and manage
threadIdx, shared memory, coalescing, and barriers by
hand. Triton is block-level SPMD, you write per-tile code indexed
by program_id, and the compiler decides the
thread mapping, allocates shared memory, coalesces loads, and
inserts synchronization. You keep control of grid and tile sizes
because that is the algorithmic choice.
3. Walk the vector-add kernel line by line.
program_id selects this block, multiplying by
BLOCK_SIZE and adding tl.arange builds a
tile of offsets, the mask offsets < n_elements
guards the ragged last block, two masked tl.loads
bring tiles of x and y on-chip, they are added as tiles, and a
masked tl.store writes the result. The grid launches
ceil(n / BLOCK_SIZE) programs, one per tile.
4. Why is the fused softmax faster than a naive PyTorch softmax?
Softmax is memory-bound, and the naive version makes several passes over DRAM, one to find the max, one for the exponentials, one for the sum, one to normalize. The fused kernel loads each row once into on-chip memory, does the max, exp, sum, and divide there with tile reductions, and writes the row once, so it moves far fewer bytes across the slow memory.
5. What is a tile, and why is it the unit of the language?
A tile is a small multi-dimensional block of values with a shape known at compile time, treated as a first-class value. Making it the unit lets the programmer express the algorithm at block granularity while giving the compiler a whole tile to map onto threads, shared memory, and tensor cores, which is where the optimization lives.
6. What are the four IR stages a kernel passes through?
Triton IR (the tt dialect, hardware-agnostic tiles),
Triton GPU IR (the ttg dialect, tiles plus layouts
and warp counts), LLVM IR, and finally PTX assembled into a cubin
on NVIDIA (or AMDGCN on AMD). The pivotal stage is TTGIR, where
layouts are chosen.
7. What is a layout in Triton GPU IR and why does it matter?
A layout is a typed attribute on a tile value that says exactly how the tile's elements are distributed across threads, warps, and shared memory. It matters because coalescing, tensor-core operand formats, and shared-memory staging all depend on it, and the compiler's job is largely choosing good layouts and inserting the fewest conversions between them.
8. What do num_warps and num_stages control?
num_warps is how many 32-thread warps run one program
instance, setting the thread block size and trading parallelism
against register pressure. num_stages is the depth of
software pipelining for loops feeding tl.dot, how many
iterations of operand tiles are prefetched into shared memory to
overlap memory with compute, at the cost of more shared memory.
9. What does autotuning actually do, and what is its cost?
@triton.autotune benchmarks a list of candidate
configs on the real hardware and caches the fastest per distinct
value of the key arguments. The cost is that the first
call for a new key runs every config, so that one call is slow,
which is why libraries pre-warm autotuned kernels.
10. Why is the first call to a kernel slow and later calls fast?
The first call with a new specialization signature triggers full
compilation, AST to TTIR to TTGIR to LLVM to PTX to cubin. The
result is cached in memory and on disk under
~/.triton/cache, keyed by argument types, constexpr
values, alignment facts, and target, so every later call with the
same signature is a cache lookup and a driver launch.
11. When would you reach for CUTLASS or hand-written CUDA instead of Triton?
When you need the last few percent of peak on a fixed shape and a single GPU generation, or a hardware feature the DSL has not surfaced. Vendor libraries and CUTLASS win top-end GEMM and convolution, and raw CUDA gives instruction-level control. Triton's advantage is custom fused kernels written fast and run on more than one vendor.
12. Why do block pointers exist when raw pointer arithmetic already works?
Raw arithmetic is fine in one dimension but error-prone for tiled
matrix traversal. tl.make_block_ptr carries the tensor
shape, strides, offsets, tile shape, and memory order, gives
automatic boundary checking on load, and slides with
tl.advance. It also maps onto bulk tensor-memory
transfers on newer hardware, so a matmul loop becomes advance,
load, tl.dot, repeat.
13. Where does the divisibility-by-16 assumption come from and why guard it?
The specializer records whether pointers are 16-byte aligned so the compiler can emit wider vectorized loads. Because that assumption is baked into the compiled kernel, it becomes part of the cache key, and a kernel compiled for aligned pointers must never be launched on an unaligned one, which is why alignment is a specialization axis rather than a runtime branch.
14. Why can print debugging not work inside a kernel?
The kernel body is compiled and executed on the GPU across a grid
of programs, so a host-side Python print in the body
runs only at trace time, not per program. You use
tl.device_print, shrink to a single block and compare
against a torch reference, or dump the generated IR to reason
about what actually ran.
Part IX: Design lessons
Raise the unit of programming, keep the strategic choice. Triton moves the programmer from the thread to the tile, automating the tedious and error-prone per-thread work while leaving the one decision that needs human judgment, the block decomposition, in human hands. Good abstractions hide the mechanism and expose the policy, the same instinct as a query planner that owns execution while you own the schema.
Make layout an explicit type. Encoding data distribution as layout attributes on tile values turns coalescing, tensor-core formats, and shared-memory staging into ordinary rewrites over typed IR. Represent the invariant in the data and composition becomes checkable, which is exactly why torchtitan makes cluster-scale sharding a DTensor placement rather than model code.
Compile lazily, cache aggressively, specialize honestly. A kernel compiles on first use per signature and is then cached in memory and on disk, so the abstraction is free after the first call, and the assumptions the compiler bakes in (alignment, constexpr values) are encoded into the cache key so they can never be silently violated. Caching is only safe when the key captures every assumption.
Search the parameters you cannot reason about. The best tile size and pipeline depth depend on shape and hardware in ways no simple formula captures, so autotuning measures instead of guessing and caches the winner. When the cost model is too complex to model, benchmark, and remember.
Keep the seams pluggable. The MLIR dialects and
the vendor backends under third_party let a new
accelerator consume Triton GPU IR without touching the front-end,
which is how one language reaches NVIDIA and AMD and why Pallas can
sit on top. A stable IR in the middle is what lets the ends move
independently.
Part X: Memorization framework
The one-sentence summary. Triton lets you write a GPU kernel
as a single program over compile-time-sized tiles, indexed by
program_id and moved with tl.load,
tl.store, and tl.dot, and its MLIR
compiler lowers those tiles through Triton IR and Triton GPU IR
(where layouts are chosen) to LLVM and PTX, caching one cubin per
signature, while autotuning searches the block sizes you should not
guess.
@triton.jit -> kernel[grid](args, BLOCK=...) -> specialize + cache key -> code gen -> Triton IR (tt) -> Triton GPU IR (ttg, layouts) -> LLVM IR -> PTX -> cubin (ptxas) -> driver launches grid of programs each program: pid = program_id -> tile of offsets + mask -> tl.load -> compute (tl.dot / reductions) -> tl.store
The chain mapped to the repository.
language triton/language/ (tl: program_id, load, store, dot, reduce) launch + cache triton/runtime/ (JITFunction: specialize, cache key, launch) autotune triton/runtime/ (autotune, Config, heuristics) code gen + driver triton/compiler/ (AST -> TTIR, then the pass pipeline) MLIR passes lib/ + include/triton/ (tt and ttg dialects, layout passes) backends third_party/nvidia, third_party/amd (TTGIR -> LLVM -> PTX/GCN) tutorials tutorials/ (01 vector-add ... 06 fused-attention)
Memorize these blocks.
- Programming model: per-block SPMD over tiles,
program_idpicks the block, masks guard the edges, the compiler owns threads, coalescing, shared memory, and sync. - Core language:
tl.program_id,tl.arange,tl.load/tl.storewithmask/other,tl.dot, reductionstl.sum/tl.max,tl.where, block pointers viatl.make_block_ptr/tl.advance. - IR stages: Triton IR (tt, agnostic) to Triton GPU IR (ttg, layouts and warps) to LLVM IR to PTX/cubin, with TTGIR the pivotal layout-choosing stage.
- Autotuning:
@triton.autotunebenchmarks configs and caches the best perkey,num_warpssets thread count,num_stagessets pipeline depth. - Cost model: compile once per signature (types, constexpr, 16-byte alignment, target), cache in memory and under
~/.triton/cache, cheap launch forever after.
Part XI: Papers and further reading
The ideas in this chapter trace back to a small set of papers and projects, and each one rewards a direct read. Where this site derives the same idea in depth, the companion link points there.
- Tillet, Kung, and Cox, Triton, An Intermediate Language and Compiler for Tiled Neural Network Computations, MAPL 2019. The original paper, which raises the unit of programming to the tile and sketches the coalescing, shared-memory, and pipelining machinery this chapter walks through.
- Lattner et al., MLIR, A Compiler Infrastructure for the End of Moore's Law, 2020. The extensible compiler framework Triton was rebuilt on, whose dialect and pass machinery carries the
ttandttgIRs of Part IV. - Lattner and Adve, LLVM, A Compilation Framework for Lifelong Program Analysis and Transformation, CGO 2004. The substrate underneath everything here, since every Triton kernel becomes LLVM IR on its way to PTX or AMDGCN.
- Dao et al., FlashAttention, Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022. The flagship algorithm for hand-written Triton kernels, derived in the FlashAttention chapter, and the Unsloth walkthrough shows kernels in that style at production scale.
- Milakov and Gimelshein, Online Normalizer Calculation for Softmax, 2018. The one-pass softmax trick that the fused-attention tutorial carries across its key and value loop, worked on the softmax page.
- Williams, Waterman, and Patterson, Roofline, An Insightful Visual Performance Model for Multicore Architectures, CACM 2009. The model that explains why fusion wins, since kernels like softmax sit on the memory-bound side of the roof, and the parallel computing class builds on it.
- Ansel et al., PyTorch 2, Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation, ASPLOS 2024. Describes TorchDynamo and TorchInductor, the compiler that emits most of the world's Triton kernels, covered in the PyTorch walkthrough.
- Ragan-Kelley et al., Halide, A Language and Compiler for Optimizing Parallelism, Locality, and Recomputation in Image Processing Pipelines, PLDI 2013. The paper that separated the algorithm from the schedule, the ancestor idea behind every kernel DSL including this one.
- Chen et al., TVM, An Automated End-to-End Optimizing Compiler for Deep Learning, 2018. A graph-level tensor compiler for whole models, a useful contrast with Triton's single-kernel scope.
- NVIDIA, CUTLASS. The template-metaprogramming route to peak GEMM that Part III weighs against Triton, covered in the CUTLASS walkthrough.
- The JAX team, Pallas, a JAX Kernel Language. A sibling DSL that lowers to Triton on GPU, living proof of the pluggable seam praised in Part IX.
Part XII: Final takeaway
If the GPU execution model underneath, warps, memory coalescing, shared memory, and occupancy, is the gap, the parallel computing class builds it from the hardware up, and the algorithm that most rewards a hand-written Triton kernel is derived in the FlashAttention chapter and the softmax page. Then come back and read the fused-attention tutorial once more, it will read like tiled NumPy with a couple of loop-carried scalars, which is the entire point.