VeloxQuant-MLX

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

VeloxQuant-MLX shrinks the KV cache of any mlx_lm model on Apple Silicon so you can run longer contexts or bigger models in the same RAM — up to 16× smaller with near-lossless quality, in three lines of code. Under the hood it's 41 compression methods (each adapted from a published paper), from zero-calibration 1-bit quantizers to token-eviction caches to cross-layer merging, plus hand-written Metal kernels that speed up the hottest path by up to 14.7×.

If you're running mlx_lm locally and hitting a context-length or memory wall on Apple Silicon, this swaps in a compressed cache with no model changes. Actively developed, with a full test suite gating every release (see badges above).

Accounting vs. resident memory: the compression ratios above are the theoretical byte count (bit-width accounting) — not what Activity Monitor will show you. Most quantization methods still store full fp16 tensors under the hood on the default mlx_lm serving path today, so process RSS won't drop by the same factor yet — see #27 for the packed-storage roadmap. Methods that do shrink resident memory today (eviction/merging, which actually drop tokens) are marked 🔻RSS in the method library below; the rest reduce accounting-only storage while staying fp16-sized in memory.

Why VeloxQuant-MLX:

• Try any of 41 compression strategies without rewriting your code — they all share one 3-line API, so switching is just changing method="..."
• The hot path runs on hand-written Metal kernels: 6.9–14.7× faster quantize, 98% less peak memory at the shape that used to OOM
• If a method had to cut a corner to work as a drop-in cache instead of a full model rewrite, we say so on its docs page — no silent approximations
• Battle-tested on 12 production models: Llama, Mistral, Qwen, Phi, Gemma 3/4, Falcon
• Works with 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, grouped by family
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

No Python? Start here — the control panel

veloxquant panel

Opens a local web UI at http://127.0.0.1:7860: pick a model and compression method, press Start Server, and point any OpenAI-compatible client (Claude Code, Cursor, the OpenAI SDK) at the URL it gives you.

Under the hood it runs veloxquant serve, which you can also use directly:

veloxquant methods --servable-only        # what can be served
veloxquant serve --model mlx-community/Llama-3.2-1B-Instruct-4bit \
                 --method turboquant_rvq --bits 2 --port 8000

Compression is currently accounting-only — byte counters measure compression fidelity, not runtime memory saved. See #27.

Full guide: docs/control-panel.md

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 drop in the same way — just set method="<id>" in KVCacheConfig. For the full comparison table, a decision tree, and per-model recommendations — mechanism, config, evidence, and honest limitations for every method — see the algorithm overview.

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)
• Hard cap on token count, fixed RAM → h2o or snapkv (eviction, reduces resident memory)

The 41 methods span three families — each links to its full docs page:

• Quantization (21 methods) — compress every token. Default: TurboQuant RVQ turboquant_rvq. Also: VecInfer, SpectralQuant, RateQuant, RaBitQ, QJL, PolarQuant, CommVQ, KIVI / KIVI-Sink, SKVQ-adapted, SVDq, Kitty, KVQuant-NUQ, NSNQuant-adapted, ZipCache-adapted, GEAR, CacheGen, AMC-adapted, A2ATS-adapted.
• Low-rank & cross-layer (6 methods) — compress across dimensions or depth. PALU, XQuant, MiniCache, xKV-adapted, AdaKV-proxy, KVTC-adapted.
• Token eviction & merging (14 methods) — drop or merge low-value tokens; these reduce resident memory today (🔻RSS), not just accounting. SnapKV-adapted, StreamingLLM-adapted, H2O-adapted, TOVA-adapted, PyramidKV-adapted, SqueezeAttention-adapted, ChunkKV-adapted, CaM-adapted, L2Norm-adapted, Q-Filters-adapted, Keyformer-adapted, MorphKV-adapted, KVzip-adapted, CurDKV-adapted, NestedKV-adapted.

Category legend used on the docs site: 🧮 won't shrink your Mac's memory usage today, only the theoretical bit count (tagged accounting_only — still stores full fp16 under the hood; see #27); 🔻RSS actually reduces memory you can see today (tagged eviction/true_latent — drops tokens or stores a genuinely smaller tensor); ⚙️ needs a bit more wiring to use (tagged standalone — doesn't subclass mlx_lm's KVCache, so it isn't plugged into the default serving path the same way). Also worth knowing: every "-adapted" method is an honest adaptation, not a 1:1 port — the cache only sees per-layer K/V, never the model's real query/attention maps, so attention-based signals fall back to a key-as-query proxy. Full per-method compression ratios, categories, and release versions are on the algorithm overview.

Metal kernels — new in 0.5.1

VecInfer's quantize_vq — the slowest step in the pipeline — now runs on the GPU instead of in Python: a 30-line Metal shader, JIT-compiled by mx.fast.metal_kernel the first time you call it. Same Python API, no code changes required to benefit.

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

Why the memory win: nothing extra ever gets written out to memory — the [N, n_centroids, sub_dim] diff tensor that the pure-MLX version has to materialize is skipped entirely, since the argmin accumulator lives in thread-local GPU registers instead. That's the whole 98% peak-memory drop.

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

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 pays a real cost every decode step: it reconstructs code * scale + zero into a full fp16 tensor, writes that to memory, then reads it back for scaled_dot_product_attention — a dequantize → DRAM → SDPA round-trip. This kernel skips the memory round-trip entirely: it reconstructs x_hat directly in GPU registers inside a FlashAttention-style online softmax (a numerically stable way to compute softmax over a stream of values without holding them all in memory at once), so no dequantized K_hat/V_hat ever touches DRAM. The win grows with context length: the fp16 K_hat the old path builds grows linearly with S_kv, while the packed codes this kernel reads directly 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

Every method runs the same three-step pipeline: rotate the K/V tensors into a friendlier basis, quantize them (optionally with a residual pass for extra precision), then pack the bits. That's why swapping method="..." just works — every quantizer plugs into the same KVCacheConfig → KVCacheBuilder → mlx_lm-compatible cache path regardless of what it does internally.

If you're curious how that's wired up: it's built on standard design patterns (Factory, Strategy, Builder, and others) plus a few custom data structures for the lower-level bit-packing work. Full pipeline diagrams (TurboQuantRVQ, VecInfer) and the complete 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:

Beyond compression: cross-model KV transfer

One capability in this repo is not a compression method and is deliberately not in the 41: cross-model KV cache transfer (veloxquant_mlx.transfer). Instead of shrinking one model's cache, it maps a source model's already-prefilled KV into a target model's format, so the receiver can skip prefill when you swap between two models in the same family. Cache size is unchanged; what you save is prefill compute.

It lives in its own subsystem rather than behind method="..." because it needs two models, an offline per-pair fit, and a multi-GB artifact — none of which the single-model cache contract can express. Adapted from Cross-Model KV Cache Transfer (NVIDIA, arXiv:2608.03893); the paper's retention and speedup figures are its own, measured on datacenter-scale pairs, and are not reproduced here. See the docs page for the caveats before relying on it.

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. The cross-model transfer subsystem above is counted separately, as it compresses nothing.

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.