Lab 09: Attention Backend Comparison Reading
Annotated code reading lab. Running code is optional.
Attention Backend Comparison Reading
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.
Attention Backend Comparison Reading
Compare math-level attention with kernel/runtime choices: SDPA, FlashAttention and PagedAttention.
Mechanism to keep in mind
- `math` names the formula.
- `kernel` names how the formula is executed.
- `runtime` names serving-time cache/request management.
Starter preview
Excerpt from code/lab-09-attention-backend-comparison/attention_backends.md. The linked starter file is the source of truth.
A Transformer block turns token ids into vectors, mixes context with attention, applies per-token nonlinear transformations, and uses residual and normalization layers to keep deep training stable.What each block is doing
- Setup / contract
- `math` names the formula.
- Main transition
- `kernel` names how the formula is executed.
- Interview hook
- `runtime` names serving-time cache/request management.
Reading checkpoints
- FlashAttention is exact attention with lower HBM traffic for supported layouts/backends.
- SDPA is an API that may choose different kernels.
- Transformers attention backend option names and defaults are version-sensitive and model-specific.
- PagedAttention is about KV cache blocks in serving.
What this lab prevents
- Do not call FlashAttention an approximation.
- Do not confuse kernel backend with serving scheduler.
How to say it out loud
Compare math-level attention with kernel/runtime choices: SDPA, FlashAttention and PagedAttention. 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.
