Triton 核心设计理念是块级编程。以区块为单位调度,不像CUDA那样按线程编程。块级调度方法让Triton抽象大量 GPU 编程细节,专注于数据移动控制的高层工作。所有底层线程细节,如并发、共享内存同步、内存合并,由编译器处理。这种抽象让不精通GPU的程序员也能轻松理解。
深度神经网络是高度可并行化的工作负载。为了有效训练和运行这些计算,我们使用大规模并行处理器(GPU,TPU)硬件,专门设计用于大规模数据执行相同计算(类似 SIMD 并行处理)。但 GPU 无法自动理解如何组织和调度并行计算。需要以特定方式与硬件交流,处理数据进行并行计算。英伟达推出 CUDA 编程模型,也是C/C++语言扩展,用于与GPU硬件通信。它提供对GPU线程、warp、内存层级、同步等细节的细致控制。CUDA 提供专家级性能,但需要具备系统编程专业知识。大多数机器学习研究者写的是Python,他们只是想尽快实验和验证他们的想法,并不关系底层运行原理。
为了解决这个问题,2010年代出现了机器学习框架,抽象了GPU编程的复杂性。TensorFlow等框架实现了高性能 GPU 计算无需编写 CUDA 代码。这些早期框架依赖静态图执行。在图执行框架中,神经网络程序先被编入计算图,然后让数据流经它。代码不会立即运行,错误只有在编译完整图后才会发现(就像C这样的编译语言)。调试缓慢晦涩难懂。研究人员抱怨由于框架的静态性,构建新模型非常困难;计算图中深层的错误难以追踪和调试,还有许多其他问题。
2016年初,Meta AI 发布了PyTorch。PyTorch 是一个 eager 执行框架,这是一种标准的脚本执行方法,代码行中的每个网络层都像 Python 程序一样依次执行。这种编程模型使 PyTorch 非常易于使用和调试,机器学习实验变得更快更直观。虽然PyTorch带来了极大的研究灵活性和可用性,但 eager 模式执行也有缺点。在 eager 模式下执行时,每个算子分别在GPU上启动,数据从全局内存读取,计算后将结果写回内存,然后再开始下一次计算。这导致更多的算子启动开销,过多的全局内存流量,硬件利用率差。如果不大量优化,会显著增加内存带宽需求,而GPU存在内存带宽瓶颈。频繁读写 HBM,导致程序执行时间和性能变慢。在底层,PyTorch 仍然依赖 CUDA 算子提升性能,这种访存模式和执行模型极大限制了 PyTorch 程序性能。所以我们只能在灵活性和最大性能之间挣扎。
Triton 是由 OpenAI 开发的一种嵌入 Python 的领域专用语言(DSL)和编译器,开发者可以直接在 Python 中编写高效的 GPU 算子。主要目标是灵活易使用:研究人员可以在 GPU 上实验算法,无需 CUDA 专业知识。Triton 通过抽象复杂的硬件细节,保留内存访问和管理控制,在易用性和高性能之间取得平衡。使得普通工程师能够直观地开发高性能GPU代码。
创建编程语言的步骤,从宏观角度来看,创建新的编程语言涉及三个主要步骤。
🥝定义语法。
🥝为源代码构建前端编译器。
🥝构建后端代码生成器。
所以可以先在纸上面定义所用语言的语法。如果不知道什么是形式语法,那就把它想象成人类语言的语法。例如,在人类语言中,你可能会说“一个句子由冠词、名词和动词构成”。然而,语言本身固有的歧义性,主要源于口语,使得这条规则失效。在形式语言中,如果语法规定“赋值是通过变量名、等号(=)和数字完成”,那么这是程序源代码中实例化变量的唯一方法,否则代码将无法编译(或者如果正在构建一种解释型语言,则会引发错误)。简而言之,形式语法定义了该语言中程序源代码必须是什么样子。它规定了程序员为了使用该语言编写程序必须遵守的语法规则。
🍑编译器前端,读取源代码并生成一些数据结构。
🍑后端代码生成器,接收前端生成的内容,创建实际可运行的代码。
所以这完全就是把初始源代码转换成一个可以运行的代码?是的,但必须是机器能够理解的那种,通常是目标机器码。这意味着需要提供汇编代码,将其提供给汇编器,汇编器与链接器/加载器协同工作,然后它们会将机器码返回。
设计新语言时,可能会决定让编译器生成 C 代码,而不直接生成机器码,然后再将 C 代码交给 C 编译器。这是一种标准做法,它并不会减少编写编译器过程中的工作:仍然需要自己构建一个编译器。
那么,编写编译器究竟要怎么什么呢?
创建编程语言编译器的步骤
为编程语言构建编译器是一项极艰巨的任务。需要由众多技术人员组成的大型团队。例如,FORTRAN 编译器的开发耗时约三年。FORTRAN 语言专为科学计算而设计,它彻底改变了该领域的格局,至今仍被广泛使用。他们的工作卓有成效,几乎完全改写了计算机发展史。
尽管编译器结构复杂,但它只有 5 个主要部分:
🥝词法分析。识别语言关键词、运算符、常量、语法定义的每一个词法单元。
🥝解析。理解词法单元数据流,每对词法单元之间的关系都被编码在一个树状数据结构中。树状结构主要描述源代码每一行中操作的含义。
🥝语义分析。可能是所有步骤中最晦涩难懂的。主要涉及理解类型,检查源代码“含义”(而不仅仅是语法)中的不一致的地方。
🥝优化。无论编译前的源代码有多好,在降级编码(直至机器代码)过程中,通常都可以进行一些优化,如内存优化、计算优化,当然还有运行时优化。
🥝代码生成。原始代码的优化版本最终被翻译成可执行代码。
尽管创建编译器的这五个步骤几十年来都没有改变,但每个步骤却发生了很大的变化。
让我们来看看每个步骤的细节:
词法分析
这一步直接源于形式语法。假设构建一种语言来表示最基本的算术运算。语法可能如下所示:
🍋🟩程序:[表达式";"]
🍋🟩表达式:赋值
🍋🟩赋值:结果“=" 操作
🍋🟩操作:变量运算符变量
🍋🟩运算符:"+" | "-" | "*" | "/"
🍋🟩变量:"1" | "2" | "3" | "4" | "5" | "6" | "7" | "8" | "9" | "0"
🍋🟩结果:"L" | "O"
上面的语法理解为:
🍋🟩程序由表达式组成,每个表达式后跟一个分号。
🍋🟩表达式只能是赋值语句。
🍋🟩赋值语句由结果后跟等号,再后跟运算语句组成。
🍋🟩运算语句由变量、运算符和另一个变量组成。
🍋🟩运算符必须是以下四个符号之一:“+”、“-”、“*”、“/”(双引号括起来的符号是终止符)。
🍋🟩变量必须是十个符号之一(即十个数字)。
🍋🟩结果必须是“L”符号或“O”符号。
用这种简单的语法只能表示两个小整数(0 到 9,包括 0 和 9)之间的加减乘除运算。这个序列也可以为空,如果不是空序列,则其中的元素必须用分号分隔。每次运算的结果只能命名为“L”或“O”。
有效的“源代码”可能如下所示
这个比较简单,用它也做不了什么。但可以用来解释词法分析的工作原理。词法分析器会遍历源代码,并根据语法识别每个词法单元及类型。所以,上面这段两行代码的词法分析结果是:
🍋🟩“O” - 结果
🍋🟩“2” - 变量
🍋🟩“+” - 运算符
🍋🟩“3” - 变量
🍋🟩“L” - 结果
🍋🟩“5” - 变量
🍋🟩“*“ - 运算符
🍋🟩“2” - 变量
这样就能理解软件工具可以对任何定义明确的语法执行词法分析。
解析器
解析器会接收标记列表(上一步的输出),创建一个树状结构,描述源代码的“含义”。对于之前的示例(简单的算术运算)中,这棵树反映了实际的运算过程:
经过简化,可以不需要解析树。仅仅知道解析树的叶子节点就可以完成编译源代码的目标。但是,每个解析器的行为可能有所不同,许多解析器都不会将实际的解析树存储在内存中。
语义分析
如果构建静态类型或特殊机制的语言,那么这一步就需要检查用户是否在进行任何无效的操作。以下Python代码:
def foo(foo): return foo + fooval = 5print(foo(val)) # >> 10def foo2(foo2): return foo + foo2foo = valprint(foo2(7)) # >> 12
显然不是好的代码。但现在要更仔细地分析一下:Python 解释器要理解这段代码会有多混乱?变量名为 foo,函数名为 foo,函数参数命名为 foo,函数 foo 的作用域内使用标识符 foo。当前作用域中绑定,是将正确的值(或整个函数体)分配给正确的标识符(如foo、val、foo2)。语义分析是理解编写的代码是否犯了错误。
优化
每个编辑器都有自己的目标。作为代码优化器,重点可能在于降低内存占用或缩短运行时间,也可能减少数据访问次数,或降低计算量或消息传递。由于专用硬件众多,各种组件和架构也多种多样,编译器在代码优化方面存在很大差异不足为奇。在这个主题中,必须特别提及数据流优化。任何有意义的源代码都会大量使用聚合数据结构,如数组,这些数据结构会在整个程序中被访问。数据流优化是改进程序员在程序中对数据结构所做的操作,提高代码整体执行效率。
代码生成
最后一步通常被称为代码生成(CodeGen),也实现了自动化。事实上,将语法分析树翻译成任何现有语言,基本上已经很成熟了。由于编译器方面的一些创新或特定后端,你可能想要构建自己的代码生成器。最常见的情况是,将语法树翻译成汇编代码,或者翻译成 C 代码,使用现有的 C 编译器进行编译。
GPU与优化基础
要理解 Triton GPU编程模型,我们需要快速了解GPU是如何执行工作的,以及内存访问如何影响性能。
GPU通过同时运行数千个轻量线程实现了高并行性。
🥝线程 Thread:GPU上执行最小的单元,处理一小块数据。每个线程都可以访问自己的私有存储,称为寄存器。
🥝Warps:一组32个线程,可以同步执行同一指令。Warps 是流式多处理器(SM)内部的基本调度单元。如果warp中的线程在控制流中出现分歧( 如 if ),性能会下降。
🥝块 Block:一组在同一SM上运行的协作线程组,能够协作并共享数据。一个 block 最多有1024个线程,最多32个warps。block 块内的所有线程都可以访问共享内存 SRAM。
🥝网格 Grid:kernel 发射的一组方块 block。网格中的所有线程都可以读取和写入全局内存和常量内存。网格内的方块 Blocks 不能直接同步。
🥝流式多处理器 SM:SM是GPU中执行 warps 的硬件单元。可以把它想象成一个小型多核处理器,负责调度 warps, 分配寄存器和共享内存,并行执行指令。
更多的SM会提升峰值计算吞吐量(FLOP)。 但实际性能很大程度上取决于GPU内存单元向这些SM提供数据的效率,这归结于内存访问模式。在深入探讨Triton算子代码之前,先简单介绍一下GPU上的这些内存单元,了解每个层级如何影响性能。
🥝寄存器:最接近计算。非常小,但速度快如闪电。每个线程存储很少的值。
🥝共享内存 / SRAM:仍然在片上,比寄存器大但也不够大。速度快到我们想在这里重复使用数据。Triton将SRAM暴露给开发者实现底层操作和控制。NVIDIA A100上,每个SM有~192 KB SRAM,带宽约19 TB/s。
🥝L2缓存: 体积更大,但速度较慢。整个GPU共享。会有帮助,但不能依赖它。它不由开发者控制,设备本身隐式地处理和管理。
🥝全局内存(HBM / VRAM):DRAM 系统。GPU上最大的内存空间,但速度最慢。A100 GPU(40–80GB,~1.5–2.0 TB/s带宽),从HBM传输数据的速度可能比从片上SRAM读取慢~15×。这里保存大张量和模型权重,因此高效的算子需要最小化HBM流量,将常用数据保存在寄存器或共享内存中以避免计算空等。
GPU性能的黄金法则:把数据往内存层级上移并尽量保持在那里。除非必要,否则不要返回 HBM。
Triton 编程模型
理解了GPU的工作执行方式,Triton 的设计变得更加清晰。Triton 保留了CUDA的强大功能,消除了CUDA带来的痛苦。Triton引入了一个更高层次的抽象,不像CUDA那样编程单个线程,引入了面向块的编程模型:
🍅 程序 Program 代替线程 Thread
Triton program 程序是一个工作单元,运行在一个数据块 block (tile) 上,不是单个元素,不同于 CUDA 中的线程。为工作块 block 编写逻辑,Triton 会自动并行启动多个程序 programs(kernel)处理全部张量。
更清晰更直观地思考输出张量被分割成块的发射的程序数量,每个启动的程序实例负责索引正确的输入数据并在输出中处理该工作块的计算。这是帮助我们清晰地可视化网格 grid 和 program 程序实例的一个方法。
通过以下方式选择 pragram 索引:
pid = tl.program_id(axis=0)
每个程序 program 计算数据的子集,比如 softmax 矩阵或向量加法中的一行。Triton 代码是从程序实例角度编写的。启动的程序 program 数量就是网格 grid。
直观地想象输出张量切分成地块就是启动的Triton程序数量。每个程序program实例选择正确的输入数据,在输出中对应的块上进行计算。这个模型可以帮助我们清晰地可视化程序 program 实例的网格 grid 如何映射到输出张量,更容易推导内存访问和并行计算。
🍅 显式内存控制
性能很大程度上来自于控制数据的流动方式。 Triton 让开发者能够直接控制内存访问模式的低层次细节。开发人员决定加载什么、何时加载、以何种模式(在哪里)。Triton 会自动补全线程并发、内存合并、SMEM 使用与同步、warp 配置和硬件映射等底层细节。这些抽象概念比较难,即使资深CUDA程序员也会感到困难。Triton 实现了底层细节处理的自动化,开发者和研究人员可以专注于并行化的高层逻辑。Triton 不会跨 SM 调度工作,重要的算法考量(如平铺、SM 间同步)交给开发者自行决定。
🍅 编译器驱动的优化
一旦写好 kernel 逻辑,Triton编译器会:
🍋🟩自动生成并调度块内线程
🍋🟩确保内存的合并访问
🍋🟩管理同步和调度
🍋🟩最佳性能的参数自动调优
开发人员专注于:这个program应该运行在哪个数据块上?
编译器处理:如何干净高效地将这些代码映射到GPU硬件?
为了更好地理解Triton编程模型,我们接下来将用2篇文章说明一下 triton 编程的语法并解说一下Triton官方的算子。