单核运行整个模型！AutoMegaKernel将Llama编译为持久协同内核，零手工CUDA代码

在当今的大语言模型推理部署领域，我们常常面临一个效率瓶颈。当你向模型输入一个提示词，等待它生成下一个词元时，传统的执行方式（如PyTorch）会为模型中的每一个算子（如线性层、注意力层、归一化层）单独启动一个GPU内核。这就像是让一支庞大的施工队，每砌一块砖都要开一次会、领一次任务，大量的时间浪费在“开会”（内核启动）和“搬运材料”（数据在高速显存中的往返传输）上，而非真正的“砌墙”（计算）本身。

这种模式对于单批次、低延迟的推理场景（比如你与聊天机器人的一次对话）尤其低效。每一次内核启动都带来CPU开销，每一次算子间的数据交换都要经过显存带宽，而单步解码本质上是一个受内存带宽限制的操作——它需要读取全部权重一次，只产生一个词元。理论上的延迟下限就是权重字节数除以显存带宽，但传统方式远未触及这个极限。

有没有可能让整个模型的前向传播，像一支训练有素、协同作业的施工队一样，只开一次工前会，然后一气呵成地完成所有工作？这就是“超级内核”的核心理念：将整个模型编译成一个持久化的、协同的内核，一次启动，全程在芯片内部协同计算，彻底消除内核启动和中间数据回写的开销。

近期，来自RightNow AI的研究者Jaber Jaber和Osama Jaber提出了AutoMegaKernel（AMK） 系统，将这一理念向前推进了一大步。AMK不仅能够自动将HuggingFace的Llama家族模型编译成单个超级内核，更重要的是，它引入了一个静态检查的调度中间表示层，能够在代码运行前就确保生成的内核不会死锁或发生数据竞争，为智能体驱动的自动优化提供了一个安全、可靠的编辑界面。

论文标题： AutoMegaKernel: A Statically-Checked Agent Harness for Self-Retargeting Megakernel Synthesis

论文链接： https://arxiv.org/pdf/2606.09682

开源仓库： https://github.com/RightNow-AI/AutoMegaKernel

研究背景：从“内核拼图”到“超级内核”

当前主流的大语言模型推理系统，如vLLM、SGLang、TensorRT-LLM等，主要优化目标是高批次的吞吐量。它们通过先进的批处理、KV缓存管理技术来提升效率，但在算子层面，通常还是采用“一个算子一个内核”或“一个引擎一个内核”的模式。CUDA Graphs技术可以捕获并重放一个算子序列，从而分摊内核启动的开销，但算子间的边界以及由此产生的显存数据往返依然存在。

对于追求极低单次响应延迟的场景（如实时对话、边缘设备），这种模式就成了瓶颈。研究者们开始探索“超级内核”的路径。例如，Mirage Persistent Kernel (MPK) 系统能够将张量程序自动转换为一个带有内核内调度的持久化内核。也有工作通过手工编写的方式，为Llama-1B模型构建了一个超级内核，移除了所有启动气泡，观察到在低延迟下，vLLM和SGLang最多只利用了H100显卡一半的带宽潜力。

然而，这些方法在可信度和可移植性上存在缺口。MPK没有提供静态的死锁/竞态安全检查门；手工编写的内核则针对单一模型和单一架构，难以扩展和维护。更重要的是，它们没有设计成可供编码智能体安全驱动的“编辑界面”。

AMK的核心创新：安全第一的编译架构

AMK建立在团队之前的工作AutoKernel之上，后者是一个通过迭代式智能体驱动搜索来优化单个GPU内核的自主系统。AMK将这种“提出→评估→保留/回退”的搜索方法论提升了一个抽象层次：从一次优化一个内核，升级为将整个模型编译并自我优化为单个持久化超级内核。

AMK最核心的贡献在于，它将正确性视为编译器的一个结构不变量，而非输出结果的属性。整个系统分为四层（如图1所示），并通过两层自动研究循环进行驱动。

图1：AMK的四层系统架构与两个自动研究循环。第0层是可信的、冻结的、手写的每架构基础；第1-2层由智能体搜索，但受到隔离的正确性检查（循环1）和静态调度验证器（循环2）的约束。

1. 静态验证的调度IR（中间表示）

这是AMK安全体系的基石。模型的前向传播被“降低”为一个类型化的调度IR，这是一个在流多处理器级别、通过单调计数器同步的任务有向无环图。一个冻结的验证器会在这个IR被加载到GPU运行之前，静态地检查其是否无死锁和无竞态。

验证器执行一系列图检查（如表1所示），例如：确保每个等待的阈值在1和生产者数量之间（避免不可满足的等待）；检查生产者-消费者图是无环的；确保每个SM的串行队列是DAG的一个线性扩展；对于有多个生产者的计数器，任何等待都必须使用阈值等于生产者总数（避免“部分等待”引发的竞态）；对于任何读操作，其写操作者都必须是拓扑序上的前驱。

图2：AMK的“构造即正确”编译流程。一个HuggingFace Llama模型被降低为类型化的调度IR（一个SM级别的任务DAG），通过静态验证门后，作为一个持久化的协同超级内核运行。同一套代码源可重定向到sm_80/sm_90/sm_120架构。

研究团队对验证器进行了严格的压力测试。在7160个对抗性构造的调度中（其中6091个被独立预言机确认为不安全），验证器实现了零误接受，并接受了所有360个真实的调度降低，其中24个与PyTorch的eager模式逐比特匹配。这意味着，一个由智能体提出的不安全调度，会在验证阶段被干净地“拒绝”，而不会导致GPU挂起。

2. 自动生成与自我重定向

AMK无需为每个模型手写CUDA代码，即可为10个支持的Llama家族模型自动生成正确的超级内核，包括三个真实的检查点（SmolLM2-135M/360M, TinyLlama-1.1B）。对于不支持的变体（如带有偏置投影的模型），系统会在导入时明确拒绝。

更重要的是，同一套源代码可以针对不同的GPU架构（如sm_80的A100、sm_90的H100、sm_120的RTX 5090）生成并运行正确的超级内核，其生成代码从运行设备动态派生。在真实的SmolLM2-135M检查点上，AMK不仅能够逐词元复现HuggingFace的贪婪解码结果，其教师强制困惑度也与原版匹配至约2.5e-7的精度。

3. 智能体可驱动的编译器与自我改进

AMK为智能体提供了一个结构化的编辑界面——ScheduleConfig对象，而非原始的内核代码。智能体可以调整分块大小、融合分组、SM分配策略、流水线深度等参数。冻结的虚拟机（VM）会确定性地将任何配置点“降低”为超级内核程序，而验证器则保证无论选择哪个点，结果都是安全的。

在此基础上，AMK集成了一个无人值守的自动调优循环。该循环在一个小的调度和内核参数网格上进行搜索，每次尝试都经过“降低→验证→正确性门控→测量→保留/回退”的流程。它采用了抗时钟漂移的逐样本交错测量和物理性能上限诚实保护机制。实验表明，这个循环能够自动将生成的超级内核在其默认配置基础上提升1.25至1.72倍。

性能评估：诚实的比较与有意义的胜利

AMK的研究者进行了一场全面而诚实的性能评估，涵盖了七款不同的GPU，并报告了所有方向的結果，包括胜利和不足。

量化支持与性能

AMK能够自动生成int8和int4权重量化的超级内核。int8路径实现了无损的贪婪解码（在SmolLM2-135M上32个词元完全一致），并且每个词元解码速度比bf16版本快1.12倍。int4路径虽然能减少2.42倍的权重数据流量，但由于简单的舍入策略会带来精度损失，速度提升相对有限（1.06倍）。这为实际部署提供了清晰的选择：int8在几乎不损失精度的情况下带来了可观的加速。

与cuBLAS的较量：一个“非对称”的胜利

最引人注目的实际硬件结果出现在与NVIDIA cuBLAS库的对比中。在针对推理任务优化的数据中心GPU（如L4、L40S、A10G）和消费级GPU（RTX 5090）上，AMK通过搜索找到的int8权重量化超级内核，在单批次解码任务中击败了使用CUDA Graphs的cuBLAS bf16实现。

• 在L4上，优势随着模型增大而增长，最高可达1.33倍。
• 在当前一代的推理旗舰L40S（864 GB/s带宽）上，优势为1.25-1.27倍。
• 在A10G上，模型达到一定规模后优势可达1.08倍。
• 在消费级的RTX 5090笔记本GPU上，优势为1.19-1.23倍。

图3：AMK在L4推理GPU上自我调优超越cuBLAS的过程，且优势随模型规模增大而增长。灰色点代表搜索中每个配置在3.5B模型上的表现（比值=cuBLAS/AMK int8）。实线为3.5B模型的最佳历史表现，虚线为2.7B模型的最佳历史表现。

研究者特别强调，这是一个精度非对称的比较：AMK使用的是int8权重、bf16激活（W8A16），而对比方是纯bf16。在同等精度的bf16对比中，AMK的超级内核仍然落后于cuBLAS（约0.76-0.88倍）。因此，int8的胜利主要来自于传输更少的权重字节数，而非其内核每字节的处理速度更快。这一胜利的适用范围也很有启发性：它不适用于高带宽的训练级GPU（如A100/H100）。在这些卡上，AMK的超级内核因每块数据所需的跨SM同步开销而成为瓶颈，导致性能落后于cuBLAS。研究者没有隐藏这一边界，而是明确地将其作为局限性进行了分析和披露。

性能天花板与差距

AMK优化后的bf16 GEMV（通用矩阵向量乘）内核在RTX 5090上达到了约460 GB/s的带宽，这大约是实测峰值带宽（731 GB/s）的63%，而同等条件下cuBLAS的天花板约为实测峰值的90%。在端到端的单次解码延迟上，AMK的bf16超级内核仍比CUDA Graphed的cuBLAS实现慢约13%。在数据中心GPU上，与默认模式的vLLM相比，AMK（使用早期GEMV版本）也处于劣势（在H100上慢约1.65倍）。

这些差距诚实地指出了AMK当前内核实现（v1 GEMV）的优化空间。研究者分析，下一步的关键是设计具有更粗粒度同步、更深内存流水线的GEMV内核，而这正是AMK所提供的安全编译架构未来可以安全吸纳的改进。

意义与展望

AutoMegaKernel的工作为大语言模型的高效、安全部署提供了一条新的路径。它的核心价值不在于目前是否在所有场景下都达到了最快的绝对速度，而在于构建了一个可信的、自动化的、智能体可驱动的超级内核合成框架。

1. 安全性前置：通过静态验证，将运行时可能发生的危险（死锁、竞态）提前到编译时拦截，这对于构建可靠的自主优化系统至关重要。
2. 自动化与可移植性：告别针对特定模型和架构的手工优化，实现“一次编写，多处运行”的自动内核生成。
3. 为智能体赋能：提供了一个结构化的、安全的界面，让AI智能体能够自主探索和优化整个模型的执行调度，开启了模型编译自动化的新可能。
4. 诚实的性能基线：研究没有回避与强大对手（如cuBLAS、vLLM）的差距，清晰地界定了当前技术的优势区间和局限性，为社区提供了可靠的参考。

未来，随着在GEMV内核质量上的持续投入（如利用Tensor Core、设计更粗粒度的同步机制），AMK有望在更广泛的硬件和场景下缩小乃至消除与手工优化内核的性能差距。而其底层“构造即正确”的架构，确保了这些未来的性能提升能够被安全、自动地整合进来。

对于从事大模型推理优化、编译器设计以及自主智能体研究的开发者和研究者来说，AutoMegaKernel不仅是一个实用的工具，更是一个展示了如何将形式化验证思想与高性能计算、机器学习系统相结合的有趣范本。

> 本文由 Intern-S2 等 AI 生成，机智流编辑部校对

-- 完 --

机智流推荐阅读：

1. ‍

2. 

3. 

4.