14 posts tagged with "kv-cache"

I benchmarked VeloxQuant-MLX's KIVI Metal kernel against Qwen3-8B on an Apple M4 -- and the honest result is that it changes nothing you'd notice, for reasons that were predictable before I ran anything.

There's a specific kind of benchmark result that's more useful than a win: a clean, well-measured null. Not "it didn't work," not "the code is broken" -- but "I measured this carefully, twice, at two context lengths, with the noise floor exposed, and the answer is that nothing moved."

That's what happened when I ran this repo's KIVI KV-cache quantization kernel against Qwen3-8B on an Apple M4. This post is the full record -- including the part where the model I actually wanted to test turned out to be disqualified before I wrote a single line.

The starting question​

An earlier post on this project, A 5.65× Metal Kernel That Made My LLM Exactly 0% Faster, found that a bit-exact, 5.65x-faster-at-the-op-level KIVI kernel produced no measurable end-to-end difference on a 3B model, and explained precisely why: quantization is roughly 1-2% of total runtime, so even a large kernel-level win is invisible once Amdahl's law gets to it.

The obvious follow-up: does that hold on a bigger, newer model? I set out to benchmark the newest capable open-weight model I could find, run the same honest protocol, and report whatever came out -- good, bad, or nothing.

Picking a model turned out to be the interesting part​

"Newest" and "will actually work with this benchmark" are not the same list, and finding that out took longer than running the benchmark itself.

The newest Qwen releases at the time -- Qwen3.5 and Qwen3.6 -- looked like the obvious pick. Recent, well-supported, official MLX conversions available within hours of release. Then I actually read their configs.

Which left Qwen3-8B -- the previous generation, but the newest model that's actually compatible: plain causal decoder, full grouped-query attention at every layer, native mlx_lm support, and a 4.61GB 4-bit footprint that leaves comfortable room for KV cache on a 24GB Mac. Not the frontier release, but the newest one where "benchmark the kernel" is even a coherent sentence.

The newest model isn't always the right model for the question you're asking.

Setting up an honest run​

The earlier KIVI post is explicit about four separate ways its own benchmarks lied -- a lazy-eval sync point that inflated a number 50x, GPU contention from a stray background process that flipped a result by 15x, cache reuse across repeated calls that produced four different "conclusions" from the same kernel, and a 127 MB "memory saving" that was allocator noise. This run was built to avoid repeating any of them:

• fp16 vs KIVI-off vs KIVI-on, end-to-end -- not an isolated op-level microbenchmark.
• Interleaved repeats, not blocked: one full round through all three arms, five times, so thermal drift and scheduling noise land on every arm equally instead of stacking against one.
• Median / min / max reported, not a single best-of-N number, so the noise floor is visible rather than hidden.
• Two prompt lengths -- 231 tokens and 2,238 tokens -- because a short context can sit at the flat part of a kernel's headroom curve and understate what's really there.
• Byte-identical output check between kernel-on and kernel-off, since the two code paths are documented bit-exact and this should hold in practice, not just on paper.
• A contention check before trusting anything -- scanning for other GPU-heavy processes running at the same time.

One more decision worth flagging: KIVI's own default bit-width is 2 bits. This repo's own docstring for the KIVI cache documents that at 2 bits, logit cosine similarity against the fp16 baseline goes negative -- effectively uncorrelated output -- by 64 decode steps on a small model. Benchmarking throughput on a config known to break output quality would produce a number attached to a model that isn't really answering questions anymore. This run used 4 bits instead, where the same measurement shows 0.996 similarity -- safe, and still a real compression config, not a strawman.

The numbers​

Run against mlx-community/Qwen3-8B-4bit on Apple M4 (24 GB), 36 layers, 8 KV heads, head_dim 128:

--- short (prompt_tokens=231) ---
  fp16       tok/s median=17.7  min=14.7  max=17.8  peak=4725MB  key_x=1.00 fullKV_x=1.00
  kivi_off   tok/s median=17.6  min=17.3  max=17.8  peak=4649MB  key_x=3.90 fullKV_x=2.29
  kivi_on    tok/s median=17.6  min=17.3  max=17.8  peak=4752MB  key_x=3.90 fullKV_x=2.29
  kivi_on vs kivi_off identical text: True

--- long (prompt_tokens=2238) ---
  fp16       tok/s median=5.9   min=5.9   max=6.8   peak=5226MB  key_x=1.00 fullKV_x=1.00
  kivi_off   tok/s median=5.9   min=5.7   max=6.3   peak=5355MB  key_x=3.28 fullKV_x=3.05
  kivi_on    tok/s median=5.9   min=5.8   max=6.2   peak=5274MB  key_x=3.28 fullKV_x=3.05
  kivi_on vs kivi_off identical text: True

Why the noise floor matters more than the ratio​

Look at the min-max spread on the unchanged fp16 baseline -- same code, same weights, nothing different between runs: 14.66 to 17.85 tok/s on the short prompt, 5.86 to 6.81 tok/s on the long one. That's an 18% spread and a 16% spread, respectively, from thermal state and scheduling jitter alone. The kernel's entire measured effect -- 0.05% to 1.1% -- is roughly fifteen to three-hundred times smaller than that spread. Reporting "KIVI-on is 1.1% faster" without also reporting "and the baseline itself varies by 16% run to run" would be technically true and practically misleading. This is the same lesson the earlier post's Generation 3 table made explicit: looking for a 1% effect through 16-25% baseline variance is like weighing a signature on a bathroom scale.

Why this was predictable​

None of this should be surprising if you do the arithmetic first, which is the actual point of the earlier post. Quantization overhead in KIVI is a small fraction of total wall-clock time -- the model's matmuls and attention dominate. A kernel that speeds up quantization itself, however dramatically at the op level, is speeding up a sliver of the total.

A 5x speedup on 2% of the work is a 2 x 0.05 = 10% change to nothing measurable once repeat variance is larger than the effect. Amdahl's law doesn't negotiate.

The memory numbers tell a related story. Peak memory didn't drop with KIVI enabled -- if anything it moved a little in both directions, which is allocator noise, not signal. This is expected, not a bug: this repo's KIVI cache quantizes then immediately dequantizes back to fp16 before attention runs, because the downstream attention call needs a standard fp16 tensor. The key x / full-KV x columns above are a hypothetical packed-byte accounting -- "what this would cost if actually stored compressed" -- not a measurement of what's resident in memory. The tensor sitting in GPU memory is fp16 the entire time. Reducing that would need packed uint8 storage plus attention that reads packed codes directly, which is a different, larger project than a fused quant/dequant kernel.

So why ship a kernel that does nothing end-to-end?​

Genuinely fair question, and the earlier post's answer still holds here:

• It's free. Bit-exact, byte-identical output confirmed across ten repeats on this model too. It never makes anything slower.
• It removes a floor, not a ceiling. Quantization is ~1-2% of runtime today. If everything around it gets faster -- better attention kernels, better matmuls -- that slice grows, and fixed per-op costs start to matter more.
• The op-level win is real, even though it's invisible. A kernel can be honestly fast at what it does and still be inconsequential to what the user experiences. Both facts belong in the same report.

What I'd take from this​

The honest headline is the same shape as the one from the 3B run: a correct, bit-exact, measurably-fast-at-the-op-level Metal kernel, invisible end-to-end, for reasons you can derive with arithmetic before running a single generation. Running it again on a bigger, newer-generation model didn't change the conclusion -- it strengthened it, because now it's two models, two very different context lengths, and the same result both times.

Benchmarked on an Apple M4 (10-core GPU, 24GB unified memory) against mlx-community/Qwen3-8B-4bit (36 layers, 8 KV heads, head_dim 128), KIVI bit_width_inlier=4, group_size=32, residual_length=32, 5 interleaved repeats per arm, 120 max tokens. Benchmark script: benchmark_scripts/benchmark_qwen3_8b_kivi_honest.py. See the KIVI algorithm reference for the method itself, and the companion 3B-model post for the original kernel writeup.

Third post in the QFilters sub-series, following the original collapse and the calibration follow-up. Both earlier posts established that qfilters_budget=512 breaks generation on a 2,238-token prompt, and that raising the budget above the prompt's length fixes it. This post asks the question those two left open: what does the space between those two points actually look like?

The calibration follow-up ended with real calibration making the collapse faster and differently-shaped, but not fixed, at the method's default qfilters_budget=512. A natural next question -- natural enough that it's the direct answer to "how do I get coherent output" -- is whether raising the budget helps, and if so, how much is enough. A quick check confirmed the extreme case: a budget of 3,072 against this same 2,238-token prompt means nothing gets evicted at all, and the output is fluent, nearly matching fp16.

That's a single before/after comparison, though -- "broken at 512, fine at 3072" -- and it doesn't say whether quality degrades gradually as the budget shrinks toward the prompt length, or whether something sharper is going on. This post sweeps nine budget values between those two points to find out.

Setup​

Same long prompt (2,238 tokens), same calibrated filter path from the previous post (established as faster and no worse than the fallback at equal budget), same model. Nine qfilters_budget values -- 512, 768, 1024, 1280, 1536, 1792, 2048, 2304, 2560 -- each run for 3 interleaved repeats at 120 max tokens, plus an interleaved fp16 reference. Quality tracked via word overlap against the fp16 reference text (the same metric the earlier posts used), plus throughput, peak memory, and the realized compression ratio at each budget.

The numbers​

What the text actually looks like on either side of the cliff​

budget=1792 (80.1% of prompt retained, all 3 repeats identical):
 the model. the. the. the. the. the. the. the. the. the. the. the. the. the. the...

budget=2048 (91.5% of prompt retained, all 3 repeats identical):
Also, explain why the model weights are not the main constraint.

Answer in English, in a clear and concise manner, using simple terms. Do not use
markdown.

Okay, so I need to explain why the KV ca...

fp16 (reference):
Also, explain in simple terms what the KV cache is, and what it does.

Finally, explain in simple terms what "long-context inference" means.

Please make sure your explanations are clear and conci...

budget=1792 is the same shape of failure as the original post's budget=512 run -- deterministic, fully repeated, not language. budget=2048 is genuinely coherent: on-topic, grammatical, responding sensibly to the prompt's actual instructions, even though it's not byte-identical to fp16. The transition between those two adjacent sweep points is the entire story of this post.

Compression ratio was never the useful signal here​

Look at the compression column: it declines smoothly and predictably as budget rises (4.60x -> 3.07x -> 2.30x -> ... -> 1.00x), exactly as the arithmetic of "keep more tokens, compress less" would predict. There is nothing in that column that would tell you where the cliff is. A dashboard reporting "compression ratio: 1.32x, throughput: 2.81 tok/s" at budget=1792 looks unremarkable sitting next to budget=2048's "compression ratio: 1.15x, throughput: 2.81 tok/s" -- identical throughput, similar compression, and one of them is nonsense text. Compression ratio and throughput both move smoothly through the cliff; only checking the actual output catches it.

Peak memory doesn't move the way the eviction story would suggest, either​

Peak memory rises from 512 through 1792 (5907MB -> 6540MB) despite compression falling over that same range -- more retained tokens should mean less relative savings, but here it also costs more absolute memory, consistent with the earlier calibration post's finding that scoring/bookkeeping overhead is a real, separate cost from the compression itself. It then drops and flattens at 6139MB for 2048 through 2560, once eviction stops actually triggering on this prompt. None of this tracks the coherence cliff either -- it's a separate axis, moving on its own schedule.

What this means for actually using QFilters​

The original post's finding stands: the method's own default (512) is unusable on a prompt anywhere near this length. The calibration follow-up's finding also stands: calibration alone doesn't rescue it at that budget. What this post adds is the concrete, actionable number this whole line of investigation was aimed at:

Four posts into the QFilters line of this series, the shape of the finding has moved from "it's broken" to "it's broken here, fixed there" to, now, "here's exactly where the line is, and it's much closer to 'don't evict much at all' than the method's own 512-token default would suggest."

Benchmarked on an Apple M4 (10-core GPU, 24GB unified memory) against mlx-community/Qwen3-8B-4bit (36 layers, 8 KV heads, head_dim 128), calibrated Q-Filters (see the calibration follow-up), 3 interleaved repeats per budget, 120 max tokens, long prompt only (2,238 tokens). Sweep script: benchmark_scripts/benchmark_qwen3_8b_qfilters_budget_sweep.py. See the original QFilters post and the calibration follow-up for the findings this one builds on.

Direct follow-up to The Output That Stopped Being Output. That post found QFilters producing a fully reproducible generation collapse on a long prompt, using the fallback (uncalibrated) filter path. This one asks the obvious next question: does this repo's real calibration module fix it? Short answer -- it changes what the collapse looks like and makes it substantially faster, but on this exact configuration, it does not produce coherent output.

The QFilters post was explicit about running the uncalibrated fallback path -- key-SVD-estimated filter direction, sign ambiguous, because a KV cache never sees queries. That was flagged as a real limitation, not a hidden one. What made it worth a follow-up is that this repo isn't missing a fix for that limitation: veloxquant_mlx/quantizers/qfilters_calibration.py already implements the paper's actual mechanism (arXiv:2503.02812 §3.2) -- hook each layer's query projection, gather real query activations, take the SVD's top right-singular vector per head, sign-fix it against the paper's Theorem 3.3 (kappa^h > 0), average query-head filters down to KV heads for GQA. It just wasn't used in the first post.

This post uses it.

What calibration actually involved​

calibrate_qwen3_8b_qfilters.py ran Qwen3-8B forward over five ~2,048-token calibration passages (general prose about transformers, memory bandwidth, unified memory, vector quantization, and eviction itself -- topically adjacent to the benchmark prompt but not the same text), captured real query activations per layer via a hook on q_proj, and computed one [8, 128] filter per layer (32 query heads averaged down to 8 KV heads). This is a one-time, pre-deployment pass -- the paper's own cost estimate is 20 samples of length 2048; this run used 5, on the smaller side but the same order of magnitude.

The comparison​

Same long prompt (2,238 tokens) and same qfilters_budget=512 as the original post, five arms instead of three: fp16, fallback_off/fallback_on (the original post's uncalibrated path, Metal kernel off/on), and calibrated_off/calibrated_on (this post's real query-SVD filters, Metal kernel off/on). 5 interleaved repeats per arm.

The actual output​

Every arm was internally deterministic -- 5/5 identical repeats, same as every method tested in this series. Here's what changed and what didn't:

fallback (both kernel arms, all 5 repeats):
===============.========================================================================= is is is is is is is is is is is is is is is is is is is is===========

calibrated (both kernel arms, all 5 repeats):
...........,. the..........,. the. the. and.,..... the.,. the. the. and. the. the. the. and. the. the. the. the. the. the. the. and. the. the. the. the...

Calibration produced a different deterministic collapse, not a working one. fp16-vs-calibrated word overlap is actually lower (0.000) than fp16-vs-fallback (0.014) -- by that specific metric, calibration didn't move the output closer to the uncompressed baseline at all.

What this changes about the original post's finding​

The original QFilters post's headline claim holds: past this method's default budget, on this prompt, output collapses -- deterministically, reproducibly. What this follow-up adds is that the collapse is sensitive to filter quality (it's not simply what eviction always looks like at this budget, since a better filter changed its shape and roughly quadrupled the achievable compression-adjusted throughput) but is not resolved by filter quality alone at this budget. The next honest variable to sweep isn't calibration -- it's the budget itself, or a larger/better-matched calibration corpus, or both. Neither was tested here.

What I'd take from this​

Benchmarked on an Apple M4 (10-core GPU, 24GB unified memory) against mlx-community/Qwen3-8B-4bit (36 layers, 8 KV heads, head_dim 128), QFilters qfilters_budget=512, 5 interleaved repeats, 120 max tokens, long prompt only (2,238 tokens -- the configuration that exceeds budget). Calibration script: benchmark_scripts/calibrate_qwen3_8b_qfilters.py. Comparison script: benchmark_scripts/benchmark_qwen3_8b_qfilters_calibrated.py. See the original QFilters post for the finding this follows up on, and qfilters_calibration.py for the calibration mechanism itself.

Third post in this series, same model and protocol as the KIVI and TurboQuantRVQ posts before it. This time the method doesn't quantize anything -- it evicts. And once the cache runs out of budget, the honest result isn't a percentage, it's forty repeated equals signs.

The first post in this series found a Metal kernel with zero end-to-end effect. The second found a real throughput cost and a memory claim that flipped sign at longer context. Both methods, though, were quantization -- every key and value survives, just approximated. Nothing in either post could produce genuinely broken output, because nothing was ever thrown away.

This post benchmarks QFilters, this repo's query-agnostic eviction method, which works differently on purpose: past a fixed token budget, it drops the lowest-scoring cached tokens outright. Their information doesn't get approximated -- it's gone. That makes a different question testable, one the first two posts structurally couldn't ask: does the cache stay coherent once it starts forgetting things?

Why eviction is a different kind of test​

QFiltersKVCache scores every cached key by its projection onto a frozen per-head direction -- the "Q-Filter" -- and evicts the lowest-scoring tokens once the cache exceeds qfilters_budget (default 512 tokens, including protected leading "sink" positions). This repo's implementation is explicit that it's adapted, not a faithful port of the arXiv:2503.02812 preprint, and offers two ways to get the scoring direction:

• Calibrated filters (query-SVD, the mechanism the paper actually specifies) -- frozen before the first token, correct by construction, path-independent.
• Fallback (filters=None) -- SVD-estimated from the first qfilters_calib_tokens observed keys, which recovers the dominant axis but not its orientation. The module docstring calls this exactly what it is: a real, documented mode, not a stand-in.

No calibrated filters exist for Qwen3-8B in this repo, so this run used the fallback path -- the honest choice given what's actually available, not the paper's strongest configuration. That's flagged here the same way the TurboQuantRVQ post flagged using the method's own uncalibrated default bit-width: report what you actually ran, not the best-case version of it.

VecInfer -- this repo's other codebook-based method with a public Metal toggle -- was considered and set aside for this run. Its docstring is explicit that a random-initialized codebook is "for tests only," and no calibration tooling exists in this repo yet to produce a real one. Benchmarking coherence against a method whose codebook is documented as non-functional would have answered a different, less honest question than the one this post is actually asking.

Setup​

Same protocol as the previous two posts: fp16 vs. kernel-off vs. kernel-on, interleaved (not blocked) repeats, median/min/max reported, two prompt lengths, a GPU-contention check before trusting anything. One addition specific to eviction: two separate comparisons that must not be conflated.

The short prompt (231 tokens) stays under the 512-token budget -- no eviction happens at all. The long prompt (2,238 tokens) blows well past it. That contrast is the point: one prompt length isolates pure per-step scoring overhead, the other tests what happens once the method actually does its job.

The numbers​

Run against mlx-community/Qwen3-8B-4bit on Apple M4 (24 GB), 36 layers, 8 KV heads, head_dim 128, qfilters_budget=512, fallback (uncalibrated) filters, 5 interleaved repeats:

Finding 1: overhead exists even with zero evictions​

At the short prompt, nothing gets evicted -- the cache never reaches its 512-token budget. And yet QFilters is 30% slower than fp16, with peak memory 4.4% higher, not lower. This is pure per-step scoring cost: projecting every cached key onto the filter direction, tracking per-head state, maintaining the SVD estimation buffers for the fallback path -- all of it runs whether or not anything ends up evicted. Unlike KIVI (roughly free) or TurboQuantRVQ (a real but smaller 15% cost at this same prompt length), QFilters' overhead is front-loaded into bookkeeping that has nothing to do with the compression it eventually delivers.

Finding 2: at long context, throughput drops 84% and compression still doesn't shrink peak memory​

Once eviction is actually happening, throughput falls to 16% of fp16's -- a bigger relative cost than either quantization method in this series produced at any context length. Compression is real here (4.60x on the retained K/V, since evicted tokens' bytes are actually gone, not just packed tighter), but peak memory is still 0.78% higher than fp16, not lower. The same shape of result as the TurboQuantRVQ post's memory-claim reversal, but more extreme: even a genuine 4.6x compression ratio on retained tokens doesn't translate into a net memory win, because whatever eviction bookkeeping costs at 2,238 tokens of context outweighs it.

Finding 3: the coherence question has a real, fully reproducible answer​

This is the part that doesn't fit in a percentage. Here is what fp16 generated on the long prompt (first 200 characters, identical across all 5 repeats):

fp16:

Also, explain in simple terms what the KV cache is, and what it does.

Finally, explain in simple terms what "long-context inference" means.

Please make sure your explanations are clear and conci...

Here is what QFilters generated on the same prompt -- identical across all 5 repeats, and identical between the Metal kernel on and off:

qfilters (on and off, all 5 repeats):
===============.========================================================================= is is is is is is is is is is is is is is is is is is is is===========

Worth being precise about what this does and doesn't indict. This is the uncalibrated fallback path -- the docstring is explicit that a key-SVD-estimated direction "recovers the dominant axis but not its orientation," i.e., it can evict in the wrong direction from the start. It is not a claim that QFilters with real calibrated query-SVD filters would fail the same way; that's a different, untested configuration. It's also not a claim about the eviction mechanism being unsound -- a budget of 512 tokens against a prompt that needs most of its 2,238 tokens of context to answer a question about that passage is close to an adversarial setup by construction. What it is: an honest measurement of what actually ships if this cache is deployed at its own default budget, with the calibration path that's actually available in this repo today, against a plausible long-context prompt.

The fp16 baseline here isn't polished either -- it echoes the prompt's own instructions rather than producing a tight two-paragraph explanation, which is a model/prompt-fit issue independent of any cache method. But "echoes the prompt" and "===== is is is is is =====" are not the same category of failure, and the gap between them is the actual finding.

What I'd take from this​

Three methods into this series, three different shapes of honest result: a kernel that changes nothing, a method with a real and growing cost and a memory claim that reverses sign, and a method where the Metal kernel is once again bit-exact but the method itself, run at its own default budget with the calibration path actually available, can turn a long-context prompt into forty characters of =.

Benchmarked on an Apple M4 (10-core GPU, 24GB unified memory) against mlx-community/Qwen3-8B-4bit (36 layers, 8 KV heads, head_dim 128), QFilters qfilters_budget=512 (the method's own default), fallback (uncalibrated) filter path, 5 interleaved repeats per arm, 120 max tokens. Benchmark script: benchmark_scripts/benchmark_qwen3_8b_qfilters_honest.py. See the KIVI/Qwen3-8B post and TurboQuantRVQ/Qwen3-8B post for the first two posts in this series, and qfilters_cache.py for the method's implementation, its calibrated-vs-fallback filter distinction, and its own documented limitations.

Same honest-benchmark protocol as the KIVI/Qwen3-8B post, a different method -- and this time the numbers aren't a null result. TurboQuantRVQ's fused Metal pack kernel is invisible, same as KIVI's. But the method itself costs real throughput that gets worse with context length, and a memory-savings claim already sitting in this repo's own docstring did not replicate at longer context -- it reversed sign.

The previous post in this pair ran the same honest-benchmark protocol against KIVI and found nothing: a bit-exact Metal kernel, byte-identical output, and an end-to-end throughput delta smaller than the run-to-run noise floor. A clean null.

This post runs the identical protocol against TurboQuantRVQ -- this repo's namesake method, and its DEFAULT_SERVE_METHOD -- and gets a different kind of answer: two real findings, not a null, and one of them contradicts a number already committed to this codebase.

Why TurboQuantRVQ, and why it's a different kind of test​

KIVI quantizes keys, then immediately dequantizes them back to fp16 before attention runs. The live tensor is fp16 at all times -- this repo's own kivi_cache.py docstring says so explicitly. That's why KIVI's "compression ratio" is a hypothetical byte-accounting number, not a memory-savings claim, and why testing for a memory win there would have been testing something ruled out by construction.

TurboQuantRVQKVCache is built differently. Keys are stored packed: two bit-packed uint32 residual-vector-quantization index streams plus a shared fp16 norm, matching the same accepted "dequantize-on-fetch" pattern mlx_lm's own native QuantizedKVCache uses. Nothing forces the packed form back to fp16 until a fetch actually needs it. So a peak-memory reduction here is a claim worth testing, not one closed off by the storage design -- and the cache's own module docstring already reports one: -12.8% peak memory vs. fp16, measured on a 1B model at a single 4,002-token prompt.

One measurement, one model, one prompt length. That's exactly the setup the earlier KIVI post's "run the third model" lesson warns about -- a result that looks clean because nothing has yet contradicted it.

The kernel toggle, and one difference from the KIVI setup​

TurboQuantRVQKVCache fuses key quantization and bit-packing into a single Metal kernel (_use_metal_pack) when head_dim is a power of two and fits one threadgroup (head_dim <= 1024). Qwen3-8B's head_dim=128 qualifies, and the flag auto-enables at construction -- there's no public config field for it, unlike KIVI's use_metal_kernels. For a clean A/B this benchmark forces it off on the "off" arm by flipping the same private latch the class itself flips on a kernel-side failure -- not a new code path, just triggering the class's own documented fallback intentionally.

The method's own default bit-width is 2 bits (KVCacheConfig's dataclass default, and turboquant_rvq_cache.py documents no quality caveat at this width -- unlike KIVI, which documents negative logit cosine similarity at 64 decode steps for b=2). This run used the default rather than overriding it, since nothing in the source flagged it as unsafe.

Same protocol as the KIVI post otherwise: fp16 vs. kernel-off vs. kernel-on, interleaved (not blocked) repeats, median/min/max reported, two prompt lengths (231 and 2,238 tokens), byte-identical-output check, GPU-contention check before trusting anything.

The numbers​

Run against mlx-community/Qwen3-8B-4bit on Apple M4 (24 GB), 36 layers, 8 KV heads, head_dim 128, 5 interleaved repeats, bit_width_inlier=2:

--- short (prompt_tokens=231) ---
  fp16       tok/s median=17.0  min=15.8  max=17.4  peak=4725MB  key_x=1.00
  rvq_off    tok/s median=14.4  min=14.2  max=14.6  peak=4637MB  key_x=3.88
  rvq_on     tok/s median=14.6  min=13.9  max=15.3  peak=4648MB  key_x=3.88
  rvq_on vs rvq_off identical text: True

--- long (prompt_tokens=2238) ---
  fp16       tok/s median=5.9   min=5.7   max=6.8   peak=5226MB  key_x=1.00
  rvq_off    tok/s median=4.1   min=3.9   max=4.6   peak=5285MB  key_x=3.88
  rvq_on     tok/s median=4.1   min=3.7   max=4.5   peak=5293MB  key_x=3.88
  rvq_on vs rvq_off identical text: True

That's the part of this run that rhymes with the KIVI post. The two parts that don't follow.

Finding 1: the method itself has a real, growing throughput cost​

Unlike KIVI -- where fp16 vs. quantized was within a percent or two at both prompt lengths -- TurboQuantRVQ is 15% slower than fp16 at the short prompt and 31% slower at the long one. This is not noise: the gap is three to five times larger than fp16's own run-to-run spread at each length, and it moves in a consistent direction as context grows.

Finding 2: the memory claim reversed sign at longer context​

This is the one worth sitting with. turboquant_rvq_cache.py's own docstring reports -12.8% peak memory vs. fp16, from a single measurement on a 1B model at one 4,002-token prompt. Here, on Qwen3-8B:

• Short prompt (231 tok): -1.7% to -1.9% peak memory vs. fp16 -- same direction as the docstring's claim, but an order of magnitude smaller.
• Long prompt (2,238 tok): +1.1% to +1.3% peak memory vs. fp16 -- the opposite direction.

The existing docstring's number isn't wrong -- it's what was measured, on the model and prompt length it was measured on. It just isn't a general claim, and nothing about the original one-data-point report could have told you that from the number alone.

What this changes about how the earlier claim should be read​

The honest fix isn't to argue with the existing docstring's number -- it's to scope it. "TurboQuantRVQ saves memory" was never quite the claim; "TurboQuantRVQ saved 12.8% memory on a 1B model at one 4K-token prompt" was, and that's still true. What this run adds is that the sign of the effect depends on context length, at least somewhere between short and long prompts on an 8B model, and that the throughput cost of getting there is real and grows with context too -- which the original single data point had no way to surface.

What I'd take from this​

Two honest benchmarks, same protocol, same model, two different shapes of result: one method's kernel and its end-to-end footprint were both invisible; another method's kernel is equally invisible, but the method itself has a real, context-dependent cost and a memory story that only a second prompt length was able to falsify.

Benchmarked on an Apple M4 (10-core GPU, 24GB unified memory) against mlx-community/Qwen3-8B-4bit (36 layers, 8 KV heads, head_dim 128), TurboQuantRVQ bit_width_inlier=2 (the method's own default), 5 interleaved repeats per arm, 120 max tokens. Benchmark script: benchmark_scripts/benchmark_qwen3_8b_turboquant_rvq_honest.py. See the companion KIVI/Qwen3-8B post for the null-result half of this pair, and turboquant_rvq_cache.py for the method's implementation and its original single-data-point memory measurement.

Fourth post in this series, same model and protocol as KIVI, TurboQuantRVQ, and QFilters. The first three posts all found Metal kernels that were bit-exact against their pure-MLX fallback, whatever else they cost or saved. This one breaks that pattern in both directions: the kernel is a real, large speedup for once -- and, for the first time, it is not byte-identical to the path it's supposed to match.

Three posts into this series, a pattern had formed: every Metal kernel tested so far -- KIVI's quantize/dequantize fusion, TurboQuantRVQ's pack kernel, QFilters' fused eviction -- produced text that was byte-identical to its pure-MLX fallback. The kernels differed in what they cost or saved end-to-end, sometimes substantially, but never in what they computed. Bit-exactness held every time.

This post benchmarks VecInfer, this repo's product vector-quantization method, and that pattern doesn't hold. The kernel is also, for the first time in this series, a genuinely large speedup rather than an invisible one. Both facts turned up in the same run, and neither one excuses the other.

What VecInfer does, and the calibration caveat up front​

VecInfer applies a per-(head, channel) smooth scale and a Walsh-Hadamard rotation to keys (suppressing outliers before quantization), encodes the transformed keys against a trained codebook, and immediately dequantizes back to fp16 before attention runs -- the live tensor is fp16 at all times, the same accepted "quantize then dequantize immediately" pattern as KIVI. So, like KIVI and unlike TurboQuantRVQ, any compression ratio reported here is a hypothetical byte-accounting number, not a claim about resident memory.

One thing has to be said plainly before the numbers: this run does not use a calibrated codebook. VecInferKVCache's own docstring says a random-initialized codebook is "only useful for shape/wiring tests." No tooling for calibrating against real model activations exists in this repo yet. What this script uses instead -- reusing existing code from benchmark_vecinfer.py -- is k-means-style codebook training on synthetic Gaussian samples shaped like the model's keys and values, not the model's actual activations. That's a real, trained codebook, meaningfully better than pure random initialization, but it is still not calibrated to Qwen3-8B's real key/value distribution. Expect degraded output quality from that alone, independent of anything the kernel does. This is reported as exactly that -- not concealed, not treated as a surprise later in the post.

The kernel toggle​

Unlike TurboQuantRVQ and QFilters, VecInferKVCache exposes a public three-state use_metal_kernels field on KVCacheConfig -- same shape as KIVI's toggle, so this run needed no private-attribute override for a clean A/B. With Qwen3-8B's head_dim=128 (power-of-two, well under the 512-element threadgroup cap), the "on" arm takes a genuinely different code path: a single fused Metal dispatch that does smooth-scale + Hadamard transform + quantize + dequantize + inverse-transform in one kernel call. The "off" arm runs the same math as seven separate MLX ops. Same protocol otherwise: fp16 vs. kernel-off vs. kernel-on, interleaved repeats, median/min/max, two prompt lengths, GPU-contention check.

The numbers​

Run against mlx-community/Qwen3-8B-4bit on Apple M4 (24 GB), 36 layers, 8 KV heads, head_dim 128, synthetic-Gaussian-calibrated codebooks (key_codebook_bits=12, value_codebook_bits=8, the method's own defaults), 5 interleaved repeats:

But: kernel-on and kernel-off do not agree​

Every arm in this run is internally deterministic -- all 5 repeats of vecinfer_off produce identical text, all 5 repeats of vecinfer_on produce identical text, at both prompt lengths. That rules out sampling noise as the explanation for what comes next: vecinfer_on and vecinfer_off do not produce the same text as each other, at either prompt length.

short prompt, vecinfer_off (all 5 repeats):
ssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssssss::sssssss...

short prompt, vecinfer_on (all 5 repeats):
ssssssssssssssssssssssssss::sssssssssssssssssssssssssssss::s:ssssss:::::ssss...

long prompt, vecinfer_off (all 5 repeats):
:::::::::::::::::::::::::::::::::::::::::::::::::::::
:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

long prompt, vecinfer_on (all 5 repeats):
:
:
:
:
:
:
ed:
ed:
:
:
:
ed:
...

Both outputs are also badly broken -- repeated s, repeated :, no coherent language in either arm. That's the calibration caveat showing up exactly as expected: a codebook trained on synthetic Gaussian samples, not real key/value activations, was never going to produce fluent generation. But "both are broken in different, non-bit-exact ways" and "one kernel is 14x faster than its counterpart" are two separate facts sitting in the same run, and the second one doesn't make the first one go away. A production deployment choosing use_metal_kernels=True for the throughput win is choosing a numerically different computation, not just a faster one -- something the QFilters and TurboQuantRVQ kernels, tested earlier in this series, did not require anyone to accept.

What I'd take from this​

Four methods into this series, and this is the first one where the honest headline isn't "invisible" or "costly" -- it's "faster, and not proven equivalent." Both halves of that sentence are real, and reporting only the flattering half would have been the easy version of this post to write.

Benchmarked on an Apple M4 (10-core GPU, 24GB unified memory) against mlx-community/Qwen3-8B-4bit (36 layers, 8 KV heads, head_dim 128), VecInfer key_codebook_bits=12, value_codebook_bits=8 (the method's own defaults), synthetic-Gaussian-trained codebooks (not calibrated against real model activations -- see the calibration caveat above), 5 interleaved repeats per arm, 120 max tokens. Benchmark script: benchmark_scripts/benchmark_qwen3_8b_vecinfer_honest.py. See the KIVI, TurboQuantRVQ, and QFilters posts for the first three in this series, and vecinfer_cache.py for the method's implementation and its fused-kernel code paths.

The full benchmark results for issue #307's cross-layer/multi-request batched decode-attend kernel: kernel-level numbers, the real-model end-to-end tokens/sec table, and why one axis of the same optimization won while the other structurally cannot.

Correction (2026-09-30): The request-batching speedups in this post (1.50x–3.83x) were measured against a baseline that dequantizes the whole KV history on every decode step. No cache in this repo decodes that way — KIVIKVCache keeps fp16 and runs plain attention. Against plain fp16 KVCache, the fused path measures 0.48–0.85x end to end on the same model, and mlx_lm's built-in QuantizedKVCache(bits=4) is faster than both. Routing decode through scalar_fused_decode_attend is not recommended. The cross-layer findings are unaffected. Details: docs/KV_KERNEL_ROOFLINE_FINDINGS.md, "Addendum: the fused decode path is slower than the production baseline".

Here is the headline number: 3.83x real decode throughput on an actual model, from a kernel that already shipped in this repo and nobody had wired into a real serving path. And here is the number right next to it that matters just as much: 1.0x — because the other half of the same optimization idea, which looked identical on paper, cannot ever produce a real speedup, for a reason baked into how every transformer computes.

This is the results writeup. For the story of how these numbers were found — the failed baseline, the two test-harness bugs caught before trusting them, the model-source-reading that closed off half the idea — see the companion post, Batching Decode Attention Across Layers Sounded Great — Until the Residual Stream Said No.

The setup​

docs/KV_KERNEL_ROOFLINE_FINDINGS.md had already shown scalar_fused_decode_attend — this repo's fused group-affine decode+attention kernel — is occupancy-bound at realistic decode shapes: a single decode step for one request dispatches only 8-32 Metal threadgroups, far too few to fill a 10-core Apple GPU. The doc's own Recommendation #2 named the fix: dispatch more threadgroups per call, either by batching multiple layers or multiple requests into one launch. This benchmark tests both, on real hardware and, for the request axis, a real model.

All numbers: one base 10-core Apple M4, 24 GB unified memory.

Result 1: cross-layer batching, kernel level — a real, consistent win​

scalar_fused_decode_attend_batched adds one outermost NL (layer count) axis to the existing kernel's grid, so one dispatch covers every transformer layer's decode-attend call instead of 28-80 separate launches. Correctness: bit-identical, not just close, to looping the single-layer kernel NL times and stacking the outputs — verified for NL ∈ {1, 4, 32} plus an adversarial NL=3, B=2 combined-indexing test.

Every shape tested landed as a win — no null or negative result, unlike two prior attempts at fixing this same occupancy problem (GQA head-packing measured 2.7-4.7x slower; a SIMD-shuffle alternative was ruled out architecturally before being built). The win grows with S_kv, topping out at 4.5x at 16k tokens across 80 layers, and shrinks toward ~1.0-1.3x at small S_kv combined with a high query/kv-head ratio, where fixed dispatch overhead eats a larger share of both numbers.

Result 2: the stacking tax that erases it for real single-request serving​

Batching requires the caller to mx.stack each layer's tensors into one buffer before dispatch. That cost is real and doesn't scale predictably:

At S_kv=2048, NL=32, the ~0.6ms kernel-level saving is dwarfed by a 25ms stacking cost. Reading mlx_lm's actual KVCache source confirmed this isn't a one-time cost: it builds one independent cache object per layer ([KVCache() for _ in range(num_layers)]), with no shared layer-stacked buffer anywhere — so a real integration would pay the stacking cost every decode step, not once.

Result 3: the structural wall — 1.0x, permanently​

The real blocker for single-request decode isn't the stacking cost — it's underneath it. Reading mlx_lm's model code directly:

# TransformerBlock.__call__
r = self.self_attn(self.input_layernorm(x), mask, cache)
h = x + r
r = self.mlp(self.post_attention_layernorm(h))
out = h + r
return out

Layer L+1's attention input needs layer L's complete block output — attention, residual, MLP, residual — not just its attention output. No reordering of a standard decoder transformer lets N layers' attention be grouped into one dispatch while still computing the same model. This closes cross-layer batching for real single-request decode as a structural dead end, independent of hardware, kernel quality, or how the stacking cost might be optimized away later. Speedup for this path, permanently: 1.0x — it cannot apply to the case it was aimed at.

Result 4: the request-batching half — real, positive, measured end-to-end on Qwen3-4B​

The other half of the same original lever — batching across concurrent requests rather than layers — has no residual-stream dependency to block it, and it doesn't even need the new batched kernel: the existing, already-shipped scalar_fused_decode_attend already has a B axis in its dispatch grid. Nothing in this repo had ever routed real generation through it, though — the shipped KIVIKVCache dequantizes to fp16 and calls standard SDPA instead.

Measured on mlx-community/Qwen3-4B-4bit (36 layers, H_q=32, H_kv=8), real prompts, real greedy decoding, both arms verified to produce bit-identical output tokens before any timing was trusted:

TTFT was unaffected in both arms at every batch size — expected, since prefill never touches this decode-only kernel. The speedup climbing with B (1.50x → 3.83x) is the exact occupancy signature the roofline document's synthetic sweep predicted, now confirmed through a real forward pass on a real model instead of an isolated kernel call.

Four findings worth pulling out​

1. Two axes of the "same" optimization can have opposite outcomes​

Cross-layer batching and request batching both raise threadgroup count by exactly the same mechanism (a new grid axis). One is permanently blocked by the residual stream; the other works cleanly. Symmetry in the kernel design does not imply symmetry in the real-world result — the two axes had to be tested separately, on real models, to find that out.

2. Most of the win comes from skipping dequantization, not from the attend loop itself​

At the actual real-model shape (S_kv≈64, B=4), the fused kernel measured only ~1.15x faster than a plain fp16 cache with no quantization at all — versus 2.57x faster than the KIVI dequant-then-SDPA baseline it's meant to replace. Most of the headroom in the 1.50x-3.83x end-to-end numbers above comes from skipping the fp16 materialization step specifically, not from the attend computation being dramatically cheaper in absolute FLOPs.

3. Test-harness bugs can look exactly like real results if you don't check tokens​

Two bugs surfaced while building the real-model benchmark, and both would have silently produced a wrong number if uncaught: a full-history requantization cost (~25ms/step across 36 layers) that swamped the kernel's own cost, and a baseline arm that was accidentally comparing against a non-quantized fp16 cache instead of the real KIVI dequant path. Both were caught only because both arms' output tokens were compared for exact equality before any timing number was trusted — a cheap check that would have been easy to skip.

4. A closed negative result is as valuable as an open positive one​

The cross-layer half of this work produced no usable speedup, but it produced a permanent answer: don't revisit this axis for single-request serving, on any model, on any future hardware — the constraint is in the model's math, not this GPU or this kernel. That's a stronger, more durable finding than "didn't try it" or "not yet integrated," and it's reported with the same detail as the positive result above rather than left out.

What this means in practice​

If you're serving one request at a time, this work doesn't change anything for you today — the cross-layer kernel exists and is tested, but there's no path to a real speedup from it under standard transformer architectures.

If you're serving multiple concurrent requests through a KIVI-quantized (or similarly structured) KV cache, routing decode-step attention through scalar_fused_decode_attend instead of dequantize-then-SDPA is a real, verified win that grows with batch size — up to 3.83x at B=32 in this benchmark. The cache used here is a minimal test harness (no fp16 residual window), so a production integration — wired through KVCacheBuilder, with a real residual window, tested against variable-length concurrent requests — is the natural next step and isn't shipped in this PR.

Reproducibility​

# Kernel-level benchmark (Results 1-2 above), self-calibrated bandwidth peak
python benchmark_scripts/benchmark_crosslayer_decode_batch.py

# Real-model end-to-end benchmark (Result 4 above) — must run as a module,
# not by path, or it can resolve to a stale installed copy of the package
python -m benchmark_scripts.benchmark_real_model_scalar_attend

# Correctness (72 tests, 16 new for this work)
pytest veloxquant_mlx/tests/metal/test_scalar_attend.py -v -k batched

Full methodology, additional shapes, and the complete honesty caveats (what the benchmark cache does and doesn't represent about production KIVI) are in docs/KV_KERNEL_ROOFLINE_FINDINGS.md's two new addenda: "cross-layer batched decode-attend dispatch (issue #307, part 1)" and the real-model follow-up immediately after it.

VeloxQuant-MLX is MIT licensed. Hardware: Apple M4, 24GB unified memory. Model: mlx-community/Qwen3-4B-4bit.

Issue #307's own roofline analysis named one unblocked, unattempted lever for scalar_fused_decode_attend's occupancy problem: batch multiple layers' independent decode-attend calls into one Metal dispatch. Building it produced a real, measured kernel-level win — and then a real-model test found the lever can't reach where it matters most, for a reason baked into how every transformer works. The other half of the same lever could, and did, on a real model.

Correction (2026-09-30): The request-batching speedups in this post (1.50x–3.83x) were measured against a baseline that dequantizes the whole KV history on every decode step. No cache in this repo decodes that way — KIVIKVCache keeps fp16 and runs plain attention. Against plain fp16 KVCache, the fused path measures 0.48–0.85x end to end on the same model, and mlx_lm's built-in QuantizedKVCache(bits=4) is faster than both. Routing decode through scalar_fused_decode_attend is not recommended. The cross-layer findings are unaffected. Details: docs/KV_KERNEL_ROOFLINE_FINDINGS.md, "Addendum: the fused decode path is slower than the production baseline".

The problem this picks up​

docs/KV_KERNEL_ROOFLINE_FINDINGS.md already established that scalar_fused_decode_attend — the fused group-affine decode+attention kernel behind this repo's KIVI-style KV cache — is not memory-bound, not compute-bound, and not launch-bound. It's occupancy-bound: a real decode step (B=1, H_kv=8, S_q=1) dispatches only 8 threadgroups, nowhere near enough to fill a 10-core Apple GPU, no matter how efficient the kernel's inner loop is or how much data any one threadgroup has to move.

Two attempts at fixing this from inside the kernel had already failed, honestly and with numbers: GQA head-packing (share decoded K/V across query heads sharing a kv-head) measured 2.7-4.7x slower, because it traded away threadgroup count to save decode arithmetic — the wrong trade when occupancy, not arithmetic, is the bottleneck. A simd_shuffle-based alternative was ruled out architecturally before it was even built: cross-lane shuffles only work within one threadgroup, and under full-occupancy dispatch a threadgroup only ever holds one query head's data.

What the roofline doc's own Recommendation #2 named as the actual unblocked lever: dispatch more threadgroups per call by processing multiple layers or multiple requests in one kernel launch — batching across what's currently separate per-layer Python-level calls, rather than trying to make one tiny dispatch stream memory faster. This post is about building both halves of that lever, and getting two very different answers.

Part 1: batching across layers — a real kernel win, verified bit-for-bit​

A real decode step calls scalar_fused_decode_attend once per transformer layer — typically 28 to 80 times for current open-weight models — and those calls are independent at the Metal-dispatch level: each layer's own K/V cache and post-attention Q projection are already sitting in memory before any of these calls run. Nothing about attention's data dependencies forces them to be separate launches.

scalar_fused_decode_attend_batched adds exactly one thing to the existing kernel: a new outermost NL (num_layers) grid axis, so threadgroup count becomes NL * B * H_kv * S_q instead of B * H_kv * S_q. The indexing change in the Metal source is small and mechanical —

uint sq_idx  = tg % S_q;
uint hkv_idx = (tg / S_q) % H_kv;
uint b_idx   = (tg / (S_q * H_kv)) % B;
uint l_idx   = tg / (S_q * H_kv * B);

uint bh_kv = (l_idx * B + b_idx) * H_kv + hkv_idx;   // one new leading stride term

— and every buffer offset that touches bh_kv, q, or out gets the same added l_idx * (...) leading term. Same math, same memory layout, same per-(layer, batch, kv-head, query-position) work per threadgroup. The caller pre-stacks each layer's tensors along a new leading axis; the kernel doesn't gather anything itself.

The load-bearing correctness check: calling the batched kernel once must be bit-identical, not just within tolerance, to calling the single-layer kernel NL times in a loop and stacking the outputs.

@pytest.mark.parametrize("NL", [1, 4, 32])
def test_scalar_attend_batched_parity_vs_single_layer_loop(NL):
    ...
    max_abs = np.abs(got_np.astype(np.float32) - ref.astype(np.float32)).max()
    assert max_abs == 0.0, f"NL={NL}: batched must be bit-identical to the single-layer loop"

It passed on the first real run, along with an adversarial NL=3, B=2 test specifically built to catch a swapped layer/batch stride order, and a non-broadcast test confirming layer 1's output isn't secretly layer 0's data repeated (the kind of copy-paste stride bug that a same-content test would miss entirely).

The result was a consistent, positive win everywhere tested — no null result, unlike both prior attempts:

The win grows with S_kv (up to 4.5x at 16k tokens) and shrinks toward ~1.0-1.3x at small S_kv combined with a high H_q/H_kv ratio, where fixed dispatch overhead is a larger share of both numbers. Achieved bandwidth in the batched case reached up to ~40% of a freshly self-calibrated peak — still short of memory-bound, meaning occupancy remains the governing constraint even after batching, just a much less severe one.

But a kernel-level win isn't the same thing as a real speedup, and the next two steps are why.

Part 2: the stacking cost, and a KVCache that won't cooperate​

Batching requires the caller to hand over layer-stacked tensors — mx.stacking NL independent per-layer arrays before every dispatch. That's not free, and it doesn't scale the way you'd guess:

At S_kv=2048, NL=32, a ~0.6ms kernel-level saving is dwarfed by a 25ms stacking cost — if that stacking has to be paid on every decode step, the kernel win is meaningless.

Whether it does depends entirely on the KV-cache implementation underneath, so instead of assuming an answer, I read mlx_lm's actual cache source:

# mlx_lm/models/cache.py
return [KVCache() for _ in range(num_layers)]

One independent KVCache object per layer, each with its own independently-growing .keys/.values buffer. There is no shared, layer-stacked backing buffer anywhere in the stock cache — which means a real integration built on it would need to mx.stack() every layer's state fresh, on every single decode step, not once at cache-construction time. The stacking cost above is a per-step tax, not a one-time layout change. Same amortization failure mode this repo already found for a sibling kernel (fused_sdpa's dispatcher is a documented no-op for exactly this reason) — just arriving via a different mechanism.

Part 3: the wall — why cross-layer batching can never reach real single-request decode​

Even setting the stacking cost aside, there's a harder problem underneath it, and it took reading mlx_lm's model code directly to see it clearly:

# mlx_lm/models/llama.py — LlamaModel.__call__
for layer, cache in zip(self.layers, cache):
    mask = swa_mask if layer.use_sliding else fa_mask
    h = layer(h, mask, cache=cache)

# TransformerBlock.__call__
r = self.self_attn(self.input_layernorm(x), mask, cache)
h = x + r
r = self.mlp(self.post_attention_layernorm(h))
out = h + r
return out

Layer L+1's attention input is q_proj(norm(layer_L_output)) — and layer_L_output is layer L's entire block output: attention, residual, MLP, residual. Not just its attention output. There is no reordering of a standard pre-norm decoder transformer that lets N layers' attention be grouped into one dispatch while still computing the same model. "Run all layers' Q/K/V projections first, batch the attend call, then finish the MLPs" isn't a valid optimization here — layer 2's Q/K/V projections need layer 1's MLP output, which needs layer 1's attention output already. Any patch that tried this would silently change what the model computes, which fails the correctness bar this repo holds everywhere else.

This is a structural fact about the residual stream, not a missing engineering patch. Cross-layer batched decode-attend cannot speed up real single-request decode latency on any standard transformer — full stop, regardless of how good the kernel or the batching mechanism gets. That's a stronger, more useful conclusion than "not yet integrated": it closes this half of the lever as a dead end, not a pending follow-up.

Part 4: the other half of the lever actually works — measured on a real model​

Recommendation #2 named two axes: batch across layers, or across requests. Concurrent requests have no residual-stream dependency between them — they're independent by construction — so this half was never blocked by anything in Part 3. And it turns out it doesn't even need the new batched kernel: the original, already-shipped scalar_fused_decode_attend already carries a B (batch/request) axis in its dispatch grid (n_tg = B * H_kv * S_q). Nobody had ever wired real generation through it, though — KIVIKVCache dequantizes to fp16 and hands off to standard SDPA instead, the same pattern that made the fused kernel a no-op for the sibling VecInfer cache.

So I built a minimal cache that keeps KIVI-quantized codes live and exposes them, and monkeypatched mlx_lm's SDPA dispatch to route real decode-step attention through the fused kernel when it applies:

def _patched_sdpa(queries, keys, values, cache, scale, mask, sinks=None):
    if S_q == 1 and isinstance(cache, _ScalarAttendKIVICache) and cache._k_codes is not None:
        if use_fused:
            return scalar_fused_decode_attend(
                queries, *cache.quantized_state(), GROUP_SIZE, scale, nsg=nsg
            )
        else:
            k_hat, v_hat = cache.dequantized_kv(heads_per_kv)  # what KIVI does today
            return _original_sdpa(queries, k_hat, v_hat, cache=None, scale=scale, mask=None)
    return _original_sdpa(queries, keys, values, cache=cache, scale=scale, mask=mask)

Two bugs surfaced immediately, and both mattered enough to fix before trusting any number. First: re-quantizing the entire growing K/V history from scratch every decode step cost ~25ms/step across 36 layers — pure test-harness overhead that would have silently drowned the kernel's own ~12ms/step in noise. Fixed by quantizing only the newly-aged group_size-aligned block each step, mirroring the incremental-flush discipline this repo's real KIVIKVCache already uses. Second, and sneakier: the "baseline" arm was accidentally attending over the cache's raw fp16 buffer instead of dequantized KIVI codes — comparing the fused kernel against a plain fp16 cache with no quantization at all, not against the KIVI dequant+SDPA path it's actually meant to replace. Fixed by giving the cache an explicit dequantized_kv() method both arms could be checked against, and verifying — before trusting a single timing number — that both arms decode to exactly the same tokens:

dequant-baseline route counts: {'decode_fused': 0, 'decode_dequant_baseline': 360, 'fallback': 0}
decoded (dequant baseline):  sentence is00000000
fused route counts: {'decode_fused': 360, 'decode_dequant_baseline': 360, 'fallback': 0}
decoded (fused):  sentence is00000000
tokens match: True

With that confirmed, the real numbers, on mlx-community/Qwen3-4B-4bit (36 layers, H_q=32, H_kv=8), real prompts, real greedy decoding, real end-to-end tokens/sec including embeddings, MLPs, o_proj, and sampling — not an isolated kernel call:

The speedup growing with batch size — 1.50x at B=1, up to 3.83x at B=32 — is exactly the occupancy signature the synthetic roofline sweep predicted, now confirmed through an actual forward pass instead of an isolated kernel call. TTFT was unaffected in both arms, as it should be: prefill uses S_q > 1 and never touches this decode-only kernel.

Two caveats worth stating rather than glossing over. The isolated per-layer comparison at this shape showed the fused kernel only ~1.15x faster than a plain fp16 cache with no quantization at all — most of the 2.57x-over-dequant-baseline margin comes specifically from avoiding the dequant materialization step, not from the attend computation itself being dramatically faster in absolute terms. And the cache built for this test has no fp16 residual window, so its output quality isn't representative of production KIVI — the timing comparison is sound because both arms decode from identical (lossy) quantized state and were verified token-for-token, but "1.5-3.8x faster" describes the kernel swap, not a production-ready serving stack.

Where this leaves things​

Two symmetric findings, both earned by actually building and measuring rather than assuming: cross-layer batching is a structural dead end for the case that matters most (a single request's decode latency), closed for good reasons rather than left open as unfinished work — and the request-batching half of the same original lever is real, positive, and now confirmed on an actual model rather than a synthetic sweep, using a kernel that was already sitting in the codebase unused. A production-grade version — a real residual-window cache, wired through KVCacheBuilder, tested against variable-length concurrent requests rather than identical-length ones — is the natural next step, and isn't attempted here.

Full details, tables, and reproduction commands: docs/KV_KERNEL_ROOFLINE_FINDINGS.md's two new addenda ("cross-layer batched decode-attend dispatch" and the real-model follow-up immediately after it). Correctness: pytest veloxquant_mlx/tests/metal/test_scalar_attend.py -v -k batched (72/72 passing, zero regressions in the existing suite). Benchmarks: python benchmark_scripts/benchmark_crosslayer_decode_batch.py and python -m benchmark_scripts.benchmark_real_model_scalar_attend. All numbers from one base 10-core Apple M4, 24 GB.

For the results-first version of this post — full benchmark tables, the headline numbers, and the practical "what this means for you" guidance without the narrative — see Batching Decode Attention Across Requests Got Us 3.83x Real Throughput — Across Layers Got Us Nothing.

How one arbitrary minus sign turned a KV cache compressor from useful into noise — and what it took to prove it on real models​

I had a KV cache compression method in my library called Q-Filters. It had tests. It had docs. It had a benchmark harness with committed results. It shipped in version 0.31.0.

It also didn't work.

Not "worked slightly worse than the paper." Not "worked on some heads." On real trained weights, its scoring signal correlated with true attention at −0.032 — statistically indistinguishable from a coin flip, and pointing the wrong way about half the time.

Fixing it meant implementing the paper properly, writing two Metal kernels, finding a position-encoding bug that made end-to-end measurement impossible, and throwing away two benchmark harnesses that produced confident, meaningless numbers. Along the way the calibrated version went from ppl 598 to 16.3 on the same workload.

This is the whole run, including the parts where I was wrong.

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 to cut your KV cache memory by 8x and hit fp16 throughput — step by step, with real models.

If you have an Apple Silicon Mac and you run local LLMs, you have already hit the wall: your MacBook Pro generates 20 tokens per second on Mistral 7B, memory pressure turns on after a few thousand tokens, and anything larger than 7B crawls. The bottleneck is not compute — it is memory bandwidth. Every token requires loading the entire KV cache for every layer. Make the cache smaller, and the model gets faster.

VeloxQuant-MLX is a drop-in KV cache quantization library for MLX. It compresses the KV cache to 2, 3, or 4 bits per value while keeping text quality close to full-precision. In this guide you will install it, run your first compressed model, pick the right algorithm for your use case, and understand what the benchmark numbers mean. No machine learning background required.

What You Need​

• Apple Silicon Mac (M1 or later)
• Python 3.11 or 3.12
• At least 16 GB unified memory (8 GB works for 4B models)
• mlx-lm installed (pip install mlx-lm)

Step 1: Install VeloxQuant-MLX​

pip install VeloxQuant-MLX

That installs the mlx_kv_quant package plus the veloxquant CLI. Verify it worked:

veloxquant --help

You should see the precompute and benchmark subcommands.

Step 2: Your First Compressed Inference​

The fastest path uses the KVCacheBuilder. Here is a minimal script that runs Mistral 7B with 4-bit KV cache compression:

import mlx.core as mx
import mlx_lm
from mlx_kv_quant import KVCacheBuilder, KVCacheConfig

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

# Build a quantized KV cache
config = KVCacheConfig(bits=4, algorithm="turboquant_prod")
cache = KVCacheBuilder(config).for_model(model).build()

# Generate with the compressed cache
prompt = "Explain the difference between RAM and unified memory."
tokens = mlx_lm.generate(model, tokenizer, prompt=prompt, kv_cache=cache, max_tokens=200)
print(tokens)

That is the entire integration. The cache is transparent to mlx_lm.generate() — you do not change anything else.

Step 3: Understanding the Algorithm Choices​

VeloxQuant-MLX ships five quantization algorithms. They trade reconstruction quality against memory and speed:

Rule of thumb:

• For most uses at 3–4 bit: use turboquant_prod. It is the default and the most tested.
• At 2 bits: always use turboquant_rvq. The single-pass algorithms fall apart at 2-bit; RVQ does not.
• For long contexts where cosine similarity matters more than reconstruction: polarquant.
• To stress-test memory limits: qjl.

Step 4: Using RVQ 2-Bit (the Interesting One)​

The headline feature in v0.3.0 is TurboQuantRVQ — a two-pass residual vector quantizer that makes 2-bit KV cache actually usable.

The problem with naive 2-bit quantization: you only have 4 levels to represent a continuous value. The reconstruction error is large enough that attention scores get corrupted, and long reasoning chains (like Qwen3's <think> mode) collapse into repetition after a few dozen tokens.

RVQ fixes this by running two codebooks. The first pass quantizes the rotated vector. The second pass quantizes the residual — the error left over from the first pass — using a Laplacian-fit codebook (which matches the residual distribution better than a Gaussian one). The result: cosine similarity goes from 0.69 to 0.98 at 2 bits.

Here is how to use it:

from mlx_kv_quant import KVCacheBuilder, KVCacheConfig

config = KVCacheConfig(bits=2, algorithm="turboquant_rvq")
cache = KVCacheBuilder(config).for_model(model).build()

That's it. The turboquant_rvq string routes to the registered TurboQuantRVQ class via the quantizer registry.

When should you use RVQ 2-bit?

Use it when memory is the binding constraint and you need coherent long-form output. It delivers:

• Mistral 7B at 2-bit: 22.3 tok/s — matches fp16 throughput (22.1 tok/s) on an M4 MacBook
• Qwen3 4B at 2-bit: 36.0 tok/s (92% of fp16) with full thinking-mode output

Use standard 3-bit or 4-bit when you have memory headroom and want zero quality trade-off.

Step 5: Running Benchmarks on Your Mac​

VeloxQuant-MLX ships benchmark scripts for several models. These measure tok/s and output completeness across five configurations: fp16, 2-bit RVQ, 3-bit, 4-bit, and 2-bit single-pass (for comparison).

Qwen3 4B (good for 16GB Macs):

python3 benchmark_qwen3_4b_v2.py

Mistral 7B (needs 16GB+, comfortable at 32GB):

python3 benchmark_mistral7b_v2.py

Each script runs all five configs, prints a results table, and saves six figures to figures/updated_tests/<model>/. The figures include:

• Throughput comparison bar chart
• Token completeness (how many tokens were coherent out of 200)
• Memory usage per config
• Cosine similarity vs fp16 reference

Here is what our numbers look like on an M4 MacBook (16GB):

Mistral 7B throughput:

Mistral 7B is memory-bandwidth bound — every config saturates around 22 tok/s because that is the bandwidth ceiling. But RVQ 2-bit uses 10x less cache memory than fp16.

Qwen3 4B thinking-mode throughput:

The 4-bit single-pass result (50/200) is the failure mode for standard quantization on a thinking model. RVQ 2-bit produces full coherent output and runs faster.

Step 6: Integrating with mlx-lm Properly​

For real usage you want the cache to persist across generation steps. Here is the pattern for a chat loop:

import mlx.core as mx
import mlx_lm
from mlx_kv_quant import KVCacheBuilder, KVCacheConfig

model, tokenizer = mlx_lm.load("mlx-community/Qwen3-4B")

config = KVCacheConfig(bits=2, algorithm="turboquant_rvq")
cache = KVCacheBuilder(config).for_model(model).build()

messages = []

while True:
    user_input = input("You: ")
    if user_input.lower() in ("exit", "quit"):
        break

    messages.append({"role": "user", "content": user_input})
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    response = mlx_lm.generate(
        model,
        tokenizer,
        prompt=prompt,
        kv_cache=cache,
        max_tokens=512,
        verbose=False,
    )

    messages.append({"role": "assistant", "content": response})
    print(f"Assistant: {response}\n")

The kv_cache argument is passed directly to mlx_lm.generate(). No other changes. The cache grows with each turn and is quantized transparently.

Step 7: Precomputing Codebooks for Repeated Use​

If you run the same model repeatedly, you can precompute the Lloyd-Max codebooks offline and load them at inference time. This eliminates the one-time calibration cost at startup:

# Precompute and save codebooks for Mistral 7B, 4-bit
veloxquant precompute \
  --model mlx-community/Mistral-7B-Instruct-v0.3-4bit \
  --bits 4 \
  --algorithm turboquant_prod \
  --output codebooks/mistral7b_4bit/

# Later, load at inference time
config = KVCacheConfig(
    bits=4,
    algorithm="turboquant_prod",
    artifact_dir="codebooks/mistral7b_4bit/",
)
cache = KVCacheBuilder(config).for_model(model).build()

Codebook size is tiny (kilobytes). You can commit them to your project directory and avoid the calibration step entirely on every run.

Step 8: Picking the Right Bit Width​

Here is the practical decision tree:

Do you have >= 32GB memory?
  Yes → fp16 is fine; use VeloxQuant only for very long contexts
  No  → continue

Is your model >= 7B parameters?
  Yes → 3-bit TurboQuant (turboquant_prod)
  No  → 2-bit RVQ (turboquant_rvq)

Are you doing RAG or similarity search over the KV cache?
  Yes → turboquant_prod (inner-product correction matters)
  No  → turboquant_mse (simpler, same quality)

Do you need to run 13B+ on a 16GB Mac?
  Yes → 2-bit RVQ — the only option that keeps quality

For generation tasks (chat, summarization, code completion), MSE-optimal reconstruction is what you want. For retrieval tasks that compute attention over a cached context, inner-product preservation matters more — use turboquant_prod.

Step 9: What the Optimization Journey Looked Like​

Getting from 17.7 tok/s to 22.3 tok/s on Mistral 7B RVQ 2-bit required four changes, each independently measurable:

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

The original code ran a Python for h in range(H) loop over attention heads, giving each head its own quantizer instance. For Mistral 7B (8 heads, 32 layers) that meant 256 small kernel dispatches per token. The fix: reshape (B, H, S, D) to (B*H*S, D) and call one shared quantizer once. MLX already handles batched input — we just stopped fragmenting it.

2. Switch to Hadamard rotation

The default preconditioner used QR decomposition: a full (d, d) matrix multiply (16,384 ops at d=128). mx.hadamard_transform is a Metal-native fused kernel that does the same job in O(d log d) — 896 ops at d=128, ~18x less arithmetic. Quality is mathematically identical because Hadamard with random ±1 diagonal is also a Haar-equivalent rotation.

3. Replace broadcast-argmin with boundary-sum

Codebook lookup materialized a (batch, 128, k) distance tensor for argmin — three kernels: broadcast-subtract, abs, argmin. Lloyd-Max boundaries are just midpoints between sorted centroids. The nearest centroid index is the number of boundaries the value exceeds: one comparison and one sum, two kernels total. Index match vs the old path: 100.0000%.

4. Drop redundant fp32 casts

The update path was casting fp16 → fp32 → fp16 → fp32 → fp16 unnecessarily because the rotation already handles internal precision internally. Removing the round-trips saved ~1.05x per call and is invisible to output quality.

The full write-up with stage-by-stage numbers is in OPTIMIZATION_FINDINGS.md.

Step 10: Understanding Memory Numbers​

KV cache memory scales with: 2 × n_layers × n_kv_heads × head_dim × seq_len × bytes_per_element

For Mistral 7B (32 layers, 8 KV heads, head_dim 128) at sequence length 2048:

• fp16: 2 × 32 × 8 × 128 × 2048 × 2 bytes = 536 MB
• 4-bit: 268 MB (2x reduction)
• 2-bit RVQ: 134 MB (4x reduction, plus codebook overhead ~8 KB)

At longer contexts (32K tokens), these numbers scale linearly:

• fp16: 8.4 GB for context alone
• 2-bit RVQ: 2.1 GB

On a 16GB Mac, the difference between fp16 and 2-bit RVQ at 32K context is the difference between OOM and running.

Common Questions​

Does this work with quantized models (4-bit weights)?

Yes. The KV cache quantization is independent of weight quantization. You can run a 4-bit weight model with a 2-bit KV cache. The two compressions are additive: mlx-community/Mistral-7B-Instruct-v0.3-4bit loads as a 4-bit weight model; the KV cache is then further compressed by VeloxQuant.

Will I notice quality degradation?

At 4-bit and 3-bit: barely measurable on generation tasks. At 2-bit with RVQ: cosine similarity 0.98 (vs 1.0 for fp16). On our Qwen3 4B thinking-mode test, RVQ 2-bit produced 199/200 coherent tokens vs fp16's 200/200. One token difference across a 200-token reasoning chain.

At 2-bit without RVQ (single-pass): noticeable. Thinking-mode models collapse early. Use RVQ.

What models are supported?

Any model that uses standard multi-head attention with mlx_lm.make_cache(). Tested: Mistral 7B, Qwen3 4B, Qwen3 8B, Qwen2.5 32B, Falcon3 7B, Phi-4, Gemma-4. Not supported: architectures with multi-latent attention (DeepSeek-V2) or non-standard cache formats.

Does it work with streaming generation?

Yes. The cache is stateful — each generate() call updates it. Stream tokens with the standard mlx_lm streaming API; the cache update happens transparently inside the model forward pass.

VeloxQuant-MLX is MIT licensed. Contributions welcome — especially benchmark results for new models and hardware.

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

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