Chapter 3: Introduction to Triton
Meet Triton: what it is and why it exists, how it compares to writing raw CUDA or staying in PyTorch, its block-oriented programming model, the vocabulary of programs / grids / tiles, and the @triton.jit decorator that turns a Python function into a GPU kernel.
- Topics
- 5
- Demo
- 1
- Videos
- 3
- Min read
- 16
Chapter Overview
You now understand the hardware. The question is how to program it without drowning in low-level detail. Triton is the answer this course is built around: a Python-embedded language and compiler, originally from OpenAI, that lets you write GPU kernels that rival hand-tuned CUDA — in a fraction of the code.
The key idea is that Triton raises the unit of programming from the thread to the block. In CUDA you write code from the viewpoint of a single thread and manually manage indexing, shared memory, and synchronization. In Triton you write code from the viewpoint of a program that operates on whole tiles of data as vectors and matrices. The compiler handles the painful parts — thread mapping, memory coalescing, shared-memory allocation, and synchronization — while leaving you in control of the decisions that actually matter for performance: how big the tiles are and how data moves between DRAM and SRAM.
This chapter is the conceptual bridge between the hardware chapters and the hands-on kernels that follow. By the end you'll be able to read a Triton kernel and understand exactly what each line does and where it runs.
This chapter covers:
- What Is Triton?: a Python DSL + compiler for GPU kernels
- Triton vs CUDA vs PyTorch: the right tool for each job
- Block-Oriented Programming: thinking in tiles, not threads
- Programs, Grids & Tiles: the core vocabulary
- The @triton.jit Decorator: how a Python function becomes a kernel
Chapter Roadmap
Click any topic to jump in
What Is Triton?
A Python DSL + compiler that emits fast GPU kernels from tile-level code.
Triton vs CUDA vs PyTorch
Control where it matters, automation where it's painful — between raw CUDA and stock PyTorch ops.
Block-Oriented Programming
Reason about tiles of data as vectors/matrices; the compiler maps them onto threads.
Programs, Grids & Tiles
A program owns a tile; the grid = ceil(work / tile) covers all the data.
The @triton.jit Decorator
Marks a kernel, enables [grid] launch, and specializes on tl.constexpr constants.
Topics
5 topics in this chapter
Chapter 2 ended with an uncomfortable conclusion: most deep-learning kernels are memory-bound, so speed comes from moving fewer bytes, staging tiles in shared memory, coalescing loads and fusing operations. Doing that in CUDA means managing thread indices, shared-memory buffers, barriers and vector widths by hand, and a fast matmul written that way runs to hundreds of lines. Most ML researchers never write one, so they are stuck with whatever kernels the framework ships. Triton was built to close that gap. It is a Python-embedded language in which a kernel is written over blocks of data, plus a compiler that turns those block operations into the thread-level code an expert would write. This topic explains what Triton actually is: a decorator, a small tl library of tile operations, and a multi-stage compiler pipeline. On an RTX 2060, a 13-line Triton vector add compiles in about 1.7 seconds on first call and then matches torch.add at about 280 GB/s. The next topic places Triton between PyTorch and CUDA, and the rest of the chapter builds its vocabulary: tiles, programs, grids and the @triton.jit boundary.
Definition
Triton is an open-source, Python-embedded domain-specific language and compiler for GPU kernels. A Triton kernel is a function decorated with @triton.jit whose body manipulates blocks of values through triton.language operations. The compiler lowers it through Triton IR, TritonGPU IR and LLVM IR to PTX and a GPU binary, deciding thread mapping, vectorization and shared-memory use itself.
A Python DSL with a Real Compiler
Triton kernels are Python functions decorated with @triton.jit. Inside, you call triton.language (tl) operations such as tl.load, tl.store, tl.arange, tl.dot and tl.sum, which act on whole tiles rather than scalars. The Python is never interpreted on the GPU. On first launch, Triton parses the function into Triton IR, lowers it to TritonGPU IR (where it picks a thread layout, vector widths and shared-memory staging), then to LLVM IR, then PTX, and finally a cubin binary that the driver loads. On the RTX 2060, a vector add yields 232 lines of PTX using 128-bit ld.global.v4 loads. You never wrote a thread index or __syncthreads.
Think of Triton as raising the abstraction level without hiding the performance model. A CUDA kernel exposes one decision per thread: which index it touches, when it synchronizes. A Triton kernel exposes one decision per tile and lets the compiler expand each tile op into SIMT code. For a 1,024-element tile with the default 4 warps, that expansion is elements per thread, issued as two 4-wide vector loads per input. You keep the high-leverage knobs (tile shape, what to load, what to fuse) and shed the low-leverage ones.
Why can a researcher prototype a custom fused kernel in Triton in an afternoon, when the equivalent CUDA kernel might take days?
The previous topic showed what Triton is, but not when to use it, and that choice matters because every option has a cost. Eager PyTorch is the default: each operator is a well-tuned, precompiled kernel and costs nothing to write. Its weakness is composition. Each op reads its inputs from DRAM and writes its result back, so a chain of four elementwise ops makes four round trips. On an RTX 2060, (x*2+1).relu()*3 on 50 million floats takes 5.9 ms eagerly and 1.5 ms as one fused Triton kernel, a 4.0x gap that is pure memory traffic. Raw CUDA sits at the other extreme: it can express anything, including warp shuffles, explicit shared-memory layouts and inline PTX, but the developer pays in lines of code and debugging time. Triton occupies the middle, giving control over tiling, fusion and data movement while automating thread-level detail. This topic lays out that spectrum, gives a rule of thumb for each tool, and notes that torch.compile itself generates Triton kernels, so much of the ecosystem already runs on it. The next topic explains the block-oriented model that makes that middle ground possible.
Definition
The flexibility-performance spectrum ranks GPU programming tools by how much of the kernel the programmer controls. Eager framework ops give no control over fusion or memory movement; Triton exposes tile shapes, loads, stores and fusion while compiling thread mapping automatically; raw CUDA exposes every thread, register and shared-memory byte at the highest engineering cost.
The Flexibility–Performance Spectrum
Eager PyTorch runs each op as a separate precompiled kernel. That is ideal for standard ops, but a chain of ops round-trips every intermediate through DRAM, and a novel fused operation cannot be expressed at all.
Raw CUDA offers total control (warp shuffles, explicit shared-memory layouts, inline PTX) at the cost of hundreds of lines and manual tuning.
Triton controls what matters for memory-bound work, namely tiling, fusion and data movement, and automates the painful thread-level detail. Reach for it when built-in ops are too slow or cannot express your fusion. Reach for CUDA when you need a hardware feature Triton does not expose.
For a chain of memory-bound elementwise ops over elements, eager execution moves roughly values through DRAM (one read and one write per op), while one fused kernel moves about . The speedup ceiling is about , and it is reached when every op is purely bandwidth-bound. Measured on an RTX 2060 for (multiply, add, ReLU, multiply) at : 5.94 ms eager versus 1.48 ms fused, a ratio of 4.01.
You have y = softmax(x) followed by a custom masking step, and profiling shows the op is memory-bound because PyTorch materializes intermediates in DRAM. What's the right tool?
Theory Exercise
Problem:
For each task, pick PyTorch eager, Triton, or raw CUDA, and justify it: (a) a one-off relu(x) on a tensor in a research script; (b) a novel fused attention variant that needs to beat the default for a paper; (c) a kernel that must squeeze out the last 2% on a specific GPU for a shipping inference server with a dedicated systems team.
Hints:
- How much does performance matter vs developer time for each?
- Can the built-in op already express it efficiently?
- Who maintains the code and how specialized is the target?
The previous topic placed Triton between PyTorch and CUDA but did not say how it automates thread-level work while still producing fast code. The answer is a change in the unit of programming. CUDA uses the single-program, multiple-data model at thread granularity: you write what one thread does to one element and launch hundreds of thousands of threads. The hardware, however, does not care about single threads. Memory moves in 32-byte sectors, warps issue 32 lanes together, and shared memory holds tiles. Thread-level code makes the programmer reassemble those blocks by hand, which is where coalescing bugs and missing barriers come from. Triton's designers flipped the model: each program instance is scalar, but the values it manipulates are blocks, so you write out = a + b over 1,024-element tiles. The Triton documentation calls this "blocked program, scalar threads" versus CUDA's "scalar program, blocked threads". This topic shows what the compiler does with a tile: on an RTX 2060, a 1,024-element add becomes 4 warps of 128 threads, each handling 8 contiguous floats with 128-bit loads. The next topic names the pieces: programs, grids and tiles.
Definition
Block-oriented programming is a single-program, multiple-data model in which each program instance operates on blocks (tiles) of values instead of single elements. The programmer specifies tile-level loads, arithmetic and stores; the compiler partitions each tile across the threads of one thread block and chooses vector widths, coalescing and shared-memory staging.
Tiles, Not Threads
In CUDA you write out[i] = a[i] + b[i] from one thread's perspective and launch thousands of threads. In Triton you write out = a + b where a, b, and out are tiles — vectors of, say, 1024 elements — and you launch a handful of programs, each responsible for one tile.
This block-oriented view has two big payoffs. First, it's how you naturally express memory movement: tl.load pulls a whole tile from DRAM into on-chip memory in one coalesced operation. Second, the compiler is free to choose the best thread-level implementation — vectorizing, coalescing, and scheduling — because you only specified the tile-level algorithm. You think in the same units the memory hierarchy cares about (blocks that fit in SRAM), which is exactly why Triton code tends to be fast by default.
If the data has elements and the tile size is , the programmer reasons about programs instead of threads. Each program is executed by threads for warps, so each thread handles elements: 8 for , . The compiler chooses that factor and the vector width, which is why the same source can run well on GPUs with different preferred access widths.
A Triton kernel declares BLOCK_SIZE = 1024 and does x = tl.load(ptr + offs) where offs is a 1024-wide range. How many threads run this, and what is the programmer reasoning about?
Related Problems on PixelBank
Block-oriented programming from the previous topic says each program instance works on a tile, but a real input is larger than one tile. Someone has to decide how many instances to launch, which slice each one owns, and what happens when the data does not divide evenly. Triton answers with three words used in every kernel. A program is one running instance of the kernel, identified by tl.program_id. The grid is the set of programs launched by one call, laid out along up to three axes. A tile is the block of data a single program owns, with a shape fixed by compile-time constants. Grid size follows directly from the other two: cover the work with tiles and round up, triton.cdiv(n, BLOCK_SIZE). Rounding up guarantees coverage but creates a ragged last tile, so the overhanging lanes must be masked off. This topic works through the arithmetic for a 5,000-element vector with 1,024-element tiles: five programs, the last of which owns only 904 valid elements. It also shows how a 2D grid gives matmul one program per output tile. The next topic covers @triton.jit, which turns the kernel into something you can launch over that grid.
Definition
In Triton, a program is one instance of a kernel, identified along each grid axis by tl.program_id(axis). The grid is the 1D, 2D or 3D set of programs launched together, normally sized as the ceiling of the work divided by the tile shape. A tile is the block of elements that one program loads, computes and stores, with compile-time dimensions.
The Core Vocabulary
A program is one running instance of your @triton.jit kernel; each gets a unique id via tl.program_id(axis). The grid is the set of all programs you launch for one call — it can be 1D, 2D, or 3D, sized to cover your output. A tile is the contiguous block of data (vector or matrix) that a single program is responsible for, with shape set by compile-time BLOCK_SIZE constants.
The pattern is always the same: pick a tile shape, compute the grid as ceil(work / tile) using triton.cdiv, and inside the kernel use tl.program_id to figure out which slice of data this program owns. A 2D grid (e.g. for matmul) gives each program two ids — one per output-tile axis.
Using guarantees coverage: the last program handles indices , of which (or all if it divides evenly) are valid. The wasted lanes are at most , negligible for large , and masking makes the overhang correct rather than dangerous.
You want to process a vector of n = 5000 elements with BLOCK_SIZE = 1024. How big is the grid, which elements does program 3 own, and what's special about the last program?
The previous topic defined programs, grids and tiles, but a Python function cannot run on a GPU as written, and Triton must decide when to compile it and what to specialize on. @triton.jit is that boundary. Decorating a function turns it into a JIT-compiled kernel object that you launch with square-bracket syntax, kernel[grid](args). Nothing is compiled when the decorator runs. On the first launch, Triton inspects the arguments, builds a specialization key, compiles through the IR pipeline described earlier in this chapter, and caches the binary in memory and on disk. On the RTX 2060 that first call took 1.68 seconds; the second, including execution, took 78 microseconds. The key includes every tl.constexpr value, such as BLOCK_SIZE, but also less obvious properties of ordinary arguments: whether an integer equals 1 and whether pointers and integers are divisible by 16. This topic explains which arguments are compile-time and which are runtime, and how many variants a real workload produces: six for one add kernel in the probe. Chapter 4 then uses all of this to write and launch a complete vector-add kernel.
Definition
@triton.jit marks a Python function as a Triton kernel compiled just in time. At first launch for a given specialization key, built from argument types, tl.constexpr values, and integer and pointer properties such as divisibility by 16, Triton compiles the body to a GPU binary and caches it. Later launches with the same key reuse that binary.
From Python Function to GPU Kernel
Decorating a function with @triton.jit tells Triton to compile it to a GPU kernel the first time it's launched (just-in-time). You launch it with the grid in square brackets: kernel[grid](args...). Inside, you can only use Triton-compatible operations — tl.* ops, control flow on tile values, and arguments.
Arguments marked tl.constexpr (like BLOCK_SIZE) are compile-time constants: the compiler bakes them in and specializes the kernel for each distinct value, enabling loop unrolling and better register allocation. Pointers (your tensors) and runtime sizes (n) are ordinary runtime arguments. The first call to a given specialization pays a one-time compile cost; subsequent calls reuse the cached machine code, so it's effectively free.
The kernel cache is keyed by argument types, every tl.constexpr value, and specialization flags on ordinary arguments: integers equal to 1 or divisible by 16, and pointers aligned to 16 bytes. Total cost is about for variants, a compile cost near 1.7 s, calls and per-call host overhead of a few microseconds. In the probe, one add kernel reached after varied lengths, two extra block sizes and a misaligned view. Because , compilation amortizes away. do_not_specialize removes a flag when variants multiply.
In add_kernel[grid](x, y, out, n, BLOCK_SIZE=1024), why is BLOCK_SIZE passed as a keyword tl.constexpr while n is a normal argument?