triton-kernel

Writing Optimized Triton GPU Kernels

Targets: Triton >= 2.1, any GPU with tl.dot support (SM70+/CDNA2+)

Core Patterns (always apply)

Kernel structure: Use @triton.jit decorator. Get block ID with tl.program_id(axis). Compute element offsets with tl.arange(0, BLOCK_SIZE). Build mask = offsets < n_elements for all loads/stores.

Block sizes: Strongly prefer powers of two (required for tl.arange; non-power-of-two may work but can reduce performance). Declare as tl.constexpr parameters. Use @triton.autotune to sweep BLOCK_SIZE_M/N/K configs per hardware.

Memory hierarchy: Keep intermediates in SRAM via block-level reductions (tl.sum, tl.max) before writing to global memory. Fuse multiple pointwise ops into one kernel to avoid DRAM round-trips.

Matmul: Use tl.dot(a, b) for tensor core operations. Always accumulate in tl.float32 when inputs are FP16. For L2 cache locality, use grouped tile ordering via group_id = pid // GROUP_SIZE.

Grid launching: Size grid dynamically: grid = lambda meta: (triton.cdiv(n, meta['BLOCK_SIZE']),).

Masking: ALWAYS mask boundary loads/stores: tl.load(ptr + offs, mask=offs < dim, other=0.0). Missing masks corrupt memory silently.

Benchmarking: Use triton.testing.Benchmark with x_names, x_vals, line_arg, line_vals to compare against PyTorch baselines.

Quick Reference Examples

Fused row-wise softmax — verified, based on official Triton tutorial:

@triton.jit
def fused_softmax(x_ptr, out_ptr, cols, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    offs = tl.arange(0, BLOCK)
    mask = offs < cols
    x = tl.load(x_ptr + row * cols + offs, mask=mask, other=-1e9)
    x_max = tl.max(x, axis=0)
    ex = tl.exp(x - x_max)
    out = ex / tl.sum(ex, axis=0)
    tl.store(out_ptr + row * cols + offs, out, mask=mask)

Seed-based dropout — verified, based on official Triton tutorial:

@triton.jit
def dropout(x_ptr, out_ptr, seed, p, n, BLOCK: tl.constexpr):
    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask)
    r = tl.rand(seed, offs)  # Philox PRNG, deterministic
    keep = r > p
    tl.store(out_ptr + offs, x * keep / (1.0 - p), mask=mask)