From 9a8b2796ca12e4d4c24717485a635a301aa6d07f Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Mon, 6 Jul 2026 09:17:36 -0500 Subject: aggregate speed bench: speed tier hf+sd 1.47x; accuracy tier hf+sd+t80+avg (0.91x, cos 0.94, avg free); compile demoted MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/agg_bench.py | 51 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 ep_run/agg_bench.py (limited to 'ep_run/agg_bench.py') 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) -- cgit v1.2.3