Implemented TurboQuant (Google's new KV cache compression paper) for MLX with fused Metal kernels.
Results on Qwen2.5-32B, M4 Pro 48GB:
- 4.6x compression, 0.98x FP16 speed, identical quality
- 16K context: 4.2GB cache → 897MB
The main challenge was speed — went from 0.28x to 0.98x FP16 through fused Metal quantize/dequantize kernels and an incremental decode buffer.
Writeup with the full optimization journey: https://medium.com/@antonrozanov/turboquant-on-mlx-4-6x-kv-cache-compression-with-custom-metal-kernels-9cdee3f7d2a2
Code: https://github.com/arozanov/turboquant-mlx
PR to mlx-lm: https://github.com/ml-explore/mlx-lm/pull/1067
[link] [comments]
