How To Write A CUDA Program The Parallel Programming Edition
How To Write A CUDA Program The Parallel Programming Edition
Stephen Jones, CUDA Architect | GTC 2025
目录
- 引言:并行编程的挑战
- CUDA的两层并行机制与可扩展性设计
- 两种并行机制的结合
- 寻求两全其美的编程模型
- 编程模型的层次结构
- 引入 cuTile:CUDA 的 Tile 编程
- 任务并行与数据分片 (Tiling
- 块与 Tile 的关系
- 执行块到数据 Tile 的灵活映射
- cuTile:结合可编程性与数据并行操作
- cuTile 弥合数据并行与设备库之间的鸿沟
- Llama-3.1 在 cuTile 中的实现
- 在任何工作流中结合使用 Tile 和 SIMT 核函数
- 这不是“非此即彼”——而是“兼而有之”
引言:并行编程的挑战
并行编程非常困难。因此,最好的策略是尽可能避免进行并行编程。
然而,在诸如NVIDIA AI和NVIDIA OMNIVERSE等复杂系统中,底层仍然需要并行编程来驱动。
CUDA 的秘密及其并行编程版图
CUDA 不仅仅是一件事物
CUDA 的秘密在于它并非单一的技术,而是一个包含多个层次的生态系统。如果某个工具或库能让你的程序在 GPU 上运行,那么你就在使用 CUDA。
CUDA 的技术栈可以大致分为以下几个层次:
- 框架(Frameworks)与领域特定语言(DSLs)
- 软件开发工具包(SDKs)
- 领域特定库(Domain-Specific Libraries)
- 加速库(Accelerated Libraries)
- 通信库(Communication Libraries)
- 设备库(Device Libraries)
- 内核编写(Kernel Authoring)
- 编译器栈(Compiler Stack)
- 主机运行时与工具(Host Runtimes & Tools)
CUDA 的并行编程版图
在 CUDA 的整个生态系统中,真正需要开发者进行并行编程的部分比想象中要少得多。
- 需要并行编程的层次:主要集中在“内核编写(Kernel Authoring)”层。
- 某种程度上相关的层次:设备库(Device Libraries)也与并行编程相关,但方式不尽相同。
- 几乎不需要并行编程的层次:在通信库(Communication Libraries)及以上的所有高层抽象中,开发者通常不需要直接处理并行编程的复杂性。
总结来说,CUDA 的设计使得开发者需要进行的并行编程远比想象的要少,这是一件好事,因为并行编程本身是困难的。
现代 CUDA 平台堆栈概览
当前的 CUDA 平台堆栈包含了数百个框架、SDK、编译器、库、语言、DSL 和运行时。
下图展示了 CUDA 平台堆栈的一部分,并列举了每个层次中的一些具体技术和工具。
堆栈的两个世界
CUDA 堆栈可以被看作两个不同的世界:
- 上层(高生产力/高性能):从通信库到顶层的框架和 DSL,这些工具旨在提供高生产力和高性能,开发者通过使用这些库和框架来利用 GPU 加速,而无需深入底层细节。
- 下层(“龙之领域”):设备库及以下的层次是更底层的领域,直接进行内核编写等工作。这个领域非常强大,但同时也极其复杂,充满了挑战,好比古代地图上标注的“Here be Dragons!”(此处理有恶龙)。
核心观点:你几乎不需要编写内核代码
通过一个(完全虚构的)饼图可以说明一个观点:在实际开发中,开发者需要思考的核心算法代码(Actual code I have to think about)只占一小部分。大部分工作是围绕着样板代码、测试、配置、文件I/O等“无聊的东西”(Boring stuff)。
这个比喻的核心思想是,正如大部分编程工作不是在写核心算法一样,大部分使用 CUDA 的工作也不需要编写底层的内核代码。
代码编写的现实:为何几乎不需要编写内核代码
幻灯片通过一系列(虚构的)饼图阐述了一个核心观点:开发者在实际工作中很少需要从头编写底层的并行计算内核。
- 大部分代码是“乏味的”: 在程序员需要考虑的所有代码中,大部分是样板代码、测试、配置、文件I/O等。真正需要深入思考的核心逻辑代码只占一小部分。
- 核心代码中大部分是串行的: 在这些核心的“真实”代码中,绝大部分是用于设置、清理、通信、内存管理等的顺序CPU代码。
- 可并行代码中大部分已有库支持: 在那小部分可用于加速的数据并行代码中,绝大多数功能(约占总代码量的99%)已经可以通过调用标准库中的函数来实现。开发者需要自己编写的并行算法(即内核代码)通常只占整体代码的极小一部分(例如1%)。
GPU加速的三个层次
为了实现GPU加速,开发者可以从不同抽象层次的工具入手,而不需要直接编写底层代码。
1. 高级框架与SDK
许多领域专用的高级框架和软件开发工具包(SDK)已经内置了GPU加速功能。用户可以在不接触底层并行编程细节的情况下,利用GPU的计算能力。这些工具涵盖了从机器学习(PyTorch, TensorFlow, JAX)到科学计算和工程仿真(Ansys, Siemens NX, LAMMPS, OpenFOAM),再到内容创作(Adobe Photoshop, Houdini)等多个领域。
2. 可直接集成的加速库
NVIDIA提供了一系列CUDA数学库,可以作为“即插即用”的加速器集成到现有程序中,为计算密集型应用(如分子动力学、计算流体动力学、医学成像等)提供基础。这些库包括:
- cuBLAS: 用于基础线性代数运算。
- cuFFT: 用于快速傅R叶变换。
- cuRAND: 用于随机数生成。
- cuSOLVER: 用于密集和稀疏直接求解器。
- cuSPARSE: 用于稀疏矩阵的BLAS。
- cuTENSOR: 用于张量线性代数。
- cuDSS: 用于直接稀疏求解器。
- CUDA Math API: 用于标准数学函数。
- AmgX: 用于模拟和隐式非结构化方法的线性求解器。
3. 从CPU调用的GPU加速并行操作
如果现有库无法满足特定需求,许多现代编程语言也提供了从CPU代码中直接调用GPU执行并行操作的机制,避免了编写和管理CUDA内核的复杂性。这些语言和库包括:
- C++17 标准并行(Standard Parallelism)
- Fortran DO CONCURRENT
- CUDA C++ Thrust
- Mathworks Matlab
- GPU Julia
努力与回报的现实曲线
在进行GPU编程和性能优化时,投入的努力与获得的回报(性能提升)之间的关系并非线性。
- 期望与现实的差距:
- 理想期望: 人们期望性能回报与投入的努力成正比,呈一条直线。
- 对框架的期望: 希望框架能以极小的努力获得峰值性能。
- 现实曲线: 实际的努力/回报关系更像一条S形曲线,初期投入大量努力但收效甚微,之后进入一个快速增长期,最终在接近峰值性能时收益递减。
- 真实的挣扎: 实际优化过程充满波折,可能会经历性能不升反降、过早兴奋、绝望低谷等阶段,最终才能逐步接近目标。
-
正确的策略:
- 应该致力于缩短初期收效甚微的时间。
- 尽快进入并停留在S曲线的陡峭部分,这是投入产出比最高的阶段。
- 避免在接近性能极限时过度投入,因为此时回报极低。
-
选择合适的工具:
不同的工具对应不同的努力/回报曲线。开发者应根据需求选择最合适的路径:- 框架 (Frameworks): 以最少的努力快速达到一个不错的性能水平,但性能上限较低。
- 库 (Libraries): 需要中等程度的努力,但可以达到更高的性能水平。
- 手写内核 (Hand-coded kernels): 需要投入巨大的努力,但有潜力达到硬件的理论峰值性能。
CUDA入门示例:SAXPY
SAXPY操作(Y = a*X + Y)是并行计算中一个经典的入门示例,用于展示从串行CPU代码到并行GPU代码的转换。
-
C语言CPU版本: 使用一个简单的
for循环来遍历数组,对每个元素执行计算。 -
CUDA C++ GPU版本:
- 内核函数 (
__global__ void saxpy(...)): 这是在GPU上并行执行的函数。每个GPU线程计算一个全局唯一的索引i,并负责处理第i个元素的计算任务 (y[i] = A * x[i] + y[i])。 - 主机代码 (
run_cuda()): 这是在CPU上执行的代码,负责启动GPU内核。通过saxpy<<<...>>>语法来指定启动的线程块(blocks)和每个块内的线程数(threads per block)。在此示例中,启动了1024个线程块,每个块包含1024个线程,总共1024 * 1024 = 1,048,576个线程,并行处理1M个元素。
- 内核函数 (
GPU如何运行并行程序
GPU通过将工作分解为包含线程子集的“块”(blocks)来运行并行程序。首先,一个大的任务被视为一整块线程集合。
随后,这个大的线程集合被分解为多个独立的块。
每个块被赋予一个唯一的标识符,例如Block 0, Block 1, Block 2, Block 3等。这些块随后会被调度到GPU的硬件上执行。下图以2012年的Kepler GK110 GPU为例,展示了其硬件架构。
在GPU架构中,执行这些块的基本单元是流式多处理器(Streaming Multiprocessor, SM)。下图放大了Kepler GK110 GPU中的一个SMX(Kepler Streaming Multiprocessor)。
这些程序块被分配到GPU中的多个SM上并行执行。例如,一个Kepler GK110 GPU拥有15个SMs,可以将不同的块同时调度到这些SM上。
GPU的调度器动态地将等待执行的块分配给可用的SM。
一个典型的并行计算任务可能包含大量的块,例如1024个块(Block 0到Block 1023)。这些块会被调度到GPU的所有SM上。对于拥有15个SM的Kepler GK110 GPU,这意味着每个SM需要依次处理多个块。
随着GPU硬件的发展,SM的数量显著增加。例如,2022年的Hopper H100 GPU拥有132个SMs,这使得它可以同时执行更多的块,从而大幅提升并行处理能力。
GPU通过超额订阅(Oversubscription)工作
GPU利用一个深度的块队列来确保其计算资源始终处于忙碌状态,避免空闲。块(Blocks)以流式方式进入GPU,填充空闲的SMs。当一个块完成其计算任务后,它会退出,从而释放其占用的SM,以便新的块可以被调度上来执行。
增加SM数量 = 更多并发 = 更高性能
增加GPU中的SM数量是提升性能的关键。更多的SM意味着可以同时执行更多的线程块,即更高的并发度。下图对比了拥有15个SM的早期GPU和拥有132个SM的Hopper H100 GPU,展示了SM数量增加带来的并发能力提升。
一个微妙但重要的观点:线程与数据
回顾之前的saxpy示例,我们处理了两个数组x和y,每个数组包含100万(具体为1048575)个元素。
我们将这个任务分解为1024个独立的块,每个块包含1024个线程。在CUDA C++代码中,这通过saxpy<<<1024, 1024>>>(...)的语法实现,其中第一个参数是块的数量,第二个参数是每个块的线程数。
然而,这里有一个需要澄清的关键点:线程和数据并非同一回事。
以H100 GPU为例,其硬件能够同时支持大约25万个线程的并发执行(132个SMs × 每个SM 2048个线程 = 总共270,336个线程)。但是,我们的数据量是100万个元素,是硬件并发能力的四倍。
尽管物理线程数远少于数据元素数,我们的程序依然可以高效运行。这是因为我们可以将代表计算任务的块以流的方式送入GPU。当一部分块完成计算并退出后,新的块会立即被调度到释放出的SM上执行。通过这种方式,GPU可以处理远超其物理并发能力的大规模数据。但这并非是用25万线程处理100万元素的唯一方法。
每个线程处理多个数据元素
简单示例:每个线程处理一个数据元素
在基础示例中,每个线程处理一个数据元素。这意味着数据索引 i 与线程索引相等,因此需要的线程总数等于数据元素的总数。
如上图所示,为了处理 1,048,576 个元素,需要启动 1,048,576 个线程,这可以配置为 1024 个线程块,每个线程块包含 1024 个线程。
优化:每个线程处理4个数据元素
我们可以通过让每个线程处理多个(例如4个)数据元素来优化。
在这种情况下,数据索引 i 变为 (线程索引) * 4。我们现在每个线程处理4个数据元素。这将启动的线程总数减少到 262,144 个(即 256 个线程块 * 1024 个线程/块)。
由于数据总量 1,048,576 不一定能被启动的线程总数(在此例中为 262,144)乘以每个线程处理的元素数(4)整除,因此必须在代码中添加边界检查 (if(i < N)),以防止越界访问。
相应的,内核启动配置也需要修改。由于不再需要超过100万个线程,线程块的数量从 1024 个减少到 256 个。
H100 GPU上的并发执行能力
NVIDIA H100 GPU 能够同时运行 262,144 个线程。其架构包含 132 个流式多处理器(SMs),每个SM最多可运行 2048 个线程,总计可并发运行 270,336 个线程。
线程与数据:一个微妙但重要的观点
让每个线程加载多个元素可以显著提升性能。这种模式被称为“网格跨步循环”(grid-stride loop)。在上述4元素代码示例中,循环被展开了4次。
如图所示,每个线程处理4个元素的性能大约是每个线程处理1个元素的两倍。其性能提升的原因包括:
- 无额外开销:无需创建额外的线程块。
- 4路内存加载批处理:可以批量处理内存加载。
- 减少块内抖动效应:降低了线程块内部的执行不确定性。
CUDA的两层并行机制与可扩展性设计
CUDA的并行模型专为可扩展性而设计。在GPU架构的演进中,主要是SM的数量发生变化,而SM本身的规模(如每个SM的最大线程数)保持不变。
- SM规模固定:每个SM的线程数是固定的。这意味着为旧GPU编写的代码,其每个线程块的工作量(work-per-block)无需为新GPU增加,程序结构也无需改变。
- SM数量逐代增加:通过增加SM数量来提升总体性能。程序的线程块数量超过SM数量(即“超额订阅”,Oversubscription)可以使其自动扩展到未来的GPU,只要保证
线程块数量 >> SM数量。
并行编程模式
软件栈与并行编程的领域
软件栈从高到低包括:框架与DSL、SDK、领域特定库、加速库和通信库。更底层的部分则更为复杂,如同进入“龙穴”(Here be Dragons)。
在设备库(Device Libraries)、内核编写(Kernel Authoring)、编译器栈(Compiler Stack)和主机运行时与工具(Host Runtimes & Tools)等底层,需要并行编程知识。
两种基本的并行模式
基本上只有两种并行模式:任务并行(Task Parallelism)和数据并行(Data Parallelism)。
- 任务并行:将独立的程序分布到不同的处理器上。它指的是在不同线程/处理器上同时运行独立的程序,也可能是同一程序在代码的不同位置的多个副本。
- 数据并行:将单个数据元素分布到不同的处理器上。它指的是一个程序在多个线程/处理器上运行,所有线程以步调一致(lockstep)的方式对不同的数据元素执行相同的操作。
CUDA的混合并行模型
CUDA同时利用了这两种并行模式:任务并行中嵌套数据并行。
在CUDA中,不同的线程块(Block 0, Block 1)可以看作是任务并行,而每个线程块内部的线程则执行数据并行操作。
并行编程在软件栈中的应用
- 数据并行 主要应用于 内核编写(Kernel Authoring) 层面。
- 任务并行 也与 设备库(Device Libraries) 相关,但方式不完全相同。
数据并行编程:从难点开始
我们将从数据并行开始,因为它通常是并行编程中更具挑战性的部分,主要集中在内核编写层面。
数据并行的基本操作符:归约(Reduction)
“归约”操作符是并行编程的基本操作符之一,一个典型的例子是计算N个数字的和。
上图展示了一个在CPU上运行的串行求和函数 sum_array,它通过一个循环来累加数组中的所有元素。
串行求和分析
对于一个包含64个元素的数组进行串行求和,完全没有并行性。sum_array 的总时间复杂度为 N 步,即需要 64 个步骤。
并行编程:分治策略
分治思想 (Divide and Conquer)
并行编程的精髓在于分治策略。以一个包含64个元素的数组求和为例,如果将其划分为4个部分,由4个并行单元分别计算部分和,最后再将这4个部分和相加,总时间消耗为 64/4 + 4 = 20。
进一步分治与性能瓶颈
如果将任务划分得更细,例如分为16个部分,计算时间为 64/16 + 16 = 20。结果显示,这样做并没有加快速度,因为虽然第一步的并行计算(64/16)时间缩短了,但第二步的串行求和(+16)开销增大了,抵消了并行带来的收益。
递归分治:并行化归约过程
为了解决上述瓶颈,可以对第二步的求和过程也进行并行化。通过再次应用分治策略,将16个部分和的求和任务也并行处理,总成本可以降低为 64/16 + 16/4 + 4 = 12。
这种递归应用分治策略最终形成一个树状的归约(Reduction)结构。
并行算法改变时间复杂度
并行算法能够改变问题的解决时间复杂度,从而解决一些过去难以处理的问题。例如,使用 N/2 个线程对 N 个元素进行求和,可以将时间复杂度从 O(N) 降低到 O(log₂(N))。
基础并行算子
数据并行程序由一系列基础算子构建而成,这些算子能够带来超线性加速。
- Map (映射): 将一个函数并行地应用于一个序列的每个元素。
- Reduction (归约): 将一个序列的数值相加(或进行其他聚合操作)。
- Scan / Prefix Sum (扫描/前缀和): 对一个序列进行累积求和。
- Sort (排序): 根据键值对值进行重新排序。
数据并行的优势与挑战
数据并行需要大量线程
数据并行在拥有大量线程的硬件(如GPU)上比在线程较少的硬件(如CPU)上更为有效。N/2个线程可以实现 O(log₂(N)) 的效率,而仅有4个线程则只能实现 O(N/4) 的效率。
线程级并行的实现
以并行归约(Parallel Reduction)为例,线程级并行能够带来显著差异。下面是一个基础的并行归约CUDA代码实现。该代码利用共享内存和线程同步来高效地完成块内归约。
高级SIMT编程的复杂性
编写高度优化的单指令多线程(SIMT)程序非常复杂,需要对并行算法和GPU硬件架构都有深入的理解。下图展示了一个经过深度优化的“忍者级”并行归约代码,其性能相比简单实现和高级实现有巨大提升(30倍加速)。
数据并行编程的难度
自己从头实现所有这些基础并行算子(如排序、映射、求和)是一项艰巨的任务。
数据并行编程尤其困难。因此,推荐的做法是不要自己从头做起,而是利用已有的高效库,因为这些工作已经有专家为你完成了。
CUDA并行算法库
CUB 和 cuda.cooperative
cub 和 cuda.cooperative 是所有并行算法的基础集合构建模块。
- CUB: 一个用于并行集体操作(collective operations)的C++库。
- cuda.cooperative: 在CUDA Python中与CUB等效的库。
- 集体操作: 指在给定范围内所有线程协同工作以实现加速的算法。
- CUB提供的集体操作涵盖四个范围:
1. 线程 (Thread)
2. 线程束 (Warp)
3. 线程块 (Block)
4. 设备 (Device)
- 这两个库都用于内核(kernel)代码内部。
MathDx 库
MathDx库是可以在CUDA内核中内联使用的数学库,提供了经过“忍者级”调优的块级(Block-wide)并行数值算法。
MathDx 和 nvmath-python 库的特性:
- 允许在内核内进行融合、定制和组合。
- 提供经过高度优化的CUDA数学库的全部性能。
- 在整个线程块范围内操作。
- 是包含在内核中的CUDA C++头文件库。
- 提供可从Python内核代码(如numba-cuda)调用的Python API。
下图展示了基于FFT的卷积在不同实现下的性能对比,以及MathDx生态系统中的核心库,如cuBLAS, cuFFT, cuSOLVER和cuRAND。
两种并行机制的比较
任务级并行(Task-level parallelism)虽然更易于编写,但缺乏线程级数据并行(thread-level data parallelism)的灵活性。下图展示了两种并行方式在NVIDIA软件栈中的位置:
- 数据并行(Data Parallelism):主要通过 CUDA C++ 和 numba-cuda 等工具在内核创作(Kernel Authoring)层面实现。
- 任务并行(Task Parallelism):主要通过 MathDx、CUB 等设备库(Device Libraries)以及 nvmath-python、cuda.cooperative 等接口实现。
两种并行机制的结合
在实际应用中,通常需要同时利用这两种并行机制。开发者应尽可能利用库来实现任务并行,并在必要时编写定制化的数据并行代码。库无法覆盖所有应用场景,因此总需要编写自定义代码。下图展示了在CUDA块(Block)内部,任务并行和数据并行可以协同工作。
寻求两全其美的编程模型
许多应用程序和算法本质上是数据并行的,它们操作的对象是数组和张量。理想情况下,我们希望有一种任务并行的编程模型,能够让我们更轻松地编写自定义代码。这意味着不仅设备库可以提供任务并行能力,内核创作层面也应能支持逻辑上的任务并行(Logical Task Parallelism)。
编程模型的层次结构
CUDA编程模型可以分为三个层次:
-
网格级模型 (Grid-level models):
- 包括并行库、DSL、框架和SDK。
- 编译器负责将数据映射到块(blocks)和线程(threads)。
-
块级模型 (Block-level models):
- 采用任务并行(Task-Parallel)编程。
- 应用程序负责将数据映射到块,编译器负责将块映射到线程。
-
线程级模型 (Thread-level models):
- 采用细粒度的并行编程。
- 应用程序负责将数据同时映射到块和线程。
引入 cuTile:CUDA 的 Tile 编程
cuTile 是一个即将推出的面向 CUDA Python 和 CUDA C++ 的新编程模型。它属于块级模型,旨在为基于数组的程序提供高生产力和高性能。
核心特性:
- Tile 编程:在规则的数据 Tile(瓦片)上进行块内协作执行(Block-wide cooperative execution)。
- 粒度:数据和操作的粒度是数组或张量(array / tensor)。
- 优势:
- 非常适合规整的、数据并行的计算问题。
- 便于实现高性能的核函数融合(kernel fusion)。
- 通过隐式使用线程简化了数组操作。
下图对比了 Tile 张量加法(数组粒度)和 SIMT 张量加法(元素粒度)。
cuTile 能够将基于数组/张量的 Tile 模型高效地映射到 Tensor Cores,从而简化代码结构并更容易地实现高性能。编译器会自动管理内存空间并将其映射到 Tensor Core 硬件上。使用 cuTile 编写的程序可以跨越当前和未来的 Tensor Core 硬件代际进行移植(但可能需要重新调优以达到峰值性能)。
任务并行与数据分片 (Tiling)
任务并行的核心思想之一是对数据进行分片(Tiling)。这与之前讨论的 saxpy 示例中对数据进行分区和粒度选择的决策过程相似。
任何并行化都需要对数据进行深思熟虑的划分。以之前的一维 saxpy 示例为例,原始数据是两个一维数组 x 和 y。
为了实现并行化,我们将数据划分到不同的计算单元(如CUDA块)中。
在数据划分时,粒度的选择至关重要。例如,可以采用“每个线程处理一个元素”的策略,也可以采用“每个线程处理四个元素”的策略,后者意味着一个块可以处理4倍的数据量。
块与 Tile 的关系
正如线程和数据不是同一个概念一样,执行单元(块)和数据单元(Tile)也不是同一个概念。我们可以将数据(如一张图片)划分为多个 Tile,然后将这些 Tile 映射到块网格(Grid of Blocks)上。一种直接的映射方式是每个块处理一个 Tile,例如将图片划分为 8x8 的 Tile,并创建一个 8x8 的块网格,共64个块。
执行块到数据 Tile 的灵活映射
应用程序可以定义执行块(Execution Blocks)与数据 Tile 之间的映射关系。与SIMT类似,一个块可以操作一个或多个 Tile,这种灵活性使得程序可以根据具体算法和硬件特性进行优化。
以下是几种不同的映射示例:
- 1x8 网格 = 8个块:每个块处理 8 个 Tile(即一行 1x4 的 BlockTiles)。
- 2x8 网格 = 16个块:每个块处理 4 个 Tile。
- 4x4 网格 = 16个块:每个块处理 4 个 Tile(排列为 2x2 结构)。
- 8x2 网格 = 16个块:每个块处理 4 个 Tile(排列为 1x4 结构)。
cuTile:结合可编程性与数据并行操作
cuTile 是一种将可编程性与数据并行操作相结合的编程模型。其核心概念是 "Tile",即一个适合协作操作的数组。
cuTile 编程模型:
* 启动方式: 启动一个由 "tile" 块组成的网格。
* 逻辑线程: 每个 Tileblock 只有一个逻辑线程。
* 执行映射: 执行到物理线程的映射是隐式的。
* 数据单元: 数据处理的基本单位是 Tile,而非单个元素。
* 集体操作: 操作是在 Tile 上进行的集体操作。
* 并行算子: 定义了许多并行算子,例如 "reduce" 和 "sort"。
* 互操作性: 可与 SIMT(单指令多线程)进行互操作,以实现细粒度的并行。
cuTile 弥合数据并行与设备库之间的鸿沟
cuTile 的设计理念是让编译器负责数据并行数组的优化,从而使编程更简单、性能更高。它在软件栈中扮演着连接高层任务并行与底层设备库的角色。
如下图所示,cuTile 使得 Tile 并行编程的方式,类似于使用为特定任务高度优化的设备库(如 CUB, MathDx等),但提供了更高的灵活性。开发者可以专注于任务并行逻辑,而将复杂的数组优化交给 cuTile 和编译器。
Llama-3.1 在 cuTile 中的实现
本页展示了使用 cuTile 编写的自定义核函数,在 PyTorch 框架下实现 Llama-3.1 模型推理的案例。
左图是 Llama-3.1 的模型架构图。右图的性能对比测试(在 Blackwell B200 平台上进行)显示:
* 与 Torch Eager CUDNN SDP 这一峰值性能基准(基于 TRT-LLM)相比,cuTile 实现的性能达到了 1.00x。
* 与 Torch Eager Flash SDP 相比,cuTile 性能是其 0.93x / 0.41x ≈ 2.26 倍。
* cuTile 的性能达到了 Torch Eager Flash SDP 的 0.93x。
* Torch Eager CUDNN SDP 的性能为峰值 E2E 性能的 1.00x。
这个案例证明了 cuTile 在复杂AI模型中实现高性能推理的潜力。
在任何工作流中结合使用 Tile 和 SIMT 核函数
并非所有算法都完全是数据并行或基于数组的,因此能够灵活选择和组合使用不同的编程模型至关重要。
- 初始版本: 将提供独立的 Tile 核函数和 SIMT 核函数。
- 未来版本: 将支持混合核函数,即在 Tile 函数内部调用 SIMT 函数。
cuTile 核函数本质上是 CUDA 核函数,因此它们可以像普通 CUDA 核函数一样在任何地方“直接运行”,无缝集成到现有的工作流(如右侧的计算图所示)中。
这不是“非此即彼”——而是“兼而有之”
并行编程本质上是困难的,因此简化编程是至关重要的。NVIDIA 的目标不是用一种模型取代另一种,而是提供一个涵盖所有需求的完整生态系统。
如下图所示,不同的并行范式在软件栈的不同层次上发挥作用:
* 隐式并行: 由框架和领域特定语言(DSL)处理。
* 任务并行: 可通过设备库或基于 Tile 的 cuTile(支持 C++ 和 Python)来实现。
* 数据并行: 可通过传统的 SIMT CUDA C++ 和 numba-cuda 来实现。
cuTile 与 SIMT CUDA 在核函数编写层面协同工作,共同为开发者提供了强大的工具,以应对不同粒度和类型的并行计算挑战。