From 75ff326dcb40cd960abd56e9c9c18a45d9e5e2c2 Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Sun, 5 Jul 2026 05:09:36 -0500 Subject: =?UTF-8?q?--sdpa:=20fused=20flash=20attention=20in=20the=20no=5Fg?= =?UTF-8?q?rad=20relax=20loop=20=E2=80=94=201.45x=20free=20phase,=20z*=20p?= =?UTF-8?q?arity=204e-7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/sdpa_gate.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) create mode 100644 ep_run/sdpa_gate.py (limited to 'ep_run/sdpa_gate.py') 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) -- cgit v1.2.3