1. What is it?#
A tensor is an n-dimensional array. Broadcasting is the rule that lets operations between different-shaped tensors work without explicit copying.
(3, 1) + (1, 4) → (3, 4) broadcasting expands bothShape fluency is the practical skill this file builds. Most of the bugs you write in model code will be shape bugs, and most of the code you read will be comprehensible the moment you can name each axis.
2. Why does it exist?#
Broadcasting exists so that “add this bias vector to every row” doesn’t require materializing a copy of the bias for every row. It saves memory and, more importantly, saves memory bandwidth — the broadcast value is read once and reused from registers.
3. Simple analogy#
A stencil. You have one stencil (the bias vector) and 10,000 sheets of paper (the batch). You don’t make 10,000 copies of the stencil; you reuse the one you have. Broadcasting is the array-programming version of reusing the stencil.
4. Tiny example#
package main
import "fmt"
func main() {
X := [][]int{
{1, 2, 3},
{4, 5, 6},
} // (2, 3)
b := []int{10, 20, 30} // (3,)
// "Broadcasting": b is reused for every row. No copy of b is ever made.
for _, row := range X {
for j := range row {
row[j] += b[j]
}
}
fmt.Println(X) // [[11 22 33] [14 25 36]]
}
Broadcasting rules, applied right to left:
Align shapes from the RIGHT. For each dimension pair:
- equal → fine
- one of them is 1 → stretch it
- otherwise → ERROR
(2, 3) (2, 3) (2, 1, 3) (5, 4)
+ ( 3) + (2, 1) + ( 4, 1) + ( 3)
--------- --------- ------------ ---------
(2, 3) (2, 3) (2, 4, 3) ERRORThe classic trap:
a := []int{1, 2, 3} // shape (3,) — a row
b := [][]int{{1}, {2}, {3}} // shape (3, 1) — a column
// A broadcasting library "adds" these by stretching both: every b[i] meets every a[j].
out := make([][]int, len(b))
for i := range b {
for j := range a {
out[i] = append(out[i], a[j]+b[i][0])
}
}
fmt.Println(out) // [[2 3 4] [3 4 5] [4 5 6]] — shape (3, 3). Probably not what you wanted!
You wanted elementwise addition of two length-3 vectors and got a 3×3 outer sum. In model code this produces a plausible-looking tensor and silently wrong results. Guard against it by asserting shapes.
5. Technical explanation#
The shapes of LLM inference (memorize these)#
input_ids (B, S) int64
embeddings (B, S, d)
───── residual stream ─────
after q_proj (B, S, h·d_head) → view → (B, S, h, d_head) → transpose → (B, h, S, d_head)
after k_proj (B, S, h_kv·d_head) → (B, h_kv, S, d_head)
KV cache (num_blocks, block_size, h_kv, d_head) [paged]
or (B, h_kv, S_max, d_head) [contiguous]
attention scores (B, h, S_q, S_k) ← never materialize this if you can help it
attention output (B, h, S, d_head) → (B, S, d)
ffn gate/up (B, S, d_ff)
ffn down (B, S, d)
final hidden (B, S, d)
logits (B, S, V) during training/prefill-all
(B, 1, V) during decode / last-position-only prefillReshape, view, transpose, permute — and which ones cost#
t = torch.randn(2, 8, 64)
t.view(2, 512) # free IF contiguous; errors otherwise
t.reshape(2, 512) # free if possible, silently COPIES if not
t.transpose(1, 2) # free — swaps strides, result is non-contiguous
t.permute(2, 0, 1) # free — generalized transpose
t.contiguous() # COPIES — full read + write of the tensor
The pattern that costs you: transpose (free) followed by view (needs contiguous) forces
a contiguous() copy. In attention code this happens per layer per step. Watch for it in
profiles as unexplained copy_ kernels.
The idiomatic attention reshape:
# (B, S, d) → (B, h, S, d_head)
q = q.view(B, S, h, d_head).transpose(1, 2)
# after attention: (B, h, S, d_head) → (B, S, d)
out = out.transpose(1, 2).contiguous().view(B, S, d) # ← this .contiguous() is a real copy
Fused attention kernels (FlashAttention) accept and return layouts that avoid some of these.
Head layout: [B, h, S, D] vs [B, S, h, D]#
This is a genuine engineering decision with performance consequences:
[B, h, S, D] — "head-major". Each head's sequence is contiguous.
Good for: attention math per head.
Standard in most reference implementations.
[B, S, h, D] — "sequence-major". All heads for one token are contiguous.
Good for: appending a new token's KV (one contiguous write),
paged KV cache blocks.
Used by many fused kernels and paged attention.vLLM’s paged KV cache uses a layout chosen so that the attention kernel’s reads are coalesced (Section VI.09) and a new token’s K and V land contiguously. This is not a detail — it is a measurable fraction of decode performance.
6. Under the hood#
Broadcasting is not implemented by copying. The kernel computes indices with stride 0 along broadcast dimensions:
b has shape (3,) with stride (1,)
broadcast to (2,3): shape (2,3), stride (0, 1)
↑
stride 0 means "don't move when i changes"So the same 3 values are read repeatedly — and because they’re read repeatedly, they stay in cache/registers. Broadcasting is genuinely cheap.
Where it stops being cheap: broadcasting into a very large intermediate. (B,1,S,S) * (B,h,1,1) materializes (B,h,S,S) — the very tensor you were trying to avoid.
7. Performance implications#
- Reshapes and transposes are free;
contiguous()is not. Eachcontiguous()on a(B,h,S,d)tensor is a full read+write of ~B·S·d·2bytes. - Broadcasting avoids materialization — good. But check that a broadcast isn’t creating a huge intermediate.
- Layout determines coalescing. The same kernel can be 2-5x faster with a friendly layout.
- Views prevent some fusions. A non-contiguous tensor may force a compiler or kernel to fall back to a slower generic path.
8. Production implications#
- Assert shapes in custom code.
assert x.shape == (B, S, d), x.shapecatches broadcasting bugs immediately instead of three layers later. - Log shapes on error paths. When a request fails, the shape is usually the clue.
- Choose the KV layout to match your attention kernel, not to match the reference implementation.
- Avoid
.contiguous()in the hot path when a kernel accepting strided input exists.
9. Common mistakes#
Silent broadcasting. (3,) + (3,1) giving (3,3). Assert shapes.
Using reshape when you meant view. reshape silently copies; you lose a performance
guarantee and gain a hidden allocation.
Transposing then viewing. Forces a copy.
Mixing up (B,S,h,D) and (B,h,S,D). Produces valid shapes and wrong results.
Forgetting the sequence dimension during decode. Decode has S=1, and code written
assuming S>1 may index incorrectly.
Materializing (B,h,S,S). At S=8192, B=8, h=32, FP16 that is 32 GB per layer.
10. Hands-on exercise#
A. Broadcasting drills. For each pair, state the result shape or ERROR without running it:
(5,4)+(1,); (5,4)+(4,); (5,4)+(5,); (15,3,5)+(15,1,5);
(8,1,6,1)+(7,1,5); (2,3)+(3,2); (B,h,S,S)*(B,1,1,S)Then verify all of them.
B. Cost of contiguous. Time t.transpose(1,2).contiguous() vs t.clone() for
t = torch.randn(8, 32, 4096, 128, device='cuda'). Are they the same cost? Why?
C. Layout benchmark. Write a simple attention-score computation for both [B,h,S,D] and
[B,S,h,D] layouts. Time them. Which is faster on your hardware and why?
D. Annotate a real model. Take a transformer implementation (e.g. nanoGPT’s model.py) and
add a shape comment to every line of the forward pass. This is the single most useful hour you
can spend before Section V.
11. Interview questions#
- State the broadcasting rules.
- Which of view/reshape/transpose/permute/contiguous copy memory?
- Why does
q.view(B,S,h,D).transpose(1,2)not copy, but the reverse operation does? - Compare
[B,h,S,D]and[B,S,h,D]KV layouts. When would you choose each? - Give the shape of every intermediate tensor in one transformer layer.
- How large is the attention score tensor at B=8, h=32, S=8192, FP16? Why does that matter?
12. Further reading#
- [REFERENCE] NumPy broadcasting documentation
- [REFERENCE] PyTorch tensor views documentation; ezyang’s “PyTorch internals”
- Next: 04 — Neural networks from scratch