add TurboQuant KV cache compression

[TheTom]

Proposed changes

Adds TurboQuant KV-cache compression (arXiv:2504.19874) behind the kvScheme parameter that landed in #230. This is a complete rework since the last push; the earlier draft's codec and kernels are gone.

• Codec: Walsh-Hadamard rotation (dense QR fallback for non-power-of-two head dims) with per-vector norms and a Lloyd-Max codebook at 2, 3, or 4 bits.
• JIT Metal kernels via MLXFast.metalKernel (source-string JIT; no Package.swift or build changes): fused encode, and a SIMD-parallel two-pass online-softmax flash decode with three K-load variants (turbo-packed, raw fp16, 8-bit affine) plus an inline sparse value pass. Per-call values travel in a parameter buffer rather than kernel templates, so each shape compiles once.
• Asymmetric scheme family: keys are far more quality-sensitive than values (softmax amplifies key error; value error averages out), so the recommended schemes keep K at fp16 (turbo0v4/v3/v2) or 8-bit affine (turbo8v4/v3/v2, scored via inline dequant) while V takes the compression. Symmetric turbo4/3/2 are available for validated use.
• Boundary-layer protection: for fragile schemes (turbo-quantized keys, 2-bit values) the first and last two attention layers use 8-bit affine instead.
• Two-phase cache: raw fp16 during prefill (zero prefill overhead), compression on the first decode step, incremental encode per token afterwards. Prompt-cache save/restore round-trips for all four state layouts. Speculative decoding and multi-turn reuse route through the compressed path. Generation settles cache state with the token eval and drains autoreleased temporaries per step, so long generations hold steady memory.

Usage

var parameters = GenerateParameters()
parameters.kvScheme = "turbo8v3"   // recommended default

Measured quality and compression

M5 Max, WikiText-2 decode-time evaluation, ctx 4096. KLD is mean full-vocab KL divergence against an fp16 cache with identical weights, so it isolates cache compression from weight quantization. Compression is measured cache bytes against the fp16 cache.

scheme	K	V	Qwen3-1.7B	Qwen2.5-7B	Phi-4-mini	Mistral-7B	Qwen3-30B-A3B (MoE)	Qwen2.5-32B
affine8 (#230)	q8	q8	.0088 / 1.88x	.0406 / 1.88x	.0004 / 1.88x	.0001 / 1.88x	.0049 / 1.88x	.0003 / 1.88x
turbo0v4	f16	tq4	.0066 / 1.58x	.0046 / 1.58x	.0050 / 1.58x	.0015 / 1.58x	.0076 / 1.58x	.0063 / 1.58x
turbo8v3	q8	tq3	.0387 / 2.72x	.0455 / 2.72x	.0203 / 2.72x	.0093 / 2.72x	.0195 / 2.72x	.0373 / 2.72x
turbo4 (sym)	tq4	tq4	.147 / 3.29x	.0596 / 3.29x	.0365 / 3.35x	.0076 / 3.35x	.0383 / 3.47x	.0376 / 3.54x

Reading the table:

• turbo8v3 is the headline: 2.72x KV compression against affine8's 1.88x at comparable quality, across six model families spanning 1.7B to 32B and a 30B MoE.
• turbo0v4 is the quality play: lower KLD than affine8 on five of six families (9x lower on Qwen2.5-7B) at 1.58x.
• Symmetric turbo4 rides on per-dimension key calibration: at compression time the cache computes a per-dimension scale from the post-rotation key variance of the prefill cache (clamped, mean-normalized, skipped when variance is already flat) and folds it into the query side, so the score kernels are untouched. Without calibration turbo4 degrades badly on qk-norm and small models (KLD 2.65 on Qwen3-1.7B, 2.76 on Phi-4-mini); with it every measured family lands between 0.008 and 0.15. The docs still steer users to start asymmetric and validate before going symmetric.

Decode throughput

Single-token decode uses one SIMD-parallel flash dispatch that reads K once in its stored basis (raw fp16 or inline-dequantized affine, no full-cache cast or GQA materialization), dequantizes V inline with a sparse gate, and applies the inverse value rotation in the reduction pass. Qwen3-1.7B, tok/s:

scheme	decode tok/s	vs fp16
fp16	150	1.00x
turbo0v4	102	0.68x
turbo0v2	111	0.74x
turbo8v3	114	0.76x
turbo4 (calibrated)	122	0.81x

Family sensitivity notes

• Mistral-class: tolerant across the board, including symmetric key compression.
• Qwen3 / qk-norm families and small models: most sensitive to key compression; asymmetric schemes stay healthy.
• Qwen2.5-class: large K outliers in the boundary layers (scores are computed in f32 for this reason); affine8 is weak on this family, turbo value schemes are not.
• Phi-family: unusually affine-friendly.
• Gemma-family: sliding-window (rotating) layers are already memory-bounded and stay fp16 with a one-time notice; the global layers, where cache memory actually grows, compress normally.

Tests

60 tests, CI-runnable, no checkpoint downloads: codec round trips and pinned codebooks, rotation orthogonality, packing, kernel-vs-reference attention parity (GQA repeat factors 1/2/3, non-power-of-two head dims), prompt-cache save/restore for all four state layouts, compressed partial trim with continued decode, speculative decoding with a turbo scheme, allocation-boundary growth, boundary-protection gating, end-to-end generation on a tiny random-weight model, calibrated-encode kernel-vs-reference exact index parity, calibration state save/restore, and mixed rotating/standard cache conversion.

Not in this PR (follow-ups)

• A single-dispatch decode merging the two flash passes was built and measured: no end-to-end win at 4K context or below (MLX pipelines the two dispatches without a sync, and the merged kernel has fewer threadgroups at decode shapes), so it is parked rather than shipped.
• Compressing rotating-cache layers themselves (their eviction layout does not fit sequential-append storage); mixed models compress their global layers today and a one-time notice lists the layers that stay fp16.

Review history

All five earlier review items are resolved in this revision: Package.swift untouched, threadgroup sizing for head dims above 128, NR0 KV-head-boundary gating, prompt-cache restore, and speculative-decoding kvScheme plumbing.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code
• I have added tests that prove my feature works
• I have updated the necessary documentation (<doc:kv-cache-quantization> article + GenerateParameters.kvScheme)

[aleroot]

You cannot change this library, I mean this needs to reference the official ml-explore mlx-swift, if you need to add there, you need to first add a pull request there and get it merged, this is no sense.

[TheTom]

This PR is still in draft, pointing to the wrong library. I would hold off on review. Defiant pointing at the wrong library

[TheTom]

This will be fixed in the next push. No intention of switching the dependencies.

[aleroot]

Out-of-bounds threadgroup memory in fused encode kernels (can corrupt results/crash for head_dim > 128).
In both dense and WHT encode kernels, shared_norm is hardcoded to 4 entries, but indexed by sg_id = d/32 and read up to num_groups = (Dim+31)/32. For Dim=256, this accesses indices 0...7.

[TheTom]

Good catch, fixed in next push. shared_norm is now sized [(Dim + 31) / 32] instead of hardcoded 4. Current models all use head_dim=128 so this wasn't hit in practice but would break on 256.

[aleroot]

NR0 path derives kv_idx from only the first query row in the group. If a group spans multiple query heads (possible when grouping is not aligned to L), some rows will read/write against the wrong KV head. Guarding only on totalQ % nr0 == 0 is insufficient for correctness.

[TheTom]

Added a queryChunkLength % nr0 == 0 in, next push, guard so the NR0 path only activates when L aligns with the group size. Falls back to per-row dispatch otherwise.

[aleroot]

TurboQuant cache cannot be restored correctly from prompt cache. TurboQuantKVCache is not represented in cacheClassName, so it is serialized as KVCache. On load it is restored as KVCacheSimple, but TurboQuant compressed state carries 3/4 arrays rather than KVCacheSimple’s required 2 arrays, leading to invalid restoration behavior. Add explicit TurboQuant class name mapping and restore path.

[TheTom]

Fixed in nex push. Added TurboQuantKVCache to cacheClassName and restoreCacheFromMetaState. metaState now carries bits, keyBits, valueBits, and seed so the cache can be reconstructed correctly on load.

[TheTom]

@aleroot I appreciate the early reviews but this is still in draft.

[aleroot]

kvScheme dropped in speculative decoding quantization.TurboQuant scheme selection is ignored for speculative generation despite being part of GenerateParameters...

[TheTom]

Yep, missed that. Threaded kvScheme through the speculative quantization closure now.

[aleroot]

The implementation is ambitious and has strong ideas, and honestly is something I wanted to work on as well... At the moment this pull request is just a prototype as it does not look like a full inner-product TurboQuant formulation ...
There are no tests as well that are really proving is working.

[TheTom]

The implementation is ambitious and has strong ideas, and honestly is something I wanted to work on as well... At the moment this pull request is just a prototype as it does not look like a full inner-product TurboQuant formulation ...
There are no tests as well that are really proving is working.

Appreciate your early feedback. Still a work in progress and know that this is an ambitious PR as is. Thank you for the draft comments!

[aleroot]

The implementation is ambitious and has strong ideas, and honestly is something I wanted to work on as well... At the moment this pull request is just a prototype as it does not look like a full inner-product TurboQuant formulation ...
There are no tests as well that are really proving is working.

Appreciate your early feedback. Still a work in progress and know that this is an ambitious PR as is. Thank you for the draft comments!

Sorry, I did not notice it was a draft .

[TheTom]

The implementation is ambitious and has strong ideas, and honestly is something I wanted to work on as well... At the moment this pull request is just a prototype as it does not look like a full inner-product TurboQuant formulation ...

There are no tests as well that are really proving is working.

Appreciate your early feedback. Still a work in progress and know that this is an ambitious PR as is. Thank you for the draft comments!

Sorry, I did not notice it was a draft .

All good let's work together on this. It's a big lift and it's hard to separate into individual PRs. Lots of work ahead.

[TheTom]

The implementation is ambitious and has strong ideas, and honestly is something I wanted to work on as well... At the moment this pull request is just a prototype as it does not look like a full inner-product TurboQuant formulation ... There are no tests as well that are really proving is working.

Follow up as i was on the road for my last comment:
Appreciate the thorough review, genuinely helpful. All the bugs you flagged are fixed locally and will be in the next push.

Still polishing this, the PPL data in the description shows it working across Llama, Qwen, Mistral, and Qwen3.5 MoE. Happy to collaborate on getting this to a state that works for the project. If you have ideas on the formulation or testing approach I'm all ears. (we can also do email or discord or whatever works)

[davidkoski]

@TheTom what is the state of this one? Trying to dig out of the backlog, so no rush :-)

[TheTom]

Thanks for the ping, just been busy with other projects. will bump it up and report back soon

[davidkoski]

I wonder if this should be an enum? I realize it is a little late (pre-existing) but potentially we could have an enum property and have this one be computed and set the other one.

[davidkoski]

Ah, it is in parameters too -- perhaps that ship has sailed.

[TheTom]

I'm flexible just tell me what you prefer as the repo owner

[davidkoski]

Let's leave it as-is. I would prefer enum, but I don't know if it is feasible and it is outside the scope of this PR. We can think about it in a followup. Maybe we only need the enum here and have the rest of the parameters stay as String -- I think it may be more important for callers to have an easy way to know what values are possible.

[TheTom]

Yeah agreed, enum would be nicer but the string is already in parameters so it's a bit baked in. Happy to do a follow-up that adds an enum for discoverability and keeps the string as the storage. For now I left it alone like you said.

[davidkoski]

I am curious about the width conversion on rawKeys (which per the call to update() are in the dtype of the activation weights). float32 rawKeys are unlikely, but what about bfloat16? The exponent may be too large and convert to (I presume) an infinity. It looks like the math in the kernels are float32, so it is really only an issue for the inputs.

Is this a real issue? Does this require bfloat16 and float32 kernels (or at least bfloat16)?

Related: I think if we don't have bfloat16 kernels, this will repeatedly cast from bfloat16 to float16, potentially leading to n^2 cost.

[TheTom]

Good catch, and it was actually worse than the overflow. That asType was also copying the whole raw K cache every decode step since the cache grows, so the n^2 you suspected was real too.

Fixed by just not casting. mlx generates the kernel signature from the runtime dtype of whatever array you pass in, so I made the K pointer type a Dtype template arg (KT, plus KScaleT/KBiasT for the affine scales/biases which had the same problem). Template args get hashed into the kernel name so each dtype gets its own compiled variant, one JIT compile per dtype ever used. Had to do it that way btw, mlx's custom kernel library cache is single slot per kernel name, so without the dtype in the name alternating dtypes would recompile every switch.

So to answer your question directly: yes it needs bf16/f32 variants and now it has them for free. Added a test that runs the raw-K decode with a bf16 cache containing a key above 65504 and checks against an f32 reference, no inf.

[davidkoski]

I think this needs a groupSize check like maybeQuantizeKVCache() has:

    func quantize(_ cache: KVCacheSimple) -> KVCache {
        let state = cache.state
        if state.count == 2 {
            let keyHeadDim = state[0].dim(3)
            let valueHeadDim = state[1].dim(3)
            guard
                resolvedKVQuantizationGroupSize(
                    requested: effectiveGroupSize,
                    keyHeadDim: keyHeadDim,
                    valueHeadDim: valueHeadDim
                ) != nil
            else {
                return cache
            }
        }
        return cache.toQuantized(groupSize: effectiveGroupSize, bits: effectiveBits)
    }

Specifically resolvedKVQuantizationGroupSize() will return nil if the headDim are incompatible with the groupSize (affineKeyGroupSize) and it won't try to quantize the cache.

[davidkoski]

QuantizedKVCache also has this:

    public func updateQuantized(keys: MLXArray, values: MLXArray) -> (
        (MLXArray, MLXArray, MLXArray?), (MLXArray, MLXArray, MLXArray?)
    ) {
...
        guard effectiveGroupSize != nil else {
            fatalError(
                "KV cache quantization requires head dimensions divisible by one of the supported group sizes (32, 64, 128). Requested group size: \(groupSize). Key head dim: \(kHeadDim). Value head dim: \(vHeadDim)."
            )
        }

I think you could only get to that case if you constructed it directly, but a similar message might be worthwhile -- it is a programmer error so a fatal is ok, but giving an actionable message is useful.

[TheTom]

Done, used resolvedKVQuantizationGroupSize exactly like maybeQuantizeKVCache does. Covers both the turbo8 schemes and the boundary-layer protection, skips the layer (stays fp16) when nothing fits, and if it resolves to something other than 64 (like 32 for head dim 96) that gets threaded through to the kernels too so they always agree. Tests at head dim 80 for both paths.

One thing I wasn't sure about: on an incompatible head dim the value side could technically still compress (the codebook path has no group size constraint), so turbo8v3 could degrade to turbo0v3 instead of skipping entirely. I went with skip to match the existing semantics but easy to flip if you'd rather degrade gracefully.

[TheTom]

Yep, added. The affine-K encode now resolves the group size up front and the fatal prints the head dim and the supported sizes, same style as updateQuantized. Only reachable by constructing the cache directly but at least now it tells you what to fix.

[TheTom]

Pushed the review fixes: native bf16/f32 dtype variants for the raw-K and affine-K kernels (no more per-step casts), the group size guard with tests at head dim 80, and the actionable fatal. Also dropped some dead experimental kernels I found while in there. 63 tests now. CI will need an approve-and-run when you get a chance, it got gated after the force push.

[davidkoski]

Change looks good, thank you!