File size: 2,732 Bytes
324ecd2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | """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 = [], []
# Each bit is exactly +/-1 in BF16. Normalized scalar clock is recovered
# as a weighted sum, without quantizing the absolute position itself.
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))
|