Fused cross-entropy — the memory kernel
The logits tensor is the single largest tensor a language model ever holds — bigger than any weight, bigger than any activation. Materializing it, its softmax, and its gradient three times over is the most wasteful thing a naive training loop does. The fused cross-entropy kernel makes all three vanish from HBM. It's the biggest single memory win in LLM training, and it's a direct application of lesson 26's GEMM plus lesson 12's online softmax.
The question this lesson answers
You profile training memory and the peak is dominated by one allocation near the loss — far bigger than you expected from the model size. Why is the output layer, not the 32 transformer blocks, the memory bottleneck, and how does one fused kernel cut it by ~4×?
Why logits are the biggest tensor in the model
The final layer projects each token's hidden state to a score over the whole vocabulary: logits ∈ ℝ(B·T)×V. Modern vocabularies are large — V = 128k–256k. Compare the shapes:
| Tensor | Shape | Size (B·T = 8192, V = 128256, d = 4096, bf16) |
|---|---|---|
| A hidden activation | (B·T)×d | 8192 · 4096 · 2 = 0.067 GB |
| logits | (B·T)×V | 8192 · 128256 · 2 = 2.10 GB |
| softmax probs (same shape) | (B·T)×V | 2.10 GB |
| dlogits (same shape) | (B·T)×V | 2.10 GB |
The logits are ~31× a hidden activation because V ≫ d. And a naive loss path holds three of them: the logits, the softmax probabilities, and the gradient dlogits — ~6.3 GB for one tensor's worth of information, on top of everything else. That is the allocation your profiler is screaming about.
The gradient that makes fusion trivial
Cross-entropy on softmax has the cleanest backward in all of deep learning. With p = softmax(logits) and one-hot label y, over N tokens:
Read the second line carefully: the gradient is the softmax itself, minus 1 at the true class, scaled by 1/N. You do not need to store p to get dlogits — you compute p and immediately overwrite it with p − y. The "probs" tensor and the "dlogits" tensor are the same buffer at two instants. That collapses three tensors to one, and the next trick removes even that one from HBM.
Trick 1 · online softmax: never materialize probs
logsumexp over V needs the max and the sum of exponentials, both of which you can compute in a single streaming pass with the running-max recurrence from lesson 12 — the same Milakov–Gimelshein update that powers FlashAttention. So one row of logits is read in tiles, reduced online to (m, ℓ), and the loss is m + log ℓ − logits[true]. No full probs row is ever stored; the gradient (p − y)/N for a tile is recomputed from the saved (m, ℓ) as exp(logits − m)/ℓ on the way out.
Trick 2 · fuse the lm-head matmul: never materialize logits
The bigger win. The logits are produced by a matmul logits = X·Wᵀ (X is (B·T)×d, W is V×d). Instead of computing the whole (B·T)×V logits then reducing, chunk over the token dimension: for a chunk of rows Xc,
- compute that chunk's logits Xc·Wᵀ (a small GEMM, lesson 26) — lives only in SMEM/registers;
- online-softmax it to the chunk's loss and the chunk's dlogitsc = (pc − yc)/N;
- immediately push dlogitsc back through the matmul backward: dXc = dlogitsc·W and accumulate dW += dlogitscᵀ·Xc (the two-GEMM rule);
- discard the chunk's logits and move on.
At no point does the full (B·T)×V logits tensor exist in HBM — only one chunk does. HBM holds the inputs (X, W), the scalar loss, and the gradients (dX same size as X, dW same size as W). The 2.1 GB tensor is gone; peak drops to the size of one chunk plus the gradients you were going to compute anyway.
Worked number: the headline win
B·T = 8192, V = 128256, d = 4096, bf16, chunk = 1024 rows.
| Path | Logit-space HBM peak | Notes |
|---|---|---|
| Naive (logits + probs + dlogits) | 3 × 2.10 = 6.3 GB | Three full (B·T)×V tensors live at once. |
| Fused softmax+CE (keep logits, drop probs/dlogits) | 2.10 GB | Trick 1 only: probs and dlogits reuse the logits buffer. |
| Fully fused (chunk the matmul) | 1 chunk = 1024·128256·2 ≈ 0.26 GB transient | Trick 1+2: full logits never materialize; ~24× smaller peak. |
On a 7B model that 6 GB → 0.26 GB swing is often the difference between fitting a sequence length and OOM — without changing the math or the result. It is the first optimization any LLM trainer applies, and it ships in libraries (Liger Kernel's "fused linear cross entropy," PyTorch's chunked CE).
The cost you pay, and why it's nearly free
Chunking the matmul means you cannot reuse a single big cuBLAS GEMM call; you issue one small GEMM per chunk, and in the backward you recompute each chunk's logits rather than reading them back (they were never stored). That is extra compute — but the lm-head GEMM is a tiny fraction of total model FLOPs, and you were memory-bound, not compute-bound, at the loss. Trading a few percent more FLOPs for a 24× memory cut is the definition of a good deal when activations decide whether the step fits (lesson 25's widget).
Numerical traps
| Trap | Symptom | Fix |
|---|---|---|
| exp without max-subtraction | inf/NaN at large logits | online running-max (logsumexp), lesson 12 |
| fp32 loss/grad accumulation skipped | loss drifts vs reference; grad noise | accumulate loss and dW in fp32 |
| label smoothing / ignore_index dropped | subtly wrong loss; padded tokens contribute | fold the smoothing term and the mask into the same fused pass |
| chunk too large | still spills to HBM / register spill | size the chunk so one chunk's logits fit on chip; autotune it |
Interactive · logit-space memory vs chunk size
Set vocabulary, tokens, and chunk size. The widget shows the naive 3×|logits| peak against the fused per-chunk transient, and the crossover where chunking stops helping (chunk ≈ full batch ⇒ back to naive).
What this sets up
You've seen the two structural moves of memory-efficient backward kernels: fuse a reduction so the intermediate never lands, and recompute on the way out instead of storing. The next lesson applies the recompute idea to the normalization layers — where the saved statistic is one scalar per row and recomputing it beats reading it back.