PyTorch KernelAgent 源码解读 ---(1)--- 原理

📅 2026/6/26 5:02:49
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 上获得了最佳的 Pass1 和 SpeedUp11 分数表明它擅长生成优化的 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-32BKevin-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-L1CUDA-L1: Improving CUDA Optimization via Contrastive Reinforcement LearningCUDA-L1 的成功建立在一个精心设计的三阶段渐进式训练流程之上旨在逐步提升 LLM 的 CUDA 编程与优化能力。三阶段像“模仿→比较→竞争”每一步都把“能跑”→“较好”→“最好”做成独立里程碑用对比学习做“护栏”让强化学习始终在高速路而不是野地里飙车最终拿到的模型既敢改代码又不会改崩。训练流水线总览如下阶段一数据增强监督微调——用 LLM 产生大量 CUDA 代码变体筛选出可执行且正确的样本对基座模型进行微调建立 CUDA 基础认知。即让模型尽可能多地见到各种 CUDA 模式和编程结构首要目标是“写出能编译、能跑通的代码”。阶段二自监督学习——模型迭代生成内核自主验证正确性与可执行性并用成功案例继续自我训练实现无监督情况下的持续改进。即模型自己生成内核→验证正确性→再用通过验证的样例继续训练无需人工标注就能大幅提升可执行性与正确率并带来适度加速。阶段三对比强化学习——引入基于执行时间的对比奖励训练模型区分“更快”与“更慢”的实现最终生成性能优异的 CUDA 内核。即用“运行时加速比”当奖励通过对比快慢两种实现让模型学会生成高性能 CUDA 代码实现显著提速。2.3 小结基于微调fine-tuning的方法在训练高性能 GPU 内核代码模型时面临多重瓶颈。训练数据质量参差不齐难以覆盖复杂的内核模式GPU 内核代码通常上下文长、细节高度复杂例如坐标计算与内存访问优化导致模型难以捕捉全局优化逻辑调优空间庞大端到端训练模型同时生成正确且高性能内核几乎不可能例如上千个 op 的图。0x03 Agent Pipeline 范式相比之下Agent Pipeline 方法通过将生成过程拆解为多阶段模块并结合编译器反馈、性能分析与工具链协作实现了更灵活的优化和验证。但是基于 Agent 的方法生成的代码很少是“纯血”Kernel总是会调用高层 APIe.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