TurboQuant QJL投影深度剖析：揭秘注意力分数背后的无偏内积估计器原理

TurboQuant 是面向 LLM 推理的 KV cache 量化方案（3-bit keys + 2-bit values），其核心秘密在于 QJL 投影：用 1 bit/维的随机符号草图（quantized Johnson–Lindenstrauss）校正 MSE 量化残差，使注意力分数的内积估计器无偏——E[估计值] = 真实内积。本文用通俗的方式带你拆解开压缩三阶段流水线、QJL 的无偏性证明，以及 quantizer.py 中的关键实现，并附可直接运行的验证脚本路径。

一、先看效果：无偏估计带来什么 🚀

在长上下文推理中，KV cache 是显存和带宽的最大消耗者。TurboQuant 把历史 KV 压成比特串后，注意力分数<q, k>不再"近似"，而是统计意义上无偏地估计真实值——这是排序质量（哪些 token 该被关注）的理论保证。

项目自带的验证测试（README.md 的 Paper Validation 一节）实测：

💡 无偏 ≠ 无损：单次估计仍有噪声，但噪声随维度增大而衰减，且不产生系统性偏移——对 softmax 排序而言，这比"有偏但低方差"的量化更稳。

二、压缩三阶段流水线：从浮点向量到比特串

一个 key 向量x（d 维，通常 d = head_dim = 128）进入 quantizer.py 的TurboQuantProd后，经历三步：

1. 归一化 + 随机旋转：先存下||x||，把x/||x||送到单位球面上，再乘随机正交矩阵 Π。旋转后的每个坐标都服从缩放 Beta 分布（高维下近似N(0, 1/d)）——坐标分布变得"人人一样"，量化器才能按坐标独立设计。旋转矩阵由 rotation.py 的 QR 分解生成（128×128 矩阵仅 64KB，每层一个固定种子）。
2. Lloyd-Max 粗量化（b−1 bit）：针对上述 Beta 分布解连续 1D k-means 得到最优质心，逐坐标查找决策边界。代码书离线预生成在 codebooks/ 目录（如 codebook_d128_b2.json），加载逻辑见 codebook.py。
3. QJL 残差草图（1 bit/维）：把粗量化重建x̃与原始x的残差r = x − x̃投到随机高斯矩阵S（rotation.py），只保留每个投影分量的正负号。d 个符号比特打包成 d/8 个字节，外加||r||一起存储。

以 3-bit key 为例，每维共 3 bit = 0.75 字节，比 bf16 的 2 字节省近2.7x；再叠加 value 的 2-bit 分组量化，就是 README 中 "3b key / 2b val" 的来源。

三、为什么单靠 MSE 量化不够？——偏差的坑

MSE 最优量化只保证"重建误差平方和最小"，并不保证：

E[x̃] ≠ x

一旦估计器带系统偏差，<q, x̃>的期望就偏离真实<q, x>。对注意力打分来说，偏差会整体平移/扭曲 logits，长上下文下误差叠加，排序失真。这正是传统低比特 KV 量化在超长上下文质量滑坡的根源之一。

TurboQuant 的解法很巧妙：不用更多比特去硬压残差，而是花 1 bit/维把残差"无偏地"表达出来——这就是 QJL（quantized Johnson–Lindenstrauss）登场的地方。

四、QJL 投影原理：1 bit 为什么能"无偏"？

核心恒等式（无偏性的全部魔法）：设g ~ N(0, I_d)各分量独立，s = sign(g)，则对任意两个向量：

E[ <g, s'> ] = √(π/2) · <r, y>

其中s' = sign(S r)是残差经随机高斯矩阵S投影后的符号向量，y是查询方向。关键直觉：

• 单个投影g·r是高斯随机变量，其符号与g的夹角余弦的期望恰为√(π/2)·⟨单位r, 单位y⟩（反正切积分的解析结果）；
• 把幅度用存好的||r||还原，得到√(π/2)/d · ||r|| · <Sᵀy, s'>，这个量的期望正好等于⟨y, r⟩；
• 维度 d 越大，该估计的相对方差越小——高维下 1 bit 草图意外地"够用"。

所以 QJL 项是一个无偏的随机估计器：单次有噪声，但期望严丝合缝，噪声还随 d 增大而衰减。

五、两阶段组合：MSE + QJL = 无偏内积估计器

TurboQuantProd 的类注释写得很直白：

反量化内积估计为<y, x̃_mse> + ||r|| · √(π/2)/d · <Sᵀ qjl_signs, y>， 且E[estimate] = <y, x>（无偏）。

分解一下这个恒等式：

也就是说：MSE 部分的偏差恰好是残差内积<y, r>，而 QJL 项的期望恰好补上这块残差——两块拼起来，偏差归零。这就是"b bit 拆成 (b−1) + 1"的深意。

量化入口quantize()（quantizer.py）产出ProdQuantized四元组：MSE 索引、QJL 符号字节、残差范数、原始范数——全部比特打包，历史 token 只占极小显存。

六、注意力分数怎么算？三个融合 Triton 核

读取路径分两种实现。朴素 PyTorch 路径在 attention_score：

1. MSE 贡献：查询先"前向旋转"q @ Πᵀ（注意：不是把 key 旋转回来！），再与查表得到的质心直接点乘——避免物化 d 维反量化 key；
2. QJL 贡献：查询预草图q @ Sᵀ，与解压出的 ±1 符号向量点乘，乘上qjl_scale · ||r||，qjl_scale = √(π/2)/d正是上面推导的无偏常数（见 quantizer.py）；
3. 两项相加得到 logits，再进 softmax。

生产路径则走 triton_kernels.py 的 3 个融合核：turboquant_mse_score（融合解包+查表+点乘）、turboquant_qjl_score（融合符号解包+草图点乘，query 草图每查询只算一次）、turboquant_fused_decode_attention（在线 softmax 一步完成 decode 注意力）。融合的意义：永不落地反量化后的 d 维向量，直接在比特流上算分，省带宽又省显存。

压缩存储由 store.py 的CompressedKVStore管理：按 chunk 追加、惰性拼接成 flat cache，每层用seed=42+层号×7生成独立 Π/S（store.py），避免层间噪声相关；最近 token 由 capture.py 的环形缓冲保留全精度，仅在历史足够长（≥16 token）时走压缩路径（score.py）。

七、实测与落地：验证脚本与集成路径 🧪

想亲手验证无偏性？仓库提供了完整工具链：

已知边界（诚实版，来自 README）：2-bit value 是质量瓶颈（cos_sim≈0.94，追求质量可切 4-bit 的 0.997）；MoE/线性注意力层的 state 不可压缩，收益打折；混合 decode 路径目前会全量反量化历史到 float32，融合核的加速尚未在该路径启用。

八、要点回顾 📌

• QJL 是 TurboQuant 无偏性的来源：1 bit/维的符号草图 +√(π/2)/d常数，把 MSE 量化的系统偏差精确补平；
• (b−1)+1 的比特拆分是精髓：粗量化管精度，QJL 管无偏，各司其职；
• 旋转让一切可预测：随机正交 Π 把坐标分布拉成 Beta 分布，Lloyd-Max 代码书才能逐维最优；
• 工程上零妥协：查询前向旋转/预草图 + 融合 Triton 核，比特流上直接算分，decode 阶段不吃额外显存。

理解了这一套"旋转 → 最优粗量化 → 无偏残差草图"的组合拳，你就拿到了 TurboQuant 3-bit KV cache 压缩的理论钥匙——下一次看到 4x 压缩比时，你知道它背后站着的不是一个更激进的量化器，而是一个数学期望意义下零偏差的估计器。