CourseTransformers: From Attention to Trained, Served and Interpreted Models · Module 2: Self-Attention Mechanics · part 18 of 111
Part 18 · Module 2: Self-Attention Mechanics

Lab: Package attention and test every claim

20 min read·13 Oct 2026
Module 2 · Lab · About 60 minutes · Needs: lessons 1 to 5 of this module, and Module 1's lab

What you are building

TINYFORMER-LAB AFTER MODULE 2 attention.py attention, masks, heads test_m02_attention 7 claims m02_kv_heads.py 4, 2, 1 kv heads runs/ record 9 measured runs v0.1.0: data.py, Module 1 git tag v0.2.0
The lab adds one package file (blue) guarded by seven tests (green), a race whose nine runs land in a record (amber), and a v0.2.0 tag on top of Module 1's state (grey).

Sol 406 at Aurora Station. The five lessons wrote attention in notebooks: naive_attention with its scale switch, five masks, MultiHeadAttention, GroupedQueryAttention and the cost formulas. Each lesson ended "the module lab saves this". Today you do: one file, tinyformer/attention.py, and one test file, tests/test_m02_attention.py, holding the checks the lessons ran by hand, from agreement with PyTorch to the causal leak test and the "same as alone" packing check. Then you answer the question Hana Aoki asked in Lesson 5, "should we keep fewer key and value heads?", with a measured race on the shared reverse task, three seeds per setting, written to a run record. You commit and tag v0.2.0. Every number is labelled measured here, cited or arithmetic.

The rest of this course is yours to keep

This course is bought on its own, once, and stays readable afterwards, including the parts added to it later.

Buy this course · ₹399See what is inside