Repair Week 1 Day 3 GQA learner contract

[skyzh]

Why

Week 1 Day 3 did not fully bind the public grouped-attention and Qwen3 attention contract. The chapter left multi-query attention, arbitrary leading dimensions, independent head dimensions, additive-mask forwarding, mask dtype, the causal S >= L domain, Q/K normalization order, and the float32 attention/cast-back boundary implicit. The supplied tests therefore allowed several learner-visible API and numerical errors to pass undetected.

Change

• Explain the MHA/GQA/MQA head relationships, arbitrary positive leading axes, independent head_dim, additive masks, requested mask dtype, and the causal requirement that key/value length is at least query length.
• Specify that Q/K RMSNorm precedes non-traditional RoPE, attention arithmetic runs in float32, and its result returns to the model dtype before the output projection.
• Make the completed grouped-attention path reject causal L > S before it can produce NaNs.
• Add deterministic public behavioral witnesses for MQA and non-divisible heads, BF16 mask dtype, causal-domain rejection, Qwen batch and dimension handling, nonuniform norm weights, additive masking, and BF16 promotion/castback against an independent public-operation oracle.

Review

• Scope from the PR base is exactly three modified existing paths: book/src/week1-03-gqa.md, src/tiny_llm_ref/attention.py, and tests_refsol/test_week_1_day_3.py (+254/-9), with no add/delete/rename/mode change.
• The completed/reference Day 3 checkpoint passes all 84 supplied cases; the untouched starter fails all 84 as intended. The local full reference suite is 527 passed / 8 skipped after the documented extension builds; the exact successor head's hosted check must be terminal green before delivery.
• Twelve causal mutations are rejected, covering MQA, non-divisible grouped heads, arbitrary leading axes across additive, causal, and unmasked attention, mask dtype, causal L > S, batch size 2, independent head_dim, norm/RoPE order, additive masks, float32 attention promotion, and castback.
• Installation, starter extension, copied-test freshness, Ruff formatting for changed Python, mdBook, sitemap, and diff checks pass. Both trees retain 331 records and two symlinks; the 328-record protected complement is unchanged.
• Supplied grading uses only public imports, APIs, MLX operations, numeric outputs, shapes, dtypes, and observable rejection. It does not parse source or grade private names, fields, loops, implementation shape, error wording/type, or whether K/V grouping is materialized.
• Starter code and public signatures, already-correct Qwen behavior, Day 1/2 and Day 4+, Week 4, tooling, lockfiles, navigation, publication configuration, and unrelated bytes are unchanged. Ordinary course merge occurs only after the written exact-head gates close and will trigger the repository's existing Pages workflow.

AI-Assisted: GPT-5.6 Sol + Forge