Interest in TurboQuant / rotating quantized KV cache support?

[RNT56]

Hi, we have been experimenting in a fork with a TurboQuant path for MLX Swift and an LM-side rotating quantized KV cache, and wanted to ask whether this is directionally interesting before trying to shape it into upstream PRs.

The implementation is currently split across:

• mlx-swift: https://github.com/RNT56/mlx-swift/tree/schtack/turboquant-kv
• mlx-swift-lm: https://github.com/RNT56/mlx-swift-lm/tree/schtack/turboquant-kv

What exists in the fork:

• A TurboQuant API in mlx-swift with packed tensor support, reference encode/decode, Metal codec kernels, compressed attention kernels, runtime capability probing, and fallback diagnostics.
• mlx-swift-lm integration with TurboQuant KV cache strategies, rotating quantized KV cache support, compressed attention routing, cache diagnostics, and tests around cache copying/serialization/fallback behavior.
• Runtime fallback behavior: the public path can fall back to MLX packed quantized lanes when the PolarQuant/QJL Metal backend is unavailable or fails the self-test. The Metal attention path is gated behind a runtime probe, so unsupported devices should keep using the safer fallback path.
• Device/profile handling currently distinguishes portable Apple GPU profiles from wider/sustained profiles using Metal availability, GPU family signals, architecture name, and recommended working-set size. The code currently treats A16/A17-class devices as portable, A18/A19 / Apple8+ as wider, and larger working sets / A19 Pro as sustained.

Measured/observable benefit so far:

• The main concrete benefit is KV-cache memory pressure reduction. The compressed representation stores K/V payloads in low-bit packed form instead of keeping full fp16/bf16 raw cache tensors. For 4-bit values, the raw value payload is about 4x smaller than fp16 before per-group scale/bias metadata; the 2.5-bit / 3.5-bit TurboQuant presets reduce payload lanes further. The LM-side compressed-attention path also has tests that verify raw-free compressed cache state when the Metal attention backend is available.
• We have not yet produced polished end-to-end tokens/sec benchmark numbers suitable for an upstream PR. If this direction is interesting, we can turn the fork into a smaller benchmarked PR series rather than sending the current fork stack wholesale.

Suggested upstreaming shape, if maintainers are interested:

1. First discuss the desired API surface and benchmark/device matrix.
2. Split the mlx-swift core support from the mlx-swift-lm cache integration.
3. Keep the fallback/probe behavior mandatory so unsupported devices remain on existing paths.
4. Add focused benchmarks for memory footprint and generation throughput before any large PR.

I also prepared a few small unrelated cleanup/fix branches separately, because those are more suitable as conventional PRs than the TurboQuant stack:

• mlx-swift linalg nuclear norm API/docs: https://github.com/RNT56/mlx-swift/tree/upstream-pr/linalg-norm-kind-nuc
• mlx-swift-lm VLM processor completions: https://github.com/RNT56/mlx-swift-lm/tree/upstream-pr/vlm-processor-completions
• mlx-swift-lm model config/runtime error hardening: https://github.com/RNT56/mlx-swift-lm/tree/upstream-pr/model-config-runtime-hardening
• mlx-swift-lm model compatibility docs: https://github.com/RNT56/mlx-swift-lm/tree/upstream-pr/model-compatibility-docs

Would you be interested in us preparing a benchmarked TurboQuant PR series, and if so, which part would you prefer to evaluate first?

[davidkoski]

Interested, but see also:

• Add TurboQuant KV cache compression (arXiv 2504.19874) mlx-swift#405 (an attempt to add it in the mlx layer via mlx-swift)
• Add TurboQuant KV cache compression on KVCacheSimple #287
• add TurboQuant KV cache compression #232
• Add TurboQuant KV cache backend #160

I have a bit of a backlog I am working though. I don't know exactly how these turboquant PRs will land since I am still reviewing and trying to get them ready for merge. I am not set on any of these, so if you think your approach is better you are welcome to submit a PR. I think ultimately I will have to select one of these to merge and it may be the one that looks the best or the one that is ready.

I think your various branches would make fine PRs.

[RNT56]

Thanks, that context helps. I reviewed the linked TurboQuant PRs and the MLX-layer guidance.

I’ll keep the current working fork/pinned dependency setup intact and prepare clean, reviewable PR branches from it rather than repurposing or force-pushing the existing integration branch.

I’ll start with the unrelated cleanup branches as separate small PRs, then prepare a focused mlx-swift-lm TurboQuant KV-cache PR covering the cache strategy/API, fallback/probe behavior, serialization/copying, and tests. After that I’ll split compressed attention, rotating-cache support, and any required lower-level Swift/Metal support into follow-up PRs.

That should make the work easier to compare with the existing TurboQuant PRs and review incrementally.

[[ml-explore/mlx-swift#412: TurboQuant Swift support primitives
ml-explore/mlx-swift-lm#304: focused LM-side TurboQuant KV cache strategy](https://github.com//pull/304)](#304)

[joelnishanth]

@davidkoski — thanks for the feedback on #405 and for linking the related PRs here. Understood that the C++ primitives belong at the mlx layer, not mlx-swift. Closing #405 was the right call.

I've been following #232 (TheTom), #160 (timonharz), and now @RNT56's fork — there's solid lm-layer work already in progress across all three, and I don't want to add a fourth competing implementation.

What I can contribute: benchmarking tooling

I built a dedicated benchmark CLI at joelnishanth/mlx-swift-turboquant that measures fp16 vs TurboQuant across short (256), medium (2K), and long (8K) token workloads. It tracks:

• tok/s (generation throughput)
• TTFT (time to first token)
• Peak memory (resident set size)
• Output quality (snippet for spot-checking)
• JSON export for reproducible comparison

None of the current PRs have a reusable benchmark suite — happy to adapt this to validate whichever implementation moves forward, or to help contributors run standardized comparisons across approaches.

On coordination

Since David mentioned he'll ultimately select one implementation to merge, it might help to align on a shared evaluation criteria — e.g. PPL across a standard model set (Llama-3.2-3B, Qwen2.5-3B, a MoE model), memory savings at 4K/8K/16K context, and decode throughput. I can set that up in the benchmark tool if there's interest.

Also worth flagging: @Landon-Molt's findings on mlx#3404 suggest codebook lookup is 3-5x slower than scalar dequant on Metal GPU, and their revised recommendation is Hadamard rotation + mx.quantize scalar format leveraging the generic quantized SDPA from mlx#3026. That could affect the kernel approach for all implementations here — might be worth discussing before anyone optimizes further down the codebook path.

Happy to help review PRs, run benchmarks, or fill gaps. Let me know where I can be most useful.

— Joel Nishanth · offlyn.AI