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

随着 AI 模型扩展其功能以解决更复杂的挑战，一种称为“ 测试时扩展 ”或“ 推理时扩展 ”的新扩展法则正在出现。该技术也称为 AI 推理 或 长时思考 技术，通过在推理过程中分配额外的计算资源来评估多种可能的结果，然后选择最佳结果 (neural network)，从而提高模型性能。这使得 AI 能够以类似于人类剖析复杂问题并单独解决这些问题以达成最终解决方案的方式，制定战略并系统化地解决复杂问题。

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

事实证明，在某些情况下，最终结果优于由技术精湛的工程师开发的优化内核。

优化注意力内核的需求和相关挑战

注意力是一个关键概念，彻底改变了大语言模型（LLM）的发展。这是一种功能强大的机制，可让 AI 模型在执行任务时，有选择地专注于与输入内容最相关的部分。通过聚焦于重要信息，注意力运算可帮助模型做出更好的预测，并找到数据中隐藏的模式。

相对于输入序列长度，注意力运算的计算复杂性呈正交增长。这促使我们需要开发经过优化的低级实现 (即 GPU 内核)，以防止简单实现产生的运行时错误 (例如内存不足的错误)，并提高计算效率。

注意力有多种变体 (因果关系、相对位置嵌入、不在场证明等)，工程师通常必须将这些变体的组合用于给定任务。

多模态模型 (例如，vision transformers) 带来了额外的一层挑战，因为它们需要专门的注意力机制 (Spatial Neighborhood Attention) 来维护计算机视觉、视频生成模型等领域中经常遇到的时空信息。

Four images of people watching ski jumps have different dilation values, from 1-4. The dilation expands the bounding box from just one skier (dilation=1) to almost the entire image (dilation=4). — *图 1. 邻域注意力在 2D 输入上的应用*

创建经优化的 GPU 内核以供关注需要大量技能和时间，即使对于经验丰富的软件工程师而言也是如此。

最近的 LLMs（如 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 线程映射也并非易事，也是一项具有挑战性的任务，通常需要进行迭代优化才能实现正确高效的内核。