Skip to main content

I Ported a Google Research Paper to Apple Silicon

· 11 min read
Rajveer Rathod
Rajveer Rathod
Author of VeloxQuant-MLX

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:

ComponentBits per dimensionTotal bytes (d=128)
MSE indicesb-1 bits × dceil(128 × 3 / 8) = 48 B
QJL signs1 bit × mceil(128 / 8) = 16 B
Residual normfp162 B
Per-vector normfp162 B
Key total68 B
Values (int8 + scale)8 bits × d + fp16130 B
Grand total198 B

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

Modelfp16 tok/s3-bit quality4-bit quality
Llama 3.2 3B47.2⚠️ Repetition loops✅ Near-lossless
Mistral 7B22.5✅ Near-lossless✅ Near-lossless
Falcon3 7B22.1✅ Near-lossless✅ Near-lossless
Qwen3 4B38.7✅ Near-lossless⚠️ Early stop
Qwen3 8B20.6⚠️ Partial⚠️ Partial
Gemma-419.3✅ Near-lossless✅ Near-lossless
Qwen2.5 32B7.1✅ Near-lossless✅ Near-lossless

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​