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: The safe everyday pick: it works out of the box with no setup
           step, and shrinks the key half of the cache by about 7.5x. It
           unpacks each value back to full precision as it is read, so this
           is a size measurement rather than a drop in live memory use.
warnings:
  - RAM is tight for a model this size. For long prompts you will get more
    out of 'Fit the longest conversation' (rabitq), which compresses the
    whole cache, or 'Never grow past a fixed memory limit' (streaming_llm),
    which caps it outright.
  - This method measures smaller but may not free much actual RAM on short
    prompts, because its default path unpacks values back to full precision
    as it reads them. The size figure is real; treat it as a measure of how
    well the data compresses, not as RAM you get back.

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.