--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, 12 insertions, 3 deletions

diff --git a/ep_run/lt_ep_train.py b/ep_run/lt_ep_train.py
index e7155d3..1307702 100644
--- a/ep_run/lt_ep_train.py
+++ b/ep_run/lt_ep_train.py
@@ -63,6 +63,9 @@ class EQBlock:
         if getattr(self, 'qknorm', False):                   # Qwen3-style q/k RMSNorm: bounds logits, tames J
             q = q * torch.rsqrt(q.pow(2).mean(-1, keepdim=True) + 1e-6)
             k = k * torch.rsqrt(k.pow(2).mean(-1, keepdim=True) + 1e-6)
+        if getattr(self, '_sdpa', False):                    # fused flash path — no_grad relax/eval only (same
+            o = F.scaled_dot_product_attention(q, k, v, is_causal=True)   # scale 1/sqrt(dh), same causal mask)
+            return o.transpose(1, 2).reshape(B, self.T, self.C) @ self.WO
         a = (q @ k.transpose(-2, -1)) / math.sqrt(self.dh)
         a = torch.softmax(a.masked_fill(~self.cmask, float('-inf')), -1)
         return (a @ v).transpose(1, 2).reshape(B, self.T, self.C) @ self.WO
@@ -127,9 +130,13 @@ def relax(blk, z, xin, steps, eps):
             for _ in range(steps):
                 z = cstep(z, xin)
         return z.detach()
-    for _ in range(steps):
-        with torch.no_grad():
-            z = z + eps * blk.force(z, xin).detach()
+    blk._sdpa = getattr(blk, 'sdpa', False)              # fused attention for the pure-forward loop only
+    try:
+        for _ in range(steps):
+            with torch.no_grad():
+                z = z + eps * blk.force(z, xin).detach()
+    finally:
+        blk._sdpa = False                                # grad paths (jvp/vjp/resreg graphs) stay on manual attn
     return z.detach()
 
 
@@ -421,6 +428,7 @@ def main():
     ap.add_argument('--navg', type=int, default=1)             # restart-averaged contrast estimates per update
     ap.add_argument('--track', action='store_true')            # common-mode-tracking AEP correction
     ap.add_argument('--holofast', action='store_true')         # exact halved-jvp track (1.55x nudged phase; parity = FD noise floor)
+    ap.add_argument('--sdpa', action='store_true')             # fused flash attention in the no_grad relax loop
     ap.add_argument('--rt_final', type=float, default=0.0)     # anneal res_target to this (0=off), 25%-75% of run
     ap.add_argument('--nudge_brake', type=float, default=0.0)  # kappa: anchor spring during nudge (Tikhonov adjoint)
     ap.add_argument('--init_ckpt', type=str, default='')       # warm-start weights from a saved ckpt
@@ -496,6 +504,7 @@ def main():
     blk.navg = cfg.navg
     blk.track = cfg.track
     blk.holofast = cfg.holofast
+    blk.sdpa = cfg.sdpa
     blk.nbrake = cfg.nudge_brake
     blk.qknorm = cfg.qknorm
     if cfg.resinit != 1.0:                          # near-identity block at init (contractive) -> stable big-width start