Highly optimized Triton snippets

Fusion and tile-shape choice are the two levers that make a Triton kernel fast. Everything else — autotuning, swizzling, pipelining — is refinement around a tile you have already justified.

Why fusion is Triton's highest-return use case

Chaining PyTorch operators is correct but wasteful. Each operator is its own CUDA kernel launch, and each one reads its inputs from HBM and writes its output back to HBM. Consider y = gelu(x + bias) written as two ops: x + bias reads x, reads bias, and writes a full-size temporary t to HBM; then gelu(t) reads t back and writes y. The intermediate t exists only to be round-tripped through memory.

The cost that matters here is arithmetic intensity: FLOPs performed per byte moved to and from HBM. Elementwise ops have tiny intensity — a GELU is a handful of FLOPs per element, but each element costs a load and a store (8 bytes round trip in fp32). These kernels are memory-bound: the ALUs sit idle waiting on HBM. The only way to speed them up is to move fewer bytes.

Fusion does exactly that. One kernel reads each input once, does all the arithmetic while the values are live in registers, and writes the result once. The intermediate never touches HBM, and you pay one launch instead of several.

Snippet 1 — fused bias + GELU

The classic transformer epilogue: add a per-column bias, then apply GELU. We fuse the add and the activation into one elementwise pass. The bias is indexed by the column of each element, so from a flat offset offs over a row-major (rows, D) tensor we recover the column with d = offs % D and gather bias[d].

import triton
import triton.language as tl

@triton.jit
def gelu(x):
    # tanh approximation of GELU (matches torch's approximate="tanh")
    # 0.5 * x * (1 + tanh( sqrt(2/pi) * (x + 0.044715 * x**3) ))
    c = 0.7978845608028654          # sqrt(2/pi)
    inner = c * (x + 0.044715 * x * x * x)
    return 0.5 * x * (1.0 + tl.tanh(inner))

@triton.jit
def bias_gelu_kernel(x_ptr, b_ptr, y_ptr,
                     n_elements, D,
                     BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n_elements

    x = tl.load(x_ptr + offs, mask=mask, other=0.0)
    d = offs % D                              # column index of each element
    b = tl.load(b_ptr + d, mask=mask, other=0.0)

    y = gelu(x + b)                           # add + activation, all in registers
    tl.store(y_ptr + offs, y, mask=mask)


def bias_gelu(x, bias):
    # x: (rows, D), bias: (D,)
    y = torch.empty_like(x)
    n = x.numel()
    D = x.shape[-1]
    grid = lambda meta: (triton.cdiv(n, meta["BLOCK"]),)
    bias_gelu_kernel[grid](x, bias, y, n, D, BLOCK=1024)
    return y

Snippet 2 — numerically stable row softmax

Softmax over the last axis. We assign one program per row, load the whole row into registers, and do the reduction there. The naive exp(x) / sum(exp(x)) overflows: a single large logit makes exp inf, and inf/inf is NaN. The fix is to subtract the row max first — mathematically a no-op (it cancels in the ratio) but it bounds every exponent at exp(0) = 1.

@triton.jit
def softmax_kernel(x_ptr, y_ptr,
                   row_stride, n_cols,
                   BLOCK: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    mask = cols < n_cols
    ptrs = x_ptr + row * row_stride + cols

    # other=-inf so masked lanes never win the max and contribute 0 to the sum
    vals = tl.load(ptrs, mask=mask, other=-float("inf"))

    m   = tl.max(vals, axis=0)                # row max, for stability
    num = tl.exp(vals - m)                    # in [0, 1], no overflow
    den = tl.sum(num, axis=0)
    out = num / den

    tl.store(y_ptr + row * row_stride + cols, out, mask=mask)

The masked-load fill is doing double duty here, and getting it right is the whole game on the edge tile (see Correctness under masking below). Filling with -inf means a padding lane can never be the max, and its exp(-inf - m) = 0 contributes nothing to the denominator. Fill with 0 instead and the padding lanes inject exp(0 - m) into the sum, silently corrupting the result.

Snippet 3 — fused RMSNorm

RMSNorm normalizes each row by its root-mean-square and scales by a learned per-feature weight: y = x * rsqrt(mean(x^2) + eps) * w. Like softmax it is a per-row reduction, so again one program per row. The subtlety is precision.

@triton.jit
def rmsnorm_kernel(x_ptr, w_ptr, y_ptr,
                   row_stride, D,
                   eps,
                   BLOCK: tl.constexpr):
    row  = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    mask = cols < D
    ptrs = x_ptr + row * row_stride + cols

    # cast to fp32 for the reduction; other=0 is sum-safe
    x  = tl.load(ptrs, mask=mask, other=0.0).to(tl.float32)
    ss = tl.sum(x * x, axis=0)                # sum of squares, fp32
    inv_rms = tl.rsqrt(ss / D + eps)          # 1 / sqrt(mean(x^2) + eps)

    w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    y = x * inv_rms * w

    tl.store(y_ptr + row * row_stride + cols, y.to(tl.float32), mask=mask)

Snippet 4 — tiled matmul with tl.dot

The canonical compute-bound kernel. Each program owns one BM×BN output tile of C = A @ B. It walks the shared K dimension in chunks of BK, loading an A tile and a B tile each step and accumulating their product into a register-resident accumulator.

@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr,
                  M, N, K,
                  stride_am, stride_ak,
                  stride_bk, stride_bn,
                  stride_cm, stride_cn,
                  BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(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)

    # 2D pointer grids into A (BM x BK) and B (BK x BN)
    a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn)

    acc = tl.zeros((BM, BN), dtype=tl.float32)         # lives in registers

    for k in range(0, K, BK):
        a = tl.load(a_ptrs,
                    mask=(offs_m[:, None] < M) & (offs_k[None, :] < (K - k)),
                    other=0.0)
        b = tl.load(b_ptrs,
                    mask=(offs_k[:, None] < (K - k)) & (offs_n[None, :] < N),
                    other=0.0)
        acc += tl.dot(a, b)                             # fp16/bf16 in, fp32 acc
        a_ptrs += BK * stride_ak
        b_ptrs += BK * stride_bk

    c = acc.to(c_ptr.dtype.element_ty)
    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)

Map this onto the four tile properties from the opening:

• What stays live: the acc tile (BM×BN fp32 values) lives in registers for the entire K loop. This is the dominant register cost and the main occupancy constraint.
• What streams: the A and B tiles flow through the loop one BK-chunk at a time. They are loaded, multiplied, and discarded.
• What is reused: each loaded A tile (BM×BK) participates in BN output columns, and each B tile in BM rows. That reuse is exactly why a big output tile raises arithmetic intensity.
• What is reduced: the K axis — which is why it is the loop, not the grid.

Snippet 5 — autotuning the tile

The algorithm above is fixed; what is not fixed is the tile shape and the launch parameters. @triton.autotune benchmarks a list of candidate configs and caches the winner per distinct key.

@triton.autotune(
    configs=[
        triton.Config({"BM": 128, "BN": 256, "BK": 64}, num_warps=8, num_stages=3),
        triton.Config({"BM": 256, "BN": 128, "BK": 64}, num_warps=8, num_stages=3),
        triton.Config({"BM": 128, "BN": 128, "BK": 32}, num_warps=4, num_stages=4),
        triton.Config({"BM": 64,  "BN": 64,  "BK": 32}, num_warps=2, num_stages=5),
    ],
    key=["M", "N", "K"],
)
@triton.jit
def matmul_kernel(...):   # same body as Snippet 4
    ...

The knobs trade off against each other:

Snippet 6 — grouped / swizzled program ordering for L2 reuse

With a 2D grid, the default program order sweeps a full row of output tiles before moving to the next row. Along that row every program loads a fresh B column tile but the same A row tile — so the B tiles a program needs were loaded by far-away programs long ago and have been evicted from L2. We can keep the math identical but change the visitation order so that programs running close together in time touch overlapping B tiles, which then stay hot in L2.

@triton.jit
def matmul_grouped(a_ptr, b_ptr, c_ptr, M, N, K, ...,
                   BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
                   GROUP_M: tl.constexpr):
    pid = tl.program_id(0)                    # 1D launch this time
    num_pid_m = tl.cdiv(M, BM)
    num_pid_n = tl.cdiv(N, BN)

    # remap a flat pid into (pid_m, pid_n) grouped along M
    num_pid_in_group = GROUP_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m
    # ... identical tile body using pid_m, pid_n ...

The software pipeline and num_stages

Inside the K loop there are two phases: load the next A/B tiles from HBM, and compute the tl.dot. Done naively they serialize — the SM stalls on the load before it can multiply. The Triton compiler can instead prefetch: issue the load for K-tile i+1 while the tensor cores are still chewing on K-tile i. That overlap hides HBM latency behind useful compute. num_stages is how many K-tiles deep this prefetch pipeline runs.

Correctness under masking

When a tensor dimension is not a multiple of the tile, the boundary tile reads and writes out-of-range addresses unless masked. Two rules:

1. Mask every load and every store. An unmasked load past the end reads garbage (or faults); an unmasked store past the end corrupts neighboring memory. The mask on the store must match the valid region of the output tile.
2. The other fill must be reduction-safe. The value masked-out lanes take has to be the identity of whatever reduction consumes it: 0.0 for a sum (and for sum-of-squares), -inf for a max, +inf for a min. A wrong fill does not crash — it quietly poisons only the edge tiles, so the kernel looks correct on power-of-two shapes and fails on the others.

Pitfalls

Snippet 7 — benchmarking correctly

A kernel is only "optimized" relative to a measured baseline, on the right unit. Use triton.testing.do_bench, which warms up the GPU and the autotune cache, then times many reps and returns the median milliseconds.

import torch, triton

def bench_matmul(M, N, K, dtype=torch.float16):
    a = torch.randn((M, K), device="cuda", dtype=dtype)
    b = torch.randn((K, N), device="cuda", dtype=dtype)

    ms_triton = triton.testing.do_bench(lambda: matmul(a, b),
                                        warmup=25, rep=100)
    ms_torch  = triton.testing.do_bench(lambda: torch.matmul(a, b),
                                        warmup=25, rep=100)

    # matmul is compute-bound: count FLOPs = 2*M*N*K (one mul + one add per term)
    flops = 2 * M * N * K
    tflops_triton = flops / (ms_triton * 1e-3) / 1e12
    tflops_torch  = flops / (ms_torch  * 1e-3) / 1e12
    print(f"triton {tflops_triton:7.1f} TFLOP/s   torch {tflops_torch:7.1f} TFLOP/s")


def bench_rmsnorm(rows, D, dtype=torch.float16):
    x = torch.randn((rows, D), device="cuda", dtype=dtype)
    w = torch.randn((D,),      device="cuda", dtype=dtype)
    ms = triton.testing.do_bench(lambda: rmsnorm(x, w), warmup=25, rep=100)

    # memory-bound: count bytes moved = read x + read w + write y
    bytes_moved = x.numel() * x.element_size() \
                + w.numel() * w.element_size() \
                + x.numel() * x.element_size()
    gbps = bytes_moved / (ms * 1e-3) / 1e9
    print(f"rmsnorm {gbps:7.1f} GB/s")

Optimization answer template

When asked to make a Triton kernel fast (in an interview or in real life), walk this in order:

1. Classify the op. Elementwise, reduction, dot (matmul/attention), or a fused epilogue on top of one of those? This decides whether you are memory-bound or compute-bound, and therefore which lever matters.
2. Choose tile ownership. Decide what one program is responsible for — a block of flat elements (elementwise), one row (per-row reduction like softmax/RMSNorm), or one BM×BN output tile (matmul). This follows directly from which axis is reduced.
3. Pick tile sizes from reuse. Reason about what stays live in registers, what streams from HBM, and what is reused, then size the tile big enough to amortize launches and feed tensor cores but small enough to stay in the register/shared budget.
4. Get correctness with masks and strides. Mask every load and store, choose reduction-safe other fills, accumulate reductions in fp32, and use real strides (or .contiguous()) so non-contiguous inputs are read correctly.
5. Autotune nearby configs. Fix the algorithm first, then search a small neighborhood of BM/BN/BK, num_warps, num_stages with @triton.autotune keyed on the dimensions that change the optimum.
6. Measure the right unit. do_bench with warmup, GB/s for memory-bound and TFLOP/s for compute-bound, compared against a same-dtype torch baseline and the hardware roofline.