docs(matrix): seq=4096 burst dissected — mamba3-chunk + grad-checkpoint-correctness fixes, measured floor

Records the 2026-06-01 seq=4096 memory investigation: two dominant terms found
and fixed (mamba3 Path-B bwd torch scratch 15->2 GB via CPPMEGA_MAMBA3_BWD_SEQ_CHUNK;
grad_checkpoint correctness — mx.checkpoint silently dropped the final MoE
layer's gradients, fixed via nn.utils.checkpoint). Measured unified free-g peaks
(forward 16 GB, fwd+bwd 73 GB, full step 87 GB, real m04 step >114 GB). Verdict:
path_b bs=1 seq=4096 does NOT fit under 70 GB — hard floor ~73 GB fwd+bwd, the
dominant residual is MLX-CUDA backward activation retention that checkpointing
does not reduce on this backend. No fake throughput. Corrects the stale
'path_b blocked on CUDA' note (was the grad-checkpoint SDPA-mask bug).

Author: apstenku123

reports/MATRIX-REVIEW-INDEX.md

@@ -36,7 +36,7 @@ gb10 1B, flag-OFF (serial `path_c`) → flag-ON (`path_c_chunked`, `CPPMEGA_PATH
 
 ## Headline 3 — Platform reachability (fail-loud, honest)
 
-- **gb10 1B:** `path_b` is **blocked** (MLX-CUDA cannot run the path_b reference); `path_c` / `path_c_chunked` run and pass the loss check (11.3 → 6.27, finite + decreasing). `nvfp4` cells all **fail-loud** (no nvfp4 training kernels yet; backward GEMM miscompiled on sm_121 — see `NVFP4-TRAINING-KERNELS.md`).
+- **gb10 1B:** `path_b` runs on CUDA at seq≤1024 and **scales** with batch + seq (156→259 batch, 156→191 seq; the earlier "path_b blocked" note was a `--grad-checkpoint` bug — `mx.checkpoint` made the SDPA mask a differentiable input — now fixed via `nn.utils.checkpoint`, see the 2026-06-01 follow-up). At seq=4096 the step does **not fit** (≥114 GB; out of scope for tok/s). `path_c` / `path_c_chunked` run and pass the loss check (11.3 → 6.27, finite + decreasing). `nvfp4` cells all **fail-loud** (no nvfp4 training kernels yet; backward GEMM miscompiled on sm_121 — see `NVFP4-TRAINING-KERNELS.md`).
 - **gb10 v4 (a–e):** only `path_a` (pure-MLX reference) and `GDN path_e` run on CUDA. `path_b` ("No Metal back-end"), `path_c` (`DLPackDeviceError` — TileLang compiled for `target=metal`, can't take CUDA arrays), `path_d` (Triton disabled), `KDA path_e` (Metal-only kernel) all **fail-loud**. → the v4 op-level Path-C/B/E dispatcher is still **Metal-target-only**; CUDA-target wiring for the v4 GDN/KDA ops is the open item (distinct from the m04-training Path-C, which already works on CUDA per Headline 1).
 - **local Metal v4 (a–e):** `path_a/b/c/e` all run; `path_d` fails loud (Triton frontend disabled). Fastest = `path_e` (vendored mlx-lm gated_delta Metal kernel, ~1 ms, 225–287 Melem/s).
 - **local Metal 1B:** at batch=1 `path_b` leads `path_c_warm` (e.g. adamw 457 vs 259) — same per-call-overhead-dominates-tiny-workload effect; not representative of the fused advantage at scale.
@@ -67,6 +67,28 @@ The like-for-like **batch=4×seq=4096** comparison the goal asked for is **not a
 
 **What it would take to reach seq=4096 fwd+bwd under 70 GB** (the remaining wall, evidence above): (a) a fused/streamed optimizer step that updates params in chunks instead of materializing all grads + AdamW m/v at once; (b) more aggressive activation checkpointing across the attention/mamba backward (the bulk of the burst); (c) the efficient MoE (now landed, env-gated) is a necessary-but-not-sufficient piece. Item (a) is the highest-leverage next step.
 
+### 2026-06-01 follow-up — seq=4096 burst dissected + two fixes landed (still does NOT fit under 70 GB)
+
+A dedicated profiling pass (`scripts/profile_bwd_mem_20260601.py`, `scripts/probe_real_step_mem_20260601.py`) located **two** distinct dominant terms in the seq=4096 backward burst and fixed both, but the full real m04 step still exceeds the 121 GB unified box. **Use `free -g` as the unified-memory truth — `mx.get_peak_memory()` on the MLX-CUDA backend under-reports by ~25–55 GB (it tracks only a subset of allocations).**
+
+**Term 1 — Mamba3 Path-B backward scratch (FIXED, env-gated `CPPMEGA_MAMBA3_BWD_SEQ_CHUNK`, default OFF).** The Path-B mamba3 MIMO bwd (Metal + the gb10 CUDA-eager kernel) allocates fp32 partials of shape `(B,SEQ,H,N,P)` plus a `(B,SEQ+1,H,P,N)` state slab → **~21 GB per mamba3 layer at seq=4096, ×3 ≈ 63 GB** (in the torch CUDA allocator, invisible to `mx.get_peak_memory`). Per-layer `mx.checkpoint` is a no-op for this — the buffers are *inside* the mamba custom_function VJP, not in the autograd-retained set. **Fix:** sequence-chunked backward carrying the scan state `h` and end-state cotangent `dh` across chunk boundaries — numerically identical (direct-kernel parity `max_rel=3–4e-7` CUDA / `3.8e-7` Metal; full-model loss `dloss=0`, grads `max_rel=1.4e-6`). **Measured gb10, real CCE loss, fwd+bwd:** torch-CUDA peak **15.1 → 1.9 GB (8×)**; unified `free -g` **105 → 87 GB** (full step w/ AdamW). Commit `7f32db4`.
+
+**Term 2 — grad_checkpoint was silently BROKEN (FIXED — correctness, RULE #1).** The decoder used `mx.checkpoint(layer)`, which checkpoints w.r.t. the call's *array inputs only* and does **not** thread the module's trainable parameters through the recompute. Under `nn.value_and_grad` this **silently dropped gradients** for many params — verified the **entire final MoE layer's expert + shared_expert weights**, plus `conv_bias` and `rope_inv_freq`, had `|grad|==0` with checkpoint vs `|grad|>0` without (loss identical → invisible to any loss-only check). It also made the additive attention mask a differentiable checkpoint input → `scaled_dot_product_attention does not support VJP w.r.t. mask` on CUDA (the *real* reason `--grad-checkpoint` path_b was "blocked" on gb10, NOT "MLX-CUDA cannot run path_b"). **Fix:** `nn.utils.checkpoint(layer, fn)` (threads params via `module.trainable_parameters()`) with mask/doc_ids closed over. Now bitwise-correct: grad_checkpoint ON == OFF, `dloss=0`, grad `max_rel=0.0`. Regression test `test_grad_checkpoint_gradients_match_non_checkpoint`. Commits `966611d`, `32aa7a0`. **Every prior `--grad-checkpoint` run trained the final MoE layer with zero gradients.**
+
+**Measured seq=4096 bs=1 unified `free -g` peaks (gb10, bf16, `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`, `CPPMEGA_MOE_EFFICIENT=1`):**
+
+| stage / config | unified peak | notes |
+| --- | ---: | --- |
+| forward only | **16 GB** | mlx_peak 4.9 GB |
+| fwd+bwd, chunk OFF | **105 GB** | mamba torch scratch 15 GB |
+| fwd+bwd, chunk ON | **73 GB** | torch scratch 2 GB — **backward adds +57 GB over forward** |
+| fwd+bwd, chunk ON + ckpt ON | **73 GB** | checkpoint does NOT lower the unified peak on MLX-CUDA |
+| fwd+bwd+AdamW, chunk ON | **87 GB** | optimizer adds +14 GB |
+| fwd+bwd+AdamW, chunk+split-eval+adam8bit | **103–109 GB** | quantize + extra-eval host roundtrips make it *worse* |
+| **real m04 step (matrix cell), chunk ON, --grad-checkpoint** | **>114 GB** | SIGTERM at 114 GB; the full m04 harness adds ~27 GB over the minimal probe |
+
+**Verdict: path_b at bs=1×seq=4096 does NOT fit under 70 GB.** Hard floor ≈ **73 GB for fwd+bwd alone** (before the optimizer); the real m04 step exceeds **114 GB**. The dominant remaining term is **backward activation retention (+57 GB)**: the MLX-CUDA eager autograd holds the full forward graph of all 13 full-width layers, and **activation checkpointing does not reduce the unified peak on this backend** (recompute adds as much working set as it frees). The one term checkpointing couldn't touch — the mamba3 torch scratch — is now chunked 15→2 GB, but it was never the unified bottleneck. bs=2 / bs=4 × seq=4096 are further out of reach. **No fake throughput: the step cannot complete under the safe budget, so there is no honest path_b tok/s at seq=4096.** Real tok/s stays bounded to seq≤1024 (166 efficient / 191 dense, bs=1). Box left clean after every run (`fuser`/`free` verified, SIGTERM-only, never `kill -9`).
+
 ## Resolved since first draft
 - **Memory-efficient MoE** ✅ landed, **env-gated** `CPPMEGA_MOE_EFFICIENT=1` (default OFF → existing paths byte-identical). `ReferenceMoE._routed_combine_sparse` gathers only the tokens routed to each expert (top_k), runs the FFN on that subset, scatter-adds back — **same math** as the dense loop. Numeric parity pinned in `tests/test_moe_efficient.py`: **forward bitwise-identical (max-diff 0.0 on BOTH Metal and CUDA)**; gradients within the fp32 reduction-order envelope (7.6e-5 Metal / 3.7e-3 CUDA — pure cuBLAS non-associativity, *not* an algorithm delta — proven by a bitwise-exact gather/scatter round-trip test). Single-layer fwd+bwd: seq=4096 **3.84→2.20 GB (1.75×), 2.7× faster**. Full-model gb10 seq=1024 bs=1: **166 tok/s, 32.8 GB, loss 11.29→5.65** (matches dense 191/32.6). *Conclusion: necessary but not sufficient for seq=4096 — the wall is the whole-model backward+optimizer burst, see Fair-comparison section.* (commits `72d8827`, `6c49768`, `60477f5`, `c19d4a8`).
 - **nvfp4 forward** ✅ works on gb10 (rel_err 0.147). **nvfp4 backward** ✅ now REAL on sm_121 via the reduced RtN recipe `NVFP4BlockScaling(disable_rht=True, disable_stochastic_rounding=True)` — dgrad 0.1465 / wgrad 0.1343 vs bf16, enabled behind a numeric gate; the DEFAULT RHT+SR recipe stays fail-loud (genuinely datacenter-only). See `NVFP4-TRAINING-KERNELS.md`, `NVFP4-NANOCHAT-RECIPE.md`, `NVFP4-GB10-MIXED-APPROACH.md` (commit `78b7b31`).