# Lab 028 lab detail: Off NVIDIA, you fix the platform yourself

The raw working record behind Lab 028 on labs.layer2c.com, derived from the author's run notes. Times are US Central.

**Boundary.** This file carries the substrate: build steps, the patch and what each change does, every measurement and failure, and the scripts by name. The patch itself is public at github.com/kltownsend/mlx (branch `sdpa-windowed-backward`). Lab 18's corpus contents and the trained adapters stay proprietary, as do the verdict's judgment calls.


Working notes for a candidate Layer2C lab. Keith rules on framing and sufficiency later through the layer2c-lab skill. Times are US Central.

Lab question (draft, Keith's to rule on): what does it take for an owner to fix the platform gap themselves, and does the fix make the Studio a stock venue for Lab 027's job?

## 2026-10-21

**Starting point.** Lab 027 found that a 30,553-token Gemma 4 31B LoRA step can't run on the Studio's stock mlx-lm. On mlx 0.32.4 the sliding-window layers get a boolean mask, and the backward falls back to the full score matrix. Measured peak memory for `mx.grad` through SDPA (`sdpa_vjp_probe.py`): causal D=512/4 KV heads 1.97 GB at 8,192 tokens and 3.66 GB at 16,384; boolean causal+1,024-window D=256/16 KV heads 12.94 GB and 50.25 GB. The full step hit a Metal OOM at a 224 GB footprint. The Lab 027 workaround is `finetune-model-first-tools/blocked_attention.py`.

**Upstream state, checked about 8:00 PM.** ml-explore/mlx#4734 (jagrit06, sliding-window forward kernels) is open, approved by zcbenz at 7:43 PM, not merged. Its `has_fused_kernel` returns "window support for backward attention is NYI" when `output_logsumexp && window_size > 0`, on Metal and CUDA. On CUDA, any `window_size` is NYI, forward included. No issue or PR for a windowed backward exists. No maintainer reply had come to Keith's #4734 comment before he deleted it.

**What the code says.** The #4563 backward is `sdpa_vjp_blocked` in `mlx/backend/metal/scaled_dot_product_attention.cpp`: a host-side loop over BQ x BK tiles (default 1,024, env `MLX_SDPA_VJP_BLOCK`), matmuls per tile, plus `sdpa_vjp_ds` in `scaled_dot_product_attention_vjp.metal` to form P and dS per element. It already skips tiles above the causal diagonal (`j0 > i0 + bq_len - 1 + diag_off`) and zeroes elements above it inside the kernel.

**Comment deleted (shortly after 8:00 PM, before the 8:08 PM canary).** Keith deleted his #4734 comment, after I flagged MLX's AGENTS.md, which bars agents from writing comments, PR descriptions or commit messages for the user, or pushing, on penalty of a contributor ban. Verified gone (API 404). The issue draft no longer links it and carries the probe inline. Open for Keith: anything posted upstream (issue, PR description, commit messages, comments) is his to write in his own words; I prepare code, tests and a technical changelog only.

**Issue draft.** `finetune-model-first-tools/mlx-issue-windowed-sdpa-backward.md`. Not posted; Keith decides.


## Plan for the build (approved by Keith; started 7:54 PM)

1. Clone ml-explore/mlx to `~/src/mlx` on the Studio, check out #4734's head, build from source into its own venv (leave the Lab 027 venv alone). Record commit SHAs and build time.
2. Baseline on that build: rerun the probe with a boolean mask and with `window_size=1024`. Expect both to fall back and go quadratic.
3. Forward: let the windowed steel attention kernels write `lse`; drop the Metal NYI check.
4. Backward: add `window_size` to `sdpa_vjp_blocked` (skip tiles entirely below the window) and to `SDPAVJPTileParams` / `sdpa_vjp_ds` (zero P where `j < i + diag_off - window_size + 1`).
5. Gradient parity against the dense boolean-mask fallback on small shapes: bf16/fp16/fp32, head dims 64 to 512, GQA ratios, windows that do and don't divide the tile size, `qL < kL`.
6. Probe goes linear for the windowed case.
7. mlx-lm: local patch so Gemma 4 sliding layers pass `window_size` instead of a boolean mask. Then the 30,553-token step on Lab 027's inputs (`~/larger-job/shared/`, `~/models/gemma-4-31B-it-mlx-bf16`).

8. Distribution (Keith's ruling): fork ml-explore/mlx under kltownsend, put the change on a clearly named branch on top of #4734 (e.g. `sdpa-windowed-backward`), keep the MIT license and upstream attribution intact, and add a PR-style description: what changed, how it was tested, the AI-assisted disclosure. Then draft a short comment for #4734 (or the issue, if Keith posts it) linking the branch. Prepare both; Keith approves the push and the comment before either happens. Open-bench rule: the code, tests and probe are public; the lab's judgment stays on labs.layer2c.com.

All of it is owned hardware, $0. Every kernel change gets a plain-language explanation here for Keith's review, since he's accountable for the code under MLX's CONTRIBUTING.

## Build log, 2026-10-21 (Studio, owned hardware, $0)

**Toolchain.** Xcode 27.1 is installed but `xcode-select` points at the Command Line Tools, so `xcrun metal` fails by default. `DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer` fixes it without sudo; cmake 4.1.1 is in `/opt/homebrew/bin`, which isn't on the non-interactive PATH. Build scripts: `~/src/build_mlx.sh` (#4734 baseline into `~/venvs/mlx-src`) and `~/src/build_mlx_win.sh` (patched worktree `~/src/mlx-win`, branch `sdpa-windowed-backward`, into `~/venvs/mlx-win`). Both build from #4734 head `fdf1d256a7dc` (2026-10-20), base `bff10220c742` on main. Each editable build took about 2 minutes. The Lab 027 venv (`~/venvs/mlx`, mlx 0.32.4) is untouched.

**Baseline: #4734 alone doesn't fix training (7:57 PM).** `sdpa_win_probe.py`, peak GB above inputs, bf16, 32 query heads:

| path | 8,192 | 16,384 | 32,768 |
|---|---|---|---|
| `"causal"`, D=512, 4 KV heads | 1.97 | 3.66 | 7.04 |
| boolean causal+1,024 mask, D=256, 16 KV heads | 12.94 | 50.25 | 198.0 |
| #4734 `window_size=1024` (no patch) | 12.94 | 50.5 | not run |
| patched `window_size=1024` | 1.24 | 2.49 | 4.74 |

So the `window_size` API in #4734 lands on the same quadratic fallback in training. With the patch, the windowed case doubles with length, like causal.

**What I changed, in plain terms** (`mlx-windowed-backward-tools/sdpa-windowed-backward.patch`, 6 files, +107/-15, uncommitted in `~/src/mlx-win`):

1. *Forward: one line removed.* #4734's Metal `has_fused_kernel` refused the fused forward whenever training needed the logsumexp (the per-row softmax normalizer the backward reuses) and a window was set. I read the two steel attention kernels in #4734 (regular and NAX): both already apply the window and both already write the logsumexp. The guard was the only thing in the way, so it's gone on Metal. CUDA keeps its guard; #4734 has no CUDA window kernel at all.
2. *Backward gets told about the window.* The backward primitive (`ScaledDotProductAttentionVJP`) only knew scale, causal and sinks. It now carries `window_size` too, in its constructor, its cache key (`state`, `is_equivalent`) and the call in `fast.cpp`, so two graphs with different windows can't share a compiled backward.
3. *Skip whole tiles outside the window.* The backward walks the attention matrix in tiles (1,024 x 1,024 by default) and already skipped tiles entirely above the causal diagonal. It now also skips tiles entirely more than `window_size` behind the diagonal, and, for a non-causal window, tiles entirely more than `window_size` ahead. That's where the memory and time savings come from: a query tile now touches two or three key tiles instead of every one before it.
4. *Mask inside the edge tiles.* The per-element kernel (`sdpa_vjp_ds`) recomputes the softmax probability P and zeroes it above the diagonal. It now also zeroes P where the key is `window_size` or more positions from the query's diagonal. One new field in the tile parameters struct, one comparison.
5. *Two bookkeeping fixes the causal code never needed.* (a) dQ for a query tile was "set" on key tile 0 and "added" on every later one. With a window, tile 0 is skipped for most queries, so the first visit was adding to uninitialized memory. Now each query tile tracks its own first visit. (b) When queries are a suffix of the keys (`qL < kL`), keys behind every query's window are never visited, so their dK and dV were never written. The backward now zeroes dK and dV first, only when a window is set and `qL < kL`.
6. *Guard and test.* CUDA's backward throws "NYI: window attention on CUDA" if it ever gets a window. New `test_sdpa_window_grad` in `python/tests/test_fast_sdpa.py` checks q/k/v gradients against the reference attention for windows 7, 128 and 1,000, causal and not, including `qL < kL` and multi-tile lengths.

**Gradient parity (about 8:00 PM).** `sdpa_win_parity.py`: 867 cases (fp32/fp16/bf16, head dims 64/128/256 plus 512 in bf16, GQA 8/8 and 8/2, lengths 300 to 3,000 including `qL < kL`, windows 1 to 5,000, causal and not), at tile sizes 256 and 1,024. Each case's error against an fp32 reference is held to within 2x the error of the already-shipped causal tiled backward on the same shape, plus a dtype floor. Result (`results/parity_b256.txt`): every window from 7 up passes in all three dtypes except one bf16 case (`qL=1000, kL=3000`, window 128, non-causal) that missed by 0.0004 against a fallback that itself sits at 0.014. Window 1 fails in fp32 and bf16. At window 1 the true dQ and dK are exactly zero, so what's measured is noise. `w1_check.py`: with `MLX_ENABLE_TF32=0` the fp32 fused gradients come out exactly 0.00000; with TF32 on (MLX's default) they're 0.006, and the shipped causal backward shows the same noise (0.001) on its one single-key row. bf16 noise (0.03) doesn't change with TF32 and matches the causal row's kind (0.011). My reading is that it's the existing backward's precision showing through a zero-valued gradient, not the window logic. Upstream's whole `test_fast_sdpa` passes on the patched build (36 tests, 2 skipped). pre-commit's clang-format then reformatted the C++ (whitespace only); the tests ran on the binary built before that.

**Model step (8:08 to 8:15 PM).** mlx-lm 0.32.2 stock plus `window_sdpa.py`, which swaps the sliding layers' boolean mask for `mask="causal", window_size=1024` on calls without a KV cache. Trainer: `studio_train_window.py`, a copy of Lab 027's `studio_train.py` with `--window-sdpa` and a forward-only `baselong` mode. Runner: `run_window.sh`. Same longest example (k=60, 30,553 tokens, 637 labeled), same shared init, all on the patched build:

| path | step | loss | grad norm | MLX peak | footprint |
|---|---|---|---|---|---|
| window patch, canary | 95.4 s | 1.55087 | 91.92 | 141.4 GB | 182.2 GB |
| Lab 027 blocked (2,048), canary | 114.8 s | 1.57930 | 96.96 | 142.5 GB | 205.1 GB |
| forward only: window patch | 27.5 s | 1.5508651853 | | 72.2 GB | 90.0 GB |
| forward only: stock boolean mask | 49.7 s | 1.5508651853 | | 72.4 GB | 91.8 GB |
| forward only: blocked | 29.8 s | 1.5498202098 | | 75.1 GB | 94.3 GB |

The blocked canary reproduces Lab 027's 0.32.4 run to the last digit (loss 1.5793049, grad norm 96.964), so the build change doesn't move Lab 027's numbers. The window patch's forward is bit-identical to stock mlx-lm's. Its training-step loss equals that forward loss; the blocked path's training forward sits 0.028 above stock. GPU power during the windowed step was about 145 to 149 W.

## fp32 gradient truth (Keith asked, 2026-10-21 evening)

**Model level, k = 53 (5,085 tokens).** Lab 027's CPU fp32 truth for this example already exists (`grad_truth_k53.safetensors`: loss 1.13935, LoRA-B grad norm 536.4). `gpu_vs_truth_window.py` is Lab 027's `gpu_vs_truth.py` with the window path added, B = 0, dropout 0, shared init, all on the patched build (`results/gvt_window_k53.log`):

| Path (Studio GPU, bf16) | Loss | Grad norm | Cosine to truth |
|---|---|---|---|
| CPU fp32 truth | 1.13935 | 536.4 | 1 |
| stock mlx-lm (boolean mask, dense fallback) | 1.13432 | 447.6 | -0.056 |
| Lab 027 blocked, 2,048 | 1.10497 | 39.6 | 0.141 |
| window patch | 1.11306 | 2,667.8 | -0.037 |

Stock and blocked match Lab 027's 0.32.4 numbers (-0.056, 0.141). All three bf16 paths are near noise against the truth at this length, as Lab 027 found, so this comparison doesn't separate them. The window patch's grad norm of 2,668 (5x truth) is the outlier. Lab 027 already showed norms swing widely between bf16 paths here (39.6 to 447.6), so I don't read it as a bug on its own. It's why I ran the layer check below.

**Layer level, k = 53.** `layer_grad_check.py` captures real q, k, v for the first six layers in an inference pass and checks each SDPA gradient against an fp32 CPU reference from the same bf16 inputs (cosine, relative L2 error; random cotangent). Gemma 4 attends at scale 1.0; head-0 scores reach 20 to 27.

| Layer | Path | Cosine dQ / dK / dV | Rel. error dQ / dK / dV |
|---|---|---|---|
| 0 to 4, sliding, D=256 | window fused (this patch) | 0.998 / 0.9995 / 0.9995 | 5.1–6.7% / 2.9–3.4% / 3.2–3.5% |
| 0 to 4, sliding | dense fallback | 0.9996 / 0.9997 / 0.9999 | 2.1–2.9% / 1.9–2.4% / 1.1–1.6% |
| 5, full, D=512 | shipped causal tiled backward | 0.9999 / 0.9998 / 0.9998 | 1.7% / 1.9% / 2.2% |
| 5, full | dense fallback | 0.9999 / 0.9999 / 1.0 | 1.3% / 1.4% / 0.9% |

So the windowed backward is pointed the right way on real activations (cosine 0.998 or better on every gradient), with about 2x the dense fallback's error (dQ up to 6.7% against 2.9%). The shipped causal path on the full layer has about 1.4x. My unverified guess at the source: the tiled backward keeps the recomputed score tile S in the input dtype (bf16) and forms P = exp(S - lse) from it. At scores near 25, a bf16 rounding step is 0.125, which is a 13% error in P for that element. Keeping S and dP in fp32 tiles would test it; that's a change to the existing upstream design, not to the window logic, and I haven't tried it.

**What this does and doesn't settle.** The fix's gradients are the kernel's own gradients, slightly noisier per layer than the dense path. Whether that matters for training can't be read from a B = 0 model-level cosine, where every bf16 path is noise. A 30,553-token CPU truth is possible on the Studio but not cheap: k = 53 took 606 s, the full layers' plain attention would need blocking to fit next to the 124 GB of fp32 weights, and I'd estimate hours. I don't expect it to separate the paths any better than k = 53 did.

## fp32 score tiles (Keith asked, about 8:25 to 8:45 PM)

**Change** (`mlx-windowed-backward-tools/fp32-tiles-experiment.patch`, on top of the window patch, behind `MLX_SDPA_VJP_FP32_TILES=1`, default off, not clang-formatted). My first try asked the GEMM kernels for bf16 inputs with fp32 output; the NAX GEMM kernels have no output-type parameter, so that doesn't compile. The version that ran upcasts q, k, v and dO to fp32 once per backward call and computes the S and dP tiles from those copies. bf16 values are exact in fp32, so this is the same as keeping S and dP in fp32. dS and P stay bf16 for the next three matmuls, so the per-element kernel now takes fp32 S and dP and writes bf16 dS and P into their own buffers.

**Layer check, k = 53, same binary, flag off vs on** (`results/lgc_k53_f32tiles{0,1}.jsonl`):

| Layer | Path | Rel. error dQ / dK / dV |
|---|---|---|
| 0 to 4, sliding | window fused, bf16 tiles | 5.1–6.7% / 2.9–3.4% / 3.2–3.5% |
| 0 to 4, sliding | window fused, fp32 tiles | 0.7–1.5% / 0.4–0.7% / 0.28% |
| 0 to 4, sliding | dense fallback | 2.1–2.9% / 1.9–2.4% / 1.1–1.6% |
| 5, full | causal tiled, bf16 tiles | 1.7% / 1.9% / 2.2% |
| 5, full | causal tiled, fp32 tiles | 0.5% / 0.5% / 0.28% |

Cosines go to 0.9999 or 1.0. So the bf16 score tile was the source of the extra error, and with fp32 tiles the tiled backward beats the dense fallback by 2x to 5x. This applies to upstream's shipped causal path too, not only the window.

**Cost.** Probe memory, peak GB above inputs: windowed 4.18 at 16K and 7.93 at 32K, against 2.49 and 4.74 with bf16 tiles; causal D=512 5.97 and 11.73, against 3.54 and 7.04. Still linear, about 1.6x higher, almost all of it the whole-tensor fp32 copies; doing the upcast per tile would cut most of that, and I haven't written it. 30,553-token canary with the flag on: 96.5 s against 95.4, footprint 187.4 GB against 182.2, same loss (1.5508652, the forward doesn't change), grad norm 92.81 against 91.92.

**Model level, k = 53, flag on** (`results/gvt_window_k53_f32.log`): window patch cosine to truth -0.037, grad norm 2,705 (flag off: -0.037, 2,668). Stock -0.058, blocked 0.138 (both use the tiled backward on their full layers, so they moved slightly too). Making attention gradients 4x more accurate per layer moved the model-level cosine by nothing measurable. So the near-zero model-level cosine, and the window path's large norm, come from elsewhere in the bf16 model, not from the attention backward. My best guess is the training forward: under the window patch, training runs the fused windowed forward (loss 1.11306 at k = 53) where stock's training path runs the dense fallback (1.13432); truth is 1.13935. I haven't tested that.

## Two tests Keith approved (8:35 PM)

**1. Is the training forward the source of the model-level gap? No.** `layer_grad_check.py` now also compares forward outputs against the fp32 CPU reference (`results/lgc_k53_fwd.jsonl`). Sliding layers 0 to 4: fused windowed forward and dense forward both at 0.16 to 0.17% relative error, cosine 1.0. Full layer 5: fused 0.16%, dense 0.47%. The fused forward is no less accurate than the dense path per layer, so the model-level loss differences at k = 53 (window 1.113, stock 1.134, blocked 1.105, truth 1.139) look like bf16 error compounding through 60 layers and the rest of the model, not an attention defect. That refutes my guess from the fp32-tiles section.

**2. Does the fix train as well as the workaround? Yes, on this measure.** (Results below.) `run_arms.sh`, started 8:38 PM: Lab 027's C2 recipe (86 steps, seed-42 init, shared order, checkpoints at 43 and 86) twice on the patched build, `c2-win` (window patch) then `c2-win-f32` (window patch plus fp32 score tiles), then Lab 027's `e1_val.py` on each final adapter in Lab 027's runtime (`~/venvs/mlx`, 0.32.4, blocked attention), so the val loss compares directly with C2's 0.1510. Baseline is C2's own log; the blocked canary reproduced it to the last digit on this build. About 3 hours per arm.

**Test 2 results (finished 2026-10-22, 1:16 AM training, validation by about 1:35 AM).** Both arms ran all 86 steps clean: no stop rule, no swap guard, rc=0 (`results/train_c2-win*.log`, `results/e1_val_c2-win*.json`). The `C2_studio` row in each val file is the arm under test.

| Run | Train time | Final-step loss | Mean / max abs. loss gap to C2, 86 steps | Val loss | Next-token match | Same prediction as C1 |
|---|---|---|---|---|---|---|
| C2 (Lab 027, blocked, 0.32.4) | 9,670 s | 0.1746 | | 0.1510 | 94.81% | 98.31% |
| window patch | 8,228 s | 0.1758 | 0.010 / 0.083 | 0.1496 | 94.77% | 98.46% |
| window patch + fp32 tiles | 8,430 s | 0.1824 | 0.010 / 0.087 | 0.1530 | 94.60% | 98.34% |
| C1, Spark (Lab 027) | | | | 0.1536 | 94.75% | 100% |
| Lab 018 original (Spark) | | | | 0.1534 | 94.75% | 98.36% |

Correction (2026-10-22): the window adapter's val 0.1496 sits slightly below the three existing adapters' range (0.1510 to 0.1536), not inside it; the fp32 arm's 0.1530 is inside. Match 94.60% to 94.81% and agreement with C1 98.3% to 98.5% across all five. The window patch trains the Lab 027 job on stock mlx-lm about 15% faster than the blocked workaround, to the same held-out quality. fp32 score tiles, despite 4x more accurate attention gradients per layer, didn't produce a better adapter: val 0.1530 against 0.1496, which is within the run-to-run spread above. Peak footprint for the whole window run was 230.9 GB, the same as C2's 230.8 GB, so the 23 GB saved on the 30,553-token canary isn't the run's binding step. GPU power median 143 W against C2's 136 W (same work, less time).

## Keith's ruling, 2026-10-22: a lab and a fork

Keith: "This solves an open problem in the community so it deserves a lab and a fork on github. I'm not getting involved in becoming a maintainer of MLX. If someone wants to pick this up, great, if not, it's a great experiment and an example of the overhead of not going with NVIDIA's ecosystem." On MLX's AGENTS.md: it governs contributions to ml-explore/mlx, not his fork; it applies again only if he opens a PR or comments upstream.

- **Lab 028 scaffolded** as `mlx-windowed-backward` (proposed). Title "Off NVIDIA, you fix the platform yourself": Keith accepted it, 2026-10-22. Spec, bench and entry are template stubs. These notes are the bench record.
- **Pushed (2026-10-22, on Keith's go, email keith@thectoadvisor.com).** Public fork https://github.com/kltownsend/mlx, branch `sdpa-windowed-backward`: `156eb0325` window backward, `2d63f1b67` fp32 tiles experiment, `0be6d8587` README. Commits re-authored to keith@thectoadvisor.com; code tree unchanged from the tested diff. README edits before push: the Lab 028 link became "will be written up as a lab on labs.layer2c.com" (not public yet), and the parity line now names the one bf16 case that missed by 0.0004. No PR to ml-explore/mlx, no upstream comment.
- **README disclosure updated and pushed (2026-10-22, Keith's go):** `dcfc3dbb8`, "I wrote the code, the tests and this README using Claude Code, Anthropic's coding agent." Matches the lab record and essay.
- **Fork branch ready, not pushed.** `~/src/mlx-win` on the Studio, branch `sdpa-windowed-backward` on #4734 head `fdf1d256a`, three commits: `4dc5dc1ab` window backward, `16fd3fe44` fp32 tiles experiment (off by default), `5cb448d86` README (`WINDOWED_SDPA_BACKWARD.md`, draft in `mlx-windowed-backward-tools/FORK-README.md`). The committed tree equals the diff that was formatted, rebuilt and tested (`test_fast_sdpa` OK with the fp32 flag off and on). Full patch: `mlx-windowed-backward-tools/sdpa-windowed-backward-full.patch`. Commits are authored "Keith Townsend <sales@thectoadvisor.com>" (the Studio's git config); Keith to confirm the public email before push.
- **Disclosure check for the lab** (wire step): Gemma 4 is Google's model and Google Cloud is a client; NVIDIA is named in the framing. Confirm both with Keith before publish.