摘要
本文介绍了 KernelFalcon，这是一种用于生成 GPU kernel的深度Agent架构，它通过把kernel生成任务分解成多层pipeline而达到更高的正确性。KernelFalcon 是已知第一个在所有 250 个 L1/L2/L3 KernelBench 任务中实现 100% 正确性的开放Agent系统。
KernelFalcon 的代码库位于 github.com/meta-pytorch/KernelAgent，其中包含文档和入门示例。
引言
编写优化的 GPU kernel仍是部署机器学习模型的瓶颈。团队很少有足够的时间为每个形状、数据类型和硬件手动调整算子。随着模型的演进，这个问题进一步加剧——适用于 ResNet 的模式无法直接映射到 Mamba 的选择性状态或 MoE 的条件路由。
现代编译器已取得显著进展，但仍难以处理长尾问题。TorchInductor 覆盖了常见模式，TVM 自动调度密集内核，XLA 则专为动态形状优化。但不寻常的算子、动态控制流和异构融合模式仍难以实现最优编译。NVIDIA 的早期工作 使用 DeepSeek-R1 和推理时缩放在 KernelBench L1/L2 上取得了出色结果，证明了基于 LLM 的方法结合验证循环可以匹配或超越传统方法——但并未处理完整的模型架构（L3）。
如果我们能够自动合成 Triton kernel，同时保留 PyTorch 语义并接近手动优化的性能，而无需扩展lib或雇用更多 GPU 专家，那将非常完美？
KernelFalcon 应运而生：这是一个代码到代码的系统，在生成优化的 Triton 内核的同时保留 PyTorch 语义。它不采用一次性生成，而是使用并行探索结合基于执行的验证——交付的内核能够在 GPU 上实际运行，并与原始模型的数值结果匹配。
为什么选择 KernelFalcon（以及为什么使用深度Agent）？
传统的静态基于图的编译器依赖于 IR 转换和每个模式的调度。追踪通常将控制流固定为单一路径，并难以处理动态形状。KernelFalcon 采取了不同的路径：
-
保留 Python 语义。我们保持在 PyTorch 的代码到代码层面，因此 if/else、while、数据依赖路由和动态形状仍有效。 -
优先验证循环。KernelAgent 编译并测试候选内核；失败反馈本地处理；我们在第一个数值正确的内核上提前退出。 -
端到端组合和验证。融合内核替换原始算子，随后进行全模型一致性检查以确保接受。
底层是一个深度代理架构——一个多阶段系统，旨在通过结构化问题来减少 LLM 失败模式：
-
显式任务分解将模糊目标转化为精确的、工具就绪的子问题。 -
确定性编排将控制逻辑保留在 Python 中，让 LLM 专注于认知。 -
具有早期停止的并行搜索高效探索多样化解决方案。 -
基于实际工具的执行在每一步使用真实编译器和硬件进行验证。 -
结构化状态持久化提示、日志和工件，以实现可审计性和恢复。
这不仅仅是一个更干净的实现；它是一种不同的范式。我们不再问“LLM 能否解决这个问题？”，而是问“我们如何塑造任务以使 LLM 更可能成功？”结果是更广泛的覆盖和更现实的性能——无需膨胀规则集或牺牲语义。
KernelFalcon 架构
图 1：KernelFalcon 的深度代理架构以 Orchestrator 为中心，协调整个工作流。Planning 处理任务分解和预算分配。Context Engineering 提供结构约束（模板、指南）。Sub-Agents 处理专化任务，包括提取融合边界、生成 Triton 内核、组合端到端模块以及执行验证。Persistent Memory 存储工件以便调试和恢复。Orchestrator 委托给专家，接收结构化错误反馈，并在整个执行过程中维护状态。
这一架构体现了深度Agent原则：
-
分层委托：Orchestrator 将高层任务（融合此模型）分解为精确的子问题（提取子图、生成内核、组合结果），并分配给专化代理。 -
确定性控制：规划和编排逻辑是显式的 Python 代码，而非 LLM 驱动——工作者生命周期、超时和成功条件均为程序化。 -
基于实际执行：每个代理均使用真实工具（Triton 编译器、PyTorch 参考、GPU 执行）进行验证，而非模拟或 LLM 判断的结果。 -
持久状态：所有中间结果、提示、日志和工件均持久化到磁盘，以实现可审计性、调试和跨会话恢复。 -
结构约束：上下文工程将规则编码为模板和策略，使正确性要求通过结构强制执行，而非依赖提示。
图 2：多阶段工作流图示，显示 PyTorch 输入流经 FuserAgent（创建可融合子图）、ExtractorAgent（生成 JSON 规范）、并行 KernelAgent 工作者（三个 Triton 框显示并发生成）以及 ComposerAgent（缝合为验证内核）。箭头显示阶段间数据流，并标注中间表示。
Pipeline详解
Pipeline由以下四个阶段组成：
-
FuserAgent – 保留 Python 语义的代码到代码融合。 -
ExtractorAgent – 形状推理和合约生成。 -
Dispatcher + KernelAgent – 协调具有验证的并行 Triton 内核合成。 -
ComposerAgent – 端到端集成和验证。
架构：逐阶段分解
阶段 1：FuserAgent – 代码到代码融合
传统编译器在融合分析期间将 PyTorch 降低为静态 IR，从而丢失信息，导致调试困难并在动态控制流上失效。FuserAgent 直接在 PyTorch 源代码上操作。Orchestrator 管理融合工作流，生成具有显式子图边界的干净 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)
过程：
-
解析和分析：提取操作序列、数据依赖和控制流边界。 -
识别融合机会：找到可融合的操作组，同时保留语义。 -
生成融合模块：创建具有显式测试的干净 PyTorch 函数。 -
逐步验证：在继续之前独立测试每个融合子图。
输出：具有子图函数的融合 PyTorch 模块
# 保留控制流的融合模块
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: # 控制流完整保留
x = self.branch(x)
return self.norm(x)
这么做为什么有效：Orchestrator 生成精确规范，下游阶段可据此执行。控制流（if x.sum() > 0）保留在 Python 中——我们从不尝试编译它。
通过保持在 Python 源代码级别，我们保留变量名、注释和完整控制流上下文。大多数传统编译器风格的融合器假设它们在优化静态数据流图，因此动态 Python 侧控制流要么被折叠，要么必须在融合前重写。
因此，与我们基于提示的方法不同，后者保留 Python if 并仅在其中插入融合子模块，传统基于编译器的融合往往在追踪期间专化为单一分支，或需要大量手动工作来显式编码控制流。当 TorchScript 降低为 SSA 形式时，您精心命名的 hidden_states 变为 t0。当 torch.fx 通过条件追踪时，未取分支简单消失。即使使用 TorchDynamo/torch.compile，虽然它通过图断点和守卫更好地处理控制流，但它仍为观察路径专化图——您的 if x.sum() > 0 变为守卫检查，要么重用缓存图，要么触发重新编译。FuserAgent 采取不同方法：我们保留 Python if 语句，但融合每个分支内的操作。您仍获得内核融合益处（每个分支内的算子成为优化的 Triton 内核），但控制流本身保留为可读 Python。
这对现代 ML 模式至关重要，这些模式如今无处不在：TreeLSTM 在解析树上的递归、早期退出网络在自信时退出、Mixture-of-Experts 路由到不同子网络。而且，当调试出错时——当您的内核产生 NaN 或融合失败时——您希望阅读 Python，而非 IR。您希望看到系统实际尝试融合的内容，以您编写它的语言。
深度代理原则：确定性控制平面
所有编排——工作者生命周期、超时、工件路径和成功提前退出——均在 Python 中实现。LLM 生成候选代码和元数据（融合模块、子图 JSON、Triton 内核、组合内核）；控制器通过执行验证输出，而非 LLM。
工作流：
-
Orchestrator 启动 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 模块
提取过程：
-
运行 Orchestrator：首先从阶段 1 获取融合代码。 -
提示 LLM：要求 LLM 识别不同的子图函数、推理形状并目录操作。 -
生成 JSON 规范：LLM 产生具有操作序列、形状和权重元数据的类型化规范。 -
去重和合并：通过稳定签名（操作 + 形状 + 权重）分组子图，聚合计数。
输出：子图规范的 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 的合约——显式、类型化和可验证。每个子图包括：
-
具有操作特定参数的操作序列。 -
输入和输出的形状合约。 -
跟踪融合和原始参数的权重元数据。 -
位置信息（模型中的位置、源模块）。 -
去重子图的计数。
Orchestrator 控制工作流；LLM 生成形状感知元数据；去重处理模型中的相同模式。
参考：Fuser/subgraph_extractor.py
阶段 3：Dispatcher + KernelAgent – 并行 Triton 生成
Dispatcher 为每个子图规范协调并行 Triton 内核生成。对于每个子图，它创建一个新的 TritonKernelAgent，并配备工作者池（默认 4 个工作者）。
图 3：KernelAgent 启动并行工作者，使用多样化采样参数生成 Triton 内核。每个候选通过验证阶段（语法、编译、数值）。失败候选仅触发其原始工作者的隔离错误反馈——无上下文污染。第一个通过所有阶段的候选立即部署并取消剩余工作者。这实现了具有隔离上下文和早期退出的并行探索。
并行方法：
使用相同提示但多样化温度设置（0.8、0.9、1.0 等）生成 N 个内核种子。启动 N 个工作者（默认 4），每个在自己的工作目录中运行隔离精炼循环。不同的温度引导工作者探索不同的优化策略——有些保守，有些探索性。
关键机制：
-
本地错误反馈防止上下文污染
每个工作者维护自己的工作目录和每轮历史。当工作者 2 遇到编译错误时，仅工作者 2 的下一迭代看到它——错误上下文保持本地。工作者将 kernel.py 和 test_kernel.py 写入自己的工作目录，通过子进程执行测试，并独立跟踪结果。其他工作者继续使用干净上下文。
-
早期终止节省计算
集中式管理器监控结果队列以捕获完成事件。一旦任何工作者报告成功（测试子进程退出码 0），管理器设置共享成功事件以信号所有工作者停止，然后加入/终止它们。第一个通过所有验证阶段的内核获胜；剩余工作者立即终止。
深度代理原则：基于实际工具的使用
工作者在隔离子进程中执行真实的 Python/Triton 代码。每个工作者生成 Triton 内核实现及其验证工具，然后作为独立子进程运行验证。Triton 的 JIT 编译器在测试工具首次调用时自动将内核编译为 PTX——编译在测试执行期间隐式发生，因此任何语法或编译错误表现为具有非零退出码的测试失败。
验证工具将内核输出与 PyTorch 参考实现比较。成功意味着子进程退出码 0；失败捕获 stderr 以供下一轮精炼。框架不判断正确性——它简单执行代码并报告发生的情况。这种基于实际执行的 grounding 消除了“模拟判断”问题，其中 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 和原始问题
组合过程：
-
提示 LLM：提供原始问题代码、紧凑子图摘要和成功内核文件。 -
生成集成：LLM 合成具有所需结构的端到端 Triton 实现。 -
可选验证：执行组合内核并通过 PASS/哨兵检测验证。 -
打包工件：将组合实现和验证元数据写入输出目录。
生成结构： LLM 产生一个完整的 Python 模块，包含：一个或多个 Triton 内核：每个用 @triton.jit 装饰，实现融合操作。例如，一个内核可能处理 conv-bn-tanh-pool 融合，而另一个处理归一化。
顶层包装函数： 名为 kernel_function(...)，匹配原始模型的输入。此包装分配输出张量、配置网格维度，并按序启动 Triton 内核，协调它们之间的数据流。
自测试工具： 测试函数种子随机数生成器、构建原始 PyTorch 参考、调用组合内核函数，并使用 torch.allclose 验证一致性，容差来自提示指导。这些数据类型特定的容差考虑了每个精度级别的舍入误差累积，匹配 PyTorch 自身的内部测试标准。成功时，它打印 “PASS” 并以退出码 0 退出。
验证过程： Composer 确保个别正确的内核正确组合——验证整体等于部分之和。Python 通过作为子进程执行组合模块并检查 stdout 中的 “PASS” 以及退出码 0 来验证。这将验证 grounding 在实际执行中，而非模拟或 LLM 判断的正确性。
输出工具: 成功记录验证状态、计时和工件路径。完整的组合模块成为最终交付物，准备好部署或进一步使用。
参考：Fuser/compose_end_to_end.py
深度Agent原则在实践中的应用
合约和接口启用自动化
核心洞见：严格合约解锁自主操作。当每个交互点具有类型化接口——操作类型、形状、数据类型、预期输出——系统可无人类监督运行。模糊性会扼杀自动化。
我们处处编码这些合约。子图规范从 extractor 流向 dispatcher 再到个别内核代理，每个都知道下一个期望什么。关键 Triton 约束烘焙到模板中，而非希望 LLM 记住。这将正确性从“希望提示有效”转向“结构上不可能违反”。
魔力发生在组合这些合约时。通过个别测试的内核仍需与其他正确组合。最终验证不仅仅是“此内核有效吗？”，而是“整个管道是否产生与原始 PyTorch 相同的输出？”这种端到端合约使系统值得信赖。
持久状态和内存
每个尝试均持久化到磁盘以便调试。我们保留日志、提示、生成代码、错误追踪和计时数据。每个工作者写入每轮快照，捕获其本地演进。Orchestrator 实时流式传输事件。这种鲁棒日志级别是运行 250 个任务过夜并理解每个任务发生情况的稳定性所需。
更重要的是，持久性启用调试。失败尝试不会丢弃——工程师可通过日志 grep 识别模式并相应调整系统。目前学习是手动的（分析日志、更新提示），尽管路线图包括未来检索成功模式。
通过隔离实现可靠性
生产系统需要边界。每个工作者在其自己的目录中运行，在子进程隔离中执行测试，并维护本地错误上下文。当工作者 2 遇到编译错误时，该错误保持本地于工作者 2 的精炼循环。工作者 0、1 和 3 继续使用干净上下文探索。
这种隔离超出错误处理。资源限制防止失控进程。超时捕获无限循环。清理处理程序确保无泄漏 GPU 内存。这些不是后期添加的功能——它们是使系统大规模工作的基础。
隔离的美妙之处：它使调试易于处理。当某事失败时，您确切知道哪个工作者、哪一轮、哪个提示导致该失败。错误未通过共享状态级联或污染其他尝试。您可重放该确切序列并理解出错原因。
无压倒性的可观测性
系统生成大量数据——提示、代码、错误、计时。诀窍是使这些数据有用而非淹没其中。我们使用结构化日志（JSONL 用于流式传输，JSON 用于快照），既人类可读又机器可解析。
您可 grep 特定事件、聚合错误类型、追踪单个工作者的历程，或计算不同温度设置下的成功率。日志讲述故事：系统如何探索解决方案空间、何有效、何无效以及为什么。
批量摘要提供 30,000 英尺视图——通过率、平均轮次、工作者利用率。详细日志在需要时提供显微镜。这种分层可观测性意味着您可快速评估整体健康，然后深入特定失败而不会迷失在噪声中。
实际有帮助的自评估
精炼并非关于通用“再试一次”提示。当工作者测试失败时，实际错误输出——编译消息、数值不匹配、超时指示——成为下一个提示的一部分。这种结构化反馈引导 LLM 向特定修复而非随机探索。
框架不抽象判断正确性。它运行代码、捕获发生情况，并反馈结果。关于张量设备的编译错误得到不同于输出数值不匹配的处理。工作者维护本地历史，因此每个精炼建立在先前尝试之上，在每个工作者上下文中创建学习轨迹。
关键是保持反馈具体且可操作。而非“内核无效”，工作者看到“编译失败于第 23 行：张量维度不匹配。”而非“错误输出”，它看到“最大误差：0.0023 于位置 [15, 7]。”这种具体性将调试从猜测转为针对性问题解决。
结论
覆盖率
我们简单测量成功：生成的代码是否实际有效？当 Triton 内核编译、无崩溃运行并在有界容差内产生数值等价于 PyTorch 的输出时，任务通过。
|
|
|
|
|
|---|---|---|---|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
性能指标（加速、延迟）将在后续文章中涵盖。本工作首先确立正确性。
准确性验证
-
每个工作者将 Triton 内核及其测试工具写入磁盘，然后在新鲜工作目录中的子进程中运行测试。对于个别内核，退出码 0 意味着成功。 -
对于完整管道，我们需要退出码 0 且输出中 “PASS”。当测试失败时，我们捕获 stderr 并再试。测试使用适合数据类型的容差与 PyTorch 比较——fp32 更紧，fp16 更松。
如何防止作弊？
强制通过 Triton 的编译器：@triton.jit 函数不能包含 PyTorch 操作——它们会编译失败。这防止内核中使用高级 API。我们也添加包装侧检查。对于测试，我们依赖随机输入和子进程执行。然而，我们信任 LLM 生成的测试工具本身——我们不静态分析它以防作弊。
尾声
验证胜过一次性生成，但架构是乘法。此系统将可执行反馈连接到每个阶段，并让第一个通过路径结束搜索。多个工作者并行探索、执行候选，并在其中一个通过时取消同行。更少上下文消耗、更低延迟、更强鲁棒性。数字重要，但循环更重要：执行、读取错误、修复、在某物有效时停止。
这不仅仅关于内核生成。它关于当您停止问“LLM 能做这个吗？”并开始问“我们如何结构化这个以使 LLM 无法失败？”时发生什么。KernelFalcon 在 KernelBench 上实现 100% 正确性，不是因为我们有更好的LLM模型，而是因为我们使用LLM构建了更好的Pipeline。
结构化问题，而非更努力提示。
