Achieving Extreme Efficiency through Specialized GPU Kernel Generation

通过专用GPU内核生成实现极致效率 | Databricks博客

跳至主要内容

数据科学与机器学习

2026年9月4日

通过专用GPU内核生成实现极致效率

可靠且经过验证的GPU内核生成

作者：李 Leo、Khudia Daya 和金乐生

摘要

• 新内核草稿的生成成本低且易于并行生产。但信任并非如此。系统的进展速度取决于我们验证这些草稿的速度。

• 一个看似奇迹般快速的数字往往是一个测量错误：来自前一次运行的残留工作、两侧未执行相同操作的比较，或在隐藏假设下才表现良好的内核。

• 上下文是一个权衡，而非需要最大化的堆叠。更多文本为模型提供更多素材，但成本更高，额外注释会使下一次尝试更容易偏离。太少则循环无法推进。

• 能够自由探索的智能体能生成更优内核。严格的外部系统必须定义反馈并决定允许发布的范围。良好的设计需要两者兼顾。

• 生成的单个Qwen 3.5 122B内核速度比vLLM中现有最佳实现快1.8-5.2倍。

传统生产推理系统依赖通用内核来处理多样化的模型和工作负载。这并非最优方案，因为GPU操作形状由静态模型参数和动态请求时因素共同决定；例如，虽然模型定义了矩阵乘法的一个维度，但另一个维度会根据每个请求的具体令牌数量波动。人们对智能体驱动的GPU内核生成日益感兴趣，近期尝试已显示出前景。我们探索了一个核心问题：如果内核生成可以自动化，为何参数规模差异巨大的模型（从10亿到1万亿参数）需要依赖相同内核？通过将内核专门化为运行时遇到的具体形状，我们可以实现极致效率。

在本文中，我们分享了使用智能体生成GPU内核的成功经验和见解。我们构建了Proteus系统，该系统专为实现极致专门化而设计，需要一个专为严格优化、验证和上下文管理定制的框架。

图1：Proteus框架的简化视图

传统编码框架在此处常会失败，因为智能体倾向于奖励黑客行为：遵循字面规则而非精神实质。如果你给智能体一个基准测试，它可能会优化基准测试本身而非预期操作。

为解决此问题，Proteus提出内核，将其与受控参考实现进行验证，对成功案例进行计时，并不断改进最佳结果。虽然流程简单直接，但其成功完全取决于解决两个基础性挑战。图1展示了我们设计的简化架构。使用我们的Proteus框架，我们生成的Qwen 3.5 122B内核速度比vLLM中现有最佳实现快1.8-5.2倍。

验证

最初我们将内核搜索视为难题：如何在不陷入无改进平台的情况下探索大量程序空间？实际上第一个问题更为基础。我们是否在测量我们以为正在测量的内容？

模型会根据你给予的评分进行优化。它并不需要依赖什么奇特的漏洞：可能只是评估过程本身存在某种假设。一个例子是注意力层中常见的旋转位置嵌入（RoPE）核函数。候选方案可能复用之前尝试遗留下来的编译代码，看起来比从零开始重新构建更便宜。另一个方案可以将一批GPU启动记录在图中（例如CUDA图），然后作为一个单元重放，而我们对比的基线方案仍然需要单独启动每个部分，因此双方实际上并没有执行相同的工作。还有一个案例在可见测试集中的输入尺寸表现优异，但在未见过的尺寸上表现较差。

因此我们在早期设计阶段投入了大量工作来完善验证器，而非提示词。对双方采用完全一致的计时方式，包括在需要交叉验证时使用多个计时器（例如CUDA事件计时器、墙钟时间和CUPTI计时器）。清除不应保留的遗留编译状态，保持初始化和清理的顺序一致，防止某一方跳过另一方仍需执行的工作。在将候选方案作为下一轮起点前，需要再次对其性能进行验证。保留部分候选方案无法看到的测试用例，防止其仅针对考试内容进行拟合。为防止通过人为夸大性能指标进行评估作弊，我们实现了自动一致性检查，用于标记理论上不可能的加速（例如超过100倍的提升），这些加速超过了物理GPU的带宽和计算限制。这可以防止过去工业案例中出现的奖励黑客攻击漏洞，当时代理程序优化的是评估指标而非真正的性能提升。没有这些约束条件时，生成更多核函数通常只会产生更多噪声。

对验证器的重视也改变了智能体核函数生成的瓶颈。在单纯的程序搜索（即通过反复生成变体进行迭代优化的系统）工作中，优质候选方案非常稀少，因此编写它们占用了主要成本。我们可以并行生成大量草稿，但无法跳过验证环节。我们必须精心设计验证过程，且验证必须在真实GPU上独立运行，并且需要多次执行。系统的运行速度取决于它能信任核函数的程度，而非它能编写核函数的速度。

图2：初始评估框架各阶段的token使用情况

图3：改进后评估框架各阶段的token使用情况

上下文管理

另一个挑战是确定核函数生成模型可以查看的内容范围。这是一个权衡问题。给模型提供更长的提示词会提供更多信息：当前最佳核函数、近期失败案例、分析器提示、早期运行的笔记。这确实有帮助，但代价更高，因为我们要为模型读取的每个token付费。随着提示词变长，下一次尝试更容易偏离方向。有用的信号会与过时的建议、相互矛盾的提示以及适用于不同输入尺寸或不同操作的细节混杂在一起。模型并不总能判断哪些句子可信，因此会遵循最响亮的建议，或对所有建议都稍作采纳。

如果提供太少信息，情况会适得其反。每次尝试都必须从零开始。相同的死胡同会反复出现。上一次运行或相关操作中的任何经验都无法继承，循环无法取得进展。

我们希望有一个知识层来解决这个问题：记住哪些方法有效，之后可以复用这些经验，而且无需人工干预。这个知识层存在另一个权衡，即存储的经验细节程度与适用范围之间的平衡。

一个非常具体的提示（“在该内核上，针对该输入规模，展开此循环”）可能正是下一次尝试所需的内容。但这种提示也容易在下一次操作、不同GPU或不同输入规模时被错误使用。一个非常通用的提示（“更好地利用片上内存”）几乎适用于所有场景，但对模型几乎没有具体指导。我们观察到了这两种失败模式：当经验过于通用时，它们只是重复了失败信息而没有给出具体行动；而当存储了更多细节时，这些经验又往往过于依赖特定运行环境，难以指导后续操作。在一次长时间运行中，模型读写操作的大部分时间都花在了获取和路由内存数据上，而非编写内核代码。内存层承担了大量工作，但并未有效提升后续候选方案的质量。图2展示了此类系统的token成本分布，其中知识层占据了主要成本。

值得保留的知识版本更加精炼，并在通用性与具体性之间取得平衡。当模型即将编写内核时，其提示信息应仅包含高可信度的上下文：将具体场景与具体行动相匹配的可操作经验（从过往修改到影响的映射中提炼得出），以及来自高度相关父运行的简洁失败记录。这些经验应通过分层标签过滤结合混合（关键词+语义）搜索进行检索，需具体到足以指导行动，同时明确适用范围。更深层次的操作，如重组和进一步提炼经验库，应属于后台任务，而非在每次尝试时都进行同步的多跳追溯。如果某条经验无法明确说明具体场景和对应行动，则不值得放入提示信息中。图3展示了修复知识层后的token成本分布，修复后大部分token消耗都用于候选方案生成。

案例研究：Gated DeltaNet打包解码

一个具体案例是Qwen 3.5 122B模型中Gated DeltaNet路径上的打包解码内核。该操作需要更新循环状态，并从打包的QKV输入、门控参数和状态索引中写入解码输出。我们使用该任务在NVIDIA B200 GPU上通过Triton后端完整测试了Proteus流程：验证任务契约、测量参考实现、请求代理生成候选内核、运行静态检查和构建、与受控参考验证正确性、仅对已验证候选进行基准测试，然后重新测量最佳候选。

图4从左到右展示流程。基准节点将基准线锚定在0.025毫秒。候选0000是安全种子：它复现了打包解码结构并通过验证，但速度慢于参考实现，因此Proteus将其作为已测量父节点保留，而非视为成功方案。从这里开始，Proteus停止对每种形状优化单一通用内核，而是将搜索拆分为形状特定路径。

图4：Proteus框架下的内核演化案例研究

Batch-1 修复路径在候选方案 012 中生成了一个形状特定的内核，在单批次解码形状上实现了 1.5 倍的加速。最强的结果出现在服务-解码路径上：候选方案 030 在 0.018 毫秒内实现了最低的内核延迟测量值，而候选方案 036 在形状加速方面表现最佳，达到 1.6 倍。这个获胜的服务片段专门针对 Batch=4、Key=128、Value=128 的布局进行优化，以 64 宽的块处理值维度，因此它是针对该特定形状的安全内核，而不是通用替代方案。

时间线中的最后一次绕道展示了轨迹的重要性。后续的 C++ 代际（而非 Triton）尝试遇到了构建和生成失败的问题，当耗尽该分支的尝试预算后，长期运行也随之结束。因此，有用的产物不仅仅是最快的候选方案。图中显示的完整路径同样重要：语义失败被拒绝，正确但较慢的内核被测量，真正性能提升的内核则被保留在使其安全组合成生产内核的形状上。

我们下一步的工作重点

我们构建的进化循环对作者来说非常严格。它经常以固定模式调用模型：采用当前最佳内核，尝试小幅修改，检查结果，重复此过程。这种方式剥夺了代理所需的自主性。它无法轻松改变结构、切换语言或放弃失败的设计。

这个循环仍然需要存在。不是为了编写内核，而是为了给代理提供一个可信的下一步提示。这个提示必须来自两个方面。

首先，与知识层的沟通：一些足够具体以采取行动的要点，同时范围明确，使我们清楚它们适用和不适用的领域。没有这些信息，每次尝试都必须从零开始。

其次，来自可信检查器的结果：代理自身未测量的正确性和时序数据。这些数字是下一步尝试的提示，也是我们唯一应该信任的评分。如果代理自行计时其工作，我们将回到未匹配的比较、可见的测试和残留缓存的老问题。

因此，我们想要的划分比“代理与循环”更精细。赋予代理对内核编写方式的自主权。保持循环作为内存和评估的通道。代理提出建议，循环返回其允许查看的内容，并确认上一个提案是否真正获胜。

Proteus 在其 Gated DeltaNet 路径（一种线性注意力风格的模块）上，为 Qwen 3.5 122B 的部分构建了专用内核，该路径运行在 NVIDIA B200 GPU 上。单个内核的加速效果在 1.8 倍到 5.2 倍之间。

教训是：生成是廉价的步骤。验证和上下文管理才是困难的部分。这正是需要精心设计时间和创新的地方。

代理式 GPU 内核生成释放了极端专业化的巨大潜力，但构建可靠、生产就绪的框架仍然是一个具有挑战性的前沿领域。我们正在解决人工智能与系统交叉处的最困难问题，并正在寻找勇敢的工程师加入我们，共同塑造高效推理的未来。如果你热衷于突破可能性的边界，我们正在招聘！

在邮箱中获取最新文章

订阅我们的博客，获取最新文章直接发送到你的邮箱。

注册

查看所有博客

slice-start id="_gatsby-scripts-1"

slice-end id="_gatsby-scripts-1"