QuantMLA：用于低位 MLA KV 缓存的功能对齐双路径量化

多头潜在注意力 (MLA) 通过紧凑的内容缓存和解耦的 RoPE 路径实现富有表现力的多头注意力，但缓存内存仍然随上下文长度和批量大小线性扩展。在这项工作中，我们建立了 MLA 双路径量化误差的系统模型，描述了它们对注意力输出失真的独特影响，并解释了 RoPE 路径误差的显着放大。在此分析的指导下，我们引入了 QuantMLA，一种用于低位双路径量化的功能对齐框架。我们推导出特定于路径的变换空间，该空间保留全精度计算，同时保持完全融合到离线模型参数中，从而消除在线变换开销。在这些空间内，QuantMLA 学习具有功能对齐目标的特定于路径的转换：注意力输出重建捕获内容路径的耦合匹配和聚合错误，而位置 QK 重建保留 RoPE 引起的注意力logits分量，并给出输出失真的理论界限。据我们所知，在四个 MLA 模型系列中，QuantMLA 首次实现了内容缓存与RoPE缓存的联合INT4量化，且精度下降最小。进一步将内容缓存压缩到 INT2，同时保留 INT4 的 RoPE键缓存，可以在具有挑战性的推理和代码基准测试中保持具有竞争力的性能。我们开发了一个原生低位 MLA 注意力内核，它将解包和反量化直接集成到注意力计算中。物理缓存布局在 128K 上下文下提供 3.59 倍的压缩，而缓存压力服务工作负载的整体作业输出吞吐量比 BF16 高 5.168 倍。代码计划在论文被接收后发布。