何之源 2025-08-08 23:14 北京
一句话加速大模型Prefilling
©作者 | 何之源
单位 | 微软亚洲研究院
研究方向 | LLM efficiency
我们近期提出了一项新工作 TriangleMix：一种可以在几乎不损失性能的前提下，大幅加速 LLM Prefilling 阶段的简洁方法。该方法实现极为简单，相较现有加速方案可谓“极简优雅”，却能带来显著的加速效果。
论文已在 arXiv 发布，代码也已开源，欢迎感兴趣的朋友复现实验、交流反馈。本文将简要介绍 TriangleMix 的核心思路，同时分享我们在设计与验证过程中遇到的一些有趣发现与思考。
论文标题：
TriangleMix: A Lossless and Efficient Attention Pattern for Long Context Prefilling
论文链接：
https://arxiv.org/pdf/2507.21526
代码链接：
https://github.com/microsoft/MInference/tree/main/TriangleMix
介绍
在 LLM 的 Prefilling 阶段，attention 的用时会和长度 N 呈现 O(N^2) 的关系，在处理长文本的时候，attention 会成为性能瓶颈。我们提出了一个简单高效的 pattern 来加速：
简单来说，我们发现，在浅层使用 full attention，在深层使用如图所示的三角形形状的 attention（Triangle attention），可以保持 LLM（几乎）无损的性能！这个 Pattern 的好处在于：
1. Triangle Attention 的计算复杂度只有 O(N)，远比 full attention 的 O(N^2) 要小。此外，它是一个静态的 pattern，运行起来也会远比动态 attention 快。
2. 在浅层，可以使用动态 attention 方法（如 MInference 或者 FlexPrefill）进一步加速。
速度
我们测试了 triangle attention kernel 的速度，可以看到这种纯静态 pattern 相比动态 attention 方法确实是更快的，相比 full attention 加速效果也比较大，如在 128K 的长度上大概有 15.3 倍的纯 attention 加速：
在端到端的 Time-to-first-token（TTFT）上，在 32K 到 128K 的长度上，可以达到 -17% 到 -58% 填充时间的效果，128K 以上加速的效果应该会更明显。
性能
我们在 LongBench 的所有任务 和 RULER的 4K-128K 长度上测试，平均来说每种 setting 的掉点都在 1% 以内，和 FlexPrefil 的结合还会轻微提升性能，详细对比数据可以从论文里找到，这里就不放了。
一个比较重要的参数是“从第几层开始，做 Triangle Attention”，我们在 RULER 64K 上测试了三个模型，效果如下：
可以看到不同的模型的情况不一样，对于 Llama-3 系列的两个模型，基本上做12层的 full attention 就足够了，剩下来的层都可以只做 Triangle Attention，对于 Qwen2.5，稀疏度会低一些，要做 20 层的 full attention。
缘由
发现这个 pattern 其实也是机缘巧合。在一次实验中，它其实是一个 baseline 方法，结果我们意外地发现，这个我们认为会掉点严重的 Baseline 方法，居然没掉点，还比我们当时实验的方法更好！（经典 Baseline 干翻主方法）
这个现象引起了我们的注意，又做了更多的实验，实验发现它确实是一个比较靠谱的方法。接着我们又进一步思考，为什么它可以保持 LLM 的性能呢？
我们发现，这个现象可能和 Training 阶段、Inferencing 阶段不同的目标有关系，我们叫它 Train-test misalignment，在现有的工作中基本没有被提到：
在LLM 训练阶段，假设输入是 4096 个 token，预测的目标是这 4096 token 各自的 next token。也就是说，对于中间的某个位置，如输入的第 2048 个 token，我们要在这个位置预测第 2049 个 token。
但在 LLM 的 inference 阶段，假设输入的是 4096 个 token，我们实际对预测中间的 token 没有任何兴趣（如预测第 2049 个 token），我们只想知道预测第 4097 个 token 是什么以及 4097 之后的 token，这就是 Train-test misalignment。
Triangle Attention 挖掉的“中间那部分 attention”，其实主要对应了“在输入的第 2048 个 token预测第 2049 个 token”的这个任务，而和 inference 阶段的任务关系不大。
为了验证这一点，我们使用了一个导数的方法，导数的 y 是输出 token 的 logit，导数的 x 是各种不同 attention 区域的得分，结果如下图所示：
可以发现，Triangle Attention 挖掉的“Middle Q-K” 这部分的 attention，在深层的时候，导数会骤减！这就印证了这部分 attention 计算实际不影响后续 token 的生成。
此外，我们还验证了，去掉这部分 attention 后，确实会影响中间位置 token 的困惑度：
如上表所示，在输入的 2048 token 的 Wikipedia 文章中，应用 TriangleMix，可以发现中间部分 （1024-1152 位置的 token）的困惑度会激增，但后面部分的 token（1920-2048位置）的困惑度变化不大。
这就证明了，我们挖掉的那部分 attention 的计算，确实会比较影响中间 token 的预测，但这在 inference 阶段无关紧要。
总的来说，这个现象确实非常有意思，欢迎大家交流！
更多阅读
#投 稿 通 道#
让你的文字被更多人看到
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。
📝 稿件基本要求：
• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注
• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题
• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道：
• 投稿邮箱：hr@paperweekly.site
• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿
△长按添加PaperWeekly小编
🔍
现在，在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
·
