The transformer, walked through with tensor shapes

What it is

A transformer is a stack of identical blocks that each do two things: mix information across positions (attention) and transform each position independently (a feed-forward network). Everything else is plumbing.

The reason to walk it in shapes rather than in prose is that every performance and capacity question you will be asked reduces to arithmetic on those shapes. "Why does context length cost quadratic memory," "how big is the KV cache," "why is prefill compute-bound and decode memory-bound," and "how much does a LoRA adapter cost" all have the same answer form: multiply out the dimensions.

The notation used throughout, with a concrete model to anchor it (roughly Llama-3-8B):

Symbol	Meaning	Example
`B`	Batch size	4
`S`	Sequence length (tokens)	2048
`d`	Model dimension (hidden size)	4096
`L`	Number of layers	32
`H`	Number of attention heads	32
`d_h`	Dimension per head, `d / H`	128
`d_ff`	Feed-forward inner dimension	14336
`V`	Vocabulary size	128256

What this is confused with: the "attention is quadratic" claim is about the attention score matrix, not about the whole model. For a typical model at moderate sequence length, the feed-forward network uses more FLOPs than attention does. The quadratic term dominates only when S gets large relative to d, and knowing where that crossover sits is the difference between reciting a fact and understanding it.

The problem it solves

Before transformers, sequence models were recurrent: process token 1, carry a hidden state, process token 2. Two consequences followed and both were fatal at scale.

No parallelism over the sequence. Token t cannot be computed until token t-1 is done, so training on a 2,000-token sequence is 2,000 sequential steps. GPUs are wide parallel machines and an RNN uses almost none of that width.

Information decay. The influence of token 1 on token 500 passes through 499 multiplications, so gradients vanish and long-range dependencies are learned badly. LSTMs and GRUs mitigated it and did not remove it.

Attention replaces the recurrence with a direct, learned, all-pairs comparison: every position looks at every other position in one operation. That gives constant path length between any two tokens and full parallelism across the sequence during training. The price is the all-pairs comparison itself, which is the S x S matrix everyone quotes.

Mechanics

One block, end to end

x                                       [B, S, d]      input
  │
  ├── RMSNorm(x)                        [B, S, d]      pre-norm
  │     └── Attention(...)              [B, S, d]
  ├── x = x + attn_out                  [B, S, d]      residual
  │
  ├── RMSNorm(x)                        [B, S, d]
  │     └── FFN(...)                    [B, S, d]
  └── x = x + ffn_out                   [B, S, d]      residual

Shape is unchanged end to end: [B, S, d] in, [B, S, d] out. That invariance is why you can stack 32 of these, and it is worth stating because it makes the whole architecture a repetition of one function.

Attention, shape by shape

# x: [B, S, d]  = [4, 2048, 4096]

# 1. Project to queries, keys, values. Three [d, d] matrices.
q = x @ W_q            # [4, 2048, 4096]
k = x @ W_k            # [4, 2048, 4096]
v = x @ W_v            # [4, 2048, 4096]

# 2. Split into heads: reshape then transpose.
q = q.view(B, S, H, d_h).transpose(1, 2)    # [4, 32, 2048, 128]
k = k.view(B, S, H, d_h).transpose(1, 2)    # [4, 32, 2048, 128]
v = v.view(B, S, H, d_h).transpose(1, 2)    # [4, 32, 2048, 128]

# 3. Scores: every query against every key. THE QUADRATIC STEP.
scores = q @ k.transpose(-2, -1) / sqrt(d_h)   # [4, 32, 2048, 2048]
#                                                        ^^^^^^^^^^^ S x S

# 4. Causal mask: position i may not see j > i.
scores = scores.masked_fill(causal_mask, float('-inf'))

# 5. Softmax over the last dim: each query's attention distribution over keys.
attn = softmax(scores, dim=-1)                 # [4, 32, 2048, 2048]

# 6. Weighted sum of values.
out = attn @ v                                 # [4, 32, 2048, 128]

# 7. Merge heads back and project out.
out = out.transpose(1, 2).reshape(B, S, d)     # [4, 2048, 4096]
out = out @ W_o                                # [4, 2048, 4096]

Step 3 is the whole "quadratic" story, and it is worth pricing.

scores tensor: B x H x S x S x 2 bytes (fp16)
             = 4 x 32 x 2048 x 2048 x 2
             = 1.07 GB

At S = 8192:   4 x 32 x 8192 x 8192 x 2 = 17.2 GB
At S = 32768:  4 x 32 x 32768 x 32768 x 2 = 275 GB

A 16x increase in sequence length is a 256x increase in that one intermediate tensor. This is what FlashAttention removes: it never materialises the full S x S matrix, computing softmax in tiles that stay in SRAM. The maths is identical; the memory behaviour is not.

The 1/sqrt(d_h) scaling exists because the dot product of two d_h-dimensional vectors with unit-variance components has variance d_h. Without the scaling, at d_h = 128 the scores have standard deviation around 11, softmax saturates, and gradients vanish. It is one line and it is load-bearing.

Grouped-query attention: the KV cache fix

Standard multi-head attention gives each head its own K and V. GQA shares K and V across groups of query heads:

H   = 32     # query heads
H_kv = 8     # key/value heads: 4 query heads share each KV head

q = x @ W_q   # [B, S, 32, 128]  -> full
k = x @ W_k   # [B, S,  8, 128]  <- 4x smaller projection
v = x @ W_v   # [B, S,  8, 128]

k = k.repeat_interleave(H // H_kv, dim=2)   # [B, S, 32, 128] for the matmul

The point is not the compute saving; it is the KV cache, which is what you carry per token during generation:

Per token, per layer:  2 (K and V) x H_kv x d_h x 2 bytes
MHA  (H_kv = 32):  2 x 32 x 128 x 2 = 16,384 bytes = 16 KB
GQA  (H_kv = 8):   2 x  8 x 128 x 2 =  4,096 bytes =  4 KB

Full model, 32 layers, 8192-token context, one sequence:
MHA:  16 KB x 32 x 8192 = 4.29 GB      per concurrent request
GQA:   4 KB x 32 x 8192 = 1.07 GB      per concurrent request

On an 80 GB GPU holding a 16 GB model, that is the difference between about 14 concurrent requests and about 59. GQA is not a small optimisation; it is what makes serving long contexts economic. See KV cache math and batching for the serving consequences.

The feed-forward network, and why it is most of the parameters

Modern models use a gated FFN (SwiGLU):

# x: [B, S, 4096]
gate = x @ W_gate          # [B, S, 14336]
up   = x @ W_up            # [B, S, 14336]
h    = silu(gate) * up     # [B, S, 14336]   elementwise gate
out  = h @ W_down          # [B, S, 4096]

Three matrices of 4096 x 14336 = 58.7 M parameters each, so 176 M per layer. Attention's four 4096 x 4096 matrices are 16.8 M each, 67 M per layer (less with GQA).

The FFN is roughly 72 percent of a layer's parameters. The counterintuitive fact worth carrying: attention gets all the attention, and the feed-forward network is where most of the model lives.

Parameter and FLOP accounting

Per layer:
  Attention (GQA, H_kv=8):
    W_q: 4096 x 4096              = 16.8 M
    W_k: 4096 x 1024              =  4.2 M    (8 heads x 128)
    W_v: 4096 x 1024              =  4.2 M
    W_o: 4096 x 4096              = 16.8 M
                                    ───────
                                     42.0 M
  FFN:
    3 x (4096 x 14336)            = 176.2 M
                                    ───────
  Per layer total                   218.2 M

x 32 layers                       = 6.98 B
Embedding (128256 x 4096)         =  0.53 B
Output head (tied or separate)    =  0.53 B
                                    ───────
Total                             ≈ 8.0 B parameters

For FLOPs, the rule of thumb that survives contact with reality:

$$\text{forward FLOPs} \approx 2 \times N_{\text{params}} \times N_{\text{tokens}}$$

$$\text{training FLOPs} \approx 6 \times N_{\text{params}} \times N_{\text{tokens}}$$

(forward is 2, backward is roughly 2x forward, so 6 total). The attention score computation adds 2 x B x H x S^2 x d_h on top, which for S = 2048 and this model is about 5 percent of the total, and at S = 32768 becomes comparable to everything else. That crossover is the honest answer to "is attention quadratic": it is, and it only dominates past roughly S ≈ d.

Prefill versus decode, in shapes

This is where the shapes explain the serving behaviour.

PREFILL (process the whole prompt at once):
  x: [B, S, d]  with S = 2048
  Every matmul is a big matrix-matrix product.
  Arithmetic intensity is high -> COMPUTE-BOUND.

DECODE (generate token S+1):
  x: [B, 1, d]  <- ONE position
  Every matmul is matrix-VECTOR.
  You read the entire model's weights (16 GB) to compute one token per sequence.
  Arithmetic intensity is terrible -> MEMORY-BANDWIDTH-BOUND.

The decode arithmetic, which is the number to remember:

Weights read per decode step:  16 GB (fp16, 8B params)
H100 memory bandwidth:         3.35 TB/s
Minimum time per step:         16 / 3350 = 4.8 ms
Maximum tokens/sec/sequence:   ~209

That ceiling is independent of batch size, because the same weights serve every sequence in the batch. Which is exactly why batching helps throughput and not per-stream latency: you read the weights once and amortise them across B sequences. See prefill vs decode.

A worked example: sizing a deployment from the shapes alone

A team needs to serve an 8B model at 8k context, targeting 200 concurrent requests, on H100 80GB GPUs. The question is how many GPUs, answered before touching any hardware.

Step 1: weights.

8.0 B params x 2 bytes (fp16) = 16.0 GB

Step 2: KV cache per request.

GQA: H_kv = 8, d_h = 128, L = 32
Per token: 2 x 8 x 128 x 2 bytes x 32 layers = 131,072 bytes = 128 KB/token

At 8192 tokens: 128 KB x 8192 = 1.07 GB per request

Step 3: what fits.

GPU memory:              80.0 GB
Weights:                -16.0 GB
Activations + overhead:  -6.0 GB   (framework, workspace, fragmentation)
                         ───────
Available for KV cache:  58.0 GB
Concurrent requests:      58.0 / 1.07 = 54 per GPU

Step 4: GPUs needed.

200 concurrent / 54 per GPU = 3.7  ->  4 GPUs

Step 5: check the throughput ceiling, which is the step teams skip.

Decode is bandwidth-bound: 16 GB weights / 3.35 TB/s = 4.8 ms/step minimum
Per GPU: ~209 steps/sec, each producing one token per sequence in the batch
With batch 54:  54 x 209 = 11,286 tokens/sec/GPU (theoretical ceiling)
Realistically ~50-60% of that: ~6,000 tokens/sec/GPU

At an average of 400 output tokens per response, that is about 15 completed responses per second per GPU, or 60 across four GPUs.

Step 6: what changes the answer.

If the model were MHA instead of GQA (H_kv = 32):
  KV per token: 512 KB, per request at 8k: 4.29 GB
  Concurrent per GPU: 58 / 4.29 = 13
  GPUs for 200 concurrent: 16    <- FOUR TIMES the hardware

If context were 32k instead of 8k:
  KV per request: 4.29 GB
  Concurrent per GPU: 13
  GPUs for 200 concurrent: 16

If weights were int8 instead of fp16:
  Weights: 8 GB, freeing 8 GB for KV
  Available for KV: 66 GB -> 61 concurrent per GPU
  And decode bandwidth halves: 8/3350 = 2.4 ms/step -> ~418 steps/sec
  GPUs for 200 concurrent: 4 (same), but ~2x the token throughput

Three findings from arithmetic alone. GQA versus MHA is a 4x hardware difference. Context length trades linearly against concurrency, so "support 32k context" is a quadrupling of the fleet at fixed concurrency, not a configuration change. And quantisation helps twice: less memory for weights and less to read per decode step.

The whole sizing took ten minutes and no GPU. That is the argument for knowing the shapes: capacity questions about LLM serving are arithmetic, and the arithmetic is this arithmetic.

Production evidence

Vaswani et al., "Attention Is All You Need" (2017) introduced the architecture with post-norm and learned positional embeddings. Essentially every production model has since moved to pre-norm (more stable training), RMSNorm (cheaper than LayerNorm, no mean subtraction), rotary position embeddings (see RoPE and ALiBi), and SwiGLU feed-forward layers. Quoting the original paper's exact configuration as "the transformer" is a dated answer.

GQA (Ainslie et al., 2023) was adopted rapidly across Llama 2 70B, Llama 3, Mistral and most subsequent open models, precisely because of the KV cache arithmetic above. Multi-query attention (MQA, one KV head) is the extreme version and loses more quality; GQA at 4 to 8 KV heads is the settled compromise.

FlashAttention (Dao et al., 2022) and FlashAttention-2 are standard in every serving stack. The insight is IO-aware tiling: compute exact attention without materialising the S x S matrix in HBM. It is not an approximation, which is why adoption was universal and immediate.

The Chinchilla paper (Hoffmann et al., 2022) established the compute-optimal parameter-to-token ratio (roughly 20 tokens per parameter) using the 6 N D FLOP approximation above, which is why that formula appears in every scaling discussion.

vLLM's PagedAttention treats the KV cache like virtual memory with paging, which matters because the naive allocation (reserve max context per request) wastes most of the cache to internal fragmentation. Their reported 2 to 4x throughput improvement is mostly recovered fragmentation.

The debate

Is attention quadratic a real problem? At S = 2048 and d = 4096, the FFN dominates FLOPs and attention is around 5 percent. At S = 32768 attention is comparable to everything else, and the S x S intermediate would be hundreds of gigabytes without FlashAttention. My position: for context up to roughly 8k, the quadratic term is a memory problem that FlashAttention solved and not a compute problem. Past 32k it becomes both. Answering "attention is quadratic so long context is expensive" without the crossover is the answer of someone who has read about it.

Linear attention alternatives (Mamba, RWKV, state-space models) trade the all-pairs comparison for a recurrent state, giving linear scaling and constant memory during generation. They are genuinely promising and, as of now, transformer quality at equivalent scale has not been matched on the tasks people care about, and hybrid designs (a few attention layers among many SSM layers) are where the practical results are. Worth naming; not worth betting an architecture on yet.

How much does architecture choice matter versus data? Less than people expect. Chinchilla's finding was that most large models of its era were badly under-trained relative to their size, and subsequent progress has come more from data quality and quantity than from architectural change. The architectural changes that did stick (GQA, RoPE, SwiGLU, RMSNorm) are mostly efficiency improvements rather than capability ones, which is itself the interesting observation.

Should you know this level of detail? For an infrastructure or platform role, yes, because the sizing arithmetic in the worked example is a routine task and doing it wrong costs real money. For an application role building on APIs, the shapes matter less than the serving behaviour they imply (prefill versus decode, KV cache versus concurrency, context length versus cost). The honest boundary: you should be able to derive why a 32k context costs 4x the memory of an 8k one, and you do not need to implement FlashAttention.

Follow-up Q&A

"Walk me through the shapes in one attention layer."

Input [B, S, d]. Three projections give Q, K, V each [B, S, d], reshaped to [B, H, S, d_h]. Scores are Q @ K^T giving [B, H, S, S], scaled by 1/sqrt(d_h), causally masked, softmaxed. Multiply by V for [B, H, S, d_h], merge heads back to [B, S, d], project through W_o. The [B, H, S, S] tensor is the quadratic one, and at B=4, H=32, S=2048 in fp16 it is 1.07 GB, which is what FlashAttention avoids materialising.

"Why divide by sqrt(d_h)?"

The dot product of two d_h-dimensional vectors with unit-variance components has variance d_h, so at d_h = 128 the raw scores have standard deviation around 11. Feed that to softmax and it saturates: one value near 1, the rest near 0, and the gradient through softmax vanishes. Scaling by 1/sqrt(d_h) brings the variance back to 1 so softmax operates in its useful range.

"How big is the KV cache and what controls it?"

2 x H_kv x d_h x bytes_per_element x L per token. For an 8B GQA model with H_kv = 8, d_h = 128, fp16, 32 layers, that is 128 KB per token, so 1.07 GB for an 8k context. The levers are H_kv (GQA versus MHA is a 4x difference), context length (linear), and KV quantisation to int8 (2x). It matters because KV cache is what limits concurrent requests: on an 80 GB GPU with a 16 GB model, it is the difference between 54 and 13 concurrent requests.

"Where do the parameters actually live?"

Mostly the FFN. With d = 4096 and d_ff = 14336, the three FFN matrices are 176 M parameters per layer against attention's 42 M under GQA, so the FFN is roughly 72 percent of each layer. Attention gets the conceptual attention and the feed-forward network is where the model's capacity mostly sits.

"Why is decode memory-bound and prefill compute-bound?"

Shapes. Prefill processes [B, S, d] with S in the thousands, so every operation is a large matrix-matrix product with high arithmetic intensity: many FLOPs per byte loaded. Decode processes [B, 1, d], so every operation is matrix-vector: you read all 16 GB of weights to produce one token per sequence. On an H100 at 3.35 TB/s that is a 4.8 ms floor per step regardless of batch size, which is why batching improves throughput and not per-stream latency.

"How would you estimate training cost for this model?"

6 x N x D FLOPs, where N is parameters and D is training tokens. For 8B parameters on 2 trillion tokens: 6 x 8e9 x 2e12 = 9.6e22 FLOPs. An H100 at roughly 1e15 achievable FLOP/s (about half of peak dense fp16, which is realistic with good utilisation) gives 9.6e7 seconds, about 1,111 GPU-days, so roughly 46 days on 24 GPUs or 3 days on 384. The approximation ignores the attention term, which is fine at moderate sequence length and understates it at long context.

Common misconceptions

"Transformers are quadratic, so long context is impossible." The quadratic term is the S x S score matrix, which FlashAttention computes without materialising. The remaining costs of long context are the KV cache (linear in S) and the attention FLOPs (quadratic, but only dominant past roughly S ≈ d).

"Attention is where the parameters are." The FFN is about 72 percent of each layer's parameters in a modern GQA model. Attention is the conceptually interesting part and the smaller one.

"More heads means more parameters." The head count partitions the same d dimension: H x d_h = d. Going from 16 to 32 heads at fixed d changes nothing about parameter count, only how the dimension is split. GQA is different, because it genuinely shrinks the K and V projections.

"The KV cache is an optimisation you can skip." Without it, generating token n recomputes attention over all n-1 previous tokens from scratch, making generation quadratic in output length. It is not optional, and its size is the primary constraint on serving concurrency.

"The original paper describes current models." Post-norm became pre-norm, LayerNorm became RMSNorm, learned positions became RoPE, ReLU FFN became SwiGLU, and MHA became GQA. The block structure survived; most of the details did not.

Interview delivery note

Say this verbatim: "The shape that matters is [B, H, S, S] for the attention scores, and 2 x H_kv x d_h x L bytes per token for the KV cache. The first is what FlashAttention removes, and the second is what limits how many concurrent requests fit on a GPU, which is usually the number the business cares about." Two shapes, each tied to a consequence, rather than a recitation of the architecture.

The senior-versus-staff separator is going from shapes to a deployment size. A senior engineer can walk the tensor shapes correctly. A staff engineer computes that an 8B GQA model at 8k context needs about 1.07 GB of KV per request, that 58 GB of usable cache on an 80 GB GPU means 54 concurrent, and that switching to MHA or to 32k context each cost 4x the fleet. Being able to size a deployment before touching hardware is what makes the shapes worth knowing.

The second signal is naming the FFN as most of the parameters, and knowing where the quadratic term actually starts to dominate (S ≈ d, so around 4k for a 4096-dimensional model). Both are places where the folk understanding and the arithmetic disagree.

The Staff & Team Lead Interview Reference