A standalone, pip-installable TurboQuant adapter for mlx-lm on Apple silicon, with a custom Metal kernel for non-uniform quantization.
TurboQuant (Zandieh, Daliri, Hadian, Mirrokni, 2025) is a data-oblivious (calibration-free) vector quantizer. Its core trick is a random rotation (a Randomized Hadamard Transform): rotating a vector spreads outliers across coordinates and turns the marginal into a concentrated, near-Gaussian distribution that low-bit scalar quantizers handle gracefully. Crucially, an orthogonal rotation preserves inner products, so attention scores computed on rotated queries and keys are unchanged — which is what makes it so effective for KV-cache quantization.
This adapter implements both regimes:
Weights (MSE regime). Rotate each weight matrix, then quantize. Rotation is the robust, always-on win; you can quantize with MLX's fast affinequantized_matmul
(default) or with a customnon-uniform Lloyd–Max LUT Metal kernel(--mode lut
).KV cache (inner-product regime). A drop-inTurboQuantKVCache
that stores keys in the rotated frame and rotates the query to match. This is where TurboQuant shines (see numbers below).
pip install mlx-turboquant # from PyPI
pip install -e .
Requires mlx>=0.31.2
and mlx-lm
on macOS/Apple silicon.
turboquant convert --model mlx-community/Qwen3-0.6B-bf16 --out ./qwen3-tq4 --bits 4
Produces a standard mlx-lm model directory (safetensors + config.json
with a
quantization_config
of quant_method: turboquant
).
import mlx_turboquant as tq
tq.register() # teach stock mlx-lm to load turboquant dirs
from mlx_lm import load, generate
model, tok = load("./qwen3-tq4")
print(generate(model, tok, prompt="Why is the sky blue?", max_tokens=128, verbose=True))
or the CLI:
turboquant generate --model ./qwen3-tq4 --prompt "Why is the sky blue?"
turboquant serve
wraps mlx_lm.server
: it installs the TurboQuant hooks (so a
TurboQuant-quantized model dir loads through the stock server), optionally swaps
in the rotated TurboQuant KV cache, then forwards every other flag straight to
mlx_lm.server
. The result is a drop-in OpenAI-compatible endpoint.
turboquant serve --model ./qwen3-tq4 --port 8080
turboquant serve --model ./qwen3-tq4 --port 8080 --kv-bits 4 --qjl
TurboQuant-specific flags (everything else — --host
, --port
,
--adapter-path
, --temp
, --max-tokens
, --trust-remote-code
, … — is parsed
by mlx_lm.server
; run turboquant serve --help
for the full list):
| flag | default | effect |
|---|---|---|
--kv-bits N |
||
| off | enable the rotated TurboQuant KV cache at N bits (omit → normal fp16 cache) | |
--kv-group-size N |
||
| 64 | KV quantization group size | |
--qjl |
||
| off | add the unbiased 1-bit QJL residual estimator (near-fp16 scores at low KV bits) |
Endpoints (served by mlx_lm.server
): POST /v1/chat/completions
,
POST /v1/completions
, GET /v1/models
.
Model id gotcha:themodel
field in each request must match the server's--model
(the exact path or HF repo id). A short/renamed id makes the server try toload that as a new modelfrom the Hub (→ 401 "Repository Not Found"). Use the same string you launched with, or omitmodel
to use the default.
curl http://127.0.0.1:8080/v1/chat/completions -H "Content-Type: application/json" -d '{
"model": "./qwen3-tq4",
"messages": [{"role": "user", "content": "Name two primary colors."}],
"max_tokens": 64, "temperature": 0.0
}'
curl -N http://127.0.0.1:8080/v1/chat/completions -H "Content-Type: application/json" -d '{
"model": "./qwen3-tq4",
"messages": [{"role": "user", "content": "Count to five."}],
"stream": true
}'
Or with the official OpenAI Python client:
from openai import OpenAI
client = OpenAI(base_url="http://127.0.0.1:8080/v1", api_key="not-needed")
resp = client.chat.completions.create(
model="./qwen3-tq4", # match the server's --model
messages=[{"role": "user", "content": "Hello!"}],
)
print(resp.choices[0].message.content)
Requires mlx-lm>=0.31.3 — the server runs each request in a worker thread, and older mlx-lm uses a single global generation stream that fails there (
RuntimeError: no Stream(gpu, 0)
); 0.31.3+ uses thread-local streams. (pip install mlx-turboquant
pulls a compatible version automatically.) The stock
mlx_lm.server ...
also works for TurboQuant weight models if you call
mlx_turboquant.register()
in the process first — turboquant serve
just does that for you and adds the KV-cache flags.The KV cache is applied at generation time to any (even unquantized) model:
import mlx_turboquant as tq
from mlx_lm import load, generate
model, tok = load("mlx-community/Qwen3-1.7B-bf16")
cache = tq.make_prompt_cache(model, kv_bits=4, kv_group_size=64) # rotated KV
cache = tq.make_prompt_cache(model, kv_bits=3, kv_group_size=64, qjl=True)
generate(model, tok, prompt=..., max_tokens=256, prompt_cache=cache, verbose=True)
The rotated-MSE key k̂
gives a biased attention score:
<Rq, k̂> = <Rq, Rk> − <Rq, r>
(it under-counts by the residual inner product
<Rq, r>
), and the bias is largest for the high-similarity keys softmax
weights most. TurboQuant fixes this with a 1-bit QJL sketch of the residual
r = Rk − k̂
: store sign(RHT₂(r))
(1 bit/channel) and ‖r‖
, then estimate
<Rq, r> ≈ √(π/2d)·‖r‖·⟨RHT₂(Rq), sign(RHT₂(r))⟩
. Adding it back yields an
unbiased score. Verified numerically (tests/test_qjl.py):
for correlated (query, key) pairs the MSE-only bias of ≈ −0.05 is removed to
≈ 0. Enable with qjl=True
.
Teacher-forced perplexity with a quantized KV cache (lower is better):
Qwen3-1.7B (fp16 reference = 2.93):
| KV bits | plain affine KV | TurboQuant KV | TurboQuant + QJL (+1 b/ch) | |---|---|---|---| | 8 | 2.91 | 2.93 | 2.91 | | 4 | 31.39 💥 | 3.16 | 3.03 ✅ | | 3 | 4625 | 77.5 | 4.82 | | 2 | 2.1e6 | 28704 | 6937 |
Two effects, both the paper's claims, reproduced:
Rotation— at 4-bit, plain affine KV collapses (ppl ~31) while TurboQuant's rotated KV stays near-neutral (3.16). Rotation preserves inner products and removes the outliers that wreck low-bit affine KV.QJL residual— the +1-bit unbiased correction closes the last gap: 4-bit KV becomes fp16-neutral (2.93), and it rescues 3-bit KV from unusable (77.5) to usable (5.55). This matches the paper's "quality-neutral at ~3.5 bits/channel".
(2-bit weight-only and ≤3-bit KV without QJL break these small models; the larger the model, the lower the bits you can push.)
Median over 50 ShareGPT prompts of varied length (13–2631 tokens, median
224), 64 decode tokens each, on Apple silicon. prepare_sharegpt.py
samples the
prompts; throughput.py
runs the grid.
| config | decode tok/s | TPOT (ms) | TTFT (ms) | weights | KV @ 2k ctx |
|---|---|---|---|---|---|
| MLX LM bf16 (baseline) | 26.6 | 37.6 | 327 | 3.44 GB | 0.23 GB |
| TurboQuant 4-bit + KV4 | |||||
| 56.6 | |||||
| 17.7 | |||||
| 330 | 1.42 GB | ||||
| 0.066 GB | |||||
| TurboQuant 4-bit + KV4 + QJL | 27.6 | 36.3 | 543 | 1.42 GB | 0.074 GB |
Recommended config — TurboQuant 4-bit weights + rotated 4-bit KV:
~2.1× faster decode than bf16 (56.6 vs 26.6 tok/s) and**~2.1× lower time-per-token**(17.7 vs 37.6 ms) — memory-bound decode loves 4-bit weights, and the rotated KV cache adds essentially no overhead.** 2.4× smaller weights**(3.44 → 1.42 GB) and**~3.5× smaller KV cache**(0.23 → 0.066 GB at 2k context) — so you fit far longer contexts in the same memory, the real constraint for on-device long-context inference.- All while staying near fp16 quality at 4-bit KV where plain affine KV collapses (see above).
QJL mode (qjl=True
) is the maximum-quality / maximum-compression option: it makes 3-bit KV usable and 4-bit KV fp16-neutral for only +1 bit/channel (0.066 → 0.074 GB). It trades decode speed for that quality (the unbiased correction adds an extra sketch dot-product per step), so reach for it when you are memory-bound at very low KV bits and want fp16-grade scores.
The KV cache — not the weights — is what grows with context and dominates memory for long sequences. Since MLX's affine KV is unusable at 4-bit (ppl ~31), the honest comparison is iso-quality: to stay near fp16, affine needs 8-bit KV while TurboQuant is neutral at 4-bit, so TurboQuant's KV cache is ~1.9× smaller at matched quality (and 3.5× smaller than fp16).
Measured on Qwen3-1.7B (benchmarks/kv_memory.py
→ plot_kv_memory.py
); the
ShareGPT prompt range is shaded, extrapolated to long context. KB/token is
measured directly from the cache arrays (verified constant across 2k/4k/8k, and
matching the throughput benchmark's independent 0.066 GB @2k), so the GB columns
are exact linear scaling — not estimates. ppl is the quality proxy from
kv_quality.py
at 509-token context (not re-measured at 32k/128k):
| KV config | quality (ppl) | KB/token | KV @ 32k | KV @ 128k |
|---|---|---|---|---|
| fp16 | 2.93 | 112.0 | 3.76 GB | 15.0 GB |
| MLX affine 8-bit | 2.91 | 59.5 | 2.00 GB | 8.0 GB |
| MLX affine 4-bit | 31.4 ✗ | 31.5 | 1.06 GB | 4.2 GB |
| TurboQuant 4-bit | ||||
| 3.16 | ||||
| 31.5 | 1.06 GB | |||
| 4.2 GB | ||||
| TurboQuant 3-bit + QJL | ||||
| 4.82 | ||||
| 28.4 | 0.95 GB | 3.8 GB |
TurboQuant 4-bit uses the same bytes as affine 4-bit but is actually usable (3.16 vs 31.4). At 128k context that's a 4.2 GB KV cache instead of 8 GB (affine, matched quality) or 15 GB (fp16) — often the difference between fitting long context on-device or not.
TurboQuant needs two operations MLX can't express with built-ins, so the adapter
ships two hand-written mx.fast.metal_kernel
s:
Non-uniform LUT dequant + matmul(kernels/qmm.py
).mx.quantized_matmul
only supports uniform/affine codebooks; TurboQuant's MSE-optimal quantizer uses anon-uniform Lloyd–Max codebook. The kernel does LUT dequant + matmul directly on packed indices (one SIMD group reduces overK
per output),bits ∈ {2,4,8}
, and cuts weight-reconstruction MSE ~15% vs affine at 2-bit. Enable with--mode lut
.Packed-sign QJL inner product(kernels/qjl_dot.py
). The unbiased KV correction computesΣ_d qproj[d]·sign_bit(d)
against the 1-bit residual sketch. The kernel reads the packeduint32
sign words directly and accumulates±qproj
per bit in fp32 — no dense unpack, no extra full matmul — powering the QJL KV path.
Rotations run on Metal via mx.hadamard_transform
.
Weights:nn.Linear → TurboQuantLinear
module swap (same pattern as mlx-lm'sbitnet_quantize
).register()
wrapsmlx_lm.utils.load_model
so turboquant dirs load through the stockmlx_lm.load
/generate
/server
.KV cache:TurboQuantKVCache
subclasses mlx-lm'sQuantizedKVCache
(inherits all mask/state logic), rotates keys, and (withqjl=True
) stores the 1-bit residual sketch;register()
patchesscaled_dot_product_attention
to rotate the query and add the QJL correction.
- A fused RHT Metal kernel (currently we reuse MLX's already-optimized
mx.hadamard_transform
). - 3-bit LUT packing (32 not divisible by 3).
pip install -e ".[test]"
pytest -q # rotation invariance, codebook MSE, kernel numerics, KV + QJL correctness
python benchmarks/kv_quality.py --model mlx-community/Qwen3-0.6B-bf16 # KV perplexity
python benchmarks/throughput.py --model mlx-community/Qwen3-1.7B-bf16 # TPOT/TTFT over ShareGPT
python benchmarks/kv_memory.py && python benchmarks/plot_kv_memory.py # memory chart
MIT licensed. Not affiliated with the TurboQuant authors or Apple.