Gunjan Dhanuka — Learning Notes

Learn: Triton & Modern Kernel Authoring

Properties1
tagsgpu, triton, torch-compile, kernels, cuda, autotuning, llm-systems

29 min read

Learn: Triton & Modern Kernel Authoring

What you're learning

How to write GPU kernels in Python at block granularity with Triton — the abstraction level where you reason about program instances and tiles instead of individual SIMT threads. By the end you should be able to write vector-add, fused-softmax, and tiled-matmul kernels from scratch, autotune them, explain how Triton lowers to PTX, and decide when to reach for torch.compile, hand-written Triton, or CUDA/CUTLASS.

This lesson assumes you’ve internalized lessons 04–06: the CUDA SIMT execution model and memory hierarchy (04), GEMM tiling through shared memory (05), and fusion / FlashAttention (06). We’ll repeatedly contrast the Triton mental model against those.

1. Learning map

graph TD
    A["Lesson 04:<br/>CUDA SIMT threads,<br/>memory hierarchy"] --> W["Why Triton:<br/>block-level programming"]
    B["Lesson 05:<br/>GEMM SMEM tiling"] --> W
    C["Lesson 06:<br/>fusion, FlashAttention"] --> W
    W --> PM["Programming model:<br/>@triton.jit, program_id,<br/>arange, load/store+mask, dot"]
    PM --> V["Ex 1: vector add<br/>(1D grid, masking)"]
    V --> S["Ex 2: fused softmax<br/>(one block per row)"]
    S --> M["Ex 3: tiled matmul<br/>(SMEM tiling → a few lines)"]
    M --> AT["Autotuning<br/>@triton.autotune"]
    AT --> CP["Compilation pipeline<br/>Python → TTIR → TTGIR<br/>→ LLVM → PTX"]
    CP --> TI["torch.compile /<br/>TorchInductor<br/>(Triton, automatically)"]
    TI --> DF["Decision framework:<br/>eager → compile → Triton<br/>→ CUTLASS → cuBLAS"]

    style W fill:#44a,color:#fff
    style DF fill:#4a4,color:#fff

Figure 1 — Dependency graph. The CUDA/GEMM/fusion background from lessons 04–06 is exactly what makes Triton’s abstractions feel motivated: Triton automates the tedious parts (coalescing, SMEM staging, pipelining) while leaving you the parts that actually need judgment (tile shapes, program structure).

2. Why Triton — a different altitude on the same hardware

In lesson 04 you wrote CUDA in the SIMT model: you author the program for one thread, then reason about how 32 of them form a warp, how warps form a block, how you hand-stage data into shared memory, __syncthreads(), and how threads must be arranged so their global loads coalesce into 128-byte transactions. Correctness and performance both live at the thread level. That’s maximal control — and maximal surface area for bugs.

Triton, introduced by Tillet, Kung & Cox (2019)1 and open-sourced by OpenAI in 20212, moves you up one level. You write the program for one block — Triton calls it a program instance — and you operate on tiles (small dense sub-tensors, e.g. a 128×64 block). The compiler owns everything inside the block:

  • Intra-block parallelism — it decides how the tile maps onto threads/warps (you only hint via num_warps).
  • Memory coalescing — a tl.load of a contiguous tile is automatically emitted as coalesced global accesses.
  • Shared memory — operands of block-level ops like tl.dot are automatically staged into SMEM, allocated and synchronized via liveness analysis. You don’t write __shared__ or __syncthreads().
  • Software pipelining — loop iterations are pipelined (double/triple buffering) via num_stages.

What you still control: the grid (how many program instances), the tile/block sizes, the program structure (loop order, which axis each program handles, masking), and coarse knobs (num_warps, num_stages).

graph TB
    subgraph CUDA["CUDA — SIMT (lesson 04)"]
        direction TB
        C0["You author: one THREAD"]
        C1["thread → warp (32) → block"]
        C2["YOU manage:<br/>__shared__ staging,<br/>__syncthreads(),<br/>coalescing layout,<br/>register/tile mapping"]
        C0 --> C1 --> C2
    end
    subgraph TRITON["Triton — block-level"]
        direction TB
        T0["You author: one BLOCK<br/>(program instance)"]
        T1["operate on TILES<br/>(e.g. 128×64)"]
        T2["COMPILER manages:<br/>SMEM staging,<br/>sync, coalescing,<br/>thread mapping, pipelining"]
        T0 --> T1 --> T2
    end

Figure 2 — Two altitudes on the same silicon. CUDA: you are the thread and manually orchestrate the block. Triton: you are the block and manipulate whole tiles; the compiler synthesizes the thread-level program. Neither is “higher level” in capability — Triton still targets tensor cores and shared memory — it’s a different granularity of control.

The mental-model shift is the whole game

The most common failure mode coming from CUDA is to keep thinking in threads. In Triton there is no threadIdx. tl.program_id(0) gives you the block index, not a thread index. A single tl.load(ptr + tl.arange(0, BLOCK_SIZE)) is not “one thread loading one element” — it is the whole block loading a BLOCK_SIZE-element tile, which the compiler spreads across all threads of the block and coalesces. If you catch yourself asking “what does thread 5 do here?”, you’re at the wrong altitude — ask “what tile does this program instance own?”

Why this abstraction pays off

The fused-softmax and FlashAttention-style wins from lesson 06 require warp-level reductions and careful SMEM choreography in CUDA. In Triton those become tl.max(row, axis=0) and tl.sum(...) on a tile — a few lines. You keep ~80–95% of hand-CUDA performance for most memory-bound and many compute-bound ops3, at a fraction of the authoring cost, and the same kernel recompiles for NVIDIA (PTX) or AMD (AMDGCN)4.

3. The programming model

The vocabulary you’ll use in every kernel:

ConstructWhat it does
@triton.jitMarks a Python function as a Triton kernel; JIT-compiled on first call, keyed on arg dtypes/constexpr values.
tl.program_id(axis)Index of this program instance along a grid axis (0/1/2). The Triton analogue of blockIdx.
tl.num_programs(axis)Total program instances along an axis (grid size).
tl.arange(0, N)A compile-time-sized [0..N) tile of offsets; N must be a power-of-two constexpr.
tl.load(ptr, mask=, other=)Gather a tile from global memory; mask guards out-of-bounds lanes, other fills them.
tl.store(ptr, val, mask=)Scatter a tile to global memory, masked.
tl.dot(a, b, acc)Tile matmul acc + a @ b; uses tensor cores when dims are multiples of 16.
tl.max/tl.sum/tl.minTile reductions along axis.
BLOCK_SIZE: tl.constexprA compile-time constant — specializes the kernel and enables unrolling/vectorization.
grid launch kernel[grid](...)grid is a tuple or a lambda meta: (...); launches that many program instances.

The constexpr distinction is load-bearing: BLOCK_SIZE is baked into the compiled kernel (different values → different compiled binaries), which is exactly what lets the compiler unroll tl.arange loops and pick vector widths. Runtime scalars like n_elements are ordinary arguments.

4. Worked example 1 — vector add

The “hello world.” One 1D grid; each program instance owns a BLOCK_SIZE-element tile. Mirrors the official vector-add tutorial.

import torch
import triton
import triton.language as tl
 
@triton.jit
def add_kernel(
    x_ptr,          # *Pointer* to first input
    y_ptr,          # *Pointer* to second input
    out_ptr,        # *Pointer* to output
    n_elements,     # runtime scalar (NOT constexpr)
    BLOCK_SIZE: tl.constexpr,   # tile size, compile-time constant
):
    # Which block am I? (analogue of CUDA blockIdx.x)
    pid = tl.program_id(axis=0)
 
    # This block owns elements [block_start, block_start + BLOCK_SIZE)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)   # a whole tile of indices
 
    # Guard the tail: the last block may run past n_elements.
    mask = offsets < n_elements
 
    # Whole-tile loads/stores; the compiler coalesces these.
    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)
 
 
def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    n_elements = out.numel()
    # grid is a lambda over the compile-time meta-params: how many blocks?
    grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
    add_kernel[grid](x, y, out, n_elements, BLOCK_SIZE=1024)
    return out
 
 
if __name__ == "__main__":
    x = torch.rand(98_432, device="cuda")
    y = torch.rand(98_432, device="cuda")
    out = add(x, y)
    torch.testing.assert_close(out, x + y)
    print("ok")

Notice what’s absent versus CUDA: no thread indexing, no blockDim, no manual coalescing. tl.arange(0, BLOCK_SIZE) is the tile; the compiler maps it onto threads. triton.cdiv is ceil-division so the grid covers a non-multiple length, and the mask cleans up the tail.

Masking is not optional — it is correctness

Because BLOCK_SIZE is a fixed constexpr and lengths rarely divide it evenly, every boundary load/store needs a mask. Forget it and the last block reads/writes out of bounds — often silently (garbage results), sometimes an illegal-memory-access crash. For reductions, also set other= to the reduction’s identity: -inf for max, 0 for sum. This is the single most frequent Triton bug for people arriving from NumPy, where slicing handles bounds for you.

5. Worked example 2 — fused softmax (the canonical Triton win)

Row-wise softmax on an M×N matrix. This is the textbook demonstration of why fusion beats eager PyTorch. From lesson 06’s memory-wall framing: a naive PyTorch softmax materializes several intermediates (x.max, x - max, exp, sum, divide), reading ~5MN and writing ~3MN elements from/to DRAM. A fused kernel reads X once, does everything in registers/SRAM, and writes Y once — roughly a 4× traffic reduction on a bandwidth-bound op.5

Strategy: one program instance per row (or a persistent program that strides over rows). Each program loads its entire row into SRAM as a single tile, reduces on-chip, writes back. This is the official fused-softmax tutorial; the CUDA equivalent needs explicit warp-shuffle reductions and synchronization — here it’s tl.max/tl.sum.

import torch
import triton
import triton.language as tl
 
@triton.jit
def softmax_kernel(
    output_ptr, input_ptr,
    input_row_stride, output_row_stride,
    n_rows, n_cols,
    BLOCK_SIZE: tl.constexpr,     # >= n_cols, next power of two
    num_stages: tl.constexpr,
):
    # Persistent programs: each one strides over multiple rows.
    row_start = tl.program_id(0)
    row_step = tl.num_programs(0)
    for row_idx in tl.range(row_start, n_rows, row_step, num_stages=num_stages):
        # Pointer to the start of this row.
        row_start_ptr = input_ptr + row_idx * input_row_stride
        col_offsets = tl.arange(0, BLOCK_SIZE)
        input_ptrs = row_start_ptr + col_offsets
 
        # BLOCK_SIZE may exceed n_cols → mask, fill padding with -inf so it
        # never wins the max and contributes 0 to the sum after exp.
        mask = col_offsets < n_cols
        row = tl.load(input_ptrs, mask=mask, other=-float("inf"))
 
        # Numerically-stable softmax, entirely on-chip.
        row_minus_max = row - tl.max(row, axis=0)
        numerator = tl.exp(row_minus_max)
        denominator = tl.sum(numerator, axis=0)
        softmax_output = numerator / denominator
 
        # Write the row back once.
        output_row_start_ptr = output_ptr + row_idx * output_row_stride
        output_ptrs = output_row_start_ptr + col_offsets
        tl.store(output_ptrs, softmax_output, mask=mask)
 
 
def softmax(x: torch.Tensor) -> torch.Tensor:
    n_rows, n_cols = x.shape
    # One block must hold a whole row → next power of two >= n_cols.
    BLOCK_SIZE = triton.next_power_of_2(n_cols)
    # More columns per row → more warps to spread the work.
    num_warps = 8 if BLOCK_SIZE >= 2048 else (4 if BLOCK_SIZE >= 512 else 2)
    num_stages = 4 if BLOCK_SIZE < 32768 else 2
 
    y = torch.empty_like(x)
    # Launch as many persistent programs as the device can co-resident; a
    # simple, robust choice is min(#SMs * occupancy, n_rows). For clarity we
    # just use n_rows here (one program per row is always correct).
    grid = (n_rows,)
    softmax_kernel[grid](
        y, x, x.stride(0), y.stride(0), n_rows, n_cols,
        BLOCK_SIZE=BLOCK_SIZE, num_stages=num_stages, num_warps=num_warps,
    )
    return y
 
 
if __name__ == "__main__":
    x = torch.randn(1823, 781, device="cuda")
    torch.testing.assert_close(softmax(x), torch.softmax(x, axis=1))
    print("ok")

The three ideas that make this fast: (1) one row = one tile, so the whole reduction happens in SRAM/registers with no round-trip; (2) other=-float("inf") makes the mask correct through both the max and the exp→sum; (3) subtracting the row max before exp is the standard numerical-stability trick — identical math to eager PyTorch, but never spilled to DRAM.

"One block per row" has a ceiling

This pattern assumes an entire row fits in one block — i.e. BLOCK_SIZE ≥ n_cols and the tile fits in registers/SRAM. For very wide rows (say n_cols in the hundreds of thousands) you run out of shared memory / registers and must switch to an online / streaming softmax that tiles the row and combines running (max, sum) — the exact trick lesson 06 used inside FlashAttention. Know which regime you’re in.

6. Worked example 3 — tiled matmul (SMEM tiling in a few lines)

Lesson 05 built a GEMM by hand: load BLOCK_M×BLOCK_K and BLOCK_K×BLOCK_N tiles into __shared__, __syncthreads(), multiply-accumulate into registers, loop over K. In Triton that entire choreography — SMEM staging, sync, and even double-buffering — is implicit. You write the tiling; the compiler stages tiles into shared memory (by inspecting tl.dot operands) and pipelines the K-loop. Based on the official matmul tutorial.

import torch
import triton
import triton.language as tl
 
@triton.autotune(
    configs=[
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 64, "GROUP_M": 8},
                      num_stages=3, num_warps=8),
        triton.Config({"BLOCK_M": 64,  "BLOCK_N": 256, "BLOCK_K": 32, "GROUP_M": 8},
                      num_stages=4, num_warps=4),
        triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "GROUP_M": 8},
                      num_stages=4, num_warps=4),
        triton.Config({"BLOCK_M": 64,  "BLOCK_N": 64,  "BLOCK_K": 32, "GROUP_M": 8},
                      num_stages=5, num_warps=2),
    ],
    key=["M", "N", "K"],   # re-tune when any of these change
)
@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,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
    GROUP_M: tl.constexpr,
):
    # ---- Grouped program ordering for L2 reuse (see note below) ----
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    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 % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m
 
    # ---- Tile offset ranges ----
    offs_am = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)) % M
    offs_bn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) % N
    offs_k = tl.arange(0, BLOCK_K)
 
    # Pointers to the first A and B tiles.
    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
 
    # ---- The K-loop: SMEM staging + pipelining are IMPLICIT ----
    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        # Mask the K tail so we don't read past K.
        k_mask = offs_k[None, :] < K - k * BLOCK_K
        a = tl.load(a_ptrs, mask=k_mask, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0)
        # Tensor-core matmul; accumulate in FP32.
        accumulator = tl.dot(a, b, accumulator)
        # Advance to the next K tile.
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk
 
    c = accumulator.to(tl.float16)   # cast after accumulating in FP32
 
    # ---- Store the output tile, masked on the M/N edges ----
    offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)
 
 
def matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    M, K = a.shape
    K2, N = b.shape
    assert K == K2
    c = torch.empty((M, N), device=a.device, dtype=torch.float16)
    grid = lambda meta: (
        triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]),
    )
    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),
    )
    return c

Three things to see. First, the entire shared-memory dance from lesson 05 collapsed into accumulator = tl.dot(a, b, accumulator) inside a plain Python for loop — the compiler stages a and b tiles into SMEM and pipelines the loop (num_stages). Second, FP32 accumulation (tl.zeros(..., tl.float32)) with a cast to FP16 only at the end — same numerics discipline as cuBLAS. Third, the grouped program ordering: instead of walking output tiles row-major, programs are grouped into GROUP_M-row super-blocks so that consecutive programs reuse the same rows of A / columns of B, dramatically improving L2 hit rate (the tutorial’s 9×9 example drops from ~90 tile loads to ~54).6

tl.dot is your tensor-core door

tl.dot compiles to mma/wgmma tensor-core instructions when the tile dims are multiples of 16 (and dtypes are supported). You never touch wmma/PTX MMA intrinsics or CUTLASS templates — but you also don’t get CUTLASS’s full control over warp specialization and epilogue fusion. For most GEMM shapes this is close to cuBLAS3; for the shapes cuBLAS/CUTLASS are hand-tuned for (large square FP16/FP8), the vendor libraries still win the last several percent. That trade-off drives the decision framework in §10.

7. Autotuning

The optimal BLOCK_*, num_warps, and num_stages depend on the problem shape and the GPU. Rather than guess, wrap the kernel in @triton.autotune: you give it a list of triton.Configs and a key; on the first call for each distinct key value, Triton benchmarks every config by wall-clock and caches the winner.7

@triton.autotune(
    configs=[
        triton.Config({"BLOCK_SIZE": 512},  num_warps=2),
        triton.Config({"BLOCK_SIZE": 1024}, num_warps=4),
        triton.Config({"BLOCK_SIZE": 2048}, num_warps=8),
        triton.Config({"BLOCK_SIZE": 4096}, num_warps=8, num_stages=3),
    ],
    key=["n_elements"],   # different sizes → re-benchmark
)
@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    tl.store(out_ptr + offsets,
             tl.load(x_ptr + offsets, mask=mask) + tl.load(y_ptr + offsets, mask=mask),
             mask=mask)

Recall the two most impactful compilation knobs (both triton.Config fields):

  • num_warps — thread-level parallelism inside a block. num_warps=8 → the compiler spreads the tile across 8 × 32 = 256 threads.
  • num_stages — software-pipelining depth of the K-loop (double/triple buffering of loads). Most useful for matmul on SM80+; set num_stages=1 to disable pipelining when debugging.

Autotuning is empirical, and it costs the first call

The autotuner doesn’t model the hardware — it just times each config and keeps the fastest. Set TRITON_PRINT_AUTOTUNING=1 to see the winner and time spent. Two practical consequences: (1) the first invocation for each new key value is slow (it benchmarks everything), so keep the config list focused; (2) put every shape parameter that changes the optimal config in key, or you’ll silently reuse a stale winner.

8. How Triton compiles

Understanding the lowering demystifies what Triton abstracts (and what it can’t). The pipeline is a progressive MLIR-based lowering (see the PyTorch blog, Triton kernel compilation stages)8:

graph LR
    P["Python<br/>@triton.jit fn"] -->|"walk AST,<br/>SSA build"| TTIR["Triton IR (TTIR)<br/>MLIR dialect<br/>machine-independent<br/>(CSE, DCE, LICM,<br/>unroll)"]
    TTIR -->|"add GPU layouts<br/>+ passes"| TTGIR["TritonGPU IR (TTGIR)<br/>MLIR dialect<br/>layout encodings,<br/>coalesce, pipeline,<br/>tensor-core lowering"]
    TTGIR --> LL["LLVM IR"]
    LL -->|"NVIDIA"| PTX["PTX → ptxas → cubin"]
    LL -->|"AMD"| GCN["AMDGCN → hsaco"]

    style TTIR fill:#44a,color:#fff
    style TTGIR fill:#44a,color:#fff

Figure 3 — Triton compilation pipeline. Both TTIR and TTGIR are MLIR dialects. The interesting hardware-specific work happens at TTGIR, where layout encodings (#blocked, #shared, nvidia_mma/amd_mfma, dot_op, slice) describe how each tile is distributed across warps and lanes, and where coalescing, shared-memory staging, and software pipelining are applied. Note: Triton’s own MLIR pipeline — not the closed-source ptxas — is ~4/5 of compile time.8

Stage by stage:

  1. Python AST → Triton IR (TTIR). The @triton.jit decorator walks the function’s AST and builds SSA-form Triton-IR. Optimizations here are hardware-independent: inlining, common-subexpression elimination, dead-code elimination, loop-invariant code motion, unrolling.
  2. TTIR → TritonGPU IR (TTGIR). GPU-specific. This is where layout encodings get attached to tensors — the compiler decides how a 128×64 tile is partitioned across warps and threads — and where the passes you actually benefit from run: coalescing, shared-memory allocation for tl.dot operands, software pipelining (num_stages), tensor-core (mma/wgmma) lowering, and vendor-specific passes (NVIDIA TMA/async-dot, AMD LDS/ping-pong).
  3. TTGIR → LLVM IR → target. Lower to LLVM IR, then to PTX (assembled by ptxas into a cubin) on NVIDIA, or AMDGCN (→ hsaco) on AMD. Artifacts (.ttir, .ttgir, .llir, .ptx, .cubin) are cached on disk; dump them with TRITON_KERNEL_DUMP=1.

So: Triton abstracts thread-level details, coalescing, SMEM management, and pipelining (the TTGIR passes). You still control block/tile sizes, program structure, and the coarse num_warps/num_stages knobs. What you cannot express is anything below the tile abstraction — bespoke warp specialization, custom async-copy schedules, or hand-laid register allocation. That’s where CUDA/CUTLASS remains necessary.

9. torch.compile / TorchInductor — Triton, generated for you

Here’s the twist that reframes everything above: most Triton kernels running in production LLM code today were not written by a human. torch.compile (PyTorch 2.x, default backend TorchInductor) generates Triton automatically (PyTorch 2 / TorchInductor)910.

The pipeline:

  1. TorchDynamo hooks CPython frame evaluation to safely capture the FX graph of your eager code.9
  2. AOTAutograd traces an ahead-of-time backward graph, and PrimTorch canonicalizes ~2000+ ATen ops down to ~250 primitives.11
  3. TorchInductor lowers that graph and, for GPU, emits Triton kernels — fusing pointwise + reduction chains, choosing persistent vs. non-persistent reductions, and picking tile sizes/autotuning automatically. On CPU it emits OpenMP C++ instead.9

The headline capability is fusion: Inductor spots memory-bound op chains and fuses them into one Triton kernel that reads inputs once, computes in registers, and writes once — exactly the DRAM-traffic win you engineered by hand in §5’s fused softmax, but derived automatically from torch.softmax. For large GEMMs, Inductor is selective: it typically defers to cuBLAS/cuDNN, and under max-autotune it benchmarks cuBLAS vs. CUTLASS vs. a Triton template and picks the winner, optionally fusing the pointwise epilogue in.12

import torch
 
@torch.compile   # TorchInductor fuses this into (typically) one Triton kernel
def gelu_bias(x, b):
    x = x + b
    return x * 0.5 * (1.0 + torch.erf(x / 2**0.5))
 
x = torch.randn(8192, 4096, device="cuda", dtype=torch.float16)
b = torch.randn(4096, device="cuda", dtype=torch.float16)
y = gelu_bias(x, b)   # first call compiles; subsequent calls hit cached kernels

To see the generated Triton, set TORCH_LOGS="output_code" (or torch._inductor.config.debug=True) — you’ll get readable @triton.jit kernels you can learn from and even copy out.

The relationship between hand-written and compiler-generated Triton is not competitive but layered:

  • Compiler-generated Triton covers the long tail of elementwise/reduction/norm fusions across a whole model, with zero authoring effort and correctness guaranteed by the tracer. It’s excellent at fusing what’s adjacent in the graph.
  • Hand-written Triton wins where Inductor’s pattern matching gives up: a novel algorithm the graph doesn’t express (FlashAttention-style online softmax, fused MoE routing, custom quantized matmul), or where you need control the heuristics won’t find. You can drop a user-defined @triton.jit kernel into torch.compile — Inductor treats it as an opaque node and still fuses around it and autotunes it.13

Real-world evidence for hand-written Triton: Liger-Kernel (LinkedIn) — a library of hand-written Triton kernels (fused RMSNorm, RoPE, SwiGLU, and a fused-linear-cross-entropy that avoids materializing logits) that cut LLM-training memory and boost throughput beyond what torch.compile finds14; and Unsloth, whose fast fine-tuning rests on hand-written Triton kernels15. Both exist precisely because there’s a gap above what Inductor generates automatically.

10. When to use what — the decision framework

A ladder from fastest-to-develop to most-control. Climb only as far as the profiler forces you.

graph TD
    E["Eager PyTorch<br/>dev speed: ★★★★★<br/>perf: baseline"] --> TC["torch.compile<br/>dev speed: ★★★★★ (1 line)<br/>perf: fuses memory-bound ops,<br/>picks GEMM backend"]
    TC --> HT["Hand-written Triton<br/>dev speed: ★★★<br/>perf: novel fusions,<br/>tensor cores via tl.dot,<br/>NV+AMD portable"]
    HT --> CU["CUDA / CUTLASS<br/>dev speed: ★★<br/>perf: full warp-spec,<br/>async copy, epilogue,<br/>last few %"]
    CU --> LIB["cuBLAS / cuDNN<br/>dev speed: ★★★★<br/>perf: peak on standard<br/>GEMM/conv, NVIDIA-only,<br/>closed, inflexible"]

Figure 4 — The kernel-authoring ladder. Each rung trades development speed for last-mile control. Most LLM systems work lives on the first two rungs; the top rung (vendor libraries) is where you land for standard GEMM/conv rather than a step you climb through.

Decision heuristics:

Reach for…When
Eager PyTorchPrototyping, research iteration, anything not yet proven hot. Correctness and dev speed dominate.
torch.compileAlways try it before hand-optimizing. Big wins on memory-bound fusion chains; free GEMM-backend selection. The default for production training/inference.
Hand-written TritonA specific hotspot Inductor can’t express or beats poorly; you need tensor cores via tl.dot with custom fusion; you want one source that runs on NVIDIA and AMD.
CUDA / CUTLASSThe last few percent on a critical GEMM/attention kernel; you need warp specialization, custom async-copy pipelines, TMA/wgmma control, or bespoke epilogues Triton won’t emit.
cuBLAS / cuDNNStandard, large GEMM/conv on NVIDIA — decades of hand-tuning you can’t beat. This is what torch.compile and even hand-Triton call into for the heavy GEMMs. Closed, NVIDIA-only, inflexible on fusion.

The through-line for a GPU engineer at a frontier lab: torch.compile handles the breadth; your hand-written Triton (and occasionally CUDA/CUTLASS) handles the handful of kernels that define the model’s efficiency frontier — attention variants, MoE, quantized matmuls, custom losses.

11. Practice problems

12. Practice & resources

Hands-on (do these in order):

  1. Work through the four official Triton tutorials end to end, typing them yourself, not copy-pasting: vector add → fused softmax → matrix multiplication → low-memory dropout. For each, run the included triton.testing.Benchmark and read the plot against the PyTorch baseline.
  2. Take one memory-bound fusion in your own model (e.g. bias → GELU → dropout, or RMSNorm) and write a single fused Triton kernel for it. Then wrap the eager version in @torch.compile, dump Inductor’s generated kernel with TORCH_LOGS="output_code", and compare — both the wall-clock (triton.testing.do_bench) and the generated code. Note where they converge and where your hand version wins.
  3. Add @triton.autotune to your kernel over a handful of BLOCK_SIZE/num_warps configs; set TRITON_PRINT_AUTOTUNING=1 and see which config wins on your GPU.
  4. Dump the IR of your matmul with TRITON_KERNEL_DUMP=1 and read the .ttgir — find the #blocked and nvidia_mma layouts and the pipelined K-loop.

Read / watch:

13. What’s next

Two threads to pull:

  1. 08-data-parallelism-and-collectives (next) — you can now write fast single-GPU kernels; next is scaling across GPUs: data/tensor/pipeline parallelism and the NCCL collectives (all-reduce, all-gather, reduce-scatter) that stitch them together.
  2. 06-memory-wall-and-flashattention (back-reference) — revisit FlashAttention now that you can read/write Triton. The online-softmax tiling you studied there is a canonical hand-written Triton kernel; try sketching its inner loop with tl.dot and running max/sum accumulators, and compare to the official fused-attention tutorial.

14. References

Part of gpu-systems-for-llms | Filed: 2026-09-02

Footnotes

  1. Philippe Tillet, H. T. Kung, David Cox, “Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations,” MAPL 2019 (3rd ACM SIGPLAN Int’l Workshop on Machine Learning and Programming Languages), DOI 10.1145/3315508.3329973. PDF: eecs.harvard.edu/~htk/publication/2019-mapl-tillet-kung-cox.pdf. Introduces the tile abstraction and block-level programming model. [established] ↩

  2. OpenAI, “Introducing Triton: Open-source GPU programming for neural networks,” July 2021, openai.com/index/triton. OpenAI’s open-sourcing and re-implementation of Triton. [established] ↩

  3. The ~80–95%-of-hand-CUDA figure is an approximate synthesis, not a single benchmarked number. Primary support: the Triton paper reports matmul “on par with cuBLAS,” reaching >90% of device peak performance on some tasks (Tillet, Kung & Cox, MAPL 2019, PDF). Exact fraction is workload- and shape-dependent. [contested] ↩ ↩2

  4. Triton supports multiple hardware backends (NVIDIA → PTX, AMD → AMDGCN) from one source; see the Triton documentation and third-party backends, triton-lang.org. [established] ↩

  5. OpenAI Triton official tutorial, “Fused Softmax,” triton-lang.org/main/getting-started/tutorials/02-fused-softmax.html. Motivates the fused, single-read/single-write kernel and its DRAM-traffic reduction over eager PyTorch. [established] ↩

  6. OpenAI Triton official tutorial, “Matrix Multiplication,” triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html. Source of the grouped-ordering L2-reuse example (9×9 grid, ~90 → ~54 tile loads). [established] ↩

  7. OpenAI Triton API reference, triton.autotune and triton.Config. Documents empirical per-key benchmarking of configs and caching of the winner. [established] ↩

  8. PyTorch blog, “Triton kernel compilation stages,” pytorch.org/blog/triton-kernel-compilation-stages. Walkthrough of the Python → TTIR → TTGIR → LLVM IR → PTX/AMDGCN MLIR-based lowering, including the observation that Triton’s own compiler passes (not ptxas) dominate compile time. [established] ↩ ↩2

  9. Jason Ansel et al., “PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation,” ASPLOS 2024, DOI 10.1145/3620665.3640366. PDF: docs.pytorch.org/assets/pytorch2-2.pdf. Describes TorchDynamo (CPython frame-evaluation graph capture), AOTAutograd, and TorchInductor (Triton for GPU, C++/OpenMP for CPU). [established] ↩ ↩2 ↩3

  10. PyTorch official tutorial, “Introduction to torch.compile,” docs.pytorch.org/tutorials/intermediate/torch_compile_tutorial.html. [established] ↩

  11. PyTorch 2.0 overview, “GET STARTED / PyTorch 2.0,” pytorch.org/get-started/pytorch-2.0: “PrimTorch canonicalizes ~2000+ PyTorch operators down to a closed set of ~250 primitive operators.” [established] ↩

  12. TorchInductor max-autotune mode benchmarks candidate backends (cuBLAS/cuDNN, CUTLASS, Triton templates) for GEMM/conv and selects the fastest, with epilogue fusion. See PyTorch docs, torch.compile / TorchInductor configuration, docs.pytorch.org/docs/stable/torch.compiler.html and the torch._inductor.config max_autotune options. [established] ↩

  13. PyTorch official recipe, “Using User-Defined Triton Kernels with torch.compile,” docs.pytorch.org/tutorials/recipes/torch_compile_user_defined_triton_kernel_tutorial.html. Inductor treats a user @triton.jit kernel as an opaque node and fuses/autotunes around it. [established] ↩

  14. Pin-Lun Hsu et al. (LinkedIn), “Liger Kernel: Efficient Triton Kernels for LLM Training,” arXiv:2410.10989 (2024); code: github.com/linkedin/Liger-Kernel. Reports ~20% average training-throughput gain and ~60% GPU-memory reduction vs. HuggingFace baselines via fused RMSNorm/RoPE/SwiGLU and fused-linear-cross-entropy. [recent] ↩

  15. Unsloth, open-source fast LLM fine-tuning built on hand-written Triton kernels, github.com/unslothai/unsloth. [recent] ↩

Graph View

Backlinks