| """Bounded attention memory with an exactly representable binary clock. |
| |
| Like beam_probe, this is an isolated mathematical probe, not a Beam checkpoint. |
| Position is stored as bipolar bits so bfloat16 does not collapse adjacent values. |
| It still assumes an explicitly supplied, bounded write counter. |
| """ |
| import math |
| import torch |
| import torch.nn.functional as F |
| from .attention_cpu import bits |
|
|
|
|
| class BinaryTimeProbe: |
| def __init__(self, max_writes=128, dtype=torch.float32): |
| if not isinstance(max_writes, int) or not 1 <= max_writes <= 4095: |
| raise ValueError('max_writes must be 1..4095') |
| self.maximum, self.dtype = max_writes, dtype |
| self.width = max_writes.bit_length() |
| self.writes = 0 |
| self.keys, self.values = [], [] |
| |
| |
| self.weights = [2**i / (2 * (max_writes+1)) for i in range(self.width)] |
| q2 = 4 + sum(w*w for w in self.weights) |
| k2 = 4 + self.width |
| self.gain = 16*(max_writes+1)*math.sqrt(q2*k2)/math.sqrt(128) |
| for address in range(16): self._append(address, 0, 0) |
|
|
| def _append(self, address, value, position): |
| key = [*(2*b-1 for b in bits(address, 4)), |
| *(2*b-1 for b in bits(position, self.width))] |
| self.keys.append(key + [0.]*(128-len(key))) |
| self.values.append(bits(value, 4) + [0.]*124) |
|
|
| def write(self, address, value): |
| if not 0 <= address < 16 or not 0 <= value < 16: raise ValueError('expected nibbles') |
| if self.writes >= self.maximum: raise ValueError('write capacity exceeded') |
| self.writes += 1 |
| self._append(address, value, self.writes) |
|
|
| @staticmethod |
| def norm(x): |
| return x * torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + 1e-6).to(x.dtype) |
|
|
| def read(self, address): |
| if not 0 <= address < 16: raise ValueError('expected nibble address') |
| q = [*(2*b-1 for b in bits(address, 4)), *self.weights] |
| query = torch.tensor(q + [0.]*(128-len(q)), dtype=self.dtype) |
| keys = torch.tensor(self.keys, dtype=self.dtype) |
| values = torch.tensor(self.values, dtype=self.dtype) |
| query = self.norm(query) * self.gain |
| keys = self.norm(keys) |
| position = len(self.keys) |
| query = query * ((1 + (position+1)/2048)**.1) |
| q = query.reshape(1, 1, 1, 128).expand(1, 6, 1, 128) |
| k = keys.reshape(1, 1, -1, 128) |
| v = values.reshape(1, 1, -1, 128) |
| out = F.scaled_dot_product_attention(q, k, v, enable_gqa=True)[0, 0, 0, :4] |
| return sum((float(x) > .5) << i for i, x in enumerate(out)) |
|
|
