triton-lang

name	triton-lang
description	Triton language skill for Python GPU kernel authoring. Use when writing Triton kernels with @triton.jit, tl.load/store, masking, atomics, benchmarking with triton.testing, or integrating kernels into PyTorch. Activates on queries about Triton, tl.constexpr, block pointers, Triton benchmarking, or PyTorch custom ops.

Triton

Purpose

Guide agents through writing GPU kernels in OpenAI Triton: the @triton.jit decorator, block-oriented tl.load/tl.store with masking, atomic operations, shared memory via tl.constexpr, benchmarking with triton.testing.Benchmark, PyTorch integration, and debugging with barriers.

When to Use

• Writing custom PyTorch ops faster than pure PyTorch but without raw CUDA
• Prototyping fused kernels (e.g., softmax + scale + bias)
• Comparing block sizes and warp counts with Triton's autotuner
• Porting NumPy-style elementwise ops to GPU
• Learning GPU programming with higher-level Python syntax
• Benchmarking kernel variants systematically

Workflow

1. Minimal Triton kernel

import torch
import triton
import triton.language as tl

@triton.jit
def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offsets < n
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)

def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    n = x.numel()
    out = torch.empty_like(x)
    grid = lambda meta: (triton.cdiv(n, meta["BLOCK"]),)
    add_kernel[grid](x, y, out, n, BLOCK=1024)
    return out

Key concepts:

• tl.program_id(0) — block index (like blockIdx.x)
• tl.arange(0, BLOCK) — vector of thread indices within block
• mask — predication for tail elements (no separate bounds kernel)
• BLOCK: tl.constexpr — compile-time constant, enables unrolling

2. Load/store and masking

@triton.jit
def masked_load_example(ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    # masked load returns 0 for masked-off lanes
    vals = tl.load(ptr + offs, mask=mask, other=0.0)
    return vals

Block pointers (Triton 2.x+) for structured 2D access:

@triton.jit
def matvec_kernel(a_ptr, x_ptr, y_ptr, M, N, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
    pid_m = tl.program_id(0)
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    acc = tl.zeros((BLOCK_M,), dtype=tl.float32)
    for start_n in range(0, N, BLOCK_N):
        offs_n = start_n + tl.arange(0, BLOCK_N)
        a = tl.load(a_ptr + offs_m[:, None] * N + offs_n[None, :])
        x = tl.load(x_ptr + offs_n)
        acc += tl.sum(a * x[None, :], axis=1)
    tl.store(y_ptr + offs_m, acc, mask=offs_m < M)

3. Atomic operations

@triton.jit
def atomic_histogram(data_ptr, hist_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    data = tl.load(data_ptr + offs, mask=mask)
    bucket = (data % 256).to(tl.int32)
    tl.atomic_add(hist_ptr + bucket, 1, mask=mask)

Use atomics sparingly — they serialize memory updates. Prefer block-level reduction then single atomic per block.

4. Shared memory via constexpr

@triton.jit
def reduce_kernel(x_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask, other=0.0)
    # Block reduction
    x = tl.sum(x, axis=0)
    tl.atomic_add(out_ptr, x)

BLOCK as tl.constexpr lets the compiler allocate shared memory and unroll loops at compile time.

5. Autotuning

@triton.autotune(
    configs=[
        triton.Config({"BLOCK": 128}, num_warps=4),
        triton.Config({"BLOCK": 256}, num_warps=4),
        triton.Config({"BLOCK": 512}, num_warps=8),
    ],
    key=["n"],
)
@triton.jit
def tuned_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    # ... kernel body ...
    pass

Autotuner benchmarks configs on first run and caches the best for each key shape.

6. Benchmarking

from triton.testing import Benchmark

def benchmark_add():
    n = 1024 * 1024
    x = torch.randn(n, device="cuda")
    y = torch.randn(n, device="cuda")

    def triton_add():
        return add(x, y)

    def torch_add():
        return x + y

    bench = Benchmark(
        x_names=["n"],
        x_vals=[2**i for i in range(10, 24)],
        line_arg="provider",
        line_vals=["triton", "torch"],
        line_names=["Triton", "PyTorch"],
        plot_name="add-bench",
        args={},
    )
    bench.run(lambda n, provider: {
        "triton": lambda: add(x[:n], y[:n]),
        "torch": lambda: x[:n] + y[:n],
    }[provider](), quantiles=[0.5, 0.9])

# Quick timing in REPL
import triton.testing as tt
ms = tt.do_bench(lambda: add(x, y))
print(f"{ms:.3f} ms")

7. PyTorch integration

import torch
from torch.library import custom_op

@custom_op("mylib::triton_add", mutates_args=())
def triton_add_op(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    return add(x, y)

@triton_add_op.register_fake
def _(x, y):
    return torch.empty_like(x)

# Use in model
class MyModule(torch.nn.Module):
    def forward(self, x, y):
        return triton_add_op(x, y)

For torch.compile compatibility, register fake/meta kernels and avoid Python side effects in the JIT function.

8. Debugging

@triton.jit
def debug_kernel(x_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < n
    x = tl.load(x_ptr + offs, mask=mask)
    # Synchronize threads within block for inspection
    tl.debug_barrier()
    tl.store(x_ptr + offs, x * 2.0, mask=mask)

# Dump generated PTX/LLVM IR
TRITON_PRINT_AUTOTUNING=1 python script.py
# Set cache dir to inspect compiled kernels
export TRITON_CACHE_DIR=/tmp/triton_cache

Compare against PyTorch reference on small inputs before scaling up.

Common Problems

Symptom	Cause	Fix
OutOfResources	BLOCK too large	Reduce BLOCK or num_warps
Wrong results on tail	Missing mask	Add mask=offs < n to load/store
Slower than PyTorch	Suboptimal BLOCK	Use @triton.autotune
CompilationError	Type mismatch	Ensure consistent dtypes; use .to(tl.float32)
NaN in output	Uninitialized masked lanes	Pass other=0.0 to masked loads
torch.compile fails	No fake kernel	Register register_fake meta function

Related Skills

• skills/gpu/cuda — underlying CUDA concepts when Triton limits are hit
• skills/gpu/cuda-profiling — Nsight profiling of Triton-compiled kernels
• skills/gpu/gpu-memory-model — coalescing and occupancy theory
• skills/gpu/hip-rocm — AMD GPU path (Triton supports ROCm)
• skills/hpc/openmp — CPU parallelism alongside GPU kernels