Prototype course-owned Metal matrix tiles

[skyzh]

Why

Prototype the course's quantized-prefill and paged-prefill Metal paths without relying on hidden mlx::steel loader or MMA primitives, while preserving the existing algorithms and measuring the cost honestly.

Change

• add one fixed-shape cooperative loader with contiguous per-thread chunks, padded threadgroup storage, tail zero-fill, and the K transpose form;
• reuse it for Week 2 activations and Week 3 contiguous Q/K/V staging;
• replace Steel BlockMMA with a small direct-simdgroup_matrix wrapper;
• add ordinary/Split-K matrix-boundary coverage plus causal contiguous/manual-cross-page coverage;
• preserve 20 raw interleaved Steel↔zero benchmark files, the accepted comparison contract, frozen L9/L65 harness, hashes, and a concise comparison report.

Starter code, chapters, navigation, workflows, publication files, public APIs, dispatch, block-table translation, manual cross-page loading, QKᵀ/PV, masking, online softmax/rescaling, int4 dequantization, Split-K policy, and reduction are unchanged.

Review note

Experimental Draft only. Base is 1078c08130fa0648fe45f11bf32ca33a7ecfdd58 / tree 164b5f9c867b51458a6dc3a85c197fd3e4322774. Frozen head is 84f6f4f9e9ae0e8b3bfee458569b063464ea0e3c / tree 43300470ae50c5b99016b78d9725a00f32c9a91e; the measured implementation commit is 69c4df4aa187088bdf6b78f3f1fc46be6107a0e7 / tree 17dd72ea5e818cb394d11647a34d030f0ba123d7 (the successor commit only adds evidence).

Validation on Apple M4 Pro / Python 3.12.13 / MLX 0.32.0:

• clean refsol Metal rebuild;
• full reference suite: 489 passed, 8 intentional model skips;
• focused Python lint and format checks pass;
• required-solution mlx::steel source matches: zero;
• all eight paged-prefill commands match dense BF16 within 0.015625 max absolute error;
• the 247 protected base entries are byte-identical, manifest SHA-256 4ec26111a26adbe84a65da69e533911f14a8f94441b1a54ce61784085d22d3b6.

Accepted interleaved command-level comparison (positive latency delta is slower):

• Week 2 SIMD/split-K: +12.73%/+12.41% at context 128; +32.51%/+31.42% at context 512. Same-run MLX drift sentinel: +0.40%/+1.50%.
• Week 3 paged prefill: +3.46% at L9 and +3.33% at L65.
• End-to-end primary metric: serving prefill throughput -13.35%. Output throughput and requests/s are each -0.38%; the unrelated direct-decode metric is noisy and not evidence for the changed paths.

Recommendation: retain this as a readable correctness prototype, but do not mainline this exact implementation until the direct-MMA/quantized-prefill gap is closed under the same contract. No Ready, merge, release, tag, or deployment is requested.

AI-Assisted: GPT-5.6 Sol + Forge