TurboQuant融合Triton内核深潜：单次遍历直接计算压缩KV注意力的3个技巧-育师

【免费下载链接】turboquant

TurboQuant: Near-optimal KV cache quantization for LLM inference (3-bit keys, 2-bit values) with Triton kernels + vLLM integration

项目地址：https://gitcode.com/gh_mirrors/tu/turboquant

点击查看免费下载

TurboQuant 是一款面向 LLM 推理的KV cache 量化工具，用3-bit keys + 2-bit values的极致压缩配合Triton 融合内核，在 vLLM 里把压缩后的 KV 直接变成注意力输出——省显存、不掉速、几乎不损精度。decode 阶段的真正瓶颈是"从压缩 KV 里算注意力分数"，本文拆解 TurboQuant 用单次遍历直接计算压缩 KV 注意力的 3 个关键技巧，帮你读懂它快在哪。

decode 为什么卡在 KV 读回：先看清 9 步冗余

大模型生成文字时是一个 token 一个 token地往外蹦（decode）。每生成一个新 token，都要拿当前 query 去和所有历史 KV做注意力。历史越长，要读的数据越多。

TurboQuant 把历史 KV 压到约 3 bit/元素，但"压得好"不等于"算得快"。如果用普通 PyTorch 一步步还原再算，一个 KV 向量要走满9 步：

阶段	普通（无融合）9 步流程	痛点
还原 key	1. 解包 MSE 索引（位移）→ 2. 查码本质心 → 3. 反向旋转(d×d 矩阵乘) → 4. 乘范数 → 5. 与 query 点积	要把每个 KV完整还原成浮点向量
残差修正	6. query 过 S 矩阵(d×d 矩阵乘) → 7. 解包 QJL 符号 → 8. 与符号点积 → 9. 乘残差范数	又要物化一整批中间张量

问题核心：N 个 KV 就要做 N 次 d×d 矩阵乘、N 次向量物化，大量中间结果反复落显存再读回。TurboQuant 的解法是——根本不还原 KV，直接在压缩位上算出注意力分数。下面 3 个技巧正是围绕这一点。

技巧 1：反转旋转方向——转 query 一次，而不是转 N 个 key

TurboQuant 的 key 是先把向量做随机正交旋转（y = x @ Pi^T）再量化的。要算注意力分数<q, key>，天然是"先把 key 旋转回去"，也就是对每个 KV 做一次y @ Pi的 d×d 矩阵乘。历史有 N 个 KV，就是 N 次矩阵乘——这就是最大的浪费。

TurboQuant 用了一个漂亮的方向反转：

既然Pi是正交矩阵，<q, Pi·y> == <Pi^T·q, y>。那就别动 key，把query 向前旋转一次（q @ Pi^T），之后对每个 KV 只需q_rot[j] × centroid[idx[j]]逐维相乘累加。

一句话总结：把"每个 key 转一次"变成"整个 query 只转一次"。d×d 矩阵乘从 N 次降到 1 次，而且从头到尾没有物化任何一个 d 维还原向量。这正是第一个融合内核turboquant_mse_score的核心，见 triton_kernels.py 里的关键注释。

技巧 2：预计算 query 草图，符号位就地展开直接点积

key 的量化分两层：MSE 主层（b-1bit，查码本）+QJL 残差层（每维 1 bit 的符号）。符号位是 8 个挤在一个字节里的，怎么用最省？

第二个内核turboquant_qjl_score做了两件事：

草图只算一次：query 过 S 矩阵得到q_sketch = q @ S^T，每个 query 只算一遍，然后对 N 个 KV 复用。
符号就地展开：内核里直接把位拆成{-1, +1}，与草图逐维相乘累加，再乘一个常数sqrt(π/2)/d。

妙处在于这套结构是无偏估计：E[估计内积] = 真实内积，也就是压缩后算出来的分数在期望上和精确值一致，从数学上保证了"压得狠但分不偏"。相关推导见 quantizer.py 的类注释。

技巧 3：在线 softmax 单趟融合，全程不落地 FP16 KV

前两个技巧算出分数，第三个技巧把分数 + softmax + 加权求值塞进同一个内核，一次遍历压缩 KV 就出最终注意力输出。这是"单次遍历"真正的落点，见 triton_kernels.py 的设计说明。

它借鉴 Flash-Attention 的在线 softmax，只维护三个状态变量：

状态	含义	更新时机
`m_i`	运行中的最大值	每扫一块 KV 就更新
`l_i`	运行中的 exp 求和	每扫一块就更新
`acc`	运行中的加权和	旧值按修正因子缩放后累加新值

扫完所有 KV 块，最后acc / l_i一次归一化就得到输出。整个过程：

读取的是压缩 KV（约 3 bit/元素），带宽占用直接降到零头的量级；
从不在显存里生成完整的 FP16 KV，省的不只是容量，还有反复读写；
值向量用分组量化（v * scale + zero）在寄存器里现解现用。

对应入口是 turboquant_fused_decode，它会先预旋转/草图 query，再一把调起内核。

三个内核如何拼成一次完整 decode 注意力

拆开看是 3 个技巧，用起来是一条流水线。顶层封装 turboquant_attention_score 把流程串起来：

对 query 做一次q @ Pi^T（技巧 1 的旋转）和q @ S^T（技巧 2 的草图）；
调内核 1得到 MSE 分数；
调内核 2把 QJL 分数原地加到同一块输出上；
需要完整输出时，交给内核 3做在线 softmax + 值聚合。

三步共用同一份"预旋转 query"，避免了重复计算，这就是"直接计算"的含义——没有任何一步把 KV 还原成原始浮点再走标准注意力。

压缩 KV 从哪来：写入路径与码本速览

只讲读取会不完整。TurboQuant 的 KV 是写入时就压好的，整条写路径分工清晰：

捕获：capture.py 用一个环形缓冲区暂存最近的精确 token，满了才把最老的一批刷成压缩块，保证 decode 热路径上没有逐 token 量化的开销。
存储：store.py 按块追加，首次读取时惰性拼成扁平缓存，读多写少时缓存命中、免拼接。
码本：codebook.py 用Lloyd-Max算法为旋转后的 Beta 分布求最优质心，结果缓存在 codebooks/ 目录（如 codebook_d128_b3.json）。
值量化：kv_cache.py 对 value 做 2-bit/4-bit 分组量化 + 位打包。
vLLM 集成：integration/vllm.py 以极小的 monkey-patch 面挂钩，支持off / capture_only / hybrid / full_tq四种模式。

实测收益：省 30GB 显存、上下文翻倍

技巧值不值，看数字。项目在 RTX 5090 上跑 Qwen3.5-27B（dense，4-bit 权重）的实测对比：

指标	基线（bf16 KV）	TurboQuant（3b key / 2b val）
Prefill tok/s（30k 上下文）	1,804	1,907（+5.7%）
Decode tok/s（30k 上下文）	1.264	1.303（+3.1%）
释放 KV 显存	—	30.0 GB（4 卡合计）
最大 token 容量	457,072	914,144（2.0x）
峰值激活显存	644.6 MB	599.2 MB（-7.0%）

要点：省内存的同时速度不降反微升，纯 dense 架构下 KV 压缩比可达约4.4x。质量上 3-bit key 的余弦相似度达 1.000000（近无损），瓶颈在 2-bit value——对质量敏感的场景可改用 4-bit value（cos_sim 0.997）。

快速上手：安装与运行基准

想亲手验证，两步即可：

git clone https://gitcode.com/gh_mirrors/tu/turboquant pip install -e .

无 GPU 也能跑论文定理验证：python proof.py（A/B 对比基准需 4× RTX 3090 + Qwen3.5-27B-AWQ）。
想深入读代码，主线是 turboquant/ 目录，从 triton_kernels.py 与 quantizer.py 切入最容易。

适用边界与局限

仅压缩 full-attention 层：线性注意力 / Mamba 类层的 state 不可压缩，MoE 混合模型收益会打折。
2-bit value 是精度瓶颈：质量敏感场景建议上 4-bit value。
hybrid 路径当前仍会全量反量化历史：完全依赖融合内核的full_tq模式仍在演进中，选型时留意 integration/vllm.py 里对各模式的说明。

TurboQuant 的精髓一句话概括：不还原 KV，直接在压缩位上算注意力——旋转方向反转、草图预计算、在线 softmax 单趟融合，三个技巧合力，把 decode 的 KV 读回从"瓶颈"变成"顺路的事"。