Kwker

Speed up PyTorch inference

You add one line to a program that already works. torch's CPU kernels for sorting, top-k, unique, quantiles, searching and indexing then run on Kwker, about 9x faster each (median of 97 real ML inputs, one thread), and give torch's results bit for bit. Your model's code does not change. About ten minutes.

Prerequisites

Step 1: switch Kwker on

PythonNeeds torch: runs on your machine.
import torch
import kwker.torch_ops as sto

sto.install()                          # from here on, torch's CPU kernels below run on Kwker
x = torch.randn(100_000)
values, indices = torch.sort(x)
top = torch.topk(x, 5)
print(top.values.shape)
Output
torch.Size([5])

install() replaces torch's CPU kernels for sort, argsort, msort, topk, kthvalue, median, nanmedian, unique, searchsorted, bucketize, quantile, sparse coalescing and the embedding weight gradients. It applies to everything in the process: your code, the libraries you call, TorchScript and torch.compile's extern kernels. Anything Kwker does not cover runs on torch's own kernel.

Step 2: check that the results did not change

Run the same call with Kwker off and compare:

PythonRuns on your machine.
fast = torch.sort(x).values
sto.uninstall()                        # torch's own kernels again
print(torch.equal(fast, torch.sort(x).values))
sto.install()
Output
True

You do not have to do this for every call. The first install() in a process checks every replaced kernel against torch's on probe inputs and leaves off any kernel that differs. The result is cached in ~/.cache/kwker/, and python -m kwker doctor shows it.

To use Kwker for one block of code only, use the context manager instead of install():

PythonRuns on your machine.
sto.uninstall()
with sto.accelerated():                # Kwker inside the block only
    q = torch.quantile(x, torch.tensor([0.1, 0.5, 0.9]))
print(q.shape)
Output
torch.Size([3])

Step 3: bfloat16 and Hugging Face models

bfloat16 models need no extra setting. install() runs their linear layers on Kwker's bfloat16 kernels, on CPUs without AMX too. Many float32 checkpoints hold bfloat16 values, because the model was trained in bfloat16; Kwker reads those weights through a bfloat16 copy, with half the bytes and the same values. For Llama-style Hugging Face models, install() also fuses each RMSNorm, rotary embedding and SiLU(gate) x up into one operator.

So a plain model.generate(...) gets faster. On 4 cores, bfloat16 models from 135M to 1.7B parameters ran their 128-token prompts 2.3-3.6x faster and their new tokens 1.3-1.7x faster, with the same greedy tokens. bfloat16 kernels add up in another order than torch's, so single values can differ by bfloat16 rounding.

Step 4: measure your own program

kwker audit runs your program several times with and without Kwker, alternating, and compares both the time and the output:

ShellOn your machine.
python -m kwker audit -- predict.py
python -m kwker audit --output preds.npy --rtol 1e-4 -- predict.py     # also compare a file the program writes

It reports wall time, CPU time, the speed-up and whether the results agree. A program gets faster only as much as these kernels are of its run time: if sorting takes 10% of a step, a 9x faster sort makes the step about 10% faster.

Turn it off

Next steps