turboquant-pytorch

TurboQuant PyTorch 是 Google TurboQuant（ICLR 2026）算法的从零 PyTorch 实现，目标是对大语言模型的 KV 缓存进行压缩。它通过两阶段向量量化，在显著降低显存占用的同时尽量保持注意力计算的准确性。

用途该技能用于压缩 LLM 的键值（KV）缓存，将每个坐标压缩到 2–4 比特。其核心思路是保留内积（注意力分数）的准确性，而非追求单个向量的还原保真度。

主要能力算法分为两个阶段：

第一阶段：随机正交旋转 + Lloyd-Max 标量量化（MSE 最优）。
第二阶段：QJL 残差校正——1 比特符号投影，使内积估计无偏。

在 Qwen2.5-3B、8K 上下文下（FP16 基线 289 MB）的压缩比：

4 比特 → 76 MB（3.8 倍）
3 比特 → 58 MB（5.0 倍），为实用最佳点
2 比特 → 40 MB（7.3 倍）

项目结构包含：lloyd_max.py（Lloyd-Max 最优标量量化器）、turboquant.py（核心：TurboQuantMSE、TurboQuantProd、TurboQuantKVCache）、compressors.py（面向真实模型张量的生产级压缩器）、test_turboquant.py（合成验证测试）、validate.py（真实模型 Qwen2.5-3B 验证）。

安装方式 git clone https://github.com/tonbistudio/turboquant-pytorch cd turboquant-pytorch pip install -r requirements.txt

CUDA PyTorch： pip install torch --index-url https://download.pytorch.org/whl/cu128

requirements.txt 包含：torch>=2.0、scipy（用于 Lloyd-Max 码本计算）、transformers、accelerate、bitsandbytes（仅用于真实模型验证）。

关键命令

运行合成算法验证（无需 GPU，但 GPU 可启用速度基准）：python -m turboquant.test_turboquant
在 Qwen2.5-3B-Instruct 上运行真实模型验证（需要 ≥6GB 显存的 CUDA GPU，首次运行下载约 2GB 模型）：python -m turboquant.validate

核心 API

Lloyd-Max 码本：build_lloyd_max_codebook(dim=128, bits=3) 返回 (boundaries, centroids)。
第一阶段 MSE 量化：TurboQuantMSE(dim=128, bits=3)，提供 quantize 与 dequantize。
第二阶段无偏内积估计：TurboQuantProd(dim=128, bits=3, proj_dim=64)，提供 compress 与 inner_product。
KV 缓存封装：TurboQuantKVCache(dim=128, bits=3, proj_dim=64)，提供 append_key、append_value、attention_scores、get_values。
生产级压缩器：TurboQuantCompressorV2（键压缩，支持非对称注意力分数计算）与 TurboQuantCompressorMSE（值压缩，仅第一阶段 MSE 重建）。

常见使用模式

压缩整个模型的 KV 缓存：逐层对 keys 使用 TurboQuantCompressorV2、对 values 使用 TurboQuantCompressorMSE，再用 asymmetric_attention_scores 直接由压缩键计算注意力分数，softmax 后对解压的值做加权求和。
验证压缩质量：对比真实注意力分布与 TurboQuant 估计分布的余弦相似度及 Top-1 匹配率。
大海捞针检索测试：检验压缩后最近邻排序是否保持。
估算内存节省：按层数、KV 头数、序列长度、head_dim、比特数计算压缩后大小与压缩比。

算法细节

随机旋转：用随机正交矩阵 R 旋转单位向量，使各坐标近似服从 N(0, 1/d)，从而支持最优的逐坐标标量量化（Lloyd-Max）。
键使用 QJL 而值不使用：键参与与查询的点积，内积估计的偏差会直接破坏注意力权重，因此 QJL 校正必不可少；值用于 softmax 后的加权求和，小的逐向量 MSE 误差会被平均掉，第一阶段 MSE 量化已足够。
proj_dim 选择：经验默认 proj_dim = head_dim // 2（head_dim=128 → 64；64 → 32；256 → 128）。

比特宽度选择指南

4 比特：3.8 倍压缩，余弦相似度 0.999，Top-1 匹配 87%，适用于质量关键任务。
3 比特：5.0 倍压缩，余弦相似度 0.995，Top-1 匹配 82%，推荐默认。
2 比特：7.3 倍压缩，余弦相似度 0.988，Top-1 匹配 66%，适用于极端内存压力。

适用场景适用于需要在有限显存下运行长上下文 LLM 推理、希望压缩 KV 缓存同时保持注意力准确性的场景，尤其适合对内存敏感、可接受轻微精度损失的部署。

限制与注意事项

真实模型验证需要 ≥6GB 显存的 CUDA GPU，首次运行需下载约 2GB 模型。
构建码本依赖 scipy，首次运行较慢（数值积分），建议按 (dim, bits) 组合缓存码本。
TurboQuantProd 期望输入为单位向量，压缩前需 L2 归一化；compressors.py 内部会处理归一化。
内积估计方差过高时可增大 proj_dim。
若注意力保真度低于预期，需检查向量是否已归一化。

参考资料：TurboQuant 论文（ICLR 2026）、QJL 论文（1 比特残差校正技术）、PolarQuant（相关极坐标方法）。

技能介绍