OpenAI 开源 Triton语言：取代英伟达的 CUDA

知名AI研究实验室OpenAI LLC今天发布了Triton；它声称，这种类似Python的专门编程语言可使开发人员能够更轻松自如地开发高速机器学习算法。

两年前，OpenAI的科学家Philippe Tillet就在一篇学术论文中介绍了Triton的第一个版本。作为今天重大发布的一部分内容，OpenAI推出了大幅升级的版本：Triton 1.0，针对企业机器学习项目进行了诸多优化。

深度神经网络已成为一种很重要的AI模型，能够在自然语言处理、计算机视觉及其他领域获得最先进的性能。这种模型的优势在于其层次结构，这种层次结构生成大量的高度并行化工作，非常适合GPU之类的多核硬件来处理。近些年来，用于通用GPU计算的框架（比如CUDA和OpenCL）使得开发团队更容易开发高性能程序。然而，GPU优化起来却仍然特别困难，一方面是由于其架构在迅速演进。

绝大多数的企业AI模型在英伟达的图形处理单元（GPU）上运行。开发人员使用英伟达提供的软件来构建这些模型。英伟达最重要的框架之一是CUDA框架，它提供了底层的软件构建模块，AI应用程序使用这些构建模块借助GPU进行计算。

OpenAI正在用Triton处理的问题是，CUDA框架被认为用起来难度很大。尤其是，面临的主要挑战是最大限度地提升AI模型的性能，以便它尽可能快地处理数据。对于使用CUDA的开发团队来说，最大限度地提升AI性能需要对代码进行复杂而精细的优化，即使开发人员拥有多年的丰富经验，做这项工作也很难。

这时OpenAI的Triton编程语言闪亮登场。据这家实验室声称，该语言可自动执行许多 AI代码优化，为开发人员节省时间。

OpenAI承诺可为软件团队带来两大好处。首先是，Triton可以加速AI项目，因为开发人员可以减少用来优化代码的时间。另一大好处是，据OpenAI声称，Triton相对简单，连没有丰富CUDA编程经验的软件团队都能够创建比其他方法更高效的算法。

OpenAI的Tillet在今天的一篇博文中解释：“有了Triton ，就可以花较少的精力获得最高的硬件性能。比如说，它可用于编写与cuBLAS性能相匹配的FP16矩阵乘法内核——这是许多GPU程序员无法做到的，仅用不到25行代码。”矩阵乘法内核是一种软件机制，机器学习算法高度依赖这种机制来执行计算。

Triton提升AI性能的诀窍是，优化在英伟达芯片上运行的机器学习算法处理数据所用的工作流程的三个核心步骤。

第一个步骤是在GPU的DRAM存储器电路和SRAM存储器电路之间移动数据这项任务。GPU在信息未频繁使用时将其存储在DRAM中，然后将信息传输到SRAM内存以执行计算。两个部件之间的数据传输速度越快，机器学习算法运行的速度就越快，这就是为什么开发人员将优化计算工作流程的这个方面视作AI项目的重要部分。

优化过程包括将从DRAM移动到SRAM 的数据块合并为较大的信息单元。OpenAI表示，Triton 自动执行这项任务，从而为开发人员节省时间。

Triton优化的第二个计算步骤是将收到的数据块分布到GPU的SRAM电路中，以便尽可能快地分析它们。

这个步骤中遇到的主要挑战之一是，避免所谓的存储体冲突（memory bank conflict）。这个术语是指这种情况：两个软件无意中试图将数据写入到同一段内?6?7?6?7存中。存储体冲突会阻碍计算，直至冲突得到解决，这意味着通过减少这类错误发生的频次，开发人员可以提升AI算法的性能。

Tillet解释道：“必须在重新使用之前将数据手动存储到SRAM中，并加以认真管理，以便最大限度地减少检索时发生的共享存储体冲突。”

Triton帮助实现自动化的第三项也是最后一项任务并不涉及GPU的内存单元，而是其CUDA核心，即负责对存储在内存中的数据进行计算的计算电路。单单一块英伟达数据中心GPU就有数千个这样的电路。它们让芯片可以同时执行大量计算。

为了最大限度地提升AI模型的性能，开发人员必须配置AI模型，以便将计算分布在多个CUDA核心上，从而可以同时执行计算，而不是一个接一个执行。Triton也可以自动完成这项工作，不过只是局部自动化。它没有自动化整个工作流程的原因是，OpenAI试图让开发人员可以地灵活地根据需要为项目手动定制流程。

Triton的源代码可从GitHub 获得（https://github.com/openai/triton）。