| """ctypes boundary; no subprocess or compilation in the request path.""" |
| import ctypes |
| import math |
| import os |
| from pathlib import Path |
|
|
|
|
| class BendReducer: |
| def __init__(self, library=None): |
| path = Path(library or os.environ.get('JULIA_ROUTER_LIBRARY') or |
| Path(__file__).parent / 'build/libjulia_router.so') |
| try: |
| self._lib = ctypes.CDLL(str(path.resolve())) |
| except OSError as error: |
| raise RuntimeError( |
| f'Cannot load Bend runtime {path.resolve()}. From Julia-1 run ' |
| 'python -m julia.router.build --native-cpu; or explicitly select ' |
| 'transformer_backend="torch".') from error |
| self._argmax = self._lib.julia_router_argmax |
| self._argmax.argtypes = [ctypes.POINTER(ctypes.c_float), ctypes.c_uint32, |
| ctypes.POINTER(ctypes.c_uint32)] |
| self._argmax.restype = ctypes.c_int |
| self._ptr = ctypes.POINTER(ctypes.c_float) |
| self._softmax = self._lib.julia_router_softmax |
| self._softmax.argtypes = [self._ptr, ctypes.c_uint32, self._ptr, ctypes.POINTER(ctypes.c_uint32)] |
| self._softmax.restype = ctypes.c_int |
| self._layernorm = self._lib.julia_router_layernorm |
| self._layernorm.argtypes = [self._ptr, self._ptr, self._ptr, ctypes.c_uint32, |
| ctypes.c_uint32, ctypes.c_float, self._ptr] |
| self._layernorm.restype = ctypes.c_int |
|
|
| def profile(self, *, enabled=False, reset=False): |
| """Process-wide packed projection timings; opt in only for diagnostics.""" |
| fn = self._lib.julia_router_profile |
| fn.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_double)] |
| fn.restype = None |
| values = (ctypes.c_double * 4)() |
| fn(enabled, reset, values) |
| return dict(zip(('pack_ms', 'execute_ms', 'unpack_ms', 'calls'), values)) |
|
|
| def argmax(self, scores): |
| scores = tuple(scores) |
| if not 1 <= len(scores) <= 1048576: |
| raise ValueError('Expected 1–1048576 scores') |
| if not all(math.isfinite(x) for x in scores): |
| raise ValueError('Scores must be finite') |
| values = (ctypes.c_float * len(scores))(*scores) |
| result = ctypes.c_uint32() |
| code = self._argmax(values, len(scores), ctypes.byref(result)) |
| if code: |
| raise ValueError(f'Bend rejected scores (status {code})') |
| return result.value |
|
|
| def softmax(self, scores): |
| import numpy as np |
| values = np.ascontiguousarray(scores, dtype=np.float32) |
| if values.ndim != 1 or not 1 <= values.size <= 1048576: |
| raise ValueError('Expected a nonempty vector of at most 1048576 scores') |
| output = np.empty_like(values) |
| index = ctypes.c_uint32() |
| code = self._softmax(values.ctypes.data_as(self._ptr), values.size, |
| output.ctypes.data_as(self._ptr), ctypes.byref(index)) |
| if code: |
| raise ValueError(f'Bend rejected scores (status {code})') |
| return index.value, output |
|
|
| def layernorm(self, values, gamma=None, beta=None, epsilon=1e-5): |
| import numpy as np |
| x = np.ascontiguousarray(values, dtype=np.float32) |
| if x.ndim < 1 or x.size == 0: |
| raise ValueError('Expected a nonempty feature array') |
| width = x.shape[-1] |
| g = np.ones(width, np.float32) if gamma is None else np.ascontiguousarray(gamma, dtype=np.float32) |
| b = np.zeros(width, np.float32) if beta is None else np.ascontiguousarray(beta, dtype=np.float32) |
| if g.shape != (width,) or b.shape != (width,): |
| raise ValueError('gamma and beta must match the last dimension') |
| if x.size > 1048576 or width > 65536: |
| raise ValueError('LayerNorm input exceeds native capacity') |
| output = np.empty_like(x) |
| code = self._layernorm(x.ctypes.data_as(self._ptr), g.ctypes.data_as(self._ptr), |
| b.ctypes.data_as(self._ptr), x.size // width, width, epsilon, |
| output.ctypes.data_as(self._ptr)) |
| if code: |
| raise ValueError(f'Bend rejected LayerNorm inputs (status {code})') |
| return output |
|
|
|
|
| class BendMatrix: |
| """Resident immutable FP32 projection weights owned by the Bend heap.""" |
| def __init__(self, weights, reducer=None): |
| import numpy as np |
| import threading |
| import weakref |
| self.reducer = reducer or BendReducer() |
| self._lock = threading.RLock() |
| values = np.ascontiguousarray(weights, dtype=np.float32) |
| if values.ndim != 2 or not all(values.shape): |
| raise ValueError('Projection weights must be a nonempty matrix') |
| self.outputs, self.inputs = values.shape |
| lib = self.reducer._lib |
| create = lib.julia_router_matrix_create |
| create.argtypes = [self.reducer._ptr, ctypes.c_uint32, ctypes.c_uint32] |
| create.restype = ctypes.c_void_p |
| release = lib.julia_router_matrix_free |
| release.argtypes = [ctypes.c_void_p] |
| release.restype = None |
| self._linear = lib.julia_router_linear |
| self._linear.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_uint32, ctypes.c_void_p] |
| self._linear.restype = ctypes.c_int |
| self._linear_inference = lib.julia_router_linear_inference |
| self._linear_inference.argtypes = self._linear.argtypes |
| self._linear_inference.restype = ctypes.c_int |
| self._handle = create(values.ctypes.data_as(self.reducer._ptr), self.inputs, self.outputs) |
| if not self._handle: |
| raise ValueError('Bend rejected projection dimensions or nonfinite weights') |
| self._finalize = weakref.finalize(self, release, self._handle) |
|
|
| def close(self): |
| with self._lock: |
| self._finalize() |
|
|
| def tensor(self, values, *, validate=True): |
| """Fill a contiguous Torch output directly, without NumPy/list/concat wrappers.""" |
| import torch |
| if values.device.type != 'cpu' or values.dtype != torch.float32 or values.shape[-1] != self.inputs: |
| raise ValueError('Bend tensor projection requires matching CPU FP32 inputs') |
| x = values.contiguous() |
| rows = x.numel() // self.inputs |
| if not rows: |
| raise ValueError('Projection input must be nonempty') |
| out = torch.empty((*x.shape[:-1], self.outputs), dtype=x.dtype, device=x.device) |
| linear = self._linear if validate else self._linear_inference |
| limit = min(1048576 // self.inputs, 1048576 // self.outputs) |
| with self._lock: |
| if not self._finalize.alive: |
| raise RuntimeError('Bend matrix has been closed') |
| for start in range(0, rows, limit): |
| code = linear(self._handle, |
| x.data_ptr() + start * self.inputs * 4, |
| min(limit, rows - start), |
| out.data_ptr() + start * self.outputs * 4) |
| if code: |
| raise ValueError(f'Bend projection failed (status {code})') |
| return out |
|
|
| def __call__(self, values): |
| import numpy as np |
| x = np.ascontiguousarray(values, dtype=np.float32) |
| if x.ndim < 1 or x.shape[-1] != self.inputs or not x.size: |
| raise ValueError('Projection input must match resident weights') |
| out = np.empty((*x.shape[:-1], self.outputs), dtype=np.float32) |
| with self._lock: |
| if not self._finalize.alive: |
| raise RuntimeError('Bend matrix has been closed') |
| code = self._linear(self._handle, x.ctypes.data_as(self.reducer._ptr), |
| x.size // self.inputs, out.ctypes.data_as(self.reducer._ptr)) |
| if code: |
| raise ValueError(f'Bend projection failed (status {code})') |
| return out |
|
|
