本文基于 KernelFalcon: Autonomous GPU Kernel Generation via Deep Agents 进行整理和拓展。
KernelFalcon 是PyTorch 提出的一个Deep Agents架构系统，该框架主要尝试利用Agent端到端 实现torch模型优化及 Triton算子 自动生成，是首个在全部 250 个 L1/L2/L3 KernelBench 任务上达到 100% 正确率的开源智能体系统。
KernelFalcon 代码库位于 github.com/meta-pytorch/KernelAgent，附带文档与入门示例。
KernelFalcon 面临的背景和挑战如下：
如果能在不扩充规则库或雇佣更多 GPU 专家的前提下，自动合成保持 PyTorch 语义且逼近手工调优性能的 Triton 内核，会怎样？
因此，PyTorch推出了 KernelFalcon：一个保持 PyTorch 语义的代码到代码系统，可生成优化的 Triton 内核。它采用并行探索与基于执行的验证，而非一次性生成——交付的内核真实运行在 GPU 上，且与原始模型数值等价。
注：长尾问题原本是指在数据分布中，少数数据（“头部”）的频率非常高，而大多数数据（“尾部”）的频率非常低的一种现象。在这里主要想表达的是AI、算子相关的工作中总是由于创新、软硬件变化等缘故，出现传统方案无法覆盖的场景，导致功能不支持或者性能下降。
KernelFalcon的核心目标是：将PyTorch 程序自动转换为经过验证的 GPU内核（kernel）。比如，它通过LLM生成 Triton 代码，并在沙盒子进程中自动验证正确性，迭代修复直到通过测试。
KernelFalcon基于 Pytorch Module 自动生成GPU Kernel，其设计理念如下：
项目分两个层次：
为何选择 KernelFalcon（为何是深度智能体）？
传统静态、基于图的编译器依赖 IR 变换和每模式调度。追踪常常将控制流冻结为单一路径，并在动态形状下失效。KernelFalcon 则走另一条路：
底层是一个深度智能体架构——多阶段系统，通过结构化问题来减少 LLM 失效模式：
这不仅是更干净的实现；这是不同范式。不再问“LLM 能否解决此问题？”，而是问“KernelFalcon 如何塑造任务，使 LLM 可能成功？”结果是更广覆盖与更现实性能——无需膨胀规则集或牺牲语义。
图 1：KernelFalcon 的深度智能体架构以 Orchestrator 为中心协调整个工作流。规划负责任务分解与预算分配。上下文工程提供结构约束（模板、指南）。子智能体处理专业任务，提取融合边界、生成 Triton 内核、组合端到端模块并执行验证。持久化内存存储产物用于调试与续跑。Orchestrator 委派给专家，接收结构化错误反馈，并在整个执行过程中保持状态。
该架构体现深度智能体原则：
图 2：多阶段工作流图显示 PyTorch 输入流经 FuserAgent（创建可融合子图）、ExtractorAgent（生成 JSON 规范）、并行 KernelAgent worker（三个 Triton 框显示并发生成）和 ComposerAgent（缝合已验证内核）。箭头表示数据流，并标注中间表示。
流水线包含四个不同阶段：
系统中的模块依赖关系如下：
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 模块，控制流保持完整
# 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)
优势: 保持Python语义，不丢失调试信息，支持动态控制流
为何有效？这是因为 Orchestrator 生成精确规范，下游阶段可执行验证。控制流（if x.sum() > 0）保留在 Python 中——KernelFalcon 从不尝试将其编译掉。
具体而言，通过停留在 Python 源码层，KernelFalcon 保留变量名、注释和完整控制流上下文。大多数传统编译器式融合器假设它们在优化静态数据流图，因此动态 Python 侧控制流要么在追踪期间被折叠，要么需要大量手动工作显式编码控制流：
因此，与KernelFalcon 保持 Python if 并在其中插入融合子模块的提示驱动方法不同，传统基于编译器的融合往往特化为追踪期间的单分支，或需要大量手动努力显式编码控制流。当 TorchScript 降级到 SSA 形式时，您精心命名的 hidden_states 变成 t0。当 torch.fx 追踪条件时，未采取的分支直接消失。即使有 TorchDynamo/torch.compile，虽然通过图中断和守卫更好地处理控制流，它仍然特化图为观察到的路径——您的 if x.sum() > 0 变成守卫检查，要么重用缓存图，要么触发重新编译。
FuserAgent 采取不同方法：KernelFalcon 保留 Python if 语句，但融合每个分支内的操作。您仍可获得内核融合收益（每个分支内的操作变成优化的 Triton 内核），但控制流本身保持可读的 Python。
这对现代 ML 模式至关重要：TreeLSTM 递归解析树、早退网络在自信时退出、混合专家路由到不同子网络。而且关键的是，当调试出错——当您的内核产生 NaN 或融合失败——您希望阅读 Python，而非 IR。您希望看到系统实际尝试融合的内容，用您编写的语言。
所有编排——worker生命周期、超时、产物路径和成功时早退——都用 Python 实现。LLM 生成候选代码与元数据（融合模块、子图 JSON、Triton 内核、组合内核）；控制器执行并验证输出。
工作流如下：
参考：Fuser/orchestrator.py、Fuser/worker.py、Fuser/runner.py、Fuser/prompting.py
本阶段是 提取器使用 LLM 分析融合代码，并识别具有形状合约的精确子图边界。
输入：来自阶段 1 的融合 PyTorch 模块
提取过程：
输出：子图规范的 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
Dispatcher 为每个子图规范协调并行 Triton 内核生成。对于每个子图，它创建一个带有worker池（默认 4 个worker）的全新 TritonKernelAgent。
图 3：KernelAgent 生成并行worker，采用多样采样参数生成 Triton 内核。每个候选者经历验证阶段（语法、编译、数值）。失败的候选者触发隔离错误反馈至其源worker——无上下文污染。首个通过所有阶段的候选者立即部署并取消剩余worker。实现并行探索与隔离上下文及早退。
用相同提示但不同温度设置（0.8、0.9、1.0 等）生成 N 个内核种子。生成 N 个worker（默认 4 个），每个在其自己的工作目录中运行隔离的精炼循环。不同温度导致worker探索不同的优化策略——有些保守，有些探索性。
每个worker维护自己的工作目录与每轮历史。当worker 2 遇到编译错误时，只有worker 2 的下一次迭代看到它——错误上下文保持本地。worker将 kernel.py 与 test_kernel.py 写入其自己的 workdir，通过子进程执行测试，并独立跟踪结果。其他worker继续以干净上下文运行。
集中管理器监控结果队列以获取完成事件。任何worker报告成功（测试子进程退出码 0）时，管理器设置共享成功事件以通知所有worker停止，然后加入/终止它们。首个通过所有验证阶段的内核获胜；剩余worker立即终止。
worker在隔离子进程中执行真实 Python/Triton 代码。每个worker生成 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
Composer 使用 LLM 获取已验证的 Triton 内核，并将其集成为完整、可测试的模块。
输入：来自阶段 3 的已验证 Triton 内核集合、subgraphs.json 与原始问题
组合过程：
生成结构：
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
KernelAgent 项目的入口点分为两个层面：
用户界面层面：通过命令行工具（pipeline.py 和 auto_agent.py）和 GUI 界面
编程接口层面：通过 TritonKernelAgent 和 AutoKernelRouter 类提供 API 访问
其中，auto_agent.py 是推荐的主要入口点，因为它会根据问题复杂度自动选择最优路径。
KernelAgent 项目提供了多个命令行入口点，主要通过 Fuser 模块访问：
管道执行入口：pipeline.py 是主要的端到端管道执行入口：
python -m Fuser.pipeline \
--problem /abs/path/to/problem.py \
--extract-model gpt-5 \
--dispatch-model o4-mini \
--compose-model o4-mini \
--workers 4 --max-iters 5 \
--verify
自动路由入口：auto_agent.py 是自动路由决策的入口：
python -m Fuser.auto_agent \
--problem /abs/path/to/KernelBench/level1/19_ReLU.py \
--verify
各个管道组件入口
不同级别的 UI 入口
TritonKernelAgent 类
agent.py 中的 TritonKernelAgent 类提供了编程接口：
from triton_kernel_agent import TritonKernelAgent
agent = TritonKernelAgent(num_workers=4, max_rounds=8, model_name="gpt-5")
result = agent.generate_kernel(
problem_description="Implement ReLU over a contiguous 1D tensor of length 1024"
)
AutoKernelRouter 类
auto_agent.py 中的 AutoKernelRouter 类提供了自动路由功能：
from Fuser.auto_agent import AutoKernelRouter
router = AutoKernelRouter()
result = router.solve(problem_path)
完整管道执行顺序
自动路由决策流程
pipeline.py 是 KernelAgent 系统中的主协调器，将复杂的多步骤过程封装为简单的端到端管道。它提供了一个高级接口，隐藏了底层的复杂性，使得用户可以轻松执行从 PyTorch 模型到优化 Triton 内核的完整转换过程。
pipeline.py 是 KernelAgent 系统中的核心管道文件，实现了“extract → dispatch → compose”的端到端工作流程。这是一个一站式管道运行器，将三个关键步骤整合到一个统一的执行流程中。
问题文件（.py）
↓
[subgraph_extractor.py] - 提取子图
↓
[dispatch_kernel_agent.py] - 分发到 KernelAgent 生成 Triton 内核
↓
[compose_end_to_end.py] - 组合最终内核
↓
最终结果
对应函数如下：
run_pipeline () # 主函数
├─ extract_subgraphs_to_json () # 提取阶段
│ ├─ extract_subgraphs_to_json ()
│ └─ orchestrator.run ()
├─ dispatch_run () # 分发阶段
│ └─ run ()
└─ compose () # 组合阶段
pipeline.py 的流程图如下。
参数处理阶段：
管道执行阶段：
具体如下：
pipeline.py 的数据流转如下：
输入数据流
problem_path（输入问题文件）
↓
subgraphs_path（提取的子图 JSON）
↓
kernels_summary_path（内核生成摘要）
↓
composition_result（组合结果）
输出数据流
return {
"run_dir": str (run_dir), # 运行目录
"subgraphs": str (subgraphs_path), # 子图路径
"kernels_summary": str (summary_path), # 内核摘要
"composition": comp_res, # 组合结果
}
功能描述
依赖模块
subgraph_extractor.py - 提取子图并转换为 JSON
orchestrator.py - 运行融合重构
prompting.py - LLM 提示构建
code_extractor.py - 代码提取
runner.py - 代码执行
输入输出
功能描述
并行处理
--dispatch-jobs auto）依赖模块
数据流转
功能描述
compose 函数验证功能
--verify 标志）依赖模块
模型选择
# 自动选择 Level 2/3 的默认模型为 GPT-5
if is_12 or is_13:
dispatch_model = "gpt-5"
else:
dispatch_model = "o4-mini"
并发控制
dispatch-jobs：控制并行处理的作业数量workers：控制提取阶段的并发数auto 模式自动匹配子图数量目录结构
管道在 .fuse// 目录下组织所有工件：
.fuse/<run_id>/
├── subgraphs.json # 子图描述
├── kernels_out/ # 生成的内核
├── summary.json # 每个子图的成功/失败状态
└── compose_out/
└── composed_kernel.py # 最终 Triton 程序
└── summary.json # 组合元数据
验证结果
subgraph_extractor.py 获取子图提取功能dispatch_kernel_agent.py 进行内核生成compose_end_to_end.py 进行最终组合# 核心模块依赖
from .subgraph_extractor import extract_subgraphs_to_json
from .dispatch_kernel_agent import run as dispatch_run
from .compose_end_to_end import compose
# 平台配置依赖
from triton_kernel_agent.platform_config import get_platform_choices
依赖关系如下：
pipeline.py
├─ subgraph_extractor.py
│ ├─ config.py
│ ├─ orchestrator.py
│ ├─ prompting.py
│ ├─ code_extractor.py
│ ├─ runner.py
│ └─ utils.providers
├─ dispatch_kernel_agent.py
│ ├─ triton_kernel_agent
│ │ ├─ agent.py
│ │ ├─ manager.py
│ │ ├─ worker.py
│ │ ├─ prompt_manager.py
│ │ └─ platform_config.py
└─ compose_end_to_end.py
├─ code_extractor.py
├─ runner.py
└─ utils.providers
手动执行
当需要对模型或并发性进行显式控制时使用：
python -m Fuser.pipeline \
--problem /path/to/problem.py \
--extract-model gpt-5 \
--dispatch-model o4-mini \
--compose-model o4-mini \
--workers 4 \
--max-iters 5 \
--verify
平台支持
--target-platform xpu）错误处理
pipeline.py 作为 KernelAgent 系统的核心管道，具有以下特点：
这种设计使得系统具有良好的可维护性和可扩展性，同时保证了端到端处理的可靠性。
KernelFalcon: Autonomous GPU Kernel Generation via Deep Agents
Ouyang A, Guo S, Arora S, et al. Kernelbench: Can llms write efficient gpu kernels?[J]. arXiv preprint arXiv:2502.10517, 2025.
Baronio, Carlo, et al. "Kevin: Multi-turn rl for generating cuda kernels."arXiv preprint arXiv:2507.11948(2025).
Li, Shangzhan, et al. "Autotriton: Automatic triton programming with reinforcement learning in llms."arXiv preprint arXiv:2507.05687(2025).
Li, Jianling, et al. "Tritonbench: Benchmarking large language model capabilities for generating triton operators."Findings of the Association for Computational Linguistics: ACL 2025. 2025.
Tjarko Lange, Robert, et al. "Towards Robust Agentic CUDA Kernel Benchmarking, Verification, and Optimization."arXiv e-prints(2025): arXiv-2509.
Chen, Wentao, et al. "CUDA-LLM: LLMs Can Write Efficient CUDA Kernels."arXiv preprint arXiv:2506.09092(2025).
Automating GPU Kernel Generation with DeepSeek-R1 and Inference Time Scaling
DeepSeek-R1自写CUDA内核跑分屠榜！斯坦福学霸狂飙GPU编程自动化挑战人类
大模型能否为不同硬件平台生成高性能内核？南大、浙大提出跨平台内核生成评测框架MultiKernelBench
AKG kernel Agent：利用multi-agent进行kernel的生成和迁移
AKG KERNEL AGENT: A MULTI-AGENT FRAMEWORK FOR CROSS-PLATFORM KERNEL SYNTHESIS
RL 猛刷 CUDA 核：CUDA-L1: Improving CUDA Optimization via Contrastive Reinforcement Learning
MultiKernelBench: A Multi-Platform Benchmark for Kernel Generation
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。