北京大学造出“变形金刚“AI芯片适配器

这项由北京大学人工智能研究院主导的研究以预印本形式发布于2026年5月论文编号为arXiv:2605.15250有兴趣深入了解的读者可通过该编号查询完整论文。当你把一台高端游戏电脑的显卡换成一张性能阉割版显卡时你会希望游戏能自动调整画质设置而不是强行以原来的高画质运行导致画面卡成幻灯片。这篇论文解决的正是大型语言模型也就是ChatGPT、文心一言这类AI背后的核心技术在不同计算芯片上运行时面临的同一个困境——现有的高效方案只为一种芯片量身定做换了芯片就摔跟头。研究团队提出了一种叫做分组查询潜在注意力GQLA的新设计核心思路极其简洁同一套训练好的模型参数在运行时能自动切换两种完全等价的工作模式让H100这类顶级芯片和H20这类受出口管制的平价芯片都能各自跑在自己的最优状态。没有重新训练没有定制内核一套权重走遍天下。一、先搞清楚KV缓存这个拦路虎要理解这篇论文在做什么先要明白AI大模型在生成文字时遭遇了怎样的瓶颈。每当AI生成一个新词它都需要回顾自己之前生成的所有内容作为参考。这些历史记录就叫做KV缓存Key-Value Cache其中K代表索引信息V代表内容信息。可以把KV缓存理解为一位速记员随时翻阅的工作笔记本。每生成一个新词速记员就要把整本笔记本翻一遍找出相关内容。笔记本越厚翻阅速度越慢。当AI生成长文本时这本笔记本会变得非常庞大每次翻阅都需要从芯片的高速内存HBM类似硬盘但比普通内存快得多里把数据搬运一次而这个搬运速度就成了整个系统最致命的瓶颈。学界为此发明了多种压缩笔记本的方案。最激进的叫MQA多查询注意力相当于所有人共用一本极简笔记本折中方案叫GQA分组查询注意力几个人共用一本笔记本而DeepSeek-V2/V3采用的MLA多头潜在注意力则更聪明——它不只是共用笔记本而是把笔记本内容先压缩成一段精华摘要存起来用的时候再临时展开。这样笔记本体积最小搬运最省力。MLA的精华摘要方案在英伟达H100芯片上几乎达到了理论最优效率。但问题在于这个方案是为H100量身裁缝的换了芯片就出了岔子。二、H100和H20同样的带宽天差地别的算力要理解芯片适配问题需要引入一个叫屋顶线模型Roofline Model的评估框架。这个框架把芯片性能比作一条双斜屋顶屋顶的左半边代表内存搬运速度是瓶颈右半边代表计算速度是瓶颈两边相交的最高点叫做屋脊点就是计算任务应该尽量命中的最优工作区域。H100芯片的屋脊点大约在每字节295次浮点运算。换句话说H100的计算能力极强只要一个任务的计算量与数据搬运量之比接近295就能把芯片利用率逼到极限。MLA的absorbed MQA模式即精华摘要模式在单次解码时的这个比值恰好约为242稍低于屋脊点正处于内存搬运略微吃紧的高效区间堪称完美匹配。H20则是另一回事。由于美国出口管制H20的计算能力被大幅削减至约148 TFLOPS但内存带宽几乎保留完整约4.0 TB/s导致H20的屋脊点只有大约37。而MLA在H20上的比值仍然是242——这个数字远远高于37的屋脊点意味着H20的计算单元被迫全速空转等待数据大量算力白白浪费。更直白地说MLA在H20上是严重的大马拉小车芯片的运算部件一直在等数据却等不到吞吐量大幅下降。论文中的数据清楚地展示了这个差距MLA在H20上每步解码需要约15.42微秒即便开启了多词预测MTP一次多生成几个词以提升效率的技术由于H20已经是计算瓶颈多生成一个词只会让等待时间成倍增加吞吐量几乎没有任何提升每秒只能生成约6.5万个词。与此同时H100运行同样的MLA却只需2.82微秒每秒可以生成约35.4万个词。除了效率问题MLA还有另外两个连带缺陷。其一是并行扩展受限MLA的精华摘要模式要求所有计算头共享同一份压缩数据这使得无法沿头轴方向把计算分散到多张显卡上这种分散技术叫做张量并行实际部署中需要把压缩数据复制到每张显卡浪费内存。其二是多词预测失效如前所述在H20这类计算受限的芯片上MLA一词一词地生成和一次生成多词的速度几乎相同多词预测完全失去意义。三、GQLA的核心设计一套参数两条路走GQLA的关键发明在于它在保留MLA精华摘要压缩能力的同时为同一套训练参数开辟了第二条执行路径。具体来说GQLA把键值的上投影矩阵可以理解为把摘要还原为完整内容的解码器不再复制给所有128个查询头而是按组索引每组8个查询头共享一个解码器。这个看似微小的改变让整个系统在数学上同时支持两种完全等价的计算方式。第一条路叫GQA路径临时把摘要展开成8组完整的键值对每组对应16个查询头然后按照普通GQA的方式运行注意力计算。这条路需要在缓存中存储展开后的数据每个词元需要约4224字节相当于LLaMA-3的标准GQA缓存大小但能让计算与数据搬运的比值降低到约38.8恰好命中H20的屋脊点。在H20上这条路配合多词预测每步解码只需9.06微秒每秒可生成约22.1万个词比MLA的6.5万提升了3.4倍。第二条路叫MQA吸收路径与MLA完全相同把解码器矩阵直接吸收进查询计算里让所有查询头直接对精华摘要做注意力缓存中只需存储摘要本身每个词元约1152字节计算比值约242完美贴合H100屋脊点。这条路在H100上单步解码只需2.82微秒每秒生成约35.4万个词。两条路产生的输出在数学上完全相同只是计算的中间步骤不同。部署时系统根据目标硬件一次性选定路径并相应地把KV缓存压缩或展开这只是部署时的一次性操作而非每步推理都要做此后无需任何改动。四、把现有模型变身GQLATransGQLA流水线从头训练一个GQLA模型代价高昂研究团队因此提出了TransGQLA——一种把已有GQA模型改造成GQLA模型的转换流程核心改动只有一行代码的差异。这个流程建立在此前提出的TransMLA工作基础上。TransMLA的做法是把GQA模型里每组共享的KV头合并成一个大的潜在向量然后把展开矩阵复制给所有查询头使模型在非压缩模式下行为像MHA完整多头注意力。TransGQLA的区别仅在于合并时不复制展开矩阵让它保持按组索引——于是合并后的模型行为依然是GQA而不是MHA。这个细节保留了沿组轴进行张量并行的能力最多支持8路零冗余张量并行。完成这一步之后TransGQLA沿用TransMLA的后续流程通过RoRoPE技术把位置编码信息从内容编码中解耦出来通过FreqFold技术对旋转位置频率进行重新排列以利于后续压缩最后通过对键和值进行归一化平衡后做联合低秩压缩把庞大的中间向量压缩到一个紧凑的潜在表示。这些步骤操作的对象是合并后的潜在向量完全不关心外层模型是GQA还是MHA解释因此可以直接复用。研究团队在LLaMA-3-8B上验证了这个流程。这个模型有32个查询头、8个KV组每个头的维度为128原始GQA缓存每词元每层需要2048个BF16元素。经过TransGQLA处理并压缩到576维潜在空间后MQA吸收路径的缓存降至原来的28.125%压缩率超过了7倍。GQA路径的缓存大小则与原版相近保留了原始GQA的传输效率。五、稀疏注意力的扩展为什么GQLA更有优势论文还讨论了一个进阶话题稀疏注意力。正常的注意力机制每生成一个新词都要看遍所有历史词当上下文很长比如几万个词时开销巨大。稀疏注意力的思路是每次只挑选最相关的一小批历史词来参考跳过无关的大部分历史从而大幅降低计算量。DeepSeek提出的DSA深度求索稀疏注意力是目前主流的稀疏方案但它与MLA组合时存在结构性缺陷。现代GPU芯片的矩阵乘法加速单元张量核心有一个硬性限制它的基本计算瓦片MMA瓦片的M维度至少需要16意味着每个KV头必须同时服务至少16个查询头才能让张量核心满负荷工作。MLA在非压缩模式下每个KV头只对应1个查询头远低于16的门槛张量核心完全发挥不出来所以稀疏MLA被迫在所有设备上都只能走MQA吸收路径无法利用GQA路径的优势。GQLA的标准配置是128个查询头配8个KV组每组恰好对应128/816个查询头与MMA瓦片的M16完美匹配。这意味着稀疏GQLA可以在GQA路径上让张量核心满负荷工作在H20这类计算受限的芯片上保持高效同时保留沿组轴的张量并行能力。记忆力受限的硬件可以切换到稀疏MQA吸收路径以最小化缓存搬运计算受限的硬件则留在稀疏GQA路径保持低算力消耗两种选择都不需要定制内核。论文还提到当上下文极长超过6.4万词时计算哪些历史词最相关本身也会变成瓶颈。另一项叫HISA的工作提出了一种分层评分方法来加速这个筛选过程与GQLA自然互补——HISA负责更快地找出相关词GQLA负责更高效地对相关词做注意力计算两者组合能从两端同时逼近硬件理论峰值。六、实验验证转换后掉了多少能力能恢复吗研究团队选取了开源的LLaMA-3-8B作为测试对象在六个常识推理基准测试上评估了转换效果包括MMLU综合知识、ARC科学推理、PIQA物理常识、HellaSwag情境理解、OpenBookQA开放知识和Winogrande代词消歧。原始LLaMA-3-8B在这六项测试的加权平均分为63.84分。经过TransGQLA转换但不做任何额外训练后平均分降至54.13分下降了约9.7分。PIQA和HellaSwag的得分与原始模型相差只有几分表明模型的基础语言理解能力保留较好主要损失集中在需要大量知识记忆的MMLU等任务上。这个损失程度对于超过7倍的缓存压缩率而言是相对温和的。研究团队还给出了预期恢复轨迹的参照。此前TransMLA的实验显示在同样的LLaMA-3-8B基础上做同样压缩后用约300亿词元的数据继续预训练平均分就能恢复到63.39分与原始模型只差0.5分——而300亿词元相比LLaMA-3本身训练用的15万亿词元仅仅是约五百分之一的训练量。由于TransGQLA和TransMLA在0训练词元时的分数完全相同两者的KV子空间压缩步骤完全一致只有前期的头合并方式不同研究团队预计TransGQLA的继续预训练曲线将与TransMLA高度一致因此可以合理期待同样数量级的训练后即可恢复到接近原始水平。这个继续预训练实验正在进行中完整结果将在后续版本中补充。七、选择(g8, sq2)还是(g4, sq1)两种H20优化方案的权衡论文详细讨论了两种在H20上达到最优效率的参数配置各有侧重。g8配合多词预测sq2的方案是研究团队的推荐默认配置。这种配置的优点是多方面的8个KV组意味着潜在空间维度为8×1281024比rkv512的压缩目标大一倍给PCA压缩留出了2倍冗余压缩质量更高支持最多8路零冗余张量并行大规模部署时扩展性好每组16个查询头恰好匹配MMA瓦片要求为后续的稀疏扩展打好基础配合多词预测H20每秒吞吐量可达22.1万词。g4不使用多词预测sq1的方案则是一个轻量替代选项。每词元缓存大小从4224字节降至2176字节只比MQA吸收路径的1152字节多一些缓存压力显著减小不需要多词预测头模型结构更简单但PCA冗余度降为1倍PCA是一种数据降维技术冗余度越高压缩时信息保留越完整张量并行上限也降至4路。在H20上同样能达到22.1万词/秒的吞吐量。关键的一点是无论选择哪种GQA路径配置MQA吸收路径的计算强度公式中根本不含g这个变量因此两种配置在H100上都以完全相同的2.82微秒每步在MQA吸收路径上运行不受影响。这正是GQLA设计的精髓所在——GQA路径的参数选择完全不会牵连H100上的表现两套硬件的优化方案可以独立调整。说到底GQLA做的事情听起来有点像魔术但背后的道理其实很朴素同一批数学参数根据不同的计算顺序可以既表现得像把笔记精简成摘要再直接用也表现得像按小组展开笔记分组处理两种算法算出来的答案一模一样只是计算过程中的数据搬运量和运算次数的比例不同。哪种比例更适合当前芯片的脾气就用哪种。这种灵活性以前从来没有人在这个问题上明确提出过而研究团队用一个极小的结构改动就实现了。归根结底这项工作的价值在于打破了一个隐含的假设——为顶级芯片优化的模型架构在平价芯片上只能凑合用。通过让同一套模型参数在数学意义上等价地适配不同硬件特性GQLA为大模型部署的硬件无关性提供了一条实际可行的路径。对于那些既需要在高端数据中心提供服务、又希望在不同算力级别的硬件上灵活扩展的团队来说这个思路值得认真参考。你可能想知道这种设计对未来的模型训练有什么影响。一个有趣的推论是如果模型从一开始就按照GQLA的结构训练那么部署时就可以完全不做任何转换直接根据目标硬件选择路径。这意味着未来DeepSeek这类机构在设计下一代模型时可以考虑把GQLA作为MLA的直接替代从训练阶段就埋下两路兼容的基因。而对于已有的大量GQA模型比如LLaMA系列TransGQLA流程提供了一条相对低成本的改造路径不必从零开始。当然论文本身也坦诚地指出了几处待完善的地方。屋顶线分析是一种理论工具现实中的芯片还有缓存层次结构、指令流水线延迟、多流调度等复杂因素实际的内核级基准测试还需要在真实H20和H100硬件上完成。TransGQLA的继续预训练结果也尚未完成目前的恢复预测是基于TransMLA的经验推断。此外所有实验都集中在LLaMA-3-8B这一个模型上在更大规模如70B参数和更多任务类型长文本检索、代码生成、数学推理等上的验证还需要后续工作来补充。如果你对大模型推理优化或硬件适配感兴趣可以通过arXiv:2605.15250找到这篇论文的完整版本里面包含详细的数学推导、完整的实验数据以及符号表。QAQ1GQLA和MLA有什么区别为什么GQLA能适配H20AMLA只有一条精华摘要执行路径这条路径的计算与数据搬运比值约为242适合H100但远超H20屋脊点37导致H20的计算单元空转浪费。GQLA在保留这条路径的同时新增了一条GQA路径通过按组展开数据让比值降至约38恰好命中H20屋脊点H20吞吐量因此从6.5万词/秒提升至22.1万词/秒提升了3.4倍。两条路径使用完全相同的训练参数输出也数学等价。Q2TransGQLA转换后模型能力会损失多少需要多少额外训练来弥补A以LLaMA-3-8B为例经过TransGQLA转换但不做任何额外训练六项常识推理基准的平均分从63.84降至54.13下降约9.7分同时KV缓存压缩至原来的28.125%。根据同类方法TransMLA的经验用约300亿词元继续预训练后平均分可恢复至63.39与原始模型仅差0.5分。300亿词元相当于LLaMA-3原始训练量的约五百分之一训练成本相对较低。Q3GQLA的两条路径切换时需要改变模型参数或重新训练吗A不需要。两条路径共享完全相同的训练参数切换时只需在部署阶段对KV缓存做一次性的压缩或展开操作之后运行时无需任何改动。MQA吸收路径复用MLA的现有内核GQA路径复用标准GQA内核均不需要定制开发新的计算内核。