清华大学、香港大学等顶尖高校联手破解AI内存瓶颈
这项由清华大学、香港大学、爱丁堡大学、中国科学院大学、香港理工大学以及美团LongCat团队共同完成的研究,以预印本形式发布于2026年5月,论文编号为arXiv:2605.19660,有兴趣深入了解的读者可以通过该编号在arXiv平台查询完整论文。
当你在和手机里的AI助手聊天时,有没有想过,这家伙是怎么"记住"你之前说的话的?答案是靠一种叫做KV缓存(Key-Value Cache)的东西。简单地说,AI每处理你说的一句话,就会在脑子里留下一份"笔记",记录关键信息,以便回应后续问题时有据可查。然而,当你和AI聊的话越来越多,处理的图片、语音越来越复杂,这份"笔记"就会变得越来越厚,最终把AI的"内存空间"塞满——就像一个书桌,越堆越满,最后连动笔的地方都没了。
这个问题并不是小麻烦。对于那些需要处理超长文章、视频、语音的高端AI应用来说,KV缓存占用的内存会迅速占据GPU(图形处理器,也就是AI运算的核心硬件)总容量的大头,严重限制了AI能同时服务多少用户、能处理多长的内容。研究人员早就意识到,压缩这份"笔记"是突破这一瓶颈的关键。
现有的主流做法是"量化",也就是把原本用16位精度存储的数字,压缩成更少位数的近似值——就好像把一张高清照片压缩成低分辨率的缩略图,用更小的空间存储相似的信息。其中,用2位整数(INT2)表示的极限压缩方案,理论上可以把内存占用压缩到原来的八分之一,但也会带来最严重的信息失真。如何在极限压缩的同时保住AI的智力水平,一直是这个领域的难题。
本文的研究团队经过细致分析,发现了一个此前被忽视的根本性问题——他们将其命名为"令牌范数失衡"(Token Norm Imbalance,TNI)。正是这个隐藏的毛病,让现有的极限压缩方案频频翻车。为了解决它,研究团队提出了一套名为OScaR(Omni-Scaled Canalized Rotation,全尺度通道化旋转)的新方案,不依赖复杂的辅助机制,只用两个优雅的步骤就把问题搞定,并在涵盖纯文字、图文混合、语音视频等多种AI模型的测试中,刷新了当前最优成绩。
一、先搞清楚:AI的"笔记本"为什么越来越难压缩
要理解OScaR解决了什么问题,首先得知道AI的"笔记本"长什么样,以及压缩它会遇到什么麻烦。
在AI处理语言(或图像、语音)时,每一层神经网络都会产生三种向量:Query(问询)、Key(键值)和Value(值)。其中,Key和Value会被存入缓存,方便后续复用,这就是KV缓存。Value向量的数值分布比较均匀,压缩起来相对容易。但Key向量就不同了——它里面有一些"异常通道",某几个维度的数值会特别大,像钉子一样突出来,而其他维度则比较平整。
这种现象在AI领域被称为"通道级异常值"(channel-wise outliers)。为了应对它,研究者们发明了"按通道量化"的策略:不是对整个Key矩阵用同一把尺子来压缩,而是给每个通道单独设定一把尺子,这样那几根"钉子"就不会让整体的尺子变得失灵。这种策略在4位甚至8位量化时效果相当好,代表作是一个叫KIVI的框架。
然而,当量化精度压到2位时,这把策略开始失效,而且失效得很厉害。研究团队通过大量实验发现了问题所在:按通道量化有一个隐含的前提假设,那就是同一个通道里,不同位置的"令牌"(可以理解为每一个词、每一帧图片、每一段语音的代表单元)的数值大小应该差不多。只要同一把尺子量下来,大家都能被合理表示。
但现实并非如此。研究团队仔细测量了Llama-2-7B、Llama-3.1-8B、Qwen3-8B等多个主流AI模型在处理文字时产生的Key向量,以及处理图文混合内容的Qwen3-VL模型,发现每一层神经网络里,都有极少数几个特殊令牌,其L2范数(你可以把它理解为"向量的整体能量大小",也就是向量有多强壮)远远低于其他普通令牌。这些"弱小"的令牌,和旁边那些"强壮"的普通令牌被迫共享同一把量化尺子,结果是灾难性的:为了容纳那些强壮令牌,尺子的量程被设得很大,而弱小令牌就在这把粗糙的大尺子上几乎辨认不清了——这就是"令牌范数失衡"的危害。
打个比方,用同一把刻度精度为1厘米的尺子,同时测量一栋200米高的大楼和一枚2厘米厚的硬币。大楼当然没问题,但硬币只能粗略地标记为"2厘米左右"。量化就是这样,共用一把尺子,精细的信息不可避免地丢失了。
研究团队不仅通过实验观测到了这个现象,还从理论上证明了它的危害程度:一个量化块的重建误差(也就是压缩前后信息失真的程度),从根本上取决于块内令牌范数的极差——也就是最强令牌和最弱令牌之间的差距越大,误差越严重。这意味着TNI是按通道量化范式的结构性弱点,而非偶然现象。
更有趣的是,这些弱小的低范数令牌是谁呢?研究团队发现,它们与一种叫做"注意力汇点"(Attention Sink)的现象高度重合。在标准的注意力机制中,softmax函数要求注意力权重加起来必须等于1。当AI处理某些不需要特别关注任何内容的情况时,它没有"我不看任何人"的选项,只能把注意力强行分配给某些无关紧要的令牌,比如句子开头的特殊标记符。这些被迫承接注意力的令牌,就成了注意力汇点。为了让这种"假装注意"的行为不产生实质影响,这些令牌的Value向量(从而也包括Key和Query向量)的范数都被压得很低——它们接收了所有注意力,但几乎不传递任何信息。这些几乎"沉默"的令牌,恰恰是TNI问题的源头。
在多模态AI模型中,TNI的表现更加复杂。除了注意力汇点型的低范数异常令牌,还出现了另外两种情况:不同模态(文字、图片、语音)的令牌之间范数差异巨大,同一模态内部各自平稳,但跨越模态边界时就会出现剧烈的范数跳变;以及与低范数令牌相反,某些令牌的范数极度高于平均水平,形成高范数异常点。这些多样化的TNI模式,让现有方法更加束手无策。
研究团队还通过定量实验精确衡量了TNI的破坏力。以LLaVA视觉语言模型为例,在2位量化下,含有低范数异常令牌的量化块,其误差比去除这些异常令牌后的块高出约35%。更惊人的是,包含多种模态令牌的混合块,比单一模态块的视觉部分量化误差高出约140%。这些数字清楚地说明了TNI并非小问题,而是极限压缩时的主要误差来源。
二、一个优雅的解法:为什么不能直接"缩放",以及OScaR如何两步解决问题
既然问题出在不同令牌之间范数差距太大,最直接的想法是:在量化之前,先把所有令牌的范数都调整到差不多的水平,不就解决了吗?按照这个思路,只需要计算每个令牌的范数,然后除以它,所有令牌的"能量"就变成1了——这叫做"令牌级缩放"。
然而,研究团队发现,单独使用这个方法不仅没有帮助,在很多情况下反而会让量化效果更差。原因出在一个他们命名为"缩放诱导异常值效应"(Scaling-Induced Outlier Artifact)的新问题上。
考虑这样一个场景:正常令牌的向量里,某一两个维度(通道)的数值特别大,其他维度都比较小——这就是前面说的"通道级异常值"。而那些低范数的异常令牌,各个维度的数值都比较小,分布均匀,就像一块平坦的草地。现在,如果把低范数令牌的范数放大到和正常令牌一样,会发生什么?低范数令牌的所有维度都被均匀地放大了,它原本平坦的"草地"变成了各处等高的"丘陵"——但在那些正常令牌异常值通道(也就是少数几个"高峰")之外的维度上,低范数令牌被放大后的数值反而比正常令牌在这些维度上的数值大得多。
这就制造了新的异常值:低范数令牌在被放大之后,成了按通道量化中新的"麻烦制造者",让原本平滑的通道量程被人为地拉宽,量化精度更差了。直接缩放治标不治本,甚至适得其反。
OScaR的妙处在于,它用两步走的组合拳彻底绕开了这个困境。
第一步叫做"通道化旋转"(Canalized Rotation)。这一步用的是一种叫Hadamard变换的数学操作(快速哈达玛变换,FHT)。你可以把它想象成搅拌机:把原来向量里各个通道的信息彻底打散混合,让原本集中在少数几个通道的异常能量均匀分布到所有通道上。经过这个操作,Key向量里那几根突出的"钉子"就不见了,各个通道的数值变得相对均匀。
这一步操作的另一个好处是:它作用于通道维度,不改变令牌之间的范数差异。也就是说,经过Hadamard变换后,强壮令牌依然强壮,弱小令牌依然弱小——TNI问题本身还在。但是,由于通道分布变得均匀了,此前令人担忧的"缩放诱导异常值效应"就消失了:低范数令牌不再有某几个维度特别小、其他维度特别大的不均匀结构,所以把它们统一放大时,不会在任何特定通道上制造新的异常值。
这就为第二步铺平了道路。第二步叫做"全令牌缩放"(Omni-Token Scaling)。在经过Hadamard变换的基础上,计算每个令牌向量的L2范数,然后将每个令牌除以自己的范数,把所有令牌的"能量"统一到同一水平。由于第一步已经消除了通道分布的不均匀性,这里的缩放操作不再会制造新的麻烦,而是安全地解决了TNI问题——所有令牌现在都在同一个量级上,共享量化尺子不再是问题。
为了在解码时恢复原始信息,缩放时使用的范数值会被单独存储起来作为"标签",解码时把存储的低精度向量乘回这个范数,就还原了原始向量的大小关系。对于Query向量,同样需要做Hadamard变换,这样才能保证Query和Key做点积(计算注意力得分的核心操作)时,数学关系保持正确——因为Hadamard变换是一个正交变换,Query和Key同时做相同的旋转,它们之间的相对关系不变。
对于Value向量,由于它没有通道级异常值问题,处理方式更简单:用一个离线的Hadamard变换(也就是在模型运行之前一次性完成,不需要每次推理都重算),把Value向量的权重矩阵预先乘以Hadamard矩阵,之后按令牌量化就好了,不需要在线额外操作,几乎零额外开销。
研究团队通过消融实验(也就是逐步去掉某个组件,看效果如何变化的测试)验证了这两步的缺一不可:单独用Hadamard旋转,通道均匀了,但令牌范数失衡依然存在,量化误差并没有显著降低;单独用令牌缩放,范数平衡了,但缩放诱导异常值效应又把效果抵消了,结果往往比什么都不做还差;只有两步组合使用,才能同时解决两个问题,实现真正有效的极限压缩。
三、工程上的精细打磨:让理论变成真正快起来的代码
提出一个好方法是一回事,把它变成在真实硬件上快速运行的代码又是另一回事。OScaR的研究团队在系统实现层面也下了大功夫,把理论上的优势转化为实实在在的速度提升。
整个OScaR系统由三个精心设计的CUDA核心(CUDA Kernel,也就是在GPU上并行运行的程序)构成,分别承担不同的职责。第一个核心负责在线执行Hadamard变换和令牌缩放——对于Key向量,融合完成Hadamard变换和范数归一化;对于Query向量,只做Hadamard变换。第二个核心高效完成Key和Value向量的量化打包。第三个核心则负责解码时的反量化、反缩放,以及最终的注意力计算。
团队在这里采用了两个关键的工程技巧来压低开销。其一,Hadamard变换的计算复杂度是O(d log d),而传统的矩阵乘法复杂度是O(d?)。对于维度d=128的头部向量,这意味着Hadamard变换的运算量大约是矩阵乘法的十几分之一,速度优势非常明显。而且,团队借助了HadaCore这个专门为Tensor Core(英伟达GPU上的矩阵运算加速单元)优化的Hadamard实现,利用了H??? = H? ? H??的Kronecker分解结构,进一步减少了标量运算压力。
其二,令牌缩放中需要计算平方和的倒数平方根(即1/√Σx??),团队选择使用GPU硬件加速的rsqrt指令,而非先计算L2范数再求倒数——这在数学上等价,但硬件层面速度更快,延迟更低。消融实验证实,rsqrt实现和直接计算L2范数的结果精度几乎相同,却具有更优的硬件效率。
在缓存组织上,系统同时维护一个"打包缓存"(存放2位量化后的Key和Value,以及对应的量化参数和令牌范数标签)和一个"残差缓存"(存放最近128个尚未满足量化块大小的令牌,以全精度FP16存储)。每当残差缓存积累到128个令牌时,就触发一次批量量化,把这批令牌打包进压缩缓存。这种设计既保证了按通道量化在块大小上的需求,又避免了频繁的小批量量化带来的开销。
四、用数字说话:OScaR在各类AI模型上的实际表现
方案再优雅,也要用测试结果来说话。研究团队在三类、共六个具有代表性的开源AI模型上进行了全面测试,覆盖纯文字、图文多模态和全模态(文字+图片+语音)三种场景。
在纯文字AI模型的测试中,团队使用了LongBench-E基准(一个专门测试长文本理解能力的基准,涵盖问答、摘要、代码补全等多类任务)以及"大海捞针"(Needle-in-a-Haystack,NIAH)测试(把一段关键信息藏在长达数万字的文章里,看AI能否找到它)。
在LongBench-E上,OScaR在Llama-3.1-8B模型上取得了41.75%的平均得分,比排名第二的方法(OTT,40.74%)高出1.01个百分点,而16位全精度基准是41.70%——OScaR的2位量化版本居然比16位基准还要略高!在Qwen3-8B上,OScaR得分48.74%,16位基准是49.56%,仅差1.7%,远比其他量化方法更接近基准。
在大海捞针测试中,OScaR的检索准确率达到了96.5%,不仅超越所有量化竞争对手(第二名TurboQuant+是92.7%),甚至略微超过了不做任何压缩的16位基准(96.0%)。这说明OScaR的量化操作不仅没有损害长文本中的信息检索能力,反而可能略微改善了某些情况下的表现。
在图文多模态模型测试中,团队使用了OCRBench(测试AI识别图片中文字的能力)和DocVQA(测试AI理解文档图片并回答问题的能力)。在OCRBench上,OScaR在三个测试模型上均居量化方法首位:LLaVA-v1.6-vicuna-7B得519分(16位基准536分),Qwen3-VL-8B得856分(16位基准858分,差距仅0.2个百分点!),Qwen3-VL-4B得838分(比第二名高2.5个百分点)。在DocVQA上,OScaR同样在所有三个模型上领先,Qwen3-VL-8B甚至以95.01%微超16位基准(94.93%)。
全模态模型的测试针对Qwen3-Omni-30B在MMAU-Pro基准(评估语音理解和指令跟随能力)上进行。OScaR在开放式问答、高质量回答比例和音频指令跟随三项指标上均居量化方法首位,且在所有三项上都超过了16位基准,分别超出1.2、2.0和1.1个百分点。
在效率方面,研究团队在单张H20 GPU(一款高端AI计算专用显卡)上测试了Qwen3-8B模型的推理速度。与16位精度的FlashDecoding-v2基准(目前主流的高效推理实现)相比,在128K超长文本场景下,OScaR的解码速度提升了3.0倍(从92.9毫秒/令牌降至30.9毫秒/令牌)。在批处理场景下(同时处理48个对话),OScaR将内存占用压缩至原来的五分之一(从28.3GB降至5.3GB),推理吞吐量(单位时间处理的令牌数量)提升了4.1倍(从331令牌/秒提升至1354令牌/秒)。
研究团队还与同类竞品TurboQuant+做了详细的延迟对比。在短文本场景(1K令牌)下,TurboQuant+因为底层实现方式(使用llama.cpp框架)的差异,延迟反而比OScaR低(7.8毫秒对25.1毫秒),但这优势随着文本变长急剧消失。在16K令牌时OScaR已经反超,到128K令牌时,TurboQuant+的延迟达到72.9毫秒,OScaR仅需30.9毫秒,差距超过一倍。更重要的是,TurboQuant+使用的是2.5位的混合精度量化(部分通道用3位,其余用2位),而OScaR全程只用2位,存储效率更高,且精度更好。
从理论复杂度分析来看,在d=4096、头维度128、上下文长度10000的典型配置下,OScaR每步解码的理论运算量约为1.23亿单位,虽然高于KIVI(约0.82亿),但远低于TurboQuant+(约2.48亿)。OScaR完全不需要查表操作(TurboQuant+需要大量查表),这对GPU执行效率非常友好,因为GPU的算术流水线远比内存随机访问高效。这种"理论上适度开销换取最高精度"的位置,就是OScaR在精度-效率帕累托前沿上占据的独特位置:比最便宜的KIVI准确得多,比最复杂的TurboQuant+便宜得多,同时精度还更高。
五、研究边界与未来:还有哪些问题有待解决
任何研究都有它的局限性,研究团队也坦诚地指出了OScaR的边界。
由于现代语言模型普遍使用RoPE(旋转位置编码)机制,Key向量在进入注意力计算前会被动态施加位置相关的旋转,这使得Hadamard变换无法像处理Value那样提前"离线"完成——它必须在每次生成新令牌时在线计算。这虽然已经通过HadaCore等工程优化大幅降低了开销,但相比完全没有旋转操作的KIVI,依然存在额外计算成本。未来如果能找到在RoPE存在条件下也能离线融合旋转的方案,整体效率还能进一步提升。
此外,当前的实验主要集中在具有标准LLM(大型语言模型)骨干网络的模型上。研究团队认为,OScaR的框架理论上也适用于其他需要KV缓存的自回归推理任务,比如流式视觉几何模型、视觉自回归模型、带有KV缓存的扩散型语言模型等。但这些模型的架构与标准LLM有较大差异,且在这些领域中KV缓存压缩本身也是新兴课题,相关实验还有待未来展开。
说到底,OScaR做的事情其实非常符合"奥卡姆剃刀"原则——同名论文的名称来源正是如此。这个中世纪哲学原则主张:在解释问题时,不必要地复杂的假设应该被剃掉,最简单能解决问题的方案才是最好的方案。对付KV缓存量化中令牌范数失衡这个核心问题,研究团队没有堆叠复杂的误差校正机制、没有引入额外的查找表、没有依赖高精度残差存储,只是对症下药地用Hadamard旋转消灭通道不均匀性,再用令牌缩放消灭范数失衡,两步走,干净利落。
这对于正在使用或即将使用大型AI模型的普通用户来说,意味着未来他们在手机、平板或云端与AI交互时,即便AI处理的是超长对话、高清图片或长篇音频,系统也能更流畅、更省电、服务更多人——而AI本身的聪明程度几乎不受影响。归根结底,好的压缩不是把AI变笨,而是让AI在更小的空间里依然聪明地工作。
感兴趣的读者可以通过arXiv编号2605.19660找到这篇论文的完整版本,其中包含了所有模型的详细实验数据、理论推导、CUDA实现细节以及丰富的可视化结果,是目前在多模态AI模型KV缓存极限量化这一方向上少有的系统性研究。
Q&A
Q1:KV缓存量化为什么在2位压缩时特别容易失效,其他量化方法有哪些局限?
A:2位量化只有4个可区分的数值级别,容错空间极小,任何数值分布的不均匀都会造成严重失真。现有方法如KIVI在应对通道异常值时表现尚可,但忽视了令牌范数失衡问题,导致不同令牌被迫共享同一粗糙的量化尺子,弱小令牌的信息几乎丢失。TurboQuant等方法虽加入了残差纠错,但引入了查表操作和复杂流水线,在长文本场景下速度反而下降明显。OScaR通过从根本上消除范数失衡,绕开了这些问题。
Q2:OScaR的Hadamard变换会不会破坏注意力机制的正确性?
A:不会。Hadamard变换是一种正交变换,数学上类似旋转,不改变向量之间的内积(点积)关系。注意力得分的计算依赖Query和Key的内积,只要Query和Key同时做相同的Hadamard变换,两者内积的结果与变换前完全一致。OScaR正是这样设计的:Key做Hadamard旋转后量化存储,Query在计算注意力时也先做同样的旋转,数学上等价于完全没有做旋转,正确性得到严格保证。
Q3:令牌范数失衡问题是否只存在于文字类AI模型中?
A:不是。研究团队在图文混合模型和全模态模型上都观察到了令牌范数失衡现象,而且多模态模型中的表现更为复杂。除了与注意力汇点相关的低范数异常令牌,多模态模型还存在跨模态范数跳变(文字令牌和图片令牌的范数水平差异显著)以及某些特殊层中的高范数异常令牌。这些多样化的失衡模式使多模态场景下的量化更为困难,OScaR的全令牌缩放方案能够统一处理这些不同的失衡模式,在多模态测试中同样取得了领先成绩。
