Align Week 1 lessons with runnable checkpoints

book/src/week1-01-attention.md

@@ -1,8 +1,21 @@
 # Week 1 Day 1: Attention and Multi-Head Attention
 
-On Day 1, we will implement basic attention and multi-head attention. An attention layer processes an input sequence and
-weighs the relevance of its different positions when producing each output. Attention is a key building block of Transformer
-models.
+The starter provides `softmax` through MLX so that Day 1 can focus on the attention data flow. Your required work is to
+complete `scaled_dot_product_attention_simple` and `SimpleMultiHeadAttention` in `src/tiny_llm/attention.py`, plus the
+`linear` helper in `src/tiny_llm/basics.py`. Other attention functions in the starter are for later days and are not part
+of this chapter.
+
+Start by running the focused Task 1 tests. The command refreshes the supplied Day 1 test in `tests/` before running it:
+
+```console
+pdm run test --week 1 --day 1 -- -k task_1
+```
+
+The supplied `softmax` cases pass in the untouched starter, while the attention cases fail. This expected red checkpoint
+shows the behavior that your attention implementation must add.
+
+An attention layer processes an input sequence and weighs the relevance of its different positions when producing each
+output. Attention is a key building block of Transformer models.
 
 [📚 Reading: Transformer Architecture](https://huggingface.co/learn/llm-course/chapter1/6)
 
@@ -57,7 +70,9 @@ output: N.. x L x D
 scale = 1/sqrt(D) if not specified
 ```
 
-You may use MLX's `softmax`; we will revisit lower-level operations in Week 2.
+Use the supplied `softmax` helper for the required exercise. As an optional, ungraded bonus, replace its MLX call with
+your own numerically stable implementation: subtract the maximum value along `axis`, exponentiate the shifted values,
+then divide by their sum along the same axis. Preserve the helper's public API and output behavior.
 
 When this function is called from multi-head attention, the tensors will usually have these shapes:
 
@@ -72,12 +87,15 @@ mask: 1 x H x L x L
 The function itself operates on the last two dimensions and must support any number of leading batch dimensions. The mask
 only needs a shape that can broadcast to the attention-score shape.
 
-At the end of this task, you should be able to pass the following tests:
+Run the Task 1 checkpoint again after implementing the function:
 
 ```
 pdm run test --week 1 --day 1 -- -k task_1
 ```
 
+When this checkpoint turns green, your attention function supports arbitrary leading batch dimensions, optional masks,
+and default or explicit scaling.
+
 ## Task 2: Implement `SimpleMultiHeadAttention`
 
 In this task, we will implement the multi-head attention layer.
@@ -101,6 +119,15 @@ First, implement the `linear` function in `basics.py`. It takes a tensor of shap
 `O x I`, and an optional bias vector of shape `O`. Its output has shape `N.. x O`, where `I` is the input dimension and
 `O` is the output dimension.
 
+Use the focused linear tests as your next checkpoint:
+
+```console
+pdm run test --week 1 --day 1 -- -k test_task_2_linear
+```
+
+Before you implement `linear`, this checkpoint is red. When it turns green, `linear` supports optional bias across the
+tested precisions and devices.
+
 For `SimpleMultiHeadAttention`, the input tensors `query`, `key`, and `value` have shape `N x L x E`, where `E` is the
 embedding dimension for one token. The Q, K, and V projections each map `E` to `H x D`: `H` heads, each with dimension
 `D`. Reshape that final projection dimension into separate `H` and `D` dimensions.
@@ -127,16 +154,24 @@ output/input: N x L x E
 w_o: E x (H x D)
 ```
 
-At the end of the task, you should be able to pass the following tests:
+Run the Task 2 checkpoint after implementing the layer:
 
-```
+```console
 pdm run test --week 1 --day 1 -- -k task_2
 ```
 
+When this checkpoint turns green, your layer projects query, key, and value tensors into independent attention heads and
+recombines their outputs through the final projection.
+
 You can run all tests for the day with:
 
-```
+```console
 pdm run test --week 1 --day 1
 ```
 
+When the full Day 1 suite turns green, you have a standalone multi-head attention layer that projects Q/K/V, evaluates
+each head independently, and recombines the result. Day 3 will generalize this attention mechanism to grouped-query
+attention for Qwen3; the `SimpleMultiHeadAttention` layer built here is a standalone exercise, not the model's exact call
+path.
+
 {{#include copyright.md}}

scripts/dev-tools.py

@@ -51,7 +51,7 @@ def test(args):
                     return status
                 targets.append(f"tests/test_week_{args.week}_day_{day}.py")
             return pytest.main(["-v", *targets] + args.remainders)
-        copy_test(args, skip_if_exists=True)
+        copy_test(args, force=True)
         return pytest.main(
             ["-v", f"tests/test_week_{args.week}_day_{args.day}.py"] + args.remainders
         )

src/tiny_llm/basics.py

@@ -3,7 +3,7 @@
 
 
 def softmax(x: mx.array, axis: int) -> mx.array:
-    # TODO: manual implementation
+    # Supplied for Day 1; a manual implementation is an optional bonus exercise.
     return mx.softmax(x, axis=axis)
 
 

src/tiny_llm_ref/basics.py

@@ -3,7 +3,7 @@
 
 
 def softmax(x: mx.array, axis: int) -> mx.array:
-    # TODO: manual implementation
+    # Supplied for Day 1; a manual implementation is an optional bonus exercise.
     return mx.softmax(x, axis=axis)
 
 

tests_refsol/test_week_1_day_1.py

@@ -21,7 +21,9 @@ def test_task_1_softmax(stream: mx.Stream, precision: mx.Dtype):
 @pytest.mark.parametrize("stream", AVAILABLE_STREAMS, ids=AVAILABLE_STREAMS_IDS)
 @pytest.mark.parametrize("precision", PRECISIONS, ids=PRECISION_IDS)
 @pytest.mark.parametrize(
-    "batch_dimension", [0, 1, 2], ids=["batch_0", "batch_1", "batch_2"]
+    "batch_dimension",
+    [0, 1, 2, 3],
+    ids=["batch_0", "batch_1", "batch_2", "batch_3"],
 )
 def test_task_1_simple_attention(
     stream: mx.Stream, precision: mx.Dtype, batch_dimension: int
@@ -34,8 +36,10 @@ def test_task_1_simple_attention(
         if batch_dimension == 0:
             BATCH_SIZE = ()
         elif batch_dimension == 1:
-            BATCH_SIZE = (2, 3)
+            BATCH_SIZE = (2,)
         elif batch_dimension == 2:
+            BATCH_SIZE = (2, 3)
+        elif batch_dimension == 3:
             BATCH_SIZE = (2, 3, 3)
         DIM_L = 4
         DIM_D = 5
@@ -64,7 +68,9 @@ def test_task_1_simple_attention(
 @pytest.mark.parametrize("stream", AVAILABLE_STREAMS, ids=AVAILABLE_STREAMS_IDS)
 @pytest.mark.parametrize("precision", PRECISIONS, ids=PRECISION_IDS)
 @pytest.mark.parametrize(
-    "batch_dimension", [0, 1, 2], ids=["batch_0", "batch_1", "batch_2"]
+    "batch_dimension",
+    [0, 1, 2, 3],
+    ids=["batch_0", "batch_1", "batch_2", "batch_3"],
 )
 def test_task_1_simple_attention_scale_mask(
     stream: mx.Stream, precision: mx.Dtype, batch_dimension: int
@@ -76,8 +82,10 @@ def test_task_1_simple_attention_scale_mask(
         if batch_dimension == 0:
             BATCH_SIZE = ()
         elif batch_dimension == 1:
-            BATCH_SIZE = (2, 3)
+            BATCH_SIZE = (2,)
         elif batch_dimension == 2:
+            BATCH_SIZE = (2, 3)
+        elif batch_dimension == 3:
             BATCH_SIZE = (2, 3, 3)
         DIM_L = 4
         DIM_D = 5
@@ -107,6 +115,27 @@ def test_task_1_simple_attention_scale_mask(
             )
             assert_allclose(user_output, reference_output, precision=precision)
 
+        if batch_dimension == 2:
+            mask = mx.array(
+                [
+                    [0.0, -0.5, -1.0, -1.5],
+                    [-1.5, 0.0, -0.5, -1.0],
+                    [-1.0, -1.5, 0.0, -0.5],
+                    [-0.5, -1.0, -1.5, 0.0],
+                ],
+                dtype=precision,
+            )
+            scores = mx.matmul(query, key.swapaxes(-2, -1)) * scale + mask
+            reference_output = mx.matmul(mx.softmax(scores, axis=-1), value)
+            user_output = scaled_dot_product_attention_simple(
+                query,
+                key,
+                value,
+                scale=scale,
+                mask=mask,
+            )
+            assert_allclose(user_output, reference_output, precision=precision)
+
 
 @pytest.mark.parametrize("stream", AVAILABLE_STREAMS, ids=AVAILABLE_STREAMS_IDS)
 @pytest.mark.parametrize("precision", PRECISIONS, ids=PRECISION_IDS)
@@ -120,12 +149,13 @@ def test_task_2_linear(stream: mx.Stream, precision: mx.Dtype):
             w = mx.random.uniform(shape=(DIM_Y, DIM_X), dtype=precision)
             b = mx.random.uniform(shape=(DIM_Y,), dtype=precision)
             user_output = linear(x, w, b)
-            if precision == mx.float16 and stream == mx.cpu:
-                # unsupported
-                break
             reference_output = mx.addmm(b, x, w.T)
             assert_allclose(user_output, reference_output, precision=precision)
 
+            user_output = linear(x, w)
+            reference_output = mx.matmul(x, w.T)
+            assert_allclose(user_output, reference_output, precision=precision)
+
 
 @pytest.mark.parametrize("stream", AVAILABLE_STREAMS, ids=AVAILABLE_STREAMS_IDS)
 @pytest.mark.parametrize("precision", PRECISIONS, ids=PRECISION_IDS)