Matmul backward in Triton — two tl.dot kernels

The CUDA track derived the rule: one forward GEMM becomes two backward GEMMs, dX = dY·Wᵀ and dW = Xᵀ·dY. This lesson is the Triton code. You will see that both are the lesson-09 tiled-matmul skeleton with two changes: swap the strides to read a transpose for free, and split the contraction for dW because its output can't fill the grid.

The question this lesson answers

You have a Triton forward GEMM Y = X·W. To train through it you must produce dX and dW from dY. Both name a transposed operand. Do you launch a transpose kernel first? (No.) And why does the dW kernel — same FLOPs as the others — show 20% SM occupancy in the profiler, and what's the Triton fix?

The two gradients, as shapes

With X ∈ ℝ^T×K (T = batch·tokens), W ∈ ℝ^K×N, Y ∈ ℝ^T×N, and upstream dY ∈ ℝ^T×N:

dX = dY · Wᵀ      shape T×K,  contracts over N
dW = Xᵀ · dY      shape K×N,  contracts over T (batch·tokens!)

Same FLOP count as the forward (2·T·K·N) for each. The difference that matters for the code is the contraction dimension and the transposed operand, summarised here so you can read the kernels against it:

The transpose is a stride swap, not a kernel

A matrix in memory is a base pointer plus two strides. "Transpose" means: when you build the pointer block, multiply the row index by the column stride and vice-versa. The bytes never move. For W ∈ ℝ^K×N stored row-major (stride_wk = N, stride_wn = 1), reading the (n, k) element of Wᵀ is just:

# W[k, n]  lives at  w_ptr + k*stride_wk + n*stride_wn
# Wᵀ[n, k] is the SAME element — index it with the strides swapped:
#   w_ptr + n*stride_wn + k*stride_wk
# In the dX kernel we want a (BN × BK) tile of Wᵀ, contracting over N:
w_block = w_ptr + offs_n[:, None]*stride_wn + (k + offs_k)[None, :]*stride_wk
#                  ^ N index on rows           ^ K index on cols
#                    (the row stride is the SMALL one — that's the "transpose")

Compare to the lesson-09 forward, which read W[k, n] with (k+offs_k)[:,None]*stride_wk + offs_n[None,:]*stride_wn. Same pointer, the two aranges have traded which stride they scale. That is the entire transpose.

Kernel 1 · dX = dY · Wᵀ

This is the lesson-09 skeleton verbatim, with A = dY (T×N), B = Wᵀ (N×K), output dX (T×K), contracting over N. Only the w load differs from forward — its strides are swapped.

@triton.jit
def dx_kernel(
    dy_ptr, w_ptr, dx_ptr,
    T, K, N,
    s_dy_t, s_dy_n,            # dY strides (T×N)
    s_w_k, s_w_n,              # W  strides (K×N) — NOT transposed in memory
    s_dx_t, s_dx_k,           # dX strides (T×K)
    BT: tl.constexpr, BK: tl.constexpr, BN: tl.constexpr,
):
    pid_t = tl.program_id(0)
    pid_k = tl.program_id(1)
    offs_t = pid_t*BT + tl.arange(0, BT)        # rows of dX  (T)
    offs_k = pid_k*BK + tl.arange(0, BK)        # cols of dX  (K)
    offs_n = tl.arange(0, BN)                    # contraction (N)

    acc = tl.zeros((BT, BK), dtype=tl.float32)   # fp32 accumulator, always
    for n in range(0, N, BN):
        nn = n + offs_n
        dy = tl.load(                            # dY tile  (BT × BN)
            dy_ptr + offs_t[:, None]*s_dy_t + nn[None, :]*s_dy_n,
            mask=(offs_t[:, None] < T) & (nn[None, :] < N), other=0.,
        )
        # Wᵀ tile (BN × BK): index N on the ROW, K on the COL —
        # i.e. multiply offs by s_w_n (rows) and s_w_k (cols). Stride swap = transpose.
        wt = tl.load(
            w_ptr + nn[:, None]*s_w_n + offs_k[None, :]*s_w_k,
            mask=(nn[:, None] < N) & (offs_k[None, :] < K), other=0.,
        )
        acc = tl.dot(dy, wt, acc)                # tensor cores fire

    tl.store(
        dx_ptr + offs_t[:, None]*s_dx_t + offs_k[None, :]*s_dx_k,
        acc.to(tl.bfloat16),
        mask=(offs_t[:, None] < T) & (offs_k[None, :] < K),
    )

Grid is (cdiv(T, BT), cdiv(K, BK)) — one program per output tile of dX. Because the output is T×K (T is huge), the grid is large and fills the GPU exactly like the forward. dX is the easy backward GEMM.

Kernel 2 · dW = Xᵀ · dY — and why it starves

Now A = Xᵀ (K×T), B = dY (T×N), output dW (K×N), contracting over T. The output is only K×N — the weight shape. Tiled at 128×128 that is ⌈K/128⌉·⌈N/128⌉ tiles. For K = 4096, N = 14336 that's 32·112 = 3584 tiles — fine. But shrink the layer (K = 1024, N = 1024 ⇒ 64 tiles) and you have fewer tiles than the ~132 SMs: most of the GPU is idle. Meanwhile the contraction T = tens of thousands is enormous. That mismatch is what split-K fixes.

The Triton move: add a third grid axis over K-splits. Each program owns one output tile and one slice of the T-contraction, computes a partial sum, and adds it into dW.

@triton.jit
def dw_kernel(
    x_ptr, dy_ptr, dw_ptr,
    T, K, N,
    s_x_t, s_x_k,             # X strides (T×K) — NOT transposed in memory
    s_dy_t, s_dy_n,           # dY strides (T×N)
    s_dw_k, s_dw_n,           # dW strides (K×N)
    SPLIT_T: tl.constexpr,     # number of T-chunks (the split-K factor)
    BK: tl.constexpr, BN: tl.constexpr, BT: tl.constexpr,
):
    pid_k = tl.program_id(0)
    pid_n = tl.program_id(1)
    pid_s = tl.program_id(2)              # which slice of T this program owns

    offs_k = pid_k*BK + tl.arange(0, BK)  # rows of dW   (K)
    offs_n = pid_n*BN + tl.arange(0, BN)  # cols of dW   (N)
    offs_t = tl.arange(0, BT)             # contraction within my T-slice

    # this program walks only its chunk of T:  [t0, t0 + chunk)
    chunk = tl.cdiv(T, SPLIT_T)
    t0 = pid_s * chunk
    t_end = tl.minimum(t0 + chunk, T)

    acc = tl.zeros((BK, BN), dtype=tl.float32)
    for t in range(t0, t_end, BT):
        tt = t + offs_t
        # Xᵀ tile (BK × BT): K on the ROW, T on the COL — strides swapped.
        xt = tl.load(
            x_ptr + offs_k[:, None]*s_x_k + tt[None, :]*s_x_t,
            mask=(offs_k[:, None] < K) & (tt[None, :] < t_end), other=0.,
        )
        dy = tl.load(                     # dY tile (BT × BN)
            dy_ptr + tt[:, None]*s_dy_t + offs_n[None, :]*s_dy_n,
            mask=(tt[:, None] < t_end) & (offs_n[None, :] < N), other=0.,
        )
        acc = tl.dot(xt, dy, acc)

    dw_block = dw_ptr + offs_k[:, None]*s_dw_k + offs_n[None, :]*s_dw_n
    mask = (offs_k[:, None] < K) & (offs_n[None, :] < N)
    if SPLIT_T == 1:
        tl.store(dw_block, acc.to(tl.float32), mask=mask)   # only writer, plain store
    else:
        tl.atomic_add(dw_block, acc, mask=mask)             # many writers per tile

Grid is (cdiv(K, BK), cdiv(N, BN), SPLIT_T). With SPLIT_T = 1 it is the plain GEMM; with SPLIT_T = 8 each output tile is computed by 8 programs that each cover one-eighth of T and atomic-add their partial in. You have turned 64 idle tiles into 512 programs — enough to fill the SMs.

Atomics buy occupancy, cost determinism

tl.atomic_add into the fp32 dW buffer is the cheap split-K reduction, but the partials land in nondeterministic order, and fp32 addition is not associative — so two runs give bit-different gradients. That breaks exact reproducibility. The deterministic alternative is a workspace: each split writes its partial to dw_partial[s] (shape SPLIT_T×K×N) with a plain store, then a second kernel reduces along the split axis:

# deterministic split-K: partials to workspace, then a fixed-order reduction
#   dw_kernel writes dw_partial[pid_s]  (no atomics)
@triton.jit
def dw_reduce_kernel(part_ptr, dw_ptr, SPLIT_T, KN, s_split, BLK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid*BLK + tl.arange(0, BLK)
    acc = tl.zeros((BLK,), dtype=tl.float32)
    for s in range(SPLIT_T):                      # fixed order ⇒ deterministic
        acc += tl.load(part_ptr + s*s_split + offs, mask=offs < KN, other=0.)
    tl.store(dw_ptr + offs, acc, mask=offs < KN)

One extra pass over a SPLIT_T×K×N workspace (HBM traffic, not FLOPs). Use atomics by default; switch to the workspace reduction only when you need bitwise reproducibility (debugging a divergence, or a determinism SLA).

Accumulate into the grad buffer, never overwrite

Weight gradients are accumulated, not assigned: across micro-batches (gradient accumulation) and across uses of a shared/tied weight, the contributions sum. So even SPLIT_T == 1 should read-modify-write when the buffer may already hold a partial:

# grad-accumulation friendly: add into whatever is already there.
prev = tl.load(dw_block, mask=mask, other=0.)
tl.store(dw_block, prev + acc, mask=mask)
# (with SPLIT_T > 1, tl.atomic_add already does += atomically.)

The buffer is fp32 even in a bf16 model: repeated accumulation into bf16 drifts. Zero it only at the optimizer-step boundary.

Worked numbers: one MLP up-projection, 7B-class

K = 4096 (d_model), N = 14336 (d_ff), T = 8192 (batch·tokens), bf16. H100: ~490 TFLOP/s sustained bf16, 3.35 TB/s HBM, ~132 SMs.

Three ~0.96 TFLOP GEMMs ≈ 2.9 TFLOP/layer/step → ~5.9 ms of GEMM at sustained rate. At this shape dW's 3584 tiles already exceed 132 SMs, so split-K isn't strictly needed — but the transpose layout conversion and the fp32 atomic epilogue still make dW the slowest of the three. Shrink the layer (attention's K=N=4096 q/k/v projections, or a small adapter) and dW's tile count drops below 132 → split-K becomes mandatory. The widget below lets you find that crossover.

The traps, collected

Interactive · find dW's split-K crossover

Set the layer shape and batch. The bars show the three GEMM FLOPs (all equal). The readout shows dW's output-tile count against the ~132 SMs, and the split-K factor needed to fill the grid. Shrink K and N (or grow T) to watch dW starve.

What's next

You can write both backward GEMMs, read a transpose as a stride swap, and split the dW contraction across the grid with the determinism trade-off in hand. Lesson 20 takes the matmul-into-a-loss case — the logits projection — where the forward GEMM and its backward must be fused, because materialising the [T, V] softmax and its gradient would blow the memory budget. Same two-GEMM rule, now with the softmax folded into the epilogue.