RotatingKVCache.to_quantized() raises NotImplementedError — blocks --kv-bits on Gemma 4's sliding-window attention layers

[chrislyons]

Symptom

mlx_lm/models/cache.py:552 (confirmed on current main, and on the 0.31.3 release):

def to_quantized(self, group_size: int = 64, bits: int = 4) -> QuantizedKVCache:
    raise NotImplementedError("RotatingKVCache Quantization NYI")

mlx_lm.server --kv-bits N crashes on the first inference request with this NotImplementedError, even though the server itself starts cleanly and reports healthy. The crash only surfaces once a request actually needs to quantize a rotating-window cache slot.

Affected models

Gemma 4 (both E4B and the 26B-A4B MoE variant) use sliding-window attention on the majority of transformer layers — 35 of 42 layers in the 26B-A4B configuration. Sliding-window layers are backed by RotatingKVCache, which has no quantized implementation. Full-attention layers use the (already-quantizable) standard KVCache, so the crash is specific to models with a mixed/hybrid attention schedule like Gemma 4, not to dense-attention models.

Repro

mlx_lm.server --model <gemma-4-26b-a4b-path> --kv-bits 8 --port 41966
# server starts, /health returns 200
# first /v1/chat/completions request → NotImplementedError: RotatingKVCache Quantization NYI

Current workaround (in production use since ~FCT093)

We monkeypatch mlx_lm.models.cache.make_prompt_cache (and mlx_lm.server.make_prompt_cache) before the server starts, forcing every cache slot — rotating or not — into an 8-bit QuantizedKVCache directly, bypassing the broken RotatingKVCache.to_quantized() path entirely:

from mlx_lm.models.cache import QuantizedKVCache, RotatingKVCache

def make_prompt_cache_quantized(model, max_kv_size=None, group_size=64, bits=8):
    return [QuantizedKVCache(group_size=group_size, bits=bits) for _ in model.layers]

import mlx_lm.models.cache
mlx_lm.models.cache.make_prompt_cache = make_prompt_cache_quantized
import mlx_lm.server
mlx_lm.server.make_prompt_cache = make_prompt_cache_quantized

This works but means --kv-bits on the CLI is silently ignored for these models — anyone hitting this for the first time gets a runtime crash instead of a clear "unsupported for this cache type" error at startup.

Ask

Either implement RotatingKVCache.to_quantized() (mirroring how QuantizedKVCache already handles the non-rotating case), or fail fast at server startup with a clear error when --kv-bits is combined with a model whose cache config includes any RotatingKVCache layers, instead of a first-request NotImplementedError.

Related

Companion MoE-specific issue on the mlx-vlm side: Blaizzy/mlx-vlm#904 (TurboQuant KV quantization crash on Gemma-4-E4B) — same underlying use case (quantized KV cache on Gemma 4), different repo/entry point.

Version: mlx-lm 0.31.3 (production), reproduced on current main (mlx_lm/models/cache.py:552, unchanged since at least 2026-04-21).

[PhilipJohnBasile]

Three things that may save you time here — a duplicate with a patch already in flight, a caveat on your workaround, and some measured numbers on what --kv-bits actually buys.

1. Same root cause as #1583, and a patch is coming

@mabaeyens filed #1583 two days after this from the BatchGenerator side — same two NotImplementedErrors, and they state they have a working patch (RotatingQuantizedKVCache + BatchRotatingQuantizedKVCache, tested against a real model) and will open a PR. Your Gemma 4 hybrid-attention case and their bounded-max_kv_size batching case are the same missing implementation, so that PR should close both. Worth watching it rather than duplicating the work.

2. Your workaround drops the bound you were trying to enforce

Worth double-checking, because I think it's load-bearing for you: your make_prompt_cache_quantized returns a plain QuantizedKVCache for every layer and never consults max_kv_size. RotatingKVCache is what enforces the window bound — so with the monkeypatch in place, the cache grows with the full conversation instead of being capped, on all 35 sliding-window layers.

If the reason for --kv-bits was to cap per-conversation memory, the patch trades that cap away for the 2× from quantization: bounded-and-fp16 becomes unbounded-and-int8, which crosses over as soon as a conversation runs past roughly twice the window. On long sessions that's a regression in the direction you were optimising. (Output correctness should still be fine — the sliding-window mask limits what's attended to regardless of what's retained — so this reads as a memory issue, not a quality one. Flagging rather than asserting; you'd know your traffic shape.)

3. What --kv-bits 8 actually costs, measured

Since you're running this in production, our numbers may be useful for expectation-setting. Qwen3-32B-4bit, group_size 64, int8 KV vs fp16, paired runs (posted upstream on mlx#3026 and mlx-lm#1587):

ctx	greedy-argmax agreement	PPL ratio	decode Δ vs fp16
0.5K	0.9804	0.9965	−7.4%
4K	0.9968	0.9966	−3.1%
16K	0.9990	0.9991	−2.7%

Quality is genuinely clean and improves with context. But decode is slower at every context we measured, because on a 4-bit dense model KV is only ~19% of decode-step bytes — weights dominate, so halving KV bandwidth cannot pay for the compose/dequant overhead. So --kv-bits 8 is a capacity tool (roughly half the KV bytes → longer context or more cache slots in the same RAM), bought at a few percent of decode speed. It is not a throughput lever.

Different model shape, different answer — a hybrid-attention MoE like Gemma 4 26B-A4B has a different byte mix than a dense 32B, so the crossover moves. But if the motivation was speed, that's worth measuring before building on it.

4. On your actual ask

Agreed that a startup-time rejection beats a first-request crash, and it stays useful even after the implementation lands, since new cache types can reappear without a to_quantized(). The cheap version is a check at server start: when kv_bits is set, build a throwaway prompt cache for the loaded model and confirm every entry implements to_quantized(), failing with the offending class name. That's independent of the RotatingQuantizedKVCache work and could go in either before or after it.

[mabaeyens]

Author of #1583 here. One correction to the status, since @PhilipJohnBasile's summary has the timeline slightly behind.

The PR is already open: #1584. It is not draft and currently mergeable, just waiting on review. It adds RotatingQuantizedKVCache and BatchRotatingQuantizedKVCache plus the to_quantized() implementations in mlx_lm/models/cache.py, with tests in test_prompt_cache.py and test_generate.py.

That is the same NotImplementedError you hit here, so this issue should be covered by it. I had not linked the two, which is why it looked pending. Linking now.

One thing I would like from you before assuming it closes your case: my testing was driven by the batching path with a bounded max_kv_size, not by Gemma 4's hybrid sliding-window layers. The blocking call is identical, but your model builds a mixed cache where only some layers are rotating, and I have not run that combination. If you are able to try the branch against Gemma 4 26B with --kv-bits 8, that would tell us whether the hybrid layout needs anything extra or whether it just works.

On @PhilipJohnBasile's second point, I think it is right and worth taking seriously if you are running that monkeypatch in production. Returning a plain QuantizedKVCache for every layer does drop the max_kv_size bound, so you get quantization but lose the window cap on the sliding-window layers. The proper implementation keeps both, which is the main reason to prefer it over the patch.

The startup-time validation idea stands on its own, independent of the cache work, and it stays useful after #1584 lands (a new cache type can always show up without to_quantized()). I would rather see it as its own small PR than folded into #1584, so the cache change stays reviewable on its own. I am not planning to pick it up myself, so it is yours if you want it.

[PhilipJohnBasile]

@mabaeyens Correction accepted, and the error is entirely mine — thank you for the link.

I checked before replying: #1584 has been open since 2026-07-18T14:25Z, non-draft and mergeable, seven days before I wrote that it was pending. What happened is that I read #1583's own closing line ("I have a working patch … and will open a PR shortly") and reported it as current status without ever checking whether the PR existed. The issue text was stale; my repeating it was the mistake, and a single gh pr list would have caught it. Sorry for putting a wrong timeline in front of @chrislyons.

On the startup validation — yes, I'll take it

Agreed it belongs in its own PR rather than folded into #1584, for exactly the reason you give: it keeps the cache change reviewable on its own, and it stays useful afterwards because a new cache type can always appear without a working to_quantized().

One thing I found while scoping it that makes the check less trivial than it sounds. The runtime failure comes from maybe_quantize_kv_cache in generate.py:

for e, c in enumerate(prompt_cache):
    if hasattr(c, "to_quantized") and c.offset >= quantized_kv_start:
        prompt_cache[e] = c.to_quantized(...)

RotatingKVCache does have to_quantized — it is defined, and it raises. So hasattr is satisfied, the call goes ahead, and the NotImplementedError surfaces on the first request that trips the threshold. The guard tests for presence, not for implementation, which is precisely why this fails late rather than at startup.

So the validation cannot be a hasattr sweep — it would pass on exactly the case it is meant to catch. What I plan instead:

• At server startup, when kv_bits is set, build a throwaway cache via make_prompt_cache(model, max_kv_size) — same arguments the serving path will use, so hybrid models produce their real mixed layout.
• Probe each entry by actually attempting to_quantized(...) on the empty cache and catching NotImplementedError. On an unused cache this is cheap, and for the rotating case it raises immediately.
• Fail with the offending class name and layer index, rather than a bare NotImplementedError three layers down a request.

That keeps the loud runtime failure intact — I would not want the fix to be "silently skip unsupported layers," since that quietly gives you an unquantized cache you asked to quantize.

I will build it against main and open it separately once #1584 has landed or settled, so the two do not race in the same file.

What I cannot help with

Your Gemma 4 hybrid question is the right one to hold the issue open on, and I am not the person for it — I do not have Gemma 4 26B locally, and my measurements in this thread are all dense Qwen3-32B. @chrislyons is the one with the model and the reproduction, so that test stays with them.

Worth adding for whoever runs it: the interesting failure mode is not just "does it crash" but whether the mixed layout quantizes the right layers — a hybrid model where only the full-attention layers end up quantized would pass a smoke test while silently doing half the work.

[PhilipJohnBasile]

I'm bored. I'll download it and test it out.

[PhilipJohnBasile]

@mabaeyens I picked up the Gemma 4 hybrid test since you hadn't run that combination. It works, including the rotation path. Branch kv-cache-quant-batching at d98dbe2, model mlx-community/gemma-4-26b-a4b-it-4bit, M5 Max, mlx 0.32.0.

The layout is genuinely hybrid, so the test isn't vacuous

Gemma 4 reaches your code by a different route than the batching path: gemma4_text.make_cache() builds the mix itself from layer_types and never consults max_kv_size. Composition before conversion:

{'RotatingKVCache': 25, 'KVCache': 5}   (n=30, sliding_window=1024)

Three checks, all pass

1. No NotImplementedError on any of the 30 layers.

2. Every layer converts, and to the right thing — this was the one I most wanted to see, since a hybrid model that quantizes only its full-attention layers would pass a smoke test while doing half the work:

before: {'RotatingKVCache': 25, 'KVCache': 5}
after : {'RotatingQuantizedKVCache': 25, 'QuantizedKVCache': 5}

25/5 matches layer_types exactly. No layer silently left unquantized.

3. The sliding-window bound survives. max_size=1024 intact on the rotating layers after conversion — so unlike the monkeypatch in this issue, you keep the cap and the quantization.

Rotation actually exercised

My first end-to-end run used a short prompt, which never approached the 1024-token window — a quantized ring buffer that is correct until it wraps isn't tested by a run that never wraps. Redone with a prompt of 3,592 tokens (3.5× the window), so the sliding layers rotate during prefill and again during generation. Greedy, identical input, fp16 vs kv_bits=8:

[fp16]     RotatingKVCache          offset=3652
[kv_bits8] RotatingQuantizedKVCache offset=3652
outputs IDENTICAL across the rotation boundary

Same 60 generated tokens, character for character.

One methodological caveat on my own first attempt

My initial E2E printed '\n\n### 3-3-3-3-3-...' and I nearly reported it as a quality problem. It wasn't: I had passed a raw string to an -it model with no chat template. The control settles it — on the raw prompt fp16 is equally degenerate (pure newlines), and with the chat template applied fp16 and int8 produce identical coherent output. Worth stating because "quantization broke generation" is exactly the kind of false alarm that costs a maintainer an afternoon.

Scope, so nobody over-reads this

• One model, one quantization (-it-4bit), kv_bits=8, greedy, single-sequence — not the batching path you already covered.
• 30 layers here; @chrislyons described a 42-layer configuration, so this MLX repack is a different cut of Gemma 4. The hybrid property is present either way, but if their crash was config-specific it's worth them confirming on their own artifact.
• Identical greedy output is a strong signal, not a guarantee across all prompts/lengths.

Scripts are committed on our side and I can paste them if useful. Happy to run other head dims or a longer generation if you want the window wrapped more times before merge.

[mabaeyens]

No need to apologise, the issue text was stale and you read it in good faith. Your hasattr analysis is right, and it is worth spelling out that #1584 does not change it: maybe_quantize_kv_cache still guards on presence (generate.py:294 on my branch, untouched), so probing by actually calling to_quantized is the correct shape for your check.

Two findings from digging into your point, one good and one that gives your PR more to do than expected.

Gemma 4 is fine, so this issue should close with #1584

gemma4_text.py:683 builds the sliding-window layers with keep=0 explicitly, and the full-attention layers get a plain KVCache. Both have working to_quantized on the branch, so @chrislyons' case should just work. Still worth the real run, for the reason you gave: the failure mode to watch is the mixed layout quantizing only the full-attention half, not a crash.

The gap your check will find: keep > 0 is still unimplemented

My RotatingKVCache.to_quantized refuses sink tokens:

if self.keep > 0:
    raise NotImplementedError(
        "Quantizing a RotatingKVCache with keep tokens is not supported."
    )

And the generic fallback in make_prompt_cache uses keep=4:

RotatingKVCache(max_size=max_kv_size, keep=4)   # cache.py:37

So every model that does not define its own make_cache, run with --max-kv-size N --kv-bits 8, still raises after #1584 lands. Gemma 4 escapes only because it defines make_cache and passes keep=0. When your startup probe fires on that combination it is a true positive, not a regression in #1584, and it is a good argument for the check existing at all: right now that failure waits until a request crosses quantized_kv_start.

Why I am not fixing keep > 0 in #1584

I looked at what it would take. The single-sequence side is not hard, roughly four changes, all mirrors of logic the float class already has: drop the guard in RotatingQuantizedKVCache.__init__, make _trim do the keep-aware split ([:keep] + [trim_size+keep:]), make _temporal_order do the three-way concat, and pass keep through from to_quantized. The rotate step already resets _idx to self.keep. None of it is numerically delicate, because mx.quantize groups along the last axis (head_dim) while every slice and concat here runs along axis 2 (sequence), so each position's row stays independently quantized. That is already what the existing _trim and _temporal_order rely on.

What stops me is the batching side. BatchRotatingKVCache has no keep concept at all, not even a constructor argument, and merge() validates max_size across caches while ignoring keep entirely. So merging RotatingKVCache(keep=4) instances into a batch already drops the sink tokens today, quantized or not. If I implement keep only in the single-sequence quantized class, --max-kv-size --kv-bits starts working single-stream while the batching path still silently cannot honour it, and a quiet wrong answer is worse than the loud error we have now.

Properly fixing it means teaching BatchRotatingKVCache about keep first, which is a bigger change than this PR and unrelated to quantization. So my preference is to leave the NotImplementedError in #1584, let your check surface it at startup with a clear message, and treat keep > 0 as its own piece of work. If a maintainer would rather see it in #1584, I will do it, but I would want the batch gap addressed in the same breath rather than shipping the asymmetry.

[PhilipJohnBasile]

Started on the startup-validation piece and ran into something that needs sorting first, because it changes what the fix should be and possibly what this issue is about.

--kv-bits does not exist in mlx_lm.server

@chrislyons, your repro is:

mlx_lm.server --model <gemma-4-26b-a4b-path> --kv-bits 8 --port 41966

I can't get that to parse on any version I can find. server.py's full flag set on current main (cf10f96) is:

--model --adapter-path --host --port --allowed-origins --draft-model
--num-draft-tokens --trust-remote-code --log-level --chat-template
--use-default-chat-template --temp --top-p --top-k --min-p --max-tokens
--chat-template-args --decode-concurrency --prompt-concurrency
--prefill-step-size --prompt-cache-size --prompt-cache-bytes --pipeline

No --kv-bits, --kv-group-size or --quantized-kv-start. Same in the released 0.31.3 you mention. And git log -S'kv-bits' --follow -- mlx_lm/server.py returns nothing across the file's entire history — with a positive control, since the same search on generate.py correctly finds 60358be "Quantized KV Cache (#1075)". The flag lives only in mlx_lm.generate and mlx_lm.cache_prompt.

So argparse should reject that command outright and the server should never come up — which doesn't match "server starts cleanly, /health returns 200, first request raises". Something else is going on:

• a fork or local patch adding the flag (your monkeypatch suggests you already patch this area);
• the crash arriving via mlx_lm.generate or a library call rather than the server;
• or the server picking kv_bits up from somewhere other than argv in your setup.

Whichever it is, it matters: @mabaeyens is treating #1584 as closing this issue, and if the failing path isn't the one #1584 fixes, that assumption should be checked before it merges. Could you paste the exact invocation and mlx_lm.__version__ from the environment that crashes?

To be clear, the underlying defect is real and #1584 addresses it — RotatingKVCache.to_quantized() raising is exactly what a hybrid model hits, and I confirmed the branch handles Gemma 4's mixed layout correctly above. This is only about whether your reproduction is the same path.

On the validation itself — there's already a precedent to copy

@mabaeyens, my "validate at server startup" framing was wrong for the reason above: there is no server-side kv_bits to validate. But generate_step already does this exact thing for a different capability, immediately after the cache is built:

if not cache.can_trim_prompt_cache(model_cache):
    types = {type(c).__name__ for c in model_cache if not c.is_trimmable()}
    raise ValueError(
        f"Speculative decoding requires a trimmable prompt cache (got {types})."
    )

backed by can_trim_prompt_cache() → c.is_trimmable() in cache.py. So the honest shape for this is the same convention applied to a second capability — a can_quantize_prompt_cache() helper plus an upfront raise in generate_step when kv_bits is not None, naming the offending classes. That covers every entry point rather than one CLI, including library callers passing kv_bits= and the server if it ever gains the flag.

It also cannot be a hasattr sweep, for the reason in my earlier comment: RotatingKVCache has to_quantized, it just raises. The capability has to be declared (an is_quantizable() alongside is_trimmable()), not inferred.

Why I'm holding the PR rather than opening it now

Doing it properly means touching every class in cache.py — which is the file #1584 rewrites, and you asked for that PR to stay reviewable on its own. Landing mine first guarantees a conflict and makes yours harder to review, which is the opposite of the point. So I'd rather wait until #1584 merges and then add is_quantizable() across the classes as they exist afterwards, including the two new ones. Say the word if you'd prefer it the other way round and I'll rebase around you instead.

[PhilipJohnBasile]

Follow-up validation is now open as draft #1618.

It is intentionally stacked on #1584 and will be rebased after that lands. The guard now:

• declares cache quantization applicability/capability explicitly rather than inferring support from hasattr;
• fails before the first model call with every offending layer index, class, and reason;
• covers generate_step, both caches in speculative generation, and generated or externally supplied BatchGenerator caches;
• accepts the Gemma-style RotatingKVCache(keep=0) + KVCache mix from Implement KV-cache quantization for the continuous-batching path (RotatingKVCache/BatchRotatingKVCache) #1584;
• rejects the remaining real keep>0 limitation immediately; and
• does not invent a nonexistent mlx_lm.server --kv-bits flag.

The M5-controlled cache/generation suites are green (64 passed with MLX_ENABLE_TF32=0), including explicit no-silent-skip tests.

For classification: this is a correctness/configuration-validation issue, not a cybersecurity vulnerability. It does not cross a trust boundary, expose data, execute untrusted input, or change authentication/authorization. The impact is a late exception or a requested memory/performance option not being honored.

[mabaeyens]

Found it. Your instinct to stop and check was right, and the answer changes what this issue can claim to close.

mlx_lm.server never had the flag, but mlx_vlm.server does

Confirming your side first. On my checkout, git log -S'kv_bits' --all --follow -- mlx_lm/server.py returns zero commits, same for 'kv-bits', with the positive control on generate.py behaving. It is not just missing from argparse: neither generation path passes it either. The BatchGenerator construction at server.py:757 passes completion_batch_size, prefill_batch_size, prefill_step_size and stream only, and the stream_generate call at server.py:922 has no kv_bits either, even though BatchGenerator.__init__ accepts kv_bits, max_kv_size, kv_group_size and quantized_kv_start. So stock mlx_lm.server cannot turn on KV quantization by any route, and the reported command cannot produce that traceback.

The flag exists in mlx-vlm:

# mlx_vlm/server/cli.py:143
parser.add_argument(
    "--kv-bits",
    type=float,
    default=None,
    help="Number of bits for KV cache quantization (e.g. 3.5 for TurboQuant).",
)

mlx_vlm.server --kv-bits 8 parses. mlx_lm.server --kv-bits 8 does not.

The important part: mlx-vlm vendors its own copy of cache.py

mlx_vlm/models/cache.py is a 90 KB fork of this file. It does not import from mlx_lm.models.cache; its only imports are typing and mlx. It carries its own RotatingKVCache, and at line 515:

def to_quantized(self, group_size: int = 64, bits: int = 4) -> QuantizedKVCache:
    raise NotImplementedError("RotatingKVCache Quantization NYI")

Byte-identical to what v0.31.3 has at mlx_lm/models/cache.py:552. BatchRotatingKVCache Quantization NYI is vendored too, at line 1304.

That also explains how @chrislyons could cite an accurate mlx-lm line number for a crash that did not come from mlx-lm. mlx-vlm requires mlx-lm>=0.31.3, the exact version in the report, so both packages are installed side by side and mlx_lm/models/cache.py:552 really does contain that line in their environment. The code citation is correct. The entry point attributed to it is not.

The companion issue in the report points the same way: Blaizzy/mlx-vlm#904 is a TurboQuant crash, TurboQuant is mlx-vlm's (mlx_vlm/turboquant.py), and it is named in the --kv-bits help text above.

What this means for #1584 and for this issue

#1584 fixes the defect in mlx-lm, and @PhilipJohnBasile's Gemma 4 run on the branch demonstrates that end to end. What it will not do is fix @chrislyons' production crash, because nothing we merge here touches mlx-vlm's vendored copy. That needs the equivalent change in Blaizzy/mlx-vlm, against mlx_vlm/models/cache.py.

So I would separate the two. The mlx-lm defect is real and #1584 closes it. I am withdrawing the suggestion that merging #1584 resolves the reproduction in the opening post, and I would not close this issue on that basis without @chrislyons confirming which stack actually crashes.

@chrislyons, one loose end that only you can settle. Your workaround patches mlx_lm.models.cache.make_prompt_cache and mlx_lm.server.make_prompt_cache, and neither would affect mlx-vlm's vendored path, so that snippet does not fit the mlx-vlm story either. Either you are running both stacks and the patch belongs to a different part of the deployment, or it was transcribed from somewhere else. The exact failing invocation and mlx_vlm.__version__ alongside mlx_lm.__version__ would settle it in one paste.

[PhilipJohnBasile]

---

Confirmed independently — git log -S'kv_bits' --all --follow -- mlx_lm/server.py
returns nothing here either, and neither generation path passes it: no
kv_bits at the BatchGenerator construction in server.py, none at the
stream_generate call. Stock mlx_lm.server has no route to the reported
crash. Good catch on mlx_vlm/models/cache.py being a standalone 90 KB fork
rather than an import — that's not a thing I'd have thought to check, and it
fully explains how a real mlx_lm/models/cache.py:552 line number ended up
attached to a report this repo can't produce.

So: withdrawing my part of the original framing along with yours. #1584 (and
#1618 on top of it) fix the defect that exists in mlx-lm. Whatever
@chrislyons hit needs the equivalent fix filed against Blaizzy/mlx-vlm
instead, and I'd rather that move as its own issue there than sit half-fixed
here — happy to open it if nobody more familiar with mlx-vlm's fork gets to it
first.

I'd leave this issue open pending @chrislyons confirming which stack actually
crashed, per your question about the workaround patching
mlx_lm.models.cache.make_prompt_cache — that snippet doesn't fit the
mlx-vlm story either, so there's a piece of the original report still
unresolved that #1584/#1618 won't touch.