| """Probe the BeamAttention RAM (beamboot/attn_ram.py) at long horizons. |
| |
| Seeded random writes/reads vs a dict. Clock = row position, T = 2^18 (or the |
| next power of two >= horizon). Reports read errors, the minimum logit margin |
| (true winner minus best other key, fp32 recompute from the real cache keys and |
| the module's own q_norm/q_proj/query scaling), rows/s of the attention forward |
| alone, and peak GPU memory. |
| |
| .venv-gpu/bin/python scripts/probe_attn_ram.py --horizons 4096 32768 262144 |
| """ |
| import argparse |
| import json |
| import random |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import torch |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| from beamboot.attn_ram import AttnRAM, HEAD |
| from transformers.models.beam.modeling_beam import compute_attention_scaling |
|
|
|
|
| def make_stream(rng, n): |
| """Writes ~70% (bursts to one address model CHIP-8 V/I/stack traffic), reads ~30%.""" |
| ev, hot = [], [rng.randrange(4096) for _ in range(96)] |
| while len(ev) < n: |
| r = rng.random() |
| if r < 0.10: |
| a = rng.choice(hot) |
| for _ in range(rng.randrange(2, 5)): ev.append(('w', a, rng.randrange(256))) |
| ev.append(('r', a)) |
| elif r < 0.70: |
| a = rng.choice(hot) if rng.random() < 0.6 else rng.randrange(4096) |
| ev.append(('w', a, rng.randrange(256))) |
| else: |
| a = rng.choice(hot) if rng.random() < 0.6 else rng.randrange(4096) |
| ev.append(('r', a)) |
| return ev[:n] |
|
|
|
|
| def probe(horizon, impl, dtype, device, chunk, seed, margins=True): |
| T = max(2 ** 18, 1 << (horizon - 1).bit_length()) |
| torch.cuda.reset_peak_memory_stats() if device == 'cuda' else None |
| ram = AttnRAM(max_clock=T, device=device, dtype=dtype, attn_impl=impl) |
| events = make_stream(random.Random(seed), horizon - 1) |
| latest, ref = {}, {} |
| errors = reads = 0 |
| min_margin, fwd_time = float('inf'), 0.0 |
| L = ram.layer |
| for start in range(0, len(events), chunk): |
| block = events[start:start + chunk] |
| past = ram.length |
| win, want, qpos, qaddr = [], [], [], [] |
| for i, e in enumerate(block): |
| p = past + i |
| if e[0] == 'w': latest[e[1]] = p; ref[e[1]] = e[2] |
| else: |
| win.append(latest.get(e[1], 0)); want.append(ref.get(e[1], 0)) |
| qpos.append(p); qaddr.append(e[1]) |
| if device == 'cuda': torch.cuda.synchronize() |
| t0 = time.perf_counter() |
| got = ram.run(block) |
| if device == 'cuda': torch.cuda.synchronize() |
| fwd_time += time.perf_counter() - t0 |
| reads += len(got) |
| errors += sum(g != w for g, w in zip(got, want)) |
| if margins and qpos: |
| x = torch.stack([ram.read_row(a) for a in qaddr]).to(device, dtype)[None] |
| with torch.inference_mode(): |
| q = L.q_norm(L.q_proj(x).view(1, -1, 1, HEAD)).transpose(1, 2).float() |
| pos = torch.tensor([qpos], device=device) |
| q = (q * compute_attention_scaling(pos))[0, 0] |
| k = ram.cache.layers[0].keys[0, 0].float() |
| s = (q @ k.T) * L.scaling |
| cols = torch.arange(k.shape[0], device=device)[None] |
| s.masked_fill_(cols > pos[0][:, None], float('-inf')) |
| w = torch.tensor(win, device=device) |
| ws = s.gather(1, w[:, None])[:, 0] |
| s.scatter_(1, w[:, None], float('-inf')) |
| min_margin = min(min_margin, float((ws - s.max(1).values).min())) |
| res = dict(horizon=horizon, max_clock=T, impl=impl, dtype=str(dtype).split('.')[-1], device=device, |
| rows=ram.length, reads=reads, errors=errors, |
| min_margin=round(min_margin, 2) if margins else None, |
| rows_per_s=round(ram.length / fwd_time, 1), fwd_seconds=round(fwd_time, 2)) |
| if device == 'cuda': res['peak_gpu_MiB'] = round(torch.cuda.max_memory_allocated() / 2 ** 20) |
| del ram |
| if device == 'cuda': torch.cuda.empty_cache() |
| return res |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--horizons', type=int, nargs='+', default=[4096, 32768, 262144]) |
| ap.add_argument('--impls', nargs='+', default=['sdpa', 'eager']) |
| ap.add_argument('--dtypes', nargs='+', default=['float32', 'bfloat16']) |
| ap.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu') |
| ap.add_argument('--chunk', type=int, default=512) |
| ap.add_argument('--seed', type=int, default=18) |
| a = ap.parse_args() |
| for dt in a.dtypes: |
| for h in a.horizons: |
| for impl in a.impls: |
| try: |
| r = probe(h, impl, getattr(torch, dt), a.device, a.chunk, a.seed) |
| except torch.OutOfMemoryError as exc: |
| r = dict(horizon=h, impl=impl, dtype=dt, error=f'OOM: {str(exc)[:80]}') |
| print(json.dumps(r), flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
