ONNX Runtime

Run an ONNX model you already have through ONNX Runtime, with Kwker underneath. Kwker takes the models it runs faster and leaves every other model to ONNX Runtime's own providers. It runs the BERT-family encoders - BERT, RoBERTa, XLM-R, MPNet, ModernBERT, NomicBERT, EuroBERT and the MiniLM, BGE, E5, GTE and Nomic embedding and reranking models - as optimum and torch.onnx export them. Results match ONNX Runtime's CPU provider: the model's float32 weights are used as they are.

Times on a 4-core Xeon without AMX, 128-token texts, medians of three runs (ONNX Runtime's CPU provider on the same file, then the same session with Kwker):

The reranker cross-encoder/ms-marco-MiniLM-L6-v2 scored 32 query-passage pairs of 128 tokens 1.8x faster (218.9 against 404.0 ms).

Prerequisites

• An x86-64 CPU with AVX2 or AVX-512.
• pip install kwker onnxruntime (ONNX Runtime 1.23 or newer; Installation).

Step 1: open a session with Kwker

Pass kwker.onnx.session_options() where you create the session. Nothing else in your code changes:

import onnxruntime as ort
import kwker.onnx

sess = ort.InferenceSession("model.onnx", sess_options=kwker.onnx.session_options())

Step 2: run it

The session runs as before. Here a tiny random BERT, exported to ONNX, stands in for a real model:

import os
import tempfile
import numpy as np
import onnxruntime as ort
import torch
from transformers import BertConfig, BertModel
import kwker.onnx

class Hidden(torch.nn.Module):  # (the hidden states alone: an ONNX graph's outputs are plain tensors)
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, input_ids, attention_mask):
        return self.model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state


model = BertModel(BertConfig(vocab_size=1000, hidden_size=128, num_hidden_layers=2, num_attention_heads=2,
                             intermediate_size=256), add_pooling_layer=False)
path = os.path.join(tempfile.mkdtemp(), "model.onnx")
ids = torch.tensor([[101, 592, 88, 102]])
torch.onnx.export(Hidden(model).eval(), (ids, torch.ones_like(ids)), path, input_names=["input_ids", "attention_mask"],
                  output_names=["last_hidden_state"], opset_version=17, dynamo=False,
                  dynamic_axes={"input_ids": {0: "batch", 1: "length"}, "attention_mask": {0: "batch", 1: "length"}})

sess = ort.InferenceSession(path, sess_options=kwker.onnx.session_options())
feeds = {"input_ids": ids.numpy(), "attention_mask": np.ones((1, 4), np.int64)}
hidden = sess.run(None, feeds)[0]
cpu = ort.InferenceSession(path, providers=["CPUExecutionProvider"]).run(None, feeds)[0]
print(hidden.shape, np.allclose(hidden, cpu, atol=1e-5))

(1, 4, 128) True

Options

session_options() takes the precision and the thread count:

import kwker.onnx

so = kwker.onnx.session_options(precision="bf16", threads=4)

• precision="preserve" (default): the file's own precision. A float model runs in float32 and matches the CPU provider to about 1e-6. An int8 file such as model_qint8_avx512.onnx runs in int8. A float16 file (model_fp16.onnx, model_O4.onnx) and a 4-bit weight file (model_q4.onnx, model_bnb4.onnx) run in float32 with the file's weights.
• precision="float32": always float32 (a quantized file's int8 weights are converted to float32).
• precision="bf16" or "int8": faster on long texts and big batches, with small differences in the output (relative error about 5e-3 in bf16 and 9e-2 in int8 on all-MiniLM-L6-v2).
• threads: how many cores Kwker uses (default: all).
• spinning=True: let ONNX Runtime's own threads spin while they wait for work, as ONNX Runtime does by default. Off by default here, because nodes left to ONNX Runtime (pooling, a classifier head) otherwise keep its threads spinning against Kwker's (about 10% on a reranker).

Notes

• Kwker takes a model's encoder or leaves the whole model to ONNX Runtime. Set KWKER_ONNX_DEBUG=1 to see which: it prints kwker onnx: taking the encoder (...) or the reason a model was left to ONNX Runtime.
• Taken today: BERT, RoBERTa, XLM-R, DistilBERT, MPNet, ModernBERT, NomicBERT and EuroBERT encoders - the plain export, optimum's model_O1 to model_O3 files and the output of ONNX Runtime's own transformer optimizer. Nodes after the encoder - mean pooling, the pooler output, a classifier head such as a reranker's - stay on ONNX Runtime's CPU provider and read Kwker's output; ONNX Runtime then logs a warning that "some nodes were not assigned to the preferred execution providers", which is expected.
• Quantized and float16 files are taken too: ONNX Runtime's dynamic int8 files (model_qint8_*.onnx, model_quint8_*.onnx), static QDQ files, optimum's float16 model_O4.onnx, the transformers.js files (model_int8, model_uint8, model_fp16, model_q4) and bitsandbytes 4-bit files (model_bnb4). On all-MiniLM-L6-v2's model_qint8_avx512.onnx, Kwker's output was closer to the float model than ONNX Runtime's own run of that file (relative error 0.12 to 0.16, against 0.20 to 0.23).
• ModernBERT, NomicBERT and EuroBERT encoders run in float32 or bf16; precision="int8" and their int8 files run in float32.
• Padded positions (where attention_mask is 0) come back as zeros. Mean, CLS and max pooling with the mask give the same embeddings as before.
• kwker.onnx.register() registers the library once and returns Kwker's ONNX Runtime devices, for code that builds its own session options.

Next steps

• Sentence embeddings: the same encoders through Hugging Face models, with pooling and int8.
• Evaluate on your machine: python -m kwker.bench --encoders times ONNX Runtime with and without Kwker.