VeloxQuant-MLX

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

VeloxQuant-MLX compresses the KV cache of any mlx_lm model on Apple Silicon — up to 16× smaller with near-lossless quality, in three lines of code. It ships 41 research-adapted compression methods, from zero-calibration 1-bit quantizers to token-eviction caches to cross-layer merging, plus hand-written Metal kernels that make the hottest path up to 14.7× faster.

Accounting vs. resident memory: compression ratios above are computed from bit-width accounting (bits actually needed to represent the values), not measured process RSS. Most quantization methods still store dequantized fp16 tensors under the hood on the default mlx_lm serving path today, so Activity Monitor / RSS won't drop by the same factor — see #27 for the packed-storage roadmap. Methods that reduce resident memory today are marked 🔻RSS in the method library below (eviction/merging methods, which actually shrink token count); others reduce accounting-only storage at equal fp16 residency.

Why VeloxQuant-MLX:

• 41 methods behind one identical 3-line API — swap method="..." and go
• Metal-accelerated hot paths: 6.9–14.7× faster quantize, 98% less peak memory at the OOM-trigger shape
• Every "-adapted" method documents its honest deviation from the source paper — no silent approximations
• Validated end-to-end on 12 production models: Llama, Mistral, Qwen, Phi, Gemma 3/4, Falcon
• Vision-language models too: patch_vlm_kv_cache wires the same caches into mlx-vlm single-prompt generation (Qwen2-VL, LLaVA, …) — docs

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 relativity simply.", max_tokens=200)

Numbers that matter

Compression ratios below are bit-width accounting, not measured RSS — see the accounting-vs-resident note above and #27. "Peak memory reduction" and "context at 8 GB" rows are Metal-kernel working-set/estimate figures, not steady-state cache RSS under default mlx_lm serving.

Table of contents

1. Installation
2. Quickstart
3. Method library — all 41 methods at a glance
4. Metal kernels
5. Benchmark results
6. What's inside
7. Architecture
8. CLI
9. Development
10. Documentation & blog posts
11. References
12. Support

Installation

pip install VeloxQuant-MLX

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

Full install guide (source install, conda/miniforge, Metal troubleshooting, verifying the install): installation guide.

Quickstart

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

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,
)

More examples, walked through step by step:

• 5-minute quickstart — same example above, plus VecInfer (16×, Metal-accelerated) as a "stronger algorithm" follow-on
• Mixed-precision guide — RateQuant automatic per-layer bit allocation via reverse-waterfilling
• mlx_lm integration guide — wiring compressed caches into any model

Method library

All 41 methods share the same 3-line integration (method="<id>" in KVCacheConfig). Each links to its full page — mechanism, config, evidence, and honest limitations — on the documentation site.

Quick decision:

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

Quantization — compress every token

Category legend: 🧮 accounting_only (stores dequantized fp16 today — bit-width savings are bookkeeping, not RSS; see #27), 🔻RSS eviction/true_latent (drops tokens or stores a genuinely smaller tensor — reduces resident memory today), ⚙️ standalone (does not subclass mlx_lm's KVCache, so it isn't wired into the default serving path via patch_model_kv_cache the same way).

Low-rank & cross-layer — compress across dimensions or depth

Token eviction & merging — drop or merge low-value tokens

Every "-adapted" method is an honest adaptation, not a faithful port — the cache wrapper sees per-layer K/V but not the model's true query/attention maps, so attention-based signals use a key-as-query proxy. Each method's docs page states its specific limitations plainly.

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.}

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.

Full kernel source and how it was built: blogs/metal-kernels.md. Usage, fallback behaviour, and debugging: docs — Metal GPU kernels.

Fused RaBitQ asymmetric pipeline

Two newer kernels form a fully GPU-resident pipeline for an asymmetric-precision cache — 1-bit packed keys scored via XOR+popcount, 4-bit codebook values — a K/V format combination fused attention kernels normally can't express:

• rabitq_encode — rotate + binarize + bit-pack + magnitude in one dispatch. Sign packing uses simd_ballot: each SIMD-group's 32 sign predicates land in a single vote mask, which is exactly 4 bytes of packed output.
• rabitq_fused_attend — scores packed keys, runs an online softmax split across 8 SIMD-groups (flash-decoding style), and accumulates codebook values — one dispatch, no dequantized K or V ever materialized.
• rabitq_pack_values — two 4-bit value indices per byte; the attend kernel reads nibbles directly (auto-detected from the shape), halving value-cache memory and bandwidth with bit-identical outputs.

Measured (Apple M4, D=128 — scripts/metal_rabitq_attend_bench.py, scripts/metal_rabitq_encode_bench.py):

Honest caveat: with unpacked (byte-per-index) values the fused attend loses at short contexts (0.65× at S_kv=512) — nibble-packing halves value bandwidth and flips that to a small win. Parity vs numpy references is covered by 63 dedicated tests (test_rabitq_attend.py, test_rabitq_encode.py, test_rabitq_values.py), including an end-to-end encode→attend test and bit-exact packed-vs-unpacked equality.

For large S_q — the multi-turn VLM case, where a new turn attends over a long compressed image-token history — rabitq_prefill_attend is the matmul-shaped companion: both Q·K̂ᵀ and W·V̂ run on 8×8 simdgroup_matrix tiles, with keys sign-decoded and values nibble-decoded inside the tile loop. It scores exact dots rather than the Hamming estimate, and is cross-attention only (no causal mask).

Fused group-affine (KIVI-style) attention — new in 0.42.0

scalar_fused_decode_attend is the scalar/group-quant analogue of the codebook fused attends above — it serves the KIVI / SKVQ / Kitty / group-quant family, where K/V are uint8 codes plus a per-group (scale, zero) pair instead of a codebook.

The pure-MLX path reconstructs code * scale + zero into a full fp16 tensor, then calls scaled_dot_product_attention — a dequantize → DRAM → SDPA round-trip paid every decode step. This kernel reconstructs x_hat in-register inside a FlashAttention-style online softmax, so no dequantized K_hat/V_hat ever reaches DRAM. The win compounds with context: the fp16 K_hat grows linearly with S_kv while the packed codes stay 16/b times smaller.

Measured (Apple M4 10-core GPU, B=1 H=32 D=128 b=2 g=32 S_q=1) vs. dequantize → MLX SDPA:

The kv axis is split flash-decoding style across nsg SIMD-groups so single-query decode shapes still fill the GPU (nsg=8 tuned on M4), and one compiled kernel serves any (S_kv, D, g). Parity max abs error is 1.2e-4 — the fp32 softmax accumulation makes it more accurate than the fp16 baseline it replaces (test_scalar_attend.py).

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:

Throughput (tok/s):

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).

Historical benchmark snapshots (throughput optimisation journey, RateQuant V2, 8-model RVQ sweep) and full methodology: BENCHMARK_RESULTS.md.

What's inside

Full module reference and API docs: docs — API reference.

Architecture

VeloxQuant-MLX pipelines each quantizer as rotate → quantize (± residual) → pack, built via a Builder/Factory/Strategy layering so every method shares the same KVCacheConfig → KVCacheBuilder → mlx_lm-compatible cache path. Ten design patterns are used throughout (Abstract Base Classes, Factory, Chain of Responsibility, Builder, Strategy, Registry + Plugin, Composite, Observer, DAO, and custom data structures like RingBuffer/MaxHeap/BitPackBuffer/VoronoiTree).

Full pipeline diagrams (TurboQuantRVQ, VecInfer) and design-pattern breakdown: docs — Core concepts.

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

# Which method should I use on my Mac? (new in 0.42.0)
python -m veloxquant_mlx recommend \
    --chip M4 --ram-gb 16 --model-class 7B --goal everyday

The recommender is accounting-aware — it reports the key compression ratio and tells you when resident RAM savings are unlikely, rather than quoting a ratio that won't show up in RSS:

method=turboquant_rvq
knobs={'bit_width_inlier': 1, 'seed': 42}
key_accounting_ratio≈7.5x
resident_savings_likely=False
kv_fp16_mb≈512.0  kv_compressed_mb_est≈68.27
rationale: Zero-calibration default. Key accounting ~7.5x at head_dim=128.
           Default path dequantizes into parent fp16 cache.
warnings:
  - Tight RAM with a mid/large model: consider goal=max_context (rabitq)
    or goal=constant_memory (eviction) for long prompts.

Goals: everyday, max_key_accounting, max_context, best_quality, constant_memory. Add --json for machine-readable output, or --seq-len / --n-layers / --n-kv-heads / --head-dim to match a specific model. Also available in the browser via the Compression Lab.

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 (includes 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 CONTRIBUTING.md for guidelines and CHANGELOG.md for release history.

Documentation & blog posts

Full docs, including per-method pages, guides, and API reference: https://veloxquant-mlx.netlify.app/

Deep-dive writeups live in blogs/ and are also published on the docs site:

References

41 methods, each adapted from a published paper with documented deviations (39 from a verified peer-reviewed venue; 2, NestedKV-adapted and AMC-adapted, from unpublished preprints as one-time, stated exceptions — see CITATIONS.md) — full bibliography (implemented methods, related work, and survey papers): CITATIONS.md.

Headline references: TurboQuant (ICLR 2026), VecInfer (2024), RaBitQ (SIGMOD 2024), CommVQ (ICML 2025), KVzip (NeurIPS 2025), KVTC (ICLR 2026), CurDKV (NeurIPS 2025), NestedKV (preprint, arXiv:2605.26678), AMC (preprint, arXiv:2607.10109), A2ATS (ACL 2025 Findings). Built on Apple MLX.

Support

VeloxQuant-MLX is free, MIT-licensed, and built nights-and-weekends — if it saves your Mac some memory (or you just want to see the 42nd method land), you can buy me a chai ☕ or tip on Ko-fi 💜. Stars, issues, and PRs are equally appreciated.

License

MIT — see LICENSE.