@AryaTschand: 我在英伟达架构研究部门的实习中，大部分时间都在思考前沿LLM应如何在我们的最新（异构）系统上进行推理……

我在英伟达架构研究团队的实习期间，大部分时间都在思考如何让前沿LLM在我们的最新（异构）系统上进行推理。

我想分享一些我从基本原理中学到的关于推理的知识。

太长不看 + 线程中的小实验！

核心来说，一切都归结为算术强度（AI）——计算出的FLOPs与移动的数据字节数之间的比率。

每个ML芯片都有其峰值FLOPs和内存带宽，而工作负载的AI告诉我们，我们是等待MMA运行还是等待数据移动。一个屋顶线模型（附上）是任何芯片的属性，它根据工作负载的AI告诉我们，你是会受计算限制还是受内存限制。

受内存限制 = AI小于芯片的“脊点”，我们将等待数据（Tensor Core将未被充分利用） 受计算限制 = AI大于芯片的“脊点”，我们将等待计算（带宽将未被充分利用）在GPU中，我们希望这种情况

LLM推理可以合理地分解为以下几个阶段：

1. 预填充 - 加载新生成的Token（工具调用的输出），非常受计算限制，这对GPU非常有利
2. 解码注意力 - 自回归地关注KV缓存中的过去Token，通常受内存限制，但通过MLA、稀疏注意力等策略，可能接近受计算限制
3. 解码前馈网络 - 每个Token加载权重并执行GEMM，通常在低批次时非常受内存限制，因此对GPU来说很不利

注意，在更新的混合模型中，我们甚至可以将这些阶段进一步细分。

随着GPU向更大的Tensor Core发展以支持大规模训练（LLM训练非常受计算限制，因此通过更低精度和更大的脉动阵列增加总FLOPs，使我们能够更快地训练GPT 6.3.2 Pro Max Codex），提高芯片的FLOPs相对比提高HBM带宽更容易，这就是为什么我们看到GPU的脊点越来越高。

这对预填充（高AI）来说很好，但对解码（低AI）来说非常不利，我们通常在解码时看到GPU利用率低，尤其是在没有多Token推测技巧来提升解码AI的情况下，比如草稿/验证推测解码、自推测或扩散。

为了解决这个差距，一些人开始制造只使用SRAM的芯片。

在不深入内存系统的情况下，目前所有GPU都使用HBM，它具有高容量，但延迟实际上比片上内存差很多。这在我们受计算限制时（并且通常需要大量内存来存储优化器状态或KV缓存）很好，但在解码中受内存限制时却有害。

SRAM是片上内存（台积电将其集成在晶圆上，而不是从美光单独购买HBM），具有更低的延迟，但容量也低得多。如果你能卸载KV缓存或保持上下文较短，这将使解码更快得多。

批次大小（同时推理的总数）与算术强度有非常重要的关系。由于我们知道AI对推理效率的重要性，它也与吞吐量和交互性有重要关系。

增加批次大小（意味着我们将更多请求打包在一起）几乎线性地增加了解码FFN（推理中受内存限制最严重的阶段）的算术强度。

这意味着，当我们增加批次大小时，吞吐量基本上会免费提升，因为我们能更好地利用Tensor Core。

那么，为什么我们会在低批次大小下运行呢？如果你用过LLM API，你可能见过/快速模式或更慢的批处理推理。除其他因素外，批次大小很可能是他们为了使其工作而调整的参数之一。

较小的批次大小意味着我们需要移动更少的KV缓存数据（每个批次需要独立的KV），从而降低整体Token延迟。这意味着Token生成得更快，但整体系统吞吐量更低。这就是为什么快速Token如此昂贵！

*关注系统推理特性分析的人可能对右侧的图表很熟悉。它展示了交互性（Token生成速度）和吞吐量（系统每秒产生的总Token数）之间的帕累托边界。在这个曲线上，没有绝对正确的答案，但它通常表明，让每个用户的Token更快，会导致系统整体Token变慢。

在解码阶段（希望我已经让你相信这非常重要，而且在GPU上解决也是一个难题），我们并不总是看到注意力和FFN之间有相同的划分。

FFN非常直接——极度受内存限制，AI随批次大小缩放，上下文长度不影响任何东西，因为激活值大小相同，而且我们对此无能为力。

然而，注意力的情况有趣得多。与FFN不同，注意力的计算成本随上下文长度平方地增加。我们喜欢长上下文，因为它意味着我们的Agent可以进行大量回合而无需压缩重要信息，但这使得注意力非常昂贵（不过好的一面是，这也使注意力更受计算限制）。

为了解决这种平方缩放问题，许多优秀的ML研究人员一直在探索次平方注意力，即不随上下文长度平方缩放的注意力。

更简单的变体如GQA和MLA减少了KV缓存的占用。更深入的方法如稀疏注意力（DeepSeek、GLM）、滑动窗口注意力（Gemma）和线性注意力（Nemotron）在更长上下文下需要更少的FLOPs，并且减少了KV的内存占用。

由于注意力随上下文长度缩放而FFN不缩放，因此在短上下文中FFN延迟占主导，而在长上下文中注意力延迟占主导。

现在我们已经理解了LLM推理中的算术强度模式，我们可以思考如何在真实系统上高效地执行它。

默认的Serving串行地使用所有GPU进行预填充，然后进行解码。正如你可能想象的，预填充期间的利用率很高，但解码期间利用率很差。

在ISCA 2024上，一些我非常尊敬的计算机架构研究人员发表了“Splitwise“论文。这是第一篇探索推理分离式Serving的论文，通过在受计算限制的系统上做预填充，在受内存限制的系统上做解码，现在称为PD分离。这种方法至今仍然很流行（尽管我们做了很多艰难的工程改进才能大规模运行）。

最近，Attn-FFN作为一种更细粒度的分离方案开始流行。由于注意力正变得更受计算限制，我们可以只将解码FFN放在受内存限制的系统上，而将预填充和解码注意力进行PD分离。然而，这显著增加了跨互连移动的数据量和移动频率。

随着我们转向具有独特AI模式的混合次平方注意力模型，以及预填充系统保持为GPU、解码系统可以是基于SRAM的芯片（能轻松处理受内存限制的工作负载）的更多异构系统，也许我们Serving的方法将继续改变……？

感谢阅读，希望有些部分能带来启发。

如有任何问题，请随时联系我——完整博客文章即将在我的网站上发布！

好想法！提高FFN的算术强度非常重要，但最近相比注意力变体研究得少得多。

某种融合来限制片上数据传输，或者为MoE找到更好的专家并行方法，都是可以开始的方向。

向最优秀的人学习

谢谢！