/triton-kernels

| name | triton-kernels | | --- | --- | | description | Provides guidance for writing and benchmarking portable Triton kernels targeting NVIDIA and AMD GPUs. Covers core DSL patterns, @triton.autotune, numerics (fp16/bf16/fp8), masked loads, reductions, tiling, benchmarking harness, correctness testing, and integration with HuggingFace Kernels Hub (get_kernel). Vendor-neutral: points to rocm-kernels and xpu-kernels for backend-specific tuning. | | disable-model-invocation | false | | user-invocable | true | | allowed-tools | Read, Grep, Glob, Bash | | argument-hint | kernel type: softmax, matmul, rmsnorm, layernorm, activation, reduction, element-wise, autotune, benchmark, correctness, get_kernel, transformers, diffusers |

Portable Triton Kernels

This skill provides patterns and guidance for developing portable, optimized Triton kernels that run on NVIDIA and AMD GPUs without modification. For backend-specific tuning, see [rocm-kernels](../rocm-kernels/SKILL.md) (AMD) and [xpu-kernels](../xpu-kernels/SKILL.md) (Intel).

When This Skill Applies

Use this skill when:

• Writing new Triton kernels for normalization, activation, attention, or linear algebra ops
• Deciding block sizes, num_warps, num_stages, and autotune configs
• Handling numerics (fp32 accumulation, bf16/fp16 input/output, masked values)
• Setting up correctness tests against a PyTorch reference
• Benchmarking kernel throughput (GB/s or TFLOPS)
• Publishing a Triton kernel to the HuggingFace Kernels Hub via get_kernel
• Fusing multiple ops into a single kernel to reduce DRAM round-trips

Hard Constraints

1. **BLOCK_SIZE for reductions must cover the full reduction dimension.** Use `triton.next_power_of_2(dim)` in the Python wrapper. Never autotune BLOCK_SIZE when it controls the reduction axis — partial rows give wrong results silently.

2. **Masked loads need a safe `other` value.** Use `other=0.0` for additive contexts (sum, dot product). Use `other=float('-inf')` for max-based reductions (softmax numerator). Using the wrong fill value is the #1 cause of subtle numerical bugs.

3. **Accumulate in fp32.** Cast inputs to `tl.float32` before reductions. Cast back to the input dtype only at the final `tl.store`. Half-precision accumulation compounds rounding errors across hundreds of additions.

4. **All tensors must be contiguous.** Assert `x.is_contiguous()` in the Python wrapper, or call `.contiguous()` before computing pointers. Non-contiguous tensors break flat-offset pointer arithmetic.

5. **tl.constexpr parameters must be known at compile time.** BLOCK_SIZE, num_warps, num_stages are compile-time constants. Pass them as kernel arguments with the `: tl.constexpr` annotation or via autotune configs.

Core DSL Patterns

Program IDs and Grid Launch

Every Triton kernel is launched as a grid of programs. Each program gets a unique ID via `tl.program_id(axis)`.

@triton.jit
def my_kernel(x_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

    x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
    # ... compute ...
    tl.store(output_ptr + offsets, result, mask=mask)

Grid sizing:

grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
my_kernel[grid](x, output, n_elements, BLOCK_SIZE=1024)

For 2D grids (e.g. matmul):

grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
# program_id(0) = row block, program_id(1) = col block

Masked Loads and Stores

Always mask when BLOCK_SIZE may exceed the actual dimension:

offsets = tl.arange(0, BLOCK_SIZE)
mask = offsets < n_cols
x = tl.load(ptr + offsets, mask=mask, other=0.0)
tl.store(out_ptr + offsets, result, mask=mask)

2D Pointer Arithmetic (Tiling)

For loading 2D tiles (matmul, attention):

row_offsets = row_start * BLOCK_M + tl.arange(0, BLOCK_M)
col_offsets = col_start * BLOCK_N + tl.arange(0, BLOCK_N)

# Broadcasting: [:, None] makes column vector, [None, :] makes row vector
ptrs = base_ptr + row_offsets[:, None] * stride_row + col_offsets[None, :]
mask = (row_offsets[:, None] < M) & (col_offsets[None, :] < N)
tile = tl.load(ptrs, mask=mask, other=0.0)

Reductions

Row-wise reduction (softmax, layernorm, rmsnorm):

# One program per row. BLOCK_SIZE >= n_cols (next power of 2).
pid = tl.program_id(0)
row_start = pid * stride_row
offsets = tl.arange(0, BLOCK_SIZE)
mask = offsets < n_cols

x = tl.load(x_ptr + row_start + offsets, mask=mask, other=0.0).to(tl.float32)
row_sum = tl.sum(x, axis=0)
row_max = tl.max(x, axis=0)

The dot Product (Tile-Level Matmul)

`tl.dot(a, b)` performs a tile-level matrix multiply that maps to tensor cores. Requires both operands to have their K dimension >= 16.

# a_tile: (BLOCK_M, BLOCK_K), b_tile: (BLOCK_K, BLOCK_N)
acc += tl.dot(a_tile, b_tile)  # acc: (BLOCK_M, BLOCK_N)

Autotune

`@triton.autotune` benchmarks multiple kernel configurations at runtime and caches the fastest one per problem shape.

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64}, num_stages=3, num_warps=8),
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 32}, num_stages=4, num_warps=4),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32}, num_stages=4, num_warps=4),
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 32}, num_stages=5, num_warps=2),
    ],
    key=['M', 'N', 'K'],  # re-tune when these values change
)
@triton.jit
def matmul_kernel(..., BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    ...

Guidelines:

• **key=** lists the runtime values that affect which config is fastest. Usually

the matrix dimensions.

• **num_warps**: More warps = more threads per program. Large tiles (128x256) need

8-16 warps. Small