Draft: Week 2 five-day integrated course (review in progress)

README.md

@@ -27,12 +27,10 @@ The course follows a four-week learning path:
 - **Week 1: From Matmul to Text.** Build a Qwen3 model directly from `mlx.core`
   array operations: attention, RoPE, GQA, RMSNorm, the MLP, sampling, and
   the autoregressive loop.
-- **Week 2: A Step Closer to vLLM.** Add a KV cache, establish a
-  synchronized MLX baseline, and let matched benchmarks choose each
-  optimization. The causal path moves from quantized decode matvec to fused
-  model kernels and SIMD-matrix prefill; decode attention is an optional
-  workload-conditioned lab, and split-K stays only where a measured short
-  shape supports it.
+- **Week 2: A Faster Single Request.** In five lessons, reuse the prefix,
+  bound request-owned KV storage, measure the same workload, pack W4 weights,
+  add SIMD matrix prefill and fused primitives, then tile dense prefill
+  attention. Run the cumulative `selected` checkpoint inside Day 5.
 - **Week 3: Build a Mini vLLM.** Introduce continuous
   batching and chunked admission, then make paged KV the canonical serving
   layout. Decode attention and FlashAttention learn to read pages directly so
@@ -109,13 +107,11 @@ one explicit byte range through the existing loop.
 | 1.5 | Load the Model | ✅ | ✅ | ✅ | ✅ |
 | 1.6 | Generate Responses (aka Decoding) | ✅ | ✅ | ✅ | ✅ |
 | 1.7 | Sampling | ✅ | ✅ | ✅ | ✅ |
-| 2.1 | KV Cache | ✅ | ✅ | ✅ | 🚧 |
-| 2.2 | Benchmarking and Profiling | ✅ | ✅ | ✅ | 🚧 |
-| 2.3 | Quantize the Model | ✅ | ✅ | ✅ | 🚧 |
-| 2.4 | Fused Model Kernels | ✅ | ✅ | ✅ | 🚧 |
-| 2.5 | SIMD-Matrix Prefill | ✅ | ✅ | ✅ | 🚧 |
-| 2.6 (optional) | Workload-Conditioned Operator Lab | ✅ | ✅ | ✅ | 🚧 |
-| 2.7 | Conditional Split-K and Final Decision | ✅ | ✅ | ✅ | 🚧 |
+| 2.1 | Cache and Measure | ✅ | ✅ | ✅ | 🚧 |
+| 2.2 | Keep W4 Packed | ✅ | ✅ | ✅ | 🚧 |
+| 2.3 | SIMD Matrix Prefill | ✅ | ✅ | ✅ | 🚧 |
+| 2.4 | Fused Model Primitives | ✅ | ✅ | ✅ | 🚧 |
+| 2.5 | Tiled Dense Prefill Attention | ✅ | ✅ | ✅ | 🚧 |
 | 3.1 | Continuous Batching | ✅ | ✅ | ✅ | 🚧 |
 | 3.2 | Chunked Prefill | ✅ | ✅ | ✅ | 🚧 |
 | 3.3 | Paged KV Cache | ✅ | ✅ | ✅ | 🚧 |

benches/bench.py

@@ -88,13 +88,14 @@ def parse_args() -> argparse.Namespace:
         "--week2-checkpoint",
         choices=(
             "kv-cache",
+            "capacity-cache",
             "quantized-matvec",
+            "simd-matmul",
             "rmsnorm",
             "rope",
             "swiglu",
-            "simd-matmul",
-            "decode-attention",
-            "split-k",
+            "tiled-prefill",
+            "selected",
         ),
         help="run one cumulative Week 2 end-to-end checkpoint",
     )

benches/bench_course_progression.py

@@ -42,59 +42,66 @@ class Throughput:
     WEEK1_VARIANT,
     Variant(
         "week2-kv-cache",
-        "2.1 KV cache",
+        "2.1 Reuse the prefix",
         "ref",
         "week2",
         ("--week2-checkpoint", "kv-cache"),
     ),
+    Variant(
+        "week2-capacity-cache",
+        "2.1 + Bound KV-cache movement",
+        "ref",
+        "week2",
+        ("--week2-checkpoint", "capacity-cache"),
+    ),
     Variant(
         "week2-quantized-matvec",
-        "2.3 Quantized matvec",
+        "2.2 Keep W4 packed",
         "ref",
         "week2",
         ("--week2-checkpoint", "quantized-matvec"),
     ),
+    Variant(
+        "week2-simd-matmul",
+        "2.3 SIMD matrix prefill",
+        "ref",
+        "week2",
+        ("--week2-checkpoint", "simd-matmul"),
+    ),
     Variant(
         "week2-rmsnorm",
-        "2.4 Fast RMSNorm",
+        "2.4 Compact RMSNorm",
         "ref",
         "week2",
         ("--week2-checkpoint", "rmsnorm"),
     ),
     Variant(
         "week2-rope",
-        "2.4 + Fast RoPE",
+        "2.4 + Compact RoPE",
         "ref",
         "week2",
         ("--week2-checkpoint", "rope"),
     ),
     Variant(
         "week2-swiglu",
-        "2.4 + Fused SwiGLU",
+        "2.4 + Compact SwiGLU",
         "ref",
         "week2",
         ("--week2-checkpoint", "swiglu"),
     ),
     Variant(
-        "week2-simd-matmul",
-        "2.5 SIMD matrix prefill",
-        "ref",
-        "week2",
-        ("--week2-checkpoint", "simd-matmul"),
-    ),
-    Variant(
-        "week2-decode-attention",
-        "2.6 Optional decode attention",
+        "week2-tiled-prefill",
+        "2.5 Tiled dense prefill attention",
         "ref",
         "week2",
-        ("--week2-checkpoint", "decode-attention"),
+        ("--week2-checkpoint", "tiled-prefill"),
     ),
     Variant(
-        "week2-split-k",
-        "2.7 Split-K prefill",
+        "week2-selected",
+        "2.5 + Run the selected inference engine",
         "ref",
         "week2",
-        ("--week2-checkpoint", "split-k"),
+        ("--week2-checkpoint", "selected"),
     ),
     MLX_VARIANT,
 )

benches/bench_week2_operators.py

@@ -172,7 +172,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--include-split-k",
         action="store_true",
-        help="also benchmark the Day 7 split-K path",
+        help="also benchmark the historical Split-K experiment (not a current checkpoint)",
     )
     parser.add_argument(
         "--json-output",

benches/profile_week2_kernels.py

@@ -25,13 +25,14 @@
 
 DEFAULT_CASES = (
     "kv-cache:decode:128",
+    "capacity-cache:decode:128",
     "quantized-matvec:decode:128",
-    "swiglu:decode:128",
     "simd-matmul:prefill:128",
-    "simd-matmul:prefill:32",
-    "decode-attention:decode:128",
-    "decode-attention:prefill:128",
-    "split-k:prefill:32",
+    "rmsnorm:decode:128",
+    "rope:decode:128",
+    "swiglu:decode:128",
+    "tiled-prefill:prefill:128",
+    "selected:prefill:128",
 )
 PROMPT_RULE = "synthetic-token-ids"
 PREFILL_LOGITS = "all"