Triton Kernel Practice: From Vector Addition to the FlashAttention Core

Summary: The default backend of torch.compile, TorchInductor, translates computation graphs into Triton kernels—to truly understand why PyTorch is fast, or to hand-write high-performance kernels for custom ops, Triton is unavoidable. This post writes three kernels from scratch: vector addition → tiled GEMM → the FlashAttention core, each with full code, performance comparisons, and tuning ideas.

0. Experiment Goals

• Master the Triton programming model: block-level programming + automatic parallelism
• Hand-write vector addition and tiled GEMM; understand tiling and loop unrolling
• Hand-write the FlashAttention core (blocked online softmax) vs native PyTorch
• Learn kernel tuning: block size, tl.constexpr, tl.static_range

1. Why Framework Engineers Should Learn Triton

How it works:

• Hand-written CUDA costs: threads/shared memory/sync all manual. Triton raises the programming model to "block" level—you describe "what each block does"; the compiler maps to threads, schedules, and shared memory.
• The GPU kernels TorchInductor generates are Triton (see KB PyTorch Compile Stack)—learning Triton = reading torch.compile's output.
• Custom ops (fused kernels, quantization, long-tail performance hotspots) in Triton take tens of lines, with performance approaching hand-written CUDA.

2. Kernel 1: Vector Addition (Minimal Runnable)

import torch, triton, triton.language as tl

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)                          # current block id
    offs = pid * BLOCK + tl.arange(0, BLOCK)        # element indices this block covers
    mask = offs < n                                 # boundary mask (n not divisible by BLOCK)
    x = tl.load(x_ptr + offs, mask=mask)
    y = tl.load(y_ptr + offs, mask=mask)
    tl.store(out_ptr + offs, x + y, mask=mask)

def add(x, y):
    n = x.numel()
    out = torch.empty_like(x)
    grid = (triton.cdiv(n, 1024),)                  # num blocks = ceil(n / BLOCK)
    add_kernel[grid](x, y, out, n, BLOCK=1024)
    return out

x = torch.randn(1_000_000, device="cuda")
y = torch.randn(1_000_000, device="cuda")
print("correctness:", torch.allclose(add(x, y), x + y))

Step-by-step:

• @triton.jit: compiles the Python function into a GPU kernel. The body is "per-block logic", not a serial loop.
• tl.program_id(0): current block ID; tl.arange(0, BLOCK): index vector [0, BLOCK)—Triton is vectorized (SIMD-within-block) programming.
• BLOCK: tl.constexpr: compile-time constant. Changing the block size triggers recompilation, so tuning is a "compile-time parameter".
• mask: boundary handling. n=1_000_000 is not divisible by 1024; the last block needs masking to avoid out-of-bounds.
• grid: number of blocks to launch. Triton automatically schedules blocks across all SMs—no manual thread management.

3. Kernel 2: Tiled GEMM

@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):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)    # rows of this tile
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)    # cols of this tile
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, K, BLOCK_K):                  # K-dimension tile loop
        rk = k + tl.arange(0, BLOCK_K)
        a = tl.load(a_ptr + rm[:, None] * stride_am + rk[None, :] * stride_ak)
        b = tl.load(b_ptr + rk[:, None] * stride_bk + rn[None, :] * stride_bn)
        acc += tl.dot(a, b)                          # block matmul (compiled to MMA)
    c = acc.to(tl.float16)
    tl.store(c_ptr + rm[:, None] * stride_cm + rn[None, :] * stride_cn, c)

def matmul(a, b):
    M, K = a.shape
    K2, N = b.shape
    c = torch.empty((M, N), device="cuda", dtype=torch.float16)
    grid = (triton.cdiv(M, 64), triton.cdiv(N, 64))
    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), BLOCK_M=64, BLOCK_N=64, BLOCK_K=32)
    return c

a = torch.randn(1024, 1024, device="cuda", dtype=torch.float16)
b = torch.randn(1024, 1024, device="cuda", dtype=torch.float16)
out = matmul(a, b)
ref = a @ b
print("correctness:", torch.allclose(out.float(), ref.float(), atol=0.5))

How it works:

• Tiling: the output is cut into BLOCK_M × BLOCK_N tiles, one kernel instance per tile. The K dimension is looped in BLOCK_K chunks—the core of high-performance GEMM: data locality + unrollable loop.
• tl.dot(a, b): compiles to GPU MMA (matrix multiply-accumulate) instructions—the key to cuBLAS-level performance.
• Why it is fast: each tile's A/B blocks are reused K times in registers/shared memory; main-memory bandwidth drops to O(MN(K/BLOCK_K)) instead of O(MNK).
• vs cuBLAS: hand-written Triton GEMM typically reaches 70-95% of cuBLAS—and for fused scenarios (e.g., inline multiply-add inside attention) it can even exceed the library.

4. Kernel 3: The FlashAttention Core (online softmax)

@triton.jit
def flashattn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, N, d,
                     BLOCK: tl.constexpr, d_BLOCK: tl.constexpr):
    pid = tl.program_id(0)                            # handle one query block
    start = pid * BLOCK
    rq = start + tl.arange(0, BLOCK)                  # query rows
    rd = tl.arange(0, d_BLOCK)                        # head dimension

    q = tl.load(q_ptr + rq[:, None] * d + rd[None, :])
    m_i = tl.full((BLOCK,), float("-inf"), tl.float32)   # running max
    l_i = tl.zeros((BLOCK,), tl.float32)                 # running exp sum
    acc = tl.zeros((BLOCK, d_BLOCK), tl.float32)

    for start_k in range(0, N, BLOCK):                # scan K/V in blocks
        rk = start_k + tl.arange(0, BLOCK)
        k = tl.load(k_ptr + rk[:, None] * d + rd[None, :])
        # attention scores (no scale, simplified)
        s = tl.dot(q, tl.trans(k))                    # (BLOCK, BLOCK)
        m_ij = tl.max(s, axis=1)
        m_new = tl.maximum(m_i, m_ij)
        p = tl.exp(s - m_new[:, None])                # numerically stable softmax
        alpha = tl.exp(m_i - m_new)
        l_i = l_i * alpha + tl.sum(p, axis=1)
        v = tl.load(v_ptr + rk[:, None] * d + rd[None, :])
        acc = acc * alpha[:, None] + tl.dot(p.to(tl.float16), v)
        m_i = m_new
    acc /= l_i[:, None]
    tl.store(o_ptr + rq[:, None] * d + rd[None, :], acc.to(tl.float16))

# Call (simplified: single head, d=64, N aligned to BLOCK)
N, d = 512, 64
q = torch.randn(N, d, device="cuda", dtype=torch.float16)
k = torch.randn(N, d, device="cuda", dtype=torch.float16)
v = torch.randn(N, d, device="cuda", dtype=torch.float16)
o = torch.empty_like(q)
flashattn_kernel[(triton.cdiv(N, 64),)](q, k, v, o, N, d, BLOCK=64, d_BLOCK=64)
print("output shape:", o.shape)

How it works (the online-softmax trio):

• The problem: standard attention must materialize the full N×N score matrix before softmax—O(N²) memory; long sequences explode.
• The fix: scan K/V in blocks (load only BLOCK×d per iteration) while maintaining three running quantities: m_i (running max), l_i (running exp sum), acc (weighted sum). Per block: update the max → rescale with alpha=exp(m_i-m_new) → update l_i and acc. Memory stays O(N) throughout—the core idea of FlashAttention.
• tl.dot reuses GEMM acceleration; tl.trans transposes the K block.
• Engineering notes: the real implementation adds scale (1/sqrt(d)), masking (causal), and a multi-head loop—this is the minimal runnable core.

5. Tuning Trio

# 1) Block size: larger BLOCK reduces block count/loops but is limited by registers and occupancy
#    GEMM common: BLOCK_M=64/128, BLOCK_N=64/128, BLOCK_K=16/32/64
# 2) Compile-time constants (tl.constexpr): changing a value recompiles—don't hardcode as Python vars
# 3) tl.static_range: when the K loop length is known, unroll it statically

Results (sample magnitude · A100 · FP16):

Kernel	Relative performance
Hand-written Triton GEMM (BLOCK=64/64/32)	~85% of cuBLAS
FlashAttention core (BLOCK=64)	~2-3× vs native PyTorch attention (long sequences, O(N²) memory saved)

Relative-magnitude reference. The point: hand-written Triton kernels reach library-level performance (GEMM 80-95%), and fused scenarios (attention) can beat native implementations—because softmax+multiply-add+output writeback fuse into one kernel, skipping the intermediate N×N tensor reads/writes.

6. Engineering Takeaways

1. Triton is "block-level" programming: describe what a block does; the compiler manages threads/shared memory/scheduling.
2. torch.compile's backend is Triton: learning it = reading the compiled output and customizing fusions.
3. High-performance GEMM = tiling + tl.dot (MMA): unroll the K loop, data locality.
4. The FlashAttention core = the online-softmax trio (m_i/l_i/acc); that is the key to O(N) memory.
5. Tuning params as tl.constexpr: compile-time constants; changing values recompiles.
6. tl.static_range statically unrolls known loops, giving the compiler unrolling opportunities.
7. Write custom ops in Triton: performance close to hand-written CUDA; integrates into the PyTorch ecosystem with zero changes (via dispatch, see Internals).

7. Further Reading

• Triton official tutorials
• FlashAttention paper
• KB: 09. Deep Learning Frameworks · PyTorch Internals · Compile Stack · Triton Learning Path