Triton Notes

Triton Notes

Triton is a programming language and compiler originally developed at OpenAI that makes custom GPU kernels easier to write. It makes efficient GPU code easier to express than in CUDA, especially for deep learning workloads.

In CUDA, you usually decide what each individual thread does and manually manage many low-level details. Triton provides a higher-level abstraction: you describe what one program does to a block, or tile, of data. Triton's compiler then maps those block operations onto GPU threads and generates the GPU code. We will see what this means later in this note.

It is also worth noting that Triton uses the same style of array programming found in PyTorch and NumPy. If that way of thinking is unfamiliar, these two introductions are useful refreshers to think in shapes:

This note is an introduction to Triton's abstractions. We will use two kernels—vector addition and matrix multiplication—throughout the note to build intuition.

#Overall Pattern of GPU Programming

At a high level, GPU programming has two sides regardless of the DSL:

  1. Host (CPU) code, which prepares the input tensors, decides how much parallel work to launch, and launches the kernel.
  2. Device (GPU) code, or the kernel, describes the computation that each GPU program performs.

So the overall flow looks like this:

  1. Create host and device memory for tensor data.
  2. Copy tensor data from host to device.
  3. Load tensor data on-chip.
  4. Execute the operation.
  5. Store the tensor back to device memory.

We will divide this into two steps and look at each:

  1. Launching a kernel.
  2. Defining a kernel.

#1. Launching a kernel

How can we run a kernel?

As in CUDA, the host (CPU) launches the kernel before any device-side work begins. A single launch creates many instances of the kernel, called programs in Triton.

We will look at this in three steps:

  1. Validation.
  2. Setting the grid size.
  3. Launching the kernel.

#Validations

Because kernels provide considerable freedom over variables and memory, it is helpful to fail fast and loud when something is wrong. We therefore validate shapes, strides, dtypes, and device-specific features, such as those introduced with Blackwell.

These assertions largely serve two purposes: ensuring correctness and enabling specialized fast paths. Many kernels branch into completely different implementations based on the results of these checks.

#Validation example

python
def vector_add(a: torch.Tensor, b: torch.Tensor):
    assert a.is_cuda and b.is_cuda
    assert a.device == b.device
    assert a.ndim == b.ndim == 1
    assert a.shape == b.shape
    assert a.dtype == b.dtype
    # ...

#Setting Grid Size

A launch grid is a 1D, 2D, or 3D tuple specifying how many program instances run along each axis. Inside the kernel, tl.program_id(axis) identifies the current program along that axis.

Poor grid mapping can reduce parallelism or cache locality, so it's important to tune it.

So how do we set the grid size? Most kernels use one of two launch styles:

#1. Non-persistent tiled launch (one program per output tile)

In a non-persistent tiled launch, we divide the output into tiles of BN elements and launch one program for each tile. Program 0 processes the first BN elements, program 1 processes the next BN elements, and so on. Therefore, the grid size is the number of tiles required to cover all N elements.

For example,

python
# ...
grid = lambda meta: (triton.cdiv(N, meta["BN"]),)
_vec_add_kernel[grid](
        # ...
    )

What is happening here?

triton.cdiv(a, b) computes ceil(a / b). This ensures that a final partial tile still receives a program. The kernel must mask any lanes that fall outside the logical tensor bounds.

When grid is a function, Triton passes it a dictionary containing the kernel's compile-time meta-parameters, such as BN. This lets the grid size depend on the selected block sizes, including values selected by autotuning.

#2. Persistent launch (multiple output tiles per program)

A persistent kernel launches a limited number of programs, often bounded by the number of output tiles and chosen relative to the GPU's SM count. Each program loops over tiles inside the kernel. See when one program processes multiple tiles for how this works in the kernel.

#What Parameters Do We Send?

#Input and output tensors

We pass the input tensors that the kernel reads and preallocated output tensors that it writes.

#Number of elements in each dimension

The number of elements in a tensor (e.g., 1,000) may not be divisible by its tile size (e.g., 32). The kernel receives the logical dimension and creates a mask such as offs_n < N, preventing lanes in a partial tile from accessing memory outside the vectors.

#Strides

Stride is the number of elements between adjacent positions along one tensor dimension. For vector addition, the stride is usually 1 because vectors are commonly contiguous, but a non-contiguous vector may have a different stride. We will talk more about strides in the two-dimensional case.

#Static arguments (tl.constexpr)

An argument annotated with tl.constexpr is known at compile time. Triton can use it to specialize generated code, size block-shaped tensors, or remove compile-time branches. Tile sizes are common autotuning targets, while tl.constexpr is also used for algorithm variants and feature flags.

#Example of launching a kernel

python
def vector_add(a: torch.Tensor, b: torch.Tensor):
    # ...
    _vec_add_kernel[grid](
        a, b, c, # a, b as input tensor, c as output tensor
        N,
        a.stride(0), b.stride(0), c.stride(0), # all 1
        BN=1024, # known at compile time; we can autotune it
    )

#End-to-end example: launching a vector add kernel

python
def vector_add(a: torch.Tensor, b: torch.Tensor):
    assert a.is_cuda and b.is_cuda
    assert a.device == b.device
    assert a.ndim == b.ndim == 1
    assert a.shape == b.shape
    assert a.dtype == b.dtype

    c = torch.empty_like(a)
    N = c.numel()

    grid = lambda meta: (triton.cdiv(N, meta["BN"]),)

    _vec_add_kernel[grid](
        a, b, c, # a, b as input tensor, c as output tensor
        N,
        a.stride(0), b.stride(0), c.stride(0), # all 1
        BN=1024, # known at compile time; we can autotune it
    )
    return c

#2. Defining a kernel

Until now, we have seen how the host launches a kernel. Before looking at how a Triton kernel is defined, let's first go over the two most important abstractions in Triton: tile and program.

#Central abstractions: Tile and Program

Consider adding two vectors: C = A + B.

If the vectors contain one million elements, a Triton kernel usually does not process all one million elements in a single program. Instead, we divide the output into smaller regions and launch multiple instances of the kernel.

TL;DR: A tile is a region of the data, while a program is an executing kernel instance that processes that region.

#Tile

A tile is a logical region of the overall problem that is processed together.

For example, one program might process a block of 1024 consecutive elements of the vector addition output C.

plaintext
tile 0: C[   0 : 1024]
tile 1: C[1024 : 2048]
tile 2: C[2048 : 3072]
...

A tile is represented by block-shaped offsets, pointers, and values inside the program.

This differs from CUDA, where the unit of code is a thread: values are scalars, and you index one element at a time.

#Tile vs Block

In Triton, tile and block are often used almost interchangeably. Both commonly refer to a group of tensor elements processed together by one program.

To be concrete,

  • Tile refers to the logical region of the overall tensor or problem.
  • Block refers to the block-shaped offsets, pointers, or values manipulated inside the program.

For example, a program may be responsible for the tile C[1024:2048]. Inside the program, that tile is represented by blocks of offsets and values containing up to 1024 elements.

In this note, we use them interchangeably.

#Program

A Triton program is one executing instance of a compiled kernel. It is roughly analogous to a CUDA thread block, also called a CTA.

The CPU kernel launch shown above creates as many program instances as specified by the grid. For example, if N = 4,096 and BN = 1024, four programs are launched.

plaintext
program 0 → tile 0 → C[   0 : 1024]
program 1 → tile 1 → C[1024 : 2048]
program 2 → tile 2 → C[2048 : 3072]
program 3 → tile 3 → C[3072 : 4096]

Each program receives a program ID (pid) and usually processes a different input or output tile. Triton launch grids have at most three dimensions, matching the grid defined by the host.

All four programs execute the same kernel code but operate on different data. In fancy terms, we call this SPMD.

Now that we've covered the two central abstractions, let's see how a kernel is defined and run on the GPU.

At a high level, each program follows these steps:

  1. Read the program ID and determine which output tile the program owns.
  2. Construct the tile’s global element offsets.
  3. Use those offsets to form input and output pointers.
  4. Load the input blocks.
  5. Compute the output block.
  6. Store the output block.

#Typical Kernel Parameters

Triton kernels receive runtime arguments and configuration metadata from the host code described in the previous section, including:

  • Input and output pointers: PyTorch tensors passed at launch become pointers to their GPU storage.
  • Logical dimensions: sizes used for iteration and masks.
  • Strides: element offsets used to map multidimensional indices to memory locations.
  • Runtime scalar values: optional values such as a dropout probability or scaling factor.
  • Compile-time meta-parameters: tl.constexpr values such as tile sizes or feature flags.

#Example

python
@triton.jit
def _vec_add_kernel(
    a_ptr, b_ptr, c_ptr,
    N,
    stride_a, stride_b, stride_c,
    BN: tl.constexpr,
):
    # ...

#1. Read the program ID and determine which output tile the program owns

#Program ID

A program ID, or pid, identifies a program along one launch-grid axis through tl.program_id(axis). The kernel uses it to determine which tile that program processes.

#Example

python
@triton.jit
def _vec_add_kernel(
    # ...
):
    pid = tl.program_id(axis=0)

Here we use axis=0 because the host launched a one-dimensional grid.

#2. Construct the tile’s global element offsets

Local offsets are indices within the tile, ranging from 0 to tile_size - 1. We simply use tl.arange (just like np.arange or torch.arange) to create the array.

python
local_offs_n = tl.arange(0, BN)

To locate the tile within the full tensor, we incorporate the program ID. In practice, kernels commonly construct global offsets directly:

python
offs_n = pid * BN + tl.arange(0, BN)

What does this mean? Since each program processes a tile of size BN, we multiply it by pid to locate that program's block. So offs_n ranges from pid * BN to (pid + 1) * BN - 1.

For example, if BN = 1024 and pid = 3, offs_n is:

plaintext
3 * 1024 + tl.arange(0, 1024) == [3072, 3073, 3074, 3075, ..., 4095]

These lines construct block-shaped tensors of output coordinates. The coordinates are later combined with base address and strides to construct pointers for A, B, and C.

#Masking

As explained, masks allow a regular block-shaped program to safely handle the final partial tile. A mask is usually applied to the global offsets.

#Example

python
@triton.jit
def _vec_add_kernel(
    # ...
):
    # ...
    offs_n = pid * BN + tl.arange(0, BN)
    n_mask = offs_n < N

This way, even if N is not a multiple of BN, the final partial tile can be masked correctly.

#3. Use those offsets to form input and output pointers

Once we have the global offsets, we map the logical indices to memory pointers. These are the final memory addresses we will use to load the tensor elements.

python
@triton.jit
def _vec_add_kernel(
    a_ptr, b_ptr, c_ptr,
    # ...
):
    # ...
    a_ptrs = a_ptr + offs_n * stride_a
    b_ptrs = b_ptr + offs_n * stride_b
    c_ptrs = c_ptr + offs_n * stride_c

#4. Load the input blocks

Input tensors normally already reside in GPU global memory before the kernel begins. A Triton kernel reads values with tl.load, computes on them using on-chip resources (typically registers), and writes results back with tl.store. The compiler and hardware manage lower-level movement through the cache hierarchy.

We use the mask we've created here as well to mask off elements outside the bounds.

The other argument supplies a value for masked-out loads. We usually use 0.0 or -float("inf") depending on the context. For example, vector addition can simply use 0.0.

Example

python
@triton.jit
def _vec_add_kernel(
    a_ptr, b_ptr,
    # ...
):
    # ...
    a_ptrs = a_ptr + offs_n * stride_a
    b_ptrs = b_ptr + offs_n * stride_b
    a = tl.load(a_ptrs, mask=n_mask, other=0.0)
    b = tl.load(b_ptrs, mask=n_mask, other=0.0)

#5. Compute the output block

In Triton, computation is pretty high-level. More complex kernels may require additional operations, but many computations are expressed using basic arithmetic or matrix multiplication (tl.dot). Vector addition is dead simple:

python
@triton.jit
def _vec_add_kernel(
    a_ptr, b_ptr,
    # ...
):
    # ...
    a_ptrs = a_ptr + offs_n * stride_a
    b_ptrs = b_ptr + offs_n * stride_b
    a = tl.load(a_ptrs, mask=n_mask, other=0.0)
    b = tl.load(b_ptrs, mask=n_mask, other=0.0)
    c = a + b

#6. Store the output block

After we've loaded and computed the values, we need to store them back in device memory. We need the pointers that specify where to store them, the value to store (c), and the mask.

python
@triton.jit
def _vec_add_kernel(
    a_ptr, b_ptr, c_ptr,
    # ...
):
    # ...
    c = a + b
    c_ptrs = c_ptr + offs_n * stride_c
    tl.store(c_ptrs, c, mask=n_mask)

#End-to-end example of a vector addition kernel

python
@triton.jit
def _vec_add_kernel(
    a_ptr, b_ptr, c_ptr,
    N,
    stride_a, stride_b, stride_c,
    BN: tl.constexpr
):
    pid = tl.program_id(axis=0)
    offs_n = pid * BN + tl.arange(0, BN)
    n_mask = offs_n < N
    a_ptrs = a_ptr + offs_n * stride_a
    b_ptrs = b_ptr + offs_n * stride_b

    a = tl.load(a_ptrs, mask=n_mask, other=0.0)
    b = tl.load(b_ptrs, mask=n_mask, other=0.0)
    c = a + b

    c_ptrs = c_ptr + offs_n * stride_c
    tl.store(c_ptrs, c, mask=n_mask)

#2D and beyond

We've seen the most basic kernel as an exercise. Now let's move one step forward and deal with two-dimensional scenarios.

Many operations in deep learning, including matmul, use tiles with two or more dimensions. We will go over the same path again while introducing additional things to consider when dealing with multi-dimensional arrays. For the parts already explained, I won't repeat myself.

#Launching a matmul kernel

As we've seen above, we validate, set grid size, and launch the kernel:

python
def matmul(a: torch.Tensor, b: torch.Tensor):
    assert a.is_cuda and b.is_cuda
    assert a.device == b.device
    assert a.ndim == 2 and b.ndim == 2
    assert a.shape[1] == b.shape[0] # K dimensions must match
    assert a.dtype == b.dtype == torch.float16

    M, K = a.shape
    _, N = b.shape
    c = torch.empty((M, N), device=a.device, dtype=a.dtype)

    grid = lambda meta: (
        triton.cdiv(M, meta["BM"]),
        triton.cdiv(N, meta["BN"]),
    )

    _matmul_kernel[grid](
        a, b, c,
        M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        BM=32, BN=32, BK=32,
    )
    return c

#Strides

Notice we pass two strides for each tensor because A, B, and C are all two-dimensional: one stride for moving between rows and another for moving between columns.

Consider a row-major matrix A with shape [M, K]. Moving one column advances by one element, while moving one row skips over K elements. Therefore, when A is contiguous, its strides are (K, 1).

Most of the time, tensors sent to kernels are created contiguously. However, a transposed tensor or another view may have a different memory layout. Passing a.stride(0) and a.stride(1) allows the kernel to use the correct row and column jumps in either case. See the diagram below.

#Defining a matmul kernel

As usual, we first receive all the parameters sent by the host launch.

python
@triton.jit
def _matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_a_m, stride_a_k,
    stride_b_k, stride_b_n,
    stride_c_m, stride_c_n,
    BM: tl.constexpr,
    BN: tl.constexpr,
    BK: tl.constexpr,
):

#1. Read the program IDs and determine which output tile the program owns

Since we launched a two-dimensional grid, program instances are also coordinated along two dimensions.

python
    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)
#Optional: When one program processes multiple tiles

The matmul kernel above uses a straightforward mapping: the host launches one program for every output tile, and pid_m and pid_n identify that tile directly.

Another option is a persistent-style launch, where the host launches fewer programs than there are output tiles. Each program then processes several tiles in a loop. This is where tl.num_programs(axis) becomes useful: it returns the number of programs launched along an axis, which determines how far each program advances through the tile list.

For a 1D launch grid, we can flatten the two-dimensional output-tile grid into a single sequence:

python
start_tile = tl.program_id(axis=0)
num_programs = tl.num_programs(axis=0)

num_pid_m = tl.cdiv(M, BM)
num_pid_n = tl.cdiv(N, BN)
num_tiles = num_pid_m * num_pid_n

for tile_id in tl.range(start_tile, num_tiles, num_programs):
    pid_m = tile_id // num_pid_n
    pid_n = tile_id % num_pid_n

    # Compute the output tile C[pid_m, pid_n].
    # ...

Program 0 starts with tile 0, program 1 starts with tile 1, and so on. After processing its first tile, each program advances by num_programs.

For example, if four programs are launched for sixteen output tiles, program 0 processes tiles 0, 4, 8, and 12. The diagram below illustrates this assignment.

The host must explicitly choose this smaller launch grid. tl.num_programs() only reports how many programs were launched; it does not make the kernel persistent by itself.

#2. Construct the tile’s global element offsets

Next, get the global offsets for each dimension:

python
    offs_m = pid_m * BM + tl.arange(0, BM)
    offs_n = pid_n * BN + tl.arange(0, BN)
    offs_k = tl.arange(0, BK)

#3. Use those offsets to form input and output pointers

Then form the addresses by adding those offsets to the pointers:

python
    a_ptrs = (
        a_ptr
        + offs_m[:, None] * stride_a_m
        + offs_k[None, :] * stride_a_k
    )
    b_ptrs = (
        b_ptr
        + offs_k[:, None] * stride_b_k
        + offs_n[None, :] * stride_b_n
    )

What do [:, None] and [None, :] mean? To understand this, we should brush up on how broadcasting works in array programming. Feel free to skip this section if you're already fluent with the notation!

#Broadcasting

Broadcasting allows tensors with different shapes to participate in the same element-wise operation. When a dimension of size 1 is matched with a larger dimension, its values behave as though they were copied across that dimension. Note that this expansion is logical—no larger tensor needs to be allocated or materialized in memory. NumPy, PyTorch, and Triton compare shapes from the rightmost dimension to the left.

Two dimensions are compatible when, comparing from the rightmost dimension:

  1. They are equal;
  2. One of them is 1; or
  3. One shape has fewer dimensions, in which case the missing leading dimensions are treated as 1.

Based on this rule, let's go through some examples:

For A of shape [3, 4]:

  • [3, 4] is compatible because the shapes are equal.
  • [1, 4] is compatible because the rightmost dimensions are equal and the leading 1 expands to 3.
  • [1, 1] is compatible because each dimension of size 1 expands to match [3, 4].
  • [1, 2] is incompatible because the rightmost dimensions, 4 and 2, are neither equal nor 1.

In a matmul kernel, broadcasting turns the M and K coordinate vectors into an M-by-K pointer tile for A:

text
offs_m[:, None] : (BM, 1)
offs_k[None, :] : (1, BK)
combined        : (BM, BK)

The same idea constructs the K-by-N pointer tile for B:

plaintext
offs_k[:, None] : (BK, 1)
offs_n[None, :] : (1, BN)
combined        : (BK, BN)

Since [B0, 1] and [1, B1] are compatible, Triton broadcasts both into [B0, B1]. No expanded index arrays need to be materialized in global memory.

#Expanding dimensions for broadcasting

In Triton, you can explicitly use tl.expand_dims(offs_m, axis) or tl.unsqueeze(offs_m, axis). However, we usually use : and None to express this:

python
offs_m[:, None]  # column shape (BM, 1)

This takes all M coordinates and inserts a new trailing dimension, producing one column.

python
offs_k[None, :]  # row shape (1, BK)

This takes all K coordinates and inserts a new leading dimension, producing one row.

#Strided Pointers

An intuitive way to understand multidimensional pointer arithmetic is to consider one axis at a time. The indices along an axis tell us how many steps to take, while its stride tells us how many memory elements each step moves. Multiplying them gives that axis’s contribution to the final memory offset. Broadcasting then combines the contributions from every axis to produce a pointer for each coordinate in the tile.

For example,

python
a_ptrs = (
    a_ptr
    + offs_m[:, None] * stride_a_m
    + offs_k[None, :] * stride_a_k
)

offs_m[:, None] * stride_a_m computes the row contribution and broadcasts it across every column. Meanwhile, offs_k[None, :] * stride_a_k computes the column contribution and broadcasts it across every row. Adding them to a_ptr produces the complete M-by-K grid of pointers.

#4. Load the input blocks

At this point, the program owns a BM × BN tile of the output. Computing that tile requires accumulation across the entire K dimension. Rather than loading all of K at once, the program processes it in chunks of BK elements while keeping the partial result in accumulator. accumulator is set to tl.float32 for numerical precision and later cast to tl.float16 after the loop.

python
    accumulator = tl.zeros((BM, BN), dtype=tl.float32)
    for k_tile in range(0, tl.cdiv(K, BK)):
        k = k_tile * BK + offs_k
        a = tl.load(
            a_ptrs,
            mask=(offs_m[:, None] < M) & (k[None, :] < K),
            other=0.0,
        )
        b = tl.load(
            b_ptrs,
            mask=(k[:, None] < K) & (offs_n[None, :] < N),
            other=0.0,
        )

Each iteration loads a BM × BK tile from A and a BK × BN tile from B. Their masks are constructed using the same broadcasting pattern as their pointers.

#5. Compute the output block

python
        accumulator = tl.dot(a, b, acc=accumulator)
        a_ptrs += BK * stride_a_k
        b_ptrs += BK * stride_b_k

tl.dot adds the current K chunk’s contribution to accumulator, after which the pointers advance to the next K chunk.

#6. Store the output block

python
    c = accumulator.to(tl.float16)
    c_ptrs = (
        c_ptr
        + offs_m[:, None] * stride_c_m
        + offs_n[None, :] * stride_c_n
    )
    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)

With these pieces in place, let's put the matmul host launch and Triton kernel together.

The host-side kernel launch looks like this:

python
def matmul(a: torch.Tensor, b: torch.Tensor):
    assert a.is_cuda and b.is_cuda
    assert a.device == b.device
    assert a.ndim == 2 and b.ndim == 2
    assert a.shape[1] == b.shape[0]
    assert a.dtype == b.dtype == torch.float16

    M, K = a.shape
    _, N = b.shape
    c = torch.empty((M, N), device=a.device, dtype=a.dtype)

    grid = lambda meta: (
        triton.cdiv(M, meta["BM"]),
        triton.cdiv(N, meta["BN"]),
    )
    _matmul_kernel[grid](
        a, b, c,
        M, N, K,
        a.stride(0), a.stride(1),
        b.stride(0), b.stride(1),
        c.stride(0), c.stride(1),
        BM=32, BN=32, BK=32,
    )
    return c

The matmul kernel (without L2 cache optimizations) looks like:

python
@triton.jit
def _matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_a_m, stride_a_k,
    stride_b_k, stride_b_n,
    stride_c_m, stride_c_n,
    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
):
    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)

    offs_m = pid_m * BM + tl.arange(0, BM)
    offs_n = pid_n * BN + tl.arange(0, BN)
    offs_k = tl.arange(0, BK)

    a_ptrs = (
        a_ptr
        + offs_m[:, None] * stride_a_m
        + offs_k[None, :] * stride_a_k
    )
    b_ptrs = (
        b_ptr
        + offs_k[:, None] * stride_b_k
        + offs_n[None, :] * stride_b_n
    )

    accumulator = tl.zeros(
        (BM, BN),
        dtype=tl.float32,
    )
    for k_tile in range(0, tl.cdiv(K, BK)):
        k = k_tile * BK + offs_k
        a = tl.load(
            a_ptrs,
            mask=(offs_m[:, None] < M) & (k[None, :] < K),
            other=0.0,
        )
        b = tl.load(
            b_ptrs,
            mask=(k[:, None] < K) & (offs_n[None, :] < N),
            other=0.0,
        )
        accumulator = tl.dot(a, b, acc=accumulator)
        a_ptrs += BK * stride_a_k
        b_ptrs += BK * stride_b_k

    c = accumulator.to(tl.float16)
    c_ptrs = (
        c_ptr
        + offs_m[:, None] * stride_c_m
        + offs_n[None, :] * stride_c_n
    )
    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)

That's it!

#Selecting Configurations

While Triton raises the level of abstraction so you don't need to deal with messy hardware, configuration selection remains hardware-adjacent. This is because performance still depends on choosing tile and launch parameters that fit the workload and target GPU. Common parameters include BM, BN, BK, num_warps, and num_stages.

There are two common ways to select configurations:

#1. autotune

The first uses @triton.autotune with a list of triton.Config objects.

Each configuration can set compile-time keyword arguments such as block sizes, along with launch options such as num_warps, num_stages, and num_ctas. See the triton.Config documentation for the full list.

For example, a kernel can use autotuning like this:

python
autotune_configs = [
    triton.Config({"BM": 128, "BN": 256, \
    "BK": 64}, num_stages=3, num_warps=8),
    triton.Config({"BM": 64, "BN": 256, \
    "BK": 32}, num_stages=4, num_warps=4),
    # ...
]

@triton.autotune(configs=autotune_configs, key=["M", "N", "K"])
@triton.jit
def _matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_a_m, stride_a_k,
    stride_b_k, stride_b_n,
    stride_c_m, stride_c_n,
    # meta-parameters
    BM: tl.constexpr, BN: tl.constexpr, \
    BK: tl.constexpr,
):
    ...

Exhaustively testing every possible configuration can take a long time. In practice, we provide a bounded set of plausible candidates. For larger search spaces, Triton can prune candidates through prune_configs_by, while @triton.heuristics or manual lookup tables can avoid runtime tuning when compilation and cold-start latency matter.

The key parameter in @triton.autotune() defines the runtime values that identify a tuning problem. Whenever one or more of those key values change, Triton evaluates the candidate configurations for the new key and caches the result.

#2. Manual Lookup Tables or Heuristics

Although autotuning flexibly covers different shapes and dtypes, heuristics can sometimes select configurations with less runtime overhead.

Manual lookup tables or heuristics can be keyed on shape, dtype, device, and feature flags. They are useful in latency-sensitive inference serving, where runtime autotuning and compilation overhead matter.

#A few more Triton operations

#tl.where

If the matmul fuses an activation such as ReLU, tl.where can be used as a per-element if-else on the output tile:

python
c = tl.where(accumulator > 0.0, accumulator, 0.0)

Both value branches of tl.where are evaluated, as noted in the API documentation. Use masks on tl.load and tl.store, rather than tl.where, when a branch would otherwise perform an unsafe memory access.

#Axis = 0 and Axis = 1

The axis argument identifies the dimension to reduce, and that dimension is removed from the result. For a two-dimensional tensor shaped [B0, B1], reducing along axis=0 produces [B1], while reducing along axis=1 produces [B0].

Example

plaintext
k = [
  [1, 2, 3, 4],
  [10, 20, 30, 40],
]
python
tl.sum(k, axis=1)

sums each row across columns:

plaintext
[1+2+3+4, 10+20+30+40]
= [10, 100]

So it turns [B0, B1] into [B0].

#Testing

Correctness comes before performance. Before benchmarking a Triton kernel, test it against a trusted reference implementation, typically PyTorch.

Try testing with:

  • A representative shape.
  • A shape that is not divisible by the tile size, to verify the tail masks.
  • Every dtype and memory layout that the kernel claims to support.

#Example

For example, a matmul test should use M, N, and K independently instead of testing only square matrices.

python
def test_matmul(M, N, K):
    torch.manual_seed(0)
    a = torch.randn((M, K), device=DEVICE, dtype=torch.float16)
    b = torch.randn((K, N), device=DEVICE, dtype=torch.float16)

    actual = matmul(a, b)
    expected = torch.matmul(a, b)

    torch.testing.assert_close(
        actual,
        expected,
        atol=1e-2,
        rtol=1e-2,
    )

if __name__ == "__main__":
    for shape in [
        (128, 128, 128),  # regular tiles
        (1, 17, 33),      # small problem
        (127, 65, 129),   # partial tiles on every axis
    ]:
        test_matmul(*shape)

torch.testing.assert_close() checks both absolute and relative error:

  • atol is the allowed absolute difference, which matters for values near zero.
  • rtol scales the allowed difference relative to the expected value.

We keep both atol and rtol because Triton and PyTorch may evaluate operations in a different order or at different intermediate precisions, producing small numerical differences.

#Performance Benchmarking

Once the kernel passes its tests, we can compare it against a meaningful baseline using the same shapes, dtypes, and device.

GPU work is asynchronous, so timing a launch with a normal CPU timer can measure only the time required to enqueue the work. triton.testing.do_bench performs warmup and repeated GPU timing and returns the runtime in milliseconds.

Example

Here is a compact benchmark for matrix multiplication:

python
def benchmark_matmul(M, N, K):
    a = torch.randn((M, K), device=DEVICE, dtype=torch.float16)
    b = torch.randn((K, N), device=DEVICE, dtype=torch.float16)

    # Check correctness and trigger compilation before reporting performance.
    torch.testing.assert_close(
        matmul(a, b),
        torch.matmul(a, b),
        atol=1e-2,
        rtol=1e-2,
    )

    triton_ms = triton.testing.do_bench(
        lambda: matmul(a, b),
        return_mode="median",
    )
    torch_ms = triton.testing.do_bench(
        lambda: torch.matmul(a, b),
        return_mode="median",
    )

    # A matrix multiplication performs approximately 2 * M * N * K FLOPs.
    flops = 2 * M * N * K
    to_tflops = lambda ms: flops / (ms * 1e-3) / 1e12

    print(
        f"shape=({M}, {N}, {K}) | "
        f"Triton={to_tflops(triton_ms):6.1f} TFLOP/s | "
        f"PyTorch={to_tflops(torch_ms):6.1f} TFLOP/s"
    )

if __name__ == "__main__":
    for shape in [
        (256, 256, 256),
        (512, 512, 512),
        (1024, 1024, 1024),
        (1025, 769, 513),  # irregular dimensions
    ]:
        benchmark_matmul(*shape)

Latency (ms or µs) is the clearest metric for small operations. Throughput (TFLOP/s) reports how quickly the matrix multiplication's floating-point work completes.

For a larger shape sweep with tables and plots, Triton also provides triton.testing.Benchmark and triton.testing.perf_report.

#Debugging

python
import os
os.environ["TRITON_INTERPRET"] = "1"

Interpreter mode simulates Triton programs sequentially on the CPU using NumPy equivalents, which makes intermediate values inspectable with Python print() or pdb.

Because programs execute sequentially on the CPU, interpreter mode does not reproduce GPU parallelism, scheduling, or performance and should not be used for benchmarking. It also has unsupported cases, including operations on bfloat16 values and indirect memory-access patterns. Use it to inspect correctness, then test the kernel on the target GPU as well.

#Wrapping Up

This was a friendly introduction to Triton. Before finishing this note, let's review the core ideas:

  1. A Triton program is one instance of a compiled kernel. It owns one tile, or sometimes a sequence of tiles, loads the corresponding values, computes on them, and stores its outputs.
  2. grid defines how many programs launch along each axis. For example, (1, 1, 1) launches one program, while (3, 2, 1) launches six. In the latter grid, tl.program_id(0) ranges from 0 to 2, and tl.program_id(1) ranges from 0 to 1.
  3. Most non-persistent kernels map one program to one logical output tile. Persistent kernels may assign several tiles to each program. See when one program processes multiple tiles.
  4. Each program uses its program ID to identify the tile it contributes to the output tensor.
  5. Before loading data, a program determines which logical tensor elements it owns and maps those coordinates to memory pointers.
  6. As a practical rule of thumb, expressions derived from program IDs, block coordinates, tl.arange(), and strides usually calculate indices, memory offsets, or pointers. The actual tensor values are obtained by reading from those locations with tl.load(), and computations on the data typically follow.