| """Attention-backed nibble memory through a stock BeamDecoderLayer. |
| |
| This probe uses constructed weights, a compact one-group Beam config and |
| host-encoded event vectors. Only memory reads are implemented in this decoder; |
| the layer's MLP is zeroed. It does not use pretrained Beam weights. |
| """ |
| import math |
| import torch |
| from transformers import BeamConfig |
| from transformers.cache_utils import DynamicCache |
| from transformers.models.beam.modeling_beam import BeamDecoderLayer, compute_attention_scaling |
| from .attention_cpu import bits |
|
|
|
|
| class DecoderMemory: |
| def __init__(self, max_writes=128): |
| if not 1 <= max_writes <= 1024: raise ValueError('max_writes must be 1..1024') |
| self.maximum = max_writes |
| self.writes = 0 |
| config = BeamConfig(hidden_size=128, intermediate_size=16, 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 = BeamDecoderLayer(config, 0).eval().requires_grad_(False) |
| attn = self.layer.self_attn |
| with torch.no_grad(): |
| for proj in (attn.q_proj, attn.k_proj, attn.v_proj, attn.o_proj, attn.gate_proj, |
| self.layer.mlp.gate_proj, self.layer.mlp.up_proj, self.layer.mlp.down_proj): |
| proj.weight.zero_() |
| for head in range(6): |
| offset = 128*head |
| for i in range(5): attn.q_proj.weight[offset+i, i] = 1 |
| attn.q_proj.weight[offset+5, 5] = 5 |
| for i in range(7): attn.k_proj.weight[i, 6+i] = 1 |
| for i in range(16): |
| attn.v_proj.weight[i, 13+i] = 1 |
| attn.o_proj.weight[32+i, i] = 1 |
| attn.gate_bias.zero_() |
| attn.q_norm.weight.fill_(16*max_writes*math.sqrt(30*6)/math.sqrt(128)) |
| attn.k_norm.weight.fill_(1) |
| self.cache = DynamicCache(config=config) |
| |
| rows = torch.stack([self._row(a, 0, 0, True) for a in range(16)]).unsqueeze(0) |
| with torch.inference_mode(): self._forward(rows) |
|
|
| def _row(self, address, value, position, valid): |
| x = torch.zeros(128) |
| x[:4] = torch.tensor([2*b-1 for b in bits(address, 4)]) |
| x[4:6] = 1 |
| x[6:10] = x[:4] |
| p = position / self.maximum |
| x[10], x[11], x[12] = p, 1 if valid else -1, math.sqrt(max(0, 1-p*p)) |
| x[13+value] = 1 |
| |
| x[-1] = math.sqrt(128-float(x.square().sum())) |
| return x |
|
|
| def _forward(self, rows): |
| old = self.cache.get_seq_length() |
| length = rows.shape[1] |
| mask = None |
| if length > 1: |
| mask = rows.new_zeros((1, 1, length, old+length)) |
| mask[0, 0, :, old:].masked_fill_(torch.ones(length, length, dtype=torch.bool).triu(1), float('-inf')) |
| positions = torch.arange(old, old+length).unsqueeze(0) |
| return self.layer(rows, position_embeddings=None, attention_mask=mask, |
| attention_scaling=compute_attention_scaling(positions), |
| past_key_values=self.cache, use_cache=True) |
|
|
| 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 |
| with torch.inference_mode(): self._forward(self._row(address, value, self.writes, True).reshape(1, 1, -1)) |
|
|
| def read(self, address): |
| if not 0 <= address < 16: raise ValueError('expected nibble') |
| with torch.inference_mode(): |
| output = self._forward(self._row(address, 0, self.writes, False).reshape(1, 1, -1)) |
| return int(output[0, 0, 32:48].argmax()) |
|
|
