| """Find concrete RAM reads in the in-weights CHIP-8 computer (beamboot/chip8_lm.py). |
| |
| Runs the constructed BeamForCausalLM on one real SUPERBEAM frame (checkpoint prompt + the frame's |
| tokens, as chip8_fast does) with eager attention. For the RAM-read layers (L4: opcode and memory |
| reads, L5: Vx/Vy reads) it lists query positions whose attention sits on a single write token while |
| older writes to the same address exist earlier in the context: the "latest-write RAM" circuit, |
| read straight out of the model's attention weights. |
| |
| .venv-gpu/bin/python scripts/attention_probe.py --frames 6 --out demo/attention_probe.json |
| """ |
| import argparse |
| import json |
| import os |
| import sys |
|
|
| import torch |
|
|
| ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| sys.path.insert(0, ROOT) |
|
|
| from beamboot import chip8_fast as F |
| from beamboot import chip8_lm as C |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--frames', type=int, default=6, help='advance this many frames first') |
| ap.add_argument('--draft-tokens', type=int, default=600) |
| ap.add_argument('--out', default=os.path.join(ROOT, 'demo', 'attention_probe.json')) |
| args = ap.parse_args() |
|
|
| rom = open(os.path.join(ROOT, 'roms', 'superbeam.ch8'), 'rb').read() |
| sched = json.load(open(os.path.join(ROOT, 'roms', 'superbeam_inputs.json'))) |
| s = F.run(rom, [e for e in sched if int(e[0]) < args.frames], frames=args.frames, device='cpu') |
| for g, k, a in sched: |
| if int(g) == args.frames: |
| s.c.set_key(int(k) & 15, a == 'down') |
| p = F.checkpoint_prompt(s.c, s.ipf) |
| d = s._draft([])[:args.draft_tokens] |
| seq = p + d |
| m = s.model |
| m.config._attn_implementation = 'eager' |
| with torch.inference_mode(): |
| out = m(input_ids=torch.tensor([seq]), output_attentions=True, logits_to_keep=len(d) + 1) |
| am = out.logits[0].argmax(-1).tolist() |
| agree = sum(1 for j, t in enumerate(d) if am[j] == t) |
| labels = [tuple(C.untok(t)) for t in seq] |
|
|
| def addr(lb): |
| return lb[0], (lb[2] if lb[0] == 'VW' and len(lb) > 2 else None) |
|
|
| mem = s.c.mem |
|
|
| def pc_at(q): |
| return next(int(labels[i][1]) for i in range(q - 1, -1, -1) if labels[i][0] == 'PCW') |
|
|
| found = [] |
| for li in (4,): |
| A = out.attentions[li][0] |
| for h in range(A.shape[0]): |
| for q in range(len(p), len(seq)): |
| if labels[q][0] != 'CLKL': |
| continue |
| w = A[h, q, :q + 1] |
| kpos = int(w.argmax()) |
| if float(w[kpos]) < 0.99 or labels[kpos][0] not in ('VW', 'PCW', 'IW'): |
| continue |
| if any(addr(labels[i]) == addr(labels[kpos]) for i in range(kpos + 1, q)): |
| continue |
| same = [i for i in range(kpos) if addr(labels[i]) == addr(labels[kpos])] |
| if len(same) < 2: |
| continue |
| pc = pc_at(q) |
| found.append(dict(layer=li + 1, head=h, query=q, pc=pc, opcode='%02X%02X' % (mem[pc], mem[pc + 1]), |
| older_vals=[str(labels[i][1]) for i in same[-6:]], query_tok=[str(x) for x in labels[q]], |
| key=kpos, key_tok=[str(x) for x in labels[kpos]], weight=float(w[kpos]), |
| older_same=len(same), older_pos=same[-20:], |
| older_weight=float(w[same].sum()) if same else 0.0, ctx=q + 1)) |
| found.sort(key=lambda f: (-f['older_same'], f['layer'])) |
|
|
| |
| |
| |
| A0 = out.attentions[4][0][0] |
| uses_vx = set('345678ABEF') - {'A', 'B'} |
| checked = correct = 0 |
| for q in range(len(p), len(seq)): |
| if labels[q][0] != 'CLKL': |
| continue |
| pc = pc_at(q) |
| op = '%02X%02X' % (mem[pc], mem[pc + 1]) |
| if op[0] not in uses_vx: |
| continue |
| x = int(op[1], 16) |
| writes = [i for i in range(q) if labels[i][0] == 'VW' and int(labels[i][2]) == x] |
| if not writes: |
| continue |
| checked += 1 |
| kpos = int(A0[q, :q + 1].argmax()) |
| correct += int(kpos == writes[-1] and float(A0[q, kpos]) > 0.99) |
| stats = dict(head='L5 head 0 (Vx read)', instructions_checked=checked, read_newest_write=correct) |
|
|
| fig = None |
| best = [f for f in found if f['head'] == 0 and f['opcode'][0] in '34' and f['older_same'] >= 5] |
| if best: |
| f = best[0] |
| q, x = f['query'], int(f['key_tok'][2]) |
| w = A0[q, :q + 1] |
| writes = [i for i in range(q) if labels[i][0] == 'VW' and int(labels[i][2]) == x] |
| fig = dict(query=q, ctx=q + 1, pc=f['pc'], opcode=f['opcode'], reg=x, |
| writes=[dict(pos=i, val=int(labels[i][1]), w=float(w[i])) for i in writes], |
| other_mass=float(w.sum() - w[writes].sum()), |
| kinds=[labels[i][0] for i in range(q + 1)]) |
| res = dict(frame=args.frames, prompt_tokens=len(p), checked_tokens=len(d), draft_agree=agree, |
| n_reads=len(found), vx_read=stats, figure=fig, examples=found[:40]) |
| print(json.dumps(stats)) |
| os.makedirs(os.path.dirname(args.out), exist_ok=True) |
| json.dump(res, open(args.out, 'w'), indent=1) |
| print(json.dumps({k: v for k, v in res.items() if k != 'examples'})) |
| for f in found[:10]: |
| print({k: v for k, v in f.items() if k != 'older_pos'}) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
