| """Causal attention CPU over a stream of typed event vectors. |
| |
| This standalone constructed decoder recomputes projected keys/values from its |
| input event stream. Its attention and MLP weights execute the CPU. The host |
| serializes model-emitted states and writes into event vectors, runs I/O and |
| stops at halt. It is not a Beam checkpoint or Hugging Face Transformers model. |
| """ |
| import torch |
| from torch import nn |
| import torch.nn.functional as F |
| from .attention_cpu import AttentionCPU, STATE_BITS, Step, bits, from_bits |
|
|
| |
| FEATURES = 19 |
| VALID, ADDRESS, DATA, POSITION = 9, 10, 14, 18 |
|
|
|
|
| def event(state=None, address=0, value=0, valid=False, position=0): |
| state = [0] * STATE_BITS if state is None else list(state) |
| if len(state) != STATE_BITS or any(x not in (0, 1) for x in state): |
| raise ValueError('state must contain nine binary values') |
| if not 0 <= address < 16 or not 0 <= value < 16 or position < 0: |
| raise ValueError('invalid event') |
| return state + [int(valid)] + bits(address, 4) + bits(value, 4) + [position] |
|
|
|
|
| class SequenceCPU(nn.Module): |
| def __init__(self, program, max_events=1024): |
| super().__init__() |
| if not 16 < max_events <= 8192: raise ValueError('max_events must be 17..8192') |
| self.max_events = max_events |
| core = AttentionCPU(program) |
| self.pc_detect, self.address = core.pc_detect, core.address |
| self.detect, self.result = core.detect, core.result |
| self.k_proj = nn.Linear(FEATURES, 6) |
| self.v_proj = nn.Linear(FEATURES, 4, bias=False) |
| with torch.no_grad(): |
| self.k_proj.weight.zero_(); self.k_proj.bias.zero_() |
| for i in range(4): |
| self.k_proj.weight[i, ADDRESS+i] = 2 |
| self.k_proj.bias[i] = -1 |
| self.k_proj.weight[4, POSITION] = 1 |
| self.k_proj.weight[5, VALID] = 2 |
| self.k_proj.bias[5] = -1 |
| self.v_proj.weight.zero_() |
| for i in range(4): self.v_proj.weight[i, DATA+i] = 1 |
| self.requires_grad_(False) |
|
|
| def forward(self, inputs, last_only=False): |
| """Input shape [batch,time,19]; output shape [batch,time,20]. |
| |
| With last_only, output time dimension is one and all supplied events |
| are in its past. Full mode uses an explicit causal attention mask. |
| """ |
| if inputs.ndim != 3 or inputs.shape[-1] != FEATURES: raise ValueError('invalid input shape') |
| if not 0 < inputs.shape[1] <= self.max_events: raise ValueError('event budget exceeded') |
| selected = inputs[:, -1:, :] if last_only else inputs |
| state = selected[..., :STATE_BITS] |
| address = self.address(torch.relu(self.pc_detect(state[..., :5]))) |
| base = 16 * (self.max_events + 1) |
| query = torch.cat(((2 * address - 1) * base, |
| torch.ones_like(address[..., :1]) * 16, |
| torch.ones_like(address[..., :1]) * base * 5), dim=-1) |
| keys, values = self.k_proj(inputs), self.v_proj(inputs) |
| memory = F.scaled_dot_product_attention(query.unsqueeze(1), keys.unsqueeze(1), values.unsqueeze(1), |
| scale=1., is_causal=not last_only).squeeze(1) |
| read_bits = 2 * (torch.relu(memory - .25) - torch.relu(memory - .75)) |
| result = self.result(torch.relu(self.detect(torch.cat((state, read_bits), dim=-1)))) |
| return torch.cat((result, address, state[..., 5:]), dim=-1) |
|
|
|
|
| class SequenceRuntime: |
| """Serializes event vectors and feeds back predicted machine state.""" |
| def __init__(self, model): |
| self.model = model |
| self.state = torch.zeros(STATE_BITS) |
| self.events = [event(address=i, valid=True) for i in range(16)] |
| self.halted = False |
|
|
| def _append(self, item): |
| if len(self.events) >= self.model.max_events: raise RuntimeError('event budget exhausted') |
| self.events.append(item) |
|
|
| def device_write(self, address, value): |
| if not isinstance(address, int) or not isinstance(value, int) or not 0 <= address < 16 or not 0 <= value < 16: |
| raise ValueError('device address and value must be 0..15') |
| self._append(event(address=address, value=value, valid=True, position=len(self.events))) |
|
|
| def step(self): |
| if self.halted: raise RuntimeError('machine is halted') |
| self._append(event(self.state.tolist(), position=len(self.events))) |
| with torch.inference_mode(): |
| output = self.model(torch.tensor([self.events], dtype=torch.float32), last_only=True)[0, 0] |
| if not torch.isfinite(output).all(): raise RuntimeError('non-finite circuit output') |
| emitted = (output > .5).to(torch.int64).tolist() |
| new_state = emitted[:STATE_BITS] |
| write, out, halt = map(bool, emitted[STATE_BITS:STATE_BITS+3]) |
| address = from_bits(emitted[STATE_BITS+3:STATE_BITS+7]) |
| value = from_bits(emitted[STATE_BITS+7:STATE_BITS+11]) |
| if write: self.device_write(address, value) |
| self.state = torch.tensor(new_state, dtype=torch.float32) |
| self.halted = halt |
| return Step(from_bits(new_state[:5]), from_bits(new_state[5:]), write, out, halt, address, value) |
|
|
