近日，斯坦福大学的Hazy Research团队公布了一项突破性的优化成果，他们成功地将开源模型Llama-3.2-1B的前向推理过程整合为一个名为“Megakernel”的巨型内核，这一创新将低延迟推理能力推向了新的极限。
在对话式AI和人类参与的交互式工作流等实时性要求极高的应用中，大语言模型的响应速度至关重要，直接关系到用户体验的好坏。然而，现有的开源推理引擎在处理这类极低延迟的单序列生成任务时，即使在顶级GPU如H100上，也往往无法充分利用其内存带宽。
Hazy团队经过深入研究后发现，限制LLM推理速度的关键瓶颈在于内存加载问题。具体来说，现有的推理引擎将Transformer模型的每一层拆解成数十甚至上百个CUDA kernel，每个kernel只执行非常小的操作，如RMS norm、注意力计算、MLP、Rotary Position Embedding等。这种拆解方式导致大量的上下文切换和等待时间，使得GPU在大部分时间里都处于“等待干活”的状态，而非“在干活”。
为了解决这个问题，Hazy团队提出了一个激进但有效的设计思路：将整个前向传播过程整合为一个单一的CUDA kernel，即Megakernel。他们基于已有的ThunderMLA架构，开发了一个轻量的GPU指令解释器系统，该系统为每个Streaming Multiprocessor（SM）预先分配一段包含多条按顺序排列指令的“执行计划”，每条指令代表Transformer模型中的一个结构单元。
这些指令包括融合RMSNorm、QKV projection、RoPE的复合指令，attention矩阵乘与缩减计算，O-projection与residual相加，MLP的RMSNorm、gate激活（SiLU）与上投影，down projection和最终residual，以及最后一层的RMSNorm加语言建模头。这些指令都基于统一的CUDA模板构建，实现了对load、store、compute的标准化封装。
为了确保高效的数据路径，解释器会将这些执行计划按模型结构静态编排，避免调度时的动态分支，从而提升吞吐与并发执行能力。同时，为了实现流水化计算并防止shared memory冲突，团队还对GPU的共享内存进行了分页管理，确保下一个计算阶段可以尽早开始预加载权重，从而最大化带宽使用率并消除“气泡”。
实验结果显示，Megakernel在H100上的推理延迟压缩至不足1毫秒，显存带宽利用率高达78%，相较于vLLM提升了2.5倍，相较于SGLang提升了1.5倍。在更先进的B200平台上，延迟进一步降低至600~680微秒，逼近理论极限。从一次完整推理的时间分布来看，Megakernel在存储激活、等待一致性与数据加载、RMSNorm与matvec等方面都表现出了卓越的性能。
Hazy团队的研究还揭示了一个关键问题：为什么现在主流的LLM推理系统在小batch、极低延迟场景下表现如此“不给力”。他们发现，像vLLM和SGLang这样的系统，在处理生成一个token这种极限情况时，GPU的显存带宽利用率非常低。核心原因是模型前向过程被拆成了太多太小的CUDA kernel，导致GPU在频繁切换kernel时产生了大量的固定成本和时间浪费。
因此，Hazy团队提出的核心解决方案是消除这些kernel边界，让GPU不再频繁切换任务。他们通过整合前向传播过程为单个Megakernel，实现了系统性优化。这一创新不仅提高了推理速度，还充分利用了GPU的显存带宽，为实时性要求极高的应用提供了强有力的支持。
Hazy团队还对CUDA异步屏障的性能进行了测量，并发现不同硬件架构上Megakernel的最佳实现路径应有所不同。例如，在Hopper架构（如H100）上，使用常规CUDA核心可能更有效；而在Blackwell架构上，Tensor Core则性能更优。这一发现为Megakernel在不同平台上的优化提供了重要指导。
总的来说，Hazy Research团队的Megakernel创新为LLM推理性能的优化提供了新的思路和方法。通过整合前向传播过程为单个巨型内核，他们成功消除了传统推理方式中的性能瓶颈，为实时性要求极高的应用提供了更高效、更可靠的解决方案。
未来，随着LLM模型的不断发展和应用领域的不断拓展，Megakernel优化方法有望在更多领域发挥重要作用，推动AI技术的进一步发展。
