Transformer follow-ups, quick review: the common ones

Generation happens one token at a time. Under a causal mask the K and V of past tokens never change, so recomputing them every step is wasted work. Without a cache, step $t$ recomputes $t$ keys, $t$ values and about $t^2/2$ scores; with a cache it computes one new K/V pair and one row of scores.

Per-step attention drops from $O(t^2)$ to $O(t)$. The price is memory that grows linearly with tokens, and every step has to read the entire cache: small-batch decoding is usually limited by the bandwidth for reading weights and the cache (memory-bound); for why a larger batch moves it toward compute-bound, see the decode roofline example. Everything that “shrinks the cache” afterwards (GQA, MLA, sliding windows) is aimed at exactly this.

$2 \times n_{\text{layers}} \times n_{\text{kv heads}} \times d_{\text{head}} \times \text{tokens} \times \text{bytes}$. Note that the number of query heads does not appear. A Llama-2-7B shape (32 layers, 32 KV heads, $d_{\text{head}}=128$, fp16) is 0.5 MB per token, so 16 GB at 32K context.

The logits produced position by position on the cached path must be numerically equivalent to a full forward pass without a cache. The easy mistakes are the mask and the positions: a query's absolute position is cache.pos + i, the mask is a non-square $(T, S)$, and the RoPE cos/sin must be sliced starting at cache.pos.

Cache size depends only on the number of KV heads, not on the number of query heads, so memory is saved entirely on the K/V side. MQA: all query heads share a single K/V head; the smallest cache, with some quality loss. GQA: groups of query heads share one K/V head, the middle ground; Llama 3 8B has 32 query heads and 8 KV heads.

MLA does not cache per-head K and V. It caches one low-dimensional latent vector per token and expands it, through an up-projection, into each head's own K and V when needed. DeepSeek's setting is a latent of $4\,d_{\text{head}}$ plus a $d_{\text{head}}/2$ RoPE key, $4.5\,d_{\text{head}}$ in total: what GQA would need with 2.25 KV heads, except that every head still gets its own K/V.

RoPE is a position-dependent rotation, so it cannot be absorbed into a position-independent low-rank projection. MLA therefore carries position on a separate, uncompressed small key (64 dims) and sends the rest of the content through the latent.

Through depth. Every extra layer extends the reach by another window, roughly $L\,(w-1)+1$ tokens after $L$ layers. Mistral 7B uses $w=4096$ over 32 layers, a theoretical reach of about 131K.

Compute goes from $O(n^2)$ to $O(n\,w)$, and the per-layer KV cache is capped at $w$. Gemma 3 interleaves five local layers with one global layer, and only the global layers need a full-length cache.

Softmax weights must sum to 1. When a head has “nothing to look at”, the model learns to dump the excess weight on the first few tokens. Evict them and that weight has nowhere to go, so the distribution falls apart. StreamingLLM keeps the first few tokens forever; gpt-oss adds a learned bias to the softmax denominator, which plays the same role.

Each pair of dimensions in $q$ and $k$ is rotated by an angle proportional to the position, every pair at its own frequency. The dot product of two rotated vectors depends only on the difference of their angles, so the score depends only on the relative distance $m-n$.

No. The relativity comes from the two rotations canceling: $\langle R_m q, R_n k\rangle = \langle q, R_{n-m} k\rangle$. Rotate only one side and what is left in the dot product is absolute position.

A larger base slows every pair except the first (whose frequency is $b^0=1$), so low-frequency pairs do not wrap around within a longer range and distant positions are less likely to alias. Changing the base alone does not guarantee good long-context behavior; training and evaluation still matter. Llama 3 uses $\theta = 500\text{k}$; Gemma 3 uses $\theta = 1\text{M}$ on its global layers.

The FFN is replaced by $N$ experts, and a router sends each token to only the top $k$. Total parameters scale with $N$; compute per token scales only with $k$. Mixtral is 8 choose 2 (about 47B total, about 13B active); DeepSeek-V3 routes to 8 of 256 experts, plus one shared expert that every token visits (671B total, 37B active).

Load imbalance: the router tends to send tokens to a few experts and the rest learn nothing. The usual fix is an auxiliary load-balancing loss; DeepSeek-V3 relies mainly on a bias term instead, keeping only a very small sequence-wise auxiliary loss.

It is not an approximation; the result is exact. What it saves is memory traffic, not arithmetic: the standard implementation writes the full $n\times n$ score matrix to slow HBM and reads it back. FlashAttention cuts Q, K and V into blocks that fit in on-chip SRAM, computes block by block, and writes back only the $n\times d$ output, with extra memory linear in $n$. For how far apart HBM and on-chip SRAM are in capacity and bandwidth, see the GPU memory hierarchy.

With an online softmax: every row keeps a running max and a running normalizer, and each new block updates both and corrects the output accumulated so far. Mathematically it is identical to normalizing the whole row at once.

Four: the norm went from post-LN LayerNorm to pre-norm RMSNorm (deep stacks train more stably); position went from absolute encodings added to the embedding to RoPE; the FFN activation went from ReLU / GELU to SwiGLU; attention went from MHA to GQA. On top of that, the encoder and cross-attention are gone, and so are the biases.

Post-norm puts the norm after the residual add, so the gradient passes through a norm at every layer and deep networks need warmup to train stably. Pre-norm moves the norm inside the residual branch, leaving an unscaled identity path along the trunk.