1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
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)
|
