[音乐]

大家好,我是Daniel,我在Cerebras负责AI和AI代理相关工作。你叫什么名字?我叫Michelle,是Gimlet Labs的联合创始人,我们正在构建一个多硅片的新云架构,以优化堆栈每一层的推理性能。我们非常兴奋地向大家介绍我们的演讲内容:为了满足大家对Token无休止的需求,我们需要做哪些工作。众所周知,目前的AI支出正在疯狂增长。今年我们计划在计算资源上花费7650亿美元,而到2031年这一数字将达到1.6万亿美元。前沿实验室正在疯狂争夺尽可能多的计算资源,向硬件供应商购买设备,既用于训练前沿模型,也服务于所有用户。所有这些计算资源都是为了满足每个人对Token的无尽需求。

他们正在预留内存槽位,但实际上并未使用。紧随其后的一项创新被称为提示缓存(Prompt Caching)。这里的想法是:如果你有一个编码代理,或者一个拥有极长系统提示词的聊天应用,你实际上不必在每次请求时都重新计算所需的 KV 缓存。你可以做的是复用为前一个请求已经计算好的 KV 缓存,用于下一个请求。这样你就可以从之前的地方继续处理。在此过程中,你可以跳过非常耗时的预填充(prefill)阶段。

你可能会问:什么是预填充?在推理过程中,你可以将发生的事情分为两个主要阶段:预填充和解码(decode)。预填充阶段包括请求到达后发生的所有事情,直到生成第一个 token 时结束。当你向 ChatGPT 输入内容并等待几秒钟它才开始返回第一个 token,之后便流畅输出时,就能看到这一点。在最初几秒没有返回任何内容的时间里,它正在为你输入的每一个提示 token 计算 K 和 V 向量,并生成第一个输出 token。

这一阶段非常受限于算力(compute bound),因为你可以对其进行极致的优化,实现高度的并行化。然而,在生成第一个 token 之后,发生的事情是:你一次解码一个 token,并复用预填充阶段刚刚构建的 KV 缓存。因此,这一阶段更依赖于内存带宽。我之所以谈论这些,是因为最近发生了一些事情,严重扰乱了许多推理管道,那就是极度消耗 token 的编码代理。

假设你有一个使用 20K、50K 甚至一百万上下文长度的编码代理,会发生什么呢?它会开始阻塞管道。当预填充运行时,没有其他预填充或解码请求能够通过。如果你有一群人发送编码请求,而只有几个人试图在聊天模式下解决数学问题,这种情况将是灾难性的。为了避免这种堵塞问题,有一项名为“分块预填充”(chunked prefill)的创新技术可以解决这个问题。

其想法相当简单:将非常长的预填充过程切分成更小的片段,以便交错处理不同的解码和预填充请求。这样就不会有一个单一任务占用大量时间并堵塞整个管道。现在,请 Michelle 告诉我们更多关于我们可以优化的其他方面。

是的,Daniel 刚才提到的所有优化都是运行时级别的优化。这解释了我们在单个工作节点内如何提高推理效率。但是,为了继续推动这一前沿领域,我们确实需要开始关注系统级别的优化,即如何将这些工作负载分布在许多不同的工作节点上。首先要认识到的是:每个请求都是不同的。例如,以摘要用例为例,你将一堆文档扔给大语言模型(LLM),只希望得到一个问题的答案。

这种工作负载将具有长提示词和短回答的特点,因此主要由预填充阶段主导,属于算力受限型。另一方面,你可能遇到相反的情况。假设你想写一个长故事或一本书,只需给它几个主题供其发挥。这将具有短提示词和长回答的特点,由于输出 token 的数量远多于输入 token,因此主要由解码阶段主导,属于内存受限型。

最后,你可能遇到介于两者之间的情况:长提示词和长回答。在这里你可以看到,预填充和解码阶段都很重要,因此算力和内存都至关重要。通过观察这些不同的例子,你可以清楚地看到硬件的资源需求根据你的工作负载类型而有很大差异。那么,如何避免自己受到算力或内存的限制呢?

一种方法是称为“解耦”(disaggregation),即将你的工作负载拆分到不同的工作节点上。最常见的形式称为“预填充-解码解耦”,即按阶段拆分和解耦这种工作负载。这允许你拥有两个池:预填充池,负责运行工作负载的预填充部分;解码池,负责运行工作负载的解码部分。以摘要用例为例,你需要扩展预填充池以避免受限于算力,同时缩减解码池,因为你不想在几乎没有太多解码工作时浪费资源。

这很好,因为现在你可以独立扩展所有这些部分,且任何一部分都不会阻塞另一部分。我们可以再进一步:之前我们说过,预填充工作负载本质上受限于算力;解码则受限于内存。快速回顾一下原因:在进行解码时,你逐个遍历这些 token,对于每一个解码步骤,你都读取整个 KV 缓存。这显然是内存受限的。

是的,Michelle 刚才提到的这一点——即这是内存受限的——可以通过硬件的实际能力清楚地看出来。GPU 将数据从内存传输到计算单元的速度受限于所谓的“内存带宽”。对于 H100 来说,这大约是每秒 3.35 TB。然而,Cerebras 拥有比 H100 高 7,000 倍的带宽,达到每秒 21,000 TB。这是因为我们的芯片上拥有非常专门的基于 SRAM(静态随机存取存储器)的内存架构。

这意味着 Cerebras 真的是为解码而构建的。我们可以以每秒 3,000 token 的速度提供服务,而 H100 的解码速度仅为每秒 600 token。我们非常兴奋能与 AWS 合作,实现这一技术的生产化部署。我们将与 AWS Trainium 合作,让他们使用 AWS Trainium 芯片来服务预填充,利用他们的 EFA(弹性Fabric适配器)服务将 KV 缓存传输到我们的 Cerebras CS-3 系统,最终提供 token。

这将使我们在相同的占地面积内获得 5 倍的高速 token 容量。这是我们正在做的一件事,旨在使快速推理变得更加普及和廉价。

开发者们对此充满期待,我们非常热衷于推进这一合作关系。而在 Gimlet 方面,我们正在将这种理念进一步泛化,不仅仅局限于训练环节以及与 Cerebras 的合作。Gimlet 正在构建的本质上是一个多硅片推理云,它接收你的工作负载,将其分解为更细粒度的部分,确定所有这些不同部分的资源特征,然后找出将这些任务映射到底层硬件的最佳方式。这里的“底层硬件”可以是任何东西:Nvidia、AMD、Trainium、Cerebras、D-Matrix 或 Intel Gaudi,名单还在不断延长。通过在最适合运行该任务的机器上执行这些任务,我们确实看到了令人惊叹的结果。

与具有相同功耗预算的同构系统相比,我们的推理速度提高了三到十倍。这不仅仅是比以往更深入的预填充与解码分离(disaggregation),其中一些技术例如基于模型而非阶段进行分离,这被称为投机解码。回顾一下解码阶段:通常你会为输出的每个解码词元执行一次前向传播,这是标准做法。而投机解码的关键在于你拥有一个草稿模型。这仅仅是完整模型的缩小版,通常可以近似地模拟出从完整模型中获得的结果。

其优点在于,由于草稿模型较小,它运行速度快得多且资源消耗更少。在投机解码中,为了获取接下来的九个词元,你会先在草稿模型上运行解码过程。它会返回所谓的候选词元。虽然你不能百分之百确定这些词元都是正确的,但由于它们通常经过微调或基于完整模型训练过,因此结果大致准确。你需要验证这些是否是你真正想要输出的词元。

为此,你会在完整模型上运行下一阶段,称为验证阶段(verify stage),以确认哪些词元是正确的。你可以不断重复这一过程,直到从草稿模型中验证所有词元为止。最终你会发现(如上图所示的标准方式),为了获得这九个词元,你需要在完整模型上运行九次前向传播。

另一方面,使用投机解码时,你实际上只需要在完整模型上运行三次。如果能让草稿模型跑得极快,那么这部分时间几乎可以忽略不计,从而显著缩短获取词元所需的时间。再看一些具体结果:这是丹尼尔之前展示的帕累托前沿,横轴是获取词元的速度,纵轴是吞吐量。

我们比较的基准案例是在同构 GPU 上进行的预填充与解码分离方案。你可以看到,吞吐量表现不错,但一旦追求更高速度,吞吐量就会急剧下降。那么,如果在这些相同的 GPU 上运行投机解码会怎样?现在你正在运行草稿模型,能够看到曲线向右移动。

原因在于你在完整模型上运行的解码次数减少了。但我们还能进一步推演,回到我们的主题:在最能胜任的硬件上运行工作负载。在使用诸如 D-Matrix 或 Cerebras 这样的设备时,由于它们能极快地运行小模型,你会看到性能的大幅提升。在这里,对于相同的吞吐量指标,你获得的交互速度是同构机器上的十倍。

我们要讨论的分离形式的最终形态称为注意力-前馈网络(Attention-FFN)分离。当考虑服务大型模型时,这一点尤其令人兴奋,比如拥有三万亿参数的模型。这些模型如此庞大,以至于仅运行一层就可能非常消耗资源。其概念本质上是相同的:针对你需要执行的特定任务,无论是注意力层还是前馈网络(FFN)层,你都在试图确定其资源特征以及最适合映射的硬件。

注意力层更受内存带宽限制,因为它涉及大量的 KV 缓存访问;而 FFN 层则进行大量密集矩阵乘法。因此,逻辑上很清晰:你应该在具有内存带宽优势的机器上运行注意力层,而在能快速执行矩阵乘法的机器上运行 FFN 层。

本质上会发生的情况是,即使在模型内部,为了从该模型获取推理结果,你也会在多台机器之间来回切换(ping-ponging),真正地将每一层都进行了分离。这里我们需要特别指出的是,不存在最佳的分离形式,也不存在最佳的硬件。这真的取决于你的权衡:根据你的模型、工作负载以及对延迟和吞吐量保证的需求而定。

展望未来,我们看到智能体(agents)正日益融入我们的日常生活和工作之中。每个人都在使用它,这持续推动着对更多、更快词元的需求。“是的,趋势确实向右上方发展。”“我告诉过丹尼尔别说这个(笑)。”

但好在的是,尽管这看起来是一个疯狂的目标,我们实际上可以利用智能体来帮助我们解决这些问题并不断推动前沿。因此,我们在 Gimlet 所做的一件事就是利用智能体来生成优化的内核(kernels)。特别是对于我们要支持如此广泛不同种类硬件的用例,我们需要弄清楚如何使这个问题更具可行性。

因为目前人们生成和优化内核的方式是:许多人坐在那里手动调整,试图找出优化这些内核的方法。问题是,这些优化的效果会根据你尝试运行的模型、运行所在的硬件甚至同一代硬件的不同世代以及工作负载形状本身而产生巨大差异。因此,我们在 Gimlet 构建了 KForge,这是一个用于生成优化内核的智能体。

这些内核接收大量信息,包括硬件配置、工作负载配置文件等,基本上是利用任何可用于优化内核的信息。我们将此作为暴力求解问题来处理。这里有一个例子:前向传播过程以及 KForge 从中优化出的内核,我们看到了相当有趣的结果。上周,我们的智能体运行时间不到一周,在此期间它能够在 AMD 硬件上将 GPT-OSS-120B 的整体速度提升 33%。值得注意的是,这是在已经优化过的内核基础上构建的。你可以将智能体推向更高水平,利用它们帮助我们构建和设计针对特定模型或工作负载进行超调优的定制硅片。

这随后可用于训练更优秀、更智能且更快的 AI,而这些新模型又可用于设计下一代定制硅片。这种循环不断持续下去,形成一个递归闭环。(向 Recursive 团队致敬。)我想在演讲结束时表示,我们在优化推理吞吐量和速度的旅程中才刚刚开始。因为在堆栈的最顶层,当谈论事件响应时,你通常不会真正考虑到推理优化。但当智能体使我们的系统更加可靠并增加正常运行时间时,我们会获得更高的吞吐量并能处理更多 Token。

或者你可以将智能体应用于模型部署的问题中,即在新模型上市的第一天就使其在各种不同的硬件上运行。你甚至可以在编译器方面做得更激进一些,想象这样一个世界:推理过程可以根据你的工作负载和运行的模型进行定制构建。Anthropic 团队已经展示了这一点,他们在两周内就构建了编译器。因此我认为,当你为支持的模型、工作负载和定制硬件构建即时定制软件时,有很多优化空间。

回顾上一张幻灯片,利用智能体继续帮助我们构建更多定制硅片并加速流片过程。最后,这部分不太涉及技术细节,但围绕推理服务的所有事项——包括所有采购、法律事务以及后台运营工作——结合智能体工作流程,都可以得到优化。因此,当我们完成所有这些工作时,我对未来充满乐观,相信我们能实现更快、更便宜地提供 Token 的目标。

请翻到下一页幻灯片。非常感谢大家参加我们的演讲。我们正在非常努力地工作,以提供更快、更便宜的 Token。如果有任何问题,请告诉我们,演讲结束后我们会留下来解答。(是的,去吃 Token 吧。)我们正在为你制造更多 Token,这样你也可以生成像这样的 GIF。是的。(笑声) [music]