QuantizedKVCache shows higher peak memory than fp16 KVCache (worse at longer contexts)

[gschaefer3]

Summary

QuantizedKVCache shows higher peak MLX memory than plain fp16 KVCache in stream_generate, at every bit width and context length tested, with the gap widening as context grows — the opposite of the expected result. Filing this as data in case it's useful alongside the compute-stays-float design gap already noted on #941.

Environment

• mlx-lm 0.31.3, mlx 0.31.2
• Apple M4 Max, 128GB unified memory, macOS 27.0
• Model: mlx-community/Llama-3.2-3B-Instruct-4bit

Reproduction

import mlx.core as mx
from mlx_lm import load, stream_generate
from mlx_lm.models.cache import make_prompt_cache, QuantizedKVCache

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

filler = ("The quick brown fox jumps over the lazy dog. " * 2000)
ids = tokenizer.encode(filler)[:32000]
prompt = tokenizer.decode(ids) + "\n\nSummarize the above in one sentence."

# fp16 case (separate process from quantized case, to rule out allocator state bias)
cache = make_prompt_cache(model)
# quantized case
base = make_prompt_cache(model)
cache = [QuantizedKVCache(group_size=64, bits=8) if type(c).__name__ == "KVCache" else c for c in base]

mx.reset_peak_memory()
for _ in stream_generate(model, tokenizer, prompt=prompt, max_tokens=30, prompt_cache=cache):
    pass
print(mx.get_peak_memory() / 1e9, "GB")

Each case run in its own fresh process (not just mx.reset_peak_memory() between cases in the same process), to rule out allocator-state bias between runs. Results were identical either way.

Results

context	case	peak MLX memory	decode speed
8,000 tok	fp16	3.46 GB	3.2 tok/s
8,000 tok	QuantizedKVCache 8-bit	4.87 GB (+1.41 GB)	2.6 tok/s
32,000 tok	fp16	4.72 GB	1.0 tok/s
32,000 tok	QuantizedKVCache 8-bit	7.10 GB (+2.38 GB)	0.7 tok/s
32,000 tok	QuantizedKVCache 4-bit	6.53 GB (+1.81 GB)	0.6 tok/s

Quantized cache used more peak memory in every case, at every bit width, and the gap widened with context length rather than shrinking. Decode throughput was also consistently lower under quantization.

Possible explanation

Per @angeloskath's comment on #941, QuantizedKVCache computes attention in float and only stores quantized — so the quantize/dequantize scratch buffers needed during prefill attention appear to cost more, in peak-memory terms, than the storage format saves, at least at these model/context sizes. mx.get_peak_memory() captures the max across the whole run (prefill included), so this may not reflect steady-state resident memory after prefill — that wasn't isolated separately here, and might tell a different story for very large contexts or larger models where cache dominates total memory more. Flagging as-is since the peak-memory number is what most callers will actually observe.

Happy to run additional configurations (different models, isolating post-prefill steady-state memory rather than run-peak, batched generation) if that would help narrow this down — let me know what would be most useful.

[PhilipJohnBasile]

Thanks for the reproducible repro — the numbers read as a strong signal, and we can reconcile them with measurements of ours that point the opposite direction in isolation:

• 8k-context: fp16 3.46 GB vs QuantizedKVCache 4.87 GB (+1.41 GB, +40.8%)
• 16k-context: fp16 4.72 GB vs quantized 7.10 GB (+2.38 GB, +50.5%)
• 32k-context: quantized path appears worse by design here as written (no absolute in your table, but from the first two points it’s a widening gap as context grows)

I checked the local implementation and this behavior is mechanically plausible from where the memory is spent:

1. generate_step calls maybe_quantize_kv_cache(...) during prefill:
   • QuantizedKVCache conversion only runs when kv_bits is explicitly set (kv_bits is not None) in generate.py (maybe_quantize_kv_cache).
   • It is an explicit/opt-in transform, not an automatic policy engine.
2. QuantizedKVCache.update_and_fetch allocates quantized storage with extra metadata and does extra quantization work per update:
   • It computes el_per_int = 8 * uint32_size // bits, then allocates quantized payload + scale + offset arrays via init_quant.
   • On growth (prev + num_steps > keys.shape), it resizes in step=256 chunks (new_steps = ceil((step+num_steps)/step) * step) and copies old cache into new buffers via expand_quant.
   • It always calls mx.quantize(keys, ...)/mx.quantize(values, ...), then writes into three quantized buffers per side (data, scale, offset) via self.keys[i][..., prev:self.offset, :] = ....
   • Compared to KVCache, this is more than 1:1 data structures (extra pack/scales/offset and temporaries) while sharing the same step-sized growth behavior.
3. Attention dispatch path for quantized cache is still multi-op + materialized intermediate:
   • scaled_dot_product_attention selects quantized_scaled_dot_product_attention when cache has .bits (mlx_lm/models/base.py).
   • That path is quantized_matmul -> softmax -> quantized_matmul, so even if cache bytes are lower, each attention step carries intermediate allocation pressure different from the fp16 fused SDPA path.

We also checked our own engine-cert benchmarks for a direct memory comparison:

• Our fused-SDPA benchmark (Qwen3-32B-4bit, group_size=64, bits=8, peak reset between phases) reports:
  • cache resident (cache_nbytes, MiB): 0.5K 2.0 → 1.1, 4K 16 → 8, 8K 32 → 17, 16K 64 → 34, 32K 128 → 68 (-46.875% at each sampled context)
  • attention-call peak (attention_call_peak_bytes, MiB): 0.5K 2.06 → 1.26, 4K 18.17 → 9.65, 8K 34.17 → 20.18, 16K 68.21 → 39.06, 32K 132.19 → 72.58 (~-42% to -47%)

That artifact isolates prebuilt cache + warmed quantized inputs with peak reset between phases, so it says: QuantizedKVCache can be lower in isolated KV/cache and call-only measurements, but your reported full generation peak can still go higher when prefill/resize timing is included.

So I’d frame this as a measurement-interval issue first, not necessarily a regression in steady-state KV footprint:

• Your mx.get_peak_memory() in stream_generate includes prompt prefill and generation in one continuous process lifetime.
• The quantized path does extra quantization work during prefill and may keep additional temporaries alive at resize boundaries.

Suggested falsifier (single script you can run):

1. Add a split run that forces three phases in separate processes or with forced barriers:
   • (A) prompt prefill only (max_tokens=0) with QuantizedKVCache.
   • (B) decode-only continuation (same initialized cache, same prompt, max_tokens=N).
   • (C) same two phases with fp16 cache.
2. In each phase do mx.reset_peak_memory() immediately before, then capture mx.get_peak_memory() right after the phase completes.
3. Repeat with prompt lengths around resize boundaries (e.g., 255/256/257 tokens, plus 8191/8192/8193) and with quantized_kv_start=0 vs quantized_kv_start=1.

A falsifier outcome I’d consider decisive:

• If prefill-phase-only peaks are the only source of the inversion while decode-only peaks stay -X% vs fp16, then it’s likely boundary/ephemeral quantization-allocation behavior in QuantizedKVCache.update_and_fetch and/or prefill scheduling.
• If full run stays inverted even when prefill and decode are separated with cache-state snapshots, then we likely have a true steady-state allocator-lifetime effect worth patching in cache growth policy first.

[pierre427]

Corroborating datapoint from an M5 (128GB), mlx 0.32.0 + mlx-lm current main, Qwen/Qwen3-0.6B (bf16), your reproduction pattern (stream_generate, max_tokens=32, peak measured from a post-load mx.reset_peak_memory()):

context	fp16 KVCache	kv_bits=8	kv_bits=4
8,192	2.78 GB	3.64 GB	3.37 GB
16,384	3.63 GB	5.74 GB	5.24 GB
32,768	5.53 GB	9.56 GB	9.09 GB

Same direction as your M4 Max numbers and slightly worse at the top end (+73% for kv8 at 32k), with the gap widening with context, so it does not look machine-specific. One related operational observation from serving on the same machine (2026-07-08, different codebase revision): served kv8 preserved exact answers on 4k/16k/32k probes but roughly halved decode throughput for long-context/short-output workloads — consistent with the compute-stays-float design gap referenced from #941.

[PhilipJohnBasile]

@pierre427 Thanks — that's a clean corroboration. Two things stand out from your table: the gap widens with context (kv8 at +73% by 32k), and it reproduces on different silicon from the original report, so it does not look machine-specific.

Worth noting the shape for anyone triaging: the quantized path is not merely failing to save memory, it is costing progressively more as context grows — which is the opposite of the property the feature exists for. That points at transient/allocation behaviour during quantize-and-concat rather than the steady-state cache footprint.

[rajanshxrma]

Ran the split prefill/decode experiment @PhilipJohnBasile proposed. Short version: it's a transient prefill/resize artifact, not a steady-state effect — decode-only peaks are actually lower for the quantized cache than fp16 at every context length I tested, including the two boundary triples.

Setup

• mlx-community/Llama-3.2-3B-Instruct-4bit, QuantizedKVCache(group_size=64, bits=8) built via to_quantized() at offset 0 (quantized from the start, same as the original repro).
• Each (cache type × context length × phase) cell is its own fresh python3 subprocess — no shared allocator state between measurements.
• Prefill-only: run the model over the full prompt once (same chunked loop generate_step uses, prefill_step_size=2048), reset peak memory before, capture right after. This populates the cache and samples the first token — no additional generation.
• Decode-only: run that same prefill unmeasured, then reset peak memory, then run 30 single-token decode steps off the populated cache, capture peak after.
• Context lengths: 255/256/257 and 8191/8192/8193 — the resize boundaries, since both KVCache and QuantizedKVCache grow in 256-token chunks (step = 256 in models/cache.py).

Results

ctx tokens	fp16 prefill GB	q8 prefill GB	prefill Δ	fp16 decode GB	q8 decode GB	decode Δ
255	2.213	2.166	-0.047 (-2.1%)	1.908	1.851	-0.057 (-3.0%)
256	2.215	2.168	-0.047 (-2.1%)	1.908	1.851	-0.057 (-3.0%)
257	2.193	2.294	+0.101 (+4.6%)	1.868	1.841	-0.028 (-1.5%)
8191	3.284	4.287	+1.003 (+30.5%)	3.160	2.521	-0.639 (-20.2%)
8192	3.284	4.288	+1.004 (+30.6%)	3.160	2.521	-0.639 (-20.2%)
8193	3.284	4.289	+1.005 (+30.6%)	2.780	2.335	-0.445 (-16.0%)

Reading it

The inversion lives entirely in the prefill numbers, and it's real — flips sign exactly at 257, not 255/256 (measured, all 6 rows check out: 255/256 negative, 257 onward positive). From generate_step's chunking behavior: it processes the prompt in chunks up to 2048 tokens, and each chunk that pushes offset past the cache's current capacity triggers a resize. 255 and 256 tokens fit in a chunking pattern that needs exactly one resize; 257 needs two resizes back-to-back in the same prefill call — one for the main chunk, one more when the final prompt token pushes past that. Per cache.py, QuantizedKVCache.update_and_fetch's resize path does a lot more than KVCache's: it's a tree_map over 6 arrays (packed codes + scales + biases, for both K and V) versus a single zero-pad concat, plus it immediately re-quantizes the incoming chunk. That's a much heavier transient allocation if I'm reading the resize path right, and it would explain why q8 prefill flips positive right where a second resize gets forced into the same call, and why it holds at a stable +30% once you're at 8k tokens, where every ~2048-token chunk forces its own resize+quantize cycle during prefill.

Decode is the opposite story, cleanly. At every context length here, quantized decode peak is below fp16 decode peak — by 1.5-3% at small context (where the fixed ~1.7GB of model weights dominates and the KV cache is a rounding error), and by 16-20% at 8k context where the cache actually matters. That's the memory saving KV quantization is supposed to deliver, showing up exactly where you'd expect it to: steady-state generation.

One thing worth flagging since it's not obvious and would confound anyone re-running this: the decode-phase peak itself depends on whether the fixed 30-token decode window happens to straddle one of the cache's own 256-token resize boundaries, independent of cache type — the fp16-side drop from 3.160GB (8191/8192) to 2.780GB (8193) is measured and real, ~0.38GB, not noise. My read of why, from replaying the resize arithmetic in update_and_fetch (worth double-checking against the actual code, not just this JSON): at ctx=255/256, cache capacity after prefill is 256, and decode's 30 steps push offset past that, forcing one resize mid-decode; at ctx=257, prefill itself already lands on a 512-capacity cache with room to spare, so decode never resizes. Same pattern at 8191/8192 vs 8193. If that's right, it's the same transient-resize effect as the prefill story above, just also present (symmetrically, for both cache types) in decode when the window happens to cross a boundary. Anyone benchmarking decode-only peak for this issue should either pick N far from a multiple of 256 relative to prompt length, or report which side of a boundary they're on.

Net: this is the transient/allocation-artifact story, not a genuine steady-state one, and it's specifically the QuantizedKVCache resize+quantize path during prefill/chunk boundaries. Full driver + worker scripts and raw JSON results available if anyone wants to re-run or extend this (more context lengths, other kv_bits, other models).

[PhilipJohnBasile]

@rajanshxrma That's the pre-registered branch (a), by the criteria as declared and with nothing needing reinterpretation: the sign flip lands exactly at 257, and decode-only quantized peaks come in below fp16 at every tested context. Your run, your find — and the confound you caught yourself (the decode window straddling a 256 boundary moves peaks for both cache types: 0.380 GB on the fp16 side, 0.186 GB on q8, by your table) is exactly the artifact-hunting this thread needed. Clean execution: fresh-process cells, boundary triples, and reporting which side of the boundary each cell sits on.

One correction to my own comment upthread first: I labeled the 4.72 → 7.10 GB row "16k-context." The issue's table says 32,000 tok. My error; the percentages were computed on the right numbers, the label was wrong.

I verified the resize-path reading against cache.py on current main this morning, and it holds as far as it goes:

• QuantizedKVCache.update_and_fetch's growth path concatenates zeros onto six arrays (packed codes + scales + biases, K and V), then mx.quantizes the incoming fp16 chunk, then does six slice-writes — vs. KVCache's two-array concat and two writes. The incoming chunk is still fp16 at that moment, so the transient holds the fp16 chunk, the quantized triples, and old+new copies of the buffers together.
• Your 257 flip matches the code: generate_step prefills prompt[:-1], so 255/256-token prompts fit the initial 256 capacity (init path only), and 257 is the first case where the held-back final token forces the expand path on populated buffers.
• One thing the code makes worse than your write-up claims: capacity grows by roughly one chunk per resize (new_steps = ceil((step+num_steps−1)/step)·step — no geometric growth), so at long context every ~2048-token prefill chunk forces a full six-array copy. That's quadratic copy traffic over prefill, which fp16 pays too — but with two arrays and no quantize step.

What your split can't yet separate — and a discriminating follow-up if you want it: quantized prefill differs from fp16 in two ways at once. Resize churn is one. The other is the attention path itself: with a quantized cache, prefill attention runs the unfused quantized_matmul → softmax → quantized_matmul route, which materializes the chunk×context score intermediates that fp16's fused SDPA never allocates — and that transient also scales with context. Your data implicates expansion (the 257 flip is unambiguous), but the stable +30.5% at your 8k cells could be either contributor or both; nothing measured between 257 and 8191 tells us the mix. The cheap discriminator: presize the quantized cache to prompt length before prefill (capacity is known at that point), which eliminates resizes entirely — whatever inversion survives is the attention path. If presizing kills most of it, the fix surface is growth policy (geometric growth, or prompt-aware presizing in generate_step); if not, it's the unfused prefill attention. Since it's your harness and your find, that experiment — and any patch that follows — is yours to lead if you want it; happy to run cells on our M5 rigs either way, and yes please on committing the scripts + raw JSON.

Two scope notes so this thread's numbers don't get misread later:

1. "Prefill-scoped transient" does not mean harmless. mx.get_peak_memory()'s number is the one that OOMs a real stream_generate call — a +30% peak at 8k (and the original report's +50% at 32k) is a real user-facing cost even though the steady-state cache is genuinely smaller.
2. Decode memory and decode speed are separate results and land in opposite directions. Your decode cells show quantized lower-memory; independently, quantized decode measures slower — this thread's own observations (3.2→2.6 tok/s above; roughly halved in the M5 serving datapoint), and our pre-registered dense-32B gate (−7.4/−3.1/−2.7% at 0.5K/4K/16K, quality clean — receipts in trailbrake's README). The byte accounting reconciles them: on a 4-bit dense 32B, KV is only ~19% of decode-step reads, so halving KV saves ~9% of bytes while the unfused multi-op path costs more than that. Smaller resident cache, slower step — both true, different questions, and the 3B memory data and 32B throughput data shouldn't be quoted as one benchmark.

[PhilipJohnBasile]

@rajanshxrma Following through on "yours to lead if you want it" with two concrete things, so the offer costs you nothing to accept:

1. The discriminating experiment needs no patch. step = 256 is a class attribute on QuantizedKVCache (cache.py), so your existing harness can force a zero-resize prefill with one line before the prefill loop:

for c in cache:
    c.step = 33024  # ≥ prompt length, rounded up to a multiple of 256

First allocation then covers the whole prompt (and the decode window), so update_and_fetch never takes the expand path. Pre-registered read, same style as before: whatever prefill inversion survives presizing is the unfused attention-score path; whatever disappears was resize churn. Your 8k cells rerun with only that line changed would settle the mix.

2. A starting-point diff for the eventual fix, yours to take, adapt, or discard — geometric growth, scoped to the quantized cache since that's where the six-array copy lives:

-            new_steps = (self.step + num_steps - 1) // self.step * self.step
-            shape = (B, n_kv_heads, new_steps)
+            current = 0 if self.keys is None else self.keys[0].shape[-2]
+            needed = prev + num_steps
+            # Grow geometrically: capacity at least doubles per resize, so a
+            # long prefill pays O(log n) six-array copies instead of one per
+            # ~2048-token chunk.
+            new_total = max((needed + self.step - 1) // self.step * self.step,
+                            2 * current)
+            shape = (B, n_kv_heads, new_total - current)

(Both branches stay correct: init sees current=0 and allocates new_total; expand appends the difference.) Known trade-off to weigh in review: doubling can overshoot resident memory by up to ~2× cache size at large context — a growth cap, or prompt-length-aware presizing in generate_step's prefill loop (the prompt length is known there), avoids that at the cost of touching two files. Maintainers may also want a view on whether fp16 KVCache gets the same treatment (same linear growth, cheaper per-copy). Your call which shape to propose.

The offer, plainly: the experiment and the PR are yours to lead — your falsifier run is what moved this issue, and the fix should carry your name. We'll co-sign, rerun cells on M5 hardware for the PR's evidence table, and write it up with the same pre-registered decision-rule format. If you'd rather hand it back, say the word and we'll drive it with you as co-author. Either way the raw JSON you offered would be great to have committed somewhere linkable.

[PhilipJohnBasile]

The M5 cells we promised — offered as input to your patch, lead unchanged.

Ran the presize discriminator on an M5 Max (mlx 0.32.0, mlx-lm 0.31.3 from PyPI, Llama-3.2-3B-Instruct-4bit, your split-phase protocol: fresh process per cell, peak reset at phase start, 8192-token prompt, 30 decode steps; presize = per-instance cache.step set to cover prompt+decode, so the expand path never runs):

cell	fp16	q8	inversion
prefill, unpatched	3.889 GB	4.809 GB	+0.920 GB
prefill, presized (zero resizes)	3.922 GB	4.808 GB	+0.885 GB
decode-only, unpatched	3.272 GB	2.832 GB	−13.4% (q8 lower)

Both pre-registered controls passed before reading the comparison: the unpatched cells reproduce the inversion this issue reports (and logged 3–4 live capacity-growth events each), and the presized cells logged zero growth events at higher capacity — the resizes were genuinely eliminated.

On this rig, the rule reads: presizing removed only 3.8% of the inversion. That's far below the more-than-half bar that would have pointed at resize churn — so here, the transient lives overwhelmingly in the unfused quantized attention path (the quantized_matmul → softmax → quantized_matmul route materializing chunk×context score intermediates that fp16's fused SDPA never allocates), not in update_and_fetch's growth behavior.

Which means a correction to my own suggestion upthread, struck rather than deleted: the geometric-growth diff I offered as a starting point for the fix — by these cells, growth policy would not materially reduce the peak this issue reports. It stands only as a copy-traffic improvement. The fix surface for the memory number looks like the prefill attention path itself (fused or blockwise quantized attention), which is a materially different — and harder — patch than the one I sketched. Better to know before anyone builds it.

Your rerun is still the one that settles it: this is one rig, one model, one mlx version, and the presize one-liner drops straight into your existing harness. If your cells agree, the patch conversation moves to the attention path with two rigs of evidence; if they don't, that disagreement is the most interesting result yet.

_{Cells run and receipted as a claim capsule (frozen rule, controls, hash-pinned receipt, rerunnable): https://github.com/PhilipJohnBasile/scrutineer — capsules/quantizedkv-prefill-discriminator.}

[rajanshxrma]

@PhilipJohnBasile Ran the presize discriminator on my machine (an M1 MacBook Pro, 16GB -- correcting a mislabel in my first post here, which wrongly said M4 Max), plus a follow-up experiment. Full harness, raw JSON, and a correctness-verified prototype are here: https://github.com/rajanshxrma/mlx-kvcache-1587

Presize discriminator — corroborates your M5 result, doesn't overturn it.

context	unpatched inversion	presized inversion (0 resize events)	% closed
8191	+1.003 GB	+0.988 GB	1.5%
8192	+1.004 GB	+0.989 GB	1.5%
8193	+1.005 GB	+0.990 GB	1.5%

Resize events instrumented directly (not just inferred from memory): 3-4 unpatched → 0 presized at every context, confirming resizes were genuinely eliminated. 1.5% closed here vs. 3.8% on your M5 — two independent rigs, same conclusion by your pre-registered rule: nowhere near the "more than half" bar, so this isn't a resize-churn effect. Agreed on striking the geometric-growth fix as the memory fix.

Follow-up: chunk-size scaling pins down the mechanism directly.

If the transient is the unfused (L_chunk × context) scores matrix in quantized_scaled_dot_product_attention, quantized prefill peak should scale with prefill_step_size at fixed context; fp16's fused kernel shouldn't. Ran both, context=8192:

chunk size	quantized peak GB	fp16 peak GB (control)
512	3.072	3.175
1024	3.821	3.290
2048	4.288	3.284

Quantized scales strongly with chunk size; fp16 stays flat. At chunk=512 the inversion disappears — quantized drops below fp16. The predicted size of one layer's scores tensor (n_kv_heads × n_repeats × L_chunk × context × 4 bytes, 8×3×128 head config for Llama-3.2-3B-Instruct-4bit) gives a predicted delta between chunk 2048 and 512 of 1.208 GB; measured delta is 1.216 GB. That's the mechanism, not just consistent with it.

Practical mitigation available today, no code change: smaller prefill_step_size trades prefill latency for lower quantized-attention peak — worth knowing for anyone hitting OOM on long-context quantized generation right now.

Blockwise/online-softmax prototype: correct, and it works, but not as a Python patch.

Wrote a tiled version of quantized_scaled_dot_product_attention (flash-attention-style, running softmax over the context dimension). Verified tensor-for-tensor against the existing quantized SDPA across causal/no-mask/boolean/additive masks and GQA (n_repeats=3) — max abs diff ≤1e-6 across all 8 cases, ALL PASSED (receipt, reproducible via python3 blockwise_quant_attn_prototype.py).

Memory at ctx=8192 (baseline unfused quantized 4.288 GB, fp16 3.284 GB):

block size	naive (lazy eval)	forced per-block mx.eval
256	5.620 GB (worse)	2.856 GB
512	5.793 GB (worse)	2.969 GB
1024	6.272 GB (worse)	3.309 GB

Naive tiling is worse than the current implementation — MLX's lazy evaluation keeps every block's intermediates alive until the final eval, so a Python loop pays the sum of all blocks, not the peak of one. Forcing mx.eval per block fixes that (2.856 GB at block=256, below fp16), but costs ~2x prefill latency from per-block dispatch overhead.

Net: tiling is a correct and sufficient fix in principle, but a pure-Python implementation isn't a mergeable patch — it needs to happen inside a fused kernel (the quantized analog of mx.fast.scaled_dot_product_attention) to get the memory benefit without the latency cost. That's a bigger lift than what either of us sketched originally, and probably belongs as a fused-kernel proposal rather than an mlx-lm-side PR. Filing this as where the evidence points rather than forcing a smaller patch that wouldn't actually fix the reported number.

Everything above — harness, all raw JSON, the chunk-size sweep, and the blockwise prototype + correctness check — is at the repo linked above if useful for anyone picking up the fused-kernel side of this.

One open item on my end I haven't resolved: my from-scratch harness (the original script wasn't saved) doesn't reproduce the small +4.6% inversion you'd have seen reported at context=257 — I get the resize-event timing exactly right there (0/0/1 expand events at 255/256/257, matching the boundary logic) but not the sign of that particular small delta. Flagging for transparency since it doesn't affect anything above (two orders of magnitude smaller than the 8k result, and orthogonal to the chunk-scaling/blockwise findings).

[PhilipJohnBasile]

@rajanshxrma That closes it, and then some. Two rigs, resize events instrumented directly on both, same answer under the rule as frozen: 1.5% and 3.8% closed — resize churn is refuted as the cause. And your chunk-size sweep does something our cells couldn't: 1.208 GB predicted from the scores-tensor arithmetic vs 1.216 GB measured turns "consistent with the attention path" into is the attention path. The capsule is now at verdict with joint credit — your original falsifier run, the M4 M1 MacBook Pro (16 GB) cross-rig cells (label corrected per your edit above -- and the correction strengthens the result: M1/16 GB to M5 Max/128 GB spans chip generations and memory classes, and reproducing a +1 GB prefill transient on a 16 GB machine makes the OOM stakes concrete), the chunk sweep, and the prototype are half its evidence: https://github.com/PhilipJohnBasile/scrutineer/blob/main/capsules/quantizedkv-prefill-discriminator.json

Three things in your comment deserve flagging so they don't get lost:

1. The prefill_step_size mitigation is the immediately useful result — anyone OOMing on long-context quantized generation today can trade prefill latency for peak memory with an existing knob. That's a two-line docs note in mlx-lm that would help people now, and it's yours to file if you want it on the record.

2. Your lazy-eval tiling result is a negative finding worth its receipt. "Naive Python tiling is worse because lazy evaluation keeps every block's intermediates alive" is exactly the kind of thing the next person would rediscover expensively — and your forced-mx.eval row (2.856 GB, below fp16, ~2× latency) bounds the fix's value before anyone writes Metal.

3. On the fused-kernel road: it may be revive-and-wire, not write-from-scratch. There's a stalled fused quantized-SDPA kernel in MLX core — ml-explore/mlx#3026 (CC-Yeh's 2-pass split-K flash structure, four quant modes). We benched it head-to-head earlier this month: correct at head_dim 128 (100% argmax parity vs fp16 and vs the unfused path), faster than the unfused composition, though never beating fp16 on dense-32B decode — which is exactly why it stalled without a use case. Your result supplies the use case it was missing: prefill peak memory, where the unfused path's materialized scores are the whole problem. Whether its kernel handles multi-row (chunked-prefill) query shapes efficiently is the open question — we'll post our head-to-head data on that PR so whoever picks it up starts from evidence.

On your 257 transparency flag: worth noting for the record that the +4.6% figure originated in your first harness's run — our comment only quoted it. Given your rewritten harness reproduces the resize-event timing exactly but not that sub-noise delta, and the 8k results carry everything, recording it as "original-harness-only, unreproduced, immaterial to the verdict" seems right — and flagging it unprompted is the conduct that makes the rest of the thread trustworthy.

[rajanshxrma]

Read the capsule. The joint-credit framing there is fair — your resize-churn falsifier is what actually pinned this to the attention path in the first place; the chunk-size sweep just turned that into a number.

On prefill_step_size: agreed, and yes, I'll file it. Worth noting the README doesn't currently document --kv-bits/--kv-group-size at all, so this'll be a new short paragraph rather than an edit to something existing — I've got it drafted on a branch on my fork already and I'll open it as its own PR, separate from this thread so it doesn't get tangled up with anything else.

Lazy-eval result — glad that one reads as useful on its own. It cost an afternoon to find and about four lines to state, which is the trade you want out of a negative finding.

On #3026, I went and read the PR directly rather than working off the summary — diff, review comments, and the rest of the thread. A few things worth putting on the record:

Current state: still open, unmerged, shows CONFLICTING against main in GitHub's own mergeability check. Reading the full history rather than just the recent end of it, it looks less like a maintainer decision against it and more like it fell off the queue: CC-Yeh said it was "ready for review" on Feb 7, then posted two unanswered pings (Feb 25, Mar 2) before looping in other maintainers by name, and after that the thread is almost entirely third parties — a couple of rebase branches, a benchmark run on an M3 Pro, a proposed multi-row dispatch patch on a fork, and your own reproduction from the 19th. No maintainer comment since Jan 29. So "stalled" is accurate, but it reads like a bandwidth problem more than a rejected design.

Interface match: genuinely close to what mlx-lm already calls. The PR adds fast::quantized_scaled_dot_product_attention(queries, keys, key_scales, key_biases, values, value_scales, value_biases, scale, mask, sinks, group_size, bits, mode, causal) in mlx core. mlx-lm's own quantized_scaled_dot_product_attention in models/base.py takes queries, q_keys, q_values, scale, mask, group_size, bits, where q_keys/q_values are already (data, scales, biases) tuples. Unpack those and it's the same argument list plus a mode/causal flag. If this ever lands in a release, wiring scaled_dot_product_attention's dispatch to call it instead of today's unfused composition looks like a call-site swap, not a rewrite.

The open question — multi-row/chunked-prefill shapes: this is the part I didn't expect going in. As written, this kernel would not fire on a prefill call at all. QuantizedScaledDotProductAttention::use_fallback() returns true — meaning it falls back to the unfused path — whenever query_sequence_length > 8 (there's also a query_sequence_length * gqa_factor > 32 cap). A chunked-prefill call passes query lengths in the hundreds to low thousands (2048 by default in mlx-lm), nowhere near that gate. And the fallback it falls back to is, line for line, the same quantized_matmul → softmax → quantized_matmul composition that produces the inversion today. So merged exactly as it sits now, this PR changes nothing about the prefill number — it only ever engages for single-token (or small-batch speculative) decode.

There has been movement on that specific gate: a fork patch posted on the same thread back in May moves the query-sequence dimension from the threadgroup z-index to the grid z-index specifically to lift that cap, up to qsl=32 in their numbers. That's real progress but still an order of magnitude short of prefill-chunk-sized queries. From reading the kernel (I haven't built or benchmarked it, to be clear) — the 2-pass split-K structure here reads as shaped for "one query, many KV blocks," and a few hundred queries against a few thousand KV positions starts to look like a differently-tiled problem — closer to a quantized flash-attention forward pass than a wider dispatch of the same kernel. I could be wrong about exactly where that line sits; that's a real-kernel-work question, not a read-the-diff one. But "revive and wire" undersells the gap between qsl=32 and what chunked prefill needs.

Net: I'd frame it as — strong evidence the decode-path kernel is correct and worth reviving (your reproduction, plus the affine-vs-fp16 decode numbers posted on the PR thread in May, which show quantized ahead at some head-dim/context combos and behind at others), and my prefill numbers are a legitimate reason to revive it. But closing this issue needs that qsl gate genuinely extended and probably re-tiled, not just relaxed. Post your head-to-head data whenever it's ready — especially curious if any of it goes past qsl=8.

On the 257 framing — "original-harness-only, unreproduced, immaterial to the verdict" says what I was reaching for more precisely than I did. I'll tighten the README to that language.

[PhilipJohnBasile]

@rajanshxrma Verified your read against the #3026 diff before repeating it anywhere: QuantizedScaledDotProductAttention::use_fallback returns true at query_sequence_length > 8 (plus the qsl > ksl and head-dim gates), and the fallback target is the same unfused composition measured in this thread. So confirmed — merged as written, that PR moves nothing on the prefill number, and my "revive and wire" line above deserved exactly the skepticism you applied to it. The open question I flagged is now closed in the direction of real kernel work: a re-tiled quantized flash-style forward for prefill-shaped queries, not a wider dispatch of the existing structure.

To answer your question directly: no — none of our head-to-head data goes past qsl=8. It's decode-shaped (qsl=1 throughout: 2×N=9 interleaved passes plus a real-model correctness pass at head_dim 128), so it speaks to the decode kernel's correctness and its decode-perf envelope, and is silent on prefill shapes. The comment we post on #3026 will carry precisely that scoping — evidence for reviving the decode path, plus this thread's evidence that the prefill fix is a separate, differently-tiled work item. Glad the docs PR is happening; that's the fastest real relief anyone in this issue gets.