Lab 01 — Transformer Architectural Variants

Phase: 02 — Model Architecture Mastery | Difficulty: ⭐⭐⭐⭐☆ | Time: 3–4 hours

Read ../HITCHHIKERS-GUIDE.md before starting.

What you build

Implement and compare four key attention variants used in production LLMs:

  • Multi-Head Attention (MHA) — original Transformer (Vaswani et al.)
  • Multi-Query Attention (MQA) — single KV head, multiple Q heads (PaLM, Falcon)
  • Grouped-Query Attention (GQA) — G KV groups (Llama 2/3, Mistral)
  • RoPE positional embeddings — rotary position encoding used in Llama/GPT-NeoX

Key concepts

ConceptWhat to understand
KV cacheMQA/GQA reduce memory for long context inference
Head groupingG groups share KV, reducing VRAM at minimal accuracy cost
RoPERotates Q and K before dot product — relative positions for free
Flash-attention patternWhy attention needs tiling for memory efficiency

Files

FilePurpose
lab.pyStubs with TODO markers
solution.pyFull working implementation + smoke test
test_lab.pypytest suite
requirements.txtDependencies

Run

pip install -r requirements.txt
python lab.py
pytest test_lab.py -v