# Triton Kernel Programming

> Hands-on implementation template and API reference for writing, tuning, debugging, and benchmarking Triton GPU kernels. Covers the full triton.language API surface, autotuning patterns, profiling workflows, and production integration.

- Skill: `mkurman/triton-kernel-programming` (Agent Skill, multi-file: 2 files)
- Install (CLI): `npx skillmds@latest add mkurman/triton-kernel-programming`
- Raw SKILL.md: https://api.skillmd.com/api/skills/mkurman/triton-kernel-programming/raw
- Safety review: pending
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: Integrations & APIs
- Author: mkurman (https://skillmd.com/u/mkurman)
- Updated: 2026-09-17
- Page: https://skillmd.com/skills/mkurman/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

```bash
# 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).

```python
@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

| Category | Operation | Description |
|----------|-----------|-------------|
| **Indexing** | `tl.program_id(axis)` | SPMD program index along axis 0, 1, or 2 |
| **Ranges** | `tl.arange(start, end)` | 1D range tensor for vectorized addressing |
| **Arithmetic** | `tl.sum`, `tl.max`, `tl.min`, `tl.argmax` | Block reduction along axis |
| **Arithmetic** | `tl.dot(a, b)` | Block matrix multiply (triggers tensor cores) |
| **Activation** | `tl.exp`, `tl.log`, `tl.sigmoid`, `tl.tanh` | Element-wise math |
| **Activation** | `tl.sqrt`, `tl.abs`, `tl.where` | Element-wise ops |
| **Memory** | `tl.load(ptr, mask=, other=)` | Vector load from global memory |
| **Memory** | `tl.store(ptr, val, mask=)` | Vector store to global memory |
| **Memory** | `tl.atomic_add(ptr, val)` | Atomic add (for reductions) |
| **Cast** | `tensor.to(tl.float16)` | Type conversion |
| **Cast** | `tl.cast(tensor, tl.float32)` | Explicit type conversion |
| **Debug** | `tl.device_print("x:", x)` | Runtime print |
| **Debug** | `tl.device_assert(cond, "msg")` | Runtime assertion |
| **Debug** | `tl.static_print(x)` | Compile-time print |
| **Debug** | `tl.static_assert(cond, "msg")` | Compile-time assert |

### Memory Operations — Masking Best Practice

```python
# 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)

```python
@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

```python
@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)

```python
@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

| Scenario | Autotune Impact |
|----------|----------------|
| Variable input shapes (VLLM, serving) | Critical — cache per shape |
| Fixed production shapes | Run once, freeze config |
| Memory-bound ops (softmax, norms) | Less critical — memory access pattern dominates |
| Compute-bound ops (GEMM) | Critical — 2–5x perf difference between configs |

### Config Design Heuristics

```python
# 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

```python
# 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)

```python
# 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

| Problem | Symptom | Fix |
|---------|---------|-----|
| Wrong output | Off-by-one in offsets | Check `mask` logic, use `%` modulo for boundaries |
| NaN output | Numerical instability | Subtract max before exp; check division by zero |
| Slow kernel (memory-bound) | Low bandwidth util | Increase tile sizes, check `_b128` in ISA |
| Slow kernel (compute-bound) | Low TFLOPS | Check tensor core usage in PTX; try `num_stages` tuning |
| Compilation error | `@triton.jit` function issue | Check for unsupported Python constructs (no dictionaries, no dynamic indexing) |
| `compute-sanitizer` errors | Out-of-bounds access | Check mask coverage for partial tiles |
| High launch overhead | CPU-side latency | Use CUDA Graphs for production inference |

## Quality Gates

| Gate | Command/Check | Expected |
|------|--------------|----------|
| Correctness | `torch.max(torch.abs(ref - triton_out))` | `< 0.01` (fp16) or `< 0.5` (fp8) |
| Autotuning | `TRITON_PRINT_AUTOTUNING=1` env var | Best config printed |
| Tensor core usage | Check PTX for `wgmma`/`mma` | Present for matmul kernels |
| Memory coalescing | Check ISA for `global_load_dwordx4` | Present in hot loop |
| LDS usage | `grep "triton_gpu.shared"` from MLIR dump | `< 64 KB` |
| Occupancy | Compute from VGPR/LDS counts | `> 50%` for compute-bound |
| Speedup | `triton.testing.do_bench` | `> 1.5x` over naive PyTorch |

## 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

| Resource | Link |
|----------|------|
| Triton Python API | https://triton-lang.org/main/python-api/ |
| Triton Autotune | https://triton-lang.org/main/python-api/generated/triton.autotune.html |
| Triton Tutorials | https://triton-lang.org/main/getting-started/tutorials/ |
| PyTorch User-Defined Triton | https://docs.pytorch.org/tutorials/recipes/torch_compile_user_defined_triton_kernel_tutorial.html |
| Triton Exercises | https://lweitkamp.github.io/triton_exercises/print.html |
| TK-GEMM (Llama3 FP8) | https://pytorch.org/blog/accelerating-llama3 |