Triton Kernel Programming

Triton Kernel Programming

Overview

This skill provides a hands-on reference for building production Triton kernels. It covers the triton.language API, autotune decorators, the @triton.jit compilation model, debugging/interpreter workflows, and triton.testing benchmarks.

When to Use

Use this skill when:

• Implementing any custom GPU compute kernel in Triton
• Optimizing inference latency for small-batch transformer operations
• Fusing operations (e.g., matmul + activation, attention with softmax)
• Porting CUDA kernels to Triton for easier maintenance

Do not use for:

• Standard PyTorch operations that already run fast (use torch.compile)
• Distributed or multi-GPU parallelism patterns
• CPU-bound workloads

Installation

# Triton ships with PyTorch ≥2.0. Install latest:
pip install -U triton

# Or build from source for latest features:
git clone https://github.com/triton-lang/triton.git
cd triton
pip install -r python/requirements.txt
pip install .

# For profiling:
pip install nvitools  # NVIDIA profiling helpers
pip install torch_tb_profiler  # PyTorch profiling

Core API Reference

@triton.jit Decorator

Compiles a Python function into a GPU kernel. All code inside must be valid Triton (subset of Python + triton.language ops).

@triton.jit
def kernel(  # ← compiled kernel
    ptr,              # runtime arguments: pointers, scalars
    BLOCK: tl.constexpr,  # constexpr: baked in at compile time
):
    pid = tl.program_id(axis=0)  # SPMD program index
    ...

triton.language (tl) — Key Operations

Memory Operations — Masking Best Practice

# Always mask loads/stores for safety:
mask = offsets < n_elements
x = tl.load(ptr + offsets, mask=mask, other=0.0)
# 'other' provides a safe default for out-of-bounds positions

# For matmul inner loop, use other=0.0 for partial tiles:
a = tl.load(a_ptrs, mask=offsets_k[None, :] < K - k, other=0.0)
b = tl.load(b_ptrs, mask=offsets_k[:, None] < K - k, other=0.0)

Complete Kernel Templates

Template 1: Element-wise Fusion (e.g., LayerNorm)

@triton.jit
def layernorm_kernel(
    input_ptr, output_ptr, weight_ptr, bias_ptr,
    row_stride, n_cols, eps,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(0)
    row_start = pid * row_stride
    offsets = row_start + tl.arange(0, BLOCK_SIZE)
    mask = tl.arange(0, BLOCK_SIZE) < n_cols

    x = tl.load(input_ptr + offsets, mask=mask, other=0.0)
    
    # Mean
    mean = tl.sum(x, axis=0) / n_cols
    # Variance
    x_shifted = x - mean
    var = tl.sum(x_shifted * x_shifted, axis=0) / n_cols
    # Normalize
    x_norm = x_shifted / tl.sqrt(var + eps)
    # Scale + shift
    w = tl.load(weight_ptr + tl.arange(0, BLOCK_SIZE), mask=mask)
    b = tl.load(bias_ptr + tl.arange(0, BLOCK_SIZE), mask=mask)
    y = x_norm * w + b
    
    tl.store(output_ptr + offsets, y, mask=mask)

Template 2: Flash Attention-Style Softmax with Online Safe Computation

@triton.jit
def fused_attention_kernel(
    q_ptr, k_ptr, v_ptr, output_ptr,
    stride_qh, stride_qd,
    stride_kh, stride_kd,
    stride_vh, stride_vd,
    stride_oh, stride_od,
    H, D,
    BLOCK_D: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    pid_h = tl.program_id(0)  # head index
    
    offs_d = tl.arange(0, BLOCK_D)
    offs_n = tl.arange(0, BLOCK_N)
    
    # Load Q block for this head
    q_ptrs = q_ptr + pid_h * stride_qh + offs_d[:, None] * stride_qd
    q = tl.load(q_ptrs)  # (BLOCK_D, 1)
    
    # Online safe softmax over KV sequence
    m_i = tl.full([BLOCK_N], -float('inf'), dtype=tl.float32)
    z_i = tl.zeros([BLOCK_N], dtype=tl.float32)
    acc = tl.zeros([BLOCK_D, BLOCK_N], dtype=tl.float32)
    
    for start_n in range(0, N, BLOCK_N):
        k_ptrs = k_ptr + pid_h * stride_kh + offs_n[None, :] * stride_kd + start_n * stride_kd
        k = tl.load(k_ptrs, mask=offs_n[None, :] < N - start_n, other=0.0)
        
        # S = Q @ K^T
        s = tl.dot(q.T, k)  # (1, BLOCK_N)
        
        # Online safe softmax
        m_ij = tl.maximum(m_i, s)
        p = tl.exp(s - m_ij)
        alpha = tl.exp(m_i - m_ij)
        acc = acc * alpha + p * k.T  # weighted accumulate
        z_i = z_i * alpha + p
        m_i = m_i * 0 + m_ij  # broadcast update
    
    output = acc / z_i
    
    # Store
    out_ptrs = output_ptr + pid_h * stride_oh + offs_d[:, None] * stride_od
    tl.store(out_ptrs, output)

Template 3: FP8 GEMM with Split-K (Inference-Optimized)

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'SPLIT_K': 4}, num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'SPLIT_K': 8}, num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 16, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 128, 'SPLIT_K': 16}, num_warps=8),
    ],
    key=['M', 'N', 'K'],
    prune_configs_by={
        'early_config_prune': lambda configs, named_args: [
            c for c in configs if c.kwargs['BLOCK_SIZE_M'] * c.kwargs['SPLIT_K'] <= 128
        ],
    },
)
@triton.jit
def fp8_gemm_splitk_kernel(
    a_ptr, b_ptr, c_ptr, partial_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    SPLIT_K: tl.constexpr,
):
    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    k_block_id = pid // num_pid_m
    pid_m = pid % num_pid_m
    
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = tl.arange(0, BLOCK_SIZE_N)
    offs_k = k_block_id * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
    
    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((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    
    for k in range(0, K // SPLIT_K, BLOCK_SIZE_K):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K // SPLIT_K - k, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K // SPLIT_K - k, other=0.0)
        acc += tl.dot(a, b)
        a_ptrs += BLOCK_SIZE_K * stride_ak
        b_ptrs += BLOCK_SIZE_K * stride_bk
    
    # Write partial sum
    partial_idx = k_block_id * M + pid_m * BLOCK_SIZE_M
    partial_ptrs = partial_ptr + partial_idx
    tl.store(partial_ptrs, tl.sum(acc, axis=1)[:, None])

Autotuning Strategy

When Autotuning Is Essential

Config Design Heuristics

# Rule of thumb: product of tile dimensions should fit in registers
# BLOCK_SIZE_M * BLOCK_SIZE_N * element_size <= register_budget

# For NVIDIA A100/H100 (fp16 matmul):
configs = [
    # Balanced: good all-around
    triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32}, num_warps=4, num_stages=3),
    # Throughput: large tiles for compute-bound
    triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 64}, num_warps=8, num_stages=4),
    # Latency: small tiles for memory-bound / small M
    triton.Config({'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32}, num_warps=4, num_stages=2),
    # AMD MI300X: use fewer warps, may need waves_per_eu
    triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 32}, num_warps=4, num_stages=0),
]

Profiling Workflow

Step-by-Step: Profile and Optimize

# 1. Warmup: run once to trigger JIT compilation
output_triton = my_kernel(x, y)

# 2. Benchmark with triton.testing
import triton.testing
ms, min_ms, max_ms = triton.testing.do_bench(
    lambda: my_kernel(x, y),
    quantiles=[0.5, 0.2, 0.8],
    warmup=100,  # iterations
    rep=100,     # measurement iterations
)

# 3. Compare to reference
ms_torch, _, _ = triton.testing.do_bench(lambda: torch.matmul(a, b))

# 4. Compute TFLOPS
tflops = lambda ms: 2 * M * N * K * 1e-12 / (ms * 1e-3)
print(f"Triton: {tflops(ms):.2f} TFLOPS | Torch: {tflops(ms_torch):.2f} TFLOPS")

CUDA Graph Integration (Production)

# After autotuning has selected the best config, capture a CUDA graph:
import torch

def capture_gemm_graph(a, b):
    # Warm up with the production shape
    _ = triton_matmul(a, b)
    torch.cuda.synchronize()
    
    # Capture graph
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph):
        c = triton_matmul(a, b)
    
    return graph, c

# Replay for inference — eliminates 1-2ms JIT overhead per launch
graph.replay()

Debugging Cheatsheet

Quality Gates

Cross-References

• triton-kernel-build-design guideline — full design patterns, memory hierarchy, and optimization reference
• Official tutorials: https://triton-lang.org/main/getting-started/tutorials/
• dataset-curation-manifest — when building data-loading kernels
• embedding-analysis — for understanding embedding compute patterns

References