Writing Triton Kernels with Autograd: A Minimal End-to-End Example

This post is the second part of my Triton series. If you haven’t read the first one, I recommend starting here:

👉 Introducing Triton — and How It Challenges CUDA and torch.compile

In the previous article, we explored what Triton is and how it compares to CUDA and TorchInductor. In this post, we’ll go one level deeper: how to actually write a custom Triton operator that integrates cleanly with PyTorch autograd.

By the end, you’ll understand a reusable pattern:

To implement a custom Triton op in PyTorch, you typically write:

• A Triton forward kernel
• A Triton backward kernel
• A torch.autograd.Function wrapper that ties everything together

We’ll walk through a minimal but representative example: a weighted sum operator.

1. Problem Setup

We define a simple operation:

y = x.dot(w)

Where:

• x: shape [N, D]
• w: shape [D]
• output y: shape [N]

In PyTorch, this is trivial:

def weighted_sum_torch(x, w):
    return (x * w).sum(dim=-1)

But this hides all the low-level execution details. Our goal is to reimplement this using Triton.

2. Triton Forward Kernel

Source: stanford.github.com

In Triton, we explicitly control:

• how work is partitioned
• how memory is accessed
• how computation is accumulated

A common pattern is:

One program instance processes a tile of rows, and loops over the feature dimension D.

import triton
import triton.language as tl

@triton.jit
def weighted_sum_fwd(
    x_ptr, w_ptr, out_ptr,
    stride_x_row, stride_x_d,
    stride_w_d,
    stride_out_row,
    N, D,
    BLOCK_N: tl.constexpr,
    BLOCK_D: tl.constexpr,
):
    pid = tl.program_id(0)
    row_offsets = pid * BLOCK_N + tl.arange(0, BLOCK_N)
    acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
    for d_start in range(0, D, BLOCK_D):
        d_offsets = d_start + tl.arange(0, BLOCK_D)
        x = tl.load(
            x_ptr + row_offsets[:, None] * stride_x_row + d_offsets[None, :] * stride_x_d,
            mask=(row_offsets[:, None] < N) & (d_offsets[None, :] < D),
            other=0.0,
        )
        w = tl.load(
            w_ptr + d_offsets * stride_w_d,
            mask=d_offsets < D,
            other=0.0,
        )
        acc += tl.sum(x * w[None, :], axis=1)
    tl.store(
        out_ptr + row_offsets * stride_out_row,
        acc,
        mask=row_offsets < N,
    )

Key ideas

• tl.program_id(0) identifies which tile this instance handles
• We loop over D because it may not fit in one block
• Computation is done as block-wise multiply + reduction

3. Backward: Deriving Gradients First

Before writing any Triton code, always derive gradients mathematically.

Using chain rule, we get:

There are two outputs:

• grad_x (same shape as x)
• grad_w (needs reduction across rows)

In PyTorch:

grad_x = grad_out[:, None] * w[None, :]
grad_w = (x * grad_out[:, None]).sum(dim=0)

In Triton, computing the gradient ∇x is simple in matrix form. One subtle but important detail in the backward implementation is how we compute the gradient with respect to the weights (grad_w).

This is a reduction over the row dimension, which introduces a challenge when parallelizing across multiple Triton program instances.

Instead of performing a full reduction inside a single kernel (which would require synchronization or atomic operations), we adopt a two-stage strategy:

1. In the Triton kernel

Each program instance computes a partial reduction over the rows it owns:

partial = tl.sum(x * grad_out[:, None], axis=0)

These partial results are written to a temporary buffer. For example, suppose:

ROWS = 100
BLOCK_N = 32
D = 64
BLOCK_D = 64
n_row_tiles = ceil(100 / 32) = 4

Then the intermediate tensor has shape:

partial.shape = [4, 64]

Compared to the original (100, 64) tensor, this reduces the data to one row per tile, significantly lowering the amount of data that must be stored and later reduced.

2. Outside the kernel (in PyTorch) We then perform a final reduction across all program instances:

grad_w = partial_grad_w.sum(dim=0)

which produces:

grad_w.shape = [64]

This pattern — performing a partial reduction within each kernel instance, followed by a final reduction outside the kernel — is common in GPU programming.

It avoids expensive synchronization or atomic operations inside the kernel, keeps each program instance simple and efficient, and still yields the correct global result.

The complete implementation is provided in Sections 4 (Step 1) and 5 (Step 2), corresponding to the two stages described above.

4. Triton Backward Kernel

@triton.jit
def weighted_sum_bwd(
    x_ptr, w_ptr, grad_out_ptr,
    grad_x_ptr, partial_w_ptr,
    stride_x_row, stride_x_d,
    stride_w_d,
    stride_go_row,
    stride_gx_row, stride_gx_d,
    stride_pw_row, stride_pw_d,
    N, D,
    BLOCK_N: tl.constexpr,
    BLOCK_D: tl.constexpr,
):
    pid = tl.program_id(0)
    row_offsets = pid * BLOCK_N + tl.arange(0, BLOCK_N)
    for d_start in range(0, D, BLOCK_D):
        d_offsets = d_start + tl.arange(0, BLOCK_D)
        grad_out = tl.load(
            grad_out_ptr + row_offsets * stride_go_row,
            mask=row_offsets < N,
            other=0.0,
        )
        w = tl.load(
            w_ptr + d_offsets * stride_w_d,
            mask=d_offsets < D,
            other=0.0,
        )
        # grad_x
        grad_x = grad_out[:, None] * w[None, :]
        tl.store(
            grad_x_ptr + row_offsets[:, None] * stride_gx_row + d_offsets[None, :] * stride_gx_d,
            grad_x,
            mask=(row_offsets[:, None] < N) & (d_offsets[None, :] < D),
        )
        x = tl.load(
            x_ptr + row_offsets[:, None] * stride_x_row + d_offsets[None, :] * stride_x_d,
            mask=(row_offsets[:, None] < N) & (d_offsets[None, :] < D),
            other=0.0,
        )
        # partial grad_w
        partial = tl.sum(x * grad_out[:, None], axis=0)
        tl.store(
            partial_w_ptr + pid * stride_pw_row + d_offsets * stride_pw_d,
            partial,
            mask=d_offsets < D,
        )

5. Wrapping with torch.autograd.Function

This is what connects Triton to PyTorch.

import torch

class WeightedSumFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, w):
        N, D = x.shape
        BLOCK_N = 32
        BLOCK_D = 128
        y = torch.empty((N,), device=x.device, dtype=x.dtype)
        grid = (triton.cdiv(N, BLOCK_N),)
        weighted_sum_fwd[grid](
            x, w, y,
            x.stride(0), x.stride(1),
            w.stride(0),
            y.stride(0),
            N, D,
            BLOCK_N, BLOCK_D,
        )
        ctx.save_for_backward(x, w)
        ctx.BLOCK_N = BLOCK_N
        ctx.BLOCK_D = BLOCK_D
        return y
    @staticmethod
    def backward(ctx, grad_out):
        x, w = ctx.saved_tensors
        N, D = x.shape
        BLOCK_N = ctx.BLOCK_N
        BLOCK_D = ctx.BLOCK_D
        grad_x = torch.empty_like(x)
        grid = (triton.cdiv(N, BLOCK_N),)
        partial_w = torch.zeros((grid[0], D), device=x.device, dtype=x.dtype)
        weighted_sum_bwd[grid](
            x, w, grad_out,
            grad_x, partial_w,
            x.stride(0), x.stride(1),
            w.stride(0),
            grad_out.stride(0),
            grad_x.stride(0), grad_x.stride(1),
            partial_w.stride(0), partial_w.stride(1),
            N, D,
            BLOCK_N, BLOCK_D,
        )
        grad_w = partial_w.sum(dim=0)
        return grad_x, grad_w

Usage:

def weighted_sum(x, w):
    return WeightedSumFunction.apply(x, w)

Now, calling weighted_sum on two PyTorch tensors x and w produces an output like:

tensor([ 0.85, -3.68, ...,-4.0192],
       device='cuda:0', grad_fn=<WeightedSumFunctionBackward>)

Notice the grad_fn attached to the tensor. PyTorch has successfully registered our custom operation in the computation graph and knows exactly which function to invoke during the backward pass.

In other words, our Triton kernel is no longer just a standalone GPU routine — it is now fully integrated into PyTorch’s autograd system.

With this, we’ve completed a full end-to-end implementation of the weighted sum operator in Triton: forward kernel, backward kernel, and seamless autograd support.

6. The Reusable Pattern

This example illustrates a general workflow you can reuse:

Step 1 — Define math clearly

Always start from the formula and derive gradients.

Step 2 — Write forward kernel

• Tile over rows
• Loop over reduction dimension
• Accumulate locally

Step 3 — Write backward kernel

• Compute elementwise gradients (grad_x)
• Handle reductions (grad_w) via partial sums

Step 4 — Wrap with autograd

• Save tensors in ctx
• Launch kernels in forward/backward
• Return gradients in correct order

Q: When is it unnecessary to implement a backward function?

In general, a custom backward implementation is only required when gradients need to be computed for training. If the operation is used purely for inference, where no gradient computation is needed, implementing a backward function is unnecessary. In such cases, the forward kernel can be executed under a torch.no_grad() context, which disables autograd and avoids any gradient-related overhead.

Final Thoughts

Triton sits in a very interesting space:

• More control than PyTorch ops
• Less complexity than CUDA
• Fully compatible with autograd

Once you understand this forward + backward + autograd wrapper pattern, you can implement a wide range of custom GPU operators.