TurboQuant-MLX：在 Apple Silicon 上将 KV Cache 压缩 4.6 倍的 MLX 扩展库

背景与起源：长上下文推理的内存瓶颈

当你用大语言模型处理一份 32K tokens 的长文档时，模型需要在每一步推理中存储所有历史 Key-Value 矩阵（KV Cache）——这部分内存占用随上下文长度线性增长，成为长上下文场景的致命瓶颈。以 Qwen 2.5-7B 为例，32K 上下文中 KV Cache 占用高达 6.21 GB，比模型权重本身还要大。

传统解法是量化 KV Cache，但粗暴地对 K（Key）和 V（Value）统一量化会引发问题：K 矩阵参与注意力分数计算，4-bit 量化带来的噪声会破坏 softmax 的数值稳定性，导致生成结果完全偏离基线。

TurboQuant-MLX 正是为解决这一矛盾而生。

核心技术：K/V 分而治之的混合精度量化

TurboQuant-MLX 的核心洞察来自论文 arXiv:2504.19874：K 和 V 的量化容忍度截然不同。

• K 矩阵：参与 softmax(score) 计算，对数值精度高度敏感。4-bit 量化下贪婪解码完全崩溃。
• V 矩阵：参与加权求和（weighted sum），噪声被平均值稀释，3-bit 量化下质量几乎无损。

基于这一发现，TurboQuant 提出了 PolarQuant 量化方案：

1. 随机 Hadamard 旋转（Randomized Hadamard Rotation）：对向量做 Walsh-Hadamard 变换（WHT），将信号"搅散"成近似高斯分布，降低极端值出现的概率。
2. Lloyd-Max 量化：在高斯化后的向量上应用最优标量量化（8 个码本中心），将浮点数映射为 3-bit 索引。
3. Bit Packing：10 个 3-bit 索引打包为一个 uint32，存储效率提升 3.3 倍。

性能实测：内存下降 18%，速度保持 72%

在 Qwen 2.5-7B、32K 上下文的测试条件下：

质量验证采用贪婪解码逐字比对，生成文本与基线 100% 匹配——这意味着在压缩 KV Cache 的同时，TurboQuant 实现了零质量损失的推理加速。

两种使用路径：开箱即用 vs 深度定制

路径一：V-only TurboQuant（推荐新手）

无需修改 mlx-lm 源码，只需安装 turboquant-mlx：

pip install -e .

from turboquant_mlx.v_only_cache import VOnlyTurboQuantCache

cache = [VOnlyTurboQuantCache(bits=3) for _ in range(n_layers)]
# 直接作为 mlx-lm 的 prompt_cache 使用

这种方式保持 K 为 FP16，仅对 V 应用 PolarQuant 3-bit，兼容性最佳。

路径二：K8+V4 混合量化（追求极致压缩）

需要配合 arozanov/mlx-lm 分支（feature/turboquant-kv-cache）使用：

pip install --force-reinstall --no-cache-dir \
  git+https://github.com/arozanov/mlx-mlx.git@feature/turboquant-kv-cache

支持 --kv-cache-quantization 8,4 服务端标志，以及磁盘持久化（--prompt-cache-dir），适合需要长时间运行推理服务的场景。

GPU 内核优化：融合 Metal Kernel 的工程细节

TurboQuant-MLX 不只是算法实现，更有深度的 GPU 优化：

• Pre-rotated Q Scoring（prerot_fused_qk_scores）：在一次 GPU 内核调用中完成 Query 预旋转 + QK 分数计算，消除中间结果的全局内存读写。
• Butterfly Pulled-out：WHT 是线性变换，sum(w * WHT(x)) = WHT(sum(w * x))，先在寄存器中累加加权向量，再做一次 butterfly，V-attention 内核提速 4.5 倍。
• SIMD-Group Reduction：QK 评分中的树形规约被 simd_sum 替代，内核提速 1.85 倍。
• GQA 感知：n_rep 参数避免在 GQA（Grouped Query Attention）模型上做 mx.repeat 内存分配。

上手门槛评估

局限与风险

• 平台锁定：核心依赖 MLX，仅支持 Apple Silicon，无法在 NVIDIA GPU 或 Intel Mac 上运行。
• 混合量化需要 Fork：官方 mlx-lm 尚未合并 TurboQuant KV Cache 分支，需跟踪上游演进。
• 长期维护风险：项目 stars 仅 113，贡献者较少，mlx-lm 分支的活跃度有待观察。
• 模型兼容性：需验证在其他模型（如 LLaMA、Mistral）上的质量损失是否同样为 0。

行业意义

随着 Apple Silicon 在开发者群体中渗透率提升，MLX 生态正在从"尝鲜玩具"走向"生产级推理引擎"。TurboQuant-MLX 代表了一个重要方向：在硬件约束下通过算法创新（而非暴力堆算力）提升推理效率。KV Cache 压缩每节省 1GB，意味着可以在同等内存下多处理 50% 的上下文长度——这对于检索增强生成（RAG）和长文档分析场景意义重大。

同时，PolarQuant 论文的思路（K/V 差异量化 + Walsh-Hadamard 旋转）也有望迁移到 CUDA 生态，推动更广泛的 KV Cache 优化实践。

一句话总结：TurboQuant-MLX 是 Apple Silicon 本地大模型推理的内存优化利器，通过 K8+V4 混合精度量化在保持输出质量的前提下节省 18-20% 内存，代价是平台锁定和一定的集成复杂度。