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

4 files changed, 89 insertions, 2 deletions

diff --git a/ep_run/agg_bench.log b/ep_run/agg_bench.log
new file mode 100644
index 0000000..35cb799
--- /dev/null
+++ b/ep_run/agg_bench.log
@@ -0,0 +1,13 @@
+/home/yurenh2/miniconda3/lib/python3.13/site-packages/torch/autograd/graph.py:865: UserWarning: Attempting to run cuBLAS, but there was no current CUDA context! Attempting to set the primary context... (Triggered internally at /pytorch/aten/src/ATen/cuda/CublasHandlePool.cpp:330.)
+  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
+            base:   4.50s/step  (res 8.2e-03)
+              hf:   3.49s/step  (res 8.2e-03)
+              sd:   4.07s/step  (res 8.2e-03)
+           hf+sd:   3.06s/step  (res 8.2e-03)
+       hf+sd+t80:   4.97s/step  (res 8.2e-03)
+   hf+sd+t80+avg:   4.97s/step  (res 8.2e-03)
+             cmp:   4.03s/step  (res 8.2e-03)
+/home/yurenh2/miniconda3/lib/python3.13/site-packages/torch/_inductor/compile_fx.py:321: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance.
+  warnings.warn(
+  FULL(cmp_sdpa):   4.77s/step  (res 8.2e-03)
+AGG_BENCH_DONE
diff --git a/ep_run/agg_bench.py b/ep_run/agg_bench.py
new file mode 100644
index 0000000..d44b5a8
--- /dev/null
+++ b/ep_run/agg_bench.py
@@ -0,0 +1,51 @@
+"""THE AGGREGATE SPEED BENCH: full ep_step wall-time for every combination of the speed levers, on a
+quiet local GPU, warm s2000 operator, B24 (production shape). Configs:
+  base           : orig track, t2sel40, eager, manual attn      (the historical default)
+  hf             : +holofast
+  sd             : +sdpa (eager relax)
+  hf+sd          : the exact-math pack at t2sel40
+  hf+sd+t80      : the accuracy pack (cos 0.89->0.94)
+  hf+sd+t80+avg  : + holoavg (trend/plateau estimator)
+  cmp            : --compile alone (manual attn in graph)
+  cmp(sdpa)+hf+t80+avg : FULL STACK (flash baked into compiled graph)
+Reports median full-step time of 3 (after 1 warmup step each) + the step's res, as parity smoke."""
+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))
+blk.track = True
+torch.manual_seed(42)
+idx, y = L.get_batch('train', 24, 256)
+
+CFG = [
+    ('base',            dict(hf=False, sd=False, t2=40,  avg=False, cmp=False)),
+    ('hf',              dict(hf=True,  sd=False, t2=40,  avg=False, cmp=False)),
+    ('sd',              dict(hf=False, sd=True,  t2=40,  avg=False, cmp=False)),
+    ('hf+sd',           dict(hf=True,  sd=True,  t2=40,  avg=False, cmp=False)),
+    ('hf+sd+t80',       dict(hf=True,  sd=True,  t2=80,  avg=False, cmp=False)),
+    ('hf+sd+t80+avg',   dict(hf=True,  sd=True,  t2=80,  avg=True,  cmp=False)),
+    ('cmp',             dict(hf=False, sd=False, t2=40,  avg=False, cmp=True)),
+    ('FULL(cmp_sdpa)',  dict(hf=True,  sd=True,  t2=80,  avg=True,  cmp=True)),
+]
+
+for name, c in CFG:
+    blk.holofast, blk.sdpa, blk.holoavg = c['hf'], c['sd'], c['avg']
+    blk._cstep = None
+    if c['cmp']:
+        _tf = blk.tforce_sdpa if c['sd'] else blk.tforce
+        blk._cstep = torch.compile(lambda z, xin, _tf=_tf: z + 0.1 * _tf(z, xin))
+    ts = []
+    for rep in range(4):                                   # rep 0 = warmup (compile/JIT)
+        torch.cuda.synchronize(); t = time.time()
+        _, res = L.ep_step(blk, idx, y, 150, 20, 0.1, 0.02, jacreg=0.1, holo=2, hr=0.02,
+                           t1max=300, res_est=1e-4, t2sel=c['t2'], corr_every=1, res_gate=0.0, resreg=0.2)
+        torch.cuda.synchronize(); ts.append(time.time() - t)
+    med = sorted(ts[1:])[1]
+    print(f"{name:>16}: {med:6.2f}s/step  (res {res:.1e})", flush=True)
+blk._cstep = None
+print("AGG_BENCH_DONE", flush=True)
diff --git a/ep_run/bp_lm.py b/ep_run/bp_lm.py
index c01cd4f..1bd12f8 100644
--- a/ep_run/bp_lm.py
+++ b/ep_run/bp_lm.py
@@ -13,7 +13,7 @@ def fwd(blk, idx):
     h1 = F.layer_norm(x, (blk.C,), blk.ln1g, blk.ln1b)
     h2 = F.layer_norm(x, (blk.C,), blk.ln2g, blk.ln2b)
     h = x + blk.attn(h1) + (F.gelu(h2 @ blk.fc + blk.fcb, approximate='tanh') @ blk.pj + blk.pjb)
-    return h @ blk.Wh
+    return h @ (blk.tok.t() if getattr(blk, 'tie', False) else blk.Wh)
 
 
 def evaluate(blk, nb=8, B=32):
@@ -36,6 +36,7 @@ def main():
     ap.add_argument('--stdinit', action='store_true')          # standard transformer init (EQBlock's is tuned for relaxation)
     ap.add_argument('--beta2', type=float, default=0.999)
     ap.add_argument('--sched', choices=['cos', 'const'], default='cos')
+    ap.add_argument('--tie', action='store_true')              # tok/Wh weight tying (standard small-LM trick)
     ap.add_argument('--log', type=int, default=200)
     ap.add_argument('--ckpt', type=str, default='runs/bp_lm.pt')
     cfg = ap.parse_args()
@@ -49,6 +50,13 @@ def main():
             for W in (blk.WO, blk.pj):
                 W.normal_(0, 0.02 / (2 ** 0.5))
             blk.pos.normal_(0, 0.01)
+    if cfg.tie:                                                # tie head to embedding: Wh := tok^T, single parameter
+        with torch.no_grad():
+            blk.tok.copy_(0.5 * (blk.tok + blk.Wh.t()))
+        blk.Wh = None                                          # fwd() will use tok.t() when tie is on
+        blk.tie = True
+        blk.allp = [p for p in blk.allp if p is not blk.Wh]
+        blk.allp = blk.block + []                              # block already contains tok; Wh dropped
     opt = torch.optim.AdamW(blk.allp, lr=cfg.lr, weight_decay=cfg.wd, betas=(0.9, cfg.beta2))
     if cfg.sched == 'cos':
         sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, cfg.steps, eta_min=cfg.lr * 0.05)
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)