Loading [MathJax]/jax/output/CommonHTML/config.js
首页
学习
活动
专区
圈层
工具
发布
首页
学习
活动
专区
圈层
工具
MCP广场
社区首页 >专栏 >在CUDA的天下,OpenAI开源GPU编程语言Triton,将同时支持N卡和A卡

在CUDA的天下,OpenAI开源GPU编程语言Triton,将同时支持N卡和A卡

作者头像
代码医生工作室
发布于 2021-08-10 06:27:52
发布于 2021-08-10 06:27:52
1.8K00
代码可运行
举报
文章被收录于专栏:相约机器人相约机器人
运行总次数:0
代码可运行

OpenAI 开源了全新的 GPU 编程语言 Triton,它能成为 CUDA 的替代品吗?

过去十年中,深度神经网络 (DNN) 已成为最重要的机器学习模型之一,创造了从自然语言处理计算机视觉、计算神经科学等许多领域的 SOTA 实现。DNN 模型的优势来自于它的层次结构,这一特征导致其计算量巨大,但也会产生大量高度并行化的工作,特别适合多核和众核处理器。

深度学习领域的新研究思路往往是结合原生框架 operator 来实现的,这种方法虽然方便,但需要创建或移动许多临时张量,因此可能会造成神经网络的性能损失。编写专门的 GPU 内核或许可以解决这个问题,但 GPU 编程的确是一件相当复杂的事。

DNN 计算潜力与 GPU 编程困难之间的矛盾由来已久。英伟达在 2007 年发布了 CUDA 的初始版本,CUDA 平台是一个软件层,使用者可以直接访问 GPU 的虚拟指令集和并行计算单元,用于执行计算内核。近年来,主流深度学习框架几乎都是基于 CUDA 进行加速,英伟达也一直在完善 CUDA 工具包,但对于一般的开发者来说,CUDA 还是「不那么容易上手」。

今天,OpenAI 正式推出 Triton 1.0,这是一种类 Python 的开源编程语言。即使没有 CUDA 经验的研究人员,也能够高效编写 GPU 代码。例如,它可以用不到 25 行代码写出与 cuBLAS 性能相匹配的 FP16 矩阵乘法内核,后者是许多专业的 GPU 编程者尚且无法做到的。此外,OpenAI 的研究者已经使用 Triton 成功生成了比 PyTorch 同类实现效率高 2 倍的内核。

代码地址:https://github.com/openai/triton

Triton 的最初想法来源于现任 OpenAI 科学家的 Philippe Tillet 2019 年在哈佛大学攻读研究生学位时发表的一篇论文,当时他的导师是 H. T. Kung 和 David Cox。

论文链接:http://www.eecs.harvard.edu/~htk/publication/2019-mapl-tillet-kung-cox.pdf

Tillet 希望解决的问题是打造一种比英伟达的 CUDA 等特定供应商库更好用的库,能够处理神经网络中涉及矩阵的各种操作,具备可移植性,且性能可与 cuDNN 或类似的供应商库相媲美。团队表示:「直接用 CUDA 进行 GPU 编程太难了,比如为 GPU 编写原生内核或函数这件事,会因为 GPU 编程的复杂性而出奇困难。」

Facebook AI 研究中心科学家 Soumith Chintala 也在推特上表达了自己对 Triton 的期待:

新发布的 Triton 可以为一些核心的神经网络任务(例如矩阵乘法)提供显著的易用性优势。「我们的目标是使其成为深度学习 CUDA 的可行替代方案,」Philippe Tillet 作为 Triton 项目负责人如此表示。

GPU 编程面临的挑战

现代 GPU 的架构大致可以分为三个主要组件:DRAM、SRAM 和 ALU。优化 CUDA 代码时,必须考虑到每一个组件:

  • 来自 DRAM 的内存传输必须合并进大型事务,以利用现代内存接口的总线位宽;
  • 必须在数据重新使用之前手动存储到 SRAM 中,并进行管理以最大限度地减少检索时共享内存库冲突;
  • 计算必须在流处理器(SM)内部或之间细致分区和调度,以促进指令 / 线程级的并行以及专用算术逻辑单元(ALU)的利用。

GPU 基础架构。

种种因素导致 GPU 编程难度骤增,即使对于具有多年经验的 CUDA 程序员也是如此。Triton 的目的是将这些优化过程自动化,以此让开发人员更专注于并行代码的高级逻辑。出于对泛用能力的考量,Triton 不会自动调度跨流处理器的工作,而是将一些重要的算法考虑因素(例如 tiling、SM 间同步)留给开发者自行决定。

CUDA vs Triton 编译器优化对比。

编程模型

在所有可用的领域专用语言和 JIT 编译器中,Triton 或许与 Numba 最相似:内核被定义为修饰过的 Python 函数,并与实例网格上不同的 program_id 的同时启动。但不同之处值得注意:如下图代码片段所示,Triton 通过对 block 的操作来展示 intra-instance 并行,此处 block 是维数为 2 的幂的数组,而不是单指令多线程(SIMT)执行模型。如此一来,Triton 高效地抽象出了与 CUDA 线程 block 内的并发相关的所有问题(比如内存合并、共享内存同步 / 冲突、张量核心调度)。

Triton 中的向量加法。

虽然这对 embarrassingly 并行(即 element-wise)计算可能没什么帮助,但是可以简化更复杂的 GPU 程序的开发。例如,在融合 softmax 核的情况下,对于每个输入张量 X∈R^M×N 来说,每个实例对给定输入张量的不同行进行归一化。这种并行化策略的标准 CUDA 实现可能难以编写,需要线程之间的显式同步,因为这种策略并发地减少 X 的同一行。而 Triton 很大程度上消除了这种复杂性,每个内核实例加载感兴趣的行,并使用类似 NumPy 的原语顺序对其进行规范化。

代码语言:javascript
代码运行次数:0
运行
AI代码解释
复制
import triton
import triton.language as tl
@triton.jit
def softmax(Y, stride_ym, stride_yn, X, stride_xm, stride_xn, M, N):
    # row index
    m = tl.program_id(0)
    # col indices
    # this specific kernel only works for matrices that 
    # have less than BLOCK_SIZE columns
    BLOCK_SIZE = 1024
    n = tl.arange(0, BLOCK_SIZE)
    # the memory address of all the elements
    # that we want to load can be computed as follows
    X = X + m * stride_xm + n * stride_xn
    # load input data; pad out-of-bounds elements with 0 
    x = tl.load(X, mask=n < N, other=-float('inf'))
    # compute numerically-stable softmax
    z = x - tl.max(x, axis=0)
    num = tl.exp(z)
    denom = tl.sum(num, axis=0)
    y = num / denom
    # write back to Y
    Y = Y + m * stride_ym + n * stride_yn
    tl.store(Y, y, mask=n < N)
import torch
# Allocate input/output tensors
X = torch.normal(0, 1, size=(583, 931), device='cuda')
Y = torch.empty_like(X)
# SPMD launch grid
grid = (X.shape[0], )
# enqueue GPU kernel
softmax[grid](Y, Y.stride(0), Y.stride(1), 
              X, X.stride(0), X.stride(1),
              X.shape[0]    , X.shape[1])

在 Triton 中融合 softmax

Triton JIT 把 X、Y 当作指针而不是张量。最重要的是,softmax 这种特殊实现方式在整个规范化过程中保持 SRAM 中 X 的行不变,从而在适用时最大限度地实现数据重用(约 32K 列)。这与 PyTorch 的内部 CUDA 代码不同,后者使用临时内存使其更通用,但速度明显变慢(见下图)。

融合 softmax、M=4096 的 A100 性能。

Torch (v1.9) JIT 较低的性能突出了从高级张量操作序列自动生成 CUDA 代码的难度。

代码语言:javascript
代码运行次数:0
运行
AI代码解释
复制
@torch.jit.script
def softmax(x):
    x_max = x.max(dim=1)[0]
    z = x - x_max[:, None]
    numerator = torch.exp(x)
    denominator = numerator.sum(dim=1)
    return numerator / denominator[:, None]

融合 softmax 与 Torch JIT

矩阵乘法

能够为元素操作(element-wise operation)和规约操作(reduction operation)编写融合内核是很重要的,但考虑到神经网络中矩阵乘法的重要性,这还不够。事实证明,Triton 在这些方面表现很好,仅用大约 25 行 Python 代码就能达到最佳性能。相比之下,CUDA 效率就没有那么高了。

Triton 中的矩阵乘法。

手写矩阵乘法内核的一个重要优点是它们可以根据需要进行定制,以适应其输入(例如切片)和输出(例如 Leaky ReLU)的融合变换。假如不存在 Triton 这样的系统,那么对于没有出色的 GPU 编程专业知识的开发人员来说,矩阵乘法内核将很难大改。

高级系统架构

Triton 的良好性能得益于以 Triton-IR 为中心的模块化系统架构。Triton-IR 是一种基于 LLVM 的中间表示,多维值块(blocks of values)是其中最重要的东西。

Triton 的高级架构。

@triton.jit 装饰器的工作原理是遍历由 Python 函数提供的抽象语法树(AST),这样一来就能使用通用的 SSA 构造算法实时生成 Triton-IR。生成的 IR 代码随后由编译器后端进行简化、优化和自动并行化,然后转换为高质量的 LLVM-IR,最终转换为 PTX,以便在最新的 NVIDIA GPU 上执行。目前 Triton 还不支持 CPU 和 AMD GPU,但团队表示对二者的支持正在开发中。

编译器后端

研究人员发现通过 Triton-IR 来使用块状程序表示,这种方法允许编译器自动执行各种重要的程序优化。例如,通过查看计算密集型块级操作(例如 tl.dot)的操作数,数据可以自动存储到共享内存中,并使用标准的活跃性分析技术进行数据的分配与同步。

Triton 编译器通过分析计算密集型操作中使用的块变量的活动范围来分配共享内存。

此外,Triton 还可以在 SM 之间以及 SM 之内高效、自动地并行化,前者通过并发执行不同的内核实例来实现,后者通过分析每个块级操作的迭代空间,并将其充分划分到不同的 SIMD 单元来实现。如下所示:

Triton 自动并行化。每个块级操作都定义了一个块级迭代空间,该空间可以自动并行化以利用 SM(Streaming Multiprocessor) 上的可用资源。

本文参与 腾讯云自媒体同步曝光计划,分享自微信公众号。
原始发表:2021-07-30,如有侵权请联系 cloudcommunity@tencent.com 删除

本文分享自 相约机器人 微信公众号,前往查看

如有侵权,请联系 cloudcommunity@tencent.com 删除。

本文参与 腾讯云自媒体同步曝光计划  ,欢迎热爱写作的你一起参与!

评论
登录后参与评论
暂无评论
推荐阅读
编辑精选文章
换一批
【BBuf的CUDA笔记】十三,OpenAI Triton 入门笔记一
2023年很多mlsys工作都是基于Triton来完成或者提供了Triton实现版本,比如现在令人熟知的FlashAttention,大模型推理框架lightllm,diffusion第三方加速库stable-fast等灯,以及很多mlsys的paper也开始使用Triton来实现比如最近刚报道的这个新一代注意力机制Lightning Attention-2:无限序列长度、恒定算力开销、更高建模精度。当然笔者由于目前由于工作需要也需要用Triton,所以就有了这系列Triton学习笔记。本篇文章开始入门一下OpenAI的Triton,然后首先是从Triton介绍博客看起,然后对triton官方实现的vector_add和fused_softmax还有Matmul教程做一个阅读,也就是 https://triton-lang.org/main/getting-started/tutorials/ 这里的前三节,熟悉一下triton编写cuda kernel的语法。
BBuf
2024/01/23
2.9K0
【BBuf的CUDA笔记】十三,OpenAI Triton 入门笔记一
英伟达CUDA太难!OpenAI出手要取代它,新语言性能相当但编程更简单
晓查 发自 凹非寺 量子位 报道 | 公众号 QbitAI 用CUDA为GPU编程实在太难了。 为了让没有CUDA编程经验的人写出和专家效率相当的GPU代码,现在OpenAI推出了一种新的语言和编译器——Triton。 它的难度比CUDA低,但是性能却可与之相媲美。 OpenAI声称: Triton只要25行代码,就能在FP16矩阵乘法shang上达到与cuBLAS相当的性能。 OpenAI的研究人员已经使用Triton,来生成比同等Torch效率高出1倍的内核。 Triton项目的负责人Philippe
量子位
2023/03/10
1.1K0
英伟达CUDA太难!OpenAI出手要取代它,新语言性能相当但编程更简单
25行代码≈SOTA!OpenAI发布Triton编程语言,比PyTorch快2倍
项目负责人Philippe Tillet表示:「我们的目标是让Triton成为深度学习中CUDA的替代品」。
新智元
2021/07/30
1K0
《PytorchConference2023 翻译系列》6-Triton编译器
https://youtu.be/AtbnRIzpwho?si=-lB1VI-SE3hEbVT4
BBuf
2023/12/20
8290
《PytorchConference2023 翻译系列》6-Triton编译器
Triton-Lang在Transformer优化加速中的实践 | 得物技术
众所周知,英伟达(Nvidia)自2006年推出CUDA以来,经过近20年的发展,尤其是经历了以卷积为代表的深度学习和近两年以Transformer为基础的LLM的推动,CUDA编程基本上成为了GPU编程的代名词。CUDA作为GPU的编程语言,不仅使用户能充分发挥Nvidia GPU的高性能的并行计算能力,也逐渐构筑了一个包括硬件、驱动、开发库和编程技巧的完备生态链,从而使CUDA成为了人工智能、高性能计算和云计算中的核心依赖。
得物技术
2025/01/14
2850
Triton-Lang在Transformer优化加速中的实践 | 得物技术
【BBuf的CUDA笔记】十四,OpenAI Triton入门笔记三 FusedAttention
继续Triton的学习,这次来到 https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html 教程。也就是如何使用Triton来实现FlashAttention V2。对于FlashAttention和FlashAttention V2网上已经有非常多的介绍了,大家如果感兴趣的话我推荐FlashAttention V1看 《图解大模型计算加速系列:FlashAttention V1,从硬件到计算逻辑》https://zhuanlan.zhihu.com/p/669926191 这篇文章的讲解 以及 FlashAttention V2 看 《图解大模型计算加速系列:Flash Attention V2,从原理到并行计算》 https://mp.weixin.qq.com/s/5K6yNj23NmNLcAQofHcT4Q ,原理和公式推导都非常清晰,不过想一口气读完还是要花一些精力的。同时你也可以在 https://github.com/BBuf/how-to-optim-algorithm-in-cuda 找到更多相关资料(此外Meagtron-LM,DeepSpeed等训练Infra框架的迅速跟进也说明了FlashAttention这个系列工作影响之大),例如:
BBuf
2024/02/29
2.3K0
【BBuf的CUDA笔记】十四,OpenAI Triton入门笔记三 FusedAttention
OpenAITriton MLIR 第一章 Triton DSL
上一章的反响还不错,很多人都私信催更想看Triton的具体优化有哪些,为什么它能够得到比cuBLAS更好的性能。大家不用急,这也是我为什么要写这一系列文章的初衷,来带着大家从Triton的DSL前端一步一步到最终的machine code生成有一个清晰的理解,从而为大家展示编译在高性能计算中所起到的作用。先来看看openai对Triton所打的广告:
BBuf
2023/08/25
1.2K0
OpenAITriton MLIR 第一章 Triton DSL
【BBuf的CUDA笔记】十四,OpenAI Triton入门笔记二
接着【BBuf的CUDA笔记】十三,OpenAI Triton 入门笔记一 继续探索和学习OpenAI Triton。这篇文章来探索使用Triton写LayerNorm/RMSNorm kernel的细节。
BBuf
2024/02/22
1.1K0
【BBuf的CUDA笔记】十四,OpenAI Triton入门笔记二
CUDA-MODE 课程笔记 第一课: 如何在 PyTorch 中 profile CUDA kernels
一直想系统看一下某个课程系统和科学的学习下 CUDA ,感觉 CUDA-MODE 这个课程能满足我的需求。这个课程是几个 PyTorch 的 Core Dev 搞的,比较系统和专业。不过由于这个课程是 Youtube 上的英语课程,所以要学习和理解这个课程还是需要花不少时间的,我这里记录一下学习这个课程的每一课的笔记,希望可以通过这个笔记帮助对这个课程以及 CUDA 感兴趣的读者更快吸收这个课程的知识。这个课程相比于以前的纯教程更加关注的是我们可以利用 CUDA 做什么事情,而不是让读者陷入到 CUDA 专业术语的细节中,那会非常痛苦。伟大无需多言,感兴趣请阅读本文件夹下的各个课程的学习笔记。
BBuf
2024/07/02
8530
CUDA-MODE 课程笔记 第一课: 如何在 PyTorch 中 profile CUDA kernels
triton 在模型推理中的应用
Triton 是 OpenAI 开发的一种编程语言和编译器,旨在简化 GPU 编程,特别是针对深度学习和高性能计算(HPC)中的自定义内核优化。它的设计目标是让没有 CUDA 经验的开发者也能高效编写高性能的 GPU 代码。以下是 Triton 的核心原理和关键特性:
liddytang
2025/06/20
1460
OpenAI/Triton MLIR 第四章: ROCm-triton配置
最近在整理python-based的benchmark代码,反过来在NV的GPU上又把Triton装了一遍,发现Triton的github repo已经给出了对应的llvm的commit id以及对应的编译细节,然后跟着走了一遍,也顺利的安装成功,只需要按照如下方式即可完成NV GPU上的安装,
BBuf
2024/02/29
1.1K0
OpenAI/Triton MLIR 第四章: ROCm-triton配置
OpenAI 开源 Triton语言:取代英伟达的 CUDA
知名AI研究实验室OpenAI LLC今天发布了Triton;它声称,这种类似Python的专门编程语言可使开发人员能够更轻松自如地开发高速机器学习算法。 两年前,OpenAI的科学家Philippe Tillet就在一篇学术论文中介绍了Triton的第一个版本。作为今天重大发布的一部分内容,OpenAI推出了大幅升级的版本:Triton 1.0,针对企业机器学习项目进行了诸多优化。 深度神经网络已成为一种很重要的AI模型,能够在自然语言处理、计算机视觉及其他领域获得最先进的性能。这种模型的优势在于其层次结
云头条
2022/03/18
2.3K0
pytorch v2.7.0震撼发布!Blackwell GPU支持+编译性能狂飙,AI开发者必看10大升级
NCCL 初始化时在 12.2 驱动下出现 CUDA “无效参数”失败 部分使用 12.2 CUDA 驱动(版本 535)的用户报告在 NCCL 或对称内存初始化过程中遇到“CUDA 驱动错误:无效参数”的问题。该问题正在调查中,详情见 #150852。如果您是从源码编译的 PyTorch,已知的解决方法是使用 CUDA 12.2 工具包重新编译 PyTorch。否则,您可以尝试升级系统中的 CUDA 驱动。
福大大架构师每日一题
2025/04/28
5810
pytorch v2.7.0震撼发布!Blackwell GPU支持+编译性能狂飙,AI开发者必看10大升级
ChatGPT专题|做出ChatGPT的OpenAI,是如何打破英伟达在机器学习领域的垄断地位的?
在机器学习领域,无论是硬件还是软件,英伟达无疑均拥有巨大优势,后者用 CUDA 建立起了一道软件的护城河。可惜的是,这家公司缺乏远见,未能利用其在机器学习硬软件方面的巨大优势,让自己成为机器学习默认的编译器。而它对可用性与易用性的忽视,让 OpenAI 与 Meta 得以趁虚而入,其主导地位正在被打破。
用户9861443
2023/02/26
7790
ChatGPT专题|做出ChatGPT的OpenAI,是如何打破英伟达在机器学习领域的垄断地位的?
Python CUDA 编程 - 6 - 共享内存
GPU的内存结构如图所示:GPU的计算核心都在Streaming Multiprocessor(SM)上,SM里有计算核心可直接访问的寄存器(Register)和共享内存(Shared Memory);多个SM可以读取显卡上的显存,包括全局内存(Global Memory)。
为为为什么
2022/08/04
1.8K0
Python CUDA 编程 - 6 - 共享内存
快来操纵你的GPU| CUDA编程入门极简教程
2006年,NVIDIA公司发布了CUDA(http://docs.nvidia.com/cuda/),CUDA是建立在NVIDIA的CPUs上的一个通用并行计算平台和编程模型,基于CUDA编程可以利用GPUs的并行计算引擎来更加高效地解决比较复杂的计算难题。近年来,GPU最成功的一个应用就是深度学习领域,基于GPU的并行计算已经成为训练深度学习模型的标配。目前,最新的CUDA版本为CUDA 9。
机器学习算法工程师
2018/07/27
5.2K0
快来操纵你的GPU| CUDA编程入门极简教程
英伟达CUDA垄断地位难保:PyTorch不断拆塔,OpenAI已在偷家
詹士 Alex 发自 凹非寺 量子位 | 公众号 QbitAI 英伟达的软件护城河正在逐渐消失。 随着PyTorch支持更多GPU厂商,再加上OpenAI的Triton搅局,英伟达手中的利器CUDA 逐渐锋芒不再。 上述观点来自Semi Analysis首席分析师Dylan Patel,相关文章已引发一波业内关注。 有网友看后评价: 英伟达沦落到此种境地,只因为了眼前利益,放弃创新。 Pytorch的作者之一Sasank Chilamkurthy还补刀: 当英伟达之前提出要收购Arm时,我就对潜在的垄断
量子位
2023/02/28
6610
英伟达CUDA垄断地位难保:PyTorch不断拆塔,OpenAI已在偷家
DeepSeek-R1自写CUDA内核跑分屠榜!斯坦福学霸狂飙GPU编程自动化挑战人类
近日,来自斯坦福和普林斯顿的研究者发现,DeepSeek-R1已经能生成自定义CUDA内核了,而且还在一众推理模型中,直接拿下了TOP 1!
新智元
2025/02/28
1370
DeepSeek-R1自写CUDA内核跑分屠榜!斯坦福学霸狂飙GPU编程自动化挑战人类
【论文解读】基于MLIR生成矩阵乘法的高性能GPU代码,性能持平cuBLAS
本文是对 https://arxiv.org/abs/2108.13191 这篇论文进行解读,学习一下如何基于MLIR编译器基础设施生成高效的GPU代码。本文的阅读的先后顺序分别为:
BBuf
2022/04/06
2.8K0
【论文解读】基于MLIR生成矩阵乘法的高性能GPU代码,性能持平cuBLAS
GPU加速03:多流和共享内存—让你的CUDA程序如虎添翼的优化技术!
阅读完前两篇文章后,相信读者应该能够将一些简单的CPU代码修改成GPU并行代码,但是对计算密集型任务,仅仅使用前文的方法还是远远不够的,GPU的并行计算能力未能充分利用。本文将主要介绍一些常用性能优化的进阶技术,这部分对编程技能和硬件知识都有更高的要求,建议读者先阅读本系列的前两篇文章,甚至阅读英伟达官方的编程手册,熟悉CUDA编程的底层知识。当然,将这些优化技巧应用之后,程序将获得更大的加速比,这对于需要跑数小时甚至数天的程序来说,收益非常之大。
PP鲁
2019/12/26
5.3K1
推荐阅读
相关推荐
【BBuf的CUDA笔记】十三,OpenAI Triton 入门笔记一
更多 >
领券
问题归档专栏文章快讯文章归档关键词归档开发者手册归档开发者手册 Section 归档