北京大学研发AI芯片适配器实现硬件灵活配置

首页

科技数码

热心网友

转载

2026-05-24

最近，一项由北京大学人工智能研究院主导的研究，在预印本平台arXiv上发布了编号为2605.15250的论文，探讨了一个大模型部署中非常实际的问题：如何让同一个模型，在不同算力的芯片上都能高效运行。

这就像你给一台高端游戏电脑换上了一张性能稍弱的显卡，你肯定希望游戏能自动调整画质，而不是强行用原来的高特效导致卡顿。目前主流的大语言模型（比如ChatGPT、文心一言背后的技术）就面临类似的困境——许多为顶级芯片（如英伟达H100）优化的高效方案，一旦换到受出口管制、算力受限的芯片（如H20）上，性能就会大打折扣。

研究团队提出的解决方案，叫做“分组查询潜在注意力”（Grouped-Query Latent Attention, GQLA）。它的核心思路非常巧妙：让同一套训练好的模型参数，在推理时能自动切换成两种数学上完全等价，但计算模式不同的“形态”。这样一来，H100可以跑在为其优化的“高速模式”，H20也能切换到为其定制的“节能模式”，无需重新训练或定制内核，真正实现“一套权重，走遍天下”。

一、先搞清楚“KV缓存”这个拦路虎

要理解GQLA的价值，得先明白大模型生成文本时的主要瓶颈在哪。模型每生成一个新词（token），都需要回顾之前生成的所有历史内容作为参考。这些历史信息，在技术术语里被称为“KV缓存”（Key-Value Cache）。

你可以把KV缓存想象成一本不断变厚的会议记录。每写一个新句子，作者都得把整本记录从头到尾翻一遍，找出相关的上下文。记录本越厚，翻阅的速度就越慢。在AI推理中，这个“翻阅”过程，本质上是把数据从芯片的高速内存（HBM）里搬出来进行计算。当生成长文本时，数据搬运的速度就成了拖慢整个系统的关键瓶颈。

为了压缩这本“记录本”，学术界想了不少办法。最激进的是MQA（多查询注意力），相当于所有人共用一本极简笔记；折中的是GQA（分组查询注意力），几个人共享一本；而像DeepSeek-V2/V3采用的MLA（多头潜在注意力）则更聪明——它先把冗长的记录压缩成一份“精华摘要”存起来，用的时候再临时展开，这样需要搬运的数据量最小。

MLA的“摘要”模式在H100这类计算能力极强的芯片上效率接近理论极限。但问题在于，这套方案是专门为H100这类芯片“量体裁衣”的，一旦换到H20上，就水土不服了。

二、H100和H20：同样的带宽，天差地别的算力

这里需要引入一个评估芯片性能的经典框架：“屋顶线模型”（Roofline Model）。它把芯片性能比作一个双斜面的屋顶，左边斜面代表性能受限于“数据搬运速度”，右边斜面代表受限于“计算速度”。两条斜线的交点，即“屋脊点”，是芯片能达到最优利用率的工作区域。

H100的屋脊点大约在每字节数据能进行295次浮点运算。这意味着H100计算能力极强，只要任务的计算/搬运比接近295，就能把芯片“喂饱”。巧的是，MLA在单步解码时的这个比值约为242，虽然略低于屋脊点，但正好处于内存搬运略微紧张的高效区间，匹配得相当完美。

H20的情况就截然不同了。由于出口管制，其计算能力被大幅削弱，但内存带宽保留得比较完整。这导致H20的屋脊点急剧下降到大约37。而MLA的比值依然是242——这个数字远高于37，意味着什么呢？意味着H20的计算单元一直在全速空转，焦急地等待数据从内存搬过来，大量算力被白白浪费。用个形象的比喻，就是“大马拉小车”，车跑不快不是因为马没力气，而是路太窄，货送不过来。

论文数据清晰地展示了这种差距：在H20上运行MLA，每生成一个词需要约15.42微秒，即使采用“一次多预测几个词”（MTP）的技术来提升效率，由于瓶颈在计算本身，吞吐量也几乎无法提升，每秒仅能生成约6.5万个词。相比之下，H100运行同样的MLA只需2.82微秒，吞吐量高达每秒35.4万个词。

除了效率问题，MLA在适配性上还有两个连带缺陷：一是“并行扩展受限”，其摘要模式难以有效地将计算分散到多张显卡上；二是在H20这类芯片上，“多词预测”技术几乎失效，失去了其加速意义。

三、GQLA的核心设计：一套参数，两条路走

GQLA的巧妙之处，在于它在保留MLA数据压缩能力的同时，为同一套模型参数开辟了第二条独立的执行路径。

具体来说，GQLA对模型结构做了一个微调：它让“键值”的解码器不再被所有查询头共享，而是改为按组索引，每8个查询头共享一个。正是这个看似微小的改变，使得整个系统在数学上同时支持两种等价的计算模式。

第一条路：GQA路径。 这条路会把压缩的“摘要”临时展开成8组完整的键值对，然后按标准的GQA方式进行计算。这样做需要在缓存中存储展开后的数据，虽然数据量稍大，但能将计算/搬运比降低到约38.8，恰好命中H20的屋脊点。实测在H20上，配合多词预测技术，走这条路径每秒可生成约22.1万个词，性能比MLA提升了3.4倍。

第二条路：MQA吸收路径。 这条路与MLA完全一样，所有查询头直接对压缩后的“摘要”进行操作，缓存数据量最小，计算/搬运比约为242，完美贴合H100的屋脊点。在H100上，其性能与MLA持平，每秒生成约35.4万个词。

最关键的是，这两条路径使用完全相同的模型参数，并且最终的计算结果在数学上严格等价，只是中间的计算步骤和数据流量不同。在部署时，系统只需根据目标硬件一次性选择好路径，并对KV缓存做相应的初始化处理，之后在推理过程中就无需任何切换或改动。

四、把现有模型“变身”GQLA：TransGQLA流水线

从头训练一个GQLA模型成本高昂。为此，研究团队提出了TransGQLA——一套将现有GQA模型（如LLaMA系列）转换为GQLA模型的流程，其核心代码改动甚至只有一行之差。

这个流程建立在之前TransMLA工作的基础上。简单理解，TransMLA是把GQA模型改造成MLA，而TransGQLA则是在改造时，刻意保留了“按组索引”的特性，使得合并后的模型行为上依然是GQA，而非MLA。这一保留至关重要，它使得模型后续能支持高效的张量并行。

研究团队在LLaMA-3-8B模型上验证了这一流程。转换后，如果选择MQA吸收路径，KV缓存大小能压缩至原来的28.125%，超过7倍的压缩率；如果选择GQA路径，则缓存大小与原版GQA相近，保留了其传输效率优势。

五、稀疏注意力的扩展：为什么GQLA更有优势

对于需要处理超长上下文（数万甚至数十万词）的场景，稀疏注意力技术是关键。它让模型每次只关注最相关的一小部分历史词，从而大幅降低计算量。

然而，主流的稀疏注意力方案（如DeepSeek的DSA）与MLA结合时存在结构性问题：现代GPU的张量核心有一个硬性要求，每个计算单元需要至少16个查询头同时工作才能满负荷运转。MLA在非压缩模式下无法满足这个条件，导致其稀疏版本在所有硬件上都只能走效率不高的MQA吸收路径。

GQLA的标准配置（128查询头，8个KV组）则天生与硬件特性匹配：每组恰好对应16个查询头，完美契合张量核心的要求。这意味着稀疏化的GQLA可以在GQA路径上充分发挥硬件算力，在H20这类芯片上保持高效，同时依然支持灵活的并行扩展。

六、实验验证：转换后掉了多少能力，能恢复吗

能力损失是任何模型压缩或转换技术都必须面对的问题。研究团队在LLaMA-3-8B上进行了测试，评估了包括MMLU、ARC在内的六项常识推理基准。

结果显示，经过TransGQLA转换但未进行任何额外训练时，模型加权平均分从63.84降至54.13，下降了约9.7分。值得注意的是，在PIQA、HellaSwag这类基础语言理解任务上，分数下降很小，主要损失集中在需要大量知识记忆的任务上。考虑到超过7倍的缓存压缩率，这个损失程度在可接受范围内。

更值得关注的是恢复潜力。根据此前TransMLA的实验经验，在同样的模型上，仅需使用约300亿词元的数据进行继续预训练（这仅是LLaMA-3原始训练量的约五百分之一），模型能力就能恢复到与原始模型仅差0.5分的水平。由于TransGQLA与TransMLA在核心压缩步骤上一致，研究团队预计GQLA版本也能通过类似的少量训练实现能力恢复。相关实验正在进行中。

七、选择(g=8, sq=2)还是(g=4, sq=1)：两种H20优化方案的权衡

论文深入探讨了两种针对H20的优化配置，各有优劣。

默认推荐配置 (g=8, sq=2)： 即设置8个KV组，并开启多词预测（一次预测2个词）。这个方案优点全面：为数据压缩提供了充足的冗余空间，压缩质量更高；支持高达8路的零冗余张量并行，扩展性好；完美匹配硬件计算单元要求；在H20上能达到约22.1万词/秒的吞吐量。

轻量替代配置 (g=4, sq=1)： 即设置4个KV组，不开启多词预测。这个方案缓存压力更小，模型结构更简单，在H20上也能达到相近的吞吐量。但代价是压缩冗余度和并行扩展能力有所降低。

这里有一个精妙的设计：无论为H20选择哪种GQA路径配置，都不会影响在H100上MQA吸收路径的性能。在H100上，两种配置都同样以2.82微秒每步的极速运行。这正是GQLA设计的精髓——针对不同硬件的优化可以完全独立进行，互不干扰。

总结与展望

说到底，GQLA做了一件听起来像魔术，但道理很朴素的事：它证明同一批模型参数，可以通过不同的计算顺序，既表现出“精读摘要”的高效，也表现出“分组详读”的适配性，而且结果分毫不差。哪种计算模式更适合当前芯片的“脾气”，就用哪种。这种为同一套模型赋予硬件自适应能力的思路，在此之前并未被明确系统地提出。

这项工作的核心价值，在于打破了“为顶级芯片优化的架构在平价芯片上只能将就”的隐含假设。它为大模型的“一次训练，随处高效部署”提供了一条切实可行的技术路径。对于需要在从云端到边缘的不同算力硬件上灵活部署服务的团队而言，这个思路极具参考价值。

展望未来，一个自然的推论是：如果下一代模型在训练之初就采用GQLA结构，那么部署时将无需任何转换，直接根据硬件选择路径即可。这意味着GQLA有潜力成为未来大模型基础架构的一个标准选项。同时，TransGQLA流程也为现存的海量GQA模型提供了一条低成本的改造通道。

当然，论文也指出了几个有待完善的方向。屋顶线模型是理论分析，实际性能还需要在真实的H20、H100硬件上进行内核级基准测试来验证。TransGQLA的继续预训练效果也有待最终实验数据的确认。此外，目前的所有验证均在LLaMA-3-8B模型上进行，在更大规模参数模型和更多样化的任务（如长文本、代码生成）上的泛化性，仍需后续研究探索。