NVIDIA与MIT联手：让超长文本AI助手的"记忆搜索"快了5倍多，他们是怎么做到的？

这项研究由英伟达（NVIDIA）、Thinking Machines Lab、字节跳动Seed部门以及麻省理工学院（MIT）的研究人员联合完成，论文于2026年6月3日发布，编号为arXiv:2606.04511，有兴趣深入了解技术细节的读者可以通过该编号查询完整论文。

当你向AI助手提问一个需要它"回忆"数万字内容的问题时，它其实正在做一件极其耗费资源的事情——从海量的"记忆"中翻找与问题相关的片段。这个过程不仅慢，还极为消耗计算机的存储空间。这项来自英伟达等机构的研究，就是为了解决这一困境而生的。

研究团队提出了一种名为**SparDA**（稀疏解耦注意力机制）的新架构，核心思想可以用图书馆管理员的工作来理解。假设你是一位图书馆管理员，每次读者来查资料，你都得在数万册书中精准找到几本最相关的。传统方式是读者站在原地等你找完，找书和读书这两件事只能排队进行。而SparDA的方案是：当你还在帮这位读者读书的时候，就已经在预测下一位读者可能需要哪些书，提前把书从仓库（CPU内存）搬到取书台（GPU显存）等着。这样一来，等待时间大幅缩短，整个图书馆的运转效率自然就提升了。实验结果表明，SparDA在128K超长上下文场景下，相比原有方案可以实现最高1.7倍的解码速度提升，并通过支持更大批次处理，最终达到高达5.3倍的吞吐量提升。

一、AI的"长期记忆"为什么越来越难处理

现代AI大语言模型正在被要求处理越来越长的对话和文档。无论是帮你分析一份几十页的合同、辅助长篇小说创作，还是在复杂的推理任务中追踪多个线索，这些场景都需要模型保持对数万乃至数十万个词的"记忆"。这份记忆在技术上被称为KV缓存（Key-Value Cache），你可以把它理解为模型的工作记事本，记录着所有之前出现过的信息。

问题在于，这份记事本会随着对话变长而无限膨胀。以目前主流的8B参数模型（约80亿参数）为例，在处理128K长度的上下文时，KV缓存可能占用几十GB的显存。而一张顶级GPU（如H100）也只有80GB显存，光是这份记事本就几乎耗尽了存储空间，根本没有余地同时处理多个用户的请求。

为了解决存储压力，一个直觉上很自然的方案是：把这份大记事本放到电脑的普通内存（CPU内存）里，需要用哪部分再临时调回GPU。这就像把不常用的文件从桌面（GPU）挪到抽屉（CPU内存）里一样。但问题随之而来：GPU和CPU之间的数据传输通道（PCIe总线）相对较慢，每次临时调取都会产生明显的等待延迟，就像每次要用文件都得先开抽屉翻找，效率大打折扣。

与此同时，"稀疏注意力"技术本身也带来了另一个开销。稀疏注意力的核心思路是：并非所有的"记忆"都对当前问题有用，每次只计算最相关的一小部分就好。这可以把原本需要逐一检查所有记忆的操作（复杂度O(T?)，随上下文长度平方增长）大幅压缩。但在"筛选出哪些记忆最相关"这一步骤本身，复杂度仍然是O(T?)——也就是说，随着上下文越来越长，光是做这个筛选就会越来越耗时，最终反而成为新的瓶颈。

这两重矛盾——调取慢和筛选慢——正是SparDA想要同时解决的核心问题。

二、预言未来的"预报员"：SparDA的核心创新

SparDA最关键的发明是在模型每一层中引入了一个全新的"预报"投影（Forecast Projection），与原有的查询（Query）、键（Key）、值（Value）三个投影并列，形成四元组结构。

回到图书馆的比喻：在传统稀疏注意力里，图书馆管理员（注意力机制）需要先根据读者的问题（Query）在书架目录上找到相关书目（稀疏筛选），再去取书阅读（注意力计算），这两步必须顺序完成，任何一步都卡在关键路径上。而SparDA做的事情，是给这位管理员配了一位"预报员"助手（Forecast），专门负责在管理员还在帮当前读者读书时，就提前预测下一位读者会需要哪些书，并悄悄把书从仓库搬出来备好。

更精妙的是，这位"预报员"是经过专门训练的。它不像之前某些方案那样粗糙地用"上一层的隐藏状态"来猜测下一层需要什么，而是通过专门的知识蒸馏训练，学会了精准预测下一层注意力机制真正会用到的记忆块（KV块）。这就是"解耦"二字的精髓：预报和计算彼此独立，互不阻塞，可以同步进行。

在技术层面，这个设计带来了两方面的协同收益。第一，由于预报员不需要像原始注意力那样处理所有的查询头（Query Heads），在分组查询注意力（GQA，即多个查询头共享一对键值头的结构）的设置中，每个KV头只需要一个预报头就够了，而非多个查询头。这意味着筛选步骤的计算量大幅降低，而且还可以跳过原本必须做的softmax归一化操作。第二，由于预报发生在当前层执行期间，预测出下一层需要的记忆块后，可以立刻在一条独立的数据流（CUDA Stream）上异步地把这些数据从CPU内存搬到GPU显存，整个搬运过程与当前层的计算完全重叠，等到下一层真正开始时，数据早已就位。

三、让"搬书"和"读书"同时进行：高效实现机制

仅有好的设计思路还不够，SparDA还在工程实现上做了精心设计，确保异步预取（提前搬运数据）真正能高效运行，而不是反过来拖累主计算。

负责数据传输的是一个持久化的UVA（统一虚拟地址）Triton内核。UVA是一种让GPU可以直接"看到"CPU内存地址的技术，就像给两间办公室打通了一道暗门，不需要通过繁琐的中间手续就能直接取用对方的东西。这个内核的"持久化"特性意味着它会一直保持一小批专属的GPU线程组（CTA，协作线程阵列）处于激活状态，持续处理数据传输任务，而不是每次传输都需要重新启动、重新初始化，从而极大降低了调度开销。

另一个精巧之处在于批次自适应的线程组分配策略。为传输分配越多的线程组，传输越快，但主计算（注意力和前馈网络）能用的线程组就越少，反而会变慢。研究团队发现，这个平衡点会随批次大小变化：批次较小时，GPU本身就没有充分利用，分配少量线程组给传输就够了；批次增大后，传输成为新的瓶颈，值得多分配一些线程组。于是SparDA采用了一个简单的自适应启发式规则：在H100 GPU上，批次小于32时用16个线程组，批次大于等于32时用32个线程组。这个简单的规则在所有批次大小下都能保持在最优配置4%以内的性能，效果相当实用。

此外还有一个边界情况的细心处理：模型的第一层在解码时没有前一层的预报可用，因此第一层的KV缓存始终保留在GPU上，而不进行卸载；第一层会用一个专门的"当层预报"来完成本层的稀疏筛选。这样的处理确保了整个流水线没有死锁的风险。

四、只训练一小部分，却能让整个模型变聪明

SparDA的另一个实用优势在于，它可以被"插入"到已经训练好的稀疏注意力模型中，而无需重新训练整个庞大的模型。在8B参数的基础模型上，新增的预报投影只有3350万个参数，仅占全部参数的0.41%。

训练目标很直接：让预报员学会模仿原始注意力机制的选择行为。具体来说，研究团队用KL散度（一种衡量两个概率分布差异的指标，可以理解为"判断预报员的猜测和真实情况差了多少"）作为损失函数，只优化这些新增的预报投影参数，而原始模型的所有参数保持冻结不变。训练数据使用了ProLong-64K数据集，训练序列长度为65536个词元，整个训练过程在32张H100 GPU上完成，8B模型大约需要48小时。

训练过程中有一个颇为有趣的发现：在构建"标准答案"（目标注意力分数）时，如果使用更细粒度的键压缩窗口（核大小为2、步长为1，相比推理时使用的核大小32、步长16细了16倍），训练出来的预报员更准确。直觉上理解，这就像让学生用高清地图学习导航，但考试时用普通地图——学生在高清地图上学到了更精细的地理知识，这些知识可以迁移到普通地图的使用上，反而比直接用普通地图训练效果更好。当然，训练时的高分辨率分数需要通过最大池化降采样到推理时使用的标准分辨率，保证两者在同一尺度上计算损失。

这种训练方式在两个8B稀疏预训练模型（MiniCPM4.1-8B和NOSA-8B）上均得到验证，分别在48小时和24小时内完成了预报投影的训练。

五、实验检验：速度提升了多少，精度又如何

研究团队在MiniCPM4.1-8B和NOSA-8B两个模型上进行了全面评测，覆盖准确率和效率两个维度。

在准确率方面，SparDA在所有测试套件的综合平均分上均达到或略好于原始稀疏基线。对于MiniCPM4.1-8B，综合平均分从61.4提升到61.7，RULER基准（测量长文本信息检索能力）从78.2提升到78.7，长推理任务从83.6提升到84.7。NOSA-8B的提升更为明显，综合均分从49.4跃升至51.7，推理任务从50.7大幅上升到57.2，RULER从72.2提升到73.9。

作为对比，InfiniGen（一种类似的无训练预取方案）在准确率上出现了明显下滑，原因在于它依赖相邻层隐藏状态的相似性来预测下一层需要什么，而这种相似性在许多实际情况下并不成立，导致预测失准、加载了错误的记忆块。SparDA通过训练学到了跨层的真实对应关系，因此不受这一限制。

在长度泛化能力上，SparDA在RULER基准的各个上下文长度（32K、64K、96K、128K）下均优于稀疏基线，且NOSA-8B上的优势随上下文增长而持续扩大（从32K时的+1.7分扩大到128K时的+4.3分），这说明训练出的预报机制具备良好的泛化能力。

在效率方面，通过逐层分解注意力时间，可以清晰地看到SparDA在预填充（处理用户输入的阶段）和解码（逐词生成回复的阶段）两个阶段的不同优化来源。预填充阶段，块稀疏注意力计算本身是主要开销，块筛选也会随上下文增长而增加，在128K时两者已接近。SparDA的精简预报索引器将块筛选开销降低了最多2.5倍，而块稀疏注意力计算几乎不受影响。解码阶段，由于每次只生成一个词元，注意力计算本身很轻，反而是块筛选成为主要瓶颈——而SparDA的索引器将解码时的筛选开销压缩到几乎不随上下文增长，在128K时节省超过2倍开销。

具体吞吐量数据方面，在H100 GPU、128K上下文下，SparDA的预填充吞吐量比稀疏基线提升了最高1.25倍；解码吞吐量在同等卸载条件下比稀疏基线提升最高1.69倍（NOSA-8B为1.40倍）。最亮眼的数字是与不卸载的稀疏基线相比，由于卸载后SparDA能支持更大的批次（不卸载时16批次就OOM，卸载后可以跑到64批次甚至128批次），最终在128K、64批次时实现了5.28倍的吞吐量提升。

值得一提的是，InfiniGen在吞吐量测试中表现远低于预期，甚至大幅落后于稀疏基线，根本原因在于它需要在CPU端先执行索引操作、找到相关块的位置，再批量传输，CPU端的聚合操作本身成了新的瓶颈；而SparDA在GPU端完成索引计算后直接触发异步传输，完全规避了这一问题。

在A100 GPU上的测试结果趋势与H100一致，预填充最高提升1.23倍，解码最高提升1.55倍，说明这种优化跨GPU架构具有良好的迁移性。

六、研究边界与未来方向

研究团队对SparDA的局限性保持了坦诚的态度。SparDA本质上是一个"插件"，它并不改变底层稀疏注意力的计算方式和稀疏模式本身，因此其精度上限受制于基础稀疏模型的质量——如果基础模型的稀疏注意力本身就不准确，SparDA也只能在这个基础上做优化，无法凭空弥补精度缺口。

目前的实验局限于8B规模的模型。DeepSeek-V3.2、GLM-5等更大规模的前沿模型采用了词元级稀疏注意力（DSA），以及DeepSeek-V4中更复杂的压缩稀疏注意力（CSA）结构，这些场景下SparDA解耦预报的设计思路同样适用——预报员预测词元级重要性分数而非块级分数——但将其适配到数十亿至数千亿参数规模的生产模型，工程挑战更为复杂，研究团队将其列为未来工作的重要方向。

归根结底，SparDA传递了一个更宏观的设计理念：稀疏注意力机制不应只被看作节省计算的工具，还应被设计成能够提前暴露自己的内存访问模式，让系统有足够的时间提前准备好所需的数据。换句话说，让"AI知道自己下一步要看什么"这件事本身，就可以成为一种可学习、可调度的信号，从而让整个推理系统的各个部件协同工作、互相隐藏延迟。这种"预见性调度"的思路，未来或许能推广到更多AI推理优化的场景中。

如果你对这项研究的具体技术细节和完整实验数据感兴趣，可以通过arXiv编号2606.04511找到完整论文。

---

Q&A

Q1：SparDA是专门重新训练一个新模型，还是可以加到现有模型上？

A：SparDA是一个轻量级插件，可以直接添加到已经训练好的稀疏注意力模型上，无需重新训练整个模型。它只训练新增的预报投影参数，在8B模型上这部分参数仅占全部参数的0.41%（约3350万参数），在32张H100 GPU上训练大约需要24到48小时。

Q2：SparDA的5.3倍吞吐量提升是怎么实现的？

A：这个5.3倍主要来自两方面的叠加效果。一方面，通过将KV缓存卸载到CPU内存，GPU可以腾出空间同时处理更多用户请求（更大的批次），从不卸载时最多支持16个并发请求扩展到64个；另一方面，SparDA的异步预取机制让CPU到GPU的数据搬运与当前层计算同步进行，减少等待时间。两者叠加，最终在128K上下文、64批次时相比不卸载的稀疏基线实现了5.28倍的吞吐量提升。

Q3：InfiniGen和SparDA都做了预取，为什么InfiniGen反而更慢？

A：InfiniGen的预取依赖CPU端操作：它先在CPU上计算出哪些KV块最相关，再把这些块批量传到GPU。CPU端的聚合和索引操作本身成了新的瓶颈，导致吞吐量大幅下降，在大批次时甚至比没有预取的稀疏基线还慢。SparDA则完全在GPU端完成索引计算，计算完成后立刻触发异步传输，CPU端只负责被动传输数据，完全规避了这一问题。