| import random |
| import unittest |
| import torch |
| from beamboot.binary_time_probe import BinaryTimeProbe |
| from beamboot.hf_binary_memory import HFBinaryMemory |
| from beamboot.causal_lm import CausalMachine |
|
|
|
|
| class BinaryPrecisionTests(unittest.TestCase): |
| def test_bf16_binary_clock_and_retention(self): |
| memory = BinaryTimeProbe(512, torch.bfloat16) |
| rng = random.Random(417) |
| reference = [0]*16 |
| for i in range(512): |
| address, value = rng.randrange(16), rng.randrange(16) |
| memory.write(address, value) |
| reference[address] = value |
| self.assertEqual(memory.read(address), value, f'write {i+1}') |
| for address, value in enumerate(reference): |
| self.assertEqual(memory.read(address), value) |
|
|
| def test_hf_beam_attention_sdpa_bf16(self): |
| memory = HFBinaryMemory(64, torch.bfloat16) |
| memory.layer.config._attn_implementation = 'sdpa' |
| rng = random.Random(417) |
| for _ in range(64): |
| address, value = rng.randrange(16), rng.randrange(16) |
| memory.write(address, value) |
| self.assertEqual(memory.read(address), value) |
|
|
| def test_known_unsafe_runtime_is_rejected(self): |
| with self.assertRaisesRegex(ValueError, 'known incorrect'): |
| CausalMachine._validate_precision(torch.bfloat16, 'eager') |
| CausalMachine._validate_precision(torch.bfloat16, 'sdpa') |
|
|
