比英伟达工程师还熟练！DeepSeek R1+测试时Scaling自动优化GPU内核

K1qDZ7VzZ

本周英伟达的一篇技术博客引发了业界震动！
英伟达的团队尝试利用DeepSeek-R1和推理时扩展实现GPU内核生成自动化，效果极佳。
随着AI模型的扩展，推理时扩展（inference-time scaling），也叫测试时扩展（test-time scaling）正闪亮登场。
推理时扩展就像是给AI模型配备了一个「智慧锦囊」。当模型进行推理时，它会额外分配计算资源，让模型有能力评估多种可能的结果，然后从中挑选出最优解。
通过这种方式，AI可以像人类一样，有条不紊地剖析复杂问题，找到最佳解决方案。

为了充分发挥DeepSeek-R1的优势，同时克服其在生成优化GPU内核时遇到的困难，英伟达的工程师们想出了一个创新的方法——将推理时扩展技术与DeepSeek-R1相结合，构建了一种全新的工作流程。
他们使用DeepSeek-R1，在推理过程中借助额外的计算能力来解决一个复杂问题：旨在自动生成数值正确且针对不同类型注意力机制进行优化的GPU注意力内核，整个过程无需任何显式编程。
在某些情况下，R1生成的内核甚至比娴熟的工程师开发出来的还要出色！
对此，网友评价道：「英伟达是在毁掉自己的护城河吗？」

优化注意力内核的挑战

深入了解推理时扩展技术如何发挥作用之前，先要认识一个LLM中至关重要的概念——注意力机制。
注意力机制就像是「聚光灯」，能让AI模型在处理任务时，有选择地聚焦在输入信息中最相关的部分，快速找到关键语句，做出更准确的预测，发现数据中的隐藏模式。
但是，注意力操作的计算复杂度与输入序列长度的平方成正比。输入的文本越长，模型处理起来就会越吃力，不仅计算量大幅增加，还可能出现运行时错误，比如内存不足的情况。
为了避免这些问题，提高计算效率，开发优化的底层实现，也就是GPU内核，十分必要。
另外，注意力机制有多种不同的变体，像因果注意力、相对位置嵌入、alibi等。工程师们在面对不同的任务时，往往需要把这些变体组合起来使用。
在多模态模型，比如视觉Transformer中，需要专门的注意力机制，像空间邻域注意力（Spatial Neighborhood Attention），来处理计算机视觉、视频生成模型中常见的时空信息。

开发一个优化的GPU注意力内核，对经验丰富的软件工程师来说，也是一项艰巨的任务，费时费力。
尽管像DeepSeek-R1这样的模型在代码生成任务中展现出了很大的潜力，但它们在第一次尝试生成优化代码时，还是会遇到不少挑战。
这使得在推理时必须使用其他策略来生成优化代码。
以下是为相对位置嵌入注意力内核输入的示例用户提示。

模型有时会产生幻觉，生成一些「不靠谱」的代码，要么在语法上有错误，要么把不同语言或框架的语法混在一起，导致无法运行或效率低下。
计算最优的GPU线程映射也不是一件容易的事，通常需要反复调整优化，才能得到一个正确又高效的内核。
DeepSeek R1与推理时扩展「强强联合」

为了攻克优化GPU注意力内核这个难题，英伟达的工程师们想出了一个巧妙的办法，他们把DeepSeek-R1模型和推理时扩展技术结合起来，创造了一种新的工作流程。

一开始，工程师会手动输入一个提示。然后，DeepSeek-R1会根据这个提示，在第一次遍历中生成GPU代码，也就是内核代码。
生成的代码会交给一个特殊的验证器，这个验证器运行在英伟达H100 GPU上，仔细分析生成的内核代码。
如果发现代码有不足的地方，验证器就会生成新的提示，再把这些提示作为输入，反馈给DeepSeek-R1。模型根据新的提示，对代码进行改进，如此循环往复。
工程师们发现，这个过程持续15分钟，就能得到一个性能更好的注意力内核。

根据斯坦福大学的KernelBench基准测试，它生成的内核在处理Level-1问题时，在数值上100%是正确的。在处理Level-2问题时，正确率也能达到96% 。
KernelBench Level-1问题解决率，是用来评估LLM为特定计算任务生成高效GPU内核能力的数值正确指标，它是一系列测试LLM GPU编程能力挑战的一部分。
推理时间预算对生成正确内核的影响也很明显。从测试结果来看，在Level-1问题中，如果每个问题分配的推理时间超过10分钟，就能为100个问题中的大多数生成数值正确的代码。
这也意味着，给模型足够的「思考时间」，它真的能给出更好的答案。

利用DeepSeek-R1模型，通过在推理时投入更多计算资源，可以生成比一些熟练工程师开发的优化内核还要好的结果，这为GPU内核的自动化生成开辟了一条新的道路。
目前这项技术还处于早期研究阶段，虽然已经取得了一些令人振奋的成果，但要想让它能稳定地产生更好的结果，还有很多工作要做。
研究者对DeepSeek-R1的最新进展及潜力感到兴奋。