| """Reproduce bounded Beam-style attention precision sweep on CPU.""" |
| import argparse |
| import random |
| import torch |
| from beamboot.beam_probe import BeamAttentionProbe, calibrated_gain |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--horizons', nargs='+', type=int, default=[128, 512, 1024]) |
| parser.add_argument('--seed', type=int, default=417) |
| args = parser.parse_args() |
| for dtype in (torch.float32, torch.float16, torch.bfloat16): |
| for horizon in args.horizons: |
| rng = random.Random(args.seed) |
| memory = BeamAttentionProbe(horizon, dtype, calibrated_gain(horizon)) |
| errors = 0 |
| for _ in range(horizon): |
| address, value = rng.randrange(16), rng.randrange(16) |
| memory.write(address, value) |
| errors += memory.read(address) != value |
| print(f'{dtype} horizon={horizon} immediate-read-errors={errors}', flush=True) |
|
|
|
|
| if __name__ == '__main__': main() |
|
|
