| """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 |
| W0 = X0 + 16 * LANES |
| 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.) |
| 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 |
| 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_() |
| |
| |
| 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') |
| |
| |
| 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) |
|
|
