QuantMLA：面向低比特 MLA KV 缓存的函数对齐双路径量化 QuantMLA: Function-Aligned Dual-Path Quantization for Low-Bit MLA KV Caching
揭示 MLA 双路径量化误差不对称，函数对齐学习实现首个内容+RoPE 联合 INT4 缓存
前置知识
MLA（Multi-Head Latent Attention，多头潜在注意力）
DeepSeek 系列模型采用的注意力架构。它把每层的 KV 压缩成一个共享低秩“内容潜在”向量 $C\in\mathbb{R}^{n\times d_c}$（同时供给内容 key 和 value），另缓存一条解耦的 RoPE key $K_P\in\mathbb{R}^{n\times d_r}$ 保留位置信息。推理时通过权重吸收（absorbed formulation），注意力得分写为 $S_h=\tau\bar{Q}_{C,h}C^\top+Q_{P,h}K_P^\top$，每 token 每层只需存 $d_c+d_r$（典型 $512+64=576$）个元素。
本文的一切设计都围绕这两条缓存路径展开：内容路径同时影响匹配与聚合，RoPE 路径只影响位置匹配，双路径的功能不对称是全文误差分析与方法构造的根基，不理解 MLA 就无法理解为什么要“分路径”处理。
RoPE（旋转位置编码）
将位置信息编码为对 query 和 key 按 2 维频率对施加平面旋转：每个频率对乘一个 $R_i\in SO(2)$，注意力得分中的 $q^\top R_{\Delta}k$ 只依赖相对位置。因为旋转是逐对独立施加的，任何“融合在 RoPE 之前”的线性变换必须与这些平面旋转可交换，否则会破坏相对位置语义。
这一可交换性约束把 RoPE 路径的合法变换限制为“成对旋转 + 成对绑定缩放”（论文命题 1），是它与内容路径变换空间不同的根源，也解释了为什么 SpinQuant/QuaRot 等通用旋转量化方法不能直接搬到 MLA 的 RoPE 缓存上。
等价变换与离线融合
在不改变全精度输出的前提下重参数化数值分布的量化技巧：如 SmoothQuant 的等比通道缩放、QuaRot 的 Hadamard 正交旋转，配合把逆变换补偿吸收进下游权重，例如 $zW=(zR S)(S^{-1}R^\top W)$，使量化推理时不引入任何在线变换开销，学到的旋转/缩放完全消失在权重里。
QuantMLA 的核心工程目标就是为 MLA 两条路径分别推导这种“可离线融合”的变换空间：内容路径靠无权重 RMSNorm 的旋转等变性融合，RoPE 路径靠与平面旋转可交换融合，两者都是对这一框架的架构特化。
KV 缓存量化与 NMSE
将 BF16 缓存压到 INT8/INT4/INT2，通常逐 token 分组（group size 64）做非对称仿射量化 $q=\mathrm{clip}(\mathrm{round}(x/s+z),0,2^b-1)$。传统方法用缓存重构误差 NMSE $=\|\hat{c}-c\|_2^2/\|c\|_2^2$ 作为校准目标与质量代理。
论文最关键的实证结论是：把两条路径的缓存扰动归一到相同 NMSE（如 1%）后，RoPE 路径引起的注意力输出失真可达内容路径的 1.94×–6.86×，即“缓存重构误差”是糟糕的功能代理——这正是函数对齐目标的出发点。
后训练量化（PTQ）校准
冻结预训练权重，用少量无标注校准文本（如 128 条 WikiText-2 序列）前向捕获激活，以某个重构损失对少量变换参数（旋转、缩放）做几百步 AdamW 梯度优化，取整操作用直通估计器（STE）处理，最后把学到的变换吸收进模型权重，得到一个“纯量化”的部署模型。
QuantMLA 的两条路径变换正是用这种轻量校准学出来的（300 步、32 条保留序列选 checkpoint、种子 0/1/2），理解该流程才能读懂其训练设置、消融实验以及与 SmoothQuant†/QuaRot† 基线的本质差别。
研究动机
MLA 已把 KV 缓存压到每 token 每层 $d_c+d_r=512+64=576$ 个元素（DeepSeek-V3 布局），但存储仍随上下文长度和 batch 规模线性增长，128K 上下文、多请求并发时缓存依然占满 GPU 显存。现有方案（FlashMLA 对 DeepSeek-V3.2 采用 FP8 内容缓存 + BF16 RoPE key，SnapMLA 同样保留高精度 RoPE key）只压内容路径，RoPE key 的低比特量化几乎无人研究，联合低比特压缩是空白。更根本的问题在方法论：以往 KV 量化普遍用缓存重构误差（NMSE）当作质量代理，本文实测发现该代理严重失真——把两条路径的缓存扰动归一到相同的 1% NMSE 后，RoPE 路径造成的注意力输出失真是内容路径的 1.941×–6.864×（覆盖 DeepSeek-V2-Lite/V3-Base/R1、Kimi-K2-Instruct、LongCat-Flash-Lite、GLM-4.7-Flash 六个模型、四个家族、16B–1T 参数、285 个注意力层）。也就是说，只看缓存误差会系统性低估 RoPE 路径的量化代价，这正是此前没人敢把 RoPE key 压到低比特的深层原因。
本文的目标是本文目标分三层。分析层：建立 MLA 双路径量化误差的系统模型，先跨六模型实证刻画匹配误差下两条路径输出失真的不对称，再用精确的有限误差通路推导和“算子-误差”增益模型解释 RoPE 路径误差为何被放大，并做受控干预验证。方法层：据此提出函数对齐的双路径量化框架 QuantMLA——为每条路径推导可完全离线融合的等价变换空间，并用与其功能误差通路对齐的目标学习变换，实现据作者所知首个内容缓存与 RoPE 缓存的联合 INT4 量化（精度接近 BF16），并把内容缓存进一步压到 INT2（C2R4）仍保持竞争力。系统层：开发原生低比特 MLA 注意力内核，把解包与反量化直接融进注意力计算，配合 sink/recent BF16 保护的混合精度缓存策略，在真实 vLLM 服务中兑现收益——128K 上下文 3.59× 持久缓存压缩、缓存压力负载下 5.168× 整任务输出吞吐。
与已有工作不同的是，独特切入角度有三。其一，“按路径定制”而非“统一变换”：作者抓住 MLA 双路径的功能不对称——共享内容潜在 $C$ 同时供养匹配（$\bar{Q}_{C,h}C^\top$）与聚合（$A_hCB_h$），而 RoPE key 只通过位置得分 $Q_{P,h}K_P^\top$ 起作用、其变换必须与 RoPE 平面旋转可交换——由此推出两条路径各自不同的合法变换空间，而不是把 SmoothQuant/QuaRot 的通用配方硬套到 MLA 上。其二，学习目标对齐“功能”而非“数值”：内容路径直接最小化注意力输出重构（自动涵盖匹配、聚合及交互项），RoPE 路径最小化位置 QK 重构并配有有限扰动下的输出失真上界（命题 2），直接回应了“缓存误差≠功能失真”的实证发现。其三，不只观察放大现象，还给出可检验的机理：增益分解 $G_b=T_bA_b+O_b$（算子敏感度、误差能量分配、方向耦合），用冻结预测器与等能量干预验证，使结论从相关性上升到机制解释，并反过来指导“哪条路径需要更精细的变换和目标”。
核心方法
直觉上，既然两条缓存路径以完全不同的方式影响输出，量化它们就不该共用一套变换和损失。QuantMLA 分三步。第一步，推导路径专属的离线可融合变换空间：内容路径用正交旋转 $R_C$ 加通道缩放 $S_C$，缓存重参数化为 $z_C=uR_CS_C$（$u$ 为无权重 RMSNorm 坐标），消费端权重补偿为 $W^f=S_C^{-1}R_C^\top\Gamma W$，因无权重 RMSNorm 旋转等变，全精度计算严格不变且可全部离线融合。RoPE 路径的变换必须与逐频率对的平面旋转可交换，故限制为成对旋转 $R_P=\mathrm{blockdiag}(R_i),\,R_i\in SO(2)$ 与成对绑定缩放 $S_P=\mathrm{blockdiag}(s_iI_2)$，query 侧互逆补偿 $Q'_P=Q_PR_PS_P^{-1},\,K'_P=K_PR_PS_P$，保证 $Q'_PK'^\top_P=Q_PK^\top_P$ 且可融进 pre-RoPE 投影。第二步，函数对齐学习：冻结权重，内容路径最小化归一化注意力输出重构损失 $\mathcal{L}_C$（式 8），RoPE 路径最小化归一化位置 QK 重构损失 $\mathcal{L}_P$（式 10），后者由命题 2 的输出失真有限上界背书。第三步，原生低比特执行：在 FlashMLA 解码路径上把 INT4 解包/反量化融进注意力内核，对 4 个 sink token 与最近 128 个 token 保留量化前 BF16，全部历史 token 共享一个全局 softmax，保持完整全历史注意力。
与已有方法的本质区别有三。第一，设计原则“路径结构决定变换空间，功能误差通路决定学习目标”。通用方法目标都是数值重构；本文从式 (3) 的精确有限误差分解出发：内容路径误差同时进入匹配项 $\Delta A_{C,h}CB_h$、聚合项 $A_h\Delta CB_h$ 与交互项 $\Delta A_{C,h}\Delta CB_h$，必须用注意力输出重构才能一网打尽——score-only 重构会漏掉聚合与交互；RoPE 路径误差只经位置 logits 起作用，命题 2 给出 $\|\Delta Y_P\|_F^2\le\tau^2\big(\sum_h\|\bar{V}_h\|_2^2\big)\big(\sum_h\|\Omega\odot\Delta Z_{P,h}\|_F^2\big)$ 的有限上界，位置 QK 重构因此成为有理论保证的输出控制代理。第二，命题 1 的 RoPE 兼容等价变换：块对角旋转与 RoPE 平面旋转可交换、$S_P$ 为成对标量，互逆因子可全部融入 pre-RoPE 投影、零在线变换——对比 QuaRot† 适配版的 post-RoPE Hadamard 无法融合、必须在线执行。第三，把“为什么放大”变成可验证的理论：$G_b=T_bA_b+O_b$ 分解说明等能量误差因落在敏感坐标与方向耦合不同而产生悬殊功能影响，预测验证 Spearman 达 0.9923/0.9614。
方法步骤详情
流程如下。(1) 误差分析定标：对 6 个模型共 285 层，用 INT4 诱导误差方向单路径扰动并归一到缓存 NMSE 0.01，测响应 $R_b=\mathrm{NMSE}_{out,b}/\mathrm{NMSE}_{cache,b}$，确认 RoPE/内容比值 1.941×–6.864×。(2) 参数化与初始化：内容旋转用 Cayley 参数化（初始单位阵），RoPE 角度直接可训练，缩放在 log 空间优化并按 99.9 分位幅度搜索初始化 $\alpha\in\{0,.125,.25,.5,.75,1\}$；量化器为逐 token 分组 64 的非对称仿射，校准与线上执行共享同一套 scale/零点/取整/截断规则。(3) 函数对齐优化：每路径 300 步 AdamW（内容旋转 LR $3\times10^{-3}$、角度/缩放 $10^{-3}$，10% warmup + cosine，梯度裁剪 1.0），训练库 128 条 WikiText-2 序列（长 2048、batch 4），验证库 32 条，每 20 步按保留集目标选 checkpoint，种子 0/1/2；学习任一路径时另一条保持精确。(4) 离线融合与等价校验：$\Gamma'=S_C$、$W^a=W_aR_C$、$W^{K'}=WK_PR_PS_P$、$W^{Q'}=W_QR_PS_P^{-1}$ 等全部吸收进现有参数，先关闭量化验证全精度等价。(5) 原生内核与混合精度：fused writer 为每个新 token 写 INT4 packed 副本并维护 sink4+recent128 的 BF16 缓冲；解码时 packed 数据与量化元数据分缓冲以支持预取流水线，重构内容 tile 在 QK 与 latent 聚合间复用；token 滑出窗口直接读已有 packed 副本、无需重量化；分区 softmax 在线合并（式 39）保证全历史单一归一化。
技术新颖性
技术新颖性体现在四个维度。架构理解：首次系统刻画 MLA 双路径量化的不对称性，指出“缓存重构误差”作为代理目标在 MLA 上系统性失效（比值最高 6.864×），并给出带机制解释的算子-误差模型；此前 FlashMLA/SnapMLA 只把内容缓存压到 FP8，既未建模误差传播也未联合优化两条路径。变换设计：为 RoPE 路径导出“与位置旋转可交换”这一被先前工作忽视的融合约束，得到成对旋转+成对绑定缩放的最小充分空间（命题 1），相比 SpinQuant/QuaRot 的自由旋转是结构化得多的受约束学习。目标函数：从重构缓存数值升级为重构下游函数（attention output / positional QK），并为 RoPE 路径目标配上不依赖一阶近似的有限扰动输出界（命题 2），这在 KV 量化文献中少见；位置 QK 重构用完整 QK 矩阵而非逐频率对误差，保留了跨频率抵消效应。系统能力：据作者所知首次实现内容+RoPE 联合 INT4 缓存；内核不展开 per-head KV、不做在线变换，物理布局精确到字节（每 64-token 页 20,480 字节），给出 128K 下 3.59× 的真实持久存储压缩与 5.168× 压力吞吐，而非仅 QDQ 模拟精度。
实验结果
逐项看。通用+抽取套件（表 1）：C4R4 下 DeepSeek-V2-Lite 十任务平均 66.54%，距 BF16 66.80% 仅 0.26 点（RTN/SmoothQuant†/QuaRot† 为 61.02/63.44/64.67）；GSM8K 35.25% vs 16.98%/24.56%/30.71%；Moonlight-16B-A3B 73.44%（BF16 73.57%），GSM8K 74.53% 与 BF16 持平。C2R4（内容 INT2）仍有 64.42%/72.70%，超最强同精度基线 7.20/10.84 点。推理+代码套件（表 2）：LongCat-Flash-Lite C4R4 平均 62.01%（超最强基线 4.82 点，LiveCodeBench 仅比 BF16 低 0.29）；GLM-4.7-Flash 53.75%（+2.70，GSM8K 仅低 0.15）；C2R4 下 LongCat 61.64%（HumanEval 71.34%、LiveCodeBench 49.00%），GLM 51.43%（+6.22）。长上下文检索（图 4）：RULER S-NIAH-1/2 平均 99.95%/99.60%，更难的 S-NIAH-3 平均 87.20%，超最强量化基线 7.40 点且所有长度领先。消融（表 3）：MMLU 上 RTN 51.90% → +内容变换 53.12%（+1.22）→ +RoPE 变换 56.31%（+3.19）→ +混合精度策略 57.95%（+1.64），BF16 57.90%，RoPE 路径贡献最大，与误差分析一致。系统（图 5、表 9/10）：128K 持久缓存压缩 3.59×（4K 为 3.22×）；注意力延迟 -7.42%–+4.29%，64K 以上差异 <0.5%；8 GPU 压力服务中 BF16 有 32 次抢占、重算 3,996,255 token、吞吐 2.614 tok/s，C4R4 零抢占、13.512 tok/s（5.168×）。误差分析（表 4–6）：等能量干预中整 token 置换使内容增益×5.918（接近各向同性 6.007×），符号随机化仅×1.048。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 通用+信息抽取十任务（DeepSeek-V2-Lite，C4R4） | 十任务平均准确率 (%) | 66.54 | BF16 66.80；最强量化基线 QuaRot† 64.67 | 距 BF16 仅 0.26 点，超最强量化基线 1.87 点 |
| 通用+信息抽取十任务（Moonlight-16B-A3B，C4R4） | 十任务平均准确率 (%) | 73.44 | BF16 73.57；QuaRot† 72.50 | 距 BF16 仅 0.13 点，超最强基线 0.94 点 |
| 十任务套件（DeepSeek-V2-Lite，C2R4 内容 INT2） | 十任务平均准确率 (%) | 64.42 | SmoothQuant† 57.22（RTN 50.16） | +7.20 点（vs RTN +14.26 点） |
| 十任务套件（Moonlight-16B-A3B，C2R4） | 十任务平均准确率 (%) | 72.70 | SmoothQuant† 61.86（RTN 61.72） | +10.84 点 |
| 推理+代码七任务（LongCat-Flash-Lite，C4R4） | 七任务平均准确率 (%) | 62.01 | BF16 64.26；SmoothQuant† 57.19 | 超最强基线 4.82 点 |
| 推理+代码七任务（GLM-4.7-Flash，C4R4） | 七任务平均准确率 (%) | 53.75 | BF16 55.32；SmoothQuant† 51.05 | 超最强基线 2.70 点 |
| RULER S-NIAH-3 长上下文检索（4K–32K，每点 500 例） | 平均准确率 (%) | 87.20 | 最强量化基线 79.80 | +7.40 点，且所有长度上领先 |
| 持久 KV 缓存存储 @128K（C4R4+sink4/local128） | 相对 BF16 压缩比 | 3.59× | BF16（每层每 token 1,152 字节 → 320 字节） | 3.59×（4K 时为 3.22×） |
| 缓存压力服务（DeepSeek-R1-0528，8 GPU，5×128K 输入+1K 输出） | 整任务输出吞吐 (tokens/s) | 13.512（0 次抢占，0 重算 token） | BF16 2.614（32 次抢占，重算 3,996,255 token） | 5.168× |
| 原生注意力延迟（batch 8/32，1K–1M） | 相对 BF16 缓存延迟变化 | -7.42% 至 +4.29% | BF16 原生分区 | 64K 以上差异 <0.5%，近似无损 |
局限与改进
作者明确承认的局限：5.168× 的吞吐提升是该缓存压力负载特有的（本质是避免容量驱动的抢占与重算），并非注意力内核加速，轻负载下不会出现；serving 对比只有一次 C4R4 运行对一条历史 BF16 基线，未报告置信区间（表 10 原文注明）；GLM-4.7-Flash 在 C2R4 下任务间波动大，说明极端内容压缩的敏感性因模型而异。我补充的观察：RoPE 路径只验证到 INT4，未给出 R2 或联合 INT2，而按其自身误差放大结论这恰是最难的部分；校准全部基于 WikiText-2，代码/数学等分布外的校准鲁棒性未报告；一阶算子模型在 cache NMSE $10^{-3}$ 的小误差下验证，两路径联合时的交叉项（式 18）只做了联合评估而未联合优化；服务评估禁用 prefix caching、并发 ≤16、仅在 H20 硬件与 vLLM 0.10.2 上验证，普适性未知；代码“接收后发布”，当前无法独立复现验证。
独立分析的弱点
独立分析的弱点与改进方向：(1) RoPE 路径的 INT2 缺口——分析已证明 RoPE 误差最易放大，但方法止步 R4；可为位置 QK 目标配更有表达力的编码（分组矢量量化、极坐标/旋转码本），或学习频率对的重要性加权，把位宽预算从低敏感对挪到高敏感对。(2) 学习目标是单层、路径隔离的代理：内容路径重构单层注意力输出、RoPE 路径重构位置 QK，两路径联合后的交叉项与跨层误差累积未进入优化，极端情况下层间漂移可能放大；可引入跨层端到端 NLL 或多层输出重构正则。(3) 保护策略超参（sink=4、recent=128）为手工常数，未随模型/任务自适应；可由命题 2 的界逐层估计敏感度来动态分配保护预算。(4) 校准在 1T 级模型上的开销未报告（逐层捕获激活可能不小），可探索层间共享变换或仅校准代表层。(5) packed 副本为保护 token 冗余存储（每请求每层固定 152,064 字节），超长上下文下虽占比小仍可优化为延迟打包；与 prefix caching/共享前缀的兼容被明确排除在验证范围外，而生产长上下文服务重度依赖它，需要设计共享前缀页的低比特所有权协议。
未来方向
作者提出的方向：将“函数对齐”定位为可推广的设计哲学，可延伸到其他含结构耦合缓存的架构（如混合线性注意力中的状态缓存）；C2R4 在推理/代码套件上已具竞争力，暗示与投机解码、稀疏注意力叠加的空间。基于其成果可自然延伸：(1) 把双路径分析用于训练时低比特（QAT 或 FP8/FP4 训练），在预训练阶段就按路径定制噪声注入与缩放；(2) 将命题 2 的有限输出界从“学习目标的正则”升级为“运行时保护策略的准则”——按界内乘子 $\tau^2\sum_h\|\bar{V}_h\|_2^2$ 逐层动态选保护 token，替代固定的 sink4/local128；(3) 算子-误差模型用于在线监测：由 $T_b$、$A_b$ 的运行统计提前预警某层量化健康度退化并触发重校准；(4) 攻克 RoPE INT2：结合 PolarQuant 类编码与成对绑定结构做码本学习；(5) 内核开源后与上游 vLLM/SGLang 集成，评估 prefix caching、大 batch、多租户真实负载下的端到端收益。
复现评估
复现条件评估：代码承诺接收后发布（当前未开源），但论文工程细节披露异常充分——附录 C.4 给出量化器公式与完整超参表（表 7：各学习率、warmup 比例、种子、组大小 64、尺度初始化搜索集），C.5 给出逐字节存储核算（每 64-token 页 20,480 字节、BF16 保护缓冲 $132\times576\times2=152{,}064$ 字节），D.2 精确描述两个基线适配（SmoothQuant† 的统计缩放 + 验证集选 $\alpha$、QuaRot† 的确定性 Hadamard），可据此重新实现。数据全部公开（WikiText-2 校准，RULER/MMLU/GSM8K 等标准基准），模型均为公开权重。算力：校准每层每路径 300 步、batch 4×2048，16B 级模型单卡 A100/H20 即可（论文未给 GPU 时数）；误差分析涉及 1T 级模型（Kimi-K2-Instruct、DeepSeek-R1）需多卡推理；系统实验需 8×H20。风险点：原生 CUDA 内核与 FlashMLA 特定布局、vLLM 0.10.2/PyTorch 2.8.0 强绑定，跨硬件可移植性未知。总体：代码发布后精度部分复现难度中低，系统部分门槛较高。
论文图表
六个 MLA 模型逐层的内容路径与 RoPE 路径响应 $R_b$ 曲线，旁边附层间 RoPE/内容响应比值的分布（菱形为中位数、须线为四分位距，虚线为等响应参考线）。显示尽管逐层波动明显，RoPE 路径响应在几乎所有层都高于内容路径。
支撑论文核心实证发现（匹配误差下输出失真不对称）的主图，直接引出函数对齐的动机。
285 个注意力层的全景热图：上半部分为每个模型每层的内容（C）与 RoPE（R）输出 NMSE（对数色标），下半部分为对应层间 RoPE/内容响应比值（白色表示等失真，酒红色表示 RoPE 更大），灰色填充表示超出该模型分析深度的位置。
以完整粒度展示不对称性的普遍性与逐层差异，补充 Figure 2 的六模型汇总视角。
四联验证图：(a)(b) DeepSeek-V2-Lite 与 GLM-4.7-Flash 的逐层局部增益及 GP/GC 比值（分别 22/27、24/47 层高于 1）；(c) 冻结预测器在 72 个验证点上的表现（Spearman 0.9923/0.9614），灰色叉为只用算子迹的对照；(d) 等能量干预（符号随机化、整 token 置换、各向同性方向）与预测值的对比。
证明算子-误差模型不只是描述性统计，而是有预测力和因果干预支撑的机制解释。
每模型 36 个验证点上三种预测器的对比：完整“算子+误差分配”预测器 log-ratio MAE 仅 0.1133/0.1173（Spearman 0.9923/0.9614），而“仅缓存误差”对照为 1.5657/0.8790、“仅算子迹”对照为 0.9129/1.4276。
证明误差放大可被准确预测，且必须同时建模算子敏感度与误差能量分配。
batch 8/32、上下文 1K–1M 下 C4R4（含 sink4/local128 保护）与 BF16 的配对注意力延迟。Δ 范围 -7.42%（8×1K）至 +4.29%（8×4K），64K 以上 |Δ|<0.5%（如 8×1M：8440.05 vs 8418.63 µs）。
证明解包/反量化融进注意力后内核速度几乎无损，是系统方案可行的前提。