5 posts tagged with "benchmarking"

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.

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.

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.

Issue #277 asked whether a hand-written Metal kernel could speed up prefill-side attention, the way this repo's decode kernels already speed up cache reads. The honest answer, measured on the hardware in hand: no — and the reason why is more interesting than a speedup would have been.

What decode kernels do that prefill can't reuse​

Every Metal kernel in VeloxQuant-MLX before this — quantize_vq, rabitq_encode, rabitq_fused_attend, scalar_fused_decode_attend — accelerates something that happens to an already-built KV cache. That's the right target for decode: one new query token against a long cache is bandwidth-bound, and Apple Silicon's unified memory is genuinely competitive there.

Prefill is a different problem. The whole prompt goes through attention in one shot — S_q ≈ S_kv, both potentially tens of thousands of tokens for an agentic coding session that re-feeds a large, mostly-unchanged codebase every turn. That's compute-bound: big batched matmuls, not a bandwidth story. And it's the regime where Apple GPUs are furthest behind CUDA — back-of-envelope estimates put an M3 Ultra around 30x slower than an RTX Pro 6000 at large-context prefill.

None of the existing kernels touch this. They all assume the cache already exists. So the question in #277 was: is there a real software gap here that a simdgroup_matrix-tiled kernel (the same technique behind rabitq_prefill_attend) could close, or is the ~30x gap purely a hardware FLOPs ceiling that no amount of kernel-writing fixes?

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

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

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

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

What already existed​

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

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

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

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

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

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

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

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

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

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

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

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

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

Building the thing anyway​

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

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

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

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

Where the time actually went​

I profiled it instead of guessing.

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

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

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

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

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

The part where fixing one problem created a worse one​

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

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

I broke down where the bytes went:

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

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

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

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

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

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

Making the tradeoff a choice instead of a default​

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

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

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

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

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

What didn't get tested, and why​

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

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

What actually shipped​

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

What this experiment is actually about​

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

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

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

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

How I fused KIVI's KV-cache quantization into two bit-exact Metal kernels, measured a 1.40×–5.65× op-level speedup, watched it vanish completely end-to-end — and the four separate times my own benchmarks gave me confidently wrong answers along the way.

There's a particular satisfaction in watching a GPU kernel you wrote beat the framework's version by 5×. There's a different feeling entirely when you plug it into the actual model and the tokens-per-second doesn't move at all.

This post is about both, and about the part in between: four occasions where my benchmarks confidently told me something false. One said the kernel was 11× slower than baseline. One said it was 28× faster. One said quantization consumed 97.83% of prefill time, when the real figure is about 1–2%. One showed a clean 127 MB memory saving that turned out to be nothing at all.

All four were my fault. All four produced a clean-looking number with a plausible story attached, which is exactly what made them dangerous.

The kernel is real, it's bit-exact, it ships, and I'd merge it again. But the honest headline is the one above, and the useful content is why both halves of it are true at the same time.

This is the long version, with the complete experimental record — every measurement generation, including the ones that were wrong.

What we're actually optimizing​

If you run a language model locally, the thing that eventually stops you isn't compute. It's the KV cache.

Every token the model has seen leaves behind a key and a value vector in every attention layer. The model needs them to attend to the past, so they stay resident for the whole generation. The cache grows linearly with context — and unlike model weights, which you load once, it grows while you're using it.

On a 7B model at 32k context that's several gigabytes, comparable to the quantized weights themselves. On a Mac with unified memory, it's the difference between a long conversation working and your machine swapping itself to death.

KIVI (Liu et al., ICML 2024) is one answer, and it's the baseline every other algorithm in this library gets measured against. The insight is that keys and values want to be quantized along different axes:

• Keys are quantized per-channel — each channel gets its own scale, computed across a group of tokens.
• Values are quantized per-token — each token gets its own scale, computed across a group of channels.

Why asymmetric? Key tensors have a few channels with consistently huge magnitudes. Quantize per-token and those outliers blow up the scale for every other channel sharing the group. Value tensors lack that structure, and per-token suits them better.

A third piece matters for everything that follows: the most recent residual_length tokens stay in fp16. They're what attention weights most heavily, and they're also the tokens whose group isn't full yet. Once enough fresh tokens accumulate, they get quantized as a batch and folded into the compressed store. That batching event is a flush, and it is the operation this entire post is about.

The quantization itself is textbook asymmetric min/max:

zero  = min(group)
scale = (max(group) - min(group)) / (2^bits - 1)
q     = round((x - zero) / scale)
recon = q * scale + zero

In MLX this is roughly eight array operations: reshape to expose the group axis, min, max, subtract, divide, round, clip, multiply, add, reshape back.

Eight operations means eight kernel launches and — the expensive part — eight round trips to memory. Every intermediate is materialized. The quantized codes, which never needed to exist as a full-size tensor, get written to RAM and read straight back.

That's the target. One fused kernel, one pass, no intermediates.

The kernel​

The layout problem that shapes everything​

KV tensors are [batch, heads, seq, head_dim], row-contiguous. Flatten batch and heads and you get [BH, S, D], where element (bh, s, d) sits at bh*S*D + s*D + d.

Which means the two modes face opposite problems:

• Values (per-token) — the group runs along D, the contiguous axis. Adjacent elements in a group are adjacent in memory.
• Keys (per-channel) — the group runs along S, the strided axis. Adjacent elements are D floats apart. Typically 128.

My first version was one kernel handling both, which meant transposing the key tensor so the token axis became contiguous, then reusing the same code path.

That transpose was the whole problem. It's a full-size materializing copy — exactly the memory traffic the kernel exists to eliminate. I'd removed eight round trips and added one large one back. On the key path, the "optimized" kernel was a net loss.

So I threw it out and wrote two kernels, one per layout.

Kernel A — per-channel keys: one thread, one whole group​

The trick is almost aggressively simple: give each thread an entire group, and don't reduce at all.

uint tid = thread_position_in_grid.x;

const uint BH = x_shape[0];
const uint S  = x_shape[1];
const uint NG = (S + GROUP_SIZE - 1u) / GROUP_SIZE;   // token groups

if (tid >= BH * NG * DHEAD) { return; }

const uint d   = tid % DHEAD;
const uint r   = tid / DHEAD;
const uint grp = r % NG;
const uint bh  = r / NG;

const uint base = bh * S * DHEAD + d;
const uint s0   = grp * GROUP_SIZE;
const uint s1   = min(s0 + GROUP_SIZE, S);

float gmin =  INFINITY;
float gmax = -INFINITY;
for (uint s = s0; s < s1; ++s) {
    float v = float(x[base + s * DHEAD]);
    gmin = min(gmin, v);
    gmax = max(gmax, v);
}

Each thread strides by DHEAD — 128 floats between consecutive reads. In isolation that looks like the worst access pattern available.

But look at the indexing. d = tid % DHEAD means consecutive threads take consecutive channels. At any step of that loop, the 32 threads in a SIMD group read 32 adjacent addresses. The warp's access is fully coalesced. The stride is per-thread; the warp moves through memory as a solid block.

And because a thread owns its group outright, there is no cross-thread reduction. No threadgroup memory, no barriers, no butterfly, no transpose. The strided-looking layout turned out to need the least machinery.

Kernel B — per-token values: one SIMD group, one quantization group​

The contiguous axis wants the mirror image. Lanes split a single group and cooperate:

uint lane = thread_position_in_threadgroup.x;
uint gid  = threadgroup_position_in_grid.x;

// Whole threadgroups exit together, so every lane still reaches the
// butterfly below — a divergent return would deadlock the shuffle.
if (gid >= x_shape[0] * S * NGD) { return; }

float gmin =  INFINITY;
float gmax = -INFINITY;
for (uint i = lane; i < GROUP_SIZE; i += 32u) {
    uint d = d0 + i;
    if (d < d1) {
        float v = float(x[row_base + d]);
        gmin = min(gmin, v);
        gmax = max(gmax, v);
    }
}

// Butterfly: after 5 XOR shuffles every lane holds the group-wide min/max.
for (uint off = 16u; off > 0u; off >>= 1u) {
    gmin = min(gmin, simd_shuffle_xor(gmin, off));
    gmax = max(gmax, simd_shuffle_xor(gmax, off));
}

The simd_shuffle_xor butterfly is the nice part. Five shuffles reduce 32 lanes to a min/max that every lane already holds — no broadcast step. And because lanes advance in lockstep, it needs no threadgroup memory and no barriers, unlike a tree reduction.

Note the comment on the bounds check. That return must be uniform across the threadgroup. If individual lanes bailed early, the survivors would shuffle against threads that no longer exist and the reduction would hang or return garbage. Whole threadgroups exit together, so every lane reaching the butterfly reaches it with all 32 partners intact.

At KIVI's default GROUP_SIZE=32 this is an exact fit: one lane per element, one butterfly, done.

Two optimizations I was sure would work, and didn't​

I assumed caching each group in registers would help — read once, use twice, skip the second global load. Controlled same-process A/B said otherwise:

In the channel kernel, a whole group is GROUP_SIZE floats per thread. The occupancy that costs outweighs the saved loads, which were hitting cache anyway. In the token kernel it's split across 32 lanes, so it's cheap — and buys nothing, for the same cache reason.

Both reverted, and the REG_SLOTS machinery deleted. I also swept threadgroup width and found the defaults (256 for channel, 32 for token) already optimal.

Three ways to be off by one bit​

My acceptance criterion was bit-exactness against the MLX reference — not "close enough," but identical output. That turned out to be the most instructive constraint in the project, because it surfaced three failure modes a tolerance test sails straight past.

1. FMA contraction​

My first parity run failed on 192 of 300 configurations — and every failure was off by exactly 1 ULP. That uniformity is a fingerprint: not an algorithm bug, a rounding difference.

The culprit was q * scale + gmin. Metal's compiler sees a multiply feeding an add and contracts it into a fused multiply-add: one instruction, one rounding. MLX does them separately, with two roundings. Same math, different result on ~0.02% of elements.

The fix is to break the pattern the optimizer looks for:

float prod = q * scale;      // NOT an fma
out[base + s * DHEAD] = T(prod + gmin);

The irony is that the fused version is more accurate — it carries more intermediate precision. But the contract is parity with the reference, not maximum accuracy, so the less accurate version is the correct one. Uncomfortable sentence; right call. 300/300 after the fix.

2. Rounding mode​

Metal's round() is half-away-from-zero. mx.round is half-to-even. They agree on everything except exact .5 codes — rare in random data, and systematically common in real quantization, because uniform grids produce exact midpoints. The fix is rint(). There's a test pinning it with inputs hand-built to land on .5.

3. Padding semantics​

When a group doesn't divide evenly, MLX pads the tail by replicating the edge value, x[..., -1:], not with zeros. Pad with zeros and you've silently dragged gmin to 0 for every ragged group, corrupting the scale. Since every pad slot holds that same value, folding it in once is equivalent to looping:

if (s1 < s0 + GROUP_SIZE) {
    float pad_val = float(x[base + (S - 1u) * DHEAD]);
    gmin = min(gmin, pad_val);
    gmax = max(gmax, pad_val);
}

All three have dedicated regression tests now. They're the kind of bug that produces plausible output — slightly different, never obviously broken. Without bit-exactness as the bar, all three would have shipped.

The bug that made generation hang forever​

My favourite failure of the project, because the fix made the code simpler and the symptom was so much worse than the cause.

mx.fast.metal_kernel JIT-compiles from source. I was passing shape constants through the header as #defines, which lets the compiler turn tid % DHEAD into a shift-and-mask instead of integer division. Good idea for DHEAD — it's the model's head dimension, fixed for the life of a cache.

I did the same for the sequence length.

The sequence length grows by one every decode step. So every token triggered a fresh shader compilation. Generation didn't crash and didn't error — it just stopped. I killed the process after several minutes with no output.

The fix was to pass shape as a runtime buffer. MLX provides x_shape for free on every input, so this meant deleting code, not adding it.

Four benchmarks that lied​

Here the post stops being about GPU programming and starts being about measurement, which is the part I'd actually want to read.

Lie #1 — "Quantization is 97.83% of prefill"​

I wanted to know how much runtime quantization accounted for, so I instrumented the call: timer before, timer after, mx.eval() in between to force the computation. Here's the raw output:

# mlx-community/Llama-3.2-3B-Instruct-4bit  layers=28

PREFILL-dominated (8k prompt, 4 new tokens):
  prefill  metal=False  wall= 20.754s  quant= 20303.3ms (97.83% of wall)  calls=224  [keys 19483.2ms / values  820.0ms]
  prefill  metal=True   wall= 20.940s  quant= 20489.1ms (97.85% of wall)  calls=224  [keys 19805.5ms / values  683.6ms]

DECODE-dominated (2k prompt, 240 new tokens):
  decode   metal=False  wall= 10.114s  quant=  4829.6ms (47.75% of wall)  calls=504  [keys  4550.0ms / values  279.6ms]
  decode   metal=True   wall=  9.952s  quant=  4682.5ms (47.05% of wall)  calls=504  [keys  4457.8ms / values  224.7ms]

Quantization was apparently the entire bottleneck. It was nonsense, and the mistake is in the description above: mx.eval() inside the measured region.

MLX is lazily evaluated. Operations build a graph; nothing computes until something forces it. By calling eval() inside my timer I wasn't measuring quantization — I was measuring every pending operation in the graph at that moment: attention, the MLP, the whole layer stack, all attributed to the one function that happened to trigger the flush.

Notice too that the numbers are self-refuting if you read them properly: metal=False and metal=True report the same 97.8% share. A measurement that can't distinguish the two arms is measuring something other than the thing you changed.

The one genuinely useful output was incidental: calls=224 across 28 layers on an 8k prompt means 8 flushes per layer — which revealed that mlx_lm chunks prefill at 2048 tokens. That number matters later.

Lie #2 — "0.09× and 28.08×"​

Fresh microbenchmark, both flush sizes:

per-flush cost (keys+values, H=8 D=128, 28 layers)

     S    off ms     on ms  speedup |  x28 layers off        on     saved
    32    1.1334   12.4505    0.09x |           31.7ms    348.6ms   -316.9ms
  2048   19.8718    0.7076   28.08x |          556.4ms     19.8ms    536.6ms

Eleven times slower at small sizes, twenty-eight times faster at large ones. I nearly wrote a whole section theorizing about launch-overhead crossovers — there's a tidy story available where fixed dispatch cost dominates at S=32 and bandwidth savings dominate at S=2048.

The story was fiction. I had left an LLM benchmark running in the background on the same GPU.

Same benchmark, idle GPU:

     S    off ms     on ms  speedup |  x28 layers off        on     saved
    32    0.5050    0.3601    1.40x |           14.1ms     10.1ms      4.1ms
  2048    3.8337    0.6785    5.65x |          107.3ms     19.0ms     88.3ms

Both processes were fighting for the same hardware, and contention landed unevenly across runs. These weren't noisy-around-the-truth — off by 15× in one direction and 5× in the other, and they looked like a coherent narrative. That's the dangerous part. Random noise looks random. Contention produces confident, structured, wrong answers.

The downstream damage is worth showing, because the wrong numbers propagated into a wrong prediction:

--- predicted end-to-end (from the CONTENDED numbers) ---
PREFILL 8k prompt: 4 chunks x 537ms saved = 2146ms of ~21000ms  ->  10.2% faster
DECODE 240 tokens: 7 flushes x -316.9ms   = -2218ms of ~10000ms -> -22.18% "faster"

--- predicted end-to-end (from the IDLE numbers) ---
PREFILL 8k prompt: 4 chunks x 88ms saved  =  353ms of ~21000ms  ->   1.7% faster
DECODE 240 tokens: 7 flushes x 4.1ms      =   28ms of ~10000ms  ->   0.28% faster

A predicted 10.2% prefill win and a 22% decode regression, versus the truth of +1.7% and +0.28%. Had I stopped there, I'd have gone hunting for a decode regression that never existed.

Lie #3 — the benchmark that gave four different answers​

Subtler, and I think the most broadly applicable. For the same kernel and the same configuration, my attempts produced 0.31×, 1.0×, 2.16×, and 3.4×. Not scatter around a value — four different conclusions, each internally consistent.

The root cause: I was calling the function repeatedly on the same input tensor.

MLX can recognize it has already computed something and reuse the result. The reference path — eight standard, individually cacheable array ops — benefits enormously. My custom kernel benefits far less. So the "baseline" was quietly handed a shortcut the kernel couldn't take, and every extra repetition widened the gap. Layering best-of-N on top amplified it further, because best-of-N systematically selects the run where caching helped most.

I got this badly wrong. I concluded the kernel was slower, set the feature flag off, and wrote that conclusion into the code and the tests. It was only when two independently-designed methods disagreed with me that I went back:

1. Interleaved A/B — alternate on/off inside one process, so drift and thermal state hit both arms equally.
2. Rotating input pool — cycle 20 distinct tensors so nothing can be reused.

Both landed at 1.36×–2.14×, agreeing with each other and disagreeing with me. I reverted the flag and corrected the tests.

The benchmark that ships in test_kivi_quant.py now does both, and its docstring explains both traps so the next person doesn't re-derive them.

Lie #4 — the 127 MB memory saving that wasn't​

This one I caught only because I ran a third model.

Llama's peak memory, from a clean interleaved run:

  peak memory (GB): fp16=2.604  off=2.726  on=2.599

A 127 MB reduction with the kernel on. And there's a beautiful explanation sitting right there: the fused kernel eliminates MLX's intermediate tensors, so of course the high-water mark drops. Mechanistically plausible, exactly the result I wanted, and it would have made this post better.

Then the other two models came back:

Two of three moved the opposite direction. It's allocator noise — MLX's memory pool responds to allocation ordering in ways unrelated to which kernel ran, and 127 MB out of 2.7 GB sits well inside that.

There's a deeper reason this had to be noise, which I'll come back to at the end: quantize-then-dequantize cannot reduce peak memory, by construction.

The full end-to-end record​

Four generations of end-to-end measurement, in the order I ran them. I'm including the early ones because their disagreement is the point.

Generation 1 — single-shot, three prompt lengths​

First real run. Single-shot timings, no repeats, Llama-3.2-3B-Instruct-4bit, 28 layers, 8 KV heads, head_dim=128:

## PREFILL (prompt tokens/sec)
prompt tok       fp16  kernel off  kernel on  on vs off  flush/layer
       559      488.2       481.6      486.0      1.01x          512
      2059      463.4       451.7      465.7      1.03x         2016
      8209      386.4       355.9      351.6      0.99x         8160

## DECODE (generation tokens/sec, 120 tokens)
config             tok/s  vs fp16   peak MB  KV comp
fp16               46.06     100%       0.0        -
off                46.00     100%       0.0    4.99x
on                 44.24      96%       0.0    4.99x

decode kernel on vs off: 0.962x

Two problems. First, peak MB reads 0.0 — a unit bug on my side: mlx_lm.stream_generate reports peak_memory in GB, and I was dividing by 1024**2 as though it were bytes. The raw JSON shows the real values hiding at 2.48e-06.

Second, and more importantly: decode at 0.962× looks like a 4% regression. Single-shot numbers on one prompt, with no repeat structure — nowhere near enough to distinguish a real regression from thermal drift. Generation 3 is what settles it.

Generation 2 — the sync-instrumented run​

Lie #1 above. Produced the 97.83% figure, which was wrong, and the 2048-token prefill chunking discovery, which was right and load-bearing.

Generation 3 — repeated runs with medians​

Same model, but now multiple repeats per configuration reporting median/min/max, so spread is visible:

## PREFILL (prompt tok/s, max_tokens=4)
  prompt=2065 tok  (~2 chunks of 2048)
    fp16  median=  468.49  min=  448.16  max=  473.22   vs fp16 100.0%
    off   median=  451.76  min=  420.91  max=  464.51   vs fp16  96.4%
    on    median=  460.48  min=  424.83  max=  470.01   vs fp16  98.3%
    -> kernel on vs off: 1.019x

  prompt=8213 tok  (~5 chunks of 2048)
    fp16  median=  300.09  min=  280.98  max=  352.22   vs fp16 100.0%
    off   median=  297.13  min=  291.43  max=  320.70   vs fp16  99.0%
    on    median=  296.32  min=  293.69  max=  309.83   vs fp16  98.7%
    -> kernel on vs off: 0.997x

## DECODE (generation tok/s, 240 tokens, 2k prompt)
  240 new tokens
    fp16  median=   38.48  min=   38.01  max=   40.82   vs fp16 100.0%
    off   median=   40.21  min=   39.72  max=   41.10   vs fp16 104.5%
    on    median=   39.78  min=   37.22  max=   41.05   vs fp16 103.4%
    -> kernel on vs off: 0.989x

  peak memory (GB): fp16=2.604  off=2.726  on=2.599
  KV compression:   off=5.02x  on=5.02x

This table is the most useful thing I measured, and not because of the ratios. Look at the min/max columns.

At 8k prompt, the unchanged fp16 baseline — same code, same model, nothing swapped — ranges from 280.98 to 352.22 tok/s. That's a ±25% spread from thermal state and system scheduling alone, on a configuration where nothing about the code changed between runs.

Now recall the prediction: +1.7% on prefill. Looking for a 1.7% effect through ±25% run-to-run variance is like weighing a signature on a bathroom scale. No amount of care in the on/off comparison fixes that; the instrument simply doesn't resolve the quantity.

Note also that KIVI itself (off, 96.4%) is slightly slower than fp16 at 2k — the quantization work is real, it's just small. And the decode ordering (off at 104.5% of fp16, i.e. faster than no quantization at all) is a tell that we're deep inside noise, since compressing the cache cannot make decode faster than not compressing it.

Generation 1's apparent 0.962× decode regression shows up here as 0.989×, with overlapping min/max ranges. It was drift.

Generation 4 — three models, interleaved, single process​

The final protocol: one process per model, configurations interleaved rather than run in blocks, output text compared byte-for-byte between arms.

### mlx-community/Llama-3.2-3B-Instruct-4bit
  layers=28 kv_heads=8 prompt=2065 tok
  -> kernel on vs off:  prefill 1.019x (2k) / 0.997x (8k)   decode 0.989x
  -> identical text on/off: True

### mlx-community/Qwen2.5-7B-Instruct-4bit
  layers=28 kv_heads=4 prompt=2064 tok
  config  prefill tok/s  decode tok/s   peak GB   KV comp
  fp16            206.0         23.61     5.105         -
  off             203.7         23.40     5.080     4.94x
  on              208.3         23.26     5.130     4.94x
  -> kernel on vs off:  prefill 1.023x   decode 0.994x
  -> identical text on/off: True

### mlx-community/Mistral-7B-Instruct-v0.3-4bit
  layers=32 kv_heads=8 prompt=2054 tok
  config  prefill tok/s  decode tok/s   peak GB   KV comp
  fp16            143.4         20.86     4.891         -
  off             140.0         20.34     4.921     4.75x
  on              140.2         20.26     4.941     4.75x
  -> kernel on vs off:  prefill 1.001x   decode 0.996x
  -> identical text on/off: True

Consolidated:

Everything within ±2%, which given ±25% baseline variance is indistinguishable from nothing.

Qwen is the most valuable row. Four KV heads instead of eight means a completely different flush geometry, exercising different bounds-check and ragged-tail paths. It still produces byte-identical output — which is the strongest evidence that the bit-exactness work held up outside the unit tests.

Why it's invisible, and why that was predictable​

The true per-flush picture, idle GPU, Apple M4, 8 KV heads × 128 head dim, keys and values together:

Recall that mlx_lm chunks prefill at 2048 tokens. That's a happy accident: prefill flushes land exactly on the kernel's strongest case, where there's enough work to amortize dispatch and memory-traffic savings dominate. Decode flushes are always small — residual_length tokens, 32 here — the weak case.

Scale it to 28 layers, an 8k prompt (4 prefill chunks), 240 decode tokens (7 flushes):

There's the whole story, available before running a single model.

A 5.65× speedup on 1.7% of the work is a 1.7% speedup. Amdahl's law doesn't care how good the kernel is.

And 1.7% is four times smaller than the noise floor of the measurement. The end-to-end result wasn't a disappointment — it was arithmetic, and I could have computed it in ten minutes before writing any Metal at all.

So why does the kernel ship?​

Given that it's invisible end-to-end, why merge it?

• It's free. Bit-exact output, no regression on any model, 84 dedicated tests, byte-identical generations across three architectures. It never makes anything slower.
• It removes a floor. Quantization is ~1–2% of runtime now. If the surrounding work gets faster — better attention kernels, better matmuls — that share grows. Fixed costs matter more as everything else shrinks.
• Op-level wins are real even when invisible. 1.40× and 5.65× are honest measurements of the operation. That the operation is a small slice of the whole is a separate fact, and both belong in the report.

But the real reason to stay clear-eyed: the kernel was never where the memory win lived.

KIVI as implemented does quantize-then-dequantize — it computes the compressed representation and immediately expands it back to fp16 for attention. The 4.75×–5.02× compression is real arithmetic, but it is an accounting result, not a storage result. The tensor sitting in memory is still fp16.

This is also the structural reason Lie #4 had to be noise: if nothing is stored in compressed form, no kernel that computes the compression faster can reduce the high-water mark. I should have known the 127 MB was suspect on those grounds alone, before the other two models contradicted it.

To actually reduce memory you need two more things:

1. Packed storage — keep quantized codes as uint8, never materialize the fp16 reconstruction.
2. Dequant-in-SDPA — teach attention to read packed codes directly, so expansion happens in registers and never in RAM.

That's where both the memory win and the real speedup live, because it doesn't fuse 1–2% of the work — it shrinks the tensors every other operation has to move.

The kernel was step one. It was worth doing. It just isn't the point.

What I'd take away​

If you're writing GPU kernels against a framework like MLX or PyTorch:

Match the thread mapping to the memory layout, not to intuition. The strided access pattern needed less machinery than the contiguous one — no reduction, no barriers, no transpose. My first instinct, transposing to make the layout "nice," was the version that lost.

Bit-exactness is a debugging tool, not just a correctness bar. FMA contraction, rounding mode, and padding semantics all produce plausible output. A tolerance test passes all three. Demanding identical output turned three silent behavioral differences into three failing tests with obvious causes.

Never specialize on a value that grows. Baking sequence length into a JIT header compiles one shader per token. The symptom was an indefinite hang; the fix deleted code.

Estimate the ceiling before you optimize. Ten minutes with Amdahl's law would have predicted the end-to-end result up front. It wouldn't have changed the decision to build it — but it would have set the expectation correctly, and I'd have spent my time on the parts that were load-bearing.

Measure your noise floor before your effect. The single most useful number in this entire project was ±25% — the run-to-run spread of an unchanged baseline. Without it, every ratio in every table is unfalsifiable.

Be most suspicious when the number is good. All four bad measurements came with satisfying stories. 97.83% "proved" the work mattered. 28.08× had a tidy overhead-crossover explanation. The 127 MB saving had a clean mechanism. Every one was wrong, and the plausible explanation is what let each survive as long as it did.

And: run the third model.

Kernels live in veloxquant_mlx/metal/src/, tests in veloxquant_mlx/tests/metal/test_kivi_quant.py. See the KIVI algorithm reference for the method itself, and Metal kernels for how kernels are dispatched library-wide. KIVI: Liu et al., ICML 2024. All measurements on an Apple M4 with MLX.