PN-Mix：用查询—键成对范数改进低比特 KV 缓存量化

导语

长输出场景中，模型需要保留越来越多的键值（KV）缓存。缓存规模增长不仅消耗显存，也会让解码阶段频繁读取数据，形成容量和带宽瓶颈。量化可以压缩缓存，但直接采用 2 bit 等低精度设置，往往会明显放大注意力计算误差。PN-Mix 关注其中的键缓存，并尝试回答一个关键问题：哪些维度的量化误差最值得优先控制？

核心思路：从“键的统计”转向“查询—键的作用关系”

已有混合精度键缓存方法通常根据键本身的统计特征评估通道重要性，再为重要通道分配更高精度。PN-Mix 的分析则进一步落到 RoPE 的通道对层面。由于旋转位置编码以通道对参与计算，论文考察了注意力 logits 如何受到这些通道对量化误差的扰动。

其核心观察是：当某个 RoPE 通道对对应的查询范数与键范数乘积较大时，该通道对的量化误差更可能显著改变注意力 logits。因此，PN-Mix 使用查询—键成对范数乘积作为显著性依据，把更高精度集中到这些通道对，而不是平均分配精度或只看键缓存自身的数值特征。

这一过程在预填充阶段完成。系统根据输入计算通道对的重要性，随后在整个解码过程中复用这一精度分配结果。这样的设计利用了当前请求的信息，同时避免在每个生成步骤反复执行重要性估计，适合长序列、长输出的推理流程。

实验结果与边界

论文报告显示，在标准推理、代码生成和 AIME 基准上，PN-Mix 的平均精度为 2.26 bit 时，平均准确率可达到或超过对比的混合精度方法；后者使用的平均精度为 2.32—2.39 bit。进一步的 Pareto 扫描表明，PN-Mix 可以逐步降低到 1.5 bit，性能下降相对平滑，而对比方法在低于 2.25 bit 后出现更明显的性能下滑。

这些结果说明，低比特量化的关键不只是选择多少 bit，也在于把有限的高精度预算放在真正影响注意力结果的位置。尤其对采用 RoPE 的 decoder-only LLM，通道对级别的误差分析提供了比单通道统计更贴近计算机制的分配依据。

意义与影响

PN-Mix 的价值主要在于把 KV 缓存压缩从静态、粗粒度的通道筛选，推进到输入自适应的通道对分配。对于长输出服务，这可能帮助系统在相近甚至更低的缓存位宽下，减少显存占用和数据搬运压力，同时保留更多模型能力。

不过，现有结论针对的是 RoPE-based decoder-only LLM，并且方法聚焦于键缓存；它是否同样适用于其他位置编码、不同注意力结构或更广泛的 KV 联合量化场景，仍需进一步验证。整体来看，PN-Mix 展示了一条清晰路径：利用注意力实际使用的交互关系，而非仅凭缓存自身统计量，指导混合精度预算。

来源：Hugging Face Daily Papers