aggregate speed bench: speed tier hf+sd 1.47x; accuracy tier hf+sd+t80+avg (0.91x, cos 0.94, avg free); compile demoted

Full-ep_step wall times on quiet A6000 (warm s2000, B24), res parity across
all 8 configs. compile only 1.12x at this shape (historical 1.46x was a
different workload split); FULL cmp_sdpa saves 4% over eager at t80 — not
worth the guard complexity. tforce_sdpa added (flash baked into compiled
graph, flag-free so grad paths never see SDPA). bp_lm --tie probe in flight.

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

1 files changed, 16 insertions, 1 deletions

diff --git a/ep_run/lt_ep_train.py b/ep_run/lt_ep_train.py
index 2807be3..76bb9bb 100644
--- a/ep_run/lt_ep_train.py
+++ b/ep_run/lt_ep_train.py
@@ -81,6 +81,20 @@ class EQBlock:
     def Emem(self, z):
         return -(F.relu(z @ self.Wm) ** 2).sum()
 
+    def tforce_sdpa(self, z, xin):                       # tforce with fused flash attention (compiled fast path only;
+        h1 = F.layer_norm(z, (self.C,), self.ln1g, self.ln1b)      # grad paths keep the manual attn)
+        h2 = F.layer_norm(z, (self.C,), self.ln2g, self.ln2b)
+        B = z.size(0)
+        q = (h1 @ self.WQ).view(B, self.T, self.H, self.dh).transpose(1, 2)
+        k = (h1 @ self.WK).view(B, self.T, self.H, self.dh).transpose(1, 2)
+        v = (h1 @ self.WV).view(B, self.T, self.H, self.dh).transpose(1, 2)
+        if getattr(self, 'qknorm', False):
+            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)
+        att = F.scaled_dot_product_attention(q, k, v, is_causal=True).transpose(1, 2).reshape(B, self.T, self.C) @ self.WO
+        ff = F.gelu(h2 @ self.fc + self.fcb, approximate='tanh') @ self.pj + self.pjb
+        return -(z - xin) + att + ff - self.c * z
+
     def tforce(self, z, xin):                            # pure thick force (no grad machinery) -> torch.compile
         h1 = F.layer_norm(z, (self.C,), self.ln1g, self.ln1b)
         h2 = F.layer_norm(z, (self.C,), self.ln2g, self.ln2b)
@@ -523,7 +537,8 @@ def main():
     blk._cstep = None
     if cfg.compile and cfg.attn_mode == 'thick':
         _ee = cfg.eps
-        blk._cstep = torch.compile(lambda z, xin: z + _ee * blk.tforce(z, xin))
+        _tf = blk.tforce_sdpa if cfg.sdpa else blk.tforce   # sdpa baked into the compiled graph, NO ambient flags
+        blk._cstep = torch.compile(lambda z, xin: z + _ee * _tf(z, xin))
     mis = None
     if cfg.wmis > 0:                                # fixed fabrication mismatch (same devices all run)
         gm = torch.Generator().manual_seed(1234)