2 posts tagged with "benchmarks"

A full 8-model benchmark of residual vector quantization for KV cache compression on Apple M4 16GB. What worked, what didn't, and the one result that surprised me.

Here is the result I did not expect: a 32-billion-parameter model running faster with the KV cache compressed than without.

Qwen2.5-32B at RVQ 2-bit: 4.2 tok/s. The same model at full fp16 precision: 3.7 tok/s. That is a 14% throughput improvement from compression — on a 16GB MacBook. Not from a faster algorithm. From freeing up memory bandwidth that the model weights were competing for.

This is the full story: what we built, how it works, and what the benchmarks across 8 models actually show.

The Memory Problem Every Mac LLM User Knows​

Every token your LLM generates writes two vectors into memory — a Key and a Value — for every attention head, across every layer. This is the KV cache. It grows linearly with context length.

For Mistral 7B (32 layers, 8 KV heads, head dimension 128) at 32K tokens, the KV cache at fp16 is:

2 × 32 layers × 8 heads × 128 dims × 32,000 tokens × 2 bytes = 8.4 GB

On a 16GB Mac, that is 8.4 GB for context alone — before model weights, before the OS, before your browser tab.

The KV cache is also the performance bottleneck. Most 7-8B models on Apple Silicon are memory-bandwidth bound: the chip can do the arithmetic faster than it can load data from unified memory. Every token requires streaming the full KV cache through the memory bus once per layer. Smaller cache = faster generation.

Why 2-Bit Usually Fails​

The obvious approach is to quantize each float16 value (2 bytes) down to 2 bits — an 8x reduction. The problem: with only 4 quantization levels, reconstruction error is large enough to corrupt attention scores.

Cosine similarity between the original and compressed key vector — which directly determines attention quality — drops to 0.69 at single-pass 2-bit. For short generation (< 100 tokens), this is often tolerable. For long reasoning chains, it is not. We measured this:

This is not a mild degradation. For these models, standard 2-bit quantization simply does not work.

The Fix: Two-Pass Residual Quantization​

The key insight: if your first codebook leaves large error, run a second codebook on the error itself.

Stage 1 fits a Gaussian Lloyd-Max codebook on the rotated key vector. The rotation (a randomized Hadamard transform) spreads information uniformly across dimensions, so the distribution after rotation looks Gaussian and the codebook fits it well.

Stage 2 takes the residual — the difference between the original and the stage-1 reconstruction — and fits a Laplacian codebook on that. The Laplacian distribution matches the residual better than a Gaussian: it is the tail of a Gaussian, peaked at zero, with heavier tails. Using the wrong distribution family for stage 2 leaves measurable error on the table.

Result: cosine similarity from 0.69 → 0.98 at 2 bits per dimension. The two-pass approach recovers information the single-pass approach discards.

Storage per vector (d=128):

• RVQ 2-bit: ceil(128 × 2 × 2 / 8) = 64 bytes per key vector vs 256 bytes at fp16 — 4x compression
• RVQ 1-bit (sign quantizer + Laplacian residual): 32 bytes per vector — 8x compression, cosine 0.917

The 1-bit variant uses a two-level sign quantizer as stage 1 ({−0.798, +0.798} — the exact Gaussian Lloyd-Max solution at 1 bit) and a Laplacian codebook on the sign-quantization error for stage 2. It achieves 7.5x KV cache compression while keeping cosine above 0.91.

The Throughput Problem (and How We Solved It)​

Compression reduces memory pressure. But if the quantization compute itself costs more than the bandwidth savings, the net result is slower generation.

Our initial implementation ran at 17.7 tok/s on Mistral 7B with RVQ 2-bit — slower than fp16's 22.1 tok/s. Four changes fixed this, each independently measurable:

1. Batch all heads into one MLX call (+22%)

The original code maintained a separate quantizer per attention head and looped over them in Python. For Mistral 7B (8 KV heads, 32 layers), that is 256 small kernel dispatches per token. The fix: reshape (B, H, S, D) → (B·H·S, D) and call one shared quantizer on the entire batch. MLX compiles the operation as a single graph node and dispatches it once.

2. Switch to Hadamard rotation

The rotation preconditioner used QR decomposition — a full d×d matrix multiply (16,384 operations at d=128). mx.hadamard_transform is a Metal-native fused kernel that achieves the same statistical effect in O(d log d) — 896 operations at d=128. The quality is mathematically identical: a Hadamard with a random ±1 diagonal is Haar-distributed, the same family as a random orthogonal matrix.

3. Boundary-sum quantization instead of broadcast-argmin

Codebook lookup was materializing a (batch, d, k) distance tensor and running argmin — three kernel launches: broadcast-subtract, absolute value, argmin. Lloyd-Max boundaries are just midpoints between sorted centroids. The nearest centroid index equals the number of boundaries the value exceeds: one comparison and one sum. Two kernels, no intermediate tensor. Verified bitwise-identical to the old path at 100.00% index match.

4. Remove redundant dtype casts

The update path was casting fp16 → fp32 → fp16 → fp32 → fp16 from accumulated unnecessary coercions. Removing them saved ~5% per call with no effect on output quality.

Net result: 17.7 → 22.3 tok/s on Mistral 7B. The compression now matches fp16 throughput.

Full 8-Model Benchmark​

Hardware: Apple M4 MacBook, 16GB unified memory. Each run: 200 tokens, one fresh Python subprocess per (model, config) to avoid MLX graph compilation bugs. 6 configs: fp16, TQ 2/3/4-bit (single-pass), RVQ 2-bit, RVQ 1-bit.

★ = RVQ configs at 7.5x compression. Bolded token counts in TQ 2-bit column = failure.

Three Findings Worth Highlighting​

1. RVQ 1-bit is the reliability result​

Across all 8 models at all 200 tokens, RVQ 1-bit with 7.5x compression produced complete, coherent output every time. TQ single-pass 2-bit failed catastrophically on 4 of 8 models.

This is the practical takeaway: if you want 2-bit compression and need your model to actually finish its output, residual quantization is not optional. The gap between 0.69 and 0.92 cosine is the gap between broken and working.

2. At 32B scale, compression beats fp16​

Qwen2.5-32B is too large for 16GB without aggressive weight quantization. With 4-bit weights (the Qwen2.5-32B-Instruct-4bit model), it fits — but barely. The model consumes ~17.5 GB including the OS and runtime, leaving almost nothing for KV cache headroom. The chip is running at its memory bandwidth ceiling.

Compressing the KV cache frees bandwidth that the weight loads were competing for. The result is a net gain:

• fp16 KV cache: 3.7 tok/s
• RVQ 2-bit KV cache: 4.2 tok/s (+14%)
• RVQ 1-bit KV cache: 3.9 tok/s (+7%)

The effect is largest at 32B because the bandwidth contention is sharpest. At 7-8B, the same effect exists but is smaller (Qwen3-8B: +3% at RVQ 1-bit). For models that fit comfortably in unified memory, the effect disappears.

3. The Qwen3 4B exception​

Qwen3 4B runs at 40.2 tok/s at fp16 — unusually fast for a model this size, because it is small and the chip can move through layers quickly. At RVQ 1-bit, it drops to 34.3 tok/s (85% of fp16).

This is the trade-off inverting. At small model sizes, fp16 is fast enough that the quantization compute overhead is a larger fraction of total runtime. At large model sizes, the weight loads dominate and the overhead is amortized. The crossover is somewhere around 8B parameters on this hardware.

If you are running a 4B model on a Mac with 32GB, use fp16. If you are running it on 16GB and memory is the constraint, RVQ 2-bit at 35.0 tok/s and 197/200 tokens is the right trade.

The VLM Extension​

We also extended quantization to Qwen2-VL-7B — a vision language model that uses bfloat16 internally (wider exponent range, needed for large-norm image patch tokens). The original code always cast to float16 before normalizing, which discarded the bfloat16 exponent range at exactly the wrong moment for image tokens.

The fix is one variable: kdtype = keys.dtype. Normalize in the actual dtype, cast to float16 only for the codebook lookup. The VLM key-distribution diagnostic across all 28 layers showed that image patch tokens and text tokens have nearly identical distributions after RMSNorm — the quantizer does not need special handling for image tokens. RVQ 2-bit cosine on real Qwen2-VL keys: 0.979 for both image and text tokens across all layers.

What This Means in Practice​

16GB Mac users:

For 7-8B models (Mistral, Llama, Falcon, Phi-4, Qwen3-8B), RVQ 1-bit gives you 7.5x KV cache compression at 94-99% of fp16 throughput. At 32K context, your KV cache drops from ~8 GB to ~1 GB. That is the difference between OOM and comfortable operation.

For Qwen2.5-32B on 16GB, RVQ is what makes the model faster than fp16. The compression is load-bearing.

For 4B models (Qwen3-4B, Gemma3-4B), use fp16 if you have memory. Use RVQ 2-bit if you do not.

Do not use single-pass 2-bit. It fails on Phi-4, Llama, Falcon, and Qwen2.5-32B. If you need 2-bit compression, you need residual quantization.

Using It​

pip install VeloxQuant-MLX

import mlx_lm
from mlx_kv_quant import KVCacheBuilder, KVCacheConfig

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

# 7.5x compression, 95% throughput
config = KVCacheConfig(bits=1, algorithm="turboquant_rvq")
cache = KVCacheBuilder(config).for_model(model).build()

response = mlx_lm.generate(
    model, tokenizer,
    prompt="Explain the difference between RAM and unified memory.",
    kv_cache=cache,
    max_tokens=500,
)
print(response)

For RVQ 2-bit: bits=2. For single-pass 4-bit (safe on all models, less compression): algorithm="turboquant_prod", bits=4.

What the Figures Show​

Each model folder in figures/2026-05-12/ contains 6 panels:

1. Throughput comparison — tok/s across all 6 configs
2. Quality vs compression — cosine similarity curve with RVQ ★ markers
3. Memory at scale — KV cache bytes vs sequence length for each config
4. Attention distortion — per-config distortion bars
5. Output comparison — tokens generated per config (the failure-detection panel)
6. Full report — all panels on one sheet

The output comparison panel (fig5) is the one to look at for reliability. Models where single-pass 2-bit generates 0 or 5 tokens show it plainly. RVQ configs all extend to the full 200-token bar.

Reproducibility​

All runs use subprocess isolation — one fresh Python process per (model, config) — to avoid MLX's graph-reuse bug that causes 2nd+ configs to generate 0 tokens in the same process when the cache type changes. If you run these scripts yourself and see anomalously low token counts on non-fp16 configs, check that you are not running multiple configs in the same Python session.

git clone https://github.com/rajveer43/veloxquant-mlx
cd veloxquant-mlx
pip install -e .
PYTHONPATH=. python3 benchmark_scripts/run_full_reports.py --models mistral7b falcon3_7b llama31_8b

Figures save to figures/2026-05-12/<model>/. Each model takes 10–15 minutes on M4.

VeloxQuant-MLX is MIT licensed. Hardware: Apple M4, 16GB unified memory. MLX 0.24.x, mlx-lm 0.21.x.

How I built VeloxQuant-MLX — KV cache compression for LLMs running on your Mac​

Every time a language model generates a token, it writes two vectors into memory: a Key and a Value. These accumulate silently in what's called the KV cache — one pair per token, per attention head, per layer.

At a short context that's fine. But at 4,000 tokens with a 7B model on a 16GB Mac, the KV cache alone can consume 4–6 GB. At 32K tokens, it can exceed the model weights themselves.

I wanted to fix this for Apple Silicon. So I spent the last few weeks porting three research algorithms from GPU-targeted papers into a native MLX library — and published it to PyPI as VeloxQuant-MLX.

Here's exactly how it works, what I had to rebuild from scratch, and what the benchmarks actually show.

The Core Idea: Why KV Vectors Can Be Compressed​

Key vectors in a transformer aren't arbitrary. After a learned linear projection, they tend to follow a roughly Gaussian distribution — especially after layer normalization. This matters because scalar quantization theory tells us the optimal codebook for a Gaussian is the Lloyd-Max codebook.

The Lloyd-Max algorithm iterates two conditions until convergence:

1. Optimal boundaries — the boundary between two quantization levels sits at the midpoint of their centroids
2. Optimal centroids — each centroid is the conditional mean of the distribution over its Voronoi cell

c_i = ∫[b_{i-1} to b_i] x·f(x)dx  /  ∫[b_{i-1} to b_i] f(x)dx

For a standard Gaussian, this converges to fixed centroids that you can precompute once and reuse for every layer, every model, every token. No per-block scale constants. No calibration data. Just a lookup table.

The problem: raw KV vectors are not standard Gaussian. They have varying magnitudes, correlated dimensions, and outlier channels that dominate the quantization error.

This is what TurboQuant solves.

Stage 1: Random Rotation​

Before quantizing, multiply each Key vector by a random orthogonal matrix Π:

k̃ = k · Πᵀ

Why? Because a random rotation spreads information uniformly across all dimensions. Outlier channels — dimensions with unusually high variance — get diluted into all other dimensions. After rotation, the vector looks much more like an isotropic Gaussian, and the Lloyd-Max codebook fits it well.

The original TurboQuant paper uses QR decomposition on a random Gaussian matrix to generate Π. That's O(d²) per rotation.

For VeloxQuant-MLX, I added a second option: the randomized Hadamard transform:

k̃ = D · diag(r) · WHT(k)

where D is a fixed scaling factor, r is a random Rademacher vector (±1), and WHT is the Walsh-Hadamard Transform. This runs in O(d log d) and maps directly to mx.hadamard_transform — Metal-accelerated on Apple Silicon.

class HadamardPreconditioner(Preconditioner):
    def apply(self, x: mx.array) -> mx.array:
        # x shape: (..., d)
        scaled = x * self._diag_scale  # diagonal randomization
        return mx.hadamard_transform(scaled) / math.sqrt(self._d)

For head dimensions that are powers of 2 (128, 256, 512) — which covers nearly every production model — the Hadamard preconditioner is automatically selected.

Stage 2: Lloyd-Max Scalar Quantization​

After rotation, normalize each vector to unit norm (store the norm as fp16 for later), then quantize each dimension independently using the precomputed codebook:

for each dimension j in rotated key:
    find nearest centroid in codebook
    store index (b-1 bits)

The codebook has 2^(b-1) centroids for a b-bit configuration. At b=4, that's 8 centroids — enough to achieve ~0.95 cosine similarity with the original key vector after dequantization.

The lloyd_max() function in the library solves this numerically using trapezoidal quadrature:

def lloyd_max(pdf_fn, support, n_levels, n_iter=100, tol=1e-8):
    centroids = np.linspace(lo, hi, n_levels)
    for _ in range(n_iter):
        boundaries = midpoints(centroids)
        centroids = conditional_means(pdf_fn, boundaries)  # numerical integration
        if converged: break
    return centroids, boundaries

You run this once at startup. The result is a tiny lookup table (~8 floats for 3-bit) that gets reused for every token in every layer for the entire generation.

Stage 3: QJL Residual Correction​

Even at 4-bit, MSE quantization leaves a residual error:

r = k - k̂_MSE

TurboQuant's key insight is that you can correct for this without storing the full residual. Instead, store only the sign sketch: apply a random Johnson-Lindenstrauss projection matrix S to the residual, then store only the signs:

z = sign(S·r) ∈ {-1, +1}^m

One bit per sketch dimension. For m=128 sketch dimensions, that's 128 bits = 16 bytes per key vector.

At attention time, the corrected inner product estimate becomes:

⟨q,k⟩ ≈ ⟨q,k̂_MSE⟩ + ‖r‖ · (π/2m) · Σ_j z_j · sign(⟨S_j, q⟩)

The residual norm ‖r‖ is stored as fp16 (2 bytes). The signs are bit-packed (1 bit each). The correction is unbiased — on expectation, it perfectly reconstructs the true inner product.

This is what makes TurboQuant different from plain quantization: the QJL stage recovers information that was thrown away, using only 1 bit per sketch dimension.

Memory Layout: What Actually Gets Stored​

For each token, per attention head, VeloxQuant-MLX stores:

Compare to fp16: 128 dimensions × 2 bytes = 256 B for keys alone.

At 4-bit, key compression is 4.27×. Including values, total KV compression is ~1.6×. With value compression (int8 per-token, already implemented), total KV compression reaches ~3.5–4×.

The Bit-Packing Problem​

MLX has no sub-byte dtype. A 3-bit index stored naively in uint8 wastes 5 bits and kills your compression ratio.

The solution is a BitPackBuffer — pack multiple indices into each byte:

3-bit example: pack [2, 5, 1, 7, 3, ...] into bytes
  byte 0: bits 0-2 = index[0], bits 3-5 = index[1], bits 6-7 = index[2][0:1]
  byte 1: bit 0 = index[2][2], bits 1-3 = index[3], ...

For unpack at attend time, the library uses vectorized numpy paths for b ∈ {1, 2, 4} and a loop fallback for b=3:

if b == 2:
    shifts = np.array([0, 2, 4, 6], dtype=np.uint8)
    vals = ((packed[:, :, None] >> shifts) & 0x3).reshape(n, -1)
    return vals[:, :d]

This runs entirely in numpy before handing off to MLX for the attention computation — keeping the Metal graph clean.

Outlier Two-Stream Cache​

Some attention heads have a handful of dimensions with dramatically higher variance than the others. These outlier channels dominate quantization error: if dimension 47 has 10× the typical variance, every other dimension's codebook entry gets contaminated.

VeloxQuant-MLX handles this with an optional two-stream cache:

1. Calibration phase — observe the first N tokens, track per-channel variance
2. Detection — identify the top-K highest-variance channels
3. Split encoding — outlier channels → int8 with per-token scale; inlier channels → TurboQuant as normal

# During append:
k[:, outlier_idx] = 0           # zero out outliers in inlier stream
encode inlier key with TurboQuant
store outlier channels as int8

# During attend:
scores = turboquant_ip(q, encoded_keys)
scores += int8_ip(q[outlier_idx], outlier_cache) * outlier_scales

The result in benchmarks: at K=8 outlier channels, output quality is indistinguishable from fp16 even at 4-bit compression. The cost is slightly higher memory usage (8 × 2 = 16 extra bytes per token per head) and a numpy↔MLX copy per token — which is the main throughput bottleneck in the current implementation.

Weight Quantization: Compressing the Model Too​

One thing missing from the original papers: they only quantize the KV cache. The model weights stay fp16.

VeloxQuant-MLX adds a QuantizedLinear layer — a drop-in nn.Module replacement that applies TurboQuant to the weight matrix itself:

from mlx_kv_quant.weight import quantize_model

model = load_model(...)
quantize_model(model, bits=4, seed=42)
# All nn.Linear layers are now QuantizedLinear

Each row of the weight matrix is normalized, rotated with the same Hadamard preconditioner, and Lloyd-Max quantized. At inference, dequantize on the fly. At 3-bit, this compresses weight matrices by ~5× with minimal accuracy loss.

Benchmark Results Across 7 Models​

I ran the full benchmark suite (fp16, 2-bit, 3-bit, 4-bit) across 7 models on an Apple M4 MacBook (16GB unified memory):

Key finding: 4-bit is near-lossless on 5 of 7 models. 3-bit is near-lossless on 4 of 7. 2-bit (11.6× compression) breaks generation on all models tested — which matches the theoretical SNR floor (~4 dB at 2-bit vs ~10 dB at 4-bit).

The Qwen3 family underperforms because its thinking-mode models generate a <think> token stream that terminates early when attention quality degrades — a more sensitive indicator than output text quality.

What I Had to Rebuild for Apple Silicon​

The TurboQuant paper targets NVIDIA GPUs with CUDA kernels. Porting to MLX required rebuilding several components from scratch:

1. No in-place mutation. MLX uses lazy evaluation — array slices return views, not copies, and array[i] = value silently fails. Every cache write goes through numpy (pre-allocate numpy arrays, write indices into numpy, hand to MLX only for compute).

2. No sub-byte dtypes. The BitPackBuffer class handles arbitrary bit-widths in pure numpy, with vectorized unpack for b ∈ {1,2,4} and a loop path for b=3.

3. Metal-accelerated Hadamard. MLX exposes mx.hadamard_transform natively. Using it instead of the paper's QR rotation reduces rotation cost from O(d²) to O(d log d) — the entire operation runs on the GPU die of the M-series chip.

4. No monkey-patching. The original turboquant-mlx research code patched mlx_lm.scaled_dot_product_attention() globally. VeloxQuant-MLX exposes a standalone cache.attend(q) method — no global side effects, clean integration with any framework.

5. Pluggable architecture. The production library uses a Factory + Strategy + Registry pattern so you can swap quantizers at runtime without changing model code:

cache = (
    KVCacheBuilder()
    .with_method("turboquant_prod")   # swap to "polar", "qjl", "turboquant_mse"
    .with_head_dim(128)
    .with_bit_width(inlier=4)
    .with_jl_dim(128)
    .build()
)

Three Algorithms, One Interface​

VeloxQuant-MLX implements three quantization algorithms behind the same KVCache interface:

TurboQuantProd — Rotation + Lloyd-Max (b-1 bits) + QJL residual (1 bit). Best quality-per-bit for production use.

TurboQuantMSE — Rotation + Lloyd-Max (b bits), no residual correction. Lower memory overhead, slightly lower quality.

PolarQuantizer — Recursive polar coordinate decomposition. Represents each vector as a sequence of angles at 4 levels, plus a final radius. Geometrically motivated — useful for very low-bit regimes.

QJLQuantizer — Pure 1-bit: sign sketch only. Extreme compression (32× vs fp16), loses most content but preserves attention topology. Useful for very long contexts where quality is secondary to fitting in RAM.

Installing and Using It​

pip install VeloxQuant-MLX

Requires Python ≥ 3.11 and an Apple Silicon Mac.

from mlx_kv_quant import KVCacheBuilder
import mlx.core as mx
import numpy as np

cache = (
    KVCacheBuilder()
    .with_method("turboquant_prod")
    .with_head_dim(128)
    .with_bit_width(inlier=4)
    .with_jl_dim(128)
    .with_seed(42)
    .build()
)

rng = np.random.default_rng(0)
for _ in range(1000):
    k = mx.array(rng.standard_normal(128).astype(np.float16))
    v = mx.array(rng.standard_normal(128).astype(np.float16))
    cache.append(k, v)

q = mx.array(rng.standard_normal(128).astype(np.float16))
output = cache.attend(q)
print(f"Memory: {cache.memory_bytes() / 1024:.1f} KB for {len(cache)} tokens")
# Memory: 193.0 KB for 1000 tokens  (vs 500.0 KB fp16)

What's Next​

The main bottleneck right now is the Python-level encode/decode loop. Every token requires a numpy↔MLX transfer per attention head per layer — that's n_layers × n_heads small copies per generation step. A fused Metal kernel would reduce this to near-zero overhead.

Other items on the roadmap:

• Perplexity benchmark on WikiText-2 for quantitative quality measurement
• Value compression — int8 per-token is implemented, needs tuning at longer contexts
• Longer context evaluation — 8K, 32K token sequences where KV memory dominates
• Fused attention kernel — integrate BitPackBuffer unpack + Lloyd-Max decode + attention dot product into a single Metal dispatch

References​

• TurboQuant (ICLR 2026) — Zandieh et al., "Online Vector Quantization with Near-optimal Distortion Rate"
• 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"
• Apple MLX
• VeloxQuant-MLX on PyPI