MLX上のTurboQuant：カスタムMetalカーネルによるKVキャッシュ圧縮で4.6倍（Qwen 32BがFP16速度の98%）

統合されたMetalカーネルを使って、MLX向けにTurboQuant（Googleの新しいKVキャッシュ圧縮に関する論文）を実装しました。

Qwen2.5-32Bでの結果、M4 Pro 48GB:

- 4.6x圧縮、0.98x FP16速度、品質は同一

- 16Kコンテキスト: 4.2GBキャッシュ → 897MB

主な課題は速度でした。統合されたMetalの量子化/復号化カーネルと、増分デコードバッファによって、0.28xから0.98x FP16へ改善しました。

最適化の全行程のまとめ: https://medium.com/@antonrozanov/turboquant-on-mlx-4-6x-kv-cache-compression-with-custom-metal-kernels-9cdee3f7d2a2

コード: https://github.com/arozanov/turboquant-mlx

mlx-lmへのPR: https://github.com/ml-explore/mlx-lm/pull/1067