| """Binary-clock memory probe through Transformers' BeamAttention implementation. |
| |
| This is a small constructed layer, not pretrained Beam. A host provides write |
| positions as bits. The eager Beam backend materializes attention scores in the |
| layer's dtype, unlike the separate SDPA probe's implementation. |
| """ |
| import math |
| import torch |
| from transformers import BeamConfig |
| from transformers.models.beam.modeling_beam import BeamAttention, compute_attention_scaling |
| from .attention_cpu import bits, from_bits |
|
|
|
|
| class HFBinaryMemory: |
| def __init__(self, max_writes=128, dtype=torch.float32): |
| if not 1 <= max_writes <= 2047: raise ValueError('max_writes must be 1..2047') |
| self.maximum, self.dtype, self.writes = max_writes, dtype, 0 |
| self.width = max_writes.bit_length() |
| self.clock = [2**i/(2*(max_writes+1)) for i in range(self.width)] |
| config = BeamConfig(hidden_size=128, num_hidden_layers=1, num_attention_heads=6, |
| num_key_value_heads=1, head_dim=128, layer_types=['full_attention'], |
| mlp_layer_types=['dense'], attention_dropout=0.) |
| config._attn_implementation = 'eager' |
| self.layer = BeamAttention(config, 0).eval().to(dtype).requires_grad_(False) |
| key_start = 6 |
| value_start = key_start + 5 + self.width |
| with torch.no_grad(): |
| for m in (self.layer.q_proj, self.layer.k_proj, self.layer.v_proj, |
| self.layer.o_proj, self.layer.gate_proj): m.weight.zero_() |
| for head in range(6): |
| off = head*128 |
| for i in range(4): self.layer.q_proj.weight[off+i, i] = 1 |
| for i, coefficient in enumerate(self.clock): self.layer.q_proj.weight[off+4+i, 4] = coefficient |
| self.layer.q_proj.weight[off+4+self.width, 4] = 2 |
| for i in range(5+self.width): self.layer.k_proj.weight[i, key_start+i] = 1 |
| for i in range(4): |
| self.layer.v_proj.weight[i, value_start+i] = 1 |
| self.layer.o_proj.weight[i, i] = 2 |
| self.layer.gate_bias.zero_() |
| q2 = 8 + sum(x*x for x in self.clock) |
| k2 = 5 + self.width |
| gain = 16*(max_writes+1)*math.sqrt(q2*k2)/math.sqrt(128) |
| self.layer.q_norm.weight.fill_(gain) |
| self.layer.k_norm.weight.fill_(1) |
| self.rows = [self._row(a, 0, 0, True) for a in range(16)] |
|
|
| def _row(self, address, value, position, valid): |
| row = torch.zeros(128, dtype=self.dtype) |
| row[:4] = torch.tensor([2*b-1 for b in bits(address, 4)], dtype=self.dtype) |
| row[4] = 1 |
| key = [*(2*b-1 for b in bits(address, 4)), |
| *(2*b-1 for b in bits(position, self.width)), 1 if valid else -1] |
| row[6:6+len(key)] = torch.tensor(key, dtype=self.dtype) |
| value_start = 6 + 5 + self.width |
| row[value_start:value_start+4] = torch.tensor(bits(value, 4), dtype=self.dtype) |
| return row |
|
|
| 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 budget exhausted') |
| self.writes += 1 |
| self.rows.append(self._row(address, value, self.writes, True)) |
|
|
| def read(self, address): |
| if not 0 <= address < 16: raise ValueError('expected nibble address') |
| inputs = torch.stack(self.rows + [self._row(address, 0, self.writes, False)]).unsqueeze(0) |
| length = inputs.shape[1] |
| mask = torch.full((1, 1, length, length), float('-inf'), dtype=self.dtype).triu(1) |
| positions = torch.arange(length).unsqueeze(0) |
| with torch.inference_mode(): |
| out, _ = self.layer(inputs, position_embeddings=None, attention_mask=mask, |
| attention_scaling=compute_attention_scaling(positions)) |
| return from_bits((out[0, -1, :4] > .5).int().tolist()) |
|
|
