| |
| """Replay a SUPERBEAM key schedule inside the constructed BeamForCausalLM, fast path |
| (beamboot/chip8_fast.py: per-frame checkpoints + speculative verification). |
| |
| .venv-gpu/bin/python scripts/run_chip8_fast.py --inputs results/beam_key_schedule.json \ |
| --frames 175 --out demo/beam_model_frames.json |
| |
| Writes the same JSON shape as scripts/run_chip8_lm.py (frames = 32 hex rows per frame) plus |
| per-frame stats, and compares every frame with the reference emulator (superbeam_gym replay). |
| """ |
| import argparse |
| import json |
| import os |
| import sys |
| import time |
|
|
| ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| sys.path.insert(0, ROOT) |
|
|
| import torch |
|
|
| from beamboot import chip8_fast as F |
| from beamboot.superbeam_gym import SuperBeamEnv, frames_from_events |
|
|
|
|
| def emulator_rows(events, n): |
| env = SuperBeamEnv(backend='emulator') |
| env.reset() |
| out = [] |
| for keys in frames_from_events(events, n): |
| env.backend.advance(keys) |
| out.append(env.backend.m.display_rows()) |
| return out |
|
|
|
|
| def save(path, s, ref, args, wall): |
| st = s.stats |
| k = 0 |
| while k < len(s.frame_rows) and s.frame_rows[k] == ref[k]: |
| k += 1 |
| secs = [x.seconds for x in st] |
| out = dict(rom='roms/superbeam.ch8', inputs_file=os.path.relpath(args.inputs, ROOT), ipf=s.ipf, |
| frames_run=len(s.frame_rows), verified=k == len(s.frame_rows), |
| frames_equal_emulator=k, mismatch=None if k == len(s.frame_rows) else k, |
| method='per-frame checkpoint + speculative verification (beamboot/chip8_fast.py)', |
| frames_all_draft_accepted=sum(x.accepted_all for x in st), |
| n_tokens=sum(x.frame_tokens for x in st), prompt_tokens_per_frame=st[-1].prompt_tokens, |
| seconds_per_frame=round(sum(secs) / len(secs), 3), max_seconds_per_frame=round(max(secs), 3), |
| wall_seconds=round(wall, 1), |
| tok_per_s=round(sum(x.frame_tokens for x in st) / sum(secs), 1), |
| device=torch.cuda.get_device_name() if torch.cuda.is_available() else 'cpu', |
| model=dict(cls=type(s.model).__name__, layers=s.model.config.num_hidden_layers, |
| hidden=s.model.config.hidden_size, vocab=s.model.config.vocab_size, |
| dtype='float32', attn='sdpa'), |
| per_frame=[dict(f=x.frame, tokens=x.frame_tokens, accepted_all=x.accepted_all, |
| first_reject=x.first_reject, s=round(x.seconds, 3)) for x in st], |
| frames=s.frame_rows) |
| with open(path + '.tmp', 'w') as f: |
| json.dump(out, f) |
| os.replace(path + '.tmp', path) |
| return out |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--inputs', default=os.path.join(ROOT, 'results', 'beam_key_schedule.json')) |
| ap.add_argument('--frames', type=int, default=None) |
| ap.add_argument('--out', default=os.path.join(ROOT, 'demo', 'beam_model_frames.json')) |
| ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu') |
| args = ap.parse_args() |
| sched = json.load(open(args.inputs)) |
| events = sched['events'] if isinstance(sched, dict) else sched |
| n = args.frames or (sched.get('frames') if isinstance(sched, dict) else None) |
| if not n: |
| n = SuperBeamEnv(backend='emulator', frame_skip=1).replay_schedule(events)['frames'] |
| rom = open(os.path.join(ROOT, 'roms', 'superbeam.ch8'), 'rb').read() |
| ref = emulator_rows(events, n) |
| t0 = time.time() |
|
|
| def on_frame(g, s): |
| x = s.stats[-1] |
| ok = s.frame_rows[-1] == ref[g] |
| print('frame %3d: %4d tokens %.2fs draft %s frame==emulator %s' |
| % (g, x.frame_tokens, x.seconds, 'accepted' if x.accepted_all else 'REJECTED at %d' % x.first_reject, |
| ok if ok else 'MISMATCH'), flush=True) |
| if (g + 1) % 10 == 0: |
| save(args.out, s, ref, args, time.time() - t0) |
|
|
| s = F.run(rom, events, frames=n, device=args.device, on_frame=on_frame) |
| out = save(args.out, s, ref, args, time.time() - t0) |
| print('done: %d frames, equal to emulator %d, all drafts accepted %d, %.2f s/frame, %.0f tok/s, wall %.0fs' |
| % (out['frames_run'], out['frames_equal_emulator'], out['frames_all_draft_accepted'], |
| out['seconds_per_frame'], out['tok_per_s'], out['wall_seconds'])) |
| return 0 if out['verified'] else 1 |
|
|
|
|
| if __name__ == '__main__': |
| raise SystemExit(main()) |
|
|
