| """Resident Bend encoder projections and optional decision-head normalization. |
| |
| The packed FP32 kernels own matrix arithmetic and coarse parallel scheduling. |
| PyTorch supplies embedding lookup, SDPA, activations and the selected-output head. |
| """ |
| import numpy as np |
| import torch |
| from torch import nn |
|
|
|
|
| class BendLayerNorm(nn.Module): |
| def __init__(self, original, reducer): |
| super().__init__() |
| self.weight, self.bias = original.weight, original.bias |
| self.eps = original.eps |
| self.normalized_shape = original.normalized_shape |
| self.reducer = reducer |
|
|
| def forward(self, x): |
| if x.device.type != 'cpu' or x.dtype != torch.float32: |
| raise ValueError('Bend LayerNorm requires CPU float32; use torch backend for CUDA') |
| if torch.is_grad_enabled(): |
| raise RuntimeError('Bend transformer kernels require torch.inference_mode()') |
| width = x.shape[-1] |
| values = x.detach().contiguous().numpy().reshape(-1, width) |
| gamma = self.weight.detach().numpy() if self.weight is not None else None |
| beta = self.bias.detach().numpy() if self.bias is not None else None |
| parts = [self.reducer.layernorm(values[start:start + 1048576 // width], gamma, beta, self.eps) |
| for start in range(0, len(values), 1048576 // width)] |
| result = parts[0] if len(parts) == 1 else np.concatenate(parts) |
| return torch.from_numpy(result.reshape(x.shape)) |
|
|
|
|
| class BendHeadLayer(nn.Module): |
| """Explicit pre-norm block; prevents fused MHA from bypassing Bend norms.""" |
| def __init__(self, layer, reducer): |
| super().__init__() |
| if not layer.norm_first: |
| raise ValueError('Julia expects pre-norm transformer layers') |
| self.layer = layer |
| self.norm1 = BendLayerNorm(layer.norm1, reducer) |
| self.norm2 = BendLayerNorm(layer.norm2, reducer) |
|
|
| def forward(self, x, src_key_padding_mask=None): |
| normalized = self.norm1(x) |
| attention = self.layer.self_attn(normalized, normalized, normalized, |
| key_padding_mask=src_key_padding_mask, need_weights=False)[0] |
| x = x + self.layer.dropout1(attention) |
| hidden = self.layer.linear2(self.layer.dropout( |
| self.layer.activation(self.layer.linear1(self.norm2(x))))) |
| return x + self.layer.dropout2(hidden) |
|
|
|
|
| def install_bend_head(model, reducer): |
| if model.head is not None: |
| model.head.layers = nn.ModuleList([BendHeadLayer(layer, reducer) for layer in model.head.layers]) |
| model.scorer[0] = BendLayerNorm(model.scorer[0], reducer) |
| model.eval() |
|
|
|
|
| class BendLinear(nn.Module): |
| """Inference-only CPU projection using resident weights and Bend arithmetic.""" |
| def __init__(self, original, reducer): |
| super().__init__() |
| from .native import BendMatrix |
| self.weight, self.bias = original.weight, original.bias |
| self.in_features, self.out_features = original.in_features, original.out_features |
| self.matrix = BendMatrix(original.weight.detach().numpy(), reducer) |
| self._weight_version = self.weight._version |
| self._weight_pointer = self.weight.data_ptr() |
|
|
| def _check_weight(self): |
| if self.weight._version != self._weight_version or self.weight.data_ptr() != self._weight_pointer: |
| raise RuntimeError('Bend resident weights changed; recreate the inference backend') |
|
|
| def forward(self, x): |
| if x.device.type != 'cpu' or x.dtype != torch.float32 or torch.is_grad_enabled(): |
| raise ValueError('Bend dense projections require CPU FP32 inference_mode') |
| self._check_weight() |
| result = self.matrix.tensor(x, validate=False) |
| return result if self.bias is None else result + self.bias |
|
|
|
|
| def install_bend_encoder(model, reducer): |
| """Route all encoder linear projections through Bend; attention remains SDPA.""" |
| count = 0 |
| def install(module): |
| nonlocal count |
| for name, child in list(module.named_children()): |
| if type(child) is nn.Linear: |
| setattr(module, name, BendLinear(child, reducer)) |
| count += 1 |
| else: |
| install(child) |
| install(model.encoder) |
| model.eval() |
| return count |
|
|
