Triton Lang

Triton

Contract

Inputs

1. Operation (required): the math the kernel should perform, or the kernel under review.
2. Triton version (required): python -c "import triton; print(triton.__version__)". Grounded current stable: 3.7.1.
3. Shapes and dtypes (required): tensor shapes, strides, and dtypes for the launch grid and masks.
4. PyTorch version (optional, for integration): custom-op registration needs a recent release.

Procedure

1. Write the kernel blockwise. tl.program_id indexes the block, tl.arange spans the block, and a mask covers the tail:

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

BLOCK: tl.constexpr is compile-time, so the compiler unrolls and sizes shared memory from it. Done when: every load and store carries the mask. 2. Give masked loads a fill value. A masked lane without other reads undefined; other=0.0 keeps reductions clean:

x = tl.load(x_ptr + offsets, mask=mask, other=0.0)

Done when: every masked load that feeds arithmetic names its other. 3. Use atomics sparingly. They serialize updates to one address; reduce within the block first, then issue one atomic per block:

@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)
    tl.atomic_add(out_ptr, tl.sum(x, axis=0))

Done when: atomic count per output element is one or the contention is justified. 4. Autotune block shape and warp count against the input size. The tuner benchmarks each config once per distinct key value and caches the winner:

@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):
    ...

Done when: the config set spans the plausible block range and key names every size-dependent argument. 5. Benchmark with the documented harness. do_bench returns milliseconds for a quick read; perf_report with Benchmark sweeps sizes and plots providers:

import triton.testing as tt

ms = tt.do_bench(lambda: add(x, y))
print(f"{ms:.3f} ms")

from triton.testing import perf_report, Benchmark

@perf_report(
    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={},
    )
)
def bench_add(n, provider):
    x = torch.randn(n, device="cuda")
    y = torch.randn(n, device="cuda")
    if provider == "triton":
        return lambda: add(x, y)
    return lambda: x + y

Compare against the PyTorch reference on small inputs before scaling up. Done when: the kernel beats or matches the reference at the shapes that matter, with quantiles recorded. 6. Register the kernel as a PyTorch custom op so torch.compile can trace it. The fake kernel needs no body, only the output shape:

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)

Keep Python side effects out of the JIT function; they break tracing. Done when: the op runs eagerly and under torch.compile. 7. Debug in layers. Compare against a PyTorch reference first. Then gate with tl.debug_barrier() to isolate sync bugs, print autotune decisions, and inspect compiled artifacts:

TRITON_PRINT_AUTOTUNING=1 python script.py
export TRITON_CACHE_DIR=/tmp/triton_cache   # inspect compiled kernels

Done when: the failing stage, kernel or integration, is named.

Failure and recovery

Output

The kernel with its mask and grid logic, the autotune or benchmark plan with measured numbers, the integration snippet, and the debug steps taken. Each claim names the Triton version it assumes.