"""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)