Fused Linear Cross-Entropy: The Math, the Memory Wall, and Two Ways Around It

All code permalinks in this post point to specific commits: linkedin/Liger-Kernel at 38c0d4f and verl-project/verl at bf4b152. Line numbers verified against those snapshots.

The setting

In any modern LLM, the LM head is a giant matmul followed by a cross-entropy loss. For one decoded token at position $t$:

• $\text{hidden}[t] \in \mathbb{R}^H$ — the model’s last-layer hidden state ($H$ = hidden size, e.g. 1024 for Qwen3-0.6B, 4096 for Qwen3-4B).
• $W \in \mathbb{R}^{V \times H}$ — the LM-head weight matrix ($V$ = vocab size, e.g. 151{,}936 for Qwen3, ~176k for some custom vocabs).
• $\text{target}[t] \in {0, \dots, V-1}$ or $-100$ (ignored).

Per-token forward:

\(z = W \cdot \text{hidden}[t] \quad\in\mathbb{R}^V\) \(p_i = \mathrm{softmax}(z)_i = \frac{e^{z_i}}{\sum_j e^{z_j}}\) \(\ell_t = -\log p_{y_t} = -z_{y_t} + \log\sum_j e^{z_j} \;=\; -z_{y_t} + \mathrm{LSE}(z)\)

Per-batch, with mean reduction (the default in F.cross_entropy(..., ignore_index=-100, reduction="mean")):

\[L = \frac{1}{n_{\text{valid}}} \sum_{t=1}^{BT} \ell_t \cdot \mathbf{1}[y_t \ne -100]\]

So the loss is a single scalar. The gradients we need are $\text{grad_input}\in\mathbb{R}^{BT\times H}$ and $\text{grad_W}\in\mathbb{R}^{V\times H}$. Both are merely input-sized — there is no reason, intrinsically, to ever materialize anything else.

The engineering wall

Plain autograd, in PyTorch, computes the forward as:

logits = hidden @ W.T          # (BT, V)  — the big tensor
loss   = F.cross_entropy(logits, target, ignore_index=-100, reduction="mean")
loss.backward()

That single line hidden @ W.T allocates a (BT, V) tensor. At a representative shape — BT=65535, H=1024, V=176245 in bf16:

Tensor	Bytes	Size
hidden	$BT \cdot H \cdot 2$	128 MiB
W	$V \cdot H \cdot 2$	344 MiB
logits (bf16)	$BT \cdot V \cdot 2$	23 GiB
logits (fp32, what softmax wants)	$BT \cdot V \cdot 4$	46 GiB
softmax(logits) saved-for-backward	$BT \cdot V \cdot 4$	46 GiB

So the entire LM head and the input together are under 1 GiB, and we’re spending 50–100 GiB on an intermediate tensor that is about to be summed away into a scalar. This is the canonical “obvious” memory waste in transformer training, and it bites everyone who tries to push context length, batch size, or vocab past a certain threshold.

The identity that makes fused fwbw possible

Differentiate the per-token loss with respect to logits:

\[\frac{\partial \ell_t}{\partial z_i} = \frac{\partial}{\partial z_i}\Big(-z_{y_t} + \log\sum_j e^{z_j}\Big) = -\delta_{iy_t} + \frac{e^{z_i}}{\sum_j e^{z_j}} = p_i - \mathbf{1}[i = y_t]\]

so

\[\boxed{\;\frac{\partial \ell_t}{\partial z} = \mathrm{softmax}(z) - \mathrm{onehot}(y_t)\;}\]

Three consequences make this the foundational identity of every fused linear-CE implementation:

1. No “saved softmax” needed. $\frac{\partial \ell_t}{\partial z}$ depends only on $(z, y_t)$ — not on anything else from the forward pass. You can throw softmax(z) away the moment you’ve used it.
2. No second loop needed. The gradient at the LM-head input also follows by the chain rule: $\frac{\partial \ell_t}{\partial W} = (p_t - \mathrm{onehot}(y_t)) \otimes \text{hidden}[t]$. This is a per-token rank-1 update; the global $\text{grad_W}$ is a sum over tokens. The sum commutes with chunking — every chunk contributes its own partial loss and partial $\text{grad_W}$, so you can do fw and bw in one pass over the data.
3. Numerical stability is free. $\mathrm{softmax}(z) = \mathrm{softmax}(z - \max z)$, and $\mathrm{LSE}(z) = \max z + \log \sum_j e^{z_j - \max z}$. Two scalars per row — $\max z_m$ and $\sum_j e^{z_{m,j} - \max z_m}$ — are enough to reproduce the entire row’s softmax later.

Together, (1) + (2) say: you can compute loss and both gradients in a single streaming pass over tokens, without ever holding a (BT, V) tensor in HBM for longer than one chunk — and even that chunk can be reused buffer space.

From dlogits to grad_input and grad_W

Consequence (2) above hides the chain rule — let’s actually write it. Stack the per-token gradients into the batched dlogits matrix:

\[\mathrm{dlogits} \;\in\; \mathbb{R}^{BT \times V}, \qquad \mathrm{dlogits}[t,:] \;=\; \frac{1}{n_{\text{valid}}} \cdot \mathbf{1}[y_t \ne -100] \cdot \big(\mathrm{softmax}(z_t) - \mathrm{onehot}(y_t)\big)\]

(The 1/n_valid and the ignore-mask come from the outer L = (1/n_valid) Σ_t ℓ_t · 1[y_t ≠ −100] reduction, by the chain rule on the scalar coefficient in front of each ℓ_t.) Then, because $Z = \text{hidden} \cdot W^\top$ is just a matmul, the chain rule gives both downstream gradients as single GEMMs:

\[\boxed{\;\mathrm{grad\_input} \;=\; \mathrm{dlogits} \cdot W \;\in\; \mathbb{R}^{BT \times H}\;}\] \[\boxed{\;\mathrm{grad\_W} \;=\; \mathrm{dlogits}^\top \cdot \mathrm{hidden} \;\in\; \mathbb{R}^{V \times H}\;}\]

And — if the LM head has a bias $b \in \mathbb{R}^V$ — one column-sum:

\[\mathrm{grad\_b} \;=\; \sum_t \mathrm{dlogits}[t,:] \;\in\; \mathbb{R}^V\]

In practice, most modern LM heads have no bias (Qwen, LLaMA, Mistral all set bias=False on the head — the projection ties to embeddings or is intentionally bias-free to keep the softmax shift-invariant). Liger keeps the grad_bias plumbing because some tasks attach a learnable per-class bias to the head, but for the fused-CE kernel it’s either None (no allocation, no compute) or a single dlogits.sum(dim=0) reduction. We pass bias=None in our wrapper and ignore grad_bias in the return.

The whole backward is therefore just two GEMMs and (optionally) one reduction — same compute as a regular linear-layer backward. The trick is just to never materialize the dlogits matrix that connects them.

Below are two production implementations of that idea, with different choices about what “one chunk” means.

Strategy 1 — Liger: token-chunked fused fw+bw

linkedin/Liger-Kernel → src/liger_kernel/ops/fused_linear_cross_entropy.py

Liger processes tokens in chunks along the M (batch) axis. Inside the chunk, the full vocab is laid out as one wide tile.

The chunk-size trick

The cute design choice is in fused_linear_cross_entropy_forward, lines 56–58:

inc_factor  = triton.cdiv(V, H)                                    # ≈ V / H
chunk_size  = triton.next_power_of_2(triton.cdiv(BT, inc_factor))  # ≈ BT * H / V
num_chunks  = triton.cdiv(BT, chunk_size)

The intent is documented in the comment block right above (lines 45–51):

if we were to achieve the same memory consumption as BT x H, then the chunk size should be: inc_factor = (V+H-1)//H, chunk_size = (BT + inc_factor - 1)//inc_factor

In words: pick the chunk size so that chunk_size × V ≈ BT × H, i.e. the per-chunk logits tile takes about as much memory as the entire input tensor. At our BT=65535, H=1024, V=176245:

Quantity	Value
inc_factor = ceil(V/H)	173
chunk_size = next_pow2(ceil(BT / inc_factor))	512
num_chunks	128
Per-chunk logits tile (fp32)	$512 \times 176245 \times 4 = 361\text{ MiB}$

That single 361 MiB buffer is allocated once and reused across all 128 chunks. It never grows with the batch size.

The per-chunk fused kernel

For each chunk, Liger runs a Triton kernel that does softmax + CE + dlogits = softmax − onehot all in shared memory, then drops a per-row loss scalar back into loss_1d. Crucially, in lines 200–211:

grad_logits_chunk = logits_chunk  # chunk_size x V
...
if input_requires_grad:
    grad_input[start_idx:end_idx] = grad_logits_chunk @ weight

if grad_weight is not None and input_requires_grad:
    grad_weight += torch.mm(grad_logits_chunk.t(), _input_chunk).float()

That grad_logits_chunk = logits_chunk is the Liger trick made literal: the same memory that held logits for the forward CE pass is now re-interpreted as dlogits for the backward pass, written in place by the kernel that just computed the loss. The (cs, V) buffer is never allocated twice. Two cuBLAS matmuls (grad_logits_chunk @ weight and grad_logits_chunk.T @ _input_chunk) finish the chain rule, and the loop moves on to the next chunk.

Memory profile

Buffer	At our shape
_input (input)	128 MiB
weight (input)	344 MiB
logits_chunk (reused across chunks)	361 MiB
grad_input	128 MiB
grad_weight (fp32 accumulator)	688 MiB
loss_1d	0.3 MiB
Peak working set above inputs/outputs	~360 MiB

Compared to ~46 GiB for the naïve (BT, V) fp32 logits + softmax, this is two orders of magnitude smaller. The compute is unchanged — three GEMMs of size proportional to $BT \cdot V \cdot H$.

The only soft caveat: the inc_factor = V/H heuristic implicitly assumes $V \gg H$. At small $V$ or very small $H$ the chunk count or memory can drift from the “$\approx BT \times H$” sweet spot; the formula is bounded but not always optimal.

Strategy 2 — verl: tiled streaming with online softmax + three backward modes

verl-project/verl → verl/utils/kernel/kernels.py

Liger keeps the chunk tile inside SMEM/registers per chunk, but the chunk is still (cs, V) — the whole vocab lives in shared memory for one tile of tokens. verl pushes the tiling one step further: it tiles both the M axis (tokens) and the N axis (vocab), so the logits tile in any kernel is only (BLOCK_M, BLOCK_N) ≈ 128 × 128 — about 64 KiB, fits comfortably in registers.

The cost of this finer tiling: per row, the softmax normalizer $\sum_j e^{z_{m,j}}$ depends on all $V$ columns, so you can’t form $\mathrm{softmax}(z_m)_n$ when you only have a (128, 128) slice. verl handles this with classic online softmax (Flash-Attention-style) — maintain per-row running $(\mu_m, \ell_m)$ as you stream over vocab tiles, then either rescale-on-the-fly (forward) or recompute the logits tile in a second sweep (backward).

Forward — one Triton kernel, online softmax across vocab splits

efficient_entropy_kernel_general_mainloop, line 188.

The grid is (ceil(BT / BLOCK_M), num_splits) where num_splits = ceil(V / vocab_per_split). Each CTA owns one (M_tile, vocab_split) and:

1. K-loop: stream $\text{hidden}[M_{\text{tile}}, K_{\text{block}}] \cdot W[N_{\text{tile}}, K_{\text{block}}]^\top$ into a (BLOCK_M, BLOCK_N) logits tile in registers.

2. Online softmax update per row $m$ in the M-tile:

   \[\mu_m^{\text{new}} = \max(\mu_m, \max_j z_{m,j} \text{ in tile}) \qquad \ell_m^{\text{new}} = \ell_m \cdot e^{\mu_m - \mu_m^{\text{new}}} + \sum_j e^{z_{m,j} - \mu_m^{\text{new}}}\]
3. When the tile contains the target column for row $m$, accumulate $z_{m, y_m}$ into the per-row log-prob output.

After the mainloop, a small epilogue reduces the per-split scalars to a single $(\mu_m, \ell_m, \text{entropy}_b\text{[m]}, \text{logprobs}[m])$ per row. The wide tensor that gets saved-for-backward is exactly that — four arrays of shape (BT,) in fp32, totaling about 1 MiB. The (BT, V) tensor exists only as transient register tiles inside the kernel.

Backward — three modes, same math, different perf

verl exposes three backward strategies (defined at BackwardEnum, line 146):

class BackwardEnum:
    _Total_Fuse_MN     = 0  # fuse d_logits & d_hidden & d_weight, no intermediate storage
    _Total_Separate    = 1  # store d_logits, no special requirements
    _Split_Dlogits_N   = 2  # split d_logits along its N dimension (vocab)
    _Split_Dlogits_M   = 3  # split d_logits along its M dimension (not implemented)

The host-side dispatch in efficient_entropy_backward (line 1515) picks which kernel and which intermediate buffer (if any) to allocate.

Mode _Total_Fuse_MN — “no intermediate buffer”

Kernel: efficient_entropy_backward_kernel_general_mainloop_MN, line 771.

Grid: (ceil(BT / BLOCK_M), ceil(V / BLOCK_N)). Per CTA:

1. K-loop #1: re-stream the logits tile $z[M_{\text{tile}}, N_{\text{tile}}]$ in registers using the saved $(\mu_m, \ell_m)$.
2. Form d_logits in registers: $\frac{e^{z - \mu}}{\ell} - \mathrm{onehot}(y)$, scaled by upstream $\text{d_logprobs}$.
3. K-loop #2: for each $K_{\text{block}}$, $\text{d_hidden}[M_{\text{tile}}, K_{\text{block}}] \mathrel{+}= d_{\text{logits}} \cdot W[N_{\text{tile}}, K_{\text{block}}]$ via tl.atomic_add, and symmetrically $\text{d_W}[N_{\text{tile}}, K_{\text{block}}] \mathrel{+}= d_{\text{logits}}^\top \cdot \text{hidden}[M_{\text{tile}}, K_{\text{block}}]$ via tl.atomic_add.

Pro: the d_logits tile is never allocated in HBM — it lives only in registers. Total extra workspace at our shape: a few MiB of per-row scalars.

Con: at large $V$, the atomic-add contention is brutal. grad_input[m, k] is a sum over every vocab tile $n$, so all ceil(V / BLOCK_N) ≈ 1378 N-tiles atomically race on the same grad_input cells. The hardware serializes atomic adds per L2 cache line, so ~1300-way parallel reduction becomes ~1300-way serial. At our LM-head shape (V = 176245) we measured this mode at 6.07 seconds vs 326 ms for _Split_Dlogits_N — same math, 18× slower. The mode is the right design for small V (RL-style entropy bonuses on action logits, say), and a trap at LM-head vocab.

Mode _Split_Dlogits_N — “small intermediate, cuBLAS matmuls”

Kernel: efficient_entropy_backward_kernel_general_d_logits_split_N, line 1388. Host loop: efficient_entropy_backward, lines 1663–1721.

Allocate one $(BT, \text{vocab_per_split})$ bf16 buffer; reuse it across $\text{ceil}(V / \text{vocab_per_split})$ splits.

for split_idx in range(num_splits):
    efficient_entropy_backward_kernel_general_d_logits_split_N[...](...)   # Triton fills _d_logits slice
    if split_idx == 0:
        torch.matmul(_d_logits, weight[slice], out=d_hidden)               # cuBLAS
    else:
        d_hidden += torch.matmul(_d_logits, weight[slice])                 # cuBLAS
    torch.matmul(_d_logits.T, hidden, out=d_weight[slice])                 # cuBLAS

Pro: cuBLAS matmuls handle the column reduction in d_hidden and the row reduction in d_weight — no atomics, full tensor-core throughput. Workspace is bounded by vocab_per_split (default 9504 upstream; we tuned it down to 4096 for LM-head V).

Con: still allocates a (BT, ~4k) slice — at our shape, ~512 MiB. Two cuBLAS GEMMs and one Triton kernel per split = ~132 kernel launches at 44 splits, vs Liger’s ~128 fused launches.

Mode _Total_Separate

Materializes the full (BT, V) d_logits tensor in HBM, then two cuBLAS matmuls. Identical to the textbook “save softmax for backward” pattern. Useful only when $V$ is small.

Memory profile

For _Split_Dlogits_N (the recommended mode at LM-head scale), at our shape:

Buffer	Size
_input (input)	128 MiB
weight (input)	344 MiB
maximum, accumulate, entropy_b, logprobs (saved-for-backward)	~1 MiB
_d_logits slice (reused across splits)	512 MiB
grad_input	128 MiB
grad_weight	344 MiB
Peak working set above inputs/outputs	~512 MiB

Comparable to Liger’s ~360 MiB. The intermediate is now sliced along vocab instead of along tokens, which turns out to interact better with cuBLAS GEMM shapes when $V$ dominates $BT$.

Why two passes (forward + backward) instead of one?

A natural question after seeing Liger’s per-chunk fwbw fusion: can verl do the whole thing in one streaming pass, online-softmax-style?

Short answer: no, because grad_W is column-owned while online softmax rescaling is row-local. When the running max $\mu_m$ updates, you can locally rescale a row’s contribution to $\text{grad_input}$ (which is indexed by $m$) — that’s exactly what FlashAttention does. But $\text{grad_W}[n, k]$ is a sum over all rows $m$, each with its own $(\mu_m, \ell_m)$ history. There’s no row-specific rescale you can apply to a quantity that has already been summed across rows. The two ways out are:

1. Don’t tile N — keep all of $V$ available per row (this is Liger’s choice; works as long as $\mathrm{chunk_size} \times V$ fits in SMEM).
2. Recompute $z$ in a second pass once $(\mu_m, \ell_m)$ are final (this is verl’s choice; pays ~33% extra backward FLOPs to keep the in-flight tile to BLOCK_M × BLOCK_N).

Liger’s bet pays off when $V$ is small relative to SMEM-per-CTA. verl’s bet pays off when $V$ is so large that even one-row-of-vocab is unwieldy.

Putting numbers on it

For one concrete LM-head shape (BT=65535, H=1024, V=176245, bf16, H100) we benchmarked all three (a Python reference, Liger, and verl’s tuned _Split_Dlogits_N):

Variant	Forward+Backward latency	Speedup vs reference
Reference (chunked Python)	1651 ms	1.00×
Liger	753 ms	2.19×
verl (tuned _Split_Dlogits_N)	356 ms	4.64×

The 2.1× margin verl has over Liger at this shape comes down to the cuBLAS-friendly GEMM shapes in the backward and the fact that the per-row aux is genuinely tiny (~1 MiB) — once you start tiling N you stop paying for the per-chunk (cs, V) allocation that Liger still owns. At smaller V, the answer typically flips.

Takeaways

The “obvious” logits = hidden @ W.T line in every LLM training script hides the biggest single memory waste in transformer training — a (BT, V) intermediate that materializes only to be reduced into a scalar. Two principles unlock the fix:

1. dL/dz = softmax(z) − onehot(y) lets backward forget the forward. The gradient at logits depends only on (z, y). No softmax needs to be saved.
2. Loss decomposes additively over tokens, so per-token contributions to both loss and grad_W can be accumulated in a streaming pass. No fancy global computation is needed — just commute the sum with the loop.

Liger and verl both apply these. The difference is how aggressively they tile:

• Liger keeps (chunk_size, V) together — chunk_size is sized so this is the same size as the input. One Triton kernel per chunk does softmax + CE + in-place dlogits + chain-rule matmuls. Memory: ~size of input.
• verl tiles down to (128, 128) register tiles, uses online softmax to thread per-row stats across vocab tiles, and offers three backward modes that pick a different point on the (intermediate-buffer, atomic-contention) trade-off. The default _Split_Dlogits_N allocates a (BT, ~4k) slice and offloads the column-reduction to cuBLAS.

Same math. Same identity. Two production answers.