Walk through how the same transformer block behaves differently during training and during autoregressive inference, without changing any weights.
Free · 2 AI evals / day
**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.
Quick reference
TL;DR
Training pushes the entire sequence through every block in one forward pass, with a causal mask preventing each position from peeking at future tokens.
Imagine the model is a teacher grading a 1000-word essay. In TRAINING mode, the teacher sees the whole essay at once and grades every word in parallel. They use a piece of cardboard (the causal mask) to cover the next word as they evaluate each one, so the grade for word 50 is fair, but the work happens all together. In INFERENCE mode, the student is WRITING the essay one word at a time. The teacher has to predict each next word as it comes. They keep a notebook of every previous word's worth of context (the KV cache) so they don't have to re-read the whole essay every time. Same teacher, same brain, very different workload. Writing one word at a time is way slower per unit work than grading the whole essay at once.
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example.
Bidirectional models (BERT-style encoders) Encoder-only models do not have an autoregressive decode loop; they always process the full sequence in one pass, both at training and at use time.
Non-autoregressive generation (parallel decoding, diffusion-based text) These approaches generate all positions in parallel and do not have the single-token decode regime; the training versus inference shape difference is much smaller.
Speculative decoding with draft and target models The strictly autoregressive single-token framing breaks when a draft model emits k tokens that the target verifies in parallel; the inference loop becomes block-parallel at the verifier and the simple regime split needs an extra layer.
Quick scan · 2–4 min
Everything important, quickly.
5 min: training input shape `(B, T, D)` + causal mask + teacher forcing + parallel matmuls + inference input shape `(B, 1, D)` + KV cache + tall-skinny matmuls + prefill versus decode + same weights, different shapes.
Side-by-side
Aspect
Training
Inference (decode step)
Input shape to block
`(B, T, D)`
`(B, 1, D)`
K and V source
Computed fresh from full sequence
Cache + fresh K,V for current token
Causal mask
Required (upper-triangular)
Not needed; future tokens absent from cache
Attention score shape
`(B, H, T, T)`
`(B, H, 1, t)` tall-skinny
FFN matmul shape
`(B, T, D) × (D, 4D)`
`(B, 1, D) × (D, 4D)`
GPU bottleneck
Compute (tensor cores)
Memory bandwidth (HBM)
Token supervision
Teacher forcing on all positions
Sampled from model's own previous output
Backward pass
Yes
No
Real-world examples
Real products, models, and research that use this idea.
1vLLM (Kwon et al., 2023) is built specifically around the prefill versus decode shape difference; it uses PagedAttention to manage KV cache memory efficiently during decode.
2FlashAttention (Dao et al., 2022) optimizes the training-style attention matmul; FlashDecoding (2023) and FlashAttention v3 (2024) extend the optimization to decode-step shapes.
3Llama 3.1 70B has 80 blocks; during training each forward pass processes a 8192-token sequence in parallel; during decode each step processes a single token with the KV cache holding the prior context.
4Speculative decoding (e.g., Leviathan et al., 2023) exploits the decode-step bandwidth bottleneck: a small draft model proposes multiple tokens, the big model verifies them in one batched forward pass that better amortizes HBM reads.
Interviewer follow-ups
What an interviewer would ask next. Try answering before peeking at the approach.
QWhy is the FFN sublayer the dominant cost during prefill but the attention sublayer the dominant cost during long-context decode?
Sign in or go Pro to see the approach for this follow-up.
QWhat changes if you train with the same shape pattern as inference (one token at a time)?
Sign in or go Pro to see the approach for this follow-up.
QDoes dropout behave the same in training versus inference?
Sign in or go Pro to see the approach for this follow-up.
Don't say this
Red flags and common mistakes that signal junior thinking. Click to expand.
Most common mistake
Saying the block 'does different things' in training versus inference. The block's weights and computation graph are identical; only the input shapes change.