Writing Custom GPU Kernels in Triton: A Hands-On Guide for ML Engineers

1. Why Triton Exists: The CUDA Tax

Every ML engineer who has tried to write a custom CUDA kernel has encountered the same initiation ritual: 300 lines of boilerplate before a single FLOP is issued. You declare thread blocks. You compute global memory offsets by hand. You manually insert __syncthreads(). You instrument shared-memory bank conflict avoidance through index arithmetic that reads like an ASCII art puzzle. And then, after all that, you benchmark it and discover that the NVIDIA-shipped kernel you were trying to beat is doing something with ldmatrix and cp.async that your code doesn’t — and never will without reading a 90-page ISA extension document.

CUDA is an assembly language in Python’s clothing. It is expressive, but it forces you to manage every level of the hardware hierarchy simultaneously: thread indexing, warp layout, bank conflicts, register pressure, async copy pipelines. This cognitive load exists because CUDA was designed for the full generality of GPU programming — graphics, scientific computing, signal processing. ML workloads don’t need that generality. They need fast tile-level matrix math and memory movement, and they need it expressible at a level of abstraction that doesn’t require a hardware architecture PhD to read.

This is the gap Triton fills. Philippe Tillet et al. introduced Triton in their 2019 paper (Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations, arXiv:1903.04474) with a core insight: ML kernels operate on blocks of elements, not individual threads. If you can write the per-block logic — load a tile, compute on it, store it back — a compiler can handle vectorization, memory coalescing, shared-memory staging, and warp scheduling automatically. Triton takes that bet and wins it for the 90% of ML kernels that fit this pattern.

The practical result: kernels that would take a skilled CUDA programmer a week of tuning can be written in Triton in an afternoon and run at 80–95% of theoretical hardware efficiency. For the remaining 10% — warp-specialized pipeline kernels, Hopper async copy graphs — you still need CUDA. But that boundary is narrow, and understanding it is what separates engineers who know Triton from engineers who use Triton.

2. The Mental Model: You Write Per-Block Logic

Triton’s execution model breaks cleanly from CUDA’s thread-centric thinking. There is no concept of thread index within a block. There are no threadIdx.x gymnastics. Instead, you write a kernel function where the unit of computation is a block of elements — a contiguous 1D or 2D tile of data. Triton broadcasts all scalar operations across that tile automatically.

The three primitives you build everything from:

tl.program_id(axis) — returns the index of the current block in the launch grid, analogous to CUDA’s blockIdx. This is how you figure out which tile of data your block is responsible for.

tl.arange(start, end) — generates a vector of integers [start, ..., end-1] within the block. Combine with program_id to build global memory pointers: if your block size is BLOCK, the i-th block operates on elements [i*BLOCK, (i+1)*BLOCK).

tl.load / tl.store — move data between HBM and registers. Critically, they accept a mask argument that disables out-of-bounds accesses without branching. Triton’s compiler decides whether this data flows through shared memory or registers based on access patterns and autotuning — you don’t specify it.

tl.dot(a, b) — issues a warpgroup-level matrix multiply (wgmma on Hopper, mma.sync on Ampere). The operands must be 2D block tensors. This is the only place where shared memory layout matters, and Triton manages it.

The implication: shared memory is not a variable you declare. It is an implementation detail of the compiler’s choice. You think in tiles; Triton thinks in warps. This division is the source of Triton’s productivity advantage — and its ceiling.

3. A Fused Softmax Kernel: End-to-End

Softmax is the pedagogically perfect Triton example. It is memory-bound by two orders of magnitude (arithmetic intensity ≈ 1.25 FLOPS/byte on H100), and its standard PyTorch implementation makes unnecessary HBM round-trips. Let’s build it from scratch and understand every choice.

The Standard PyTorch Baseline

A naive row-wise softmax over a matrix X ∈ ℝ^{M×N}:

def softmax_pytorch(x):
    x_max = x.max(dim=-1, keepdim=True).values   # HBM read 1
    x_exp = (x - x_max).exp()                    # HBM read 2, write 1
    return x_exp / x_exp.sum(dim=-1, keepdim=True)  # HBM read 3+4, write 2

Four HBM passes. Each one loads and stores the full matrix. For a large sequence — say, the attention score matrix at N = 8192 — this is hundreds of milliseconds of pure bandwidth tax.

The Triton Fused Version

import triton
import triton.language as tl

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE': 128}, num_warps=4),
        triton.Config({'BLOCK_SIZE': 256}, num_warps=8),
        triton.Config({'BLOCK_SIZE': 512}, num_warps=8),
        triton.Config({'BLOCK_SIZE': 1024}, num_warps=16),
        triton.Config({'BLOCK_SIZE': 2048}, num_warps=16),
    ],
    key=['N'],
)
@triton.jit
def fused_softmax_kernel(
    X_ptr, Y_ptr,
    M, N,
    stride_xm, stride_xn,
    stride_ym, stride_yn,
    BLOCK_SIZE: tl.constexpr,
):
    # Each program instance handles one row
    row_idx = tl.program_id(0)

    # Build column offset vector for this block
    col_offsets = tl.arange(0, BLOCK_SIZE)
    mask = col_offsets < N

    # Pointer arithmetic: jump to the correct row
    x_ptrs = X_ptr + row_idx * stride_xm + col_offsets * stride_xn

    # Load the row from HBM into registers (one read)
    row = tl.load(x_ptrs, mask=mask, other=-float('inf'))

    # Online softmax: max subtraction for numerical stability
    row_max = tl.max(row, axis=0)
    row = row - row_max          # broadcast scalar across block
    row_exp = tl.exp(row)
    row_sum = tl.sum(row_exp, axis=0)
    row_out = row_exp / row_sum

    # Store result (one write)
    y_ptrs = Y_ptr + row_idx * stride_ym + col_offsets * stride_yn
    tl.store(y_ptrs, row_out, mask=mask)


def fused_softmax(x: torch.Tensor) -> torch.Tensor:
    M, N = x.shape
    y = torch.empty_like(x)
    # One program per row
    grid = (M,)
    fused_softmax_kernel[grid](
        x, y, M, N,
        x.stride(0), x.stride(1),
        y.stride(0), y.stride(1),
    )
    return y

Anatomy of every decision in this kernel:

BLOCK_SIZE: tl.constexpr — this is not a runtime variable. tl.constexpr means the value is known at compile time and is used to size all internal arrays statically. Triton generates separate compiled code for each BLOCK_SIZE. This is mandatory because Triton’s SRAM allocation is static; you cannot have dynamic tile sizes at runtime.

tl.arange(0, BLOCK_SIZE) — generates the per-element column indices [0, 1, ..., BLOCK_SIZE-1]. Combined with row_idx * stride_xm, this points to the correct row in HBM. The stride encoding means this generalizes to non-contiguous tensors automatically.

mask = col_offsets < N — when N is not a multiple of BLOCK_SIZE, the last columns of the tile are out-of-bounds. The mask argument to tl.load ensures those lanes load the sentinel value (-float('inf')) instead of reading garbage memory. tl.store with a mask silently drops out-of-range writes. This is the clean alternative to CUDA’s if (col < N) branches, which break warp efficiency.

tl.max, tl.sum — these are warp-level reductions. Triton compiles them to tree-reduction sequences across the BLOCK_SIZE elements. The online-softmax pattern (subtract max, then normalize) eliminates numerical overflow without a second pass.

@triton.autotune — launches the kernel with each Config on a calibration run and caches the fastest one, keyed by N. The configs sweep BLOCK_SIZE from 128 to 2048 and num_warps from 4 to 16. The optimal config depends on N: small N fits in a tight block, large N benefits from wider tiles and more warps to hide memory latency.

Block size reasoning: every row must fit entirely within registers (not SRAM — Triton routes small block loads to registers automatically). At BLOCK_SIZE = 1024, FP32 data per row is 4 KB — well within H100 register capacity. At BLOCK_SIZE = 2048, you’re at 8 KB and register pressure begins to limit occupancy. The autotuner finds this empirically; you provide the search space.

4. Why Fused Softmax is 2–3× Faster Than torch.softmax

The roofline model gives the answer immediately. Softmax’s arithmetic intensity is approximately 1.25 FLOPS/byte — roughly 236× below the H100’s compute ridge. This operation will always be memory-bound on any GPU. The only thing that changes performance is how many times you touch HBM.

The PyTorch kernel chain makes four HBM passes: read for max, read-write for exp, read-read-write for divide. The Triton fused kernel makes exactly two: one read, one write. For a 4096-row × 4096-column FP16 matrix, that difference is:

Naive:  4 × (4096 × 4096 × 2 bytes) = 128 MB of HBM traffic
Fused:  2 × (4096 × 4096 × 2 bytes) = 64 MB of HBM traffic
Speedup = 2× from bandwidth reduction alone

The additional speedup (up to 3×) comes from kernel launch overhead elimination — instead of four separate CUDA kernel launches with their synchronization barriers, you have one. At N = 4096, the kernel occupancy is high enough that memory latency is nearly fully hidden. This is the proof-of-concept for Horace He’s “Making Deep Learning Go Brrr” thesis: fusing pointwise operations is free compute, bounded only by the single bandwidth bottleneck of moving the data once.

The same principle scales. GELU after a linear layer, LayerNorm after an attention block, residual adds at every layer boundary — each of these is a separate HBM round-trip in unfused PyTorch. A Triton kernel that fuses them all into the epilogue of the preceding matmul pays bandwidth for the data exactly once. Libraries like Liger Kernel and xformers are essentially catalogs of such fusions, each one described by a trivial roofline calculation.

5. Where Triton Hits Its Ceiling

Triton is a block-level abstraction. When the optimal kernel requires warp-level asymmetry — different groups of warps doing structurally different things simultaneously — Triton cannot express it natively.

Warp specialization on Hopper: H100’s wgmma instruction and TMA (Tensor Memory Accelerator) enable a producer-consumer pipeline where one warpgroup issues asynchronous HBM→SRAM copies while another warpgroup issues wgmma operations on the previously loaded tile. This overlap is what pushes Hopper matmuls to 95%+ MFU. Implementing it requires explicit warpgroup role assignment (__nv_is_warp_group_0() in CUDA), pipeline barrier management (cuda::barrier), and TMA descriptor setup — none of which Triton exposes. Triton’s pipelining (num_stages) approximates this with software prefetch, but cannot beat a hand-tuned warp-specialized kernel by more than 5–10% on large-tile matmuls.

cp.async pipelines: Hopper’s TMA allows HBM→SRAM transfers that bypass the warp scheduler entirely. A raw CUDA kernel using cp.async.bulk can achieve zero-overhead pipelining; Triton’s num_stages parameter inserts equivalent prefetch instructions but with more conservative scheduling. The gap is small on compute-bound shapes and larger on bandwidth-bound ones.

Bank conflicts on unusual tile layouts: Triton automatically inserts shared-memory padding for common layouts, but for exotic access patterns (e.g., the transposed-B tile in column-major layouts), a human inspecting the PTX can often find better padding. The Triton compiler doesn’t have that context.

Bottom line: Triton vs. CUDA is not a performance question for 80% of custom ML kernels. It becomes one only when you need either warp-specialized pipelines (hand-written matmul epilogues, Hopper-native FA3) or register-layout-aware kernels for specific ISA instructions. If you’re writing that, you already know you need CUDA. For everything else — softmax fusions, custom activations, sparse attention patterns, quantization kernels — Triton is the correct tool.

6. Flash Attention 2 in Triton: The Reference Implementation

Flash Attention 2 (Dao, arXiv:2307.08691) is the architectural proof that Triton can reach production quality on the most demanding kernel in modern ML. The original FA1 paper (arXiv:2205.14135) was implemented in CUDA. FA2 ships a Triton implementation that, on most attention shapes, matches the CUDA version within a few percent and substantially outperforms the unfused baseline.

The FA2 Triton kernel fuses the entire attention computation — QKᵀ matmul, online softmax across the key dimension, and PV accumulation — into a single kernel that never materializes the N×N score matrix in HBM. The key loop structure:

@triton.jit
def flash_attention_fwd(
    Q, K, V, Out,
    Lse,                   # log-sum-exp for backward pass
    stride_qz, stride_qh, stride_qm, stride_qk,
    # ... K, V, O strides ...
    Z, H, N_CTX, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
    BLOCK_DMODEL: tl.constexpr, IS_CAUSAL: tl.constexpr,
):
    # Each CTA handles a BLOCK_M × BLOCK_DMODEL output tile
    start_m = tl.program_id(0)
    off_hz  = tl.program_id(1)
    # ...

    # Initialize online softmax state
    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float('inf')
    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
    acc  = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)

    # Load Q tile once — stays in SRAM for the entire inner loop
    q = tl.load(q_ptrs)

    # Inner loop over key blocks
    for start_n in range(0, (start_m + 1) * BLOCK_M if IS_CAUSAL else N_CTX, BLOCK_N):
        k = tl.load(k_ptrs)
        qk = tl.dot(q, tl.trans(k))        # [BLOCK_M, BLOCK_N] score tile
        qk *= sm_scale

        # Causal mask
        if IS_CAUSAL:
            qk += tl.where(causal_mask, 0, float('-inf'))

        # Online softmax update — Dao et al.'s Algorithm 1
        m_ij   = tl.max(qk, 1)             # row-wise max of this block
        p      = tl.exp(qk - m_ij[:, None])
        l_ij   = tl.sum(p, 1)
        m_new  = tl.maximum(m_i, m_ij)
        alpha  = tl.exp(m_i - m_new)
        beta   = tl.exp(m_ij - m_new)
        l_new  = alpha * l_i + beta * l_ij

        # Rescale accumulator and add new contribution
        acc    = acc * alpha[:, None]
        v      = tl.load(v_ptrs)
        acc   += beta[:, None] * tl.dot(p.to(tl.float16), v)

        m_i, l_i = m_new, l_new
        # advance k_ptrs, v_ptrs

    # Normalize and store
    acc = acc / l_i[:, None]
    tl.store(out_ptrs, acc.to(tl.float16))

Three things make this kernel structurally important to study:

First, q is loaded into registers (or SRAM) once and reused across all N/BLOCK_N iterations of the inner loop. The Q tile’s bandwidth cost is amortized over all key blocks — this is the tiling argument from the FA1 paper, now visible as a literal Triton loop.

Second, the online softmax state (m_i, l_i) and the accumulator acc are maintained entirely in registers across the loop. There is no intermediate SRAM write of the score matrix. The running max and sum are updated at each key block using the numerically stable recurrence from Algorithm 1 of Dao et al. This is the realization that softmax can be decomposed into a streaming reduction — the algorithmic key that makes FA possible.

Third, BLOCK_M, BLOCK_N, and BLOCK_DMODEL are all tl.constexpr, enabling autotuning over the full tile-size space. For d = 64, the community-validated defaults are BLOCK_M = BLOCK_N = 64, fitting 3 × 64 × 64 × 2 = 24 KB of Q/K/V tiles in SRAM per CTA. For d = 128, BLOCK_M = BLOCK_N = 128 places 96 KB in SRAM — just inside H100’s 228 KB budget, with room for the FP32 accumulator in registers via wgmma.

The FA2 Triton implementation reaches approximately 73% of the H100 FP16 theoretical throughput at N = 4096, d = 128 in the forward pass — within 8% of the FA2 CUDA reference on the same shapes, and 2–4× faster than PyTorch’s unfused scaled_dot_product_attention fallback. For the backward pass, the Triton implementation is approximately 15% slower than hand-tuned CUDA due to the warp specialization gap described above — a known tradeoff and an active area of development.

7. Practical Workflow Summary

The engineering workflow for Triton kernels is short enough to internalize completely:

Step 1: Roofline first. Compute arithmetic intensity. If it’s below 100 FLOPS/byte, your kernel is memory-bound. The optimization strategy is to reduce HBM passes — which means fusion. If it’s above 100, you’re compute-bound, and tiling + occupancy become the levers.

Step 2: Write the per-block logic. Forget threads. Forget warps. Write the code that processes one tile of data: load, compute, store. Use tl.constexpr for all tile dimensions. Use mask everywhere N % BLOCK_SIZE != 0.

Step 3: Autotune over the right space. BLOCK_SIZE, num_warps, num_stages are the three primary axes. Sweep BLOCK_SIZE over powers of two from 32 to 2048. Set num_stages to 3–4 for bandwidth-bound kernels (deeper pipeline = more HBM latency hidden). Key autotune on the shapes that vary in production.

Step 4: Inspect PTX when performance plateaus. triton.compile(..., target='ptx') dumps the generated PTX. Look for unexpected global memory instructions where you expect register accesses, bank conflict patterns in SRAM, and the depth of the async copy pipeline. The PTX is your ground truth.

Step 5: Fall back to CUDA only for warp specialization. If you need TMA descriptors, producer-consumer warpgroup split, or register-layout-aware wgmma operand staging, write CUDA. Otherwise, you are solving a performance problem that doesn’t exist on your workload.

Closing: The Inversion That Matters

CUDA made the GPU programmable by exposing every degree of freedom simultaneously. Triton makes GPU programming productive by hiding the degrees of freedom that don’t matter for ML kernels and exposing — through tl.constexpr tile sizing and @triton.autotune — the ones that do. The result is a language where an ML engineer can look at a Flash Attention loop and see the algorithm clearly: load Q once, iterate over key blocks, maintain an online softmax state, accumulate the output. The hardware choreography — SRAM staging, warp scheduling, async prefetch — is not absent; it’s just not your problem.

The kernels that push H100s to peak MFU are not mysterious. They are roofline calculations made executable. Triton is the tool that closes the gap between the math and the metal for the overwhelming majority of ML engineers who should be spending their time on the math.

References

• Tillet, P., et al. (2019). Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations. arXiv:1903.04474.
• Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• NVIDIA. (2022). NVIDIA H100 Tensor Core GPU Architecture Whitepaper.
• He, H. (2022). Making Deep Learning Go Brrr From First Principles. https://horace.io/brrr_intro.html
• Luo, W., et al. (2024). Dissecting the NVIDIA Hopper Architecture through Microbenchmarking and Multiple Level Analysis. arXiv:2501.12084.

@article{fp4-2606012,
  title   = {Writing Custom GPU Kernels in Triton: A Hands-On Guide for ML Engineers},
  author  = {fp4 editorial desk},
  year    = {2026},
  url     = {https://fp4.dev/system/triton-kernel-guide/},
  journal = {fp4}
}