Same block, two regimes: training parallelism versus inference autoregression

**Training: one big parallel pass.** During training, the entire input sequence (say 4096 tokens) goes through every block in a single forward pass. The block input has shape `(batch, 4096, d_model)`. The attention sublayer computes Q, K, and V for all 4096 positions in parallel via three big matmuls. It then computes an attention score matrix of shape `(batch, heads, 4096, 4096)` and applies a **causal mask** to prevent each position from attending to positions to its right. The FFN sublayer runs as a single matmul over all 4096 positions. This is called **teacher forcing**: the model sees the correct token at every position and is asked to predict the next one at every position. Because the entire sequence is processed in parallel, training utilizes GPU tensor cores efficiently. **Inference: one token at a time, with a KV cache.** When the model is generating, it produces tokens one by one. At step `t`, the block receives input of shape `(batch, 1, d_model)`. Only the new token's Q is freshly computed. K and V for all previous tokens come from the **KV cache**, a persistent memory of K and V vectors from prior steps. The attention sublayer becomes a 'tall-skinny' matmul: `(1, d_head) × (d_head, t)` for the scores, then `(1, t) × (t, d_head)` for the values. The FFN still runs over each position independently, but now there is only one position. **Same weights.** Crucially, no weight changes between the two regimes. The block's parameters are identical. What differs is the shape of the data flowing through and whether K/V is computed fresh or read from cache.

**The block is one computation graph; the two regimes feed it different shapes.** **During training**, the input to the block is `(batch, seq_len, d_model)` for the full sequence length (commonly 4096 or 8192 tokens). The attention sublayer: - Projects to Q, K, V: three matmuls of shape `(batch, seq_len, d_model) × (d_model, d_model) → (batch, seq_len, d_model)` each. - Computes scores: `Q × K^T` produces shape `(batch, heads, seq_len, seq_len)`. - Applies a causal mask: the upper triangle is set to `-inf` so each position cannot attend to future positions. - Softmaxes and multiplies by V to produce attention output of shape `(batch, seq_len, d_model)`. The FFN runs as `(batch, seq_len, d_model) → (batch, seq_len, ffn_hidden) → (batch, seq_len, d_model)`. Two big matmuls over all positions. Both sublayers are doing one large matmul each per block; the GPU tensor cores stay saturated. Training compute scales as `O(seq_len^2)` for attention and `O(seq_len * d_model^2)` for FFN. **This is teacher forcing**: every target token at every position is supervised simultaneously from a single forward pass. **During inference**, the same block is now generating autoregressively. At each decode step `t`: - Input to the block is `(batch, 1, d_model)`: just the new token's residual stream. - Q is freshly computed: `(batch, 1, d_model) × (d_model, d_model) → (batch, 1, d_model)`. - K and V for the new token are computed and APPENDED to the KV cache, which already holds K and V from all previous tokens. - Attention scores: `Q × K_cache^T` produces shape `(batch, heads, 1, t)`. Tall-skinny. - Softmax and multiply by `V_cache` to produce attention output of shape `(batch, 1, d_model)`. The FFN still runs position-wise but on a single position: two matmuls of shape `(batch, 1, ffn_hidden)`. **The shape difference is the whole story.** Training matmuls are square or close to it; they hit FLOP peak on GPU.