随着 AI 模型将其能力扩展到解决更复杂挑战，一种新的扩展定律，即所谓的 测试时扩展 或 推理时扩展 正在兴起。这项技术也称为 AI 推理 或 长思考，通过在推理期间分配额外的计算资源来评估多种可能的结果，然后选择最佳结果，即神经网络，从而提高模型性能。这使 AI 能够像人类剖析复杂问题并单独解决它们以得出最终解决方案一样，制定策略并系统地解决复杂问题。

在这篇文章中，我们讨论了 NVIDIA 工程师进行的一项实验，他们使用最新的开源模型之一 DeepSeek-R1 模型，以及推理期间的额外计算能力来解决一个复杂的问题。该实验旨在使用自动化的方式生成在数值上正确且针对不同 Attention 变体优化的 GPU Attention Kernel，而无需任何显式编程。

结果表明，在某些情况下，其效果甚至优于经验丰富的工程师开发的优化 Kernel。

对优化 Attention Kernel 的需求以及相关挑战

Attention 是彻底革新大型语言模型 (LLM) 发展的关键概念。它是一种强大的机制，使 AI 模型能够在执行任务时有选择地专注于输入中最相关的部分。通过关注重要信息，Attention 操作有助于模型做出更好的预测，并在数据中找到隐藏的模式。

Attention 操作的计算复杂度与输入序列长度呈二次方增长关系。这就促使人们需要开发优化的底层实现（即 GPU Kernel），以防止简单实现（例如，内存不足错误）导致运行时错误，并提高计算效率。

Attention 有多种变体（因果、相对位置嵌入、Alibi 等），工程师通常必须针对给定的任务使用这些变体的组合。

多模态模型（例如，Vision Transformer）引入了额外的挑战层，因为它们需要专门的 Attention 机制（空间邻域 Attention）来维护在计算机视觉、视频生成模型等中经常遇到的时空信息。

Image 25: 四张人们观看跳台滑雪的图像具有不同的膨胀值，从 1-4。膨胀从仅一个滑雪者（膨胀=1）扩展到几乎整个图像（膨胀=4）。

图 1. 2D 输入上的邻域 Attention

即使对于经验丰富的软件工程师来说，创建优化的 GPU Kernel 也需要大量的技能和时间。

最近的 LLM（如 DeepSeek-R1）在代码生成任务中表现出了很大的潜力，但它们在首次尝试时创建优化的代码仍然面临挑战。这使得在推理时使用其他策略来生成优化的代码成为必要。

以下 Prompt 是用于相对位置嵌入 Attention Kernel 的示例用户输入。

Please write a GPU attention kernel to support relative position encodings. Implement the relative positional encoding on the fly within the kernel. The complete code should be returned, including the necessary modifications.

Use the following function to compute the relative positional encoding:

def relative\_positional(score, b, h, q\_idx, kv\_idx):

    return score + (q\_idx - kv\_idx)

When implementing the kernel, keep in mind that a constant scaling factor 1.44269504 should be applied to the relative positional encoding due to qk\_scale = sm\_scale \* 1.44269504. The PyTorch reference does not need to scale the relative positional encoding, but in the GPU kernel, use:

qk = qk \* qk\_scale + rel\_pos \* 1.44269504

Please provide the complete updated kernel code that incorporates these changes, ensuring that the relative positional encoding is applied efficiently within the kernel operations.

LLM 有时会产生幻觉代码或混合来自不同语言或框架的语法，从而导致立即出现代码错误或效率低下。计算最佳 GPU 线程映射也是一项非常复杂且具有挑战性的任务，通常需要迭代改进才能获得正确且高效的 Kernel。

用于生成优化 GPU Kernel 的推理时扩展

为了使用优化的 Attention Kernel 获得最佳结果，NVIDIA 工程师创建了一个新的工作流程，该工作流程在预定的持续时间内，以闭环方式在推理期间包含了一个特殊的验证器以及 DeepSeek-R1 模型。

Image 26: 一个流程图显示了生成 Attention Kernel 的初始 Prompt。DeepSeek-R1 创建 GPU 代码，然后进行验证。如果未满足标准，Hopper GPU 会优化并格式化 Prompt。最终结果是 GPU 优化的 Kernel。

图 2. 在 NVIDIA Hopper 平台上使用 DeepSeek-R1 进行推理时扩展

该工作流程首先通过手动 Prompt 初始化，DeepSeek-R1 模型在第一遍中生成 GPU 代码（即 Kernel）。验证器在 NVIDIA H100 GPU 上运行。它分析生成的 Kernel 并创建新的 Prompt，这些 Prompt 作为输入提供给 DeepSeek-R1 模型。

这种闭环方法通过每次以不同的方式引导代码生成过程，使其变得更好。该团队发现，让这个过程持续 15 分钟可以改进 Attention Kernel。

Image 27: 一个柱状图显示了 Hopper GPU 上 Attention Kernel 平均加速比，比较了两种方法之间不同 Attention Kernel 类型的加速比：“PyTorch API (Flex Attention)”（橙色）和“NVIDIA 工作流程与 DeepSeek-R1”（绿色）。PyTorch API 为所有 Kernel 保持 1 倍的基线，而 NVIDIA 工作流程与 DeepSeek-R1 在因果掩码和文档掩码上实现了 1.1 倍的加速，在相对位置上实现了 1.5 倍的加速，在 Alibi 偏差和完整掩码上实现了 1.6 倍的加速，在 Softcap 上实现了 2.1 倍的加速。

图 3. 使用 Flex Attention 自动生成的优化 Attention Kernel 的性能

该工作流程为 100% 的 Level-1 问题和 96% 的 Level-2 问题生成了数值正确的 Kernel，正如 斯坦福大学的 KernelBench 基准测试所测试的那样。

KernelBench 中的 Level-1 解决率是指用于评估 LLM 为特定计算任务生成高效 GPU Kernel 能力的数值正确指标。此测试是一系列挑战的一部分，旨在测试最新的 LLM 在 GPU 编程方面的能力。

图 4 显示了推理时预算如何影响 Agent 的解决率。在 Level-1 类别中，为每个问题分配超过 10 分钟的时间，该工作流能够为 100 个问题中的大多数问题生成数值正确的代码。

Image 28: 一条折线图显示了随时间生成的数值正确的 Kernel 数量。该线在约 10 分钟时接近 95%，在 20 分钟时接近 100%。

图 4. 推理时扩展在优化 GPU Kernel 中的结果

DeepSeek-R1 上的优化 GPU Kernel

这些结果表明，您可以通过在推理时使用更多计算能力，使用最新的 DeepSeek-R1 模型来提供更好的 GPU Kernel。这仍然是一个新的研究领域，在自动生成有效的 Attention Kernel 的有希望的方法上取得了初步成果。

虽然我们有了一个良好的开端，但还需要做更多的工作，以便为更广泛的问题持续生成更好的结果。我们对 DeepSeek-R1 的最新进展及其潜力感到兴奋。