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

二、低延迟LLM推理的重要性
想象一下：你和AI聊天时，每句话都要等它3秒——血压是不是瞬间飙升？低延迟LLM推理，就是专门针对这个问题的解决办法。

三、传统推理引擎的问题
（一）GPU内核运行方式
通常情况下，人们在GPU上运行代码的方式是启动一个「内核（kernel）」—— 一个执行明确定义操作的小型程序（例如RMS归一化、MLP等）。当GPU运行大模型时，通常会把任务拆成上百个小内核，目前所有AI工作负载都是如此，表现为一系列相对较短小的内核。

（二）现有方法的弊端
内核空闲
GPU内核按严格顺序启动，前一个内核中的所有线程块完全结束后，后一个内核中的线程块才能开始执行。例如，在B200上运行Llama - 1B降维投影（down projection），若一个内核运行512个线程块，而B200上只有148个流式多处理器（streaming multiprocessors），到内核执行最后阶段，就会出现80个空闲的SM，造成资源浪费。
内核开销
每次内核（kernel）的启动和关闭都会带来开销。NVIDIA的CUDA图（CUDA graphs）虽能在一定程度上隐藏这些开销，但仍有不少资源浪费。例如，在H100上运行一个简单的「假内核」测试，通过普通的CUDA流（stream）来运行，这个内核的启动开销大约是2.1微秒；使用CUDA图，启动开销虽然下降，但也只有降到约1.3微秒。
内核等待
即便启动了下一个内核，仍然需要等待权重（weights）和激活值（activations）加载完成之后，计算才能真正开始。这些等待带来的延迟会让GPU空闲上成千上万个周期。NVIDIA提供的Programmatic Dependent Launch（PDL）机制虽允许在前一个内核还在运行的同时，就提前为下一个内核做准备，但它的同步机制太粗粒度，仍会引入不必要的停顿。
四、Megakernel的设计与创新
（一）设计目标
将整个LLaMA前向传递过程融合到一个单一内核中，解决传统推理引擎存在的问题。

（二）关键问题的解决
融合大爆炸
传统内核融合通常只合并两到三个操作，而这里需要融合大约一百个操作。Megakernel基于一个在GPU上的解释器设计，使得GPU中的每个流式多处理器（SM）都能接收一连串的指令（每个指令都使用相同的CUDA模板实现）并加以执行。在Python端提前安排好每个SM的指令序列，值得注意的是，每个调度可以被重用于数百次前向传递。对于端到端的Llama前向传递Megakernel，定义了融合的RMS归一化、QKV和RoPE指令、注意力计算指令等多种指令集，并使用通用的CUDA模板来实现每条指令，从而在解释器框架内促进互操作性。
共享内存以消除内存气泡
Megakernel通过指令间进行内存加载流水线操作，解释器一旦能够开始加载某条指令的模型权重，即使前一条指令仍在完成阶段（例如将结果存储到全局内存），它也会立即开始加载，最大限度地减少了因启动多个内核而可能出现的内存气泡。但共享内存是每个SM上的稀缺资源，若前一个指令占用了全部共享内存，就无法为新指令启动新的加载操作。Megakernel通过分页共享内存来解决这个问题，将H100上的前213KB共享内存划分为13个16KiB的页面，并将剩余的共享内存用于特殊用途，指令必须显式地从解释器请求并释放页面，解释器会自动将已释放的页面传递给下一条指令，允许它们在共享内存可用后尽早开始发出内存加载操作。
同步
在常规的多kernel执行模型中，性能限制在于直到之前所有kernel中的线程块都完成后，下一个kernel中的线程块才能开始执行，这使得管理数据依赖关系变得简单。而Megakernel没有这样的保障，当一个SM开始执行新指令时，其输入可能尚未就绪。Megakernel内部显式地对指令进行同步，通过在GPU全局内存中初始化一个计数器数组（即整数数组），初始值为零，每当一条指令完成时，它会增加其中一个计数器的值，同样，每当新指令开始时，它必须等待其中某些计数器达到目标值，这表明其所有依赖项均已执行完毕。这一优化在Llama - 1B的大型多层感知机（MLP）中得以实现。
五、Megakernel的性能表现
（一）速度对比
H100平台
H100上提速1.5倍，带宽利用率飙到78%，运行速度几乎是vLLM的2.5倍，比SGLang快超过1.5倍。
B200平台
B200上仅需0.00068秒（人类眨1/3眼的时间！），将这一时间进一步缩短至每次前向传播不到680微秒，与vLLM的差距扩大到3.5倍以上，仍然比SGLang快1.5倍以上。
（二）性能分解（以B200为例）
存储激活值、等待一致性以及加载这些激活值：花费250微秒，比简单模型预测的结果高出约20％，由于每条指令都依赖于前一条指令，需要支付两次加载延迟（检查就绪状态，然后加载激活值）和两次存储延迟（存储激活值，然后标记为就绪）的开销，每条指令都是如此。以每次加载/存储大约500纳秒的延迟来计算，这将带来约200微秒的开销，剩余的大约50微秒中，有一部分来自在全局内存中处理原子操作所花费的时间。
实际运行RMS归一化和矩阵向量计算：花费了200微秒，这部分时间中约有95%用于矩阵向量运算。在Blackwell上，发现使用张量核心对此帮助不大；而在Hopper上，直接在CUDA核心上运行效果更好，这种差异的原因在于两种GPU的CUDA核心性能相对接近，但Blackwell的张量核心要快得多。
等待全局内存中的权重：花费30微秒，其中40%的时间花费在LM头部，这是整个Megakernel中流水线效率最高的部分，因为其具有高度一致性和庞大的规模。
各个线程束（warp）之间的低层次的同步开销：花费40微秒，这里的一个关键问题是，即使在「通过」状态时，CUDA的异步屏障操作速度也相对较慢，每次都需要大约60纳秒的时间。
设置和各种其他开销：花费80微秒，例如，通过指令屏障、将页面标记为完成等。
六、Megakernel的意义与展望
本次突破明确展示了减少内核切换、优化内存流水线和精细同步的重要性，这也预示着低延迟推理技术的进一步发展潜力。虽然Megakernel性能优于vLLM和SGLang基线，但距离B200上理论极限（大约每秒3,000次前向计算）仍有相当大的差距，部分原因在于该理论极限纯粹基于内存带宽——但仍需等待加载激活值。