Lab 03 — Transformer Performance Bug Hunt

Phase: 07 — Profiling & Performance Engineering | Difficulty: ⭐⭐⭐⭐⭐ | Time: 3–4 hours

Diagnose and fix 5 common Transformer performance anti-patterns seen in production code.

What you build

Identify and fix these anti-patterns in a buggy transformer implementation:

  1. Incorrect attention mask causing O(T²) work even for padding tokens
  2. Missing .contiguous() before reshape causing unnecessary copies
  3. CPU-GPU sync via .item() inside the forward loop
  4. Redundant recomputation of sqrt(d_k) on every call
  5. Inefficient KV projection using separate linear layers instead of one

Key concepts

ConceptWhat to understand
Memory contiguityNon-contiguous tensors require copy before CUDA kernels
.item() syncBlocks GPU pipeline — never call in inference loop
Flash-attention maskCausal mask should be precomputed and registered as buffer
GEMM batchingOne [d, 3*d_k] projection is faster than three [d, d_k] projections

Files

FilePurpose
lab.pyBuggy transformer implementation with TODO: fix markers
solution.pyFixed implementation with explanations
test_lab.pypytest suite checking correctness and performance
requirements.txtDependencies

Run

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