PyTorch KernelAgent 源码解读 ---（1）--- 原理

背景知识

1.1 需求

随着AI模型快速迭代与AI软硬件不断演进，AI产业对高质量算子的需求愈发强烈：

• 模型侧：大模型（LLM）领域，稠密模型、MOE、MLA、多模态场景众多；稀疏、量化、KVCache压缩等技术又为算子生成带来更多复杂性与多样性。推荐类、 CV类模型同样在持续演进，与LLM截然不同的算子场景同样存在强烈的优化落地需求；
• 硬件侧：Nvidia GPU的不断升级，国产AI芯片的持续不断演进，硬件架构、特性、参数各有不同；为提高AI芯片可用性，构建性能优势，定制化算子优化方案需求强烈；

1.2 内核范式

在深度学习模型的推理与训练过程中，绝大部分计算都依赖于底层计算内核（Kernel）来执行。计算内核是运行在硬件加速器（如 GPU、NPU、TPU）上的 “小型高性能程序”，它负责完成矩阵乘法、卷积、归一化等深度学习的核心算子运算。

高性能 GPU 内核主要依赖两种范式：手工编写和自动调优。

• 手工编写性能高，但门槛高，效率低。当前，这些内核通常由开发者使用 CUDA、AscendC、Pallas 等硬件专用并行编程语言手工编写 —— 这要求开发者具备精湛的性能调优技巧，并对底层硬件架构有深入理解。如今，传统人工优化内核的方式，在效率上已经不足以应对大量涌现的AI模型架构和硬件平台。
• 自动调优提升了部分效率，但仍受搜索空间与成本限制。如何在性能与开发效率之间取得平衡，成为关键问题。

1.3 Kernel Engineer

为什么 AI Infra 里还在大量招聘 GPU Kernel Engineer？ 论述了为何还需要Kernel Engineerr去手撸CUDA，具体如下：

• 通用性库难以适配特定场景：cuBLAS和CUTLASS 的成熟是建立在通用性基础上的，通用往往意味着妥协，难以适配特定shape、算子、存储墙等。比如，面对特定算子，想要业务上线快、性能好，只能自己下场写。
• Kernel Engineer依然极度稀缺：Kernel Engineer的高级形态是算法-硬件协同设计（Co-design）：只要硬件架构、模型架构不断演进，就需要重新适配和压榨每代硬件。
• 护城河高：这条路真正的护城河是对计算机体系结构的深刻理解，对体系结构的洞察力，对对编译技术的关注、跨栈的理解能力至关重要。

1.4 LLM 优化

虽然对 Kernel Engineer 的需求依旧强烈，但是人们也在思考：是否也能够借助LLM来模拟AI工程师的工作流程，凭借编译器反馈、硬件规格等丰富的信息，自动编写出准确且经过优化的内核代码呢？

实际上，LLM代码能力日益提升，为高性能 GPU 内核代码自动生成提供了一种极具潜力的方法。已有研究表明，现有 LLM 已具备一定的 GPU 内核生成能力。例如，英伟达工程师基于 DeepSeek-R1 设计了一套工作流程，在简单的 CUDA 内核生成任务中，该流程生成的内核在数值上全部正确，达到了 100% 的通过率。而且，在相关研究中，在完整的、复杂的机器学习架构上，平均加速达到了 50.8 倍！这验证了计算机体系结构的一个基本推论：系统的复杂性越高，自动优化的潜力就越大，因为人类工程师无法穷尽所有可能的优化组合。这里的性能提升是非线性的。

因此，基于LLM的算子生成技术逐渐成为业界研究热点。自2025年始，随着大模型的代码生成能力日益完善，用LLM来编写代码已在各类CodeAgent、AI IDE 中落地部署；在算子生成领域，各大知名高校、厂商纷纷开始投身到LLM生成算子的探索工作当中。

1.5 LLM 任务

LLM生成算子的过程，就是给定一个PyTorch程序，让模型对其优化，然后生成一个包含自定义CUDA内核的PyTorch版本。在此期间中，模型可以自由决定优化哪些操作，以提高计算效率，其目标是自动生成正确且高性能的 GPU Kernel。

任务输入

KernelBench是一个开源框架，旨在评估LLM在编写GPU内核方面的能力。KernelBench包含250个任务，涵盖了各种AI工作负载，并且易于扩展到新的工作负载。

下面给出 KernelBench 中的一个示例任务。每个任务都被封装在一个名为 Model 的类中。一个任务在 Model 类里主要包含两个核心函数：__init__和 forward；如有需要，还会提供辅助函数。AI算法通常在大型张量数据上进行操作。工作负载的最优内核取决于张量的大小和数据类型（如BF16、FP8），可以固定输入的 shape，仅通过随机生成的张量来改变数值。为此，我们提供了两个函数：get_inputs 和 get_init_inputs，分别用于生成初始化模型和运行前向传播时所需的随机参数。

任务输出

下面给出某个模型针对上述任务规范进行优化后的示例输出。该模型不仅要生成 CUDA 内核代码，还要生成将其接入 PyTorch 框架所需的外围代码。评测框架会像调用普通 PyTorch 算子一样执行模型的 forward 函数，因此常见的做法是把 CUDA 代码直接内联到 Python 中。

1.6 挑战

算子开发领域主要存在四个挑战：

• 算子开发效率低，门槛高，依赖专业知识。手动方式来实现和优化算子时，开发者需同时掌握算法逻辑与硬件细节（内存架构、并行模型、指令集等），门槛极高，同时优化算子过程中，需通过反复实验调优线程块大小、tileing策略等参数，耗时且繁琐。

• 单一 LLM 生成算子存在正确性与性能缺陷。但是如果仅仅靠通用的LLM辅助开发的模式的话，不增加专业的软硬件知识和流程保证，易出现编译错误、数值不准确或性能不佳，难以满足生产需求。

• 跨平台与多 DSL 适配难度大。新的算子开发工具如triton/tilelang，目标是支持多种硬件平台，但是实际上，不同芯片类型的硬件架构差异还是很大的，很多算子实现还是会与硬件微架构耦合，迁移时需重新设计，可移植性差。

• 性能与开发效率/可移植性难以兼顾。追求开发效率与可移植性时，易牺牲算子性能，与人工优化后的算子性能差距大；如果要追求性能，现有优化方案多为单轮生成，无法系统性探索优化空间，性能提升有限。

1.7 KernelBench

我们从 KernelBench 入手来再深入。KernelBench 是让整个GPU编程自动化的起始催化剂。

能力

下图展示了KernelBench评估语言模型（LM）生成高性能GPU内核的能力。KernelBench要求语言模型为给定的目标PyTorch模型架构生成优化的CUDA内核，并进行自动化评估。

这些任务根据包含的基本操作或PyTorch库函数的数量分为三个级别。

• Level 1包含100个单个基本操作，如卷积、矩阵乘法等AI基础构建块。

• Level 2包含100个操作序列，如卷积、ReLU和Bias的组合，这些操作可以融合成一个内核以提高性能。

• Level 3包含50个完整的机器学习架构，如AlexNet和MiniGPT等，这些架构在运行训练和推理时对内核的性能要求极高。

结果

在一次性基线评估中，LLM生成的内核平均在不到20%的任务中比PyTorch Eager更快。这表明，仅靠简单提示，LLM很难在性能上超越传统的PyTorch内核。另外，这些模型生成的内核存在大量执行错误、功能正确性问题，并且无法进行特定平台的优化。

具体来说，研究者发现：

• 对模型而言，编写功能正确的内核仍然具有挑战性；
• 模型通过优化展示了生成高性能内核的潜力；
• 利用反馈对于减少执行错误和发现更快的方案很重要。
• 只有更强大的模型，会偶尔表现出利用这些优化的能力。

我们再用论文 MultiKernelBench: A Multi-Platform Benchmark for Kernel Generation 作为印证，对于 GPU，论文得出了以下结论：

• CUDA 在功能正确性方面对 LLM 来说具有挑战性：尽管 CUDA 文档丰富，LLM 在生成功能正确的内核时仍面临困难，主要是由于输出不匹配、输出形状不匹配和 CUDA 运行时错误。
• GPT-4o 在 CUDA 编译成功率方面表现出色：GPT-4o 在 CUDA 上的编译成功率最高（97.5%），表明它与 CUDA 约定具有很强的语法一致性。
• DeepSeek-R1-0528 在 CUDA 优化能力方面表现突出：该模型在 CUDA 上获得了最佳的 Pass@1 和 SpeedUp1@1 分数，表明它擅长生成优化的 CUDA 代码。
• 类别感知的一次性提示能显著提高 LLM 在 CUDA 上的性能：当LLM 收到与目标任务相同类别的示例时，它们的性能会显著提升，尤其是在像矩阵乘法这样原本难以正确生成内核的类别中。

1.8 方向

目前，基于 LLM 的 GPU 内核代码生成方法大致可分为两类：

• 微调范式（Fine-tuned LLM）：通过 SFT / RL 微调，使大模型具备端到端的 GPU Kernel 生成与优化能力，其核心逻辑是数据驱动。
• Agent Pipeline 范式：强调多阶段推理与工具协作，通常将 Kernel 代码生成过程拆解为多个模块化步骤，结合编译器反馈、性能分析工具和运行结果实现闭环优化。

目前猜测的几条趋势（可能有些已经过时）

• 从单agent走向多agent 编排

  • 单LLM"一次生成一个内核“已经触顶；主流系统会引入路由/分解/分发/合成的多角色协作（KernelAgent、STARK 都属此列）

• 从“采样N次“走向“多轮RL+工具反馈”

  • 早期是best-of-k并行采样；现在SOTA 几乎都是反馈闭环：编译错误、运行错误、profiler
    计数器、ncu/rocprof 输出全部回灌给模型。Kevin-32B、ByteDance Agent 都是RL路线。

• 从“求正确“走向“求实测加速”

  • 奖励信号从“通过单测"升级为“实测wall-clock加速比"；这要求hardware-in-the-loop（真GPU 上跑、跑很多次、跑稳定）。

• 严格沙箱与防reward hacking成为标配

  • Sakana事件之后，所有正经系统都强制：子进程隔离 + 真PyTorch参考输出对比 + 禁止torch.*
    fallback。

• 领域专用语言（DSL）正在分化

  • Triton 现在是agent的“甜蜜区"：比手写 CUDA 容易、又能拿到80%+性能；NVIDIA CUTLASS Python、AMD Composable Kernel、Intel XPU Triton 都在往同一个方向走（高层tile DSL + autotune）。

• “编译器 × LLM"协同设计兴起

  • 大家发现纯LLM fine-tune 接近饱和，.下一波收益来自：让agent理解IR、调autotune、读
    profiler，而不是只写源码。可以把KernelAgent的Composer阶段看作这条路的早期形态。

• 基准开始“反通胀”

  • METR等独立评测在加新题、加frontier workloads，防止benchmark 被记忆/过拟合。未来一年端到端LLM训练/推理工作负载（attention变体、MoE、量化）会成为主战场。

0x02 微调范式

我们选择几个微调范式的示例来深入学习。

2.1 Kevin-32B

Kevin-32B（ https://arxiv.org/pdf/2507.11948 ）由Stanford Cognition AI实验室提出，生成语言为 CUDA，基座模型是 QwQ-32B，在 KernelBench L1 和 L2 级别展开实验。

论文作者利用 KernelBench L1 和 L2 上的 180 个题构建 RL 流程直接微调基座模型。使用的 RL 算法为 GRPO。训练流程分为了单轮训练（Single-Turn Training ）和多轮训练（Multi-Turn Training ）。

• Single-Turn Training 就是每个时间步纯并行采样，每步采样 16 次作为 GRPO 的一组，然后根据正确性和性能生成 Reward，反馈更新模型参数。
• Multi-Turn Training 扩展了 Single-Turn Training，总共训练 40 Step （80 次梯度更新）：
  • 增加采样次数（m 条轨迹，每条轨迹 n 次细化，细化就是拿到前面几次的结果 + CoT 作为本轮的输入 prompt）；
  • 提高样本利用率。拆了轨迹，用样本点训；
  • 增加历史信息和 CoT。保留前面几轮的 CoT + Eval 信息作为当前时间步的输入 prompt。

2.2 CUDA-L1

CUDA-L1: Improving CUDA Optimization via Contrastive Reinforcement Learning

CUDA-L1 的成功建立在一个精心设计的三阶段渐进式训练流程之上，旨在逐步提升 LLM 的 CUDA 编程与优化能力。三阶段像“模仿→比较→竞争”，每一步都把“能跑”→“较好”→“最好”做成独立里程碑；用对比学习做“护栏”，让强化学习始终在高速路而不是野地里飙车，最终拿到的模型既敢改代码，又不会改崩。

训练流水线总览如下：

• 阶段一：数据增强监督微调——用 LLM 产生大量 CUDA 代码变体，筛选出可执行且正确的样本，对基座模型进行微调，建立 CUDA 基础认知。即，让模型尽可能多地见到各种 CUDA 模式和编程结构，首要目标是“写出能编译、能跑通的代码”。
• 阶段二：自监督学习——模型迭代生成内核，自主验证正确性与可执行性，并用成功案例继续自我训练，实现无监督情况下的持续改进。即，模型自己生成内核→验证正确性→再用通过验证的样例继续训练，无需人工标注，就能大幅提升可执行性与正确率，并带来适度加速。
• 阶段三：对比强化学习——引入基于执行时间的对比奖励，训练模型区分“更快”与“更慢”的实现，最终生成性能优异的 CUDA 内核。即，用“运行时加速比”当奖励，通过对比快慢两种实现，让模型学会生成高性能 CUDA 代码，实现显著提速。

2.3 小结

基于微调（fine-tuning）的方法，在训练高性能 GPU 内核代码模型时面临多重瓶颈。

1. 训练数据质量参差不齐，难以覆盖复杂的内核模式；
2. GPU 内核代码通常上下文长、细节高度复杂（例如坐标计算与内存访问优化），导致模型难以捕捉全局优化逻辑；
3. 调优空间庞大，端到端训练模型同时生成正确且高性能内核几乎不可能（例如上千个 op 的图）。

0x03 Agent Pipeline 范式

相比之下，Agent Pipeline 方法通过将生成过程拆解为多阶段模块，并结合编译器反馈、性能分析与工具链协作，实现了更灵活的优化和验证。但是，基于 Agent 的方法生成的代码很少是“纯血”Kernel（总是会调用高层 API，e.g. torch API）。

3.1 需求

手工编写优化 GPU attention 内核既耗时又需要深厚经验。近期 LLM（如 DeepSeek-R1）在代码生成上表现亮眼，但仍难以一次性产出优化代码。因为把“写 GPU 内核”这件事直接扔给一次性的 LLM ，就像让一位刚毕业的学生闭眼一口气写完 500 行 CUDA——可能跑通，也可能带来灾难性后果。所以，需在推理阶段采用额外策略。

下面这段 prompt 是“relative positional embeddings attention kernel”任务里，用户扔给模型的示例输入。然而，LLM 偶尔会把不同语言或框架的语法“拌沙拉”，产出的代码要么编译不过，要么慢得离谱；而 GPU 线程怎么排布才算最优更是一门玄学，往往得“写-测-改”多轮才能拿到既对又快的内核。

Please write a GPU attention kernel to support relative position enco