Automating GPU Kernel Generation with DeepSeek-R1 and Inference Time Scaling
As AI models extend their capabilities to solve more sophisticated challenges, a new
scaling law
known as
test-time scaling
or
inference-time scaling
is emerging. Also known as
AI reasoning
or
long-thinking
, this technique improves model performance by allocating additional computational resources during inference to evaluate multiple possible outcomes and then selecting the best one, neural network. This enables AI to strategize and systematically solve complex problems in a similar fashion to how humans dissect complex problems and solve them individually to arrive at a final solution.
In this post, we talk about an experiment done by NVIDIA engineers who used one of the newest open-source models, the DeepSeek-R1 model, together with additional computing power during inference to solve a complex problem. The experiment was to automatically generate GPU attention kernels that were numerically correct and optimized for different flavors of attention without any explicit programming.
The results turned out to be better than the optimized kernels developed by skilled engineers in some cases.
The need for optimized attention kernels and associated challenges
Attention is a key concept that revolutionized the development of the large language model (LLM). It’s a powerful mechanism that enables AI models to focus selectively on the most relevant parts of input when performing tasks. By focusing on important information, the attention operation helps the models make better predictions and find hidden patterns in the data.
The computational complexity of the attention operation grows quadratically in relation to the input sequence length. This motivates the need for developing an optimized lower-level implementation (that is, a GPU kernel) to prevent runtime errors arising from simple implementations (for example, out-of-memory errors) and for computational efficiency purposes.
There are multiple variants of attention (causal, relative positional embeddings, alibi, and so on) and often engineers must use a combination of these variants for a given task. ‌
Multi-modal models (for example, vision transformers) introduce an additional layer of challenges as they require specialized attention mechanisms (Spatial Neighborhood Attention) for maintaining spatio-temporal information often encountered in computer vision, video generation models, and so on.
Figure 1. Neighborhood attention on 2D inputs
Creating an optimized GPU kernel for attention takes a lot of skill and time, even for experienced software engineers. ‌
Recent LLMs like DeepSeek-R1 have shown a lot of promise in code generation tasks, but they still face challenges creating optimized code on the first try. This makes it necessary to use other strategies at inference time to generate optimized code.
The following prompt is sample user input for a relative positional embeddings 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.
LLMs can occasionally produce hallucinated code or mix syntax from different languages or frameworks, causing immediate code errors or inefficiencies. Computing the optimal GPU thread mapping is also non-trivial and a challenging task, often requiring iterative refinement to achieve a correct and efficient kernel.
Inference-time scaling for generating optimized GPU Kernels
To get the best results with optimized attention kernels, NVIDIA engineers created a new workflow that includes a special verifier along with the DeepSeek-R1 model during inference in a closed-loop fashion for a predetermined duration.
Figure 2. Inference-time scaling with DeepSeek-R1 on the NVIDIA Hopper platform
The workflow is first initialized by a manual prompt and the DeepSeek-R1 model generates the GPU code (that is, the kernel) in the first pass. The verifier runs on an NVIDIA H100 GPU. It analyzes the generated kernel and creates new prompts that are provided as ‌input to the DeepSeek-R1 model.
This closed-loop approach makes the code generation process better by guiding it in a different way each time. The team found that by letting this process continue for 15 minutes resulted in an improved attention kernel.
Figure 3. Performance of automatically generated optimized attention kernels with flex attention
This workflow produced numerically correct kernels for 100% of Level-1 problems and 96% of Level-2 problems, as tested by
Stanford’s KernelBench
benchmark. ‌
The Level-1 solving rate in KernelBench refers to the numerical correct metric used to evaluate the ability of LLMs to generate efficient GPU kernels for specific computational tasks. This test is part of a series of challenges to test the latest LLMs’ abilities in GPU programming.
Figure 4 shows how the inference-time budget affects the agent’s solving rate. Allocating more than 10 minutes per problem in the Level-1 category enables the workflow to produce numerical correct code for most of the 100 problems.
Figure 4. Inference-time scaling results in optimized GPU kernels
Optimized GPU kernels on DeepSeek-R1
These results show how you can use the latest DeepSeek-R1 model to give better GPU kernels by using more computing power during inference time. This is still a new research area with early results on a promising approach that automatically generates effective attention kernels.
While we are off to a good start, more work is needed to generate better results consistently for a wider variety of problems. We’re excited about the recent developments in DeepSeek-R1 and its potential.
For more information or to get started, see the
DeepSeek-R1 NIM microservice
, now available on
build.nvidia.com
.
使用 DeepSeek-R1 和推理时间缩放实现自动化 GPU 内核生成
随着 AI 模型扩展其功能以解决更复杂的挑战，一种称为“
测试时扩展
”或“
推理时扩展
”的新扩展法则正在出现。该技术也称为
AI 推理
或
长时思考
技术，通过在推理过程中分配额外的计算资源来评估多种可能的结果，然后选择最佳结果 (neural network)，从而提高模型性能。这使得 AI 能够以类似于人类剖析复杂问题并单独解决这些问题以达成最终解决方案的方式，制定战略并系统化地解决复杂问题。
在本文中，我们将介绍 NVIDIA 工程师完成的一项实验，他们在推理过程中使用最新的开源模型之一 DeepSeek-R1 模型以及额外的计算能力来解决复杂的问题。该实验旨在自动生成 GPU 注意力内核，这些内核在数值上是正确的，并针对不同的注意力类型进行了优化，而无需任何显式编程。
事实证明，在某些情况下，最终结果优于由技术精湛的工程师开发的优化内核。
优化注意力内核的需求和相关挑战
注意力是一个关键概念，彻底改变了大语言模型（LLM）的发展。这是一种功能强大的机制，可让 AI 模型在执行任务时，有选择地专注于与输入内容最相关的部分。通过聚焦于重要信息，注意力运算可帮助模型做出更好的预测，并找到数据中隐藏的模式。
相对于输入序列长度，注意力运算的计算复杂性呈正交增长。这促使我们需要开发经过优化的低级实现 (即 GPU 内核)，以防止简单实现产生的运行时错误 (例如内存不足的错误)，并提高计算效率。
注意力有多种变体 (因果关系、相对位置嵌入、不在场证明等)，工程师通常必须将这些变体的组合用于给定任务。
多模态模型 (例如，vision transformers) 带来了额外的一层挑战，因为它们需要专门的注意力机制 (Spatial Neighborhood Attention) 来维护计算机视觉、视频生成模型等领域中经常遇到的时空信息。
图 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 线程映射也并非易事，也是一项具有挑战性的任务，通常需要进行迭代优化才能实现正确高效的内核。
用于生成经过优化的 GPU 内核的推理时间扩展
为了通过优化的注意力内核获得最佳结果，NVIDIA 工程师创建了一个新的工作流程，其中包括一个特殊的验证器以及 DeepSeek-R1 模型，在预先设定的时间内以闭环方式进行推理。
图 2、在 NVIDIA Hopper 平台上使用 DeepSeek-R1 实现推理时间扩展
该工作流程首先通过手动提示进行初始化，然后 DeepSeek-R1 模型会在首次通道中生成 GPU 代码(即核函数)。该验证器在 NVIDIA H100 GPU 上运行。它会分析生成的核函数并创建新的提示，这些提示以 input 的形式提供给 DeepSeek-R1 模型。
这种闭环方法每次都以不同的方式指导代码生成过程，从而改进代码生成过程。该团队发现，让此过程持续 15 分钟可以改进注意力核函数。
图 3、具有 Flex Attention 的自动生成优化注意力内核的性能
此工作流程为 100%的 1 级问题和 96%的 2 级问题生成了数值正确的内核，测试对象为
斯坦福大学的 KernelBench
基准测试。*
KernelBench 中的 1 级求解率是指用于评估 LLM 为特定计算任务生成高效 GPU 内核的能力的数字正确指标。本次测试属于一系列挑战，旨在测试最新 LLM 在 GPU 编程中的能力。
图 4 显示了推理时间预算如何影响智能体的求解率。在 Level-1 类别中为每个问题分配超过 10 分钟的时间，使工作流程能够为 100 个问题中的大多数生成正确的数字代码。
图 4、在优化的 GPU 内核中实现推理 – 时间扩展
DeepSeek-R1 上经过优化的 GPU 内核
这些结果展示了如何使用最新的 DeepSeek-R1 模型，通过在推理期间使用更强大的计算能力来提供更好的 GPU 内核。这仍然是一个新的研究领域，在自动生成有效注意力内核的前景良好的方法方面取得了早期成果。
虽然我们已经有了一个良好的开端，但我们需要做更多的工作，以便为更广泛的问题持续提供更好的结果。我们对 DeepSeek-R1 的最新进展及其潜力感到兴奋。
如需了解更多信息或入门，请参阅
DeepSeek-R1 NIM 微服务
(现已在 build.nvidia.com 上提供)。