| """Experimental CPU transition inside a stock BeamDecoderLayer. |
| |
| Both RAM retrieval and ROM-specific transitions execute in stock Beam decoder |
| attention and SwiGLU MLP weights. The host still selects the ROM operand/address |
| from the current PC and serializes discrete state and write events. This is not |
| an autoregressive LM, a trained checkpoint, or an embedded 501B Beam model. |
| """ |
| import math |
| import torch |
| from transformers.models.beam.modeling_beam import BeamMLP |
| from .attention_cpu import bits, from_bits |
| from .hf_decoder_memory import DecoderMemory |
| from .large_cpu import validate, transition, EventResult |
|
|
|
|
| class DecoderCPU(DecoderMemory): |
| def __init__(self, program, max_writes=128): |
| validate(program) |
| self.program = list(program) |
| super().__init__(max_writes) |
| |
| |
| n = 65536 |
| mlp = BeamMLP(self.layer.self_attn.config, intermediate_size=n) |
| ids = torch.arange(n, dtype=torch.int64) |
| pc = ids & 255 |
| acc = (ids >> 8) & 15 |
| mem = (ids >> 12) & 15 |
| state = (((ids[:, None] >> torch.arange(12)) & 1)*2-1).float() |
| self.layer.mlp = mlp.eval().requires_grad_(False) |
| |
| |
| |
| alpha = 1 / math.sqrt(2) |
| temperature = 32.0 |
| transitions = [] |
| for i in range(n): |
| npc, nacc, write, output, halt = transition(program, int(pc[i]), int(acc[i]), int(mem[i])) |
| transitions.append(bits(npc, 8) + bits(nacc, 4) + [write, output, halt]) |
| transitions = torch.tensor(transitions, dtype=torch.int64) |
| with torch.no_grad(): |
| mlp.gate_proj.weight.zero_(); mlp.up_proj.weight.zero_(); mlp.down_proj.weight.zero_() |
| mlp.gate_proj.weight[:, 48:60] = state * (temperature / alpha) |
| mlp.gate_proj.weight[ids, 32 + mem] = temperature / 8 |
| mlp.gate_proj.weight[:, 5] = -12.5 * temperature / alpha |
| mlp.up_proj.weight[:, 5] = 1 / alpha |
| |
| for i in range(15): |
| mlp.down_proj.weight[60 + 2*i, ids] = (transitions[:, i] == 0).float() |
| mlp.down_proj.weight[61 + 2*i, ids] = (transitions[:, i] == 1).float() |
| self.pc = 0 |
| self.acc = 0 |
| self.halted = False |
|
|
| def _cpu_row(self, address, pc, acc): |
| x = self._row(address, 0, self.writes, False) |
| state = bits(pc, 8) + bits(acc, 4) |
| x[48:60] = torch.tensor([2*b-1 for b in state], dtype=torch.float32) |
| x[-1] = 0 |
| x[-1] = math.sqrt(128-float(x.square().sum())) |
| return x |
|
|
| def device_write(self, address, value): |
| self.write(address, value) |
|
|
| def step(self): |
| if self.halted: raise RuntimeError('machine is halted') |
| |
| |
| op, operand = self.program[self.pc] if self.pc < len(self.program) else ('HALT', 0) |
| address = operand if op in {'LOAD', 'STORE', 'ADDM', 'SUBM'} else 0 |
| old_acc = self.acc |
| with torch.inference_mode(): |
| output = self._forward(self._cpu_row(address, self.pc, self.acc).reshape(1, 1, -1))[0, 0] |
| pairs = output[60:90].reshape(15, 2) |
| if not torch.isfinite(pairs).all(): raise RuntimeError('non-finite output') |
| decoded = pairs.argmax(dim=-1).tolist() |
| self.pc, self.acc = from_bits(decoded[:8]), from_bits(decoded[8:12]) |
| write, out, self.halted = map(bool, decoded[12:]) |
| if write: self.write(address, old_acc) |
| return EventResult(self.pc, self.acc, write, out, self.halted, address, old_acc) |
|
|
