博客

KernelFalcon:通过深度代理实现自主 GPU 内核生成

总结

我们推出了 KernelFalcon,这是一种用于生成 GPU 内核的深度智能体架构。它结合了层级任务分解与委派、带有“抢先获胜”并行搜索的确定性控制平面、扎实的工具使用以及持久化内存/可观测性。KernelFalcon 是首个已知的开源智能体系统,在全部 250 个 L1/L2/L3 KernelBench 任务中实现了 100% 的正确率。

KernelFalcon 的代码库位于 github.com/meta-pytorch/KernelAgent,随附入门文档和示例。

引言

编写优化的 GPU 内核仍然是部署机器学习模型的瓶颈。团队几乎没有精力为每种形状(shape)、数据类型(dtype)和硬件一代手动调优算子。随着模型的演进,这个问题变得更加复杂——适用于 ResNet 的模式无法直接映射到 Mamba 的选择性状态(selective states)或 MoE 的条件路由中。

现代编译器虽然取得了真正的进步,但在处理“长尾”问题时依然吃力。TorchInductor 可以覆盖常见模式,TVM 可自动调度密集算子内核,XLA 擅长处理动态形状。然而,非典型算子、动态控制流和异构融合模式往往无法实现最优编译。NVIDIA 今年早些时候的工作利用 DeepSeek-R1 和推理时扩展(inference-time scaling),在 KernelBench L1/L2 上取得了显著成果,证明了基于 LLM 并带有验证循环的方法可以匹配甚至超越传统方法——但尚未触及完整的模型架构(L3)。

如果我们能自动合成既保留 PyTorch 语义又接近手写性能的 Triton 内核,而无需扩展规则库或雇佣更多的 GPU 专家,那会怎样?

KernelFalcon 应运而生:这是一个代码到代码(code-to-code)的系统,在生成优化 Triton 内核的同时保留了 PyTorch 语义。它不再采用一次性生成,而是使用基于执行验证的并行探索,从而交付能够在 GPU 上实际运行并匹配原始模型数值特性的内核。

为何选择 KernelFalcon(以及为何选择深度智能体)?

传统的基于静态图的编译器依赖于 IR 转换和基于模式的调度。跟踪(Tracing)通常会将控制流冻结在单一路径上,且难以处理动态形状。KernelFalcon 采取了不同的路径:

保留 Python 语义。 我们在 PyTorch 中保持代码到代码的转换,因此 if/elsewhile、数据依赖路由和动态形状依然有效。

验证优先循环。 KernelAgent 编译并测试候选内核;失败反馈会在本地处理;我们会在得到第一个数值正确的内核时立即停止搜索并退出。

端到端组合与验证。 融合后的内核可替代原始算子,并在被采纳前进行全模型一致性检查。

其核心是一个深度智能体架构——这是一个多阶段系统,旨在通过结构化问题来减少 LLM 的失效模式:

  • 显式任务分解:将模糊的目标转化为精确的、工具可执行的子问题。
  • 确定性编排:将控制逻辑保留在 Python 中,让 LLM 专注于认知任务。
  • 带早期停止的并行搜索:高效地探索多种解决方案。
  • 扎实的工具使用:通过真实编译器和硬件验证每一步。
  • 结构化状态:持久化提示词、日志和产物,以便审计和恢复。

这不仅仅是一个更简洁的实现,更是一种不同的范式。我们不再问“LLM 能解决这个问题吗?”,而是问“我们该如何塑造任务,使 LLM 更容易成功?”其结果是更广泛的覆盖范围和更真实的性能表现——无需扩充庞大的规则集,也不必牺牲语义。

KernelFalcon 架构

图 1: KernelFalcon 的深度智能体架构以协调整个工作流的编排器(Orchestrator)为中心。规划器处理任务分解和预算分配。上下文工程提供结构性约束(模板、指南)。子智能体处理专业化任务,提取融合边界、生成 Triton 内核、组合端到端模块并执行验证。持久化内存存储产物以便调试和恢复。编排器将任务委派给专家,接收结构化的错误反馈,并在整个执行过程中维护状态。

该架构体现了深度智能体的原则:

  1. 层级委派: 编排器将高级任务(如:融合该模型)分解为精确的子问题(提取子图、生成内核、组合结果),并分配给对应的专业智能体。
  2. 确定性控制: 规划和编排逻辑是显式的 Python 代码,而非由 LLM 驱动——工作生命周期、超时控制和成功条件均以程序化方式设定。
  3. 扎实执行: 每个智能体均使用真实工具(Triton 编译器、PyTorch 参考实现、GPU 执行)进行验证,而非基于模拟或 LLM 判断的结果。
  4. 持久化状态: 所有中间结果、提示词、日志和产物均持久化至磁盘,以实现可审计性、调试和跨会话恢复。
  5. 结构性约束: 上下文工程将规则编码为模板和策略,使正确性要求在结构上得到强制执行,而非依赖提示词。

图 2: 多阶段工作流图,显示 PyTorch 输入流经 FuserAgent(创建可融合子图)、ExtractorAgent(生成 JSON 规范)、并行 KernelAgent 工人(三个展示并发生成的 Triton 方框)和 ComposerAgent(缝合生成已验证内核)。箭头指示了阶段之间的数据流,并标注了中间表示。

流水线:数据如何在系统中流动

该流水线实现了四个不同的阶段:

  1. FuserAgent – 保留 Python 语义的代码到代码融合
  2. ExtractorAgent – 形状推断与契约生成
  3. Dispatcher + KernelAgent – 协调并行 Triton 内核合成与验证
  4. ComposerAgent – 端到端集成与验证

架构:阶段性分解

阶段 1:FuserAgent – 代码到代码融合

传统编译器在融合分析期间将 PyTorch 降低为静态 IR,导致信息丢失,使得调试变得困难,并在遇到动态控制流时会中断。FuserAgent 直接在 PyTorch 源代码上操作。编排器管理融合工作流,生成带有明确子图边界的干净 PyTorch 模块。

输入: 具有任意复杂度的原始 PyTorch 模型

class Model(nn.Module):
    def forward(self, x):
        if x.sum() > 0: 
            x = self.conv(x)
            x = self.bn(x)
            x = torch.tanh(x)
            x = F.max_pool2d(x, 2)
        return self.norm(x)

处理过程:

  1. 解析与分析: 提取操作序列、数据依赖和控制流边界。
  2. 识别融合机会: 寻找在保留语义的前提下可以融合的操作组合。
  3. 生成融合模块: 创建带有显式测试的干净 PyTorch 函数。
  4. 增量验证: 在继续之前独立测试每个融合的子图。

输出: 带有子图函数的融合 PyTorch 模块。

# Fused module with control flow preserved
class FusedModel(nn.Module):
    def __init__(self, channels: int):
        super().__init__()
        self.branch = ConvBnTanhMaxPool(channels=channels)
        self.norm = ChannelwiseNorm(channels=channels)
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if x.sum() > 0:  # Control flow intact
            x = self.branch(x)
        return self.norm(x)

为何有效: 编排器生成了下游阶段可以执行的精确规范。控制流(if x.sum() > 0)保留在 Python 中——我们从不试图将其编译掉。

通过保持在 Python 源码级别,我们保留了变量名、注释和完整的控制流上下文。大多数传统的类编译器融合器假设它们正在优化一个静态数据流图,因此动态的 Python 端控制流要么被折叠,要么必须在融合前被重写。

因此,与我们这种保持 Python if 并仅在其中插入融合子模块的提示驱动方法不同,传统的基于编译器的融合倾向于在跟踪过程中特化为单一分支,或者需要大量的人工努力来显式编码控制流。当 TorchScript 降低为 SSA 形式时,你仔细命名的 hidden_states 会变成 t0。当 torch.fx 跟踪条件语句时,未执行的分支直接消失了。即使使用 TorchDynamo/torch.compile,虽然它通过图断点(graph breaks)和守卫(guards)更好地处理了控制流,但它仍然针对观察到的路径特化图——你的 if x.sum() > 0 会变成一个守卫检查,要么重用缓存的图,要么触发重新编译。FuserAgent 采取了不同的方法:我们保留 Python 的 if 语句,但融合每个分支内的操作。你依然能获得内核融合的好处(分支内的操作变成优化后的 Triton 内核),但控制流本身依然是可读的 Python。

这对于现在随处可见的现代 ML 模式至关重要:递归遍历解析树的 TreeLSTM、在有把握时提前退出的早期退出网络、路由到不同子网络的混合专家模型(MoE)。更重要的是,当调试出错时——当你的内核产生 NaN 或融合失败时——你想要读的是 Python,而不是 IR。你想要在你自己编写的语言中看到系统到底试图融合什么。

深度智能体原则:确定性控制平面

所有编排工作——工人生命周期、超时设置、产物路径、成功时的提前退出——均在 Python 中实现。LLM 生成候选代码和元数据(融合模块、子图 JSON、Triton 内核、组合内核);控制器执行并验证输出,而不是由 LLM 执行。

工作流:

  • 编排器 启动 N 个使用类型化 WorkerConfig 的工作进程,流式传输日志,在队列中等待获胜者,取消其他进程,并打包产物。
  • 工作进程 进行迭代:渲染提示词 → 流式传输给 LLM → 提取 Python 代码块 → 按 SHA 去重 → 执行候选代码 → 如果通过,发送获胜信号;否则保存错误并重试。
  • 无需人工 AST 解析或基于规则的融合检测——LLM 直接通过提示词提出融合代码,然后 Python 通过执行进行验证。

参考: Fuser/orchestrator.py, Fuser/worker.py, Fuser/runner.py, Fuser/prompting.py

阶段 2:ExtractorAgent – 子图边界推断

Extractor 使用 LLM 来分析融合代码,并识别带有形状契约的精确子图边界。

输入: 来自阶段 1 的融合 PyTorch 模块。

提取过程:

  1. 运行编排器: 首先,获取阶段 1 的融合代码。
  2. 提示 LLM: 要求 LLM 识别不同的子图函数、推断形状并编目操作。
  3. 生成 JSON 规范: LLM 生成带有操作序列、形状和权重元数据的类型化规范。
  4. 去重与合并: 按稳定签名(操作 + 形状 + 权重)对子图进行分组,并汇总计数。

输出: 子图规范的 JSON 数组。

[
  {
    "id": "sg_conv_bn_tanh_pool_1",
    "type": "Conv2d_BN_Tanh_MaxPool",
    "data_layout": "NCHW",
    "dtype": "float32",
    "ops": [
      {"op": "conv2d", "kernel_size": [3, 3], "stride": [1, 1], "padding": [1, 
1], "dilation": [1, 1], "groups": 1, "bias": false},
      {"op": "batch_norm", "eps": 1e-5, "momentum": 0.1},
      {"op": "tanh"},
      {"op": "max_pool2d", "kernel_size": [2, 2], "stride": [2, 2]}
    ],
    "input_shape": ["B", "C_in", "H", "W"],
    "output_shape": ["B", "C_out", "H_out", "W_out"],
    "weights_original": {
      "conv.weight": ["C_out", "C_in", 3, 3],
      "batch_norm.weight": ["C_out"],
      "batch_norm.bias": ["C_out"],
      "running_mean": ["C_out"],
      "running_var": ["C_out"]
    },
    "weights_fused": null,
    "count": 1,
    "where": "Model.forward conditional branch",
    "source": {
      "module": "FusedConvBnTanhPool",
      "code": "def forward(self, x):\n    x = F.conv2d(x, self.conv_w, 
stride=1, padding=1)\n    x = F.batch_norm(x, self.bn_rm, self.bn_rv, self.bn_w, 
self.bn_b, training=False, eps=self.eps)\n    x = torch.tanh(x)\n    
return F.max_pool2d(x, 2)"
    }
  }
]

这个 JSON 成为 KernelAgent 的契约——显式、类型化且可验证。每个子图包含:

  • 操作序列 以及特定于操作的参数。
  • 形状契约(输入和输出)。
  • 权重元数据(跟踪融合的和原始的参数)。
  • 位置信息(在模型中的位置、源模块)。
  • 计数(用于去重的子图)。

编排器控制工作流;LLM 生成形状感知的元数据;去重处理整个模型中相同的模式。

参考: Fuser/subgraph_extractor.py

阶段 3:Dispatcher + KernelAgent – 并行 Triton 生成

Dispatcher 为每个子图规范协调并行 Triton 内核生成。对于每个子图,它会创建一个带有工人池(默认 4 个工人)的全新 TritonKernelAgent。

图 3: KernelAgent 启动并行工作进程,使用不同的采样参数生成 Triton 内核。每个候选者经历验证阶段(语法、编译、数值)。失败的候选者只会向其所属的特定工作进程触发隔离的错误反馈——不会污染上下文。第一个通过所有阶段的候选者立即部署并取消剩余工作进程。这实现了带有隔离上下文和提前退出的并行探索。

并行方法:

使用相同的提示词但不同的温度设置(0.8, 0.9, 1.0 等)生成 N 个内核种子。启动 N 个工作进程(默认 4 个),每个进程在自己的工作目录中运行隔离的细化循环。不同的温度引导工作进程探索不同的优化策略——有的保守,有的更具探索性。

关键机制:

1. 本地错误反馈防止上下文污染

每个工作进程维护自己的工作目录和每轮历史记录。当工作进程 2 遇到编译错误时,只有它下一次迭代会看到该错误——错误上下文保持在本地。工作进程将 kernel.pytest_kernel.py 写入各自的工作目录,通过子进程执行测试,并独立跟踪结果。其他工作进程在干净的上下文中继续运行。

2. 提前终止节省计算资源

集中式管理器监视结果队列中的完成事件。一旦任何工作进程报告成功(测试子进程退出码为 0),管理器设置共享成功事件,通知所有工作进程停止,然后汇总/终止它们。第一个通过所有验证阶段的内核获胜;剩余工作进程立即被终止。

深度智能体原则:扎实的工具使用

工作进程在隔离的子进程中执行真实的 Python/Triton 代码。每个工作进程生成 Triton 内核实现及其验证工具,然后作为独立子进程运行验证。Triton 的 JIT 编译器在测试工具首次调用内核时将其自动编译为 PTX——编译在测试执行期间隐式发生,因此任何语法或编译错误都会以非零退出码的测试失败形式呈现。

验证工具将内核输出与 PyTorch 参考实现进行比较。成功意味着子进程返回退出码 0;失败则捕获 stderr 以用于下一次细化轮次。框架本身不评判正确性——它只是执行代码并报告发生了什么。这种对实际执行的扎根消除了“模拟评审”问题,即 LLM 可能会幻觉出错误代码能够运行的错觉。

参考: Fuser/dispatch_kernel_agent.py, triton_kernel_agent/manager.py, triton_kernel_agent/worker.py, triton_kernel_agent/agent.py

阶段 4:ComposerAgent – 端到端内核缝合

Composer 使用 LLM 接收已验证的 Triton 内核,并将它们集成到一个完整、可测试的模块中。

输入: 来自阶段 3 的已验证 Triton 内核集合、subgraphs.json 以及原始问题。

组合过程:

  1. 提示 LLM: 提供原始问题代码、紧凑的子图摘要以及成功的内核文件。
  2. 生成集成: LLM 合成具有所需结构的端到端 Triton 实现。
  3. 可选验证: 执行组合后的内核,并通过 PASS/哨兵检测进行验证。
  4. 打包产物: 将组合后的实现和验证元数据写入输出目录。

生成结构:

LLM 生成一个完整的 Python 模块,包含:

一个或多个 Triton 内核: 每个都带有 @triton.jit 修饰符,实现融合操作。例如,一个内核可能处理 conv-bn-tanh-pool 融合,而另一个处理归一化。

顶层包装函数: 名为 kernel_function(...),匹配原始模型的输入。该包装器分配输出张量、配置网格维度,并按顺序启动 Triton 内核,协调它们之间的数据流。

自测试工具: 测试函数播种随机数生成器,构建原始 PyTorch 参考,调用组合内核函数,并使用提示词指南中的容差执行 torch.allclose 来验证一致性。这些特定于数据类型的容差考虑了每个精度等级固有的舍入误差累积,匹配 PyTorch 内部测试标准。成功时,它打印“PASS”并返回退出码 0。

验证过程:

Composer 确保单个正确的内核能正确组合——验证整体等于部分之和。Python 通过将组合模块作为子进程执行并检查 stdout 中的“PASS”以及退出码 0 来进行验证。这使得验证扎根于实际执行,而非模拟或 LLM 评审的正确性。

输出产物:

成功记录包含验证状态、耗时和产物路径。完整的组合模块成为最终交付物,准备部署或进一步处理。

参考: Fuser/compose_end_to_end.py

深度智能体原则在实践中

契约与接口促成自动化

核心洞察:严格的契约解锁了自主操作。当每个交互点都有一个类型化接口——操作类型、形状、数据类型、期望输出——系统无需人工干预即可运行。模糊性是自动化的死敌。

我们将这些契约编码在每一处。子图规范从提取器流向调度器,再到各个内核智能体,每个阶段都知道下一阶段预期的是什么。关键的 Triton 约束被固化到模板中,而不是寄希望于 LLM 能记住它们。这使得正确性从“希望提示词有效”转变为“在结构上不可能违规”。

当组合这些契约时,奇迹便发生了。一个通过了自身测试的内核仍需与其他部分正确组合。最终验证不仅仅是“这个内核有效吗?”,而是“整个流水线产生的结果是否与原始 PyTorch 一致?”这种端到端契约使得系统值得信赖。

持久状态与记忆

每次尝试都持久化到磁盘以供调试。我们保留日志、提示词、生成的代码、错误跟踪和耗时数据。每个工作进程写入每轮快照,捕获其本地演变。编排器实时流式传输事件。这种鲁棒的日志记录水平是通宵运行 250 个任务并了解每个任务状况的必要前提。

更重要的是,持久性使调试成为可能。失败的尝试不会被丢弃——工程师可以 grep 日志来识别模式并相应地调整系统。目前学习过程是手动的(分析日志、更新提示词),尽管未来的路线图包括对成功模式的自动检索。

通过隔离实现可靠性

生产系统需要边界。每个工作进程在自己的目录中运行,在子进程隔离中执行测试,并维护本地错误上下文。当工作进程 2 遇到编译错误时,该错误仅限于工作进程 2 的细化循环。工作进程 0、1 和 3 在干净的上下文中继续探索。

这种隔离扩展到了错误处理之外。资源限制防止失控进程。超时捕获死循环。清理处理程序确保没有 GPU 内存泄漏。这些不是后来才添加的功能——它们是使系统能够大规模工作的根本保障。

隔离的妙处在于:它使调试变得可控。当某事失败时,你准确地知道是哪个工作进程、哪一轮、哪个提示词导致了该故障。错误不会通过共享状态蔓延或污染其他尝试。你可以重放该确切序列并理解出了什么问题。

不过度干扰的可观测性

系统会产生大量数据——提示词、代码、错误、耗时。诀窍是在不被淹没的情况下使数据变得有用。我们使用结构化日志(用于流式传输的 JSONL,用于快照的 JSON),这些日志既可由人类阅读,也可由机器解析。

你可以 grep 特定事件、汇总错误类型、跟踪单个工作进程的历程,或计算不同温度设置下的成功率。日志讲述了一个故事:系统如何探索解决方案空间,什么有效,什么无效,以及原因。

批次摘要提供了宏观视角——成功率、平均轮次、工作进程利用率。详细日志则在你需要时提供了显微镜。这种分层可观测性意味着你可以快速评估整体健康状况,然后在不被噪音干扰的情况下深入排查特定故障。

真正有效的自我评估

细化不仅仅是通用的“再试一次”提示词。当工作进程的测试失败时,实际的错误输出——编译信息、数值不匹配、超时指标——会成为下一次提示词的一部分。这种结构化反馈引导 LLM 针对性地进行修复,而不是随机探索。

框架不会抽象地评判正确性。它运行代码,捕获结果,并将现实反馈回去。关于张量设备的编译错误与输出中的数值不匹配会得到不同的处理。工作进程维护本地历史,因此每次细化都建立在先前尝试的基础上,在每个工作进程的上下文中创建学习轨迹。

关键在于保持反馈的具体性和可操作性。工作进程看到的不是“内核无效”,而是“第 23 行编译失败:张量维度不匹配”。它看到的不是“错误输出”,而是“位置 [15, 7] 处最大误差:0.0023”。这种特殊性将调试从猜测变成了有针对性的问题解决。

结果

覆盖范围

我们简单地衡量成功:生成的代码真的有效吗?当 Triton 内核编译成功、运行不崩溃,并在限定的容差内产生与 PyTorch 数值等效的结果时,任务即为通过。

级别 任务 描述 覆盖范围
L1 100 单一算子 ✅ 100%
L2 100 融合模式 ✅ 100%
L3 50 完整架构 ✅ 100%

性能指标(加速比、延迟)将在我们后续文章中涵盖。此工作首先建立正确性。

我们如何验证

没有模拟。没有“LLM 认为这有效”。我们执行真正的代码。

两个级别的验证:

每个工作进程都会将 Triton 内核及其测试工具写入磁盘,然后在全新的工作目录中的子进程内运行测试。对于单个内核,退出码 0 意味着成功。对于完整流水线,我们需要退出码 0 且输出中有“PASS”。当测试失败时,我们捕获 stderr 并重试。测试根据数据类型使用适当的容差与 PyTorch 进行比较——fp32 更严格,fp16 更宽松。

什么能防止作弊?

强制执行主要通过 Triton 编译器实现:@triton.jit 函数不能包含 PyTorch 操作——否则编译将失败。这防止了在内核中使用高级 API。我们也添加了包装器端的检查。对于测试,我们依赖随机输入和子进程执行。然而,我们信任 LLM 生成的测试工具本身——我们不对其进行静态分析以查找作弊行为。

组合验证至关重要:单个内核的正确性不能保证整个流水线工作。集成会暴露单个测试漏掉的形状不匹配、累积误差和语义误解。

每个人都在学习同样的教训

验证胜过一次性生成,但架构是乘数。该系统将可执行反馈连接到每个阶段,并让第一个通过的路径结束搜索。多个工作进程并行探索、执行候选方案,并在一个通过的瞬间取消同伴。减少了上下文消耗,降低了延迟,提高了鲁棒性。数字很重要,但循环更重要:执行,阅读错误,修复,一旦成功立即停止。

这不仅仅关于内核生成。这是关于当你停止问“LLM 能做到吗?”并开始问“我们如何构建它以使 LLM 无法失败?”时会发生什么。KernelFalcon 在 KernelBench 上实现了 100% 的正确率,不是因为我们有一个更好的模型,而是因为我们构建了一个更好的循环。

结构化问题,而不是更努力地提示。 那才是真正的教训。深度智能体不是关于更深的模型——它们是关于更深的系统。

致谢

此工作是与 Sijia Chen, Bert Maher, Joe Isaacson, Lu Fang, Wenyuan Chi, Jie Liu, Alec Hammond, Zacharias Fisches, Mark Saroufim, Warren Hunt, Richard Li, Jacob Kahn, Emad El-Haraty 和 Ajit Mathews 合作完成的。

引用

请将此工作引用为

Wang, Laura et al., "KernelFalcon: Deep Agent Architecture for Autonomous GPU Kernel Generation", PyTorch Blog, Nov 2025.