Writing Custom Triton Kernels for PyTorch (Without Breaking torch.compile)

Most PyTorch performance work never needs a custom kernel. torch.compile fuses pointwise chains, BF16 doubles matmul throughput, and a fixed data loader beats a hand-written kernel every time. But there is a real class of workloads where the bottleneck is an operation PyTorch does not have: a fused loss with an unusual reduction, a custom attention variant, a quantized dequant-matmul pattern, a sparse gather that eager mode turns into six kernel launches and three temporaries. That is where Triton earns its place.

This tutorial shows the whole path: proving you need a kernel, writing one, testing it for correctness, benchmarking it honestly, and registering it so torch.compile and autograd treat it as a first-class operator instead of a graph break.

Step 0: prove the kernel is the bottleneck

Do not start here. Start with a profile. Run the region under torch.profiler and look at the kernel table:

from torch.profiler import profile, ProfilerActivity

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    for _ in range(10):
        loss = model(batch).sum()
        loss.backward()
        torch.cuda.synchronize()

print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=20))

You are looking for one of three signatures:

• Many tiny kernels. Dozens of elementwise_kernel launches around one real operation. This is fusion work, and torch.compile usually solves it for free. Try that first.
• One kernel dominating that is memory-bound. Compute its achieved bandwidth: bytes moved divided by kernel time, compared against your GPU's HBM bandwidth. If you are at 20% of peak on a memory-bound op, there is headroom.
• An operation PyTorch materializes a huge intermediate for. The classic case: a [batch, seq, vocab] logits tensor that only exists to be reduced into a scalar loss. Fusing the reduction into the matmul removes gigabytes of traffic and is often the difference between OOM and not.

If your profile shows none of these, stop. Write a kernel only when the arithmetic says a kernel can win. Our post on profiling PyTorch training covers the measurement discipline in more depth.

Why Triton rather than CUDA C++

Triton is a Python-embedded DSL that compiles to GPU code. You write per-block programs over tensors; the compiler handles thread assignment, vectorization, shared-memory staging and (largely) memory coalescing. In exchange for giving up some low-level control, you get kernels that are 10-30 lines instead of 300, readable by the ML engineers who own the model, and portable across NVIDIA and (increasingly) AMD hardware.

It is also what torch.compile already emits: the Inductor backend generates Triton for GPU fusions. Writing Triton by hand means working in the same language your compiler works in, which makes the two compose well.

Reach for CUDA C++ instead when you need warp-level primitives Triton does not expose, when you are chasing the last 10% on a matmul that cuBLAS already does well, or when you must integrate with an existing C++ codebase.

A first kernel: fused RMSNorm

RMSNorm is a good example. In eager PyTorch it is a square, a mean, a rsqrt, a multiply and a weight multiply — five kernels reading and writing the same activation tensor. It is entirely memory-bound, so collapsing it to one read and one write is close to a 4-5x win on the op.

import torch
import triton
import triton.language as tl


@triton.jit
def _rmsnorm_fwd(x_ptr, w_ptr, out_ptr, rstd_ptr,
                 stride_row, n_cols, eps,
                 BLOCK_SIZE: tl.constexpr):
    row = tl.program_id(0)
    x_ptr += row * stride_row
    out_ptr += row * stride_row

    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < n_cols
    x = tl.load(x_ptr + cols, mask=mask, other=0.0).to(tl.float32)

    mean_sq = tl.sum(x * x, axis=0) / n_cols
    rstd = 1.0 / tl.sqrt(mean_sq + eps)
    tl.store(rstd_ptr + row, rstd)

    w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    y = x * rstd * w
    tl.store(out_ptr + cols, y.to(tl.float32), mask=mask)


def rmsnorm_fwd(x, weight, eps=1e-6):
    x2 = x.reshape(-1, x.shape[-1]).contiguous()
    n_rows, n_cols = x2.shape
    out = torch.empty_like(x2)
    rstd = torch.empty(n_rows, dtype=torch.float32, device=x.device)
    BLOCK_SIZE = triton.next_power_of_2(n_cols)
    _rmsnorm_fwd[(n_rows,)](
        x2, weight, out, rstd,
        x2.stride(0), n_cols, eps,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=8 if BLOCK_SIZE >= 4096 else 4,
    )
    return out.reshape(x.shape), rstd

Four things in that kernel are the whole idiom, and they recur in every Triton kernel you will write:

1. tl.program_id(0) is your block index. One program handles one row here; the grid (n_rows,) launches one program per row.
2. tl.arange plus a mask is how you handle non-power-of-two sizes. Loads outside the range return other, stores are dropped.
3. Accumulate reductions in FP32 even when the tensor is BF16. Reductions in low precision are where numerical differences against eager PyTorch come from.
4. num_warps and BLOCK_SIZE are the tuning knobs. Start with the heuristic above, then autotune.

For sizes that do not fit one block (a hidden dimension of 32k, say) you need a two-pass kernel that loops over column tiles. Do not silently allocate a BLOCK_SIZE bigger than shared memory allows; assert on it.

Autotuning

Rather than guessing block sizes, let Triton search:

@triton.autotune(
    configs=[
        triton.Config({"BLOCK_SIZE": bs}, num_warps=w)
        for bs in (1024, 2048, 4096, 8192)
        for w in (2, 4, 8, 16)
    ],
    key=["n_cols"],
)
@triton.jit
def _rmsnorm_fwd(...):
    ...

The key argument controls cache invalidation: a new n_cols triggers a fresh search, the same one reuses the winner. Autotuning costs seconds on the first call of each shape, so warm it up outside your timing loop and be aware of it in short-lived serving processes. Persist the cache (TRITON_CACHE_DIR) in your container image if cold-start latency matters.

Correctness before speed

A fast wrong kernel is a very expensive bug, because it shows up as a slightly worse eval score three weeks later rather than as a crash. Test before you benchmark:

import pytest

@pytest.mark.parametrize("shape", [(4, 512), (17, 4096), (2, 8, 1023)])
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
def test_rmsnorm_matches_eager(shape, dtype):
    x = torch.randn(*shape, device="cuda", dtype=dtype)
    w = torch.randn(shape[-1], device="cuda", dtype=dtype)

    ref = torch.nn.functional.rms_norm(x, (shape[-1],), w, eps=1e-6)
    out, _ = rmsnorm_fwd(x, w, eps=1e-6)

    torch.testing.assert_close(out.to(ref.dtype), ref, rtol=1e-2, atol=1e-2)

Include awkward shapes: non-power-of-two columns, a single row, a batch of one, and a non-contiguous input (then assert or .contiguous() explicitly — a kernel that assumes contiguity and does not check it will read garbage). If the op has a backward pass, check it with torch.autograd.gradcheck in double precision on a small input, and compare gradients against the eager implementation on realistic inputs.

Benchmark honestly

time.time() around a kernel launch measures nothing; CUDA is asynchronous. Use Triton's benchmark helper, which handles warmup, L2 cache flushing between reps and quantiles:

ms = triton.testing.do_bench(lambda: rmsnorm_fwd(x, w), warmup=25, rep=100)
gbps = (2 * x.numel() * x.element_size()) / (ms * 1e-3) / 1e9
print(f"{ms:.3f} ms, {gbps:.0f} GB/s")

Compare that GB/s against your GPU's spec bandwidth. A memory-bound kernel above 80% of peak is done; there is nothing left to win. Compare against three baselines, not one: eager PyTorch, torch.compiled eager PyTorch, and any library kernel that already exists. It is common and slightly deflating to discover that torch.compile already generates something within 5% of your hand-written kernel — which is a good outcome, because it means you get to delete code.

Making it a real PyTorch operator

This is the step teams skip, and it is what separates a demo from something you can put in a model. A raw Python function calling a Triton kernel causes a graph break under torch.compile, has no autograd support, and will not work with torch.export, ExecuTorch or FSDP2's meta-device init. Register it with torch.library instead:

@torch.library.custom_op("mylib::rmsnorm", mutates_args=())
def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
    out, _ = rmsnorm_fwd(x, weight, eps)
    return out


@rmsnorm.register_fake
def _(x, weight, eps):
    return torch.empty_like(x)


def _backward(ctx, grad):
    x, weight = ctx.saved_tensors
    return rmsnorm_bwd(grad, x, weight, ctx.eps), None, None


def _setup_context(ctx, inputs, output):
    x, weight, eps = inputs
    ctx.save_for_backward(x, weight)
    ctx.eps = eps


torch.library.register_autograd("mylib::rmsnorm", _backward, setup_context=_setup_context)

The fake (meta) implementation is what lets Dynamo trace shapes without running the kernel; without it torch.compile cannot reason about your op. mutates_args tells the compiler whether the op writes to its inputs — get this wrong and you will get silent reordering bugs. Finally, verify the registration rather than assuming it:

torch.library.opcheck(torch.ops.mylib.rmsnorm, (x, w, 1e-6))

compiled = torch.compile(model, fullgraph=True)  # fails loudly on any graph break

fullgraph=True in CI is the cheapest guard against a future refactor silently reintroducing a break. See torch.compile in practice for how the rest of the graph behaves around custom ops.

When not to write the kernel

Before committing an engineer to this, check that the work has not already been done:

• Attention variants — torch.nn.attention.flex_attention expresses masks, ALiBi, sliding windows and custom score modifications as a Python function and compiles them into a fused kernel. It covers most of what people used to hand-write.
• Fused losses, RMSNorm, SwiGLU, RoPE, cross-entropy — Liger Kernel ships tested Triton implementations for the standard transformer stack.
• Quantized matmuls — torchao already has INT8, INT4 and FP8 paths that are maintained and benchmarked.
• torch.compile itself — for pointwise and reduction chains, try it before writing anything.

A custom kernel is a permanent maintenance liability: it must be re-benchmarked on every new GPU generation, re-tested on every PyTorch upgrade, and understood by whoever inherits the model. Write one when the profile justifies it and the alternatives do not exist, then wrap it in tests so the next person can trust it.

A checklist

1. Profile; confirm the target op is a real, measurable share of step time.
2. Compute the roofline bound so you know what winning looks like.
3. Try torch.compile, FlexAttention or an existing library kernel first.
4. Write the forward kernel; accumulate reductions in FP32.
5. Test against eager on awkward shapes and both dtypes before timing anything.
6. Autotune, then benchmark with do_bench against all three baselines.
7. Write the backward, check it with gradcheck.
8. Register with torch.library, add a fake implementation, run opcheck.
9. Add a fullgraph=True compile test and a performance regression test to CI.
10. Record the GPU, PyTorch and Triton versions the numbers came from.

Kernel-level optimization is one of the higher-leverage things a PyTorch team can do, and one of the easiest to do badly. If you have a training or inference workload that looks kernel-bound, our PyTorch performance optimization and LLM inference optimization practices do this work end to end, including the profiling that tells you whether it is worth doing at all. Contact us with your profile trace and we will tell you what we see.