How to Learn Triton, from Your First Kernel to FlashAttention

How do you learn Triton?

You learn Triton by writing small GPU kernels in a fixed order, starting with one block of elements per program and masked loads, then moving to reductions and softmax, tiled matrix multiplication, fused kernels, autotuning and finally attention. TensorTonic's Triton Kernels study plan has 30 problems in six sections that follow this path, and every kernel runs on real GPU hardware from your browser and is compared with a PyTorch reference. This guide arranges those 30 problems into eight steps and explains what each step teaches.

This guide is the Triton companion to the CUDA guides on TensorTonic. The guide on how to learn CUDA explains the GPU execution model and memory hierarchy at the thread level, and the CUDA practice problems guide orders the 35 CUDA problems. The steps below focus on what changes when you write the same kinds of kernels in Triton.

The study plan groups the same problems into six sections, from foundations and the block model to performance patterns. This guide splits some of those sections into smaller steps and moves a few problems, such as dropout and the autotuned matmul, so that each step adds one new idea.

What is Triton, and how is it different from CUDA?

Triton is an open-source language and compiler for writing GPU kernels in Python, where each running instance of a kernel works on a whole block of data at once. In CUDA you write the code for a single thread and manage how threads share memory and wait for each other. In Triton you write the code for one program, which loads a block of values, computes on it as a small tensor and stores the result, and the compiler decides how that block is spread over the threads of the GPU.

The choices that decide speed stay with you in Triton. You still pick the block sizes, the shape of the launch grid, the order in which programs read memory and which operations to fuse. The compiler takes over thread indexing inside a block, shared-memory buffers and synchronisation, which in CUDA you would write and debug yourself.

What do you need to know before learning Triton?

You need comfortable Python, PyTorch tensors with their shapes and strides, and basic linear algebra before learning Triton, and you do not need any CUDA experience. Triton kernels take PyTorch tensors as arguments and work with raw pointers into their memory, so knowing how a tensor is laid out matters more than any GPU detail.

• For PyTorch, you should know what shape, dtype and stride mean and why a transposed tensor is not contiguous, which the PyTorch study plan covers.
• For linear algebra, you should be able to compute a matrix multiplication by hand and know what a dot product and a norm are.
• For the GPU itself, it is enough to know that global memory is large and slow while on-chip memory is small and fast, and the memory step of the CUDA learning guide explains why.

You also need access to a GPU to run kernels at full speed. TensorTonic runs CUDA and Triton kernels in the browser on several NVIDIA GPUs, from a T4 to an H200, so a laptop without a GPU is enough to work through this guide.

In what order should you learn Triton?

Learn Triton in the order the kernels depend on each other, because each step reuses the pattern from the one before it. The first three steps share one structure for elementwise work, the next two add reductions and two-dimensional tiles, and the last three use those pieces to build fused, tuned and attention kernels.

Step 1: How do program IDs and blocks work in Triton?

A Triton kernel runs once for every program in a launch grid, and tl.program_id tells each program which block of the data belongs to it. The program turns that ID into a block of offsets with tl.arange, adds the offsets to the base pointer of each input, and loads the whole block in one call. The grid is usually the number of elements divided by the block size, rounded up with triton.cdiv.

import triton
import triton.language as tl

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)

grid = lambda meta: (triton.cdiv(n, meta['BLOCK_SIZE']),)
add_kernel[grid](x, y, out, n, BLOCK_SIZE=1024)

The block size is marked tl.constexpr, which means it is fixed when the kernel is compiled and can be used to size tensors inside the kernel. The grid is a function of those compile-time values, so the same launch code still works when the block size changes. Vector Addition in Triton asks for exactly this kernel, and it is worth writing from memory before moving on.

Step 2: How do masked loads and stores work?

A mask is a block of true and false values that tells tl.load and tl.store which positions hold real elements, and it does the job of the bounds check you would write in CUDA. Because the grid is rounded up to whole blocks and tl.arange needs a power-of-two length, the last program usually covers positions past the end of the tensor. Loading or storing those positions without a mask reads or writes memory that belongs to something else.

A masked load can also take an other value, which is what the masked-out positions contain after the load. Elementwise kernels rarely need it, because those positions are never stored. Reductions depend on it, which is why the next steps keep coming back to masks. Fused Multiply-Add in Triton practises the same masked load and store pattern on a fused operation, and Masked Triton GELU Kernel from the Build LLM from Scratch plan gives you a second one.

Step 3: How do you write elementwise kernels in Triton?

Elementwise kernels apply the same function to every element, so the structure from the first two steps stays the same and only the math between the load and the store changes. ReLU needs tl.where or tl.maximum, and GELU and SiLU need exponentials or the error function. Work through ReLU in Triton, GELU in Triton and SiLU in Triton quickly, since their value is in repetition.

Dropout (Inverted Scaling) in Triton adds random numbers. Triton generates them with tl.rand from a seed and the element offsets, so the same seed always produces the same mask and the backward pass can rebuild it without storing it. The kept values are divided by one minus the drop probability so the expected output matches the input.

Step 4: How do reductions and softmax work in Triton?

A reduction in Triton is a single call such as tl.sum or tl.max along one axis of a block, so the hard part is deciding what each program reduces and how the partial results are combined. When a row fits in one block, the simplest design launches one program per row and reduces the whole row inside it. When a vector is too long for one block, each program reduces its own block and the partial results are combined with an atomic add or a second small kernel.

Masks matter more here than anywhere else. The other value for masked positions must be the neutral value of the reduction, which is zero for a sum and negative infinity for a maximum, or the padding changes the answer. Start with Vector Sum Reduction in Triton, Vector Max Reduction in Triton and L2 Vector Norm in Triton, then solve Single-Pass Mean and Variance in Triton, which computes two statistics from one read of the data.

Softmax joins two reductions with elementwise work. A stable softmax subtracts the row maximum before taking exponentials, divides by the sum of those exponentials, and keeps the row in on-chip memory for the whole computation so it is read from global memory only once. Fused Row-Wise Softmax in Triton asks for this kernel, and Row-Wise LogSumExp in Triton uses the same maximum trick to compute the log of the sum of exponentials without overflow.

Step 5: How do you write a tiled matrix multiplication in Triton?

Each program in a Triton matrix multiplication computes one tile of the output, loops over the shared dimension one block at a time, and adds the result of tl.dot on the two input tiles to an accumulator. The tile offsets are two-dimensional, built by broadcasting a column of row offsets against a row of column offsets, and each offset is multiplied by the stride of its dimension. Passing strides into the kernel lets the same code read a matrix whether or not it is transposed.

The accumulator is usually kept in float32 even when the inputs are float16 or bfloat16, and it is converted to the output type only at the store. Build up to the full kernel with Tiled Transpose in Triton, which practises two-dimensional offsets and masks, and GEMV: Matrix Vector Product in Triton, which practises a reduction along one axis of a tile. Then write Tiled Matrix Multiplication in Triton, where tiles at the edges of the matrix need masks on both dimensions.

Step 6: How does kernel fusion work in Triton?

Fusion means doing several operations inside one kernel so the intermediate results stay on the chip and never make a round trip through global memory. In Triton the simplest form is an epilogue, where you apply a bias and an activation to the matmul accumulator before storing it. Fused Matmul + Bias + ReLU in Triton asks for exactly that, and the CUDA learning guide explains why fusion is often the largest speed-up for memory-bound ML operations.

Normalization layers are the other common fused kernels. LayerNorm Forward in Triton and RMSNorm Forward in Triton use one program per row to compute the statistics and apply the scale in the same pass, and Cross Entropy Loss (Mean Reduction) in Triton combines the log-sum-exp from step 4 with a lookup of the target class. LayerNorm Backward in Triton is the hardest problem in this step, because the gradients for the weight and bias are sums across all rows, so programs that each handle one row have to combine their partial sums.

Step 7: How does autotuning work in Triton?

Autotuning lets Triton try several block sizes and launch settings for a kernel and keep the fastest one for each input shape. You decorate the kernel with triton.autotune, give it a list of triton.Config entries with block sizes, a number of warps and a number of pipeline stages, and name the arguments whose values should trigger a new search. The first call with a new shape benchmarks every config and caches the winner, and later calls with that shape reuse it.

Start with Autotuned Matrix Multiplication in Triton, then work through the performance patterns that production matmuls use.

• Grouped Program-ID Matmul in Triton reorders which output tiles run at the same time, so that programs running together reuse the same input tiles while they are still in the L2 cache.
• Split-K Matmul in Triton divides the shared dimension across several programs and adds their partial results, which helps when the output is small and the shared dimension is long.
• Block-Pointer Matmul in Triton uses tl.make_block_ptr and tl.advance to describe each tile by its shape, strides and position instead of building offsets by hand.
• Vectorized Vector Add in Triton returns to the first kernel and gives the compiler the alignment and contiguity hints it needs to issue wider memory loads.

Step 8: How do you write attention kernels in Triton?

Attention kernels combine everything before them, since scaled dot-product attention is a matrix multiplication, a row-wise softmax and a second matrix multiplication. Scaled Dot-Product Attention in Triton asks you to write that sequence directly, which is a good check that steps 4 and 5 are solid.

FlashAttention changes the order of work so the full table of attention scores is never written to global memory. Each program takes a block of queries, loops over blocks of keys and values, and keeps a running maximum, a running sum and a partial output, rescaling all three whenever a new block contains a larger score. Solve Online Softmax Kernel in Triton first to learn the rescaling on its own, then write FlashAttention Forward Kernel in Triton. The lesson on how FlashAttention works explains the memory argument in more detail.

Two smaller LLM kernels finish the step. Rotary Position Embedding in Triton rotates pairs of query and key features by position-dependent angles, and KV Cache Append in Triton writes new keys and values into a cache at the right positions. Both are memory-bound, so their speed depends almost entirely on how they read and write memory.

How do you check that a Triton kernel is correct and fast?

You check a Triton kernel by comparing its output with the PyTorch operation it replaces within a floating-point tolerance, and you time it only after it passes. On TensorTonic your kernel runs on real GPU hardware and its output is compared with a PyTorch reference implementation of the same operation. Triton submissions also report timings across input sizes on different GPU models, and each section of the study plan ends with a short quiz on the concepts behind the kernels.

• Compare against PyTorch with torch.testing.assert_close, which applies a tolerance suited to the data type.
• Test sizes that are not a multiple of the block size, along with very small inputs, because those cases exercise the masks.
• Time kernels with triton.testing.do_bench, which runs warm-up iterations and repeats the measurement for you.
• Ignore the first call when timing by hand, since Triton compiles the kernel the first time it sees a new combination of argument types and compile-time values.
• Benchmark against the PyTorch operation on the same GPU and at several sizes, because small inputs mostly measure launch overhead.

To judge whether a kernel is limited by memory bandwidth or by arithmetic, use the roofline reasoning from the profiling step of the CUDA learning guide, which applies to Triton kernels unchanged.

Should you learn Triton or CUDA first?

Learn Triton first if you mainly write Python and PyTorch and want faster ML code, and learn CUDA first if you want to see exactly how the GPU runs threads or you are preparing for GPU infrastructure interviews. The CUDA learning guide has a short decision table for this choice, and most people who keep writing kernels end up using both languages.

Solving the same operation in both languages shows what Triton handles for you. Tiled Matrix Multiplication in the CUDA plan makes you load shared-memory tiles and synchronise threads by hand, while the Triton version expresses the same tile loop with tl.load and tl.dot. The CUDA practice problems guide lists the CUDA problems in an order that lines up with the steps above.

What should you build after learning Triton?

After the Triton basics, the most useful next steps are building FlashAttention end to end, using your kernels to serve language models, and placing them inside a full LLM build. Each direction has a page to read and a project or study plan where you write the code.

The FlashAttention project computes attention in tiles with a running softmax in PyTorch, counts the memory traffic each version causes, and finishes with the same algorithm as a Triton kernel checked against PyTorch. The Build LLM from Scratch study plan has 62 PyTorch, NumPy, CUDA and Triton problems across 17 sections, including Tiled Triton Row Sum and Fused Tiled Matmul and ReLU.

What mistakes do people make when learning Triton?

• Leaving the mask off a load or store passes the tests on sizes that divide evenly into blocks and fails on every other size.
• Using the default other value in a reduction lets padding enter the result, so masked positions in a maximum must load negative infinity and in a sum must load zero.
• Choosing a block size that is not a power of two makes tl.arange fail to compile, so round the size up and mask the extra positions.
• Assuming a tensor is contiguous breaks the kernel on transposed or sliced inputs, so pass strides in or call contiguous before the launch.
• Accumulating a matmul in float16 loses precision over a long shared dimension, which a float32 accumulator avoids.
• Timing the first call measures compilation, which makes a kernel look far slower than it is.
• Forgetting to rescale the running sum and output in online softmax gives correct results on easy inputs and wrong ones when a later block has a larger score.
• Autotuning before the kernel is correct wastes time, because a faster config of a wrong kernel is still wrong.

Practice this

Then work through all 30 problems in the Triton Kernels study plan, or build the full kernel in the FlashAttention project.

Related guides

• How to learn CUDA: a path from C to fused kernels
• CUDA practice problems: from vector add to attention
• Inference engineering: what it is and how to learn it
• How to build an LLM from scratch, one component at a time