![]()
这项由首尔国立大学计算机科学与工程学院团队完成的研究,以预印本形式发布于2026年5月,论文编号为arXiv:2605.16839。感兴趣的读者可以通过该编号在arXiv平台找到完整论文。
当你打开一本厚厚的小说,想要快速找到某个线索时,你不会从头到尾逐字阅读,而是会翻目录、跳章节,把注意力集中在最关键的段落上。现在的大型语言模型(也就是我们常说的AI大模型,比如ChatGPT背后的技术)面临的挑战,和你读长篇小说时遇到的困难惊人地相似——只不过它们处理的"书"可能长达数十万个词,甚至更长。
随着AI技术的飞速发展,越来越多的应用场景要求大模型能够处理极长的文本:帮你分析几十页的合同、理解一个完整的代码库、对一本厚厚的研究报告做出回应。但问题是,处理这么长的文本需要极其庞大的计算资源,就像要求一个人同时记住一本百科全书的每一个字并随时查找关联,即便是超级计算机也会被压垮。为此,研究人员想出了一种叫做"分块预填充"(chunked prefill)的方法,把长文本切成一段一段地分批处理,就像是流水线作业一样。这种方法被主流的AI服务框架广泛采用,但新的问题随之浮现:每处理一段新文本时,系统都需要"回头看"之前所有处理过的内容,这个"回头看"的过程极其耗时,成为了整个流程的瓶颈。
首尔国立大学的研究团队正是针对这个具体问题提出了他们的解决方案,取名为"CompactAttention"。这项研究的核心目标是:在不牺牲AI回答质量的前提下,让分块处理长文本时的"回头看"过程尽可能地快、尽可能地省力。在实验中,他们在LLaMA-3.1-8B-Instruct这款主流开源大模型上,将处理128K长度文本(大约10万个汉字)时的注意力计算速度提升了2.72倍,同时保持了与原始模型几乎相同的回答准确率。
一、AI读长文本为什么这么累
要理解这项研究解决了什么问题,我们需要先弄清楚大模型是怎么"读"文本的。大模型处理文本的方式,依赖一种叫做"注意力机制"的核心技术。简单来说,当模型处理某个词时,它会同时"扫视"文本中所有其他的词,判断哪些词跟当前这个词最相关,然后把注意力集中在最重要的部分。这个过程就像你在听一段对话时,脑子里会自动把当前听到的内容和之前的关键信息联系起来。
但这种机制有一个严重的代价:文本越长,"扫视"的范围就越大,计算量会以平方的速度增长。处理1000个词需要的计算量,大约是处理100个词的100倍,而不是10倍。当文本长度达到10万个词时,这种计算量已经到了让人咋舌的程度。
为了应对这个问题,"分块预填充"技术把长文本切成固定大小的片段,一块一块地处理。每处理一块新文本,模型不仅要理解这块新内容本身,还要把它与之前所有处理过的内容(存储在所谓的"KV缓存"里,可以把它理解为模型的"工作记忆")进行关联。随着处理的块数增多,这个"工作记忆"越来越大,每次"回头看"的代价也越来越高。
为了加快这个"回头看"的过程,研究人员早就想到了一个聪明的思路:既然并非工作记忆里的所有内容都对当前这块新文本有用,那能不能只看最重要的部分,跳过不重要的部分?这就是所谓的"稀疏注意力"(sparse attention)技术,把"工作记忆"中大部分不相关的内容直接跳过,只精读少数关键部分。
然而,把稀疏注意力技术直接用在分块预填充场景下,却碰到了两个棘手的问题。
第一个问题可以用"游泳池跳水"来类比。现有的稀疏注意力技术,是为一次性处理整段长文本设计的,这时查询的文本(Q)和被查询的历史记忆(KV)长度相近,好比在一个大游泳池里跳水,池子够深够宽,水阻对跳水的影响可以忽略不计。但在分块预填充时,每次只处理一小块文本(比如1024个词),但历史记忆可能已经积累了十几万个词,这就相当于在一个巨大水池旁边的小水桶里跳水——水桶太小,入水的冲击完全发挥不出来。研究团队的实验数据显示,在64K长度的历史记忆和90%的稀疏率下,传统稀疏注意力在分块模式下只能达到理论加速比的30%到45%,远低于一次性处理时的表现。
第二个问题是"重复搜索"的代价。在开始计算之前,稀疏注意力需要先扫描历史记忆,找出哪些部分最重要,这个搜索本身也需要时间。在一次性处理时,只需要搜索一次。但在分块处理时,每处理一个新块就要重新搜索一遍,搜索次数与块的数量成正比。以处理64K长文本、块大小为1024为例,需要重复搜索约64次。研究团队测量发现,像XAttention这类方法的累计搜索时间高达65.4毫秒每层,而搜索本身比实际计算还要耗时。
已有一种方法(QUOKA)试图绕开稀疏注意力的执行低效问题,做法是随机抽取少量查询词,用这些"代表"来评判历史记忆中哪些内容最重要,然后只对选出的部分做密集计算。这种方法就像在图书馆里不是每本书都看,而是让一个小组的人各自代读几本,再汇总反馈。问题在于,如果某本书只对特定的一两个人有用,而他们不在这个小组里,这本书就会被忽视。实验结果也印证了这一点:在需要同时查找多个分散信息的任务上,QUOKA的准确率明显下滑,比如在"多键针中草"(Multi-key NIAH)任务上准确率仅有11.2%,而密集注意力可以达到92%。此外,QUOKA还需要把选出的内容单独复制到一块新的内存区域才能计算,这个复制操作本身随着文本长度和批次大小的增加,开销越来越不可忽视。
二、一份精准的阅卷指南:CompactAttention的核心思路
CompactAttention的思路,可以用图书馆的"索引卡片"系统来理解。传统稀疏注意力方法不仅要生成一份"这些章节值得阅读"的清单,还要按照这份清单去执行一套复杂的"跳跃式翻书"操作,而这套跳跃式翻书本身在小书签数量的情况下效率很低。CompactAttention说:没关系,你只管生成那份清单就好,翻书这件事交给我,我有更好的方式。
具体来说,CompactAttention把整个流程拆成两个清晰的阶段:**选择**和**执行**。
在选择阶段,它可以复用任何现有的轻量级稀疏注意力方法来产生一份"哪些历史记忆块重要"的粗略清单。这份清单是一个二维的表格,横轴代表当前这块文本里的每个小查询块,纵轴代表历史记忆里的每个存储块,表格里的每个格子标记"0"或"1",表示这对组合是否值得关注。
接下来,CompactAttention做了一个巧妙的转换,叫做"块联合"(Block-Union)。这个过程分两步走。
第一步叫"Q块联合"(Q-block union)。当前这块文本里可能有多个查询小块,每个查询小块都有自己认为重要的历史记忆块。CompactAttention把所有查询小块的选择结果取"并集"——只要任何一个查询小块认为某个历史记忆块重要,这个块就进入最终选择名单。这就像一个班级里有四个组,每个组各自觉得某些参考书有用,最后老师把所有组认为有用的书都放到公共书架上,确保每个组的需求都被满足。
第二步叫"组内联合"(intra-group union)。现代大模型通常采用一种叫"分组查询注意力"(GQA)的技术,把多个查询头分成若干组,每组共享同一套历史记忆。CompactAttention因此也对每组内的所有查询头的选择结果再次取并集,得到每个组唯一的一份历史记忆块清单。对于GQA分组比例很大的模型(比如每8个查询头共享一套历史记忆),CompactAttention会把大组进一步细分为4人一组的小组,避免联合操作把太多不相关的块纳入进来,导致精简效果打折扣。
这两步操作的数学表达相当简洁。假设用M(b,h,i,j)表示批次b、查询头h、查询块i、历史记忆块j的重要性(0或1),Q块联合就是对所有i取或运算,得到M(b,h,j);组内联合再对组H(g)内所有查询头h取或运算,得到G(b,g,j);最终每个组的历史记忆块清单T(b,g)就是所有G(b,g,j)=1的j的集合。这个清单有一个重要的性质:在执行约束下,它是保留所有被选中块的最小清单,没有任何冗余。
这样做的确会比原始稀疏清单保留更多的块(因为取了并集),但研究团队发现,通过把初始的稀疏清单设得更"严苛"一些(也就是更激进地删减不重要的块),经过两步联合之后的最终结果,仍然可以和直接使用原始清单时的稀疏度相当。实验数据表明,FlashPrefill以α=0.01参数生成的初始清单稀疏度为69.8%,而CompactAttention以α=0.06生成的初始清单稀疏度虽然高达89.8%,经过两步联合后降为70.2%,与FlashPrefill旗鼓相当,同时准确率保持一致。
三、"零复制"翻书术:如何让执行同样高效
生成了精准的历史记忆块清单之后,如何高效执行才是决定最终速度的关键。这里面有一个深藏在系统底层的设计问题,就像图书馆书架的摆放方式决定了找书的效率。
传统的AI内存布局把KV缓存(历史记忆)按照"序列优先"的方式存储,就好比图书馆里的书是按照"一套丛书的所有卷次放在一起"的方式摆放的,每套丛书对应一个批次里的一段序列,每个"块"里同时包含了所有KV头的内容。这种摆法有个问题:当不同的查询组需要查阅不同的书块时,系统只能为所有组提供同一张书架位置清单,没办法让不同组各自只取自己需要的书块。
CompactAttention改变了这个摆放方式,采用"KV头优先"的布局,把KV缓存按照"同一个KV头的所有序列块放在一起"的顺序存储。这样一来,每个(批次, KV头, 块)的组合就对应了内存中一块连续的区域,可以像图书馆里的独立页码一样被单独寻址。不同的执行组就可以各自持有一份独立的"页码清单",指向原始KV缓存中的不同位置,无需把数据搬来搬去。
这就是所谓的"零复制"(zero-copy)执行。CompactAttention把每个组的历史记忆块清单转换成两个轻量级的元数据数组(kv_indptr和kv_indices),直接传给底层的分页注意力计算引擎(使用FlashInfer 0.6.9框架),引擎按照这份元数据在原始内存位置上直接计算,不需要把选中的块复制到新的内存区域。这个转换由一个高效的CUDA并行计算程序完成,速度极快。
相比之下,如果用"复制再计算"的方式(CompactAttention-FP Copy变体),需要把选中的块物理搬移到一块新的连续内存,然后再调用密集注意力计算。研究团队的消融实验(即控制变量实验)专门对比了三种执行策略:直接用块稀疏核执行(传统方法)、复制选中块再密集计算、以及CompactAttention的零复制分页执行。在相同的块选择结果下,三种策略的每层每块的延迟分别是15.64毫秒、6.54毫秒和5.03毫秒,零复制分页执行是最快的,即便它本身也有构建元数据的少量开销(约0.13到0.47毫秒,随批次大小增长极缓慢),仍明显优于前两者。
此外,当前正在处理的这块新文本本身始终保持全量计算,不做任何稀疏跳过。原因是:如果对当前块也稀疏处理,因果注意力的掩码(控制每个词只能看到它之前的词)就需要在"压缩后的位置"上施加,而不是原始位置,这会导致语义错乱,破坏模型的基本推理能力。
四、实验数据说话:快了多少,准了多少
研究团队在两款主流开源大模型上进行了全面测试,分别是Meta公司的LLaMA-3.1-8B-Instruct(支持128K上下文)和阿里云的Qwen3-30B-A3B-Instruct-2507(一款混合专家架构模型,支持256K上下文)。测试硬件覆盖两种主流GPU:96GB显存的RTX PRO 6000(基于Blackwell架构)和141GB显存的H200 SXM(基于Hopper架构)。准确率评测使用了长上下文基准测试RULER和LongBench V2。
速度方面,CompactAttention-FP(使用FlashPrefill作为初始块选择方法)在H200上、批次大小为8、块大小为1024的配置下,处理128K文本时注意力计算速度达到密集注意力的2.72倍,端到端(包括所有非注意力计算)也达到了1.96倍加速。在RTX PRO 6000上、批次大小为4、块大小为512的配置下,128K文本的注意力速度达到密集注意力的2.74倍,端到端加速1.75倍。具体延迟数值来看,在H200上,密集注意力处理128K文本总计需要55678.4毫秒,而CompactAttention-FP只需20489.7毫秒;端到端分别是71015.1毫秒和36265.4毫秒。
作为对比,QUOKA在128K场景下只能实现约1.09倍的注意力加速,因为它的复制开销抵消了减少计算带来的收益;XAttention和SeerAttention甚至比密集注意力更慢,因为它们的搜索开销和稀疏核执行低效叠加后超过了节省的计算量;FlashPrefill作为最强的块稀疏基线方法,在H200上128K时能达到2.24倍注意力加速,但仍低于CompactAttention-FP的2.72倍。
两个版本的CompactAttention都随着文本长度增长而加速效果越来越显著:在32K时,两者的速度提升还不够明显,甚至偶有负优化(因为搜索和元数据开销尚未被分摊),但到了128K时,优势已经相当可观。这符合设计预期:只有历史记忆足够大,精选少数块的价值才能充分体现。
块大小也对加速比有微妙影响。当块大小从512增至2048时,CompactAttention-FP的加速比从2.85倍降至2.38倍,因为更大的块意味着Q块联合时每块内部的查询小块更多,并集的范围变大,最终保留的历史记忆块比例也随之提高,有效稀疏度下降。
在准确率方面,RULER基准测试覆盖多种长文本理解任务,CompactAttention在LLaMA-3.1-8B-Instruct上的128K准确率分别为:CompactAttention-SA(使用SeerAttention作为初始选择方法)74.28%,CompactAttention-FP 74.17%,而密集注意力基线是76.59%。二者的差距在2个百分点左右,而QUOKA的差距达到了6.15个百分点(70.44% vs 76.59%)。在Qwen3-30B-A3B-Instruct-2507上,CompactAttention-FP的128K准确率为87.03%,密集注意力为87.81%,差距约0.78个百分点,略优于FlashPrefill(86.77%)。
LongBench V2是一个更具挑战性的综合推理基准,CompactAttention-FP的总体准确率与密集注意力完全持平(均为30.4%),并在"长文本"子任务上略有超出(31.5% vs 30.6%)。QUOKA则下降到26.8%,尤其在"难题"类别中仅有25.4%(密集注意力为29.3%)。这进一步印证了覆盖所有查询块的重要性:需要同时检索多处信息的复杂任务,对"不遗漏任何查询块"的要求更高。
在Qwen3这款GQA比例为8:1的大模型上,子KV组联合(每4个查询头一组)相比全组联合,在128K文本下能将有效稀疏度从59.88%提升至70.72%,对最终加速效果有实质性贡献。CompactAttention-FP在Qwen3上的256K文本下达到1.64倍注意力加速,而QUOKA只有1.15倍,FlashPrefill甚至只有0.95倍(比密集注意力还慢)。
五、这项研究对普通人意味着什么
CompactAttention本质上是一项基础设施层面的优化,不是一个普通用户能直接在手机上点击使用的功能。但它对普通人的间接影响却相当直接。
当你在使用AI助手帮你分析一份长达数十页的报告、或者让AI记住你之前几十次对话的内容、又或者要求AI通读一个复杂代码库时,所有这些功能的背后都离不开高效的长文本处理能力。目前,这类功能要么响应速度很慢,要么需要巨大的计算资源(意味着更高的服务成本),要么准确率不够理想。
CompactAttention提供了一条在速度、成本、准确率三者之间找到更好平衡点的具体路径。从速度角度,同样的硬件可以更快地给出回应;从成本角度,处理相同长度文本所需的GPU时间减少了,服务提供商的算力账单可以降低;从准确率角度,相比之前同类方法,它更好地保留了对长文本的理解质量。
说到底,AI大模型能否处理真正意义上的"超长"文本,不仅取决于模型的参数规模,更取决于推理时的计算效率。CompactAttention给出的答案是:与其花大力气优化"选哪些内容看",不如同时认真思考"选好之后怎么看",后者的工程细节同样决定了最终能跑多快。
这项研究同时也揭示了一个更普遍的道理:在AI工程中,选择和执行的解耦,往往能打开意想不到的优化空间。未来,随着更快的块选择方法、更高效的分页注意力核、以及更多硬件平台的涌现,CompactAttention框架有望进一步受益,因为选择和执行这两个模块可以独立升级,互不干扰。
如果你对这项研究的技术细节感兴趣,或者希望了解更多关于分块预填充和稀疏注意力领域的进展,可以在arXiv平台搜索论文编号2605.16839获取完整原文和代码仓库地址。
Q&A
Q1:CompactAttention和普通的稀疏注意力方法有什么根本区别?
A:普通稀疏注意力方法既负责"选哪些历史内容重要",也负责"按稀疏方式执行计算",在分块预填充场景下这套执行方式效率很低。CompactAttention则把这两件事分开:它复用现有方法来"选",但把选出的结果转换成分页内存元数据,交给高效的密集注意力引擎来"执行",彻底避开了稀疏执行核在短查询长度下的低效问题。
Q2:分块预填充处理长文本和一次性处理有什么区别,为什么前者更难优化?
A:一次性处理时查询长度和历史记忆长度相近,稀疏计算可以充分利用并行度。分块预填充时每次只处理一小块新文本(比如1024个词),但历史记忆可能有几十万词,查询极短而历史极长,稀疏核没有足够的查询块来分摊固定的寻址开销,实际加速远低于理论值,而且每块都要重新搜索一遍最重要的历史内容,搜索开销会累积叠加。
Q3:CompactAttention会不会漏掉某些重要的历史内容导致AI回答出错?
A:CompactAttention通过"块联合"操作,保证了原始块选择清单里每个被认为重要的历史块,都会出现在最终执行的清单中,不会遗漏。它的精度损失来源于初始块选择方法本身可能存在的遗漏,而不是块联合操作引入的,所以CompactAttention的准确率与对应块稀疏基线方法非常接近。





京公网安备 11011402013531号