通过 DeepSeek-R1 和推理时间扩展自动生成 GPU Kernel [译]
随着 AI 模型在解决更复杂问题方面的能力不断扩展，一种新的扩展法则应运而生，这种方法被称为推理时间扩展（inference-time scaling）或测试时间扩展（test-time scaling）。这种技术也被称为 AI 推理（AI reasoning）或长时间思考（long-thinking），通过在推理期间分配额外的计算资源，此技巧能够评估多个可能的结果并选择最优解。此举使 AI 能够像人类一样采用系统性方式剖析并逐步解决复杂问题，从而得出最终答案。
在本文中，我们将介绍 NVIDIA 工程师的一项实验：他们使用最新的开源模型之一——DeepSeek-R1，结合推理阶段的额外计算能力，来解决一个复杂问题。这次实验的目标是自动生成数值正确且针对不同 Attention 模式优化的 GPU Attention Kernel，而无需显式编程。
结果显示，在某些情况下，这些自动生成的 Kernel 优化效果甚至优于熟练工程师开发的优化 Kernel。
对优化的 Attention Kernel 的需求及其相关挑战
Attention 是一个关键概念，它革新了大型语言模型（Large Language Model，LLM）的发展。它是一种强大的机制，能够让 AI 模型在执行任务时有选择性地专注于输入中最相关的部分。通过关注关键信息，Attention 操作能够帮助模型作出更好的预测，并发现数据中的隐藏模式。
Attention 的计算复杂度与输入序列长度呈二次方增长。这就需要开发优化的底层实现（即 GPU Kernel），以防止简单实现引发的运行时错误（例如内存不足错误）以及提高计算效率。
Attention 存在多种变体（例如因果 Attention、相对位置嵌入、Alibi 等），工程师经常需要结合这些变体来完成特定任务。
多模态模型（例如视觉 Transformer）引入了额外的挑战，因为它们需要专门的 Attention 机制（如空间邻域 Attention），以维持计算机视觉、视频生成模型等中常见的时空信息。
图 1. 2D 输入的邻域 Attention
即使是经验丰富的软件工程师，创建一个优化的 GPU Attention Kernel 也需要耗费大量时间和技术。
像 DeepSeek-R1 这样的最新 LLM 虽然在代码生成任务中表现出很大的潜力，但在首次尝试中生成优化代码仍然面临挑战。因此，在推理阶段引入其他策略生成优化代码是必要的。
以下是一个相对位置嵌入 Attention Kernel 的用户输入示例：
请编写一个支持相对位置编码的 GPU Attention Kernel。在 Kernel 内动态实现相对位置编码。应返回完整代码，包括必要的修改。
使用以下函数计算相对位置编码：
def relative_positional(score, b, h, q_idx, kv_idx):
return score + (q_idx - kv_idx)
在实现 Kernel 的过程中，需记住对相对位置编码应用一个常数缩放因子 1.44269504，该缩放因子公式为：qk_scale = sm_scale * 1.44269504。在 PyTorch 参考实现中不需要缩放相对位置编码，但在 GPU Kernel 中，请使用：
qk = qk * qk_scale + rel_pos * 1.44269504
请提供包括这些更改的完整更新版 Kernel 代码，确保在 Kernel 操作中高效应用相对位置编码。
LLM 有时会生成幻觉代码或混用不同语言或框架的语法，导致代码错误或低效。计算 GPU 线程的最佳映射也并不是一项简单任务，通常需要反复改进以实现正确且高效的 Kernel。
推理时间扩展技术用于生成优化 GPU Kernel
为了在优化 Attention Kernel 的生成过程中获得最佳结果，NVIDIA 工程师创建了一种新的工作流，该工作流在推理过程中将 DeepSeek-R1 模型与一个特殊的验证器结合以实现闭环优化，并持续预设时长运行。
图 2. 在 NVIDIA Hopper 平台上使用 DeepSeek-R1 的推理时间扩展
这个工作流首先通过一个手工输入初始化，DeepSeek-R1 模型在初次迭代中生成 GPU 代码（即 Kernel）。随后，一个验证器在 NVIDIA H100 GPU 上运行，分析生成的 Kernel，创建新的输入提示并反馈至 DeepSeek-R1 模型。
这种闭环方法通过每次以不同策略引导代码生成过程来不断改进代码。团队发现，如果这个过程持续 15 分钟，Attention Kernel 的质量会显著提高。
图 3. 使用灵活 Attention 自动生成的优化 Attention Kernel 性能
这一工作流在 KernelBench 基准测试中对 Level-1 问题的数值正确率达到了 100%，对于 Level-2 问题则达到了 96%。
KernelBench 的 Level-1 解决率是一项衡量 LLM 为特定计算任务生成高效 GPU Kernel 的数值正确性的指标。这项测试是最新 LLM 在 GPU 编程能力评估中的重要环节之一。
图 4 显示了推理时间预算如何影响代理的解决率。在 Level-1 类别中，针对每个问题分配超过 10 分钟的推理时间，能够使工作流生成数值正确的代码覆盖大多数的 100 个问题。
图 4. 推理时间扩展结果：优化的 GPU Kernel
基于 DeepSeek-R1 的优化 GPU Kernel
这些结果展示了如何利用最新的 DeepSeek-R1 模型，在推理阶段增加计算能力以生成更佳的 GPU Kernel。这仍然是一个新兴研究领域，当前的结果表明其在自动生成高效 Attention Kernel 方面的潜力。
尽管我们刚刚迈出第一步，但要为更多种类的问题持续生成更优结果还需付出更多努力。我们对 DeepSeek-R1 最近的发展及其潜力感到兴奋。
要了解更多信息或立即开始使用，请参阅 DeepSeek-R1 NIM microservice，现已在 build.nvidia.com 提供。