"""Beam-layer vector unit: weights are data in the KV cache, matvec runs in one decoder layer. Spike toward a QwenPU-style design (a neural net executing inside the language model's own forward passes). One stock ``BeamDecoderLayer`` with constructed weights computes an exact int4 x int4 dot product of length ``LANES``: * Weight rows are *tokens* in the KV cache. Row ``r`` is stored as bipolar bits (offset-binary int4 per lane) keyed by its four address bits. * A query token carries the row address plus the activation vector, one-hot per lane. * Attention retrieves the row (hard softmax, fixed-norm payload). * The SwiGLU MLP multiplies: neuron (lane, k) fires when activation==k and emits k * weight, summed by ``down_proj``. * The layer RMS-norms its MLP output, which would erase magnitude, so the result is read as a ratio against a constant anchor dimension. Nothing is trained. Not pretrained Beam; a constructed single-layer unit. """ import math import torch from transformers import BeamConfig from transformers.cache_utils import DynamicCache from transformers.models.beam.modeling_beam import BeamDecoderLayer, BeamMLP, compute_attention_scaling LANES = 8 H = 256 ROWS = 16 X0 = 16 # activation one-hot: LANES * 16 dims W0 = X0 + 16 * LANES # retrieved weight bits: LANES * 4 dims CONST = W0 + 4 * LANES OUT_D, OUT_A = 200, 201 ANCHOR = 1000.0 TEMP = 60.0 def _bits4(v): return [(v >> b) & 1 for b in range(4)] class VectorUnit: def __init__(self, weights): """weights: up to ROWS rows of LANES ints in -8..7.""" if len(weights) > ROWS or any(len(r) != LANES for r in weights): raise ValueError(f'expected <= {ROWS} rows of {LANES} int4 weights') if any(not -8 <= w <= 7 for r in weights for w in r): raise ValueError('weights must be int4') self.weights = [list(r) for r in weights] n_neurons = 16 * LANES + 1 config = BeamConfig(hidden_size=H, intermediate_size=n_neurons, 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.config = config 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): proj.weight.zero_() attn.gate_bias.fill_(30.) # sigmoid gate ~ 1 for head in range(6): for i in range(4): attn.q_proj.weight[head*128+i, i] = 1 for i in range(4): attn.k_proj.weight[i, i] = 1 for i in range(4*LANES): attn.v_proj.weight[i, W0 + i] = 1 # key rows carry their payload in W0 dims attn.o_proj.weight[W0 + i, i] = 1 attn.q_norm.weight.fill_(48.) attn.k_norm.weight.fill_(1.) self.cache = DynamicCache(config=config) self._calibrate_mlp(n_neurons) self.cache = DynamicCache(config=config) rows = torch.stack([self._row(a, weights=self.weights[a] if a < len(self.weights) else [0]*LANES) for a in range(ROWS)]).unsqueeze(0) with torch.inference_mode(): self._forward(rows) def _pad(self, x): x[CONST] = 1 x[-1] = math.sqrt(H - float(x.square().sum())) return x def _row(self, address, weights=None, activations=None): x = torch.zeros(H) x[:4] = torch.tensor([2*b-1 for b in _bits4(address)], dtype=torch.float32) if weights is not None: for lane, w in enumerate(weights): for b, bit in enumerate(_bits4(w + 8)): x[W0 + 4*lane + b] = 2*bit-1 if activations is not None: for lane, a in enumerate(activations): x[X0 + 16*lane + a] = 1 return self._pad(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 _calibrate_mlp(self, n_neurons): """Measure what the MLP input looks like after attention + norms, then set weights.""" seen = {} hook = self.layer.pre_feedforward_layernorm.register_forward_hook( lambda m, i, o: seen.__setitem__('v', o.detach().clone())) self.cache = DynamicCache(config=self.config) key = self._row(0, weights=[3]*LANES) probe = self._row(0, activations=[5]*LANES) with torch.inference_mode(): self._forward(key.reshape(1, 1, -1)) self._forward(probe.reshape(1, 1, -1)) hook.remove() v = seen['v'][0, 0] x_active = float(v[X0 + 5]) const_n = float(v[CONST]) bit_scale = float(v[W0:W0+4*LANES].abs().mean()) mlp = BeamMLP(self.config, intermediate_size=n_neurons) with torch.no_grad(): for proj in (mlp.gate_proj, mlp.up_proj, mlp.down_proj): proj.weight.zero_() # Neuron (lane, k): gate = TEMP*(x/x_active - 1/2) -> silu ~ TEMP/2 when active, ~0 otherwise. # up = k * (sum 2^(b-1) s_b - 1/2) / (TEMP/2), so gate*up = k * weight. for lane in range(LANES): for k in range(16): n = lane*16 + k mlp.gate_proj.weight[n, X0 + 16*lane + k] = TEMP / x_active mlp.gate_proj.weight[n, CONST] = -TEMP / 2 / const_n for b in range(4): mlp.up_proj.weight[n, W0 + 4*lane + b] = k * 2**(b-1) / bit_scale / (TEMP/2) mlp.up_proj.weight[n, CONST] = -k * 0.5 / const_n / (TEMP/2) mlp.down_proj.weight[OUT_D, n] = 1 a = n_neurons - 1 mlp.gate_proj.weight[a, CONST] = TEMP / 2 / const_n mlp.up_proj.weight[a, CONST] = ANCHOR / (TEMP/2) / const_n mlp.down_proj.weight[OUT_A, a] = 1 self.layer.mlp = mlp.eval().requires_grad_(False) def dot(self, row, activations): """Exact dot(weights[row], activations); activations are 0..15 ints.""" if not 0 <= row < ROWS: raise ValueError('row out of range') if len(activations) != LANES or any(not 0 <= a <= 15 for a in activations): raise ValueError(f'expected {LANES} activations in 0..15') # A query must not persist: its key matches the stored row's address, and # leftover zero-payload queries would dilute later retrievals toward the RMSNorm epsilon. with torch.inference_mode(): out = self._forward(self._row(row, activations=activations).reshape(1, 1, -1)) self.cache.crop(-1) out = out[0, 0] return round(float(out[OUT_D] / out[OUT_A]) * ANCHOR)