当前位置：首页 > news >正文

PyTorch KernelAgent 源码解读 ---（2）--- 总体流程

news 2026/5/13 21:25:46

PyTorch KernelAgent 源码解读 ---（2）--- 总体流程

PyTorch KernelAgent 源码解读 ---（2）--- 总体流程
- 0x00 摘要
- 0x01 引言
  - 1.1 背景与挑战
  - 1.2 设计理念
- 0x02 架构
  - 2.1 KernelFalcon 架构
  - 2.2 流水线：数据如何流经系统
  - 2.3 模块依赖关系图
  - 2.4 KernelAgent 调用时序图
- 0x03 架构：分阶段详解
  - 3.1 阶段 1：FuserAgent – 代码到代码融合
    - 算子融合
    - 为何有效
    - 深度智能体原则：确定性控制平面
    - Signature去重
  - 3.2 阶段 2：ExtractorAgent – 子图边界推理
  - 3.3 阶段 3：Dispatcher + KernelAgent – 并行 Triton 生成
    - 并行方法
    - 关键机制
    - 深度智能体原则：基于真实工具的执行
  - 3.4 阶段 4：ComposerAgent – 端到端内核缝合
- 0x04 KernelAgent 项目入口点
  - 4.1 命令行入口点
  - 4.2 UI 入口点
  - 4.3 编程接口入口点
  - 4.4 管道执行流程
  - 4.5 系统架构入口
- 0x05 pipeline.py 的作用分析
  - 5.1 执行流程
    - 流程图
    - 数据流转
  - 5.2 三大核心步骤
    - 提取阶段（Extract）
    - 分派阶段（Dispatch）
    - 组合阶段（Compose）
    - 参数配置
    - 输出管理
  - 5.3 依赖关系
  - 5.4 使用场景
  - 5.5 总结
- 0xFF 参考

0x00 摘要

本文基于 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，附带文档与入门示例。

0x01 引言

1.1 背景与挑战

KernelFalcon 面临的背景和挑战如下：

手写优化GPU kernel是部署瓶颈。编写优化的 GPU 内核仍是部署机器学习模型的瓶颈。团队很少有为每种形状、数据类型和硬件代际手工调优算子的带宽。随着模型演进，问题加剧：适用于 ResNet 的模式无法直接映射到 Mamba 的选择性状态或 MoE 的条件路由。
传统编译器（TorchInductor、TVM、XLA）难以处理长尾场景。现代编译器取得实质进展，但仍对长尾场景束手无策。TorchInductor 覆盖常见模式，TVM 自动调度稠密内核，XLA 针对动态形状做特化。然而，异常算子、动态控制流和异构融合模式仍逃逸最优编译。NVIDIA 今年年初使用 DeepSeek-R1 配合推理时缩放，在 KernelBench L1/L2 上取得强劲结果，证明基于 LLM 的方法配合验证循环可媲美或超越传统方法——但未触及完整模型架构（L3）。
LLM方法需要更好的架构设计。

如果能在不扩充规则库或雇佣更多 GPU 专家的前提下，自动合成保持 PyTorch 语义且逼近手工调优性能的 Triton 内核，会怎样？

因此，PyTorch推出了 KernelFalcon：一个保持 PyTorch 语义的代码到代码系统，可生成优化的 Triton 内核。它采用并行探索与基于执行的验证，而非一次性生成——交付的内核真实运行在 GPU 上，且与原始模型数值等价。

注：长尾问题原本是指在数据分布中，少数数据（“头部”）的频率非常高，而大多数数据（“尾部”）的频率非常低的一种现象。在这里主要想表达的是AI、算子相关的工作中总是由于创新、软硬件变化等缘故，出现传统方案无法覆盖的场景，导致功能不支持或者性能下降。

1.2 设计理念

KernelFalcon的核心目标是：将PyTorch 程序自动转换为经过验证的 GPU内核（kernel）。比如，它通过LLM生成 Triton 代码，并在沙盒子进程中自动验证正确性，迭代修复直到通过测试。

KernelFalcon基于 Pytorch Module 自动生成GPU Kernel，其设计理念如下：

保持Python语义（支持if/else、while、动态shape）
verification 优先循环（编译测试候选kernel，失败反馈）
端到端组合 verification（算子融合kernel替换原ops，全模型等价性检查）

项目分两个层次：

Fuser/（编排层）：面向完整PyTorch问题进行编排，包括AST分析路由、子图提取、内核分发、端到端组合。
triton_kernel_agent/（单内核合成工作池）：负责单个Triton内核的生成、验证和迭代改进，是底层工作单元。

0x02 架构

为何选择 KernelFalcon（为何是深度智能体）？

传统静态、基于图的编译器依赖 IR 变换和每模式调度。追踪常常将控制流冻结为单一路径，并在动态形状下失效。KernelFalcon 则走另一条路：

保持 Python 语义。KernelFalcon 停留在 PyTorch 代码到代码层面，因此 if/else、while、数据依赖路由和动态形状依然有效。
验证器优先循环。KernelAgent 编译并测试候选内核；失败结果本地反馈；KernelFalcon 在首个数值正确的内核上提前退出。
端到端组合与验证。融合内核替换原始算子后，进行整模型等价性检查才被接受。

底层是一个深度智能体架构——多阶段系统，通过结构化问题来减少 LLM 失效模式：

显式任务分解将模糊目标转化为精确、工具就绪的子问题，将复杂任务分解为子任务并分配给专门Agent处理
确定性编排将控制逻辑留在 Python，让 LLM 专注于认知，而非让LLM决定工作流程
并行搜索与早停高效探索多样解：多个worker并行探索不同kernel实现，一旦某个worker生成并通过verification的kernel，立即终止其他worker以节省计算资源
基于真实工具的每一步验证都针对真实编译器和硬件（使用Triton编译器、GPU执行等真实工具而非模拟）
结构化状态持久化提示、日志与产物，用于可审计与断点续跑

这不仅是更干净的实现；这是不同范式。不再问“LLM 能否解决此问题？”，而是问“KernelFalcon 如何塑造任务，使 LLM 可能成功？”结果是更广覆盖与更现实性能——无需膨胀规则集或牺牲语义。

2.1 KernelFalcon 架构

图 1：KernelFalcon 的深度智能体架构以 Orchestrator 为中心协调整个工作流。规划负责任务分解与预算分配。上下文工程提供结构约束（模板、指南）。子智能体处理专业任务，提取融合边界、生成 Triton 内核、组合端到端模块并执行验证。持久化内存存储产物用于调试与续跑。Orchestrator 委派给专家，接收结构化错误反馈，并在整个执行过程中保持状态。

该架构体现深度智能体原则：

分层委派：Orchestrator 将高级任务（融合此模型）分解为精确子问题（提取子图、生成内核、组合结果）并分配给专家智能体
确定性控制：规划与编排逻辑是显式 Python 代码，而非 LLM 驱动——worker生命周期、超时和成功条件均以编程方式定义
基于真实执行：每个智能体都针对真实工具（Triton 编译器、PyTorch 参考、GPU 执行）验证，而非模拟或 LLM 判断结果
持久化状态：所有中间结果、提示、日志和产物持久化到磁盘，用于可审计、调试和跨会话续跑
结构约束：上下文工程将规则编码为模板和策略，使正确性要求结构强制执行，而非依赖提示

2.2 流水线：数据如何流经系统

图 2：多阶段工作流图显示 PyTorch 输入流经 FuserAgent（创建可融合子图）、ExtractorAgent（生成 JSON 规范）、并行 KernelAgent worker（三个 Triton 框显示并发生成）和 ComposerAgent（缝合已验证内核）。箭头表示数据流，并标注中间表示。

流水线包含四个不同阶段：

FuserAgent – 保持 Python 语义的代码到代码融合
ExtractorAgent – 形状推理与合约生成
Dispatcher + KernelAgent – 并行 Triton 内核合成与验证
ComposerAgent – 端到端集成与验证

2.3 模块依赖关系图

系统中的模块依赖关系如下：

2.4 KernelAgent 调用时序图

0x03 架构：分阶段详解

3.1 阶段 1：FuserAgent – 代码到代码融合

算子融合

FuserAgent 其实是在做算子融合。

直接在PyTorch源代码上操作
保持控制流和Python语义
输出：带明确子图边界的干净PyTorch模块

传统编译器在融合分析期间将 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 模块，控制流保持完整

# 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 intactx = 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 内核、组合内核）；控制器执行并验证输出。

工作流如下：

Orchestrator 生成 N 个worker，携带类型化 WorkerConfig，流式传输日志，等待队列上的获胜者，取消其他worker，并打包产物
worker迭代：渲染提示 → 流式 LLM → 提取 Python 块 → 按 SHA 去重 → 执行候选 → 若通过，则发出获胜信号；否则保存错误并重试
无需手动 AST 解析或基于规则的融合检测——LLM 直接通过prompts 来融合代码，然后 Python 通过执行验证

Signature去重

构建stable signature：ops序列 + shapes + weights（包含input shape、output shape、权重、layout、数据类型）
相同signature合并，累加count
避免重复生成相同配置的kernel

参考：Fuser/orchestrator.py、Fuser/worker.py、Fuser/runner.py、Fuser/prompting.py

3.2 阶段 2：ExtractorAgent – 子图边界推理

本阶段是提取器使用 LLM 分析融合代码，并识别具有形状合约的精确子图边界。

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

提取过程：

运行 orchestrator：首先，从阶段 1 获取融合代码，分析算子融合后的PyTorch代码，识别每个唯一的子图
- 提取每个子图的shape信息（input/output/weights shape）
- 为生成Triton kernel提供精确的shape信息
提示 LLM：要求 LLM 识别不同的子图函数，推断形状，并编目操作
生成 JSON 规范：LLM 生成带类型规范，包含操作序列、形状和权重元数据
去重与合并：按稳定签名（算子 + 形状 + 权重）对子图分组，聚合计数，即构建stable signature (ops + shapes + weights) 用于去重，相同shape的子图只需生成一次kernel

输出：子图规范的 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.3 阶段 3：Dispatcher + KernelAgent – 并行 Triton 生成

Dispatcher 为每个子图规范协调并行 Triton 内核生成。对于每个子图，它创建一个带有worker池（默认 4 个worker）的全新 TritonKernelAgent。

协调多个并行worker
每个worker独立生成kernel并进行verification
Early Stop：第一个成功的worker停止其他worker

图 3：KernelAgent 生成并行worker，采用多样采样参数生成 Triton 内核。每个候选者经历验证阶段（语法、编译、数值）。失败的候选者触发隔离错误反馈至其源worker——无上下文污染。首个通过所有阶段的候选者立即部署并取消剩余worker。实现并行探索与隔离上下文及早退。

并行方法

用相同提示但不同温度设置（0.8、0.9、1.0 等）生成 N 个内核种子。生成 N 个worker（默认 4 个），每个在其自己的工作目录中运行隔离的精炼循环。不同温度导致worker探索不同的优化策略——有些保守，有些探索性。

多个worker并行探索不同解决方案
第一个成功的worker立即停止其他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

3.4 阶段 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 来验证。这种基于真实执行的基础而非模拟或 LLM 判断的正确性。

输出产物：

成功记录验证状态、计时与产物路径。完整的组合模块成为最终交付物，准备部署或进一步

参考：Fuser/compose_end_to_end.py

0x04 KernelAgent 项目入口点

KernelAgent 项目的入口点分为两个层面：

用户界面层面：通过命令行工具（pipeline.py 和 auto_agent.py）和 GUI 界面
编程接口层面：通过 TritonKernelAgent 和 AutoKernelRouter 类提供 API 访问

其中，auto_agent.py 是推荐的主要入口点，因为它会根据问题复杂度自动选择最优路径。

4.1 命令行入口点

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

各个管道组件入口

subgraph_extractor.py - 子图提取
dispatch_kernel_agent.py - 内核调度
compose_end_to_end.py - 端到端合成

4.2 UI 入口点

不同级别的 UI 入口

Triton KernelAgent UI: kernel-agent 或 python scripts/triton_ui.py
Fuser 编排 UI: fuser-ui 或 python scripts/fuser_ui
完整管道 UI: pipeline-ui 或 python scripts/pipeline_ui

4.3 编程接口入口点

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)

4.4 管道执行流程

完整管道执行顺序

提取阶段：subgraph_extractor.py - 从问题文件提取子图
分派阶段：dispatch_kernel_agent.py - 为每个子图生成 Triton 内核
合成阶段：compose_end_to_end.py - 将子内核合成为完整的端到端内核

自动路由决策流程

静态分析：解析问题文件的 AST，检查操作符模式
复杂度评估：评估是否存在难以融合的操作（如注意力机制、转置卷积等）
决策：根据复杂度决定使用 KernelAgent 直接路径还是完整的 Fuser 管道
执行：执行选定的路径，并在失败时可选择回退到另一条路径

4.5 系统架构入口

Fuser 编排器：orchestrator.py 管理融合重构过程，是 Fuser 系统的核心编排组件。
验证运行器：runner.py 负责安全地执行候选程序并验证其正确性。

0x05 pipeline.py 的作用分析

pipeline.py 是 KernelAgent 系统中的主协调器，将复杂的多步骤过程封装为简单的端到端管道。它提供了一个高级接口，隐藏了底层的复杂性，使得用户可以轻松执行从 PyTorch 模型到优化 Triton 内核的完整转换过程。

5.1 执行流程

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 的流程图如下。

参数处理阶段：
- 解析输入参数（问题路径、模型名称、工作进程数等）
- 如果未指定 dispatch_model，则根据问题级别自动选择模型
管道执行阶段：
- 调用 extract_subgraphs_to_json 提取子图
- 调用 dispatch_run 分发子图到 KernelAgent 生成内核
- 调用 compose 合成最终内核
- 返回包含所有结果的字典

具体如下：

数据流转

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,          # 组合结果
}

5.2 三大核心步骤

提取阶段（Extract）

功能描述

从输入的问题文件中提取子图
使用 extract_subgraphs_to_json 函数
生成包含子图信息的 JSON 文件

依赖模块

subgraph_extractor.py - 提取子图并转换为 JSON
orchestrator.py - 运行融合重构
prompting.py - LLM 提示构建
code_extractor.py - 代码提取
runner.py - 代码执行

输入输出

输入：问题文件路径
输出：子图描述 JSON 文件和运行目录

分派阶段（Dispatch）

功能描述

将提取的子图分派给 TritonKernelAgent
使用 dispatch_run 函数
并行生成 Triton 内核

并行处理

支持并发处理多个子图
可以自动匹配子图数量（--dispatch-jobs auto）

依赖模块

dispatch_kernel_agent.py - 调度子图到 KernelAgent
triton_kernel_agent - Triton 内核生成引擎
platform_config.py - 平台配置

数据流转

输入：subgraphs_path（来自提取阶段）
输出：summary_path（内核生成摘要）

组合阶段（Compose）

功能描述

将验证过的内核组合成单个 Triton 程序
使用 compose 函数
创建最终的端到端内核

验证功能

可选的验证步骤（--verify 标志）
确保组合后的内核功能正确

依赖模块

compose_end_to_end.py - 组合 Triton 内核
code_extractor.py - 代码提取
runner.py - 代码执行验证

参数配置

模型选择

# 自动选择 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   # 组合元数据

验证结果

验证通过时输出验证日志
确保最终结果的正确性

5.3 依赖关系

上游：从 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

5.4 使用场景

手动执行

当需要对模型或并发性进行显式控制时使用：

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

平台支持

支持 CUDA 平台（默认）
支持 XPU 平台（--target-platform xpu）

错误处理

统一的异常处理机制
系统退出码管理
详细的错误信息输出

5.5 总结

pipeline.py 作为 KernelAgent 系统的核心管道，具有以下特点：

顺序依赖：三个阶段严格按顺序执行，后一阶段依赖前一阶段的输出
模块化设计：各阶段独立封装，通过标准化接口交互
错误传播：异常向上逐层传递，最终影响返回状态
数据一致性：通过路径参数确保数据在各阶段间正确流转
配置统一：共享模型选择、超时设置等配置参数

这种设计使得系统具有良好的可维护性和可扩展性，同时保证了端到端处理的可靠性。

0xFF 参考

KernelFalcon: Autonomous GPU Kernel Generation via Deep Agents

基于 LLM 的 GPU 内核代码自动生成相关工作

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编程自动化挑战人类

CUDA、Triton 内核生成现状追踪

大模型能否为不同硬件平台生成高性能内核？南大、浙大提出跨平台内核生成评测框架MultiKernelBench

AKG kernel Agent：利用multi-agent进行kernel的生成和迁移

AKG KERNEL AGENT: A MULTI-AGENT FRAMEWORK FOR CROSS-PLATFORM KERNEL SYNTHESIS

AIKG -- 基于AI驱动的算子生成器

RL 猛刷 CUDA 核：CUDA-L1: Improving CUDA Optimization via Contrastive Reinforcement Learning

MultiKernelBench: A Multi-Platform Benchmark for Kernel Generation

查看全文