使用 DeepSeek-R1 和推理时间缩放自动生成 GPU 内核

大家好，这里是架构资源栈！点击上方关注，添加“星标”，一起学习大厂前沿架构！

随着人工智能模型不断扩展其能力以解决更复杂的挑战，一种称为_测试时间缩放_或_推理时间缩放的新解决_缩放解决方案正在出现。这种技术也称为_人工智能推理_或_长期思考_，通过在推理过程中分配额外的计算资源来评估可能的多种结果，然后选择最佳结果（神经网络），从而提高模型性能。这使得人工智能能够以人类剖析复杂问题并单独解决它们，并以最终方案的方式制定策略并系统地复杂问题。

在这篇文章中，我们讨论了 NVIDIA 工程师所做的一个实验，他们使用最新的开源模型之一 DeepSeek-R1 模型以及推理过程中的额外计算能力来解决一个复杂的问题。该实验旨在自动生成 GPU 注意力内核，这些内核的分数正确，并针对不同类型的注意力进行了优化，而无需任何显着式编程。

在某些情况下，结果比熟练的工程师开发的优化内核更好。

优化注意力核心的需求以及相关挑战

注意力机制是彻底改变大型语言模型（LLM）发展的关键概念。它是一种强大的机制，使 AI模型能够在执行任务时选择有地关注输入中最相关的部分。通过关注重要信息，注意力机制可帮助模型做出更好的预测并发现数据中的模式。

注意力操作的计算复杂度与输入序列长度呈二次增长方。这促使我们需要开发一个优化的底层实现（即GPU内核），以防止由简单实现引起的运行时错误（如内存不足错误）并提高计算效率。

焦点有多种变体（因果、相对位置嵌入、不在场证明等），工程师通常给定的任务使用这些变体的组合。

多模态模型（例如视觉变换器）带来了额外的挑战，因为它们需要专门的注意力机制（空间邻域注意力）来维护计算机、视频生成模型等经常遇到的时空信息。

四张观看滑雪跳跃者的图像具有不同的血管扩张值，从1到4。血管扩张将边界框从仅一个滑雪者（血管扩张=1）分裂几乎整个图像（血管扩张=4）。

图1. 二维输入上的邻域视角

即使对于经验丰富的软件工程师来说，创建针对注意力机制来优化 GPU 内核也需要大量的技巧和时间。

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

以下提示是相对位置嵌入注意核对的用户输入的示例。

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线程映射也是一项艰巨而艰巨的任务，通常需要迭代细化才能获得正确且高效的内核。

用于生成 GPU 内核优化的推理时间扩展

为了通过优化注意力内核获得最佳效果，NVIDIA 工程师创建了一种新的工作流程，其中包括特殊的验证器以及 DeepSeek-R1 模型，在预定期限以闭环方式进行推理。

流程图显示了生成注意内核的初始提示。DeepSeek-R1 创建经过验证的 GPU 代码。如果不满足标准，Hopper GPU 会改进并清理提示。最终结果是 GPU 优化的内核。

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

工作流程首先由手动提示初始化，DeepSeek-R1模型在第一遍中生成GPU代码（即内核）。验证器在NVIDIA H100 GPU上运行。它分析生成的内核并创建新的提示，作为 DeepSeek-R1模型的输入。

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

条形图显示了 Hopper GPU 上的平均注意力内核加速，比较了两种方法中不同注意力内核类型的加速：橙色表示“PyTorch API（Flex Attention）”，绿色表示“搭载 DeepSeek-R1 的 NVIDIA Workflow”。PyTorch API 为所有内核保持 1 倍的注意力，而搭载 DeepSeek-R1 的 NVIDIA Workflow 为 Causal Mask 和 Document Mask 实现了 1.1倍的加速，为相对位置实现了 1.5 倍的加速，为 Alibi Bias 和 Full Mask 实现了 1.6 倍的加速，为 Softcap 实现了 2.1 倍的加速。