triton-kernel

---
name: triton-kernel
description: Write optimized Triton GPU kernels for deep learning operations. Covers the full spectrum from basic vector ops to Flash Attention, persistent matmul, fused normalization, quantized GEMM, and memory-efficient patterns.
user-invocable: true
---

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

## Performance Bottleneck Quick-Reference

When optimizing an existing kernel, classify the bottleneck first (profile with `ncu`):

| Bottleneck | Diagnosis | Fix |
|------------|-----------|-----|
| **Memory-bound** | DRAM throughput > 60% of peak, compute < 30% | PID swizzle, TMA, fuse ops to reduce loads |
| **Compute-bound** | Tensor core utilization > 60%, DRAM < 40% | Persistent kernels, increase `num_stages`, warp specialization |
| **Underutilized** | Both < 60%, high stall metrics | Reduce register pressure, increase `num_warps`, autotune |

See `triton-gpu-kernel-optimization.md` for specific NCU metric names and detailed strategies.

## Specialized Topics

Read these files for detailed guidance when the task involves these areas:

| Task | File to read |
|------|-------------|
| Flash Attention / fused self-attention | `triton-flash-attention-v2.md` |
| Persistent kernels, warp specialization, TMA | `triton-persistent-warp-matmul.md` |
| LayerNorm, RMSNorm, GroupNorm (fwd + bwd) | `triton-fused-normalizations.md` |
| FP4/FP8 quantized matmul, block scaling | `triton-quantized-block-scaled-gemm.md` |
| Kernel fusion, Philox dropout, recomputation | `triton-memory-efficient-patterns.md` |
| General tiled GEMM, autotune, benchmarking | `triton-gpu-kernel-optimization.md` |
| Fusing normalization/gating/residual into attention or matmul epilogue | `triton-fused-epilogue-kernels.md` |
| Sequential stateful processing (LRU routing, mutable register state) | `triton-sequential-stateful-blocks.md` |
| Launcher tile selection, num_stages/num_warps heuristics | `triton-dynamic-launcher-tiling.md` |

**When to read specialized files:** Only read the relevant file when the user's task specifically involves that topic. The core patterns above are sufficient for basic kernels (vector ops, elementwise fusion, simple reductions).

## Other references

- `triton-opt.md`: For general optimization techniques while writing triton kernels.

More agent context in vipshop/cache-dit

6 other files this repository gives its agents.