Lab 02: Q/K/V Shape Walkthrough
Annotated code reading lab. Running code is optional.
Q/K/V Shape Walkthrough
This lab maps directly to the handbook section. Read the related handbook section first, then use the lab page and starter file to connect the concept to concrete variables, shapes, APIs, and interview-ready explanations.
Q/K/V Shape Walkthrough
Follow one hidden tensor through Q/K/V projection and head reshape. This prevents most attention shape bugs.
Mechanism to keep in mind
- `x_shape` starts as `(B,S,D)`.
- `qkv_shape` keeps D before splitting into heads.
- `head_shape` makes `H` and `Dh` explicit for score computation.
How shapes connect to attention
Q, K and V are projections of the hidden tensor. QK^T forms a score matrix over query and key sequence positions; dividing by sqrt(Dh) keeps score scale stable before softmax.
Starter preview
Excerpt from code/lab-02-qkv-shape-walkthrough/qkv_shapes.py. The linked starter file is the source of truth.
# Q/K/V Shape Walkthrough
# Annotated reading material. Running this file is optional.
# Source-of-truth focus: Follow one hidden tensor through Q/K/V projection and head reshape. This prevents most attention shape bugs.
batch, seq, hidden, heads = 2, 4, 16, 4
head_dim = hidden // heads
x_shape = (batch, seq, hidden)
q_shape = (batch, seq, hidden) # after Wq
q_heads = (batch, heads, seq, head_dim)
scores = (batch, heads, seq, seq) # Q @ K^T over head_dim
# What to explain while reading:
# - x_shape starts as (B,S,D).
# - qkv_shape keeps D before splitting into heads.
# - head_shape makes H and Dh explicit for score computation.
#
# Common traps:
# - Do not multiply parameter count by heads twice.
# - Batch, head and sequence order varies by library.
What each block is doing
- Setup / contract
- `x_shape` starts as `(B,S,D)`.
- Main transition
- `qkv_shape` keeps D before splitting into heads.
- Interview hook
- `head_shape` makes `H` and `Dh` explicit for score computation.
Reading checkpoints
- Head count partitions hidden width.
- Score shape contains both query and key sequence axes.
- GQA/MQA changes KV heads, not the basic Q reads K/V story.
What this lab prevents
- Do not multiply parameter count by heads twice.
- Batch, head and sequence order varies by library.
How to say it out loud
Follow one hidden tensor through Q/K/V projection and head reshape. This prevents most attention shape bugs. Then explain the code by naming the state being transformed, the axis or shape that matters, and the tradeoff that would appear in a real system.
Additional intuition
- Use official docs and papers for API behavior and factual claims; use blogs only to improve the mental picture.
- If support matrices, performance behavior or backend choices are version-sensitive, check current docs before repeating them.
- A strong interview answer names the state object, the shape or axis it changes, and the tradeoff it creates.