| import random |
| import unittest |
| import torch |
| from beamboot.rope_memory import RotaryCacheMemory |
|
|
|
|
| class RotaryMemoryTests(unittest.TestCase): |
| def check_memory(self, dtype): |
| memory = RotaryCacheMemory(dtype) |
| reference = [0]*16 |
| rng = random.Random(112) |
| for _ in range(200): |
| address, value = rng.randrange(16), rng.randrange(16) |
| memory.write(address, value) |
| reference[address] = value |
| self.assertEqual(memory.read(address), value) |
| for address, value in enumerate(reference): |
| self.assertEqual(memory.read(address), value) |
| self.assertEqual(memory.position, 432) |
|
|
| def test_float32(self): self.check_memory(torch.float32) |
| def test_bfloat16(self): self.check_memory(torch.bfloat16) |
|
|
| def test_evicted_write_is_not_reliably_readable(self): |
| memory = RotaryCacheMemory(window=32) |
| memory.write(0, 7) |
| for _ in range(32): memory.write(1, 3) |
| self.assertNotEqual(memory.read(0), 7) |
| self.assertEqual(memory.read(1), 3) |
|
|
