Fused cross-entropy in Triton — the memory kernel

The logits tensor is the biggest thing a language model ever holds, and the naive loss path keeps three copies of it. This lesson is the Triton mechanics of making all three vanish: an @triton.jit kernel that online-softmaxes a row of logits, writes dlogits in place over the same buffer, and — the bigger win — never materializes the logits at all by chunking the lm-head matmul over tokens.

The question this lesson answers

You know from CUDA lesson 27 that the loss is the memory bottleneck and that the fused kernel cuts it ~24×. But what does the kernel actually look like in Triton? How do you online-softmax a 128k-wide row in BLOCK_V tiles, write the gradient back over the probabilities buffer, and fold in ignore_index and label smoothing without a second pass — all while keeping the loss accumulation in fp32 so it doesn't drift?

Why this is a memory kernel, not a compute kernel

Recall the shapes (CUDA lesson 27). With B·T = 8192, V = 128256, d = 4096, bf16:

Three full (B·T)×V tensors = 6.3 GB for one tensor's worth of information, because V ≫ d. The arithmetic per element is trivial (one exp, one subtract); every byte is read and written about once. This kernel's roofline is HBM bandwidth, not the tensor cores. The whole game is keeping bytes off HBM.

The identity that lets probs and dlogits share a buffer

Cross-entropy on softmax has the cleanest backward in deep learning. With p = softmax(logits), one-hot label y, over N counted tokens:

loss = logsumexp(logits) − logits[true]   =   (m + log ℓ) − logits[true]
dlogits = (p − y) / N

The gradient is the softmax, minus 1 at the true class, scaled by 1/N. So you never need a separate dlogits tensor: compute p into the logits buffer, then subtract the one-hot in place. probs and dlogits are the same buffer at two instants. Three tensors collapse to one — and trick 2 removes even that one.

The Triton kernel — one program per row

The first kernel assumes the logits already exist in HBM (you'll drop that assumption with trick 2). Each program owns one row of length V, walks it in BLOCK_V tiles, runs the lesson-10 recurrence to get (m, ℓ), then walks the row a second time writing dlogits in place. Loss and the running sum stay in fp32.

import torch, triton, triton.language as tl

@triton.jit
def ce_fused_kernel(
    logits_ptr,                 # [N_ROWS, V]  in/out: logits in, dlogits out (same buffer)
    labels_ptr,                 # [N_ROWS]     int64 true class, or ignore_index
    loss_ptr,                   # [N_ROWS]     fp32 per-row loss out
    stride_row,                 # row stride of logits (in elements)
    V, n_valid,                 # vocab size; number of non-ignored tokens (the 1/N divisor)
    ignore_index,
    smoothing,                  # label smoothing eps in [0, 1)
    BLOCK_V: tl.constexpr,
):
    row = tl.program_id(0)
    base = logits_ptr + row * stride_row
    label = tl.load(labels_ptr + row)

    # ignored row: zero its gradient, contribute no loss, return early
    if label == ignore_index:
        for off in range(0, V, BLOCK_V):
            cols = off + tl.arange(0, BLOCK_V)
            tl.store(base + cols, tl.zeros([BLOCK_V], tl.float32), mask=cols < V)
        tl.store(loss_ptr + row, 0.0)
        return

    # ---- pass 1: online softmax over the V dimension (lesson 10 recurrence) ----
    m = -float('inf')           # running max   (fp32)
    l = 0.0                     # running sum-of-exp at the current max (fp32)
    for off in range(0, V, BLOCK_V):
        cols = off + tl.arange(0, BLOCK_V)
        x = tl.load(base + cols, mask=cols < V, other=-float('inf')).to(tl.float32)
        cm = tl.max(x, axis=0)
        new_m = tl.maximum(m, cm)
        l = l * tl.exp(m - new_m) + tl.sum(tl.exp(x - new_m), axis=0)
        m = new_m

    lse = m + tl.log(l)         # logsumexp, fully fp32 stable

    # true-class logit (re-read just that one element)
    x_true = tl.load(base + label).to(tl.float32)

    # ---- loss, with label smoothing folded in ----
    # hard CE: lse - x_true.  smoothing mixes in the mean logit: see note below.
    loss = lse - x_true
    if smoothing > 0.0:
        # smoothed target puts (1-eps) on true, eps/V on every class:
        #   loss = (1-eps)*(lse - x_true) + eps*(lse - mean_logit)
        # mean_logit needs a sum of logits; fold it into pass 1 in practice.
        pass
    tl.store(loss_ptr + row, loss / n_valid)

    # ---- pass 2: write dlogits = (p - onehot)/N in place over the logits buffer ----
    inv_l = 1.0 / l
    for off in range(0, V, BLOCK_V):
        cols = off + tl.arange(0, BLOCK_V)
        x = tl.load(base + cols, mask=cols < V, other=0.0).to(tl.float32)
        p = tl.exp(x - m) * inv_l                       # softmax, recomputed from (m, l)
        g = p - tl.where(cols == label, 1.0, 0.0)       # subtract the one-hot
        g = g / n_valid                                 # scale by 1/N
        tl.store(base + cols, g, mask=cols < V)          # in place: overwrite logits

Four things an interviewer will look for:

• other=-float('inf') on the pass-1 load. The lesson-04 boundary rule: masked lanes past V must not win the max. Use the identity of max (−∞) so the recurrence is unaffected.
• fp32 accumulators. Loads are .to(tl.float32) before any exp or reduction; m and l are fp32. A bf16 running sum over 128k terms loses the small contributions and the loss drifts vs the reference.
• The gradient is recomputed, not stored. Pass 1 keeps only (m, ℓ) — two scalars. Pass 2 re-derives p = exp(x − m)/ℓ from them and writes the gradient straight over the logits. No probs tensor ever exists.
• other=0.0 on the pass-2 load. Now we're computing exp values that get masked-stored anyway; 0.0 is harmless and avoids an exp(-inf) in lanes we won't write.

def fused_ce(logits, labels, ignore_index=-100, smoothing=0.0):
    N_ROWS, V = logits.shape
    n_valid = (labels != ignore_index).sum().item()    # the 1/N divisor (counted once)
    loss = torch.empty(N_ROWS, device=logits.device, dtype=torch.float32)
    BLOCK_V = 16384 if V > 16384 else triton.next_power_of_2(V)
    ce_fused_kernel[(N_ROWS,)](
        logits, labels, loss, logits.stride(0),
        V, n_valid, ignore_index, smoothing,
        BLOCK_V=BLOCK_V, num_warps=8,
    )
    return loss.sum(), logits        # logits buffer now holds dlogits

ignore_index and label smoothing in the same pass

Both fold into the single fused pass — that's the point of fusing. ignore_index is the early-return branch: a padded/ignored row writes a zero gradient row and contributes 0 to the loss, and it is excluded from n_valid so it doesn't dilute the average. Label smoothing with eps spreads eps/V probability mass onto every class and 1 − eps on the true one. The loss becomes

loss = (1 − eps)·(lse − x_true) + eps·(lse − mean_logit)

where mean_logit = (1/V)·Σ logits — one extra running sum you accumulate alongside ℓ in pass 1. The gradient picks up a uniform −eps/V on every class plus the usual −(1 − eps) at the true index. Same two passes, a couple more fp32 scalars carried between them; no second materialization.

Trick 2 · chunk the lm-head matmul so logits never materialize

The kernel above still needs the full (B·T)×V logits in HBM. The bigger win (CUDA lesson 27) is to never compute them as one tensor. The logits come from a matmul logits = X·Wᵀ with X of shape (B·T)×d and W of shape V×d. Loop over token chunks: for a chunk of rows X_c, compute just that chunk's logits on chip, turn them into loss + gradient, push the gradient straight through the lesson-19 two-GEMM backward, and throw the chunk away.

At the algorithm level, the fused-linear-cross-entropy structure is a Python loop over chunks calling small Triton kernels — the per-chunk logits are too transient to be one giant @triton.jit, so libraries (Liger's fused_linear_cross_entropy, PyTorch's chunked CE) drive the loop in Python and fuse each chunk's softmax+CE. Per chunk:

# X: [B*T, d], W: [V, d] (lm-head, weight-tied or not). CHUNK rows at a time.
dX = torch.empty_like(X)
dW = torch.zeros_like(W, dtype=torch.float32)     # fp32 grad accumulator
total_loss = 0.0
for c0 in range(0, X.shape[0], CHUNK):
    Xc = X[c0:c0+CHUNK]                            # [CHUNK, d]
    logits_c = Xc @ W.T                            # [CHUNK, V] transient — never the full tensor
    # fused CE on this chunk: writes dlogits_c IN PLACE over logits_c, returns chunk loss
    loss_c = ce_fused_kernel[(Xc.shape[0],)](logits_c, labels[c0:c0+CHUNK], ...)
    total_loss += loss_c
    # lesson 19 two-GEMM backward, on the chunk:
    dX[c0:c0+CHUNK] = logits_c @ W                 # dXc = dlogits_c · W
    dW += logits_c.T.to(torch.float32) @ Xc        # dW += dlogits_cᵀ · Xc  (accumulate fp32)
    del logits_c                                   # the only big buffer — freed each iteration

The full (B·T)×V tensor never exists. HBM holds X, W, the scalar loss, and the gradients dX (size of X) and dW (size of W) you were going to compute anyway. The only transient is one chunk's logits, CHUNK·V·2 bytes. You recompute each chunk's logits on the way out instead of reading them back — extra FLOPs, but the lm-head is a tiny slice of total model FLOPs and you were memory-bound here, so it's nearly free.

Worked numbers — H100

B·T = 8192, V = 128256, d = 4096, bf16, CHUNK = 1024 rows. H100 HBM is 3.35 TB/s.

Bandwidth check on the in-place kernel: it reads the 2.10 GB logits twice (pass 1 + pass 2) and writes 2.10 GB once ≈ 6.3 GB of HBM traffic, so ≈ 6.3 / 3350 ≈ 1.9 ms at roofline — pure bandwidth, no tensor cores. On a 7B model the 6.3 GB → 0.26 GB swing is routinely the difference between fitting a sequence length and an OOM, with bit-identical math.

Traps

Interactive · logit-space HBM vs the knobs

Drag vocab, tokens, and chunk size. The bars are the naive 3×|logits| peak, the in-place 1×|logits| peak, and the fully-fused per-chunk transient. Watch the crossover: push the chunk up to the full batch and you're back to materializing everything.

What's next

You've now seen the two structural moves of memory-efficient backward kernels in Triton: fuse a reduction so the intermediate never lands (probs/dlogits share a buffer; the online recurrence keeps only two scalars), and recompute on the way out instead of storing (chunk the matmul, recompute each chunk's logits in the backward). Lesson 21 applies the recompute idea to normalization: recompute the per-row mean/rstd instead of saving them, and handle the dgamma/dbeta reduction across tokens — the column-reduction problem that atomics or a two-pass split both solve.