--sdpa: fused flash attention in the no_grad relax loop — 1.45x free phase, z* parity 4e-7

Scoped via blk._sdpa set only inside relax()'s loop (grad paths jvp/vjp/resreg
keep the manual attention: no forward-mode-through-flash risk). Combined with
--holofast: ~1.51x full-step exact-math tier.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn

1 files changed, 32 insertions, 0 deletions

diff --git a/ep_run/sdpa_gate.py b/ep_run/sdpa_gate.py
new file mode 100644
index 0000000..0ea4f60
--- /dev/null
+++ b/ep_run/sdpa_gate.py
@@ -0,0 +1,32 @@
+"""Ship-gate for --sdpa (fused flash attention in the no_grad relax loop): z* parity + res + val + timing.
+Grad paths untouched by construction (the _sdpa flag is scoped to relax's loop), so no BPTT gate needed."""
+import time, torch
+import lt_ep_train as L
+
+torch.manual_seed(0)
+blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick'); blk.qknorm = True
+ck = torch.load('runs/redx_traj/s2000.pt', map_location=L.dev)
+with torch.no_grad():
+    for p, w in zip(blk.allp, ck['allp']):
+        p.copy_(w.to(L.dev))
+torch.manual_seed(42)
+idx, _ = L.get_batch('train', 24, 256)
+xin = blk.embed(idx).detach()
+
+out = {}
+for name in ('manual', 'sdpa'):
+    blk.sdpa = (name == 'sdpa')
+    z = L.relax(blk, xin.clone(), xin, 150, 0.1)          # warmup + result
+    res = (L.relax(blk, z, xin, 1, 0.1) - z).norm().item()
+    val = L.evaluate(blk, 150, 0.1, nb=4)
+    ts = []
+    for _ in range(3):
+        torch.cuda.synchronize(); t = time.time()
+        L.relax(blk, xin.clone(), xin, 150, 0.1)
+        torch.cuda.synchronize(); ts.append(time.time() - t)
+    out[name] = (z, res, val, min(ts))
+    print(f"{name:>6}: res={res:.3e} val={val:.4f} relax150={min(ts):.3f}s", flush=True)
+
+zd = ((out['sdpa'][0] - out['manual'][0]).norm() / (out['manual'][0].norm() + 1e-12)).item()
+print(f"z* rel-diff={zd:.2e}  speed={out['manual'][3]/out['sdpa'][3]:.2f}x", flush=True)
+print("SDPA_GATE_DONE", flush=True)