Autograd resident backward and fp16 fix

CLAUDE.md

@@ -121,3 +121,115 @@ session_recall("<topic>") → consult("cto", "<question>") → session_log("deci
 # Training experiment
 session_log("experiment", ...) → run_smoke_test → session_log("result", ...)
 ```
+
+
+---
+
+# Resident Autograd (branch: `autograd-resident-backward`)
+
+> Added June 2026. The sections above describe the Python/Vulkan framework. This
+> section covers the **C++ resident autograd engine** ? a reverse-mode autograd
+> that runs the whole training step on-GPU (forward outputs and gradient buffers
+> never leave VRAM). **`AUTOGRAD_STATE.md` is the detailed source of truth**;
+> this is the orientation map.
+>
+> **Status (June 2026):** the resident forward op set is COMPLETE --
+> embedding / rmsnorm / linear / mingru / swiglu all run resident and are
+> parity-verified, and the full resident gradcheck passes. The NEXT step is the
+> single-tape full-trunk integration + linking cubby's trunk onto it; the precise
+> plan is in `AUTOGRAD_STATE.md` ("NEXT: single-tape full-trunk integration").
+
+## Where the work lives
+
+- **C++ engine**: `cpp/src/autograd.cpp` + `cpp/include/grilly/autograd/`
+  (`BackwardEngine`, `TapeContext`, `BufferRegistry`).
+- **Compiled Python binding surface**: `cpp/python/bindings_autograd.cpp`.
+  **`cpp/python/bindings.cpp` is DEAD CODE** ? it is NOT in the build. Editing it
+  does nothing. Confirm against `CMakeLists.txt` before touching any binding.
+- **Tests / milestones**: `experimental/resident_train/`
+  - `train_full.py` ? composed Cubby block, **numpy forward + resident backward**.
+    Stable baseline. Online training (fresh random batch/step), `gradcheck` mode
+    vs finite-diff. ~0.945 train / ~0.891 test acc (chance 0.25).
+  - `train_full_resident.py` ? same block, **fully resident** (forward also on-GPU
+    via `forward_rmsnorm`/`forward_linear`/`forward_swiglu`/`forward_mingru`). Same accuracy.
+  - `test_backward_*.py` (8 files in the REPO ROOT: linear, chain, ce, silu, fanout, rmsnorm,
+    mingru, swiglu) ? per-op gradient unit tests. Keep all green.
+  - `ttr_mingru_fwd.py` / `ttr_embedding_fwd.py` -- parity tests for resident
+    forward_mingru (1.5e-7 vs numpy) and forward_embedding (exact gather).
+  - `ttr3.py` ? minimal resident-forward repro (was the crash repro; passes now).
+  - `ttr8.py` ? proof the multi-tape t/t2/t3 backward runs clean.
+
+## Build & run (Windows, AMD RX 6750 XT / RADV)
+
+There is **no CUDA**. The engine is C++ compiled to `grilly_core.cp312-win_amd64.pyd`.
+
+```powershell
+# Rebuild after C++/header changes (skip shader recompile):
+Get-Process python | Stop-Process -Force      # a live python LOCKS the .pyd and blocks the copy
+powershell -NoProfile -ExecutionPolicy Bypass -File ".\rebuild.ps1" -SkipShaders
+# Omit -SkipShaders only when a .glsl changed. Build prints "Build OK" then copies the .pyd.
+
+# Run against the engine with the cubby-lm venv python (grilly editable-installed there):
+& "C:\Users\grill\Documents\GitHub\cubby-lm\.venv\Scripts\python.exe" experimental\resident_train\train_full_resident.py
+```
+
+Crash exit codes: `0` ok, `1` Python exception, `-1073741819` (0xC0000005) =
+access violation / heap corruption (C++ level). For crashes, redirect to a temp
+file (`> $env:TEMP\out.txt 2>&1`) and check `$LASTEXITCODE` ? piping through
+Select-String can drop the tail.
+
+## The Vulkan validation layer is the primary debugger
+
+Installed, no rebuild needed. Enable per-run:
+
+```powershell
+$env:VK_INSTANCE_LAYERS="VK_LAYER_KHRONOS_validation"
+```
+
+It reports invalid/destroyed `VkBuffer` handles, descriptor mismatches, and
+push-constant range violations with the offending handle. **Use it before
+guessing.** A `VkBuffer` handle that decodes to ASCII (e.g. a value ending
+`...2eadc26e` = `"part."`) means a C++ struct field was overwritten by string
+data ? a **dangling reference / use-after-free**, not a GPU bug.
+
+## Invariants ? do not regress these
+
+- **`BufferRegistry` stores entries in a `std::deque`, NOT `std::vector`**
+  (`cpp/include/grilly/autograd/buffer_registry.h`). `resolve()` hands out
+  `GrillyBuffer&` references held across later `alloc()` calls. A vector would
+  reallocate on `push_back` and dangle those references ? 0xC0000005 during
+  backward once enough buffers are registered. This was the root-cause crash;
+  deque keeps element references stable. (Commit f74e58e.)
+- **Barrier discipline in `BackwardEngine::backward`**: a node reads
+  `grad_output_buffer` that may have been built by accumulation writes from
+  downstream nodes. Handlers that `fillZero` (a TRANSFER write) need
+  `transferComputeBarrier` (not plain `barrier`) so the zero is ordered against a
+  downstream shader read. Two real bugs lived here (MinGRU?Linear race;
+  three-branch RMSNorm accumulation race). Don't weaken these barriers.
+- **`getOrCreate` caches pipelines by NAME only** (`cpp/src/pipeline_cache.cpp`):
+  the first creation's `numBuffers`/`pushConstSize` wins; later calls with
+  different sizes are ignored. Keep each shader's push-constant struct size
+  consistent across every call site.
+
+## Known open issue
+
+- A **non-fatal push-constant validation warning** (`VUID-vkCmdPushConstants-offset-01795`):
+  a 20-byte push (rms-norm `RMSNormParams`) is flagged against a layout reported
+  as having no COMPUTE range, even though `rms-norm` is created with
+  `pushConstSize=20`. Forward is numerically correct (gradchecks pass), so it's
+  latent, not breaking ? RADV tolerates it. Under investigation in
+  `CommandBatch::dispatch` (`cpp/src/command_batch.cpp` ~L94): the order of
+  `vkCmdBindPipeline` / `vkCmdBindDescriptorSets` / `vkCmdPushConstants` and
+  which layout the push uses.
+- **Investigation note**: the temporary `[PIPE-CREATE]` getOrCreate instrumentation
+  has been REMOVED (build is clean). It confirmed `rms-norm` IS created with
+  `pushConstSize=20` and is the only 20-byte push in ttr3, so the layout is correct
+  -- the warning comes from the dispatch passing the wrong/stale layout for the push
+  (look at `CommandBatch::dispatch`). Still open, still non-fatal.
+
+## Relationship to cubby-lm
+
+This branch builds the **resident GPU training path** that the sibling `cubby-lm`
+repo needs for its `0.0.1 perf #2` (resident activations ? the ~25 ms/dispatch
+floor). `cubby-lm` imports grilly as an editable path source and loads this
+`.pyd` automatically. See `../cubby-lm/CLAUDE.md`.

CMakeLists.txt

@@ -262,13 +262,30 @@ target_include_directories(grilly_core_lib PUBLIC
     "${VMA_INCLUDE_DIR}"
 )
 
-# Find shaderc for JIT shader fusion (ships with Vulkan SDK)
-find_library(SHADERC_LIB shaderc_combined
-    HINTS "$ENV{VULKAN_SDK}/Lib" "$ENV{VULKAN_SDK}/lib")
-if(SHADERC_LIB)
-    message(STATUS "Found shaderc: ${SHADERC_LIB}")
+# Find shaderc for JIT shader fusion (ships with Vulkan SDK only).
+#
+# IMPORTANT: We intentionally restrict the search to $VULKAN_SDK and refuse
+# to fall through to system paths. On Linux, distro-packaged
+# /usr/lib64/libshaderc_combined.a is typically built by Google's upstream
+# against libc++ (Clang). Linking it into a libstdc++ (GCC) build produces a
+# .so that references std::__1::* symbols with no matching definition —
+# manifesting at import time as:
+#   ImportError: undefined symbol: _ZTT...basic_istringstream...
+# The LunarG SDK ships a GCC/libstdc++-compatible build, so we require it
+# explicitly. Shader JIT fusion is optional; if VULKAN_SDK is unset we just
+# disable GRILLY_HAS_SHADERC rather than risk an ABI mismatch.
+if(DEFINED ENV{VULKAN_SDK})
+    find_library(SHADERC_LIB shaderc_combined
+        HINTS "$ENV{VULKAN_SDK}/Lib" "$ENV{VULKAN_SDK}/lib"
+        NO_DEFAULT_PATH)
+    if(SHADERC_LIB)
+        message(STATUS "Found shaderc: ${SHADERC_LIB}")
+    else()
+        message(WARNING "shaderc not found in VULKAN_SDK — JIT shader fusion disabled")
+    endif()
 else()
-    message(WARNING "shaderc not found — JIT shader fusion disabled")
+    message(STATUS "VULKAN_SDK not set — JIT shader fusion disabled "
+                   "(install LunarG Vulkan SDK and export VULKAN_SDK to enable)")
 endif()
 
 target_link_libraries(grilly_core_lib PUBLIC
@@ -318,6 +335,8 @@ pybind11_add_module(grilly_core
     cpp/python/bindings_mingru.cpp
     cpp/python/bindings_bandit.cpp
     cpp/python/bindings_eggroll.cpp
+    cpp/python/bindings_autograd.cpp
+    cpp/python/bindings_distillation.cpp
 )
 target_link_libraries(grilly_core PRIVATE grilly_core_lib)
 

FP16_BACKWARD_SPEC.md

@@ -0,0 +1,189 @@
+# Spec — fp16 cooperative-matrix backward for resident Linear
+
+**Goal:** unblock coopmat in the resident *training* path. Forward coopmat
+already exists (`ops::linear` + `gemm-coopmat-shared`); the blocker is backward.
+`ops::linearBackward` hard-rejects `elemSize != 4` ("fp16 backward needs a
+cooperative matrix backward shader — TODO"), and `BackwardEngine::backward_linear`
+(autograd.cpp) is hardcoded to fp32 `gemm_tiled`. This spec defines exactly what
+to build.
+
+## Key finding (REVISED 2026-06-16 after fixing + validating the kernel)
+
+The forward coopmat kernel `gemm-coopmat-shared` is now **fixed and numerically
+validated** (was returning the per-tile transpose; see GEMM_AUTOTUNE.md
+correction). Crucially, the fix means the kernel now computes **`C = A · Bᵀ`**
+where B is the weight matrix stored `(N, K)` — i.e. it matches the linear op's
+`y = x·Wᵀ` convention, NOT a plain `A·B`. Validated parity vs fp32: ~2.9e-04
+across all trunk shapes.
+
+This changes the backward reuse story. The two backward GEMMs do NOT share a
+convention:
+
+```
+grad_input  = g @ W      # g (BS,out), W stored (out,in).
+                         #   = g · W with W as (K=out, N=in) but stored (out,in)
+                         #   = EXACTLY the kernel's A·Bᵀ form (B=W is the (N,K)=(in?,..))
+                         #   -> reuses the FIXED kernel directly. ✓
+grad_weight = gᵀ @ x     # gᵀ (out,BS), x (BS,in), plain row-major A·B.
+                         #   x is a DATA matrix, not weights -> needs plain A·B,
+                         #   which the fixed kernel NO LONGER does. ✗
+```
+
+So:
+- **grad_input** can reuse the fixed `gemm-coopmat-shared` as-is (its `x·Wᵀ`
+  convention is exactly `g·W` with W stored row-major (out,in)). Confirm with a
+  parity probe before trusting this — the index algebra is plausible but unproven
+  for the backward operand shapes.
+- **grad_weight** needs a *plain* `A·B` coopmat (no Bᵀ). Options:
+  (a) a second kernel variant `gemm-coopmat-shared-nt` that stages B without the
+      transpose (the ORIGINAL staging, before the fix), or
+  (b) materialize `xᵀ` and feed `gᵀ · (xᵀ)ᵀ` through the Bᵀ kernel — wasteful,
+      two transposes, not worth it, or
+  (c) parameterize the kernel with a push-constant `transpose_b` flag selecting
+      which staging index to use. **(c) is cleanest** — one kernel, one bit.
+
+Recommended: **(c)** — add a `uint transpose_b` push constant to
+`gemm-coopmat-shared`; `transpose_b=1` (weights, `x·Wᵀ`) for forward + grad_input,
+`transpose_b=0` (plain `A·B`) for grad_weight. The two staging index expressions
+already exist (pre-fix = plain, post-fix = transposed); the flag just selects.
+
+Helper shaders still needed (written + compiled):
+1. `cast-f32-f16.glsl` — elementwise fp32 → fp16 copy. ✓ compiled.
+2. `transpose-f32-f16.glsl` — fused transpose+cast for building `gᵀ` (fp16). ✓
+   compiled. (Still needed for grad_weight's gᵀ regardless of the flag approach.)
+
+## The three backward quantities (unchanged math)
+
+Linear: `y = x @ W^T (+b)`, W is `(outputDim, inputDim)`, x is
+`(batchSeq, inputDim)`, grad_output `g` is `(batchSeq, outputDim)`.
+
+```
+grad_input  = g   @ W            # (BS,out)·(out,in)        -> (BS,in)
+grad_weight = g^T @ x            # (out,BS)·(BS,in)         -> (out,in)
+grad_bias   = sum(g, axis=0)     # (out,)
+```
+
+Note W is stored `(out,in)` = exactly the `(K,N)` operand `grad_input` needs
+(K=out, N=in) with NO transpose. `grad_weight` needs `g^T` (transpose of g).
+
+## Dispatch plan (fp16 path), per pass
+
+All coopmat GEMMs require M%16, K%16, N%64 on the **operand as the kernel sees
+it** (M=rows of A/C, K=contraction, N=cols of B/C). Cast inputs to fp16 first.
+
+**Pass 0 — grad_input = g @ W**
+- A = g (BS×out) fp16, B = W (out×in) fp16, C = grad_input (BS×in) fp32.
+- M=BS, K=out, N=in. Coopmat-legal iff BS%16, out%16, in%64.
+- Prep: `cast-f32-f16` on g and W (W cast once per step, cacheable).
+- Dispatch `gemm-coopmat-shared` gx=in/64, gy=BS/16.
+
+**Pass 1 — grad_weight = g^T @ x**
+- Need g^T (out×BS). Use `transpose-f32-f16` on g → gT_f16 (out×BS), fp16.
+- Cast x → x_f16 (BS×in).
+- A = gT (out×BS), B = x (BS×in), C = grad_weight (out×in) fp32.
+- M=out, K=BS, N=in. Coopmat-legal iff out%16, BS%16, in%64.
+- Dispatch gx=in/64, gy=out/16.
+
+**Pass 2 — grad_bias** = unchanged. Reuse the existing fp32 reduction
+(fnn-linear-backward pass 2, or a standalone column-sum). Bias stays fp32
+everywhere (matches forward). No coopmat.
+
+## Alignment fallback (REQUIRED — correctness, not just perf)
+
+`gemm-coopmat-shared` uses floor division grids (`N/64`, `M/16`) and the loop
+`for k in 0..K step 16`. If a dimension isn't aligned it **silently drops the
+remainder** — wrong results, not a crash. So the dispatch wrapper MUST check
+M%16==0 && K%16==0 && N%64==0 for the pass's operand layout and fall back to the
+fp32 `gemm_tiled` (or the fnn-linear-backward scalar pass) when unaligned.
+
+Cubby trunk dims are friendly: BS=512, in/out ∈ {1024, 4096, 8192} are all
+%64. The danger dims are the vocab head (out=65536 in grad_weight: 65536%16 ok,
+but it's over the autotuner cap — keep it fp32 tiled) and any non-multiple batch
+(B*S). Gate per-shape; do NOT assume.
+
+## C++ wiring (where each piece lands)
+
+`ops::linearBackward` (linear.cpp): add an fp16 branch parallel to the existing
+fp32 one, selected when `p.elemSize == 2 && device.hasCooperativeMatrix() &&
+aligned`. Allocate fp16 staging for g/x/gT + W; emit cast/transpose dispatches
+into the same batched command buffer before the coopmat dispatch; keep the fp32
+path as the `else`.
+
+`BackwardEngine::backward_linear` (autograd.cpp): this is the resident hot path
+and currently emits `gemm_tiled` directly (not via `ops::linearBackward`). Two
+options:
+- **(A, smaller)** leave it fp32 — only flip `ops::linearBackward` for non-
+  resident callers. Lowest risk, but does NOT speed up resident training (the
+  thing we care about). 
+- **(B, the real win)** mirror the cast→coopmat chain inline in
+  `backward_linear`'s pass 0 / pass 1 blocks, guarded by the same alignment +
+  autotuner-table check. This is where the 1.3× actually lands. Higher risk
+  (hot path, fp16 grads). Recommended only after (A) proves parity.
+
+Autotuner table feeds the per-shape guard: only take the coopmat branch where
+`gemm_autotune` says `coopmat` AND aligned AND under the dim cap.
+
+## Numerical-parity gate (MUST pass before enabling)
+
+fp16 inputs lose mantissa; grad_weight accumulates over K=BS=512 terms, the
+riskiest for fp16 round-off (accumulation is fp32 in the coopmat accumulator,
+which helps). Gate each shape:
+- Compare fp16-coopmat grad_input / grad_weight against the fp32 `gemm_tiled`
+  result on the same inputs.
+- Threshold: relative L2 error < 1e-2 (fp16 GEMM with fp32 accumulate typically
+  lands ~1e-3). If a shape exceeds it, keep that shape fp32.
+- Add to test suite as a per-shape parametrized parity test, same spirit as
+  test_checkpoint's logits-parity check.
+
+## Build order (smallest safe increments)
+
+1. ✓ DONE — `cast-f32-f16.glsl` + `transpose-f32-f16.glsl` written + compiled
+   (SPVs in shaders/spv). ✓ DONE — fixed + validated `gemm-coopmat-shared`
+   forward (parity ~2.9e-04).
+2. ✓ DONE — added the `transpose_b` push-constant flag to `gemm-coopmat-shared`
+   (transpose_b=1 -> A·Bᵀ weights/forward; transpose_b=0 -> plain A·B).
+   Push grew 12→16 bytes; updated both C++ callers (`ops::linear` sets =1,
+   `bench_gemm` sets =0 with a separate 12-byte pc3 for gemm_tiled) and rebuilt
+   (Build OK). Forward RE-VALIDATED post-rebuild: ~2.9e-04 across all trunk
+   shapes — the flag did not regress the forward path.
+   NOTE: the transpose_b=0 *code path* itself is compiled but not yet exercised
+   from Python (no binding sets the flag). Its plain-A·B math is validated
+   indirectly (forward with pre-transposed W gives 2.9e-04), but its first true
+   test is step 3, where grad_weight invokes it. Gate it there.
+3. ✓ DONE — added `linearBackwardCoopmat` (fp32 in/out, fp16 coopmat GEMMs
+   internally) + binding `linear_backward_coopmat`. Both backward GEMMs use
+   transpose_b=0 (plain A·B): grad_input = g·W (B=W read plain as (out,in)),
+   grad_weight = gᵀ·x (gᵀ built fp16 via transpose-f32-f16). grad_bias fp32 via
+   fnn-linear-backward pass 2. Parity vs numpy ref, 6 shapes (incl. MinGRU-proj
+   512x1024x3072, SwiGLU-down 512x4096x1024): grad_input/grad_weight ~7.7e-04,
+   grad_bias EXACT — ALL PASS <1e-2. No autograd/hot-path change. NOTE: this
+   means grad_input also uses transpose_b=0, NOT transpose_b=1 as the earlier
+   key-finding guessed — because W is stored (out,in) which IS already the plain
+   (K=out,N=in) operand, so no Báµ€ needed. The transpose_b=1 path remains the
+   forward-only case.
+   BONUS FIX: found + fixed a latent bug in the EXISTING fp32 `ops::linearBackward`
+   — grad_bias pass 2 dispatched ceil(out/256) workgroups but the shader is a
+   16-wide workgroup, so only the first 64 of 1024 bias entries were ever
+   computed (15/16 left zero). Now ceil(out/16). This bug only affected the
+   `ops::linearBackward` binding, not the resident autograd path (which has its
+   own bias handling).
+4. Wire `backward_linear` (autograd.cpp) option (B) behind the autotuner table,
+   one shape at a time, parity-gated. Measure END-TO-END step time, not µbench —
+   the cast+transpose dispatches eat into the kernel win.
+5. Only widen to more shapes after net step-time improvement is confirmed.
+
+## Honest risks
+
+- **Hot path + fp16 grads.** Backward feeds AdamW; bad fp16 grads degrade
+  training silently (loss drifts, not a crash). The parity gate is the guard;
+  do not skip it.
+- **Cast/transpose overhead.** Two extra dispatches + fp16 staging per linear.
+  On small GEMMs this can erase the coopmat win — that's WHY it's autotuner-
+  gated and end-to-end-measured, not applied blanket.
+- **GPU stability.** Same coopmat kernel family that TDR'd the bench on a
+  vocab-sized M. Keep the dim cap; never dispatch coopmat on the 65536 head.
+- **Coverage floor.** The floor-division grid means alignment fallback is
+  mandatory and must be tested with a deliberately unaligned shape.
+- **Backward shader is reused, not new** — lower risk than a hand-written
+  backward kernel, but it inherits `gemm-coopmat-shared`'s exact alignment
+  contract, so the wrapper owns correctness.