File size: 7,840 Bytes
31f7037 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | """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
|