妈妈再也不用担心延迟了！斯坦福手搓Llama超级内核，推理仅需0.00068秒

【导读】斯坦福Hazy实验室推出新一代低延迟推理引擎「Megakernel」，将Llama-1B模型前向传播完整融合进单一GPU内核，实现推理时间低于1毫秒。在B200上每次推理仅需680微秒，比vLLM快3.5倍。

想象一下：你和AI聊天时，每句话都要等它3秒——血压是不是瞬间飙升？

低延迟LLM推理，就是专门针对这个问题的解决办法。

博客地址：https://hazyresearch.stanford.edu/blog/2025-05-27-no-bubbles

最近斯坦福Hazy实验室「整了个大活」：他们手搓了个叫做「Megakernel」超级玩具（推理引擎），把Llama-1B的前向传播整个塞进单个GPU内核！

结果直接炸场：

• H100上提速1.5倍，带宽利用率飙到78%

• B200上仅需0.00068秒（人类眨1/3眼的时间！）

• 比vLLM快3.5倍，把SGLang也甩出尾气

网友辣评：「传统推理引擎还在骑马，Megakernel已经开上战斗机！」

速度！使用32个token的提示词生成128个token的结果，未使用推测机制。

Fwd/s是衡量模型推理速度的一个指标，表示模型每秒可以执行多少次前向传播，数值越高，说明模型处理速度越快。

传统推理引擎：GPU在「摸鱼」

通常情况下，人们在GPU上运行代码的方式是启动一个「内核（kernel）」—— 一个执行明确定义操作的小型程序（例如RMS归一化、MLP等）。

当GPU运行大模型时，通常会把任务拆成上百个小内核，目前所有AI工作负载都是如此，表现为一系列相对较短小的内核。

比如先算RMS归一化 → 再搞注意力 → 接着MLP层...像流水线工人反复交接。

为了有一个初步认识，来看看Llama-1B Transformer模块中的运算，以及它们可能被划分的一些示例内核边界。

LLaMA-1B Transformer 模块的一组示例内核边界。红色框表示各个内核执行的工作

使用Llama-1B解码单个序列时，是一种纯粹受内存限制的工作负载：性能取决于是否能够持续不断地从GPU的全局内存中加载权重。

那么，为什么现有的方法距离充分利用GPU的全部带宽还差得远呢？

关键问题就是，当前基于内核的模型运行方法引入了停顿，阻碍了持续加载内存：

内核空闲

首先GPU内核是按严格顺序启动的，因此前一个内核中的所有线程块完全结束后，后一个内核中的线程块才能开始执行。

每次启动一个新的内核时，都必须等待前一个内核中所有落后的线程块完成。

例如，如果一个内核运行512个线程块（如LLaMA-1B降维投影（down projection）），但在B200上只有148个流式多处理器（streaming multiprocessors），那么到内核执行的最后阶段，就会出现80个空闲的SM——

148 -（512 - 148 * 3）= 80

因为大部分线程块已经运行完，只有少数几个还在运行，这些少数线程块占用了全部的SM，而其他SM就只能空等，造成资源浪费。

内核开销

其次，正如之前提到的，每次内核（kernel）的启动和关闭都会带来开销。

理论上，NVIDIA的CUDA图（CUDA graphs）可以在一定程度上隐藏这些开销，但根据测量结果来看，仍然有不少资源浪费。

各种内核开销的风格化甘特图。有时这些开销可以忽略不计，但很多时候并非如此！

举个例子，在H100上运行了一个简单的「假内核」测试（这个内核只记录开始时间、休眠一段时间、然后记录结束时间）。结果发现：

• 如果通过普通的CUDA流（stream）来运行，这个内核的启动开销大约是2.1微秒；

• 如果使用CUDA图，启动开销虽然下降了，但也只有降到约1.3微秒。

这段时间里，GPU并没有做任何有用的计算，只是在准备或等待，对整体性能来说是浪费。

所有优化目标是：让GPU的每一微秒都用在真正有意义的工作上。

内核等待

最后，即便启动了下一个内核，仍然需要等待权重（weights）和激活值（activations）加载完成之后，计算才能真正开始。

这些等待带来的延迟会让GPU空闲上成千上万个周期！

理想情况下，希望在执行前一个内核的计算和数据存储时，就能开始加载下一个内核所需的权重。为此，NVIDIA 提供了一种机制，叫做Programmatic Dependent Launch（PDL），它允许在前一个内核还在运行的同时，就提前为下一个内核做准备。

但是，PDL仍然会引入不必要的停顿，这是因为它的同步机制（cudaGridDependencySynchronize）太粗粒度了。举个例子，它要求所有的query、key和value全部计算完毕后，attention才能开始，而不能在某个head准备好时就立即开始计算。

稍后会在Llama-1B的一个具体案例中展示，这种机制在哪些情况下会限制性能。

综合来看，这些形成了标题中提到的「内存流水线气泡」——这也是并非始终从内存加载数据的一个关键原因。

对于短时操作来说，这些暂停累积起来会浪费大量的潜在带宽。

部分原因在于，Llama-1B（实际为1.24B参数）在批量大小为1时实在太过……微小：如果每个操作本身非常快，那么操作之间的间隔时间就开始变得至关重要。

为了说明问题的严重性：在单个H100上以16位精度进行单序列生成时，内存限制为3.35TB/s / 2.48GB = 每秒约 1350 次前向传递。

但每层需要7次内核启动，共16层，即使每次内核阻塞时间乐观估计为5微秒（包括尾部任务、内核启动和内存延迟），生成速度也将仅约为每秒770次前向传递。

实际情况往往更差。在低延迟工作负载下，GPU只有很少一部分时间真正用于执行有用的工作！

虽然CUDA确实提供了一些现有功能（例如图、流、PDL）来部分解决这些问题，但研究团队想看看是否有一种不同的方法可以解决所有这些问题，即将整个模型的前向计算融合成一个单一的内核。

如何设计Megakernel

如何将整个LLaMA前向传递过程融合到一个单一内核中，需要解决三个关键问题：

1 融合数十个操作从头开始做起来很难。需要一种机制来在Megakernel中执行这些操作。

2 为了在相同硬件上重叠多个操作，需要防止对有限资源（例如共享内存）的竞争。

3 在传统内核模型中，GPU会在每个内核之后进行同步。没有了内核，必须自己手动对GPU进行同步！

问题1：融合大爆炸

传统内核融合通常只合并两到三个操作。

但是这里需要融合大约一百个操作。

因此，需要一个合理的抽象方式，来对Megakernel进行编程。

一种方法是基于一个在GPU上的解释器——本质上是ThunderMLA底层架构的一个更复杂版本的解释器设计使得 GPU 中的每个流式多处理器（SM）都能接收一连串的指令（每个指令都使用相同的 CUDA 模板实现）并加以执行。

在Python端提前安排好每个SM的指令序列，值得注意的是，每个调度可以被重用于数百次前向传递！

对于端到端的Llama前向传递Megakernel，定义了以下指令集：

• 融合的RMS归一化、QKV和RoPE指令。

• 一个注意力计算指令。

• 一种注意力缩减指令（用ThunderGQA在长序列上的处理）。

• 一个O投影加残差指令。

• 融合的RMS归一化、上行门控和SiLU指令。

• 一个下投影加残差指令。

• 一个RMS归一化和语言建模头部指令，用于计算最终的tokenlogits。

• 使用一个通用的CUDA模板（包含加载、存储、计算样板函数）来实现每条这些指令，从而在的解释器框架内促进互操作性。

问题2：共享内存以消除内存气泡

指令与解释器结构使能够清晰地组织Megakernel。

然而，尚未解决一个关键问题：确保模型权重始终按顺序加载，以最大化内存带宽利用率。

Megakernel之所以能让解决此问题，是因为可以在指令之间进行内存加载流水线操作：解释器一旦能够开始加载某条指令的模型权重，即使前一条指令仍在完成阶段（例如将结果存储到全局内存），它也会立即开始加载。

正是这种紧密的指令间切换，最大限度地减少了因启动多个内核而可能出现的内存气泡。

然而，这里有个问题：如果下一个指令没有空间存放已加载的数据，那么从全局内存加载权重并不会带来太大好处！

更准确地说，所有的权重矩阵都是从GPU全局内存加载到SM的「共享内存」中——这是NVIDIA对每个SM上快速内存的称呼。

共享内存是每个SM上的稀缺资源，如果前一个指令占用了全部共享内存，就无法为新指令启动新的加载操作。

这就需要一种方法来跟踪哪个指令正在使用哪一部分共享内存，并在当前指令完成时迅速将共享内存过渡给下一个指令使用。

通过分页共享内存来实现这一点。

首先将H100上的前213KB共享内存划分为13个16KiB的页面，并将剩余的共享内存用于特殊用途，例如存储指令参数。

要使用这些页面之一，指令必须显式地从解释器请求并释放它们。

解释器会自动将已释放的页面传递给下一条指令，允许它们在共享内存可用后尽早开始发出内存加载操作。

问题3：同步

虽然Megakernels能够帮助最大限度地减少流水线气泡，但它们也引入了一个新的问题：同步。

在常规的多kernel执行模型中，性能限制在于，直到之前所有kernel中的线程块都完成后，下一个kernel中的线程块才能开始执行。

然而，正是这一特性使得管理数据依赖关系变得简单。当一个kernel启动时，CUDA保证该kernel所需的所有输入张量已经生成，并且可以安全地立即读取。

使用Megakernel时，没有这样的保障：当一个SM开始执行新指令时，其输入可能尚未就绪！

为了解决这个问题，在Megakernel内部显式地对指令进行同步。

通过一个简单的计数器系统来实现这一点。

在Megakernel启动之前，在GPU全局内存中初始化一个计数器数组（即整数数组），初始值为零。

每当一条指令完成时，它会增加其中一个计数器的值。

同样，每当新指令开始时，它必须等待其中某些计数器达到目标值，这表明其所有依赖项均已执行完毕。

这一优化在Llama-1B的大型多层感知机（MLP）中得以实现。

• 在使用PDL的朴素实现中，必须等待整个隐藏状态完成之后，才能开始下投影矩阵的乘法运算。

• 改为将中间状态划分为四个块进行处理，每个块都有各自的计数器。这样一来，针对下投影的指令只需等待其对应的输入块完成即可。

整合所有内容

据研究团队所知，H100 Megakernel代表了有人首次在GPU上实现以16位精度运行参数超过10亿的语言模型的前向传播时间低于一毫秒。

而的B200实现更是将这一时间进一步缩短至每次前向传播不到680微秒！

如文章开头的图片所示，Megakernel性能优于vLLM和SGLang基线（它们使用CUDA图和Torch编译）：

• 在H100上，Megakernel运行速度几乎是vLLM的2.5倍，比SGLang快超过1.5倍。

• 在B200上，与vLLM的差距扩大到3.5倍以上，仍然比SGLang快1.5倍以上。

距离B200上理论极限（大约每秒3,000次前向计算）仍有相当大的差距。

部分原因在于，该理论极限纯粹基于内存带宽——但仍需等待加载激活值。尽管这些激活值体积较小（不会占用大量带宽），但加载它们时仍然存在无法隐藏的延迟。

以下是当前B200前向计算运行时间的分解（总运行时间600微秒）：

• 存储激活值、等待一致性以及加载这些激活值需要花费250微秒。

这比简单模型预测的结果高出约20％：由于每条指令都依赖于前一条指令，需要支付两次加载延迟（检查就绪状态，然后加载激活值）和两次存储延迟（存储激活值，然后标记为就绪）的开销，每条指令都是如此。

以每次加载/存储大约500纳秒的延迟来计算，这将带来约200微秒的开销。（怀疑剩余的大约50微秒中，有一部分来自在全局内存中处理原子操作所花费的时间。）

• 实际运行RMS归一化和矩阵向量计算花费了200微秒。

这部分时间中约有95%用于矩阵向量运算。在Blackwell上，发现使用张量核心对此帮助不大；而在Hopper上，直接在CUDA核心上运行效果更好。这种差异的原因在于两种GPU的CUDA核心性能相对接近，但Blackwell的张量核心要快得多。

• 30微秒用于等待全局内存中的权重（流水线工作正常！）。

其中，40%的时间花费在LM头部，这是整个Megakernel中流水线效率最高的部分，因为其具有高度一致性和庞大的规模。

• 在各个线程束（warp）之间，有40微秒花费在低层次的同步开销上。

这里的一个关键问题是，即使在「通过」状态时，CUDA的异步屏障操作速度也相对较慢，每次都需要大约60纳秒的时间。

• 80微秒用于设置和各种其他开销。

例如，通过指令屏障、将页面标记为完成等。

本次突破明确展示了减少内核切换、优化内存流水线和精细同步的重要性，这也预示着低延迟推理技术的进一步发展潜力。

参考资料：

https://news.ycombinator.com/item?id=44111673

https://hazyresearch.stanford.edu/blog/2025-05-27-no-bubbles