| |
| """Hand-built ONNX graphs for Audio8-ASR-Infinite (see audio8_infinite_ref.py for the IO contract). |
| |
| The graphs are emitted node-by-node straight from the (memmapped, bf16) safetensors, one weight at |
| a time, with every initializer streamed into the external-data file as it is produced. Nothing ever |
| holds the whole fp32 model in RAM, which is what makes a 4 B-parameter export possible on a box |
| with single-digit free GB -- `torch.onnx.export` would need the full fp32 module plus its trace. |
| |
| Precisions (`--precision`): |
| fp32 MatMul fp32 weights (parity reference; ~16 GB, not shipped) |
| fp16 whole graph fp16 (RMSNorm / softmax / mel in fp32), KV caches fp16, I/O otherwise fp32 |
| int8 weight-only MatMulNBits (8-bit, asymmetric, block 32, accuracy_level 4), fp32 activations |
| int8dq (comparison only, not shipped) per-channel int8 weights + DynamicQuantizeLinear / |
| MatMulInteger, the quantize_dynamic pattern: ONE activation scale per tensor, so the |
| result depends on how many frames share a call and Qwen2.5's activation outliers flatten |
| the other channels -- measurably worse transcripts than either NBits tier |
| int4 weight-only MatMulNBits (4-bit, asymmetric, block 32, accuracy_level 4), fp32 activations |
| """ |
| from __future__ import annotations |
|
|
| import math |
| import os |
| from pathlib import Path |
|
|
| import numpy as np |
| import onnx |
| from onnx import TensorProto, helper |
|
|
| import audio8_infinite_ref as a8i |
|
|
| OPSET = 17 |
| F32, F16 = TensorProto.FLOAT, TensorProto.FLOAT16 |
| NP = {F32: np.float32, F16: np.float16} |
|
|
|
|
| class Graph: |
| def __init__(self, path: Path, act: int): |
| self.path = Path(path) |
| self.data_name = self.path.name + ".data" |
| self.data = open(self.path.with_name(self.data_name), "wb") |
| self.off = 0 |
| self.act = act |
| self.nodes, self.inits, self.inputs, self.outputs = [], [], [], [] |
| self.n = 0 |
| self._consts: dict = {} |
|
|
| |
| def uid(self, p="t"): |
| self.n += 1 |
| return f"{p}_{self.n}" |
|
|
| def init(self, name: str, arr: np.ndarray, *, onnx_type=None) -> str: |
| arr = np.ascontiguousarray(arr) |
| if arr.nbytes < 4096: |
| t = onnx.numpy_helper.from_array(arr, name) |
| self.inits.append(t) |
| return name |
| pad = (-self.off) % 4096 |
| if pad: |
| self.data.write(b"\0" * pad) |
| self.off += pad |
| buf = arr.tobytes() |
| self.data.write(buf) |
| t = TensorProto() |
| t.name = name |
| t.data_type = onnx_type if onnx_type is not None else helper.np_dtype_to_tensor_dtype(arr.dtype) |
| t.dims.extend(arr.shape) |
| t.data_location = TensorProto.EXTERNAL |
| for k, v in (("location", self.data_name), ("offset", str(self.off)), ("length", str(len(buf)))): |
| e = t.external_data.add() |
| e.key, e.value = k, v |
| self.off += len(buf) |
| self.inits.append(t) |
| return name |
|
|
| def const(self, value, dtype=np.int64) -> str: |
| arr = np.asarray(value, dtype=dtype) |
| key = (arr.dtype.str, arr.shape, arr.tobytes()) |
| if key not in self._consts: |
| self._consts[key] = self.init(self.uid("c"), arr) |
| return self._consts[key] |
|
|
| def fconst(self, value) -> str: |
| return self.const(value, NP[self.act]) |
|
|
| def op(self, op_type, inputs, n_out=1, domain=None, **attrs): |
| outs = [self.uid(op_type.lower()) for _ in range(n_out)] |
| kw = {"domain": domain} if domain else {} |
| self.nodes.append(helper.make_node(op_type, list(inputs), outs, name=self.uid("n_" + op_type), **kw, **attrs)) |
| return outs[0] if n_out == 1 else outs |
|
|
| def inp(self, name, dtype, shape): |
| self.inputs.append(helper.make_tensor_value_info(name, dtype, shape)) |
| return name |
|
|
| def out(self, src, name, dtype, shape): |
| self.nodes.append(helper.make_node("Identity", [src], [name], name=self.uid("n_out"))) |
| self.outputs.append(helper.make_tensor_value_info(name, dtype, shape)) |
|
|
| def finish(self, meta: dict): |
| self.data.close() |
| g = helper.make_graph(self.nodes, self.path.stem, self.inputs, self.outputs, self.inits) |
| opsets = [helper.make_opsetid("", OPSET), helper.make_opsetid("com.microsoft", 1)] |
| m = helper.make_model(g, opset_imports=opsets, producer_name="winstt-audio8-infinite-export") |
| m.ir_version = 8 |
| for k, v in meta.items(): |
| e = m.metadata_props.add() |
| e.key, e.value = k, str(v) |
| onnx.save(m, str(self.path)) |
|
|
| |
| def cast(self, x, to): |
| return self.op("Cast", [x], to=to) |
|
|
| def rms(self, x, w: np.ndarray, eps: float, name: str): |
| """x / sqrt(mean(x^2) + eps) * w, computed in fp32 (SimplifiedLayerNorm fusion pattern).""" |
| xf = self.cast(x, F32) if self.act != F32 else x |
| sq = self.op("Pow", [xf, self.const(2.0, np.float32)]) |
| mean = self.op("ReduceMean", [sq], axes=[-1], keepdims=1) |
| den = self.op("Sqrt", [self.op("Add", [mean, self.const(eps, np.float32)])]) |
| y = self.op("Div", [xf, den]) |
| y = self.op("Mul", [y, self.init(name, w.astype(np.float32))]) |
| return self.cast(y, self.act) if self.act != F32 else y |
|
|
| def gelu(self, x): |
| xf = x |
| e = self.op("Erf", [self.op("Div", [xf, self.fconst(math.sqrt(2.0))])]) |
| return self.op("Mul", [self.op("Mul", [xf, self.op("Add", [e, self.fconst(1.0)])]), self.fconst(0.5)]) |
|
|
| def silu(self, x): |
| return self.op("Mul", [x, self.op("Sigmoid", [x])]) |
|
|
| def rope(self, x, cos, sin): |
| a, b = self.op("Split", [x], n_out=2, axis=-1) |
| rot = self.op("Concat", [self.op("Neg", [b]), a], axis=-1) |
| return self.op("Add", [self.op("Mul", [x, cos]), self.op("Mul", [rot, sin])]) |
|
|
| def shape_dim(self, x, i): |
| return self.op("Gather", [self.op("Shape", [x]), self.const([i])], axis=0) |
|
|
|
|
| class Linear: |
| """Emits y = x @ W.T (+ b) in the requested precision.""" |
|
|
| def __init__(self, g: Graph, precision: str, block: int = 32, accuracy_level: int = 4): |
| self.g, self.prec, self.block, self.acc = g, precision, block, accuracy_level |
| self.count = 0 |
|
|
| def __call__(self, x, W: np.ndarray, b: np.ndarray | None, name: str, quantize: bool = True): |
| g = self.g |
| N, K = W.shape |
| prec = self.prec if quantize else ("fp16" if self.prec == "fp16" else "fp32") |
| self.count += 1 |
| if prec in ("fp32", "fp16"): |
| dt = np.float16 if prec == "fp16" else np.float32 |
| y = g.op("MatMul", [x, g.init(name + ".weight_t", W.T.astype(dt))]) |
| elif prec == "int8dq": |
| |
| |
| Wt = W.T.astype(np.float32) |
| scale = np.maximum(np.abs(Wt).max(0), 1e-12) / 127.0 |
| q = np.clip(np.rint(Wt / scale), -127, 127).astype(np.int8) |
| xq, xs, xz = g.op("DynamicQuantizeLinear", [x], n_out=3) |
| mm = g.op("MatMulInteger", [xq, g.init(name + ".weight_q8", q), xz, |
| g.init(name + ".weight_zp8", np.zeros(N, np.int8))]) |
| s = g.op("Mul", [xs, g.init(name + ".weight_scale", scale.astype(np.float32))]) |
| y = g.op("Mul", [g.op("Cast", [mm], to=F32), s]) |
| elif prec == "int4": |
| from onnxruntime.capi._pybind_state import quantize_matmul_4bits |
|
|
| Wt = np.ascontiguousarray(W.T.astype(np.float32)) |
| kb = (K + self.block - 1) // self.block |
| assert K % self.block == 0, (name, K) |
| packed = np.zeros((N, kb, self.block // 2), np.uint8) |
| scales = np.zeros((N, kb), np.float32) |
| zp = np.zeros((N, (kb + 1) // 2), np.uint8) |
| quantize_matmul_4bits(packed, Wt, scales, zp, self.block, N, K, False) |
| y = g.op("MatMulNBits", [x, g.init(name + ".weight_q4", packed), g.init(name + ".weight_scales", scales), |
| g.init(name + ".weight_zp4", zp)], domain="com.microsoft", |
| K=K, N=N, bits=4, block_size=self.block, accuracy_level=self.acc) |
| elif prec == "int8": |
| |
| |
| |
| |
| from onnxruntime.capi._pybind_state import quantize_matmul_8bits |
|
|
| Wt = np.ascontiguousarray(W.T.astype(np.float32)) |
| kb = (K + self.block - 1) // self.block |
| assert K % self.block == 0, (name, K) |
| packed = np.zeros((N, kb, self.block), np.uint8) |
| scales = np.zeros((N, kb), np.float32) |
| zp = np.zeros((N, kb), np.uint8) |
| quantize_matmul_8bits(packed, Wt, scales, zp, self.block, N, K, False) |
| y = g.op("MatMulNBits", [x, g.init(name + ".weight_q8", packed), g.init(name + ".weight_scales", scales), |
| g.init(name + ".weight_zp8", zp)], domain="com.microsoft", |
| K=K, N=N, bits=8, block_size=self.block, accuracy_level=self.acc) |
| else: |
| raise ValueError(prec) |
| if b is not None: |
| y = g.op("Add", [y, g.init(name + ".bias", b.astype(NP[g.act]))]) |
| return y |
|
|
|
|
| def attention(g: Graph, q, k_new, v_new, past_k, past_v, bias_f32, *, kv_heads: int, head_dim: int, groups: int): |
| """Split-score attention: never concatenates the (large) past cache with the new keys. |
| |
| q [1,H,T,D]; k_new/v_new [1,KV,T,D]; past_k/past_v [1,KV,C,D]; bias [T, C+T] (fp32). |
| Returns [1,KV,G*T,D] (== [1,H,T,D] after a reshape when G > 1). |
| """ |
| q = g.op("Mul", [q, g.fconst(head_dim ** -0.5)]) |
| |
| qg = g.op("Reshape", [q, g.const([1, kv_heads, -1, head_dim])]) if groups > 1 else q |
| sp = g.op("MatMul", [qg, g.op("Transpose", [past_k], perm=[0, 1, 3, 2])]) |
| sn = g.op("MatMul", [qg, g.op("Transpose", [k_new], perm=[0, 1, 3, 2])]) |
| s = g.op("Concat", [sp, sn], axis=-1) |
| if g.act != F32: |
| s = g.cast(s, F32) |
| if groups > 1: |
| |
| s5 = g.op("Reshape", [s, g.op("Concat", [g.const([1, kv_heads, groups]), g.op("Shape", [bias_f32])], axis=0)]) |
| p = g.op("Softmax", [g.op("Add", [s5, bias_f32])], axis=-1) |
| p = g.op("Reshape", [p, g.op("Shape", [s])]) |
| else: |
| p = g.op("Softmax", [g.op("Add", [s, bias_f32])], axis=-1) |
| if g.act != F32: |
| p = g.cast(p, g.act) |
| C = g.shape_dim(past_k, 2) |
| pp = g.op("Slice", [p, g.const([0]), C, g.const([-1])]) |
| pn = g.op("Slice", [p, C, g.const([1 << 40]), g.const([-1])]) |
| return g.op("Add", [g.op("MatMul", [pp, past_v]), g.op("MatMul", [pn, v_new])]) |
|
|
|
|
| def heads(g: Graph, x, n_heads: int, head_dim: int): |
| """[1,T,H*D] -> [1,H,T,D]""" |
| r = g.op("Reshape", [x, g.const([1, -1, n_heads, head_dim])]) |
| return g.op("Transpose", [r], perm=[0, 2, 1, 3]) |
|
|
|
|
| def merge(g: Graph, x, width: int): |
| """[1,H,T,D] -> [1,T,H*D]""" |
| t = g.op("Transpose", [x], perm=[0, 2, 1, 3]) |
| return g.op("Reshape", [t, g.const([1, -1, width])]) |
|
|
|
|
| |
| def build_encoder(W, dims: a8i.Dims, precision: str, path: Path, block: int = 32): |
| act = F16 if precision == "fp16" else F32 |
| kv = act |
| g = Graph(path, act) |
| lin = Linear(g, precision, block) |
| d = dims |
| audio = g.inp("audio", F32, ["batch", "samples"]) |
| frame_len = g.inp("frame_len", TensorProto.INT64, [1]) |
| cos = g.inp("cos", F32, ["frames", d.a_head_dim]) |
| sin = g.inp("sin", F32, ["frames", d.a_head_dim]) |
| bias = g.inp("attn_bias", F32, ["frames", "past_plus_frames"]) |
| past = [] |
| for i in range(d.a_layers): |
| pk = g.inp(f"past_key_{i}", kv, [1, d.a_heads, "past", d.a_head_dim]) |
| pv = g.inp(f"past_value_{i}", kv, [1, d.a_heads, "past", d.a_head_dim]) |
| past.append((pk, pv)) |
| if act != F32: |
| cos_a, sin_a = g.cast(cos, act), g.cast(sin, act) |
| else: |
| cos_a, sin_a = cos, sin |
|
|
| |
| a3 = g.op("Unsqueeze", [audio, g.const([1])]) |
| padded = g.op("Pad", [a3, g.const([0, 0, N_FFT_HALF, 0, 0, N_FFT_HALF])], mode="reflect") |
| spec = g.op("Conv", [padded, g.init("frontend.dft_basis", a8i.dft_basis())], strides=[a8i.HOP]) |
| F_ = g.shape_dim(spec, 2) |
| spec = g.op("Slice", [spec, g.const([0]), g.op("Sub", [F_, g.const([1])]), g.const([2])]) |
| sq = g.op("Mul", [spec, spec]) |
| re, im = g.op("Split", [sq], n_out=2, axis=1) |
| power = g.op("Add", [re, im]) |
| mel = g.op("MatMul", [g.init("frontend.mel_filters_t", a8i.mel_filters().T.copy()), power]) |
| lg = g.op("Div", [g.op("Log", [g.op("Max", [mel, g.const(1e-10, np.float32)])]), g.const(math.log(10.0), np.float32)]) |
| lg = g.op("Max", [lg, g.const(a8i.LOG_MEL_FLOOR, np.float32)]) |
| feats = g.op("Div", [g.op("Add", [lg, g.const(4.0, np.float32)]), g.const(4.0, np.float32)]) |
| if act != F32: |
| feats = g.cast(feats, act) |
| dt = NP[act] |
| c1 = g.op("Conv", [feats, g.init("embedder.conv1.weight", W["audio_tower.embedder.conv1.weight"].astype(dt)), |
| g.init("embedder.conv1.bias", W["audio_tower.embedder.conv1.bias"].astype(dt))], pads=[2, 0]) |
| c1 = g.gelu(c1) |
| c2 = g.op("Conv", [c1, g.init("embedder.conv2.weight", W["audio_tower.embedder.conv2.weight"].astype(dt)), |
| g.init("embedder.conv2.bias", W["audio_tower.embedder.conv2.bias"].astype(dt))], |
| pads=[1, 0], strides=[2]) |
| c2 = g.gelu(c2) |
| x = g.op("Transpose", [c2], perm=[0, 2, 1]) |
| B = g.shape_dim(audio, 0) |
| T = g.shape_dim(cos, 0) |
| Kw = g.op("Div", [T, B]) |
| x = g.op("Slice", [x, g.op("Neg", [Kw]), g.const([1 << 40]), g.const([1])]) |
| x = g.op("Reshape", [x, g.const([1, -1, d.a_hidden])]) |
|
|
| width = d.a_heads * d.a_head_dim |
| for i in range(d.a_layers): |
| p = f"audio_tower.layers.{i}." |
| h = g.rms(x, W[p + "self_attn_layer_norm.weight"], d.a_eps, f"enc.{i}.ln1") |
| q = heads(g, lin(h, W[p + "self_attn.q_proj.weight"], W[p + "self_attn.q_proj.bias"], f"enc.{i}.q"), d.a_heads, d.a_head_dim) |
| k = heads(g, lin(h, W[p + "self_attn.k_proj.weight"], None, f"enc.{i}.k"), d.a_heads, d.a_head_dim) |
| v = heads(g, lin(h, W[p + "self_attn.v_proj.weight"], W[p + "self_attn.v_proj.bias"], f"enc.{i}.v"), d.a_heads, d.a_head_dim) |
| q, k = g.rope(q, cos_a, sin_a), g.rope(k, cos_a, sin_a) |
| g.out(k, f"key_{i}", kv, [1, d.a_heads, "frames", d.a_head_dim]) |
| g.out(v, f"value_{i}", kv, [1, d.a_heads, "frames", d.a_head_dim]) |
| o = attention(g, q, k, v, past[i][0], past[i][1], bias, kv_heads=d.a_heads, head_dim=d.a_head_dim, groups=1) |
| o = merge(g, o, width) |
| x = g.op("Add", [x, lin(o, W[p + "self_attn.o_proj.weight"], W[p + "self_attn.o_proj.bias"], f"enc.{i}.o")]) |
| h = g.rms(x, W[p + "final_layer_norm.weight"], d.a_eps, f"enc.{i}.ln2") |
| gt = lin(h, W[p + "mlp.gate_proj.weight"], None, f"enc.{i}.gate") |
| up = lin(h, W[p + "mlp.up_proj.weight"], None, f"enc.{i}.up") |
| x = g.op("Add", [x, lin(g.op("Mul", [g.silu(gt), up]), W[p + "mlp.down_proj.weight"], W[p + "mlp.down_proj.bias"], f"enc.{i}.down")]) |
| x = g.rms(x, W["audio_tower.norm.weight"], d.a_eps, "enc.norm") |
| |
| grp = g.op("Reshape", [x, g.op("Concat", [g.const([-1]), g.op("Mul", [frame_len, g.const([d.a_hidden])])], axis=0)]) |
| padw = g.op("Sub", [g.const([d.projection]), g.op("Mul", [frame_len, g.const([d.a_hidden])])]) |
| grp = g.op("Pad", [grp, g.op("Concat", [g.const([0, 0, 0]), padw], axis=0)]) |
| e = lin(grp, W["multi_modal_projector.linear_1.weight"], None, "proj.1") |
| e = lin(g.gelu(e), W["multi_modal_projector.linear_2.weight"], None, "proj.2") |
| e = g.op("Unsqueeze", [e, g.const([0])]) |
| if act != F32: |
| e = g.cast(e, F32) |
| g.out(e, "audio_embeds", F32, [1, "tokens", d.t_hidden]) |
| g.finish({"winstt_audio8_infinite": "audio_encoder", "precision": precision, |
| "sliding_window": d.a_window, "rope_theta": d.a_theta, "head_dim": d.a_head_dim}) |
| return lin.count |
|
|
|
|
| N_FFT_HALF = a8i.N_FFT // 2 |
|
|
|
|
| def build_decoder(W, dims: a8i.Dims, precision: str, path: Path, block: int = 32, layers=None, head=True): |
| """layers=(start, end) builds a shard (for fp32 parity on low-RAM boxes): a shard that does not |
| start at 0 takes `hidden` instead of `inputs_embeds`; one without `head` outputs `hidden`.""" |
| act = F16 if precision == "fp16" else F32 |
| kv = act |
| g = Graph(path, act) |
| lin = Linear(g, precision, block) |
| d = dims |
| l0, l1 = layers or (0, d.t_layers) |
| x_in = g.inp("inputs_embeds" if l0 == 0 else "hidden", F32, [1, "n", d.t_hidden]) |
| ada = g.inp("ada_scale", F32, [d.t_layers, d.t_hidden]) |
| cos = g.inp("cos", F32, ["n", d.t_head_dim]) |
| sin = g.inp("sin", F32, ["n", d.t_head_dim]) |
| bias = g.inp("attn_bias", F32, ["n", "past_plus_n"]) |
| past = {} |
| for i in range(l0, l1): |
| past[i] = (g.inp(f"past_key_{i}", kv, [1, d.t_kv, "past", d.t_head_dim]), |
| g.inp(f"past_value_{i}", kv, [1, d.t_kv, "past", d.t_head_dim])) |
| x = g.cast(x_in, act) if act != F32 else x_in |
| cos_a, sin_a = (g.cast(cos, act), g.cast(sin, act)) if act != F32 else (cos, sin) |
| ada_a = g.cast(ada, act) if act != F32 else ada |
| groups = d.t_heads // d.t_kv |
| width = d.t_heads * d.t_head_dim |
| for i in range(l0, l1): |
| p = f"language_model.model.layers.{i}." |
| h = g.rms(x, W[p + "input_layernorm.weight"], d.t_eps, f"dec.{i}.ln1") |
| q = heads(g, lin(h, W[p + "self_attn.q_proj.weight"], W[p + "self_attn.q_proj.bias"], f"dec.{i}.q"), d.t_heads, d.t_head_dim) |
| k = heads(g, lin(h, W[p + "self_attn.k_proj.weight"], W[p + "self_attn.k_proj.bias"], f"dec.{i}.k"), d.t_kv, d.t_head_dim) |
| v = heads(g, lin(h, W[p + "self_attn.v_proj.weight"], W[p + "self_attn.v_proj.bias"], f"dec.{i}.v"), d.t_kv, d.t_head_dim) |
| q, k = g.rope(q, cos_a, sin_a), g.rope(k, cos_a, sin_a) |
| g.out(k, f"key_{i}", kv, [1, d.t_kv, "n", d.t_head_dim]) |
| g.out(v, f"value_{i}", kv, [1, d.t_kv, "n", d.t_head_dim]) |
| o = attention(g, q, k, v, past[i][0], past[i][1], bias, kv_heads=d.t_kv, head_dim=d.t_head_dim, groups=groups) |
| if groups > 1: |
| o = g.op("Reshape", [o, g.const([1, d.t_heads, -1, d.t_head_dim])]) |
| o = merge(g, o, width) |
| x = g.op("Add", [x, lin(o, W[p + "self_attn.o_proj.weight"], None, f"dec.{i}.o")]) |
| h = g.rms(x, W[p + "post_attention_layernorm.weight"], d.t_eps, f"dec.{i}.ln2") |
| row = g.op("Gather", [ada_a, g.const(i)], axis=0) |
| h = g.op("Mul", [h, row]) |
| gt = lin(h, W[p + "mlp.gate_proj.weight"], None, f"dec.{i}.gate") |
| up = lin(h, W[p + "mlp.up_proj.weight"], None, f"dec.{i}.up") |
| x = g.op("Add", [x, lin(g.op("Mul", [g.silu(gt), up]), W[p + "mlp.down_proj.weight"], None, f"dec.{i}.down")]) |
| if head: |
| x = g.rms(x, W["language_model.model.norm.weight"], d.t_eps, "dec.norm") |
| n = g.shape_dim(x, 1) |
| last = g.op("Slice", [x, g.op("Sub", [n, g.const([1])]), n, g.const([1])]) |
| last = g.op("Reshape", [last, g.const([1, d.t_hidden])]) |
| emb = W["language_model.model.embed_tokens.weight"] |
| logits = lin(last, emb, None, "lm_head") |
| if act != F32: |
| logits = g.cast(logits, F32) |
| g.out(logits, "logits", F32, [1, d.vocab]) |
| vw = np.concatenate([W[f"semantic_vad_heads.{j}.weight"] for j in range(d.vad_heads)], 0) |
| vb = np.concatenate([W[f"semantic_vad_heads.{j}.bias"] for j in range(d.vad_heads)], 0) |
| lastf = g.cast(last, F32) if act != F32 else last |
| gv = g.op("Add", [g.op("MatMul", [lastf, g.init("vad.weight_t", vw.T.astype(np.float32))]), |
| g.init("vad.bias", vb.astype(np.float32))]) |
| g.out(g.op("Reshape", [gv, g.const([1, d.vad_heads, d.vad_classes])]), "vad_logits", F32, |
| [1, d.vad_heads, d.vad_classes]) |
| else: |
| xo = g.cast(x, F32) if act != F32 else x |
| g.out(xo, "hidden", F32, [1, "n", d.t_hidden]) |
| g.finish({"winstt_audio8_infinite": "decoder", "precision": precision, "layers": f"{l0}:{l1}", |
| "rope_theta": d.t_theta, "head_dim": d.t_head_dim}) |
| return lin.count |
|
|
|
|
| |
| class OrtBackend: |
| """Session backend over the exported graphs (numpy in/out).""" |
|
|
| def __init__(self, enc_path, dec_paths, providers=("CPUExecutionProvider",), threads=None, opt_dir=None): |
| import onnxruntime as ort |
|
|
| so = ort.SessionOptions() |
| if threads: |
| so.intra_op_num_threads = threads |
| so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| self.enc = ort.InferenceSession(str(enc_path), so, providers=list(providers)) if enc_path else None |
| self.decs = [ort.InferenceSession(str(p), so, providers=list(providers)) for p in (dec_paths or [])] |
| self.kv_dtype = {} |
| for s in ([self.enc] if self.enc else []) + self.decs: |
| for i in s.get_inputs(): |
| if i.name.startswith("past_key_"): |
| self.kv_dtype[id(s)] = np.float16 if "float16" in i.type else np.float32 |
| break |
|
|
| def encoder(self, audio, frame_len, cos, sin, bias, past_k, past_v): |
| dt = self.kv_dtype[id(self.enc)] |
| feed = {"audio": audio, "frame_len": np.array([frame_len], np.int64), "cos": cos, "sin": sin, "attn_bias": bias} |
| L = len(past_k) |
| for i in range(L): |
| feed[f"past_key_{i}"] = past_k[i].astype(dt, copy=False) |
| feed[f"past_value_{i}"] = past_v[i].astype(dt, copy=False) |
| names = ["audio_embeds"] + [f"key_{i}" for i in range(L)] + [f"value_{i}" for i in range(L)] |
| outs = self.enc.run(names, feed) |
| return outs[0], [o.astype(np.float32) for o in outs[1:1 + L]], [o.astype(np.float32) for o in outs[1 + L:]] |
|
|
| def decoder(self, embeds, ada, cos, sin, bias, past_k, past_v, all_logits=False): |
| assert not all_logits, "ORT decoder returns last-position logits only" |
| x = embeds.astype(np.float32) |
| nk, nv = [], [] |
| logits = vad = None |
| for s in self.decs: |
| dt = self.kv_dtype[id(s)] |
| names_in = {i.name for i in s.get_inputs()} |
| feed = {("inputs_embeds" if "inputs_embeds" in names_in else "hidden"): x, "ada_scale": ada, |
| "cos": cos, "sin": sin, "attn_bias": bias} |
| layers = sorted(int(n.split("_")[-1]) for n in names_in if n.startswith("past_key_")) |
| for i in layers: |
| feed[f"past_key_{i}"] = past_k[i].astype(dt, copy=False) |
| feed[f"past_value_{i}"] = past_v[i].astype(dt, copy=False) |
| out_names = [o.name for o in s.get_outputs()] |
| res = dict(zip(out_names, s.run(out_names, feed))) |
| for i in layers: |
| nk.append(res[f"key_{i}"].astype(np.float32)) |
| nv.append(res[f"value_{i}"].astype(np.float32)) |
| if "hidden" in res: |
| x = res["hidden"] |
| else: |
| logits, vad = res["logits"], res["vad_logits"] |
| return logits, vad, nk, nv |
|
|
