MixKVQ: Query-Aware Mixed-Precision KV Cache Quantization for Long-Context Reasoning
MixKVQ: Query-Aware Mixed-Precision KV Cache Quantization for Long-Context Reasoning
发表时间: 2026-07 · ACL 2026
原文: https://aclanthology.org/2026.acl-long.326
作者/机构:Tao Zhang, Ziqian Zeng, Hao Peng, Huiping Zhuang, Cen Chen (华南理工大学, 北京航空航天大学, 琶洲实验室)
速读
一句话结论
MixKVQ 是一种查询感知混合精度 KV Cache 量化方法,通过结合 Key 通道的量化难度与 Query 的相关性来动态分配位宽,在极低显存占用下实现了与全精度模型相当的长上下文复杂推理能力。
要解决什么问题
大语言模型在生成长思维链时,自回归机制会导致 KV Cache 的显存占用随序列长度线性暴增,造成严重的显存容量和访存带宽瓶颈。现有的低比特(如 2-bit)KV Cache 量化方法在处理复杂推理任务时会导致精度严重崩塌,因为长序列极易产生量化误差累积,哪怕破坏了一个关键数值也会导致整条逻辑链失效。原有做法的卡点在于:固定精度量化无法处理 Key Cache 中存在极端离群值的通道;而现有的混合精度量化(如基于层或通道的启发式方法)仅根据 Key 或 Value 的数值量化误差(即缩放因子大小)来分配位宽。这种策略的致命缺陷是,如果一个 Key 通道的激活值很大,但与之计算注意力的 Query 激活值很小,那么它对最终注意力分数的实际影响微乎其微。仅凭数值误差来分配高精度,会把宝贵的显存位宽浪费在对注意力计算不重要的通道上,导致真正影响推理逻辑链的关键通道被过度压缩。
怎么做的
核心思路是将位宽分配的决定权从“单看 Key 的误差”转变为“综合评估 Key 的量化难度与 Query 的相关性”。由于 Value Cache 的误差分布较均匀且缺乏明显的离群值,方法对其直接采用统一的 2-bit 逐 Token 量化;而对于包含大量离群值的 Key Cache,则采用逐通道的混合精度量化。方法引入了一个轻量级的查询感知启发式算法,定义了显著性得分 $\mathcal{A}_d$ 来评估第 $d$ 个通道的重要性,它由两个低成本的统计量相乘构成:
第一个部件是重要性得分 $\mathcal{I}_d$,用于衡量 Query 的相关性,计算为该通道在 Query 序列上的平均绝对幅值:
效果如何
实验在单张 80GB A800 显卡上进行,评估了 DeepSeek-R1-Distill 系列(Qwen-14B/32B、LLaMA-8B)、Llama-3.1-8B 和 Mistral-7B 等模型。对比基线包括全精度 BF16,以及代表非对称量化的 KIVI 和 KVQuant、代表旋转变换处理离群值的 RotateKV、代表滑动窗口量化的 SKVQ、代表层级混合精度的 KVTuner。在 AIME、MATH 500、GPQA 和 LiveCodeBench 等复杂推理任务中,MixKVQ 显著优于其他基线。在 Qwen-32B 模型上,MixKVQ 在平均等效位宽仅为 2.3-bit 的情况下,取得了 66.04% 的平均准确率,极其逼近全精度 BF16 的 67.84%;相比之下,4-bit 的强基线 RotateKV 准确率降至 64.51%,而 2-bit 的 KIVI 则暴跌至 58.89%。在吞吐量测试中,该方法在相近显存峰值下,支持的 Batch Size 提升了 2.25 倍,吞吐量达到了全精度基线的 2.63 到 2.81 倍。该方法的局限性在于:尽管有惰性更新机制,张量转换的计算开销依然不可忽视;解码阶段管理不连续的混合精度内存块可能会引入额外的延迟;此外,该方法目前针对的是分组查询注意力,尚未覆盖多头潜在注意力等差异较大的新型机制。
主要贡献
大型语言模型(LLMs)通过生成长思维链(CoT)在复杂推理能力上取得了显著进展,但这种进展伴随着由不断扩展的键值(KV)缓存带来的巨大内存和延迟开销。在自回归生成过程中,KV缓存的内存占用随序列长度呈线性增长,并造成严重的内存带宽瓶颈。虽然KV缓存量化是一种极具前景的压缩技术,但现有的低比特(如2-bit)量化方法在复杂推理任务上往往表现出严重的性能退化。具体而言,固定精度量化难以处理键(Key)缓存中的异常通道,而现有的混合精度策略通常仅基于量化误差来分配位宽,未能准确识别需要高精度表示的关键组件。
本文的核心研究目标是在大幅降低内存占用的同时,不损害复杂推理任务中的注意力计算保真度。作者发现,一个有效的低比特KV缓存量化策略必须同时考虑两个因素:键通道的内在量化难度以及它与查询(Query)的相关性。基于这一洞察,本文提出了MixKVQ,这是一种新颖的量化方法,引入了一种轻量级的、感知查询的启发式算法,以识别并以更高精度保留关键的键通道,同时对值(Value)缓存应用逐token(per-token)的量化。在复杂推理数据集上的实验表明,MixKVQ显著优于现有的低比特方法,在大幅减少内存占用的情况下实现了与全精度基线相当的性能。
背景知识与关键观察
Transformer与KV缓存机制。Transformer层的核心是自注意力机制。第$l$层查询向量$Q_i^l$的输出$o_i^l$通过关注键和值向量序列计算得出:$\mathbf{a}_i^l = \text{softmax}\left(\frac{Q_i^l(K^l)^\top}{\sqrt{D}}\right), \quad \mathbf{o}_i^l = \mathbf{a}_i^l\mathbf{V}^l$,其中$D$是模型隐藏层大小,$\mathbf{a}_i^l$表示第$l$层第$i$个token的注意力权重向量。在自回归生成期间,矩阵$\mathbf{K}^l$和$\mathbf{V}^l$包含所有先前时间步的键和值向量。为了避免昂贵的重新计算,这些矩阵存储在KV缓存中,其大小随序列长度线性增长。
KV缓存量化原理。KV缓存的内存占用可能成为推理瓶颈。量化通过以低比特格式存储缓存来解决此问题。采用$B$-bit非对称量化方案,用反量化表示$\tilde{X}$近似张量$X$:$Q(X) = \text{round}\left(\frac{X - z}{s}\right)$,$\tilde{\mathcal{X}} = \mathcal{Q}(\mathcal{X}) \cdot s + z$。参数包括零点$z = \min(X)$和缩放因子$s = (\max(\mathbf{X}) - \min(\mathbf{X})) / (2^B - 1)$。由张量动态范围决定的缩放因子$s$是其量化敏感性的直接代理。较大的$s$意味着更宽的值分布被压缩到$2^B$个离散级别中,导致粒度更粗。具体而言,舍入操作为每个$x_i \in X$引入了量化误差$|x_i - \tilde{x}_i|$。该误差在数学上受限于$|x_i - \tilde{x}_i| \leq s / 2$。单个异常值可能会放大$s$,从而显著增加此误差界限并降低所有其他元素的近似质量。
自回归大语言模型中的量化误差累积。在自回归LLM中,KV缓存量化误差在模型深度(层到层)和序列长度(token到token)两个维度上累积。虽然单个token或层的量化误差可能微不足道,但其在数千个token上的累积效应可能很大(【索引编号[16],Kvtuner: Sensitivity-aware layer-wise mixed-precision kv cache quantization for efficient and nearly lossless llm inference+2025+ICML】)。这会导致token翻转和级联生成错误等现象。这种误差累积在数学推理任务中尤其有害,其中单个关键值的损坏可能会使整个逻辑链失效,从而浪费大量计算资源。
键缓存(Key Cache)通常更重要。对于Qwen2.5-14B-Instruct模型(第0层,第2个头)的键和值缓存的2-bit绝对量化误差进行可视化。键缓存中的某些通道表现出明显更大的量化误差。相比之下,值缓存的误差分布更均匀,缺乏明显的异常值。实验结果证实,为键缓存保留更高的精度对于维持模型性能更为关键。这与先前的工作(【索引编号[19],KIVI: A tuning-free asymmetric 2bit quantization for KV cache+2024+ICML】等)一致。因此,核心挑战是开发一种在受限内存预算下最小化注意力分数损失的键缓存精度分配策略。
固定精度方法难以处理异常值。现有的固定精度方法在2-bit等低比特宽下表现不佳。虽然这些方法可以充分压缩值缓存,但它们难以处理具有显著异常值的键缓存通道。这导致了巨大的量化误差以及在推理任务上的严重失败。
现有混合精度方法在位宽分配上存在缺陷。当前的逐层和逐通道方法在位宽分配方面表现出缺陷。由于其静态特性,逐层方法容易对敏感层进行过度激进的量化。逐通道启发式方法将更高的位宽分配给具有更大缩放因子的键通道。然而,Query激活的幅度与Key缩放因子之间几乎没有相关性,皮尔逊相关系数仅为0.16。由于高Query激活表明对保持注意力保真度具有更大的重要性,仅依赖缩放因子的方法可能会以更高精度保留实际上对准确计算注意力分数并不关键的键通道。此外,仅靠缩放因子的判别能力有限:$s$的分布高度集中,第0头中80%的值集中在[2.80, 4.46]的狭窄范围内,使得区分真正关键的通道变得困难。综合这些观察结果表明,单独的$s$不足以识别显著通道。
引入查询信息的必要性。KV缓存量化的目标是在保持注意力计算保真度的同时减少KV缓存内存占用。分析认为,如果对应的查询激活很小,具有大幅度激活的键通道可能并不关键。保留此类通道带来的收益递减,并且是对有限比特预算的低效使用。这促使需要一种包含Query信息以实现有效精度分配的更具信息量的指标,特别是在低比特状态下。
方法细节
MixKVQ方法与目标。基于上述见解,本文提出了MixKVQ,这是一种新颖的方法,引入了一种轻量级、感知查询(query-aware)的启发式算法,以识别并以更高精度保留关键通道。鉴于值缓存可以被充分压缩,该方法的主要目标是减轻注意力分数中引入的量化误差。
注意力分数量化误差的数学定义。令$\mathbf{Q} \in \mathbb{R}^{L_q \times D}$和$\mathbf{K} \in \mathbb{R}^{L_k \times D}$分别表示查询和键矩阵,其中$D$是隐藏维度。当$\mathbf{K}$被量化为$\tilde{\mathbf{K}}$时,softmax之前的注意力分数误差记为$\mathbf{E}_{attn}$,定义为:$\mathbf{E}_{attn} = \mathbf{Q}(\mathbf{K} - \tilde{\mathbf{K}})^T$。考虑第$i$个查询token和第$j$个键token的误差项,记为$E_{i,j}$。令$\epsilon_{j,d} = \mathbf{K}_{j,d} - \tilde{\mathbf{K}}_{j,d}$表示键在token $j$和通道$d$处的量化噪声。对数误差$E_{i,j}$聚合了跨通道维度的误差贡献:$E_{i,j} = \sum_{d=1}^D \mathbf{Q}_{i,d} \cdot \epsilon_{j,d}$。
显著性得分(Salience Score)的推导与计算。为了最小化对注意力机制的影响,目标是识别并保留具有最高预期误差贡献幅度$\mathbb{E}[|\mathbf{Q}_{i,d} \cdot \epsilon_{j,d}|]$的通道。本文将此期望近似为其各自预期幅度的乘积:$\mathbb{E}[|\mathbf{Q}_{i,d}|] \cdot \mathbb{E}[|\epsilon_{j,d}|]$。通过一种启发式指标,即显著性得分(Salience Score, $\mathcal{A}_d$)来实现这一点,该得分源自两个低成本统计量:
* 重要性得分(Importance Score, $\mathcal{I}_d$):使用序列长度$L_q$上查询通道$d$的平均绝对幅度来估计$\mathbb{E}[|\mathbf{Q}_{i,d}|]$:$\mathcal{I}_d = \frac{1}{L_q} \sum_{i=1}^{L_q} |\mathbf{Q}_{i,d}|$。
* 敏感度得分(Sensitivity Score, $\mathcal{S}_d$):对于位宽为$B$的均匀量化器,量化误差$\epsilon_{j,d}$受限于缩放因子的一半(即$|\epsilon_{j,d}| \le \frac{\mathcal{S}_d}{2}$)。因此,使用通道$d$的缩放因子$\mathcal{S}_d$来估计$\mathbb{E}[|\epsilon_{j,d}|]$:$\mathcal{S}_d = \frac{\max(\mathbf{k}_d) - \min(\mathbf{k}_d)}{2^B - 1}$,其中$\mathbf{k}_d$表示跨当前token的通道$d$中的键值向量。
最终,显著性得分$\mathcal{A}_d$定义为这些组件的乘积:$\mathcal{A}_d = \mathcal{I}_d \cdot \mathcal{S}_d$。具有高$\mathcal{A}_d$的通道被认为是关键通道,因为它们表现出高查询相关性和高量化敏感性,从而保证了高精度的保留。
三层混合精度策略的设计。本文设计了一种三层混合精度策略。虽然先前的研究表明4-bit量化在模型性能和内存占用之间提供了良好的平衡,但MixKVQ通过基于通道显著性分配位宽来进一步优化这种平衡。为了保持生成质量,最关键的通道以全精度(BF16)保留。同时,为了最大化压缩,将中等关键通道量化为UINT4,并将非关键通道积极压缩为UINT2。形式上,引入了两个应用于显著性得分$\mathcal{A}_d$的阈值$\tau_{BF16}$和$\tau_{UINT4}$:
* 高精度(BF16):$\mathcal{A}_d > \tau_{BF16}$的通道被识别为高度关键并保留在BF16中。
* 中等精度(UINT4):满足$\tau_{UINT4} < \mathcal{A}_d \leq \tau_{BF16}$的通道被认为是中等关键的,并量化为UINT4。
* 低精度(UINT2):具有$\mathcal{A}_d \le \tau_{UINT4}$的通道被认为是非关键的,并压缩为UINT2。
MixKVQ的整体工作流程。MixKVQ的整体工作流程如图4所示。对Key缓存应用逐通道(per-channel)量化,对Value缓存应用逐token(per-token)量化。大小为$R$的缓冲区以全精度存储token,直到达到缓冲区大小,之后对它们进行组量化。对于Key缓存,参数$\mathcal{I}_d$和$\mathcal{S}_d$会定期更新(每$R$个token)。在更新期间,$\mathcal{I}_d$是通过计算当前窗口内Query激活的绝对幅度的平均值来得出的。相反,Value缓存经历均匀的2-bit逐token量化。
实验环境
- 数据集:用于数学推理的AIME 2024-2025和MATH-500;用于研究生级科学推理的GPQADiamond;用于代码生成的LiveCodeBench(使用2025年1月1日至4月6日的子集);用于长上下文生成的LongBench;用于运行时效率评估的ShareGPT。
- 模型架构:Deepseek-R1-Distill-Qwen-14B/32B,Deepseek-R1-Distill-LLaMA-8B,Llama-3.1-8B-Instruct,Mistral-7B-Instruct-v0.3,Llama2-13B-chat,DeepSeek-R1-Distill-Qwen-7B。
- 硬件配置:单张NVIDIA A800 GPU (80GB)。
- 软件配置:所有推理评估采用采样温度0.6,top-$p$为0.95。为了确保公平比较,所有量化方法(包括MixKVQ)统一使用组大小$G = 32$和残差长度$R = 128$。
实验结果
- 复杂推理任务评估:在AIME 2024-2025, MATH 500, GPQADiamond和LiveCodeBench数据集上,MixKVQ在所有评估模型中始终优于现有方法(KVquant, KIVI, KVTuner, RotateKV)。例如,在Qwen-32B模型上,MixKVQ实现了66.04%的平均准确率,紧随BF16基线的67.84%。相比之下,4-bit基线RotateKV的平均准确率为64.51%(下降3.33%),而2-bit基线KIVI降至58.89%(显著下降8.95%)。这表明MixKVQ可以有效减轻通常会损害推理能力的量化误差。(图表引用:Table 3)
- 长上下文生成准确率:在LongBench上对Mistral-7B和Llama-3.1进行了评估。与BF16和具有竞争力的基线相比,MixKVQ将有效位宽降至2.70 bits,性能下降微乎其微。这一结果突显了混合精度策略的有效性,在提供高保真模型性能和内存效率方面超越了量化基线。(图表引用:Table 4)
- 运行时效率比较:严格遵循vLLM评估设置,使用从ShareGPT合成的工作负载评估MixKVQ的运行时效率。将批处理大小推至内存饱和,并比较Llama2-13B-chat上MixKVQ(残差长度32/128)与FP16基线的吞吐量和峰值内存使用量。在最大内存使用量相似的情况下,MixKVQ可实现高达2.25倍的批处理大小,并提供2.63倍至2.81倍的吞吐量。预计这些性能增益将随着序列长度的延长而放大。(图表引用:Figure 5)
结论
本文提出了一种无需任何微调的新颖KV缓存量化算法MixKVQ。MixKVQ根据键通道的内在量化难度及其与查询的动态相关性这两个因素来分配精度。基于该策略,MixKVQ对键缓存进行逐通道混合精度量化,对值缓存进行逐token量化。评估表明,MixKVQ在复杂推理任务的极低位宽下实现了最佳性能。未来工作计划将MixKVQ集成到vLLM等高性能服务框架中,以解锁大量额外的吞吐量增益。
局限性:尽管MixKVQ实现了大幅压缩,但张量转换期间的计算开销仍不可忽视。解码期间管理不连续的混合精度内存块可能会引入延迟。该研究未涵盖所有注意力机制,特别排除了多头潜在注意力(MLA)。此外,延迟分析主要集中在受内存限制的生成阶段,对提示处理阶段(尤其是跨多个KV序列的批处理压缩操作期间)的计算瓶颈探索不足。
附录与补充细节
量化误差界限的数学推导。非对称量化定义为:$q_i = \text{round}\left(\frac{x_i - z}{s}\right)$,$\tilde{x}_i = q_i \cdot s + z$。绝对量化误差为$|x_i - \tilde{x}_i| = |x_i - (q_i \cdot s + z)| = |(x_i - z) - q_i \cdot s|$。提取缩放因子$s$可得:$|x_i - \tilde{x}_i| = \left|s \cdot \left(\frac{x_i - z}{s} - q_i\right)\right| = s \cdot \left|\frac{x_i - z}{s} - q_i\right|$。令$y_i = \frac{x_i - z}{s}$,则$q_i = \text{round}(y_i)$。由于舍入操作映射到最近的整数,绝对差值最多为0.5,即$|y_i - \text{round}(y_i)| \leq \frac{1}{2}$。将此不等式代回误差方程:$|x_i - \tilde{x}_i| \leq s \cdot \frac{1}{2} = \frac{s}{2}$。因此,任何元素的量化误差都受限于缩放因子的一半。
KVTuner层级策略的局限性。KVTuner通过利用校准数据静态识别被认为“不太关键”的层。为了满足目标内存预算,这些层随后受到激进的K2V2量化。然而,这种层级方法存在根本限制。即使在这些所谓的非关键层中,也存在由于异常值而难以量化的特定维度。对整个层统一应用激进的量化策略会导致这些关键维度的大量信息丢失,引入大量误差,从而降低模型在复杂多步推理任务上的性能。
更深层中查询与键缩放的相关性。计算了Qwen-2.5-14B所有层中查询和键缩放之间的皮尔逊相关系数。结果(第0层0.16,第16层0.15,第31层0.25,平均0.21)证实,这种低相关性在更深的层中依然成立。
双目标优化的阈值搜索。使用OPTUNA框架在[0.1, 2.0]搜索空间内联合搜索$\tau_{BF16}$和$\tau_{INT4}$。优化目标为:1. 最大化GSM8K基准上的准确率,以确保量化过程保留复杂的思维链推理能力;2. 最小化每参数的平均有效位宽$B_{eff} = \frac{1}{N} \sum_{i=1}^N b_i(\tau_{BF16}, \tau_{INT4})$。使用TPE采样器执行30次试验构建帕累托前沿。最优阈值因架构而异:R1-Llama-8B为(1.44, 0.79)达到2.7 bits;R1-Qwen-7B为(0.63, 0.41)达到3.4 bits;较大的模型表现出更强的压缩鲁棒性,R1-Qwen-14B和32B模型分别为(1.52, 1.60)和(1.85, 1.58)达到2.3 bits。
MixKVQ的系统级实现细节。为了确保与现代LLM架构兼容,量化粒度与注意力机制严格对齐。对于采用分组查询注意力(GQA)的模型,重要性得分在KV头组级别计算,即聚合对应于共享KV头的所有查询头的查询幅度。存储布局分为三部分:1. 量化存储$Q(\mathbf{X}_{K/V})$:以低位连续张量(UINT4/UINT2)打包历史状态以最大化吞吐量;2. 稀疏异常值存储$\mathbf{X}_{K_{BF16}}$:以全BF16精度保留显著通道,使用稀疏格式;3. 高精度残差缓冲区$\mathbf{X}_R$:全精度暂存最近的token。
通过延迟更新进行摊销调度。为了减轻频繁重新量化的计算开销,采用了由残差长度$R$控制的延迟更新策略。解码期间,新生成的状态最初附加到全精度缓冲区$\mathbf{X}_R$。仅当缓冲区长度达到$R$时,才触发量化、通道选择、异常值提取和位打包等密集型任务。处理后合并到主缓存并重置缓冲区。这不仅在$R$个解码步骤中摊销了开销,而且充当了时间稳定窗口。由于通道显著性在最近token上表现出瞬时波动,将这些状态隔离在$\mathbf{X}_R$中可推迟量化决策,直到token移出局部窗口。
高效的在线显著性估计。计算KV通道的累积重要性通常需要扫描整个查询历史,这在计算上是禁止的。为了解决这个问题,在缓存中维护了查询幅度的运行累加器。在每个解码步骤,将当前查询token的幅度添加到该累加器。此外,对于包含旋转位置嵌入(RoPE)的架构,显著性评估在旋转变换之后进行,以确保指标准确反映旋转空间中的注意力分布。
超参数与查询感知组件的消融研究。研究了组大小$G$和残差长度$R$。固定残差长度为128,改变组大小为32、64和128。发现困惑度(PPL)随组大小增加而降低,组大小的选择极大影响长输入下的KV缓存压缩效果。固定组大小为32,改变残差长度为32、64、96、128和256。结果显示残差长度与模型准确率之间没有一致的模式,但足够大的残差长度对困难任务仍然至关重要。为了验证复合指标$\mathcal{A}_d = \mathcal{I}_d \times \mathcal{S}_d$的有效性,与仅依赖通道幅度决定精度($\mathcal{A}_d = \mathcal{S}_d$)的Error-only基线进行了比较。完整的MixKVQ在AIME 2024-2025基准测试中优于该基线,证实了明确纳入依赖于查询的重要性对于在复杂推理任务中保持模型保真度是不可或缺的。
操作级细分与额外实验。在R1-Qwen-7B上,通道选择占每层执行时间的2.17%,但实现了超过79%的KV缓存内存节省(从16-bit压缩到3.4-bit)。在DeepSeek-R1-Distill-Qwen-7B上的额外基准测试表明,MixKVQ(有效位宽3.4)在AIME、MATH 500等任务上的平均准确率(48.49%)优于KVquant、KIVI、RotateKV等方法。
MixKVQ预填充与解码算法伪代码。
# Algorithm 1: The MixKVQ Prefill & Decoding Algorithm
# parameter: group size G, residual length R, τBF16, τUINT4
procedure Prefill:
Input: X ∈ R^(l_prompt×d)
X_K = X W_K, X_V = X W_V
r = (l_prompt - F) % R
X_K_q = X_K[:l_prompt - r], X_K_r = X_K[l_prompt - r:]
X_V_g = X_V[:l_prompt - r], X_V_r = X_V[l_prompt - r:]
Q(X_V_g) <- GroupQuant(X_V_g, dim=token, numGroup = d//G)
Q(X_K_UINT4), Q(X_K_UINT2), X_K_BF16 <- KeyQuant(X_K_q, τBF16, τUINT4)
KV cache <- Q(X_K_UINT4), Q(X_K_UINT2), X_K_BF16, X_K_r, Q(X_V_g), X_V_r
return X_K, X_V
procedure Decoding:
Input: KV cache, t ∈ R^(1×d)
t_Q = t W_Q, t_K = t W_K, t_V = t W_V
Q(X_K_UINT4), Q(X_K_UINT2), X_K_BF16, X_K_r, Q(X_V_g), X_V_r <- KV cache
X_K_r <- Concat([X_K_r, t_K], dim=token)
X_V_r <- Concat([X_V_r, t_V], dim=token)
if len(X_K_r) == R then
Q(X_K_UINT4_new), Q(X_K_UINT2_new), X_K_new <- KeyQuant(X_K_r)
Q(X_K_UINT4) <- Concat([Q(X_K_UINT4), Q(X_K_UINT4_new)], dim=token)
Q(X_K_UINT2) <- Concat([Q(X_K_UINT2), Q(X_K_UINT2_new)], dim=token)
X_K_BF16 <- Append([X_K_BF16, X_K_new])
X_K_r <- empty tensor
Q(X_V_r) <- GroupQuant(X_V_r, dim=token, numGroup = ...)
Q(X_V_g) <- Concat([Q(X_V_g), Q(X_V_r)], dim=token)
X_V_r <- empty tensor
A <- Concat([t_Q Q(X_K_UINT4) + t_Q Q(X_K_UINT2) + t_Q X_K_BF16, t_Q X_K_r], dim=token)
A_g = Softmax(A)[:-R], A_r = Softmax(A)[-R:]
t_O <- A_g Q(X_V_g) + A_r X_V_r
KV cache <- Q(X_K_UINT4), Q(X_K_UINT2), X_K_BF16, X_K_r, Q(X_V_g), X_V_r
return t_O
function KeyQuant(X_K ∈ R^(l×d), τBF16, τUINT4):
X_K_UINT4, X_K_UINT2, X_K_BF16 <- OutlierReserve(X_K, τBF16, τUINT4)
Q(X_K_UINT4) <- GroupQuant(X_K_UINT4, dim=channel, numGroup = l//G)
Q(X_K_UINT2) <- GroupQuant(X_K_UINT2, dim=channel, numGroup = l//G)
return Q(X_K_UINT4), Q(X_K_UINT2), X_K_BF16
💬 评论讨论
欢迎在这里分享您的想法和见解!