File size: 5,263 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | """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
# Input event vector: PC[5], ACC[4], write-valid, address[4], data[4], position.
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)
|