4 posts tagged with "quantization"

What happens when you take a thought experiment about burning LLM weights into silicon, try to build the closest real thing in software, and let the machine tell you which parts of the idea survive contact with MLX's actual copy semantics.

The idea started as the kind of thing you think about in the shower: what if Apple burned a frozen LLM's weights into a ROM chip sitting right next to RAM? Read-only. Never re-derived. Never fully materialized in general-purpose memory. A model that's just there, the way a calculator's multiplication table is just there.

You can't ship that. This is a software library, not a silicon fab. But strip the hardware away and look at what's actually being asked for — weights that are read-only, computed once, and shared across processes without each one paying full price — and it stops being science fiction. It starts looking like a real, specific gap in how VeloxQuant-MLX loads models today.

So we built the closest real thing, measured it against the actual machine it would run on, and let two of our own assumptions get overturned by the data along the way. This is the record of that — including the part where a "fix" I proposed made things worse, and the part where I had to walk back a claim about compression that turned out to be mathematically impossible.

What already existed​

VeloxQuant-MLX already had most of the ingredients, just not assembled this way.

QuantizedLinear compresses a weight matrix by normalizing each row, rotating it (Hadamard transform or a QR-derived rotation, depending on the dimension), and mapping it onto a small Lloyd-Max codebook — 2 to 4 bits per weight instead of 16. quantize_model() walks an entire mlx-lm model and replaces every nn.Linear with one of these. It works, and the compression ratios are real.

But the compressed result only ever lives as an mx.array in one process's memory. Every time you load the model — a new server worker, a restarted process, a second experiment running alongside the first — you pay the full cost again: dequantize the source weights, rotate them, run nearest-centroid search against the codebook, for every layer. Nothing from the last time you did this is reused.

That's the gap. Not "we lack compression" — we had that. "We recompute the compression from scratch on every single load."

Two questions that had to be answered before writing any code​

The ROM framing implies something specific: that multiple processes could share one physical copy of the weights, the way the OS page cache lets multiple processes reading the same file share pages in RAM. Before designing a file format around that idea, it needed to actually be true for MLX. So, two direct tests against the library itself rather than its documentation.

Does mx.load() already give us this for free? I saved a 500 MB tensor to .safetensors, then launched two independent processes that both called mx.load() on it, and read each process's physical footprint with vmmap:

PID 15970: Physical footprint 514.3M
PID 15971: Physical footprint 514.3M

Each process paid the full 514 MB independently — 1.03 GB combined for one 500 MB file. vmmap showed no mapped-file region for the safetensors path at all. The loader reads and copies into a private heap allocation. No sharing, no mmap(), nothing.

Fine — what if I build the mmap myself? I mmap'd a raw binary file with np.memmap (confirmed lazy — RSS didn't move) and wrapped it with mx.array():

RSS after np.memmap (lazy, no touch):        1534.5 MB   (unchanged)
RSS after mx.array() wrap (no eval):         2542.3 MB   (+1008 MB — 2× the file size)
RSS after mx.eval(x):                        2542.3 MB   (no further change)

mx.array() copies the entire buffer immediately, before mx.eval() is even called, and the jump is roughly double the source size — there's a staging copy in there somewhere before the data lands in MLX's own unified-memory arena. There is no zero-copy constructor from a buffer in this version of MLX's Python API.

That killed the strongest version of the pitch. Processes on this machine are not going to share physical pages of model weights through anything MLX gives you today. If a future MLX version adds a real zero-copy buffer path, this is the test to rerun. Until then, the honest framing is narrower: skip the recomputation, not the RAM.

Building the thing anyway​

What survives the two negative findings above is still worth having: a file format that holds a model's already-quantized weights, persisted once, loadable without repeating the compression work. Call it a reservoir instead of a ROM — closer to what it actually is.

The first version was almost embarrassingly simple. Serialize each QuantizedLinear layer's compressed indices and per-row norms into a flat, page-aligned binary file — one blob for the 2–4 bit indices, one for the fp32 norms, a small JSON header recording each layer's shape and the seed used to derive its rotation and codebook. On load, reconstruct each layer from the header and drop the persisted indices straight in, skipping the whole dequantize-rotate-quantize pipeline.

I benchmarked it against quantize_model() on Qwen2.5-0.5B-Instruct-4bit — 168 linear layers, a small enough model to iterate on quickly. The baseline took 81.6 seconds. My new reservoir loader took 66.6 seconds.

That's not a win. That's barely a rounding error.

Where the time actually went​

I profiled it instead of guessing.

ncalls  tottime  cumtime  function
   168    0.126   66.587  QuantizedLinear.__init__
    24    2.079   57.442  make_rotation_matrix
    24   54.693   55.147  numpy.linalg.qr
   168    0.001    8.907  CodebookFactory.create
   168    4.016    8.899  lloyd_max

There it was. 55 of the 66.6 seconds were spent inside np.linalg.qr — the routine that derives a rotation matrix for any layer whose input dimension doesn't cleanly fit MLX's Hadamard transform (a dimension like 4864, which is neither a clean power of two nor a multiple of the special constants MLX's fast transform supports). Another 8.9 seconds went to fitting Lloyd-Max codebooks.

The reservoir file looked like it had skipped requantization. What it actually skipped was the nearest-centroid search — the smallest part of the cost. Every layer's rotation matrix and codebook were being silently recomputed from the stored seed, on every load, exactly as expensive as before. The format was a placebo.

The fix was to stop being clever about "everything is derivable from the seed" and just persist the actual rotation matrices and codebooks, then reconstruct each layer by bypassing QuantizedLinear.__init__'s expensive branches entirely — direct field assignment onto a bare module instead of calling a constructor that recomputes things you already have on disk.

That dropped the load time to somewhere between 0.36 and 9.2 seconds, against the same 81.6-second baseline. Even the conservative end of that range is a 9× speedup; the isolated, repeatable measurement (0.36s, matching cProfile's CPU-time reading almost exactly) is closer to 230×. The gap between those two numbers is real and I haven't root-caused it — plausibly Metal shader warm-up, plausibly GPU dispatch contention from running back-to-back with the baseline in the same process. I'm reporting the range rather than picking the number that looks better.

The part where fixing one problem created a worse one​

Feeling good about the speedup, I checked the file size. The source model is 265 MB on disk. My reservoir file was 2.5 gigabytes.

Nine and a half times larger than the model it was supposedly compressing.

I broke down where the bytes went:

Eighty-six percent of the file was rotation matrices. Specifically, the 24 layers that fall back to QR rotation instead of the fast Hadamard path — each one has in_features = 4864, and a 4864 × 4864 matrix at 4 bytes per float is about 95 MB. Twenty-four of those is 2.17 GB, and that's before touching a single actual weight.

The instinct here is "surely you can store that more compactly." I had the same instinct, and I want to walk through why it doesn't work, because the wrong answer is genuinely tempting.

np.linalg.qr has a mode='raw' option that returns the underlying Householder reflectors instead of the assembled orthogonal matrix — sounds exactly like what you'd want, a compact representation of the same rotation. I checked what shape it actually returns:

h, tau = np.linalg.qr(G, mode="raw")
# h.shape == (4864, 4864)   — 180.5 MB, even bigger than the dense matrix
# tau.shape == (4864,)      — 0.037 MB

Still a full d × d array. The reason isn't an API limitation — it's that a Haar-random orthogonal d × d matrix genuinely contains d(d-1)/2 degrees of freedom. That's Θ(d²) information, full stop. There is no encoding, clever or otherwise, that gets a truly random rotation below quadratic storage, because the matrix doesn't have any structure to exploit. It's not compressible the way the weights are compressible — the weights have statistical structure a codebook can exploit; a random rotation, by construction, doesn't.

So the earlier plan — "store the rotation as reflectors, get a smaller file" — was wrong. Not underexplored. Wrong, provably, in about ten minutes of checking.

Making the tradeoff a choice instead of a default​

Once "compress the rotation matrix" was off the table, what was left was a genuine, irreducible tradeoff: you can have a fast load (persist the rotation matrix, pay the disk space) or a small file (don't persist it, pay the QR cost again on every load). Not both, for any layer that needs QR fallback.

The fix was to stop pretending there was a single right answer and expose the choice:

save_reservoir(model, path, persist_rotation=False)  # default: small file
save_reservoir(model, path, persist_rotation=True)  # fast load, large file

With the default off, the reservoir file for the same model comes out to 349.7 MB — 1.3× the source model, essentially just the compressed weights plus some structural overhead — and load time goes back up to about 59 seconds, since QR-fallback layers regenerate their rotation from the stored seed exactly as the original quantize_model() path does. With it on, you're back to sub-second loads and a 2.5 GB file. Hadamard-compatible layers — 144 of the 168 in this model — are unaffected either way, since their rotation is just a d-length sign vector, cheap to store regardless.

Neither setting is wrong. What was wrong was shipping one of them silently as the only option and calling the result "a smaller reservoir."

What didn't get tested, and why​

The original plan called for benchmarking Qwen3-4B-4bit — a more realistic model size, and specifically what the ideation phase had proposed. quantize_model() on that model reliably killed the process with SIGKILL on this 24 GB machine. I checked the peak memory footprint before the kill: 21.2 GB, against roughly 14–15 GB of actually available memory at the time.

This isn't a bug I introduced. It's the same headroom constraint documented elsewhere in this repo (docs/MEMORY_CONSTRAINT_FINDINGS.md) for a 32B model on the same hardware — quantize_model()'s dequantize-rotate-quantize pipeline makes several full-size intermediate copies rather than working in place, and a 4B model's pipeline apparently needs more headroom than this particular machine has free right now. It's a real, separate problem, and it blocked the concurrent-process memory benchmark I'd wanted to run — the one that would have measured whether four processes loading the reservoir simultaneously actually show smaller RSS than four processes each running quantize_model() from scratch. That benchmark still needs to happen, on a model this machine can actually load, or on a machine with more headroom.

What actually shipped​

• save_reservoir() / load_reservoir() / graft_reservoir() in weight/reservoir.py — a flat, page-aligned binary format, with persist_rotation as an explicit, documented tradeoff rather than a hidden default.
• A _fast_quantized_linear() construction path that builds a QuantizedLinear from persisted state without touching __init__'s QR or Lloyd-Max branches.
• Eleven tests, including a deliberately non-Hadamard-compatible layer (in_features=50) to exercise the QR fallback path directly, and bit-exact round-trip checks in both persist_rotation modes — the loaded model's forward pass output is byte-identical to the freshly quantized one, not just "close."
• A results file with the actual numbers, not just the ones that made the feature look good.

What this experiment is actually about​

None of the interesting findings here came from writing code that worked on the first try. They came from measuring something, getting an answer that contradicted an assumption, and following that instead of the original plan:

• The "shared ROM" framing died the moment two processes independently paid full memory cost for the same file — a five-minute test with vmmap, run before any format design.
• The first reservoir format's near-zero speedup came from profiling instead of trusting that "we skip requantization" was true because the code was structured to look like it should be.
• The Householder-reflector idea died in about ten minutes of checking mode='raw''s actual return shape — a claim that sounded plausible enough to write into a planning doc, and would have stayed there if nobody had gone and checked.

The honest version of this project isn't "we built a ROM chip in software." It's: cross-process sharing doesn't work with MLX's current copy semantics, skipping recomputation is a real and large win once you persist the right things, and file size versus load time is a fundamental tradeoff for anything that needs a truly random rotation — not a bug to be engineered away, just a dial to expose honestly instead of hiding.

That's a smaller claim than the one I started with. It's also the one that's actually true.

TL;DR — VeloxQuant-MLX 0.8.0 adds KIVI (Liu, Yuan et al., ICML 2024, arXiv:2402.02750), a faithful re-implementation of the most-cited KV-cache quantization baseline. It's asymmetric (per-channel keys, per-token values), keeps a small fp16 residual window, and is fully deterministic. On an Apple M4, across Llama-3.2-3B, Qwen2.5-7B, and Mistral-7B (4-bit, ~2.2–2.4k-token prompts), KIVI-2bit measured 5.8× key / ~4.0× full-KV compression at roughly fp16 throughput. On Apple Silicon the win is memory, not speed — and we report the throughput as measured, not as hoped. Every number below traces to a committed results.json.

The wall, again​

If you run local LLMs on a Mac, you've met the unified-memory wall. The CPU, GPU, and Neural Engine all share one pool of RAM, and the KV cache — the keys and values the model stores for every past token — grows linearly with context length until it crowds out everything else. Weight quantization (GGUF, GPTQ, AWQ) is an offline, one-time compression of the model's parameters; it does nothing for the cache, which is rebuilt token by token at inference time.

VeloxQuant-MLX exists to compress that cache. It already ships a suite of methods — TurboQuant, RVQ, VecInfer, SpectralQuant, RaBitQ, CommVQ, QJL, PolarQuant, RateQuant. With 0.8.0 we added the one method a reviewer asks about first, and which the library conspicuously lacked: KIVI.

What KIVI is​

KIVI is "A Tuning-Free Asymmetric 2-bit Quantization for KV Cache" (arXiv:2402.02750, ICML 2024). Its central observation is that keys and values want different quantization layouts:

• Keys → per channel. Key distributions have a few high-variance channels. Quantizing along the token axis, one (scale, zero) per channel-group, keeps those channels accurate.
• Values → per token. Value distributions are flatter across channels but vary token to token, so the group runs along the channel axis instead.
• A residual window. The most recently generated tokens dominate attention and are cheap to keep exact, so KIVI holds the last R tokens in fp16 and only quantizes tokens once they age out of that window.

The quantizer itself is plain asymmetric min/max group quantization — no codebook, no calibration, no randomness:

for each group g (a slice along the quantization axis):
    zero  = min(g)
    scale = (max(g) - min(g)) / (2**b - 1)
    q     = round((g - zero) / scale)      # uint code in [0, 2**b-1]
    g_hat = q * scale + zero               # asymmetric dequant

Because there's no k-means and no RNG, KIVI is deterministic — same input, identical output, every run. That matters in this codebase: our vector-quantization methods train codebooks and can vary run to run; KIVI adds none of that.

Why we added it (and what we didn't invent)​

To be clear about credit: KIVI is Liu, Yuan, and colleagues' algorithm. We ported it to MLX. There is no novelty claim here.

The value is comparative. KIVI is the field's reference baseline — nearly every KV-cache paper measures against it. Until 0.8.0, VeloxQuant-MLX couldn't answer "how does your method compare to KIVI?" Now every other method in the library has a recognized point of comparison, in the same framework, on the same Apple-Silicon hardware. (The full reasoning, including why we chose KIVI over KVQuant, GEAR, and ZipCache, is in paper/NEW_METHOD_SURVEY.md: KIVI was the highest-value missing baseline, deterministic, and a clean architectural fit that needs no RoPE or attention-score hooks.)

How it plugs in​

Three lines, and mlx_lm.generate runs unchanged:

import mlx_lm
from veloxquant_mlx import KVCacheConfig, KVCacheBuilder

model, tokenizer = mlx_lm.load("mlx-community/Llama-3.2-3B-Instruct-4bit")

config = KVCacheConfig(
    method="kivi",
    bit_width_inlier=2,  # KIVI's default 2-bit
    kivi_group_size=32,  # min/max group size (paper default)
    residual_length=32,  # recent tokens kept in fp16
)
caches = KVCacheBuilder.for_model(model, config)
model.make_cache = lambda *_a, **_k: caches

response = mlx_lm.generate(
    model,
    tokenizer,
    prompt="Summarize the attention mechanism in three sentences.",
    max_tokens=300,
)

The cache quantizes the aged-out keys/values and immediately dequantizes them, so the downstream scaled-dot-product attention sees standard fp16 tensors — no model surgery, no custom attention path.

Results​

All numbers below come from figures/kivi/results_summary.json (aggregated from the per-model figures/kivi/<model>/results.json). Conditions, identical across runs: Apple M4, 24 GB, 4-bit mlx-community models, group_size=32, residual_length=32, max_tokens=120, prompts of ~2.2–2.4k tokens (a long prompt is required for KIVI to actually exercise the quantized path rather than sit inside the fp16 residual window).

KIVI-2bit, across three models:

Bit-width sweep (Llama-3.2-3B):

Two honest observations from this data:

1. Throughput is flat, not faster. KIVI here runs at 100–105% of fp16 — i.e. it does not slow generation down at these sizes, but it also doesn't speed it up. The published KIVI speedups come from a fused CUDA kernel; that kernel does not port to Metal. On Apple Silicon the deliverable is memory compression (the 4× full-KV figure), and the throughput column is there so you can see it costs you nothing, not so you can claim a speedup.
2. Full-KV compression is lower than key-only, and deliberately so. The full-KV figure includes the fp16 residual window. At residual_length=32 over a 120-token generation, a meaningful slice stays in fp16; that drags the end-to-end ratio below the key-only number. We report both rather than quoting the flattering one.

The implementation is covered by 25 passing tests (veloxquant_mlx/tests/quantizers/test_kivi.py and tests/cache/test_kivi_cache.py), including reconstruction-fidelity bounds, the per-token/per-channel asymmetry, the residual-window behavior, and a determinism check.

Honest limitations​

• Memory win, not a speed win on Metal. As above — the CUDA kernel fusion from the paper isn't available here. If you need raw throughput, this isn't your lever.
• Quality is measured as reconstruction fidelity, not task accuracy. Our tests check cosine similarity / MSE against the fp16 cache on synthetic and real key distributions. We have not run LongBench or a perplexity sweep across configs, so we do not claim "no quality loss" on downstream tasks.
• 2-bit is genuinely lossy. On unit-norm synthetic keys, KIVI-2bit reconstruction cosine sits around 0.93 — which is exactly why KIVI keeps an fp16 residual window. If you push residual_length to 0 you'll feel it. The defaults exist for a reason.
• Single chip, short generations. Everything above is one M4 at ~120 generated tokens. Behavior across other M-series tiers and very long generations isn't characterized yet.

Where KIVI fits​

• Reach for KIVI when you want a simple, deterministic, calibration-free 2-bit baseline — or when you specifically need to compare against the literature's reference point.
• Reach for RVQ when you want stronger compression-per-bit with zero calibration and near-fp16 throughput (its analytical codebooks do better than scalar min/max at low bit-rates).
• Reach for VecInfer when you want the most aggressive key compression (up to 16× key-only) and have ~2 minutes for codebook calibration.

KIVI's job in the suite isn't to win every axis; it's to be the honest yardstick the others are measured against.

Try it​

pip install VeloxQuant-MLX==0.8.0

• Docs: /docs/algorithms/kivi/
• GitHub: https://github.com/rajveer43/VeloxQuant-MLX
• Paper: KIVI, arXiv:2402.02750 (ICML 2024)

What was measured vs. not: all compression, throughput, peak-memory, and token-count figures are from committed figures/kivi/*/results.json on a single Apple M4 (24 GB) at ~120 generated tokens with ~2.2–2.4k-token prompts; correctness is from 25 passing unit tests. We did not measure downstream-task accuracy (e.g. LongBench, perplexity sweeps) — "quality" here means reconstruction fidelity against the fp16 cache.

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

TL;DR: The KV cache is the last major unoptimised bottleneck for local LLM inference on Apple Silicon. llama.cpp, Ollama, and LM Studio don't compress it. I built VeloxQuant-MLX to fix that — 9 quantization algorithms, custom Metal GPU kernels, up to 16× compression, plug into mlx_lm with 3 lines.

The wall every Mac LLM user hits​

You've got a MacBook with an M-series chip. You've downloaded llama.cpp, or Ollama, or LM Studio. You load a 7B model, start a conversation, and for the first few hundred tokens everything feels fast.

Then you push it. A long document. A complex coding task. A multi-turn conversation that's been going for a while. And suddenly — generation slows to a crawl, the fans spin up, or the process just dies.

This isn't a bug in those tools. It's a fundamental memory problem that none of them solve. And it lives in a part of the inference stack that almost nobody talks about: the KV cache.

What the KV cache actually is​

To understand the problem, you need to understand what happens inside a transformer when it generates text.

Every token in your context window requires the model to compute two matrices at every layer — a key and a value. These represent what that token "means" in context. Without caching, generating each new token would require recomputing these for every prior token — quadratic cost that makes long-context generation completely impractical.

So instead, the model stores them. Every time a token is processed, its key and value matrices are written to a cache. Future tokens look them up instead of recomputing them.

This is the KV cache. And it grows with every token you generate.

The math is simple:

memory = num_layers × num_kv_heads × head_dim × seq_len × 2 (K+V) × 2 bytes (fp16)

For Llama-3.1-8B at 8k context: 4.2 GB. At 32k context: 16.8 GB. On a machine with 16 GB of total unified memory.

Why Apple Silicon makes this worse, not better​

On a discrete GPU setup — say, an NVIDIA card with 24 GB VRAM — the GPU has its own dedicated memory pool. The KV cache lives there, separate from your system RAM.

Apple Silicon doesn't work that way. M-series chips use unified memory: the CPU, GPU, and Neural Engine all share the same physical memory pool. Your model weights, your KV cache, your browser tabs, your IDE, your OS — all competing for the same 16 or 24 GB.

A 7B model at 4-bit quantisation takes roughly 4–5 GB. Add OS and background processes (~4 GB), Chrome with a few tabs (~2 GB), your IDE (~1 GB). You're at 11–12 GB before you've generated a single token. On a 16 GB machine, you have 4 GB left for the KV cache — which runs out around 6k context for a 7B model.

On a 24 GB machine it's better, but not solved. During the development of VeloxQuant-MLX, I ran a benchmark sweep across 8 models. Qwen2.5-32B failed every quantization config because the weights alone (~17.5 GB) plus OS and development tools left negative headroom for activations and cache buffers. The memory watchdog had to kill the process before it triggered a kernel panic:

[07:28:09] free=62MB   inactive=1094MB pressure=unknown
[07:28:19] free=63MB   inactive=891MB  pressure=unknown
[07:28:19] CRITICAL: only 954MB available — killing PID 21438

System recovered to 10.2 GB free immediately after the kill, with no GPU stall.

What existing tools do about this: essentially nothing​

This is the part that surprised me most when I started digging.

llama.cpp has done extraordinary work on attention kernel optimisation, weight quantisation, batching, and speculative decoding. But the KV cache is stored at fp16 by default. There is experimental support (--cache-type-k q8_0) but it's not the default, not Metal-optimised, and the quantisation methods are basic scalar quantisation with no adaptation to actual key distributions.

Ollama builds on top of llama.cpp and inherits all of this. The UX is excellent. The KV cache situation is unchanged.

LM Studio adds a polished GUI and nice model management. Same story on the cache.

mlx_lm — Apple's own MLX-based inference library — is the most natural fit for Apple Silicon and does excellent work. But its default KV cache is full fp16. No built-in compression.

All of these tools have optimised everything around the KV cache. None of them have solved the KV cache itself.

The insight: compress the cache, not just the weights​

Most quantisation work in the LLM ecosystem targets weight quantisation — compressing model parameters from fp16 down to 8-bit, 4-bit, or lower. This is what GGUF, GPTQ, AWQ, and most Ollama models use.

Weight quantisation is a one-time operation. You compress once, save, load the compressed version.

KV cache quantisation is different. It happens at inference time, on every token, for every layer. The keys and values are different every generation — you can't pre-compute anything.

The key insight is that while weight distributions are relatively uniform, KV cache distributions are not. Keys tend to follow Gaussian or Laplacian distributions after a rotation transform. This structure can be exploited with much more aggressive compression than weight quantisation — down to 1 bit per dimension with surprisingly low quality loss.

The reason this works at such extreme bit rates: the attention computation only needs an approximation of the inner product q · k. Small errors in the reconstructed key vector translate to small errors in the attention score, which average out across many keys in the softmax.

What I built: VeloxQuant-MLX​

VeloxQuant-MLX is a KV cache compression library for Apple Silicon, built on top of MLX. It implements nine quantisation algorithms — from zero-calibration 1-bit methods to mixed-precision allocators — all backed by custom Metal GPU kernels compiled at runtime.

The algorithms​

TurboQuant RVQ — the recommended default​

Zero calibration. Works on any model immediately. Uses Residual Vector Quantisation with analytical Gaussian and Laplacian codebooks precomputed from distribution theory.

Two residual passes at 1 bit each → 7.5× compression with cosine similarity above 0.97.

import mlx_lm
from veloxquant_mlx import KVCacheBuilder, KVCacheConfig

model, tokenizer = mlx_lm.load("mlx-community/Llama-3.2-3B-Instruct-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="Write a 3000-word analysis of the transformer architecture.",
    max_tokens=3000,
)

VecInfer — 16× with Metal acceleration​

Trades a 2-minute calibration step for 16× compression via Product Vector Quantisation. Per-channel smooth scaling handles outlier dimensions that defeat standard VQ. The hot path runs through a Metal GPU kernel that is 13× faster than equivalent MLX Python ops.

Standout result: Qwen2.5-7B VecInfer-1bit exceeds fp16 throughput at 16× compression, likely due to its strong GQA ratio reducing the number of KV heads.

RateQuant — best accuracy per bit​

Runs a 90-second sensitivity probe across all transformer layers, learns which layers are most sensitive to quantisation noise, then uses reverse-waterfilling to allocate more bits to sensitive layers and fewer to insensitive ones.

At 2.0 average bits, RateQuant achieves 2.7× lower perplexity degradation than uniform 2-bit quantisation at identical memory cost.

SpectralQuant — best quality at long context​

Keys in transformer models concentrate ~96% of their variance in just 3–4% of dimensions. SpectralQuant rotates keys into their eigenvector basis via SVD, then applies separate codebooks to the "signal" dimensions (high variance) and "noise" dimensions (low variance).

At 16k context the rotation trick matters enormously — quantisation errors accumulate over long sequences, and aligning the codec to the actual variance structure cuts that accumulation dramatically.

RaBitQ — maximum context length​

1-bit key compression via randomised Hadamard transform + binary sign packing + IVF clustering. Pairs with 4-bit MSE scalar quantisation on values for 6× full KV compression.

6× more context in the same RAM budget.

CommVQ — RoPE-compatible exact VQ (ICML 2025)​

Standard VQ breaks with RoPE positional encodings because quantize(rotate(x)) ≠ rotate(quantize(x)). CommVQ (Apple ML Research, ICML 2025) trains codebooks on pre-RoPE keys with a commutativity projection in the EM M-step. RoPE is applied exactly at decode time. 64× key compression with exact positional encoding.

The Metal kernels​

All the compression math in the world doesn't help if the quantisation itself is slow. VeloxQuant-MLX compiles nine Metal GPU kernels at runtime using mx.fast.metal_kernel:

The fused SDPA kernel is the most impactful. Without fusion, dequantisation creates a full fp16 key matrix in memory before attention — which defeats much of the point of compression. The fused path keeps the cache compressed all the way until attention scores are computed.

The 30-line Metal kernel that powers VecInfer's 13× speedup:

// One thread per sub-vector. Argmin lives in registers — no diff tensor.
uint vec_idx = thread_position_in_grid.x;
float best_dist = INFINITY;
uint  best_idx  = 0;

for (uint c = 0; c < n_centroids; ++c) {
    float dist = 0.0f;
    for (uint i = 0; i < sub_dim; ++i) {
        float d = float(x[x_base + i]) - float(codebook[cb_base + i]);
        dist += d * d;
    }
    if (dist < best_dist) { best_dist = dist; best_idx = c; }
}
out[vec_idx] = best_idx;

The key insight: the [N, n_centroids, sub_dim] diff tensor is never materialised. The argmin accumulator lives entirely in thread-local registers — 98% peak memory reduction at long context shapes.

Real numbers across 10 models​

End-to-end mlx_lm.generate, 200-token prompt, 120-token generation, Apple M-series:

RVQ-1bit is the safe default — within 5% of fp16 on most 7–8B models with zero calibration. VecInfer-1bit wins on compression (always 16×) and throughput on models with strong Grouped Query Attention ratios.

What this means in practice​

On a 16 GB MacBook M3:

• Before: 7B model maxes out around 6k context before generation slows
• After with RVQ 1-bit: that same cache now fits in ~450 MB, leaving headroom for 50k+ token contexts

On a 24 GB MacBook M3 Pro:

• Before: 13B models are borderline; long conversations risk OOM
• After: 13B at 32k context fits comfortably; the cache that would have consumed 8 GB uses 500 MB

Quick decision guide​

Try it​

pip install VeloxQuant-MLX

Requires macOS on Apple M1 or later · Python 3.11+ · MLX ≥ 0.18

📖 Documentation: https://veloxquant-mlx.netlify.app/docs/ 📦 PyPI: https://pypi.org/project/VeloxQuant-MLX/ ⭐ GitHub: https://github.com/rajveer43/turboquant_mac_implementation

What's next​

• Vision-language model support — visual tokens have fundamentally different KV distributions; heavy outliers require adapted algorithms
• Per-head RateQuant granularity — the paper allocates bits per layer-head group; current implementation is per-layer, leaving ~30% accuracy improvement on the table
• Gradient-based sensitivity for RateQuant — more accurate than activation-based, at the cost of a slightly longer calibration pass

The KV cache is the last major unoptimised bottleneck for local LLM inference on Apple Silicon. llama.cpp, Ollama, and LM Studio have built excellent tools — but none of them have solved this. VeloxQuant-MLX is my attempt.

If you're running models locally on a Mac and hitting memory walls — at any context length, on any model — I'd genuinely like to hear about it.

MIT License · Built for Apple Silicon · Engineered for speed