VeloxQuant-MLX 0.24.1

Fast KV cache quantization for Apple Silicon — TurboQuant, RVQ, VecInfer (with Metal kernels), RateQuant, PolarQuant, SpectralQuant, CommVQ, RaBitQ, QJL, KIVI, KIVI-Sink, SVDq, Kitty, AdaKV-proxy, XQuant, KVQuant-NUQ, PALU, CacheGen, MiniCache, GEAR, ZipCache-adapted, SnapKV-adapted, StreamingLLM-adapted, H2O-adapted, TOVA-adapted, PyramidKV-adapted, and SqueezeAttention-adapted in MLX

VeloxQuant-MLX

Fast KV Cache Quantization for Apple Silicon
TurboQuant · RVQ · VecInfer · RateQuant · PolarQuant · QJL · SpectralQuant · CommVQ · RaBitQ — in MLX

---

A KV-cache compression library for mlx_lm that compresses the Key tensor up to 16× with near-lossless quality on Apple M-series chips. Ships twenty-seven compression strategies — from zero-calibration 1-bit RVQ to RaBitQ (1-bit keys + MSE-b4 values) which achieves 6× full KV compression and fits 6× more context in the same RAM budget on Falcon3-7B, through six token-eviction caches (SnapKV, StreamingLLM, H2O, TOVA, PyramidKV, SqueezeAttention) — plus a hand-written Metal compute kernel that makes the VecInfer quantize hot path 6.9–14.7× faster (13× at S=2048) and 98% lighter on peak memory at the OOM-trigger shape. (The companion dequant kernel is at MLX mx.take parity — the speedup is on the quantize path.) Plug it in with three lines; mlx_lm.generate runs unchanged.

---

Numbers that matter

Metric	Value	Notes
Max key cache compression	16×	VecInfer-1bit, head_dim=128
Metal kernel speedup	13×	quantize_vq at S=2048 (range 6.9–14.7× over S=128–8192)
Peak memory reduction	98%	729 MB → 12 MB, Falcon3-7B shape
RVQ-1bit compression	7.5×	Near-zero throughput cost
FP16 throughput retained	100%	Qwen2.5-7B at 16× compression
SpectralQuant compression	5.33×	per-model measured (Qwen2.5-0.5B / Gemma-4-4B), same bit-width
SpectralQuant cosine sim	+3pp	over TurboQuant on Qwen2.5-0.5B
RaBitQ full KV compression	6×	1-bit keys + MSE-b4 values, Falcon3-7B
RaBitQ context at 8 GB	~103k tokens (est.)	KV-only linear extrapolation from measured memory rows; vs ~17k fp16 — 6× more context
CommVQ key compression	64×	RoPE-commutative VQ, D=128, n_cb=4
KIVI-2bit key compression	5.8×	per-channel keys / per-token values; measured on Llama-3.2-3B, Qwen2.5-7B, Mistral-7B
KIVI-2bit full-KV compression	~4×	incl. fp16 residual window (32 tokens); 100–106% of fp16 throughput
Production models validated	12	Llama, Mistral, Qwen, Phi, Gemma 3/4, Falcon

---

Table of contents

1. Installation
2. Quickstart
3. RaBitQ — new in 0.7.0
4. CommVQ — RoPE-commutative VQ
5. SpectralQuant — new in 0.6.0
6. RateQuant — per-layer mixed precision
7. VecInfer — 16× product VQ
8. Metal kernels
9. Benchmark results
10. Algorithm guide
11. What's inside
12. Architecture
13. CLI
14. Development
15. Blog posts
16. References

---

Installation

pip install VeloxQuant-MLX

Requirements: Apple Silicon M1+, Python ≥ 3.11, MLX ≥ 0.18, NumPy ≥ 1.26.

Install from source

git clone https://github.com/rajveer43/VeloxQuant-MLX
cd VeloxQuant-MLX
pip install -e ".[dev]"

---

Quickstart

RVQ 1-bit — 7.5× compression, no calibration (recommended)

import mlx_lm
from veloxquant_mlx import KVCacheBuilder, KVCacheConfig

model, tokenizer = mlx_lm.load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")

config = KVCacheConfig(method="turboquant_rvq", bit_width_inlier=1, seed=42)
caches = KVCacheBuilder.for_model(model, config)
model.make_cache = lambda *_a, **_k: caches

response = mlx_lm.generate(model, tokenizer,
    prompt="Explain the theory of relativity in simple terms.",
    max_tokens=200,
)

VecInfer 1-bit — 16× compression, Metal kernels auto-detected

import mlx_lm
from veloxquant_mlx import KVCacheConfig, KVCacheFactory
from veloxquant_mlx.allocators.vecinfer import calibrate_smooth_factors, train_codebook

model, tokenizer = mlx_lm.load("mlx-community/Qwen2.5-7B-Instruct-4bit")

# One-time offline calibration — save and reuse
smooth   = calibrate_smooth_factors(sample_keys)
codebook = train_codebook(sample_keys_flat, n_centroids=256, sub_dim=8)

config = KVCacheConfig(
    method="vecinfer",
    head_dim=128,
    key_codebook_bits=8,
    key_sub_dim=8,
    smooth_factors=smooth,
    key_codebook=codebook,
    use_metal_kernels=None,   # None=auto-detect, True=require, False=forbid
)
caches = KVCacheFactory.create_for_model(model, config)

response = mlx_lm.generate(model, tokenizer,
    prompt="Write a 5,000-word analysis of the RLHF literature.",
    max_tokens=5000,
    prompt_cache=caches,
)

RateQuant — mixed precision per layer

from veloxquant_mlx import (
    KVCacheBuilder, KVCacheConfig,
    calibrate_layer_sensitivities,
    allocate_bits_ratequant,
)

# Step 1 — 1.6s one-time probe on real activations
weights = calibrate_layer_sensitivities(model, tokenizer)

# Step 2 — closed-form reverse-waterfilling allocation
alloc = allocate_bits_ratequant(weights, target_avg_bits=1.5, beta=3.5)
# e.g. [1, 2, 1, 1, 3, 1, 2, ...]  — one int per layer

# Step 3 — build per-layer caches
config = KVCacheConfig(method="turboquant_rvq", bit_width_inlier=alloc, seed=42)
caches = KVCacheBuilder.for_model(model, config)

---

RaBitQ — new in 0.7.0

RaBitQ (SIGMOD 2024, adapted from Ascend-RaBitQ arXiv:2605.16007) is the first method in VeloxQuant-MLX to compress both keys and values, achieving 6× full KV compression on Falcon3-7B-Instruct-4bit.

How it works:

1. IVF clustering — K-Means partitions keys into nlist clusters; only nprobe are searched per query
2. Randomised Hadamard rotation — reuses mx.hadamard_transform + make_hadamard_diagonal, O(D log D)
3. 1-bit sign quantization — sign(rotated_residual) packed into D/8 uint8 bytes per key (11.6× key compression at D=256)
4. Metal Hamming kernel — rabitq_hamming_score computes XOR + popcount distance for all candidates in one GPU dispatch
5. TurboQuantMSE b=4 values — scalar MSE-optimal codebook on values adds 4× value compression

Results on Falcon3-7B-Instruct-4bit (28 layers, 4 KV heads, D=256):

Method	KV Memory @ 1024 tok	Compression	Context @ 8 GB
fp16 baseline	117.4 MB	1×	~17k tokens
RaBitQ keys + fp16 values	63.8 MB	1.8×	~31k tokens
RaBitQ keys + MSE-b4 values	19.7 MB	6×	~103k tokens

from veloxquant_mlx.quantizers.rabitq import RaBitQQuantizer
from veloxquant_mlx.quantizers.turboquant_mse import TurboQuantMSE
import mlx.core as mx, numpy as np

# Keys: RaBitQ 1-bit  (11.6× compression on key tensors)
q_key = RaBitQQuantizer(d=256, nlist=64, nprobe=8, rerank=32, seed=42)
q_key.fit(mx.array(calibration_keys))

# Values: MSE-b4 scalar quantization (4× compression)
q_val = TurboQuantMSE(d=256, b=4, use_hadamard=True)

# Encode KV at each decode step
ev_k = q_key.encode(keys)   # [N, D//8] uint8 sign bits + IVF meta
ev_v = q_val.encode(values)  # [N, D//4] uint8 scalar indices

# Decode for attention
k_hat = q_key.decode(ev_k)  # [N, D] fp16 — approx reconstructed keys
v_hat = q_val.decode(ev_v)  # [N, D] fp16 — approx reconstructed values

What grows with context: both memory and decode latency scale linearly with T. The 6× compression slope means you can sustain 6× longer contexts before hitting any RAM limit. At 32k tokens fp16 needs 3.76 GB; RaBitQ+MSE4v needs only 631 MB.

Benchmark figures: figures/RaBitQ/falcon/ · Metal kernel figures: figures/RaBitQ/kernel/

---

CommVQ — RoPE-commutative VQ

CommVQ (arXiv:2506.18879, Apple ML Research, ICML 2025) solves the fundamental incompatibility between vector quantization and RoPE positional encodings:

The problem: Standard VQ applied after RoPE fails because quantize(rotate(x)) ≠ rotate(quantize(x)). The positional encoding rotates the keys differently for each position, so a codebook trained at position 0 gives wrong reconstructions at position T.

The fix: Train codebooks on pre-RoPE keys (at position 0). After each K-Means M-step, project every centroid onto the RoPE-commuting subspace — each pair of dimensions (2i, 2i+1) is symmetrised to mean_val = (a+b)/2 for same-sign pairs. RoPE is then applied exactly at decode time using stored positions.

from veloxquant_mlx.quantizers.comm_vq import CommVQQuantizer
import mlx.core as mx, numpy as np

q = CommVQQuantizer(d=128, b=8, n_codebooks=4, seed=42)
q.fit(mx.array(pre_rope_keys))          # train on pre-RoPE keys

# Encode: stores residual VQ indices + positions
ev = q.encode(keys_pre_rope, positions=position_ids)

# Decode: gathers centroids, applies RoPE in one step
k_hat = q.decode(ev)                    # [N, D] fp16, post-RoPE

# Approximate inner product (for attention scoring)
scores = q.estimate_inner_product(query, ev)  # [N]

Config	Compression	RoPE compatible
D=128, n_cb=4, b=8	64× vs fp16	✓ exact
D=128, n_cb=4, b=4	64× vs fp16	✓ exact

---

KIVI — tuning-free asymmetric 2-bit baseline

KIVI is a re-implementation of "KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache" (Liu, Yuan et al., ICML 2024) — the most widely-cited KV-cache quantization baseline. It is included so every other method in this library can be measured against the field's reference point.

The asymmetry (KIVI's core idea):

1. Keys are quantized per channel (group-wise min/max along the token axis) — key distributions have a few high-variance channels, so per-channel scales keep them accurate.
2. Values are quantized per token (group-wise along the channel axis).
3. The most recent residual_length tokens are kept in fp16 — newly generated tokens dominate attention and are cheap to keep exact; they are quantized only once they age out of the residual window.

KIVI is fully deterministic (min/max group quantization, no codebook training, no RNG), so it adds no run-to-run variance.

import mlx_lm
from veloxquant_mlx import KVCacheConfig, KVCacheBuilder

model, tokenizer = mlx_lm.load("mlx-community/Llama-3.2-3B-Instruct-4bit")

config = KVCacheConfig(method="kivi", bit_width_inlier=2,
                       kivi_group_size=32, residual_length=32)
caches = KVCacheBuilder.for_model(model, config)
model.make_cache = lambda *_a, **_k: caches

response = mlx_lm.generate(model, tokenizer,
    prompt="...", max_tokens=120)

Measured results (Apple M4, max_tokens≈120, residual_length=32; source: figures/kivi/<model>/results.json):

Model	KIVI-2bit key comp.	full-KV comp. (incl. fp16 residual)	throughput vs fp16
Llama-3.2-3B-4bit	5.79×	3.98×	16.3 vs 16.0 tok/s (102%)
Qwen2.5-7B-4bit	5.78×	3.98×	7.6 vs 7.6 tok/s (100%)
Mistral-7B-4bit	5.76×	4.03×	6.8 vs 6.4 tok/s (106%)

Honest scope:

• KIVI's published speedup comes from a CUDA kernel that does not port to Metal. On Apple Silicon the win is memory; throughput here is at-or-near fp16 because the per-channel/per-token min/max arithmetic is cheap on a memory-bound decode path.
• Compression only manifests once context exceeds the residual window — at short prompts the entire prefill stays fp16 and the realized ratio is 1.0× (this is correct behavior, not a bug). The numbers above use a long-context prompt.
• Peak runtime memory is not reduced (often marginally higher): like every method here, keys are dequantized to fp16 before SDPA, so the compression is in cache-storage accounting, not the peak fp16 working set.
• At 2 bits, raw-key reconstruction cosine on synthetic unit-norm Gaussian keys is ~0.93 — KIVI 2-bit is genuinely lossy, which is exactly why the fp16 residual window exists. VecInfer-2bit compresses harder (8× vs 5.8× keys); KIVI's value is being the recognized, calibration-free baseline. See figures/kivi/fig4_vs_existing.png.

KVSink-adapted sink protection — new in 0.9.0

method="kivi_sink" layers dynamic attention-sink protection on top of KIVI. Inspired by, not a faithful port of, KVSink (Su & Yuan, COLM 2025): the paper detects sinks via hidden-state outlier channels at a model-specific emergence layer, which cache wrappers cannot see — this implementation uses the cache-observable proxy of anomalously high key L2-norm (running top-k of token positions, mean over KV heads). Selected tokens are kept fp16 and excluded from quantization-parameter calibration — the detail the paper insists on: without it, a high-magnitude sink inflates its group's min/max scale and ruins every neighbor even though the sink itself is restored (our tests reproduce that failure mode).

config = KVCacheConfig(method="kivi_sink", bit_width_inlier=2,
                       kivi_group_size=32, residual_length=32,
                       n_sink_tokens=5)   # top-k high-key-norm tokens kept fp16

Evidence (unit tests on synthetic planted-sink data — tests/cache/test_sink_cache.py, 9 passing): planted sinks preserved bit-exact while neighbors quantize; sink-protected MSE < plain KIVI at equal bit-width; dynamic selection MSE < Preserve-First-N at equal fp16 budget when sinks are not all at the front (the paper's central claim, reproduced at cache level); n_sink_tokens=0 reproduces plain KIVI bit-for-bit.

Not yet benchmarked end-to-end: benchmark_scripts/benchmark_sink.py is ready but has not been run — no throughput or compression figures are claimed for this method until its results.json is committed. Known v1 limitation: sink selection is prefill-dominant (tokens already quantized are not retroactively restored).

PALU — true low-rank latent storage — new in 0.15.0

method="palu" is the first method where the KV cache itself stays low-rank. Inspired by, not a faithful port of, PALU (Chang et al., ICLR 2025): at prefill it partitions the attention heads into groups, fits one shared projection per group via group-head SVD (PALU's G-LRD), and stores the projected latent codes [S, r] directly — full fp16 keys/values are reconstructed only at attend time. Both keys and values are compressed (unlike SVDq, which is keys-only and reconstructs full fp16 so its win is bandwidth accounting). The latents are mixed-bit quantized (top-25% of channels by singular value at 4-bit, the rest at 2-bit) for a full-KV effective rate below 1 bit/element on low-rank data. Zero calibration.

config = KVCacheConfig(method="palu", head_dim=128,
                       palu_energy_threshold=0.90,   # rank from singular-value energy
                       palu_n_head_groups=4,         # heads share a projection
                       palu_hi_bit=4, palu_lo_bit=2, # mixed-bit latents
                       palu_quantize_values=True)    # False → low-rank-only (fp16 latents)

Evidence (unit tests on synthetic low-rank data — tests/cache/test_palu_cache.py (13) + tests/quantizers/test_palu.py (9), all passing): the cache stores [S, r] latents and the parent fp16 ring buffer is never populated (cache.keys is None); PALU reconstruction MSE < naive 2-bit on both keys and values; both tensors compress vs fp16; group-head SVD recovers a planted rank-r subspace; assigned_avg_bits < 2.0; deterministic. The offline harness in benchmark_palu.py reports key MSE 1.54 vs 2.37 and value MSE 1.54 vs 2.52 (naive 2-bit) at r=16/D=128 — synthetic, not model-level.

Not yet benchmarked end-to-end: benchmark_scripts/benchmark_palu.py is ready but has not been run — no throughput or compression figures are claimed until its results.json is committed. Known limitation: PALU's fused low-rank-reconstruction attention kernel is not ported (we reconstruct fp16 then call MLX SDPA), so peak memory during attention is not reduced — only the stored cache size.

CacheGen — entropy-coded KV cache — new in 0.16.0

method="cachegen" is the first method in the suite to entropy-code the quantized KV. Inspired by, not a faithful port of, CacheGen (Liu et al., SIGCOMM 2024): every other method packs codes at a fixed bit-width; CacheGen exploits token-wise locality (adjacent tokens' KV are similar) by applying a reversible token-delta transform to the codes and compressing the low-entropy residual toward its Shannon entropy. Reconstruction is identical to plain group quant (lossless over the codes) — the win is storage.

config = KVCacheConfig(method="cachegen", head_dim=128,
                       cachegen_bits=4, cachegen_group_size=32,
                       cachegen_use_delta=True)

We do not ship a serial range codec (it would bottleneck MLX's parallel decode); the entropy-coded byte size is modelled from the measured symbol entropy and capped at the fixed-width packed size, so savings are never negative.

Evidence (tests/cache/test_cachegen_cache.py (12) + tests/quantizers/test_cachegen.py (9), all passing): reconstruction byte-identical to _group_quant_dequant; token-delta reversible; delta entropy < raw entropy on correlated data; positive savings on correlated, exactly 0% (never negative) on iid; entropy primitives correct (0 for constants, 1 bit for 50/50). Offline harness: ~17% savings on correlated 3-bit, 0% on iid — synthetic, not model-level.

Not yet benchmarked end-to-end: benchmark_scripts/benchmark_cachegen.py is ready but has not been run. Known limitation: storage-only win (codes dequant to fp16 for SDPA) — does not reduce attend-time working set.

MiniCache — cross-layer depth merge — new in 0.16.0

method="minicache" compresses across network depth. Inspired by MiniCache (Liu et al., NeurIPS 2024): adjacent middle-to-deep layers have nearly identical KV directions, so a pair is merged into one shared SLERP-interpolated direction plus each layer's own per-token magnitude (a pair costs ~one layer). High-divergence token pairs are kept unmerged (the retention set). A different route to inter-layer redundancy than XQuant: XQuant reuses codes, MiniCache merges the tensors.

config = KVCacheConfig(method="minicache", head_dim=128,
                       minicache_start_frac=0.5,           # only merge past mid-depth
                       minicache_retention_threshold=0.9,  # keep divergent pairs
                       minicache_slerp_t=0.5)
caches = KVCacheBuilder.for_model(model, config)  # requires for_model (shared coordinator)

Evidence (tests/cache/test_minicache_cache.py (11) + tests/quantizers/test_minicache.py (11), all passing): role assignment (early all primary, deep has merge); SLERP unit-norm + endpoint correctness; similar layers (cosine 0.9995) merge MSE < 2e-4 with 0% retention; opposite directions 100% retained and reconstructed exactly; magnitude preserved through the shared direction; degenerate (no-coordinator) primary is a lossless passthrough. Offline harness confirms the merge-vs-retain split — synthetic, not model-level.

Not yet benchmarked end-to-end: benchmark_scripts/benchmark_minicache.py is ready but has not been run. Known limitation: MiniCache merges fp16 directions (no extra quantization) and the merge happens on reconstructed tensors, so attend-time working set is not reduced — the win is stored cache size.

GEAR — error-feedback KV cache — new in 0.17.0

method="gear" is the first method in the suite on the error-feedback axis. Inspired by, not a faithful port of, GEAR (Kang et al., arXiv:2403.05527): every other method picks a bit-width or a cache layout and lives with the quantization error; GEAR makes any ultra-low-bit base quantizer near-lossless by reconstructing what it threw away — X ≈ Quant_b(X) + L·R + S: an ultra-low-bit base quant, a low-rank approximation of the quantization residual, and a sparse matrix correcting the outlier entries the low-rank term cannot absorb. Unlike CacheGen (reconstruction identical to group quant), GEAR's reconstruction genuinely recovers quality the base bit-width loses.

config = KVCacheConfig(method="gear", head_dim=128,
                       gear_bits=2,                # ultra-low base bit-width
                       gear_rank=8,                # residual low-rank (keep small)
                       gear_sparse_fraction=0.005, # top-|residual| kept exact
                       gear_quantize_values=True)
caches = KVCacheBuilder.for_model(model, config)

Evidence (tests/cache/test_gear_cache.py (10) + tests/quantizers/test_gear.py (13), all passing): GEAR reconstruction MSE strictly below base-quant-alone on low-rank+outlier data; low-rank-alone and sparse-alone each help; rank=0, sparse=0 collapses exactly to base group quant; a rank-r residual recovered to < eps; sparse selection picks true top-magnitude entries; byte-accounting ordering base_only ≤ compressed ≤ fp16 at realistic head dim; error_recovery_ratio in (0,1]; values-off path keeps values fp16; build via both create and for_model. Offline harness reports 11–22% MSE improvement on synthetic low-rank data — synthetic, not model-level.

Not yet benchmarked end-to-end: benchmark_scripts/benchmark_gear.py is an offline-synthetic harness (loads no model) and has not been run on hardware for committed numbers. Known limitation: the stored cache shrinks but reconstruction is fp16 for SDPA, so attend-time working set is not reduced; the low-rank/sparse factors are overhead, so keep the rank low relative to the head dim.

H2O-adapted — cumulative attention-mass heavy-hitter oracle eviction — new in 0.21.0

method="h2o" is the library's third eviction axis and the first based on cumulative per-token attention mass. Inspired by, not a faithful port of, H2O (Zhang et al., ICLR 2024, arXiv:2306.14048): on every step (prefill and decode alike), each incoming token's approximate softmax attention distribution over the existing cache is computed — using the new key vector as a proxy query (true queries are not visible at cache-wrapper level). The weights are accumulated into a per-token cumulative importance score. When the cache exceeds h2o_budget, the lowest-score non-sink token is permanently evicted. The cache is bounded at all times to h2o_budget positions.

Eviction axis	When it fires	Score signal	Memory shape
SnapKV-adapted	Once at prefill end	Key-as-query attention proxy	Grows during decode
StreamingLLM-adapted	Every token	Position (recency + sink)	Constant
H2O-adapted	Every token (budget exceeded)	Cumulative attention mass	Constant (≤ budget)

config = KVCacheConfig(
    method="h2o",
    head_dim=128,
    h2o_budget=512,   # max tokens retained at any time
    h2o_n_sink=4,     # initial positions never evicted (attention sinks)
)
caches = KVCacheBuilder.for_model(model, config)

Evidence: 18 quantizer tests + 15 cache tests — all 33 passing. Budget-never-exceeded verified across a 30-step decode stress test. Sink protection verified. Score non-negativity and accumulation across steps verified. Deterministic.

Honest limitation: key-as-query proxy approximates the paper's true query attention signal; no RoPE position-ID remapping after eviction; uniform budget across all heads. No model-level benchmark run yet — benchmark_scripts/benchmark_h2o.py is an offline-synthetic harness.

Inspired by H2O (arXiv:2306.14048, ICLR 2024, Zhang et al.) — not a faithful port.

SqueezeAttention-adapted — 2D layer×token data-driven budget eviction — new in 0.24.0

method="squeeze" is the library's sixth eviction configuration, the first 2D (layer × token) budget method, and the first with a data-driven per-layer budget. Inspired by, not a faithful port of, SqueezeAttention (Wang et al., 2024, arXiv:2404.04793): it is H2O-adapted's cumulative-attention-mass eviction with a per-layer budget that is measured rather than assumed. Where PyramidKV imposes a fixed positional taper, SqueezeAttention measures each layer's attention concentration during prefill and reallocates a fixed total budget toward broad layers and away from concentrated ones. When squeeze_strength=0.0 it reduces exactly to uniform H2O.

Why measure instead of assume: PyramidKV's pyramid is calibration-free but blind to the actual prompt. SqueezeAttention reads the geometry of each layer's key set (mean pairwise cosine dispersion, an attention-free entropy proxy) and puts budget where this prompt's attention actually spreads. It is the data-driven sibling of PyramidKV.

Eviction axis	Score signal	Budget
SnapKV-adapted	Key-as-query attention proxy	Uniform
StreamingLLM-adapted	Position (recency + sink)	Uniform
H2O-adapted	Cumulative attention mass	Uniform
TOVA-adapted	Current-step attention weight	Uniform
PyramidKV-adapted	Cumulative attention mass	Per-layer fixed pyramid
SqueezeAttention-adapted	Cumulative attention mass	Per-layer data-driven (measured)

config = KVCacheConfig(
    method="squeeze",
    head_dim=128,
    squeeze_budget=512,     # AVERAGE budget across layers (uniform-H2O baseline)
    squeeze_n_sink=4,       # initial positions never evicted (attention sinks)
    squeeze_strength=1.0,   # reallocation: 0.0 = uniform (== H2O), 1.0 = full inverse-concentration
)
caches = KVCacheBuilder.for_model(model, config)   # a shared coordinator re-budgets after prefill

How the re-budget works: a single shared SqueezeCoordinator (injected by for_model) collects each layer's concentration during prefill, computes the schedule once at the prefill boundary with squeeze_budgets(...), and publishes each layer's resolved budget; over-budget layers are trimmed by H2O score. This is the first eviction method with a runtime re-budgeting step, but — unlike XQuant/MiniCache — it exchanges only per-layer scalars and runs its allocation exactly once. Single-cache construction (no coordinator) falls back to squeeze_budget and behaves as one uniform H2O layer.

Evidence: 28 quantizer tests + 19 cache tests — all 47 passing. Concentration proxy verified (1.0 for identical keys, 0.0 for orthogonal, scale-invariant). Allocator verified: strength=0 == uniform, mean within 5% of average, budgets monotone in concentration, floored at n_sink+1, strength bounds enforced. for_model verified to produce data-driven budgets where the broad early layer keeps more tokens than the concentrated deep layer on the same sequence. Coordinator finalises once-only and is idempotent per layer. Deterministic.

Honest limitation: cosine-dispersion proxy for attention entropy (paper reads actual attention maps); key-as-query proxy for eviction (same as H2O); one-shot re-budget at prefill, frozen for decode; no RoPE remapping; uniform budget across heads within a layer. No model-level (perplexity/throughput) benchmark run — benchmark_scripts/benchmark_squeeze.py is an offline-synthetic harness; its results are committed in benchmark_scripts/squeeze_benchmark_results.json (Apple Silicon) and confirm strength=0 == uniform, strength>0 reallocates broad→more/concentrated→less, mean ≈ average.

Inspired by SqueezeAttention (arXiv:2404.04793, Wang et al., 2024) — not a faithful port.

PyramidKV-adapted — layer-adaptive budget attention-mass eviction — new in 0.23.0

method="pyramidkv" is the library's fifth eviction configuration and the first with a per-layer budget. Inspired by, not a faithful port of, PyramidKV (Cai et al., 2024, arXiv:2406.02069): it is H2O-adapted's cumulative-attention-mass eviction wearing a pyramid of budgets instead of a single global one. Early layers (broad attention) get a large budget; deep layers (concentrated attention) get a small one; the average is held fixed so total memory matches uniform H2O. When the pyramid is flat (pyramid_beta=1.0) it reduces exactly to H2O.

Why a pyramid — pyramidal information funneling: attention in early layers is broad and near-uniform, while deep layers concentrate on a few tokens. A uniform budget starves early layers and over-provisions deep ones; redistributing the same total into a pyramid puts memory where attention actually spreads.

Eviction axis	Score signal	Budget
SnapKV-adapted	Key-as-query attention proxy	Uniform
StreamingLLM-adapted	Position (recency + sink)	Uniform
H2O-adapted	Cumulative attention mass	Uniform
TOVA-adapted	Current-step attention weight	Uniform
PyramidKV-adapted	Cumulative attention mass	Per-layer pyramid

config = KVCacheConfig(
    method="pyramidkv",
    head_dim=128,
    pyramid_budget=512,   # AVERAGE budget across layers (uniform-H2O baseline)
    pyramid_n_sink=4,     # initial positions never evicted (attention sinks)
    pyramid_beta=1.5,     # pyramid steepness: 1.0 = flat (== H2O), larger = steeper
)
caches = KVCacheBuilder.for_model(model, config)   # the pyramid needs layer context

The schedule for 12 layers, avg_budget=512, beta=2.0 is [1019, 927, …, 97, 5] — mean exactly 512. The pyramid takes effect only via for_model (which knows each layer's index); single-cache construction falls back to the average budget and behaves as one uniform H2O layer. No runtime coordinator — layers never exchange data during generation (unlike XQuant/MiniCache).

Evidence: 24 quantizer tests + 19 cache tests — all 43 passing. Allocator verified: monotonically decreasing, mean within 5% of average, beta=1.0 == uniform, budgets floored at n_sink+1. for_model verified to produce a decreasing pyramid where early-layer caches retain more tokens than deep-layer caches on the same sequence. Budget-never-exceeded across a 30-step stress test. Deterministic.

Honest limitation: fixed monotone (linear) budget schedule rather than the paper's prefill-entropy-derived allocation — funneling shape preserved, exact per-layer values not data-driven; key-as-query proxy for eviction (same as H2O); no RoPE remapping; uniform budget across heads within a layer. No model-level (perplexity/throughput) benchmark run — benchmark_scripts/benchmark_pyramidkv.py is an offline-synthetic harness; its results are committed in benchmark_scripts/pyramidkv_benchmark_results.json (24 configs on Apple Silicon) and confirm flat-beta == uniform, pyramid-beta strictly decreasing, mean == average.

Inspired by PyramidKV (arXiv:2406.02069, Cai et al., 2024) — not a faithful port.

TOVA-adapted — current-step attention-weight eviction (memoryless) — new in 0.22.0

method="tova" is the library's fourth eviction axis and the first memoryless one. Inspired by, not a faithful port of, TOVA / "Transformers are Multi-State RNNs" (Oren et al., 2024, arXiv:2401.06104): on every step (prefill and decode alike), the new key vector as a proxy query attends to the post-append cache; when the cache exceeds tova_budget, the non-sink token with the lowest current-step attention weight is permanently evicted. The weights are then discarded — no score is carried across steps. The cache is bounded at all times to tova_budget positions.

The key contrast with H2O: H2O accumulates attention mass, so a past heavy hitter resists eviction (inertial). TOVA scores by the present step only, so a token that stops being attended to is dropped even if it dominated earlier (memoryless, reactive). Neither policy dominates — that is exactly why both are provided.

Eviction axis	When it fires	Score signal	Memory shape
SnapKV-adapted	Once at prefill end	Key-as-query attention proxy	Grows during decode
StreamingLLM-adapted	Every token	Position (recency + sink)	Constant
H2O-adapted	Every token (budget exceeded)	Cumulative attention mass (inertial)	Constant (≤ budget)
TOVA-adapted	Every token (budget exceeded)	Current-step attention weight (memoryless)	Constant (≤ budget)

config = KVCacheConfig(
    method="tova",
    head_dim=128,
    tova_budget=512,   # max tokens retained at any time
    tova_n_sink=4,     # initial positions never evicted (attention sinks)
)
caches = KVCacheBuilder.for_model(model, config)

Evidence: 19 quantizer tests + 15 cache tests — all 34 passing. Budget-never-exceeded verified across a 30-step decode stress test. Sink protection verified. Memorylessness (no scores carried across steps) verified. Current-step eviction correctness verified with axis-aligned test vectors (an orthogonal token is dropped over an aligned one). Deterministic.

Honest limitation: key-as-query proxy approximates the paper's true query attention signal; no RoPE position-ID remapping after eviction; uniform budget across all heads. No model-level (perplexity/throughput) benchmark run — benchmark_scripts/benchmark_tova.py is an offline-synthetic harness. Its results are committed in benchmark_scripts/tova_benchmark_results.json (28 configs on Apple Silicon); measured compression ratio equals seq_len / budget exactly across every config.

Inspired by TOVA (arXiv:2401.06104, Oren et al., 2024) — not a faithful port.

StreamingLLM-adapted — sink + recency-window constant-memory eviction — new in 0.20.0

method="streaming_llm" is the repo's first constant-memory cache and first structural positional eviction method. Inspired by, not a faithful port of, StreamingLLM (Xiao et al., ICLR 2024, arXiv:2309.17453): keep only the first stream_n_sink token positions (frozen attention sinks) and the most recent stream_window_size tokens (rolling FIFO). All other positions are permanently evicted. The cache never grows beyond stream_n_sink + stream_window_size positions regardless of generation length — constant decode memory.

config = KVCacheConfig(
    method="streaming_llm",
    head_dim=128,
    stream_n_sink=4,           # initial positions frozen as attention sinks
    stream_window_size=512,    # FIFO capacity for recent tokens
)
caches = KVCacheBuilder.for_model(model, config)

Evidence: 17 quantizer tests + 15 cache tests — all 32 passing. Constant-memory guarantee verified (30-step decode stress test). FIFO trimming and sink-first ordering verified. Deterministic.

Honest limitation: no attention mask adjustment (model sees all returned K/V positions); no RoPE position-ID remapping. No model-level benchmark run yet — cite only after results_streaming_llm.json is committed from a hardware run.

Inspired by StreamingLLM (arXiv:2309.17453, ICLR 2024, Xiao et al.) — not a faithful port.

SnapKV-adapted — prefill observation-window token eviction — new in 0.19.0

method="snapkv" is the repo's first token eviction cache and the first where the paper's actual signal (attention scores) is computable at the cache level without model interception. Inspired by, not a faithful port of, SnapKV (Yuan et al., ICLR 2025, arXiv:2404.14469): during prefill the last snap_obs_window key rows attend to all prefix positions via softmax; the pooled scores select the top-snap_budget tokens to keep. All others are dropped. Decode tokens are never evicted.

config = KVCacheConfig(
    method="snapkv",
    head_dim=128,
    snap_budget=512,        # hard cap on prefill tokens retained
    snap_obs_window=32,     # trailing key rows as proxy queries
    snap_n_sink=4,          # always-keep initial sink positions
)
caches = KVCacheBuilder.for_model(model, config)

Evidence: 18 quantizer tests + 13 cache tests all pass. Eviction ratio > 1 (always smaller than full fp16). Decode accumulation correct. Deterministic.

Honest limitation: the paper uses actual prompt query vectors for the observation window; we substitute key vectors (key-as-query proxy). No max-pool smoothing. No model-level benchmark run yet — cite only after results_snapkv.json is committed from a hardware run.

Inspired by SnapKV (arXiv:2404.14469, ICLR 2025, Yuan et al.) — not a faithful port.

ZipCache-adapted — saliency-adaptive per-token mixed precision — new in 0.18.0

method="zipcache" is the repo's first per-token mixed bit-width cache. Inspired by, not a faithful port of, ZipCache (He et al., NeurIPS 2024, arXiv:2405.14256): the top hi_fraction of tokens by key L2-norm are quantized at hi_bits; the rest at lo_bits. Both groups remain quantized — not fp16. This is the fourth use of the key-norm proxy in the repo, but with a different decision: bit-width routing rather than fp16 protection (KIVI-Sink) or head budgeting (AdaKV-proxy).

config = KVCacheConfig(method="zipcache", head_dim=128,
                       zipcache_hi_bits=4,        # salient tokens get 4-bit
                       zipcache_lo_bits=2,         # rest get 2-bit
                       zipcache_hi_fraction=0.20,  # top 20% by key-norm
                       zipcache_group_size=32,
                       zipcache_quantize_values=True)
caches = KVCacheBuilder.for_model(model, config)

Evidence (tests/cache/test_zipcache_cache.py (11) + tests/quantizers/test_zipcache.py (16), all passing): saliency mask selects top-fraction by key-norm exactly; 4-bit channel quant cosine > 0.995; 2-bit cosine > 0.8; compress/reconstruct preserves shape and fp16 dtype; hi_fraction=1.0 beats hi_fraction=0.0; byte ordering compressed ≤ fp16, mixed-bit ≥ all-lo baseline; effective avg bits in [lo_bits, hi_bits]; values-off passthrough; decode accumulation; determinism; for_model config propagation.

Honest limitation: the proxy (key L2-norm) is weaker than true attention scores. No model-level benchmark yet — benchmark_scripts/benchmark_zipcache.py is an offline-synthetic harness.

SpectralQuant — new in 0.6.0

SpectralQuant implements "3% Is All You Need: Breaking TurboQuant's Compression Limit via Spectral Structure". The key insight: KV cache keys concentrate ~96% of their variance in just 3–4% of dimensions universally across all transformer architectures. SpectralQuant exploits this by rotating keys into their eigenvector basis before quantization — no more wasting bits on noise dimensions.

Three changes over TurboQuant:

1. Eigenvector rotation instead of random Hadamard — aligns signal dimensions first
2. Separate codebooks for signal dims (d_s ≈ 4) and noise dims (d − d_s)
3. No QJL on noise dims — applying QJL there injects variance without reducing bias, hurting quality

from mlx_lm import load
from veloxquant_mlx.spectral import calibrate_spectral_rotation
from veloxquant_mlx.cache.spectral_cache import SpectralQuantKVCache
from veloxquant_mlx.cache.base import KVCacheConfig

model, tokenizer = load("mlx-community/Llama-3.1-8B-Instruct-4bit")

# One-time calibration (~5s on 512 tokens)
import mlx.core as mx
tokens = mx.array(tokenizer.encode(calibration_text)[:512])[None]
rotations = calibrate_spectral_rotation(model, tokens, model_name="llama31_8b")

# Build one calibrated cache per layer
import mlx_lm
cfg = KVCacheConfig(method="spectral", head_dim=128, bit_width_inlier=3)
caches = [SpectralQuantKVCache(cfg) for _ in range(model.args.num_hidden_layers)]
for i, cache in enumerate(caches):
    if i in rotations:
        cache.calibrate(rotations[i])

response = mlx_lm.generate(model, tokenizer,
    prompt="Explain the transformer architecture.",
    max_tokens=500,
)

Results on real models (3-bit, d_s=auto-calibrated):

Model	SpectralQuant noQJL	TurboQuant 3-bit	Δ cosim	SQ ratio
Qwen2.5-0.5B	0.9072	0.8329	+7.4pp	5.33×
Gemma 4 4B	0.8625	0.7581	+10.4pp	5.33×

Calibration required — a one-time ~5–30s pass over 512 representative tokens. Save and reuse with save_rotations / load_cached_rotations. Run python scripts/run_spectral_quant_eval.py --model <name> to generate all benchmark figures.

---

RateQuant — per-layer mixed precision

RateQuant (arxiv:2605.06675) allocates more bits to high-sensitivity layers and fewer to low-sensitivity ones via Theorem 2 reverse-waterfilling, with the average held at a user-chosen target. A sensitivity ratio above ~2× indicates measurable gains over uniform allocation.

Model	Sensitivity ratio	Allocation	Result
Falcon3-7B (28 layers, head_dim=256)	6.48×	14 × b=2, 14 × b=1	100% fp16 at 5.22× compression
Gemma3-4B (34 layers, head_dim=256)	14.39×	3 × b=3, 11 × b=2, 20 × b=1	91% fp16 at 5.22× compression

What's not yet implemented from the paper: per-head allocation, gradient-based sensitivity, K/V separation. Per-layer already captures most of the benefit at ≥1.5 bits.

---

VecInfer — 16× product VQ

VecInfer (arxiv:2510.06175) (Yao et al. 2025) applies a dual transform to keys before product VQ encoding:

1. Smooth scaling — per-channel λ = √(max|K|) suppresses outlier magnitudes
2. Walsh-Hadamard rotation — spreads energy uniformly across all dims
3. K-means product VQ — encode sub-vectors against a calibrated codebook

The inverse transform is absorbed into queries so q @ K.T is preserved exactly. At 1 bit/elem a 128-dim key becomes 16 bytes instead of 256 — 16× compression.

Standout result: Qwen2.5-7B VecInfer-1bit exceeds fp16 throughput at 16× compression, likely due to its strong GQA ratio (28q/4kv heads).

---

Metal kernels — new in 0.5.1

The VecInfer quantize_vq hot path is now a 30-line Metal Shading Language shader, JIT-compiled by mx.fast.metal_kernel on first use. Same Python API — no changes required.

_{Benchmarked on Apple Silicon GPU. Left: quantize latency. Center: speedup factor. Right: peak memory.}

Metric	Pure MLX	Metal kernel	Delta
Quantize latency (S=8192)	228 ms	15.6 ms	14.7× faster
Peak memory (Falcon3-7B shape)	729 MB	12 MB	98% reduction
API change required	—	None	use_metal_kernels=None auto-detects

Why the memory win: the [N, n_centroids, sub_dim] diff tensor is never materialised — the argmin accumulator lives entirely in thread-local registers.

Honest caveat: the kernel pays a ~50–200 µs launch overhead per call. On tiny models (SmolLM2-135M, ~60 launches/token) that overhead can exceed the savings. Built for the regime that needs it: 7B+ models at realistic context lengths.

The full 30-line Metal kernel

// One thread per sub-vector. Argmin lives in registers — no diff tensor.
uint vec_idx  = thread_position_in_grid.x;
uint N_total  = x_shape[0];
if (vec_idx >= N_total) { return; }

uint n_centroids = codebook_shape[0];
uint sub_dim     = codebook_shape[1];
uint x_base      = vec_idx * sub_dim;

float best_dist = INFINITY;
uint  best_idx  = 0;

for (uint c = 0; c < n_centroids; ++c) {
    uint  cb_base = c * sub_dim;
    float dist    = 0.0f;
    for (uint i = 0; i < sub_dim; ++i) {
        float d = float(x[x_base + i]) - float(codebook[cb_base + i]);
        dist += d * d;
    }
    if (dist < best_dist) { best_dist = dist; best_idx = c; }
}

out[vec_idx] = best_idx;

Read the full writeup: blogs/metal-kernels.md

---

Benchmark results

10-model comparative study — VecInfer vs RVQ (v0.5.0)

_{End-to-end mlx_lm.generate · 200-token prompt · 120-token generation · Apple M-series unified memory}

Compression ratio:

Model	RVQ-1bit	VecInfer-1bit
SmolLM2-135M	7.1×	16×
Llama-3.2-1B	7.1×	16×
Llama-3.2-3B	7.5×	16×
Llama-3.1-8B	7.5×	16×
Mistral-7B	7.5×	16×
Qwen2.5-7B	7.5×	16×
Qwen3-8B	7.5×	16×
Phi-4	7.5×	16×
Falcon3-7B	7.8×	16×
gemma-3-4b	7.8×	16×

Throughput (tok/s):

Model	fp16	RVQ-1bit	VecInfer-1bit
SmolLM2-135M	250.4	188.5	175.8
Llama-3.2-1B	105.4	104.3	91.2
Llama-3.2-3B	47.6	46.2	40.2
Llama-3.1-8B	20.5	20.6	19.6
Mistral-7B	23.6	22.8	9.8
Qwen2.5-7B	21.0	20.7	21.5 ⬆ exceeds fp16 at 16×
Qwen3-8B	20.3	19.6	2.4
Phi-4	10.4	8.1	4.0
Falcon3-7B	17.3	21.7	17.0
gemma-3-4b	26.0	24.2	22.6

RVQ-1bit is the safe default — within 5% of fp16 on most 7–8B models with zero calibration. VecInfer-1bit wins on memory (always 16×) and throughput on strong-GQA models (Qwen2.5, Gemma).

Throughput optimisation journey (v0.3.0)

Four sequential changes to lift quantized throughput to fp16 parity:

Stage	Mistral-7B RVQ-2bit	Qwen3-4B RVQ-2bit
0. Original (per-head Python loop)	17.7 tok/s	24.8 tok/s
1. Batch heads (B,H,S,D) → (B·H·S,D)	21.5 tok/s	34.0 tok/s
2. Hadamard rotation by default	20.0 tok/s	—
3. Boundary-sum quantize (replaces argmin)	22.4 tok/s	—
4. Drop redundant fp32↔fp16 casts	22.3 tok/s	36.0 tok/s

Full writeup: OPTIMIZATION_FINDINGS.md

RateQuant V2 mixed-precision results (v0.3.5)

Per-layer allocation at target b̄=1.5, measured on Apple M4 24 GB.

Model	fp16	RVQ-1bit	RVQ + RateQuant V2	Sens. ratio
Falcon3-7B	22.9	23.1 (101%)	22.8 (100%) at 5.22×	6.48×
Gemma3-4B	39.8	37.8 (95%)	36.3 (91%) at 5.22×	14.39×

Source figures: figures/2026-05-16/

RVQ 1-bit 8-model sweep (v0.3.4)

All on Apple M4 MacBook 16/24 GB. Prompt: 200-token explanation of relativity.

Model	fp16 tok/s	RVQ-1bit tok/s	vs fp16
Mistral-7B v0.3	23.3	22.2	95%
Falcon3-7B	24.0	23.1	96%
Phi-4	11.9	11.8	99%
Qwen3-4B	40.2	34.3	85%
Qwen3-8B	20.5	21.1	103%
Llama-3.1-8B	22.0	21.5	98%
Gemma3-4B	32.5	30.5	94%

Source figures: figures/outlier_token_ratequant/

---

Algorithm guide

Method	Bits/dim	Compression	Quality (cosine)	Calibration	Best for
turboquant_mse	b	~9× @ 2b	0.86 @ 3b	None	Lowest overhead at 3–4 bit
turboquant_prod	b	~9× @ 2b	0.95 @ 4b	None	Unbiased IP estimator at 3–4 bit
turboquant_rvq @ b=1	2	7.5×	0.92	None	Default — full output on all 12 tested models
turboquant_rvq @ b=2	4	3.9×	0.98	None	2-bit with near-lossless quality
turboquant_rvq + RateQuant	1.5 avg	5.2×	≈0.96	1.6s	Heterogeneous layer sensitivity
vecinfer @ 1-bit	1	16×	model-dependent	Codebook	Max compression, strong-GQA models
spectral @ b=3	3	5.33×	0.91 (Qwen2.5)	~5s once	Best quality-per-bit, any model
comm_vq	1 (uint8 idx)	64× keys	RoPE-exact	EM training	RoPE-compatible VQ, ICML 2025
rabitq keys + MSE-b4 vals	1 + 4	6× full KV	approx	IVF fit	Max context length, same RAM
polar	b×levels	varies	medium	None	Geometric key distributions
qjl	1	~16×	0.62	None	Ranking-only retrieval, extreme compression

Quick decision:

• No calibration, best default → turboquant_rvq b=1
• Max compression, Qwen2.5/Gemma → vecinfer 1-bit
• Best quality at moderate compression → spectral b=3 (requires ~5s calibration)
• Heterogeneous layers (sens. ratio >2×) → RateQuant on top of RVQ
• 2-bit, near-lossless → turboquant_rvq b=2
• Max context length, fixed RAM → rabitq keys + MSE-b4 values (6× full KV)
• RoPE-compatible exact VQ → comm_vq (ICML 2025, 64× key compression)

---

What's inside

Module	Purpose
veloxquant_mlx/spectral/spectral_quant	SpectralQuantizer — eigenvector rotation + signal/noise codebooks, b=3
veloxquant_mlx/spectral/calibrate	calibrate_spectral_rotation, calibrate_from_vectors, on-disk rotation cache
veloxquant_mlx/spectral/bit_allocator	water_fill_bits — water-filling bit allocation per eigenvalue
veloxquant_mlx/spectral/participation_ratio	compute_participation_ratio, compute_spectral_gap
veloxquant_mlx/quantizers/rabitq	RaBitQQuantizer — IVF + randomised Hadamard + 1-bit sign packing + Metal Hamming search
veloxquant_mlx/quantizers/comm_vq	CommVQQuantizer — RoPE-commutative residual VQ, commutativity projection in EM M-step
veloxquant_mlx/metal/_rabitq	rabitq_hamming_score — Metal XOR+popcount Hamming distance kernel
veloxquant_mlx/metal/_comm_vq	comm_vq_decode_metal — fused centroid gather + RoPE Metal kernel
veloxquant_mlx/quantizers/turboquant_rvq	Two-pass scalar RVQ — Gaussian + Laplacian codebooks, b=1/2/3+
veloxquant_mlx/quantizers/turboquant_prod	Rotation + Lloyd-Max + QJL residual (b-1 + 1 bits)
veloxquant_mlx/quantizers/turboquant_mse	Rotation + Lloyd-Max, no residual correction
veloxquant_mlx/quantizers/polarquant	Recursive polar coordinate decomposition
veloxquant_mlx/quantizers/qjl	Pure 1-bit Johnson-Lindenstrauss sign sketch
veloxquant_mlx/cache/vecinfer_cache	VecInferKVCache — smooth + Hadamard + product VQ
veloxquant_mlx/cache/turboquant_rvq_cache	TurboQuantRVQKVCache — mlx_lm-compatible wrapper
veloxquant_mlx/allocators/vecinfer	calibrate_smooth_factors, train_codebook, quantize_vq
veloxquant_mlx/allocators	allocate_bits_ratequant, calibrate_layer_sensitivities
veloxquant_mlx/metal	Hand-written Metal MSL kernels, JIT via mx.fast.metal_kernel
veloxquant_mlx/preconditioners	RotationPreconditioner (QR), HadamardPreconditioner
veloxquant_mlx/observers	DistortionObserver, LatencyObserver, MemoryObserver, KeyNormObserver
veloxquant_mlx/codebooks	ScalarCodebook, Lloyd-Max strategies, AdaptiveScalarCodebook
veloxquant_mlx/dsa/bit_pack	Sub-byte index packing
veloxquant_mlx/outlier	Two-stream cache for high-variance channels
veloxquant_mlx/weight	QuantizedLinear for model weight quantization

---

Architecture

Pipeline diagrams & design patterns

TurboQuantRVQ pipeline:

x (fp16, batch × d)
     │
Rotate (Π)
     │
Stage-1 quantize  (Gaussian Lloyd-Max, b bits)  →  idx₁
     │
Compute residual  r₁ = y − ŷ₁
     │
Stage-2 quantize  (Laplacian Lloyd-Max, b bits) →  idx₂
     │
EncodedVector(idx₁, idx₂)
     │
Decode: ŷ = ŷ₁ + ŷ₂  →  unrotate

VecInfer pipeline:

x (fp16, B × H × S × D)
     │
Smooth scale  (λᵢ = √max|Kᵢ|, per channel)
     │
Walsh-Hadamard rotation  O(d log d)
     │
K-means product VQ  (sub-vectors against codebook)
     │
Packed indices  →  16× smaller than fp16 keys

Design patterns used (10): Abstract Base Classes, Factory, Chain of Responsibility, Builder, Strategy, Registry + Plugin, Composite, Observer, DAO, Custom DSA (RingBuffer, MaxHeap, BitPackBuffer, VoronoiTree).

---

CLI

# Precompute rotation matrices, JL matrices, codebooks
python -m veloxquant_mlx precompute \
    --head_dim 128 --bits 1 2 3 4 --jl_dim 128 --seed 42 \
    --output_dir ./artifacts/

# Synthetic benchmark — single config
python -m veloxquant_mlx benchmark \
    --method turboquant_rvq --head_dim 128 --bits 2 --seq_len 1000

# End-to-end model benchmarks
python benchmark_scripts/benchmark_vecinfer.py   # VecInfer 10-model sweep
python benchmark_scripts/run_outlier_ratequant.py # RateQuant mixed-precision

Load precomputed artifacts to skip re-computation at runtime:

from veloxquant_mlx.artifacts import NpyArtifactStore

cache = (KVCacheBuilder()
    .with_method("turboquant_rvq")
    .with_head_dim(128).with_bit_width(inlier=2)
    .with_artifact_store(NpyArtifactStore("./artifacts/"))
    .build())

---

Development

# Full test suite (212 tests, includes 7 Metal parity tests)
pytest veloxquant_mlx/tests/ -v

# 2-bit improvement validation — fast synthetic run
python test_2bit_improvements.py

# Generate optimization-journey figure
python scripts/plot_optimization_journey.py

Contributions welcome — please open an issue first for anything beyond a small bugfix. See CHANGELOG.md for release history.

---

Blog posts

All blog posts live in the blogs/ directory and are published at https://veloxquant-mlx.netlify.app/docs/blog/.

File	Description	Live
blogs/overview.md	High-level overview of VeloxQuant-MLX and its goals	↗
blogs/10-model-study.md	End-to-end benchmark study across 10 production models	↗
blogs/hands-on.md	Hands-on tutorial: compressing your first model	↗
blogs/kivi.md	Deep dive into the KIVI asymmetric quantization baseline	↗
blogs/metal-kernels.md	How the Metal compute kernel cuts quantize latency 13×	↗
blogs/results.md	Detailed benchmark results and analysis	↗
blogs/tensorops-research.md	TensorOps research notes and findings	↗
blogs/turboquant-metal-kernels.md	TurboQuant + Metal kernels: combined writeup	↗

---

References

Papers implemented in this library

• RaBitQ (SIGMOD 2024) — Gao et al., "RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search" — 1-bit randomised Hadamard quantization with formal error guarantees
• Ascend-RaBitQ (2026) — He et al., "Ascend-RaBitQ: Heterogeneous NPU-CPU Acceleration of Billion-Scale Similarity Search with 1-bit Quantization" — heterogeneous pipeline inspiration for key+value joint compression
• CommVQ (ICML 2025) — Apple ML Research, "CommVQ: Commutative Vector Quantization for KV Cache Compression" — RoPE-commutative additive codebook VQ
• SpectralQuant (2026) — "3% Is All You Need: Breaking TurboQuant's Compression Limit via Spectral Structure" — eigenvector PCA rotation + signal/noise codebooks, 5.95× at higher quality than TurboQuant
• TurboQuant (ICLR 2026) — Zandieh et al., "Online Vector Quantization with Near-optimal Distortion Rate"
• RateQuant (2025) — "RateQuant: Mixed-Precision KV Cache Quantization via Rate-Distortion Theory"
• VecInfer (2024) — Yao et al., "Efficient LLM Inference with Low-Bit KV Cache via Outlier-Suppressed Vector Quantization"
• PolarQuant (AISTATS 2026) — "PolarQuant: Quantizing KV Caches with Polar Transformation"
• QJL (2024) — Zandieh et al., "QJL: 1-Bit Quantized JL Transform for KV Cache Quantization"
• PALU (ICLR 2025) — Chang et al., "Palu: Compressing KV-Cache with Low-Rank Projection" — group-head low-rank projection of keys and values (true latent storage)
• CacheGen (SIGCOMM 2024) — Liu et al., "CacheGen: KV Cache Compression and Streaming for Fast Large Language Model Serving" — entropy coding of the quantized KV via token-wise locality
• MiniCache (NeurIPS 2024) — Liu et al., "MiniCache: KV Cache Compression in Depth Dimension for Large Language Models" — cross-layer SLERP merge of adjacent layers' KV directions
• GEAR (arXiv:2403.05527) — Kang et al., "GEAR: An Efficient KV Cache Compression Recipe for Near-Lossless Generative Inference of LLM" — error feedback: low-rank residual + sparse outlier correction over a base quantizer
• ZipCache (NeurIPS 2024) — He et al., "ZipCache: Accurate and Efficient KV Cache Quantization with Salient Token Identification" — per-token mixed bit-width via saliency-adaptive allocation (adapted: key-norm proxy for the attention-score saliency signal)
• SnapKV (ICLR 2025) — Yuan et al., "SnapKV: LLM Knows What You are Looking for Before Generation" — token eviction via prefill observation-window attention scoring (adapted: key-as-query proxy for the prompt query vectors)
• StreamingLLM (ICLR 2024) — Xiao et al., "Efficient Streaming Language Models with Attention Sinks" — structural positional eviction: first N sinks + rolling recency window; constant-memory streaming (adapted: no attention mask adjustment, no RoPE position-ID remapping)
• H2O (ICLR 2024) — Zhang et al., "H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models" — cumulative attention-mass token eviction with sink protection; budget-bounded constant-memory cache (adapted: key-as-query proxy for attention weights, no RoPE remapping)
• TOVA (arXiv:2401.06104) — Oren et al., "Transformers are Multi-State RNNs" — memoryless current-step attention-weight token eviction; budget-bounded constant-memory cache, reactive counterpart to H2O's inertial cumulative scoring (adapted: key-as-query proxy for attention weights, no RoPE remapping)
• PyramidKV (arXiv:2406.02069) — Cai et al., "PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling" — layer-adaptive KV budget (large early, small deep, fixed average) over H2O-style cumulative attention-mass eviction; funnels memory to where attention spreads (adapted: fixed linear budget schedule instead of prefill-entropy allocation, key-as-query proxy, no RoPE remapping)
• SqueezeAttention (arXiv:2404.04793) — Wang et al., "SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget" — 2D (layer × token) budget: measures each layer's attention concentration and reallocates a fixed total budget over H2O-style cumulative attention-mass eviction; the data-driven sibling of PyramidKV (adapted: cosine-dispersion concentration proxy instead of observed attention maps, one-shot re-budget at prefill, key-as-query proxy, no RoPE remapping)

Related work

Quantization:

• KIVI (ICML 2024) — Liu et al., "A Tuning-Free Asymmetric 2-Bit Quantization for KV Cache"
• KVQuant (NeurIPS 2024) — Hooper et al., "Towards 10 Million Context Length LLM Inference with KV Cache Quantization"
• Coupled Quantization (NeurIPS 2024) — Zhang et al., "KV Cache is 1 Bit Per Channel"
• KVTuner (ICML 2025) — Li et al., "Sensitivity-Aware Layer-Wise Mixed-Precision KV Cache Quantization"
• MixKVQ (2024) — Zhang et al., "Query-Aware Mixed-Precision KV Cache Quantization"
• FibQuant (2025) — "Universal Vector Quantization for Random-Access KV-Cache Compression"

Token eviction & sparse attention:

• SnapKV (2024) — Li et al., "LLM Knows What You are Looking for Before Generation"
• PyramidKV (2024) — Cai et al., "Dynamic KV Cache Compression based on Pyramidal Information Funneling"
• SqueezeAttention (2024) — Wang et al., "2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget"
• RocketKV (ICML 2025) — Behnam et al., "Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression"
• MagicPIG (ICLR 2025 Spotlight) — Chen et al., "LSH Sampling for Efficient LLM Generation"

Low-rank & cross-layer:

• xKV (2025) — Chang et al., "Cross-Layer SVD for KV-Cache Compression"
• KVPress (2024) — "KV Cache Compression by Estimating Attention from Future Queries Distribution"

Survey:

• KV Cache Management Survey (2024) — "A Survey on LLM Acceleration based on KV Cache Management"

Framework: Apple MLX

---

License

MIT — see LICENSE.

---

_{Built for Apple Silicon · Engineered for speed · MIT License
Landing page · Issues · Blog: 10-model study · Blog: Metal kernels v1 · Blog: TurboQuant Metal kernels}