7. CUDA

GPU

  • CPU 与 GPU 的组织架构区别

    Comparing the relative capabilities of the basic elements of CPU and GPU architectures - 其中,绿色代表计算;金色代表指令处理;紫色代表 L1 缓存;蓝色代表更高级别的缓存;橙色代表内存(DRAM,其实际容量应该是缓存的数千倍) - CPU 的设计核心是低延迟、强控制流,包含分支预测/乱序执行/缓存等核心机制,单线程性能很强发,其是为「复杂逻辑 + 少量并发」而生的 - 而 GPU 的设计核心是高吞吐、大规模并行、简单控制流,用并行来掩盖内存延迟,很适合大规模数值计算、规则的数据访问、相同算子作用在大量数据上,其是为「简单逻辑 + 海量并行」而生的 - 深度学习的核心操作是矩阵乘法、卷积、逐元素运算、归约等,这些操作的共同特征是数据量大、计算规则统一、控制流很简单、并行度很高,所以深度学习的计算形态正好符合 GPU 的设计假设

  • GPU 的计算层次结构:线程(Thread) → Warp → SM

    image.png - 线程(Thread) - 线程是最小执行单位,它可以执行一段标量指令,有自己的私有寄存器 - GPU 上的线程不会单独调度

    • Warp(真正的调度单位)

      • 线程被分组为 Warp,每个 Warp 包含 32 个线程

      • Warp 中的所有线程同步,它们执行同一条指令,但处理不同的数据

      • 即 SIMT —— Single Instruction, Multiple Threads

      • 对于 if / else 分支,不同的线程会走不同的分支,会被串行执行(warp divergence),因此 GPU 极其不适合复杂的控制流

      Threads in the same warp cannot execute statements in the if-block in parallel with statements in the else-block. The timeline shows that only one group of threads or the other can be active.

    • SM(并行计算核心)

      • GPU 的基本计算单元是流式多处理器(Streaming Multiprocessor, SM)

      • GPU 由多个 SM 组成,每个 SM 同时驻留多个 Warp

      • 每个 SM 包含若干流处理器 / 核心(core),例如 NVIDIA H100 GPU 包含 132 个 SM,每个 SM 包含 128 个核心,总共 16,896 个核心

      • SM 中的每个核心可以同时处理多个线程(threads)

      • GPU的高性能效果,实际上来自于用并行掩盖延迟:当一个 Warp 等待内存时,SM 立刻切换到另一个 Warp

      • 在 GPU 上,一次核函数调用由一个或多个流式多处理器(SM)执行

    • GPU 的使用受制于各层内存大小、warp 内线程数和 block 数量,合理分配这些资源,可实现 GPU 的高效利用

  • 内存层次结构

    image.png - GPU 内存同样高度分层,包含多层缓存和内存

    内存层级
    作用
    共享对象
    特点

    寄存器(Registers)

    最小单元,线程私有

    线程私有

    访问最快,容量最小

    共享内存(Shared Memory)

    SM 内线程共享

    同一个 SM 内的所有线程

    速度快,容量中等

    L1 Cache

    SM 内线程共享

    同一个 SM 内的所有线程

    速度快,容量中等

    L2 Cache

    所有 SM 共享

    全 GPU

    较大、访问比 L1 慢

    全局内存(Global Memory)

    GPU 最大内存

    全 GPU

    容量大(H100 为 80GB),访问最慢

    • 使用 GPU 时的目标是利用计算/内存资源的层级组织,尽可能多的工作负载并行运行于可用核心上

    • GPU 的算子性能,80% 是内存访问决定的,不是算力不够,而是数据移动瓶颈、memory access 不连续、cache 命中率低

  • Tesla V100 的内存组织架构

Local to each SM in the GPU is a 256kB register file, a 128kB unified data cache split between L1 and shared memory, and 64kB of constant caches. All SMs share a 64MB L2 cache and 32GB of global memory.

  • 片上存储(SRAM)

    • 寄存器文件(Register File)

      • 直接为 CUDA 核心提供数据的专属存储区域,其硬件架构被划分为 32 个存储体(bank),与线程束(warp)的 32 个线程数量对应

      • 可以将寄存器文件理解为一个由 4 字节数据单元构成的大型矩阵,包含若干行与 32 列;线程束以整行为单位进行数据操作,在某一行内,每个线程(对应一个 CUDA 核心)分别对不同列(对应不同存储体)的数据执行运算

    • 一级缓存(L1 Cache)

      • 即常规的片上高速缓存,可为近期从主存(RAM)读取或向主存写入的数据提供高速访问能力

      • 当活跃数据量超出 SM 寄存器文件的存储上限时,L1 缓存会作为寄存器溢出(register spilling)的兜底溢出区域

      • 在该场景下,L1 缓存中的缓存行与溢出寄存器会沿用寄存器文件的设计逻辑,同样按存储体划分组织

    • 共享内存(Shared Memory)

      • 共享内存与 L1 缓存物理上同属一块片上存储区域,核心差异在于:共享内存中的所有数据可被一个线程块(thread block)内的任意线程访问,这为线程间的通信与数据共享提供了硬件支撑

      • 共享内存中的变量需由应用程序显式声明,且开发者可通过程序配置调整 L1 缓存与共享内存的存储空间划分比例

    • 常量缓存(Constant Caches)

      • 针对全局内存中声明为只读常量的变量设计的专用高速缓存,这类常量变量可被一个线程块内的任意线程读取

      • 该缓存的核心且最优的应用场景为:向一个线程束内的所有线程广播同一个常量值,借助硬件特性实现高效的批量数据分发

  • 离 CUDA Core 较远的存储结构

    • L2 Cache

      • 片上高阶缓存,用于暂存流多处理器(SM)与主存之间交互的各类数据副本

      • 可以视为设备内存的中介缓存 / 代理缓存,设备内存的所有访问操作均经由 L2 Cache 完成,当需要为新数据腾出空间时,缓存中最久未被访问的数据会根据缓存淘汰策略,被回写至设备内存中

      • 与 L1 Cache 相同,L2 Cache 的设计初衷是为后续数据重加载提供加速;但与独立的各 L1 缓存不同,GPU 中仅配备一个全局 L2 Cache,由所有流多处理器(SM)共享使用

      • L2 Cache 还部署在设备通过 PCIe 或 NVLink 进行片外数据传入 / 传出的核心数据通路上,是设备与外部交互的必经缓存层

    • 全局内存(Global Memory)

      • 全局内存是设备主存的核心组成部分,对应 CPU 架构中的随机存取存储器(RAM)

      • 为了性能优化,不同型号的 GPU 都搭载了专属高速主存:如 Tesla V100 配备了高性能的第二代高带宽内存(HBM2),而 Quadro RTX 5000 则采用了高速图形双倍数据率 6 代内存(GDDR6)

    • 纹理内存与常量内存(Texture and Constant Memory)

      • 设备主存中被定义为只读属性的存储区域,所有 SM 都可以访问

      • 当这类内存中的数据被调取至 SM 后,以 “纹理” 或 “常量” 关键字声明的变量,可以被核函数(kernel)内的任意线程访问,相当于功能扩展型的共享内存

      • 其中,纹理内存的缓存由一级缓存(L1)实现,而常量内存则依托专用的常量缓存完成数据缓存

    • 本地内存(Local Memory)

      • 主存中为每个 SM 单独映射的专属存储区域,仅每个 SM 单独可访问

      • 当某一 SM 发生寄存器溢出且数据量超出一级缓存(L1)承载上限时,多余数据会被进一步转存至二级缓存(L2),最终写入本地内存

      • 寄存器溢出数据的重加载性能损耗会随数据调取所需遍历的存储层级增加而显著增大

    • 全局数据的一致性(write-through 策略)

      • 当 SM 对全局数据执行写操作时,L1 缓存中对应的数据行会立即写入 L2(写透策略),这样 L2 总是保持全局数据的最新版本

      • 因为 L2 是共享的,不同 SM 访问全局数据时,都会通过 L2 获取一致的值,GPU 不需要像 CPU 那样使用复杂的互连来维持各个 L1 之间的一致性

    • 本地数据的一致性(write-back 策略)

      • 当 SM 对本地数据执行写操作时,L1 缓存中的数据不必立即写回 L2(写回策略),因为 L2 中的本地内存是预留给单个 SM 的,不会被其他 SM 访问

      • 因此,L1 可以暂时保留修改,不必立即同步到 L2,这也可以提高性能,因为减少了写 L2 的频率

      • 当 L1 缓存行需要腾出空间时,才会写回 L2

  • 缓存行的大小设计

    • 4 字节操作数在 GPU 的缓存、寄存器和 CUDA 核心之间传输时,如果按 32 个为一组进行整体传输,会带来显著的性能优势

    • 这是因为一个包含 32 个操作数的分组,正好可以一次性为一个由 32 个线程组成的线程束(warp)提供所需的数据,因此 GPU 的缓存行通常设计为 32 × 4 字节,即 128 字节

    • 如果缓存行设计过小,例如 64 字节,一个线程束的 32 个操作数就需要分两次传输,从而增加内存访问延迟,显著降低并行计算效率

    • 而如果缓存行设计过大,则会造成缓存空间浪费,降低缓存命中率

CUDA

  • CUDA(Compute Unified Device Architecture)

    • CUDA 是 NVIDIA 提供的一套并行计算平台和编程模型,用于在 GPU 上运行高性能计算任务

    • 核心思想:让开发者可以使用类似 C/C++ 的语言编写程序,在 GPU 上并行执行大量线程

    • 功能:提供线程管理、内存管理、硬件加速指令等接口,让 GPU 能高效执行计算密集型任务(矩阵运算、卷积、模拟等)

  • nvcc 是 CUDA 的编译器(NVIDIA CUDA Compiler)

    • 作用:将 CUDA 代码(包含 CPU 端和 GPU 端代码)编译成可在 CPU 和 GPU 上运行的二进制文件

    • 将 GPU 核心代码(kernel)编译成 PTX 或 SASS,然后由 GPU 执行

    • 将 CPU 端代码编译成普通可执行程序

    • 例如,写了 .cu 文件,运行 nvcc my_kernel.cu -o my_program,就会生成可以调用 GPU 的程序

  • CUDA Toolkit 是 NVIDIA 提供的一整套开发工具包,包括 nvcc 编译器、GPU 库(cuBLAS、cuDNN、cuFFT 等高性能计算库)、调试工具、分析工具、示例代码和文档

  • NVIDIA 驱动是 GPU 与操作系统之间的桥梁

    • 提供底层硬件访问接口(管理 GPU 资源、执行 GPU 指令),让 CUDA 程序可以实际在 GPU 上运行

    • 即使安装了 CUDA Toolkit,没有安装兼容的 NVIDIA 驱动,GPU 也不能工作

    • 驱动版本必须与 CUDA Toolkit 兼容,否则可能报错或无法使用 GPU

  • PyTorch 与 CUDA

    • PyTorch 可以使用 CUDA Toolkit 提供的库(如 cuBLAS、cuDNN)来执行 GPU 运算

    • PyTorch 通过 CUDA API 调用 GPU,GPU 通过 NVIDIA 驱动执行具体指令,nvcc/Toolkit 用于在本地编译扩展 CUDA 内核(例如自定义 ops)

  • CUDA 的执行流程

    CPU invokes a kernel, which is structured as a 6-by-2 grid of blocks on the GPU; CPU can continue while grid executes on the GPU. - 一个 CUDA 程序依赖两个主要组件:主机(host)和一个或多个 GPU 设备(device) - 其中,主机负责执行主机代码(仅在 CPU 上运行的代码),GPU 设备负责执行设备代码(在 GPU 上运行的代码) - 典型的 CUDA 程序通常从主机代码开始执行,在某个时刻调用核函数(kernel),来告诉 GPU 开始执行设备代码 - 这里,设备代码的执行是异步的,也就是说,在核函数调用之后,主机代码会继续执行而不会等待 GPU 完成计算;为此,CUDA 提供了显式和隐式的主机与设备同步机制 - 通过调用核函数,CUDA 程序会在 GPU 上启动大量线程,这些线程会并行执行核函数,这些线程被组织成线程块(block),多个线程块又组成网格(grid) - 一次核函数调用只会启动一个网格,其内部的块和线程可以逻辑上排列成多维结构,每个线程块中的线程数量以及每个网格中的线程块数量,受 GPU 设备和其计算能力(Compute Capability) 的限制

  • cuBLAS(CUDA Basic Linear Algebra Subprograms)

    • cuBLAS 的本质是:一组针对 GPU 架构高度特化的线性代数 kernel 集合,以及一套在运行时选择合适 kernel 的调度逻辑

    • 它解决的核心技术问题是,在 GPU 的执行模型(SIMT、层级存储、有限寄存器、SM 并行)下,如何把以下操作做到接近硬件极限吞吐:

      • 向量运算(Level-1 BLAS)

      • 矩阵–向量运算(Level-2 BLAS)

      • 矩阵–矩阵运算,尤其是 GEMM(Level-3 BLAS)

    • 其中,GEMM 是绝对的核心,因为它满足三个条件:计算密集、数据复用强、可规则分块,非常适合 GPU

    • 在实现层面,cuBLAS 的工作可以拆解为几件非常具体的事情:

      • 把矩阵乘法拆分成 block / warp / thread 多级 tiling

      • 明确哪些数据放在 global memory、shared memory、register

      • 控制 shared memory 访问模式以避免 bank conflict

      • 通过增加并发 warp 数量来隐藏 global memory latency

      • 在支持的硬件上,将 FP16 / BF16 的矩阵乘法映射到 Tensor Core 指令

      • 根据矩阵尺寸(如是否是方阵、是否是小 batch)选择完全不同的 kernel 实现

    • 也就是说,cuBLAS 并不是“一个矩阵乘法实现”,而是几十甚至上百个不同实现 + 选择逻辑

    • 而在 API 层面调用的是同一个 cublasGemmEx,但在底层,走到的是完全不同的执行路径

    • 一个重要的技术事实是:cuBLAS 的性能高度依赖输入张量的布局和数据类型,例如:

      • 行主序 / 列主序是否需要转置

      • 是否是对齐的 leading dimension

      • batch GEMM 是否能合并

      • 是否满足 Tensor Core 对 shape 和对齐的要求

    • 这也是为什么在框架层面,经常通过改变 tensor layout、强制 contiguous、调整 reshape 顺序来“优化 cuBLAS 性能”

  • cuDNN(CUDA Deep Neural Network library)

    • cuDNN 并不是“比 cuBLAS 更高级”,而是针对一类特定计算模式提供了更直接的算子封装

    • 它解决的问题是:对于卷积、归一化、激活等结构化算子,如何在不同参数组合下选择合适的实现算法

    • 以卷积为例,从技术上看,卷积并不是一个单一算法,而是一族算法:

      • 直接卷积(direct convolution)

      • im2col + GEMM(本质上调用 cuBLAS)

      • FFT-based convolution

      • Winograd convolution

    • 这些算法在以下维度上的表现完全不同:

      • FLOPs 数量

      • 对 memory bandwidth 的压力

      • 中间 buffer 的大小

      • 数值误差特性

      • 对 kernel size / stride / dilation 的适配范围

    • cuDNN 的技术职责就是:

      • 为同一个“卷积算子语义”维护多种算法实现

      • 根据输入 shape、参数配置、数据类型、GPU 架构

      • 在运行前或运行时选择一个可用且性能最优的实现

    • 这里有一个很重要但常被忽略的点:cuDNN 并不保证选择“全局最优”的算法,它只保证在约束条件下选择一个可行且高性能的算法,例如:

      • workspace 内存受限时,会放弃更快但占用内存大的算法

      • 某些算法在 backward 上数值不稳定,会被禁用

      • 某些 shape 下 autotune 本身的代价过高

    • 从实现关系上讲,cuDNN 的很多路径最终会落到 cuBLAS,例如 im2col + GEMM 路径,本质是:

      Convmatrix reshapeGEMM\text{Conv} \rightarrow \text{matrix reshape} \rightarrow \text{GEMM}
    • cuDNN 负责 reshape 和算法选择,cuBLAS 负责真正的矩阵乘法 kernel

CUDA Device 与 Context

  • device 的定义

    • device 是一个计算位置的抽象描述,是一个可独立执行 kernel 的 GPU 实体,它决定 Tensor 的物理存储位置和运算 kernel 运行在哪类硬件上

    • 关键原则是,一个算子,只能作用在同一 device 上的 Tensor

    • Tensor 出生在哪个 device,就一直待在那里,PyTorch 不会自动移动数据的位置,且 model / data / intermediate 必须在同一个 device

    • cuda:0, cuda:1, cuda:2, ... 等只是一个逻辑编号,不等于进程/内存空间/执行上下文

  • device 的职责:执行 kernel、管理 device memory、提供 stream / event 等执行资源

  • CUDA Context

    • CUDA Context 是一个进程在某个 device 上的完整执行环境,它包含:

      • 已分配的 GPU 内存

      • kernel module

      • stream / event

      • cuBLAS / cuDNN 状态

      • cache / allocator 状态

    • Context ≈ GPU 版的进程地址空间

    • Context 在第一次 CUDA API 调用时创建,比如:

      • tensor.cuda()

      • model.to("cuda")

      • torch.cuda.current_stream()

    • 例如

  • 进程与 GPU

    • 一个进程在一个 device,对应一个 context,如

    • 但当一个进程多 GPU = 多 Context 时,Context 的切换代价昂贵、内存 allocator 分离、cuBLAS / cuDNN 状态不共享、stream 管理复杂

    • 因此,PyTorch 的设计是让一个进程只绑定一个 CUDA context,强制「一进程一张卡」

  • device 与 context 的映射

    • 基于「一进程一张卡」的设定,可以为当前进程设置其 rank,即 torch.cuda.set_device(local_rank)

    • 以上代码设置了当前进程默认使用的 CUDA device,在之后执行的操作都会自动分配到 rank 对应的 GPU 上

      • .cuda() → cuda:local_rank

      • torch.empty(..., device="cuda") → cuda:local_rank

    • 在实践中,需要尽早进行 set_device,这是因为如果在 set_device 之前创建了 CUDA Tensor,这时 Context 已经绑定到 cuda:0,之后再进行 set_device时,会拥有多个 context,可能会导致显存碎片、性能下降、DDP 行为异常等问题

  • 既然不同的 rank 对应的进程只控制一张 GPU,其 Context 完全隔离,那么梯度等如何同步?

    • 通过 NCCL、GPU 直接通信(NVLink / PCIe)等方式

    • 即,Context 不共享内存, 但可以通过通信原语交换数据

  • Context 与 Stream 的关系

    • Stream 是隶属于 Context 的,也就是说不同 context 的 stream 永远不能交互,所有异步行为都发生在同一个 context 内

    • 这就是为什么一进程一 context 是前提,否则 stream / event 会发生错误

  • 迁移模型:model.to(device) = 把 model 的所有参数 + buffer 移到 device

  • 迁移数据:DataLoader 生成的数据永远都是在 CPU 上,必须在训练循环中手动迁移到 GPU 中

  • 什么时候在 CPU 和 GPU 之前迁移数据

    • 每个 batch 的 x / y 一定在 forward 之前移动至 GPU 中,不能在 Dataset / collate_fn 中,因为 Dataset / DataLoader 是多进程,GPU context 不安全

    • GPU Tensor 在以下情况中回到 CPU 中

    • .item() 会强制同步 GPU,所以不要在训练主路径里调用 .item(),而日志可以 batch / step 间隔做,这也是训练慢的经典原因之一

  • 多卡训练前的基本约束(DDP)

    • DDP 的硬约束:一张卡一个进程,不能一个进程用多张 GPU,也不能一个 GPU 跑多个 DDP rank

    • device 绑定必须明确

    • 不能出现的行为

      • 在 Dataset 里 .cuda()

      • 在 forward 里新建 CPU Tensor

      • 在 loss 里混用 CPU / GPU Tensor

      • 在 rank != 0 做 IO(保存、打印)

    • 所有 rank 的计算图必须一致——每个 rank 必须执行完全相同的 forward / backward 结构,否则 DDP 的 AllReduce 会卡死,程序不报错,但会一直 hang

CUDA 执行模型

核函数(Kernels)

  • 核函数是运行在 GPU 核心上的代码

  • 可以用高级语言编写,如 CUDA 或 Triton,例如在 CUDA C++ 中

  • 编译后生成并行线程执行 Parallel Thread Execution(PTX),即 NVIDIA GPU 的低级汇编

  • 运行核函数还需要主机代码(host code),这些代码在 CPU 上执行,负责数据分配、加载数据和代码

  • kernel 与普通函数有两个本质区别,首先其不能返回值,其次会被大量线程并行执行

Each thread works on a different index (threadID) in an array of inputs to produce an array of outputs

CUDA 的执行层级

  • CUDA 里不是“一个程序顺序执行”,而是同一个 kernel 被大量线程并行执行,每个线程通过自己的索引(ID)知道应该运算哪一块数据

  • CUDA 的执行层级是一个树形结构

    Thread Configuration

  • Grid:一次 kernel launch 的全局并行空间

    • Grid 是一次 kernel 调用中,所有线程的集合

    • 当 launch kernel 时,会指定 <<<blocksPerGrid, threadsPerBlock>>>,例如 my_kernel<<<1024, 256>>>(); 的含义是这次 kernel 调用启动了 1024 个 block,每个 block 256 个 thread

  • Block:并行 + 协作的最小单位

    • Block 的作用:线程可以同步、线程可以共享内存

    • GPU 实际用于执行内核调用的 SM 数量受限于调用中指定的 block 数量

    • 以 kernel fun<<<M, N>>>(x, y, z) 为例,最多有 M 个 block 可以分配给不同的 SM

    • 一个 block 不能拆分到不同的 SM 中,如果 block 的数量多于可用的 SM 数量,那么可能会有多个 block 分配给同一个 SM,通过这种方式分配 block,GPU 可以在不同的 SM 上并行运行独立的 block

    • grid 和 block 的维度不能无限大,它们的最大值取决于设备的计算能力(Compute Capability,CC)

      • grid 最大维度通常为 (2^31−1) × 65535 × 65535

      • block 最大维度为 1024 × 1024 × 64,但由于每个 block 的最大线程数为 1024,因此实际可用维度通常远小,例如 32 × 32 × 1

  • Thread:真正执行指令的实体

    • Thread 是最小执行单元,拥有自己的寄存器、program counter、局部变量

    • Thread ≠ 独立核心

  • Warp:执行层面的真实最小单位

    • Warp 是 32 个 thread 组成的执行束,GPU 实际上是以 warp 为单位发指令,不是 thread,也不是 block

    • Warp 的执行特性是所有 thread 执行同一条指令,不同数据(SIMT)

    • 如果出现不同 thread 执行的指令不同,就会导致 warp divergence(分支发散),这时GPU 会先执行 A 分支,再执行 B 分支,而被 mask 掉的 thread 只能空转

    • block size 通常是 256 / 512,因为 256 = 8 warps、512 = 16 warps,便于调度,隐藏访存延迟,提高 SM 利用率

  • SM

    • SM 是 GPU 上可以独立执行多个 warp 的计算单元,一个 SM 包含多个 warp scheduler、ALU / Tensor Core、寄存器文件、shared memory

    • SM 决定了同时驻留多少 block,同时跑多少 warp 以及 shared memory / register 怎么划分

    • SM 的内部工作流程是 warp scheduler 选一个 warp,给这个 warp 发一条指令,32 个 thread 锁步执行,所以一个 warp ≈ 一条 SIMD 指令,而多个 warp 在 SM 内轮流执行(时间复用,从而隐藏延迟)

    • 例如,当一个 warp 等待 global memory 时,调度器会切换到另一个 ready warp,所以 GPU 性能关键是让 SM 里始终有足够多的 warp

Thread 的索引体系

  • 索引体系中的维度,和 GPU 的物理结构没有一一对应关系,而是逻辑索引空间

  • 实际上,Grid / Block / Thread → 程序员看到的逻辑并行空间,而 SM / warp / core → 硬件执行结构

  • 例如下列中的二维逻辑坐标系统,意味着 grid 是一个 8 × 4 的 block 网,每个 block 是 16 × 16 = 256 个 thread

  • 在实践中,数据通常是多维的,例如

    • 图像是 $(H, W, C)$,很自然地,blockIdx.x → 宽度方向、blockIdx.y → 高度方向、blockIdx.z → 通道方向

    • 矩阵乘法中,输出矩阵是 $(M, N)$,每个 block 负责一个 tile

    • 3D 数据(体数据、视频)的维度为 $(T, H, W)$​,那么 grid.z 直接对应时间 / 深度

  • 在这里,维度并不代表执行顺序,CUDA 不保证 block 的执行顺序,不同的 block 可以并行、可以乱序、可以延迟执行

  • block 内部索引:threadIdx

    • 线程在 block 内的编号,表示当前线程在它所属 block 里的位置

    • 是一个三维向量,包括 threadIdx.xthreadIdx.ythreadIdx.z

    • threadIdx.x:block 内 x 方向的线程编号,取值范围为 $0 \le \text{threadIdx.x} < \text{blockDim.x}$

  • block 维度大小:blockDim

    • 一个 block 有多少线程,表示一个 block 的尺寸

    • 也是三维向量,包括 blockDim.xblockDim.yblockDim.z

    • blockDim.x:block 在 x 方向有多少线程,是在 launch kernel 时指定的

  • block 在 grid 中的索引:blockIdx

    • 这是第几个 block,表示当前 block 在整个 grid 里的编号

    • 也是三维向量,包括 blockIdx.xblockIdx.yblockIdx.z

    • blockIdx.x:表示 grid 中第几个 block(x 方向)

  • grid 维度大小:gridDim

    • 一个 grid 有多少个 Block,表示一个 grid 的尺寸,即这个 Kernel 中包含多少个 Block

    • 也是三维向量,包括 gridDim.xgridDim.ygridDim.z

  • 全局 idx 计算

  • 如果数据大小在编译时未知,或数据元素数量超过最大线程数(无法一一对应),那么单线程对应单数据元素的策略就会失败,为了保证核函数鲁棒性,可以使用网格跨步循环(Grid-Stride Loops)

  • 网格跨步循环允许每个线程处理多个数据元素,每个元素的间隔为网格中线程总数。例如,在一维情况中:

  • 这样,每个线程可以处理多个数据,保证无论数据大小如何,都能完整覆盖所有输入元素

  • 索引体系的工程原则

    • block 内:threadIdx 连续 → 内存访问连续

    • warp 内:threadIdx.x 连续 → coalesced memory access

    • grid 维度:对齐数据形状,减少 if boundary check

线程映射举例

  • 下面的例子展示如何使用二维 grid 和 block 调用 CUDA 内核函数来计算矩阵转置

  • 原始的 CPU C 语言程序,用于转置一个方阵

  • 第一步优化,在 GPU 中,矩阵分配存在几个问题:

    • 多级指针访问效率低:如果使用 int **matrix,GPU 访问 [0][0] 时需要访问两次内存(先访问指针,再访问元素),效率低下

    • 数据连续性:GPU 访问内存时,会一次性读取一块固定大小的数据。非连续存储会导致多次访问,降低性能

    • 因此,通常将矩阵展平为一维数组:matrix[i][j] → matrix[i * WIDTH + j]

  • 第一步优化后结果

  • 注意:行是连续存储的,因此访问原矩阵的行是高效的;但列是跨行存储的,因此写入转置矩阵时仍然有一定开销,这是无法避免的

  • 在 GPU 中,可以使用二维线程块处理矩阵转置,每个线程负责一个矩阵元素

  • 运行后会发现转置矩阵全是 0,assert 失败。原因是 GPU 每个 block 最大线程数为 1024,而这里配置 1024×1024,超出了限制

  • 解决方案:将矩阵划分为 32×32 的子块,每个 block 处理一个子块,网格为 32×32 个 block,正好覆盖 1024×1024 矩阵

  • 这样每个线程块都只处理 32×32 的元素,符合 GPU 最大线程数限制

  • 进一步思考与挑战

    • 如何修改代码,使其能处理非方阵?

    • 是否有其他 grid/block 配置可以运行更快?

    • 最终版本的代码仅支持矩阵大小 ≤ 1024,如何修改,使其支持任意合理大小?

    • 如何实现 grid-stride loop,让每个线程处理多个元素?无论矩阵大小或 block/grid 配置如何,程序都能顺利运行

CUDA 内存模型

CUDA 的内存模型

  • 示意图如下

Memory Architecture

  • 在图中,双向箭头表示对内存的 读(R)和写(W)能力。箭头指向某个内存组件,表示具备写入能力;箭头从内存组件指向外部,表示具备读取能力

  • 通过观察可以发现,GPU 线程在某种程度上可以访问所有层级的内存;而 主机(CPU)只能访问全局内存(global memory)和常量内存(constant memory)。因此,当需要将数据加载到 GPU 设备上时,必须由主机参与,将数据拷贝到全局内存或常量内存中

  • 从上面的示意图可以看出,同一个线程块(block)中的线程可以通过共享内存(shared memory)相互通信;然而,某个线程块的共享内存无法被其他线程块中的线程访问

  • 即使是在同一个线程块内部,线程之间也不能访问彼此的寄存器(register)和本地内存(local memory)

  • 这里的“本地内存”并不一定对应一块物理上独立的内存区域,而是一个概念上的线程私有存储空间:当寄存器空间不足以容纳线程的私有变量时,这些变量会被放入本地内存中

  • 本地内存的实际存储位置可能位于 L1 缓存,并在必要时进一步延伸到 全局内存(global memory)

    内存位置
    作用范围(Scope)
    访问权限

    寄存器(Register)

    单线程(per-thread)

    读 / 写

    本地内存(Local memory)

    单线程(per-thread)

    读 / 写

    共享内存(Shared memory)

    线程块(per-block)

    读 / 写

    全局内存(Global memory)

    网格(per-grid)

    读 / 写

    常量内存(Constant memory)

    网格(per-grid)

    只读

带宽和容量

  • 不同的 GPU 内存组件在容量和带宽方面存在显著差异,全局内存(global memory)具有最大容量但最低带宽

  • 类似于 CPU 架构中的主存(RAM),全局内存的数据可以被缓存到片上 L1 和 L2 缓存中,以降低访问延迟

  • 与之相比,常量内存(constant memory)在所有计算能力(Compute Capability)下的容量固定为 64 KB;顾名思义,设备端对常量内存的访问是只读的,常量内存最初存储在全局内存中,但会被缓存在片上的常量缓存(constant cache)中;在特定的访问模式下(例如所有线程访问相同地址),常量内存的访问速度甚至可以接近寄存器

  • 共享内存(shared memory)的访问速度明显快于全局内存和未命中缓存的常量内存,并且可被同一线程块(block)内的所有线程访问。GPU 的计算能力决定了:

    • 每个流式多处理器(SM)可用的共享内存总量

    • 每个线程块可使用的最大共享内存大小

  • 例如,在 Quadro RTX 5000 GPU 上,每个 SM 的最大共享内存为 64 KB,同时每个线程块的最大共享内存也为 64 KB。当一个 SM 上同时驻留多个线程块时,这部分共享内存会在这些线程块之间进行划分

  • 每个线程都可以访问其私有的寄存器和本地内存(local memory)

  • 其中,寄存器(register)位于芯片内部,是所有内存层级中速度最快但容量最小的一种。GPU 的计算能力决定了:

    • 每个线程可使用的最大寄存器数量

    • 每个线程块的寄存器上限

    • 每个 SM 的寄存器总量

  • 以 Quadro RTX 5000 为例:

    • 每个线程最多可使用 256 个寄存器(约 1 KB)

    • 每个线程块最多 65,536 个寄存器

    • 每个 SM 总计 65,536 个寄存器(约 256 KB)

  • 当寄存器数量不足以容纳线程所需的数据时,部分变量会发生寄存器溢出(register spilling),并被存放到本地内存中

  • 本地内存的访问速度明显慢于寄存器,因为它最终位于全局内存中,尽管在访问过程中可能被缓存到 L1 或 L2 缓存中;本地内存是线程私有的,不会被其他线程访问

限定符

  • CUDA 支持多种内存空间限定符(memory space specifiers),用于在核函数中显式指定变量应当放置在哪一类内存中

  • 不同内存组件中的变量,不仅存储位置和访问速度不同,其作用范围(scope)和生命周期(lifespan)也存在显著差异

    变量声明
    存储位置
    作用范围
    生命周期
    访问权限

    int Var; 自动变量

    寄存器

    线程

    线程

    读 / 写

    int ArrayVar[N]; 自动数组

    本地内存(可能被缓存)

    线程

    线程

    读 / 写

    __shared__ int SharedVar;

    共享内存

    线程块

    线程块

    读 / 写

    __device__ int GlobalVar;

    全局内存

    网格

    应用程序

    读 / 写

    __constant__ int ConstVar;

    常量内存

    网格

    应用程序

    只读

  • Automatic Variables(自动变量)

    • 自动变量是不带任何内存空间限定符声明的变量

    • 当设备端代码声明一个自动变量时,网格中的每个线程都会拥有该变量的一份独立副本

    • 它位于每个线程的寄存器中,且仅能被该线程访问

    • 自动变量的生命周期持续到内核结束

    • 声明方式如下:

  • Automatic Arrays(自动数组)

    • 自动数组默认存放在每个线程的本地内存(local memory)中,同样只对当前线程可见

    • 本地内存是一个逻辑概念,实际上它最终位于全局内存中,因此访问速度可能较慢,但也可能会被缓存

    • 如果数组足够小,并且所有访问都使用编译期可确定的常量索引,编译器可能会将其优化并放入寄存器中

    • 自动数组的生命周期从线程创建持续到 kernel 结束

    • 声明方式如下:

  • Shared Memory:__shared__

    • 使用 __shared__ 内存空间限定符声明的变量,会被明确放入共享内存中,并由同一线程块(block)内的所有线程共享

    • 共享内存变量的生命周期从线程块创建持续到 kernel 结束

    • 以下两种声明方式是等价的:

  • Global Memory:__device__

    • 当变量仅使用 __device__ 修饰时,表示该变量存放在全局内存中,为整个网格中的所有线程可见

    • 生命周期:从程序开始到程序结束

    • 声明方式如下:

  • Constant Memory:__constant__

    • 使用 __constant__ 修饰的变量表示一个存放在常量内存中的只读变量,为整个网格中的所有线程可见

    • 存储位置:常量内存(底层在全局内存中,有专用常量缓存)

    • 生命周期:从程序开始到程序结束

    • 常量变量必须在函数体外声明,可选的 __device__ 修饰符不会改变语义

    • 以下两种声明方式等价:

  • Assign by Address(通过地址赋值)

    • 声明在共享内存、常量内存或全局内存中的变量,其地址可以被赋值给一个指针变量,用于间接访问,例如:

    • 自动变量(寄存器变量)通常无法安全地获取其地址

    • 通过指针访问全局或共享内存时,需要格外注意并发访问与同步问题

内存管理

  • 内存管理(Memory Management)指的是如何分配内存空间,以及如何在主机(host,CPU)和设备(device,GPU)之间传输数据

  • 在整体思想上,CUDA 的内存管理方式与 CPU 程序类似,但有一个非常关键的区别:主机代码负责管理所有需要在主机和设备之间共享的数据

  • 因此,一个典型的 CUDA 程序中,主机代码需要完成以下工作:

    • 在主机和 GPU 设备上分配内存

    • 将数据从主机传输到设备

    • 调用 kernel 在设备上进行计算

    • 在 kernel 执行结束后,将结果从设备传回主机

    • 释放之前分配的内存

  • 内存声明与分配(Memory Declaration and Allocation)

    • 主机内存分配

      • 在主机端,通常使用 malloc() 分配内存,并进行必要的数据初始化

      • malloc() 的特点是返回分配内存的起始地址,由调用者直接接收并保存到指针中

    • 设备内存分配

      • 在 GPU 设备端,使用 CUDA API 函数 cudaMalloc() 分配内存:

      • cudaMalloc() 不返回内存地址,它返回的是一个错误码

      • 分配得到的设备指针是通过参数传回的,通常将指针强制转换为 void **,以便适用于任意数据类型

  • 固定内存(Pinned Memory Allocation)

    • 默认情况下,通过 malloc() 分配的主机内存是可分页内存(pageable memory),即可能被操作系统换出到磁盘,当 GPU 要从该内存拷贝数据时,系统需要先确保数据在物理内存中,这会带来额外的 CPU 开销

    • 为提升主机与设备之间的数据传输速度,可以使用固定内存(pinned memory),即禁止该内存被换出到磁盘,从而允许 GPU 使用 DMA(Direct Memory Access)直接传输数据

    • CUDA 提供了 cudaMallocHost() 来分配并固定主机内存:

    • 这行代码替代了 malloc(),同时完成主机内存分配和内存固定

    • 固定内存能显著提升传输速度,但不宜滥用,大量固定内存可能降低系统整体性能

    • 如果主机端使用的是普通 malloc() 内存,CUDA 通常会先把数据复制到一个内部的固定缓冲区,再从该缓冲区传输到设备,这一步会引入额外的 CPU 开销

  • Kernel 调用前后的内存传输

    • 在内存分配并初始化完成后,需要使用 cudaMemcpy() 在主机与设备之间传输数据

    • cudaMemcpy 的参数

      • dst:目标指针

      • src:源指针

      • size:拷贝字节数

      • direction:传输方向,常用的传输方向包括:

        • cudaMemcpyHostToHost

        • cudaMemcpyHostToDevice

        • cudaMemcpyDeviceToHost

        • cudaMemcpyDeviceToDevice

    • 示例:主机 → 设备 → 主机

    • 非常重要的特性:cudaMemcpy() 是阻塞式(blocking)的,主机代码会一直等待,直到设备完成数据拷贝后才继续执行

  • 内存释放(Memory Deallocation)

    • 在 kernel 执行完成、数据传回主机之后,需要释放之前分配的内存

    • 主机端内存:使用 free()

    • 设备端或 CUDA API 分配的内存:使用 cudaFree()

    • 凡是通过 CUDA API 分配的内存(如 cudaMalloccudaMallocHost),都必须使用 cudaFree 释放

共享内存

  • 在设备端代码中,每个 kernel 函数负责管理其自身所需的设备内存,用于存储内部变量

  • 正如前面所讨论的,这些内部变量大多是自动变量和自动数组,也就是线程私有(thread-local)变量,其作用域仅限于 kernel 函数内部,通常存放在寄存器或本地内存中

  • 共享内存(shared memory)是设备内存中的一个特殊类型,CUDA 提供了明确的机制,用于规定同一线程块(block)内的线程如何使用共享内存

  • 与 kernel 作用域内定义的其他变量一样,共享内存变量和数组在 kernel 函数退出后不会继续存在

  • 在 CUDA 中,共享内存指的是一种位于芯片上的内存,对单个线程块私有,可被该线程块内的所有线程访问

  • 从线程访问角度看,共享内存从全局内存加载数据时,理论上可以快约两个数量级(~100 倍)

  • 共享内存的总容量以及每个线程块可用的最大共享内存大小,取决于 GPU 的计算能力(Compute Capability)

  • 静态共享内存分配(Static Shared Memory)

    • 如果共享内存数组的大小在编译期已知,可以在 kernel 函数体内使用 __shared__ 关键字进行静态声明:

    • 这种方式下数组大小在编译期确定,每个线程块都会分配一份大小为 SIZE 的共享内存数组

  • 动态共享内存分配(Dynamic Shared Memory)

    • 如果共享内存数组的大小在编译期未知,则可以使用 CUDA 提供的动态共享内存机制

    • 在 kernel 中,通过 extern __shared__ 声明一个不定长数组(或等价地,一个指针):

    • 此时,共享内存的实际大小需要在主机端调用 kernel 时指定,通过执行配置 <<<...>>> 中的第三个参数给出,单位是字节

    • 例如,为 sharememory 分配 nint

    • 为什么大小必须在 kernel 启动时指定?这是因为 kernel 函数体会被每个线程独立执行,如果允许每个线程自行决定共享内存大小,将导致线程块内的线程看到不一致的内存布局,因此共享内存的大小必须在线程块启动之前统一确定

  • 多个动态共享数组的实现方式

    • 动态共享内存看似有一个限制:只能声明一个不定长的 __shared__ 数组

    • 如果需要多个数组,可以采用以下做法:声明一个大的、连续的共享内存数组,然后使用指针在其中手动划分不同的子数组

    • nInF 作为 kernel 参数传入,用于告知所有线程各个数组的边界位置和长度,保证线程块内的线程对共享内存布局有一致的理解

    • 当不同数组类型的元素大小不一致(例如 charintfloat 混用)时,必须特别小心对齐(alignment)问题,否则可能导致性能下降甚至非法访问

统一内存

  • 从 Compute Capability 6.0 开始,CUDA 引入了 Unified Memory,也称为 Unified Virtual Memory(UVM)

  • Unified Memory 允许主机(CPU)和设备(GPU)使用同一块内存地址空间,而无需显式地在主机与设备之间拷贝数据(这并不意味着数据会同时驻留在主机内存和设备内存中)

  • 在 Unified Memory 模型下,数据会根据访问需求,在主机和设备之间按需双向迁移(on-demand migration)

  • Unified Memory 的工作机制:按需迁移 + 页错误

    • 在 Unified Memory 模型中,数据在任意时刻可能位于主机内存或设备(GPU)内存

    • 当主机或设备试图访问当前不在其内存空间中的数据时,会发生一次页错误(page fault),该页错误会触发 CUDA 运行时系统,将包含该数据的内存页(page)自动迁移到请求访问的一方

    • 页大小由操作系统决定,通常为 4 KB,因为迁移是以“页”为单位进行的,相邻数据也会一并被迁移,这有助于减少后续页错误的次数,从而降低整体迁移开销

    • 从编程者的角度来看,主机和设备仿佛“共享”同一块内存,不再需要显式调用 cudaMemcpy,不需要维护一对 host/device 指针,从而极大简化内存管理,只在数据真正被设备使用时才发生迁移,也更灵活,因为未被 GPU 使用的数据不会被拷贝到设备

    • 然而,页错误和按需迁移可能引入不可忽略的运行时开销;在某些高性能场景下,显式内存管理(手动 cudaMemcpy)仍可能更快、更可控

  • 下面示例展示了如何使用 Unified Memory 分配一个包含 N 个整数的数组,其中 x 是 同一个指针,主机代码和设备代码都可以直接使用它,不需要区分 h_x / d_x

  • UVA 与 UVM 的关系

    • 在较早的 CUDA 版本中,主机指针和设备指针是不同的,主机代码需要通过额外的 API 获取设备指针,再传给 kernel

    • 统一虚拟地址(Unified Virtual Addressing, UVA)解决了地址映射问题,主机和设备共享同一个虚拟地址空间,指针值在两端一致

    • Unified Memory(UVM)则是在 UVA 的基础上进一步自动化了数据迁移,通过分页机制在后台完成主机 ↔ 设备的数据移动

    • UVA 解决“地址一致性”,UVM 解决“数据在哪儿”

  • 页大小的影响

    • 在部分 TACC 系统中,内存页大小并非 4 KB,如 Vista 的 GH200 节点页大小为 64 KB,Frontera 的 RTX 5000 节点页大小为 4 KB

    • 较大的页大小意味着,更少的页错误(page faults),但每次迁移的数据量更大

    • 在 GH200 Superchip 上,CPU 与 GPU 通过 NVLink-C2C 高速互连,大页大小反而更有优势

    • NVIDIA 的 Unified Memory 驱动还会进行“密度预取(density prefetching)”,即多个页会被合并迁移,实际上传输的最小单位就是主机页大小

    • 此外,开发者也可以手动触发大规模预取,例如 cudaMemPrefetchAsync(ptr, size, device);,用于提前将 Unified Memory 迁移到指定设备,减少运行时页错误

CUDA Stream

CUDA Stream

  • GPU 的执行模式

    • 潜意识中,可能会认为 GPU 的执行流程是这样的:CPU 调用 → GPU 立刻执行 → GPU 算完 → 返回 CPU

    • 但实际上,CPU 不会等待 GPU 执行完成,CPU 创建任务并提交给 GPU 后就会立即执行其他流程,而 GPU 则是异步执行设备,它自己对 CPU 提交的任务进行排队和执行

  • CUDA Stream

    • CUDA 的设计目标是尽可能发挥并行能力,它要求开发者通过流(stream)来定义高层次的并行操作

    • CUDA Stream 是 GPU 上的一条指令队列(command queue),这条队列里可以包含

      • 计算任务:kernel launch

      • 数据拷贝:memcpy(H2D / D2H / D2D)

      • 通信任务:AllReduce

      • 事件:event

      • 清空内存:memset

    • 每次内核(kernel)调用都会被放入某个 stream 中,即使没有显式指定 stream,它也会使用默认流(default stream,又称 NULL 流或 stream 0)

    • 当一个指令被放入 stream 中时,它会进入该 stream 的先进先出(FIFO)队列,GPU 会先执行队列中第一条指令的线程块,执行完毕后,再执行下一条指令,直到队列中的所有指令执行完毕

    • 不同的 stream 默认没有顺序关系,可以并发执行(如果资源允许),这也是设计 Stream 的初衷,即允许 GPU 并行调度不同任务流,使得计算 A、同时拷贝 B、同时通信 C 可以并行执行,从而大幅提升 GPU 利用率——这也是一种计算与通信、数据拷贝中间重叠的一种关键实现

    • 一个 stream 的例子

    Streams

  • 使用 stream 的步骤如下:

    • 声明一个 cudaStream_t 类型的变量

    • cudaStreamCreate() 初始化 stream

    • 调用内核时,将流变量作为执行配置的第四个参数传入(第三个参数是内核动态分配的共享内存大小)

    • 程序结束时,用 cudaStreamDestroy() 销毁 stream

    • 示例代码

  • Default stream

    • 没有显式指定 stream 的操作,都会进入 Default stream(stream 0,NULL 流)

    • 例如 x = torch.matmul(a, b) 在 PyTorch 中,本质是会把 kernel launch → default stream

    • 早期设计中,Default stream 是特殊 stream,任何 stream 都要等待 Default stream,Default stream 也要等待所有 stream,也即 Default stream 被设计成了一个全局同步机制,这使得看似可以并行的多个 stream 变成了串行执行

    • 而在现在的设计中,改为了 Per-thread Default Stream,每个 CPU 线程拥有独立的 Default stream,不在自动同步其他 stream

    • default stream 的示例

    Streams

  • 强制同步

    • 不同 stream 的内核完成顺序是不确定的,如果需要控制不同 stream 中内核的执行顺序,CUDA 提供了同步机制(如 cudaStreamSynchronize、事件(event)等)

    • .item().cpu()torch.cuda.synchronize() 都会强制同步 GPU

    • 前面提到的 Dataloader 的 pin_memory=Truenon_blocking=True,就是为了让 memcpy 操作能够进入 Stream 中并行执行,否则 CPU 会同步等待

  • 异步内存传输

    • CUDA 的内存传输模型中,cudaMemcpy 占用主机 CPU 和默认流,这会导致一次只能进行一个同步内存传输

    • 如果使用非默认流,则可以使用异步内存传输(cudaMemcpyAsync),实现 GPU 与 CPU 的并行工作

    • 异步传输的关键是

      • 使用固定内存(pinned memory),即不会被交换到磁盘,使得系统可以使用 DMA(Direct Memory Access)直接在主机和设备间传输数据

      • 通过非默认流执行内存传输,避免阻塞默认流和其他非默认流,从而 CPU 和 GPU 可以同时继续执行其他任务

    • 示例代码

    • 异步内存传输的示例如下

      Streams

  • Stream 的一些误区

    • Stream ≠ 线程,CPU thread 可以控制多个 stream,而 stream 是 GPU 队列

    • Stream 不能强制并发、自动优化性能、避免资源冲突,它只是允许 GPU 有机会并发

    • 在自定义算子优化、通信计算重叠、流水线并行、数据预取等场景中,就需要自己写 Stream 来实现

    • 实际上,Stream 只定义依赖关系和执行顺序,而 CUDA Runtime / Driver 将任务翻译为 GPU 可执行任务,由 GPU Hardware Scheduler 来进行 SM 资源分配、Warp 调度、执行重叠、延迟隐藏、冲突回避(部分)

    • 因此,开发者就需要写出容易被 GPU 高效执行的代码

CUDA 同步

  • 在 CUDA 程序中存在许多并行组件,因此自然会有多种同步方法对应不同组件,主要有三类同步方法:

    • 主机端同步(Host-Side Synchronization)

    • 设备端同步(Device-Side Synchronization)

    • 线程级同步(Thread-Level Synchronization)

  • 同步的定义

    • 同步指的是将两个或多个进程带到已知的执行点,当一个进程必须等待另一个进程到达某个执行点后才能继续时,前者就被阻塞

    • 在 CUDA 中,由于存在大量并行进程,因此需要明确到底是哪一类进程在同步

  • 主机端同步

    • 主机端同步指的是主机(CPU)等待设备(GPU)或特定流完成任务

      • 主机在读取 GPU 输出之前,需要等待所有内核完成

      • 主机在启动新的内核前,需要确保设备完成某些任务

    • 主机端同步的方法:

      • cudaDeviceSynchronize():等待设备上所有流中的内核完成

      • cudaMemcpy(...)cudaFree(...):虽然本身用于内存操作,但它们隐式阻塞主机,等待设备完成当前队列中的所有内核

      • cudaStreamSynchronize(stream):只等待指定流完成

    • 一个例子

    Streams

    • 在上面的例子中,主机调用 cudaMemcpy,所有流被隐式同步;在调用 cudaDeviceSynchronize() 后,主机和所有流必须等待最慢的流完成

  • 设备端同步

    • 设备端同步发生在 GPU 内部,它允许非默认流与其他非默认流同步,默认流与所有流同步

    • 设备端同步不影响主机的行为(主机端同步可以影响设备)

    • 在默认流中启动内核时,会阻塞默认流直到其他流的活动内核完成;之后在默认流中启动的内核,也会被阻塞,直到默认流完成

    • 类似行为也适用于 cudaMalloccudaMemcpy

    • CUDA 提供了事件(event)机制来在流间同步,事件类型:cudaEvent_t,相关函数有

      • cudaEventCreate():创建事件

      • cudaEventRecord():在流中记录事件

      • cudaStreamWaitEvent():让另一个流等待事件

    • 一个示例如下:

    Streams

    • 上述例子的执行流程:

      • 主机创建事件 end_2 并初始化

      • 主机在 stream2 的内核执行到某个点时调用 cudaEventRecord(end_2, stream2)

      • stream3 调用 cudaStreamWaitEvent(stream3, end_2),阻塞自己直到事件被触发

      • stream2 内核执行到事件点,记录事件,stream3 收到信号后继续执行内核

      • 程序结束后,需要调用 cudaEventDestroy() 销毁事件

    • 伪代码示例:

  • 线程级同步

    • 在 CUDA 中,同一个线程块(block)内的线程通常不会同时完成任务,为了确保线程块内的所有线程在某个应用点同步,需要使用同步屏障(synchronization barrier)—— __syncthreads()

    • 注意:__syncthreads() 只作用于同一个线程块内的线程,其作用范围比之前介绍的其他 CUDA 同步函数(如 cudaDeviceSynchronize()cudaStreamSynchronize())更有限

    • __syncthreads() 最常用于防止竞态条件(race condition)

    • 竞态条件的典型场景是:线程先把计算结果写入共享内存(shared memory),然后又尝试读取这些值,如果线程没有同步,就可能出现“读取到还未计算或存储的值”的情况

    • 因此,如果一个线程使用了另一个线程更新的变量,就必须指定线程的执行顺序,否则结果不可预测

    • 示例代码:

    • 在这个例子中,如果没有 __syncthreads(),线程可能在共享内存 val 中访问还未被其他线程计算或写入的元素,导致计算结果错误

    • 线程同步通常会带来性能损失:

      • 最早到达 __syncthreads() 的线程必须等待最后一个线程

      • 如果某些线程执行时间较长,资源会被占用在同步点上

      • 同样的性能损失也存在于其他 CUDA 同步函数中

      • 尽管如此,为了保证正确性,在必要时使用 __syncthreads() 是必须的;开发者可以考虑是否有其他无需同步就能计算的方式,以避免同步开销

    • 使用 __syncthreads() 的注意事项

      • 所有线程必须执行到同一个 __syncthreads(),如果某些线程由于分支语句(if/else)跳过了同步点,而其他线程到达了同步点,就会造成死锁(deadlock)——死锁的线程永远无法继续执行

      • 正确的编程规则是,对于每个线程块,每个 __syncthreads() 要么被块内所有线程执行,要么都不执行,避免在条件分支中单独执行同步屏障,除非保证条件一致

使用示例

  • 下面的例子演示了如何使用流(streams)和异步内存传输(asynchronous memory transfer)

  • 原始代码实现了一个矩阵转置内核,每个线程负责矩阵的一个元素:

  • main() 中,原始矩阵在主机(host)分配、初始化后复制到设备(device),调用内核执行转置,然后将结果复制回主机并验证

  • 为了展示流,需要不止一个内核调用,这里将原来的一次转置改成两次,分别操作两个不同的矩阵:

  • 使用流

  • 使用流可以让两个内核同时执行,从而实现并行化

  • 为了实现异步传输,需要:

    • 固定内存(pinned memory):使用 cudaMallocHost 替代 malloc

    • 流(stream):将异步内存传输与流绑定

    • 设备同步:cudaDeviceSynchronize() 确保主机在验证结果前等待数据传回

  • 最终实现

  • 注意:即使没有使用 pinned memory,cudaMemcpyAsync 也不会报错,只是会退化为同步传输行为

内核编写

  • 如果想增加一个新的操作,而该操作还没有经过优化后的 GPU 内核实现

  • 或者想加速现有的 PyTorch 函数,那么自己实现高性能 CUDA 内核

  • 看似直接,但实际上门槛很高:需要丰富的经验,学习曲线陡峭

  • 通常,更好的方法是利用 torch.compile,它可以动态优化 PyTorch 代码,捕获操作,并使用 Triton 作为后端生成底层高性能内核

  • 一个简单的例子:ELU 激活函数

    • ELU 激活函数的定义为:

      ELU(x)={α(ex1),x<0 x,x0\text{ELU}(x) = \begin{cases} \alpha (e^x - 1), & x < 0 \ x, & x \ge 0 \end{cases}
    • 可以先用 PyTorch 写一个普通实现:

    • 然后只需加上 @torch.compile

    • 这样 PyTorch 会自动生成更高性能的内核,实测性能提升显著,尤其是在大规模矩阵上

    • 可以设置环境变量,来查看 torch.compile 生成的 Triton 内核

    • 然后运行 Python 脚本,PyTorch 会输出对应的 Triton 内核,例如:

    • 为了可读性,可以重命名变量、加注释:

    • 对 Triton 内核代码的解释:

      • tl.program_id(0) 返回 block ID,用来确定每个 block 处理数据的区间

      • block_indices 表示该 block 内的数据索引

      • valid_mask 确保不会越界

      • ELU 函数通过 tl.where 应用

      • 结果通过 tl.store 写回 GPU 内存

    • 通过 Triton,甚至可以在小规模数据上得到比 torch.compile 更快的性能,因为编译开销低

  • 性能优化工具对比

    工具
    难度
    性能
    灵活性

    PyTorch 原生

    容易

    @torch.compile

    容易

    中等

    Triton

    中等

    更快

    CUDA

    最难

    最快

    最高

  • 上述是一个简单的内核例子,而一个完整的 Kernel 编写一般分为以下步骤

    • 并行映射设计:每个 thread 到底负责算什么

      • 最基本的映射是逐元素映射:比如矩阵乘法中,每个 thread 负责计算结果矩阵中的一个元素,这是最理想的情况,因为其完美符合并行设计:完全数据并行、没有依赖关系、访存连续

      • 多维映射设计:如对于图像而言,其维度为 H × W × C,可以直接利用内存组织结构的三维进行映射

      • 映射设计的核心目标是保证并行度足够(否则 GPU 会空闲)、保证访存规则(否则访存碎片化)

    • Block 设计策略

      • block 内线程可以共享 shared memory、同步、共同完成 tile 计算

      • 一般而言,block 大小为 128 ~ 512 threads,通常是 32 的倍数

      • block 主要影响 Occupancy,block 太大,其 shared memory 不够、register 不够;block 太小,其 warp 不够,延迟无法隐藏

  • Kernel 如何优化

    • 如何组织内存访问:减少访问 global memory

    • 如何提高计算密度:在等数据的时候让 GPU 不会空闲

    • 如何划分并行任务:让访问更并行

  • 而这正好和 GPU 的性能瓶颈模型一一对应

GPU 性能瓶颈模型

  • GPU 的性能,本质上取决于:

    • 计算资源(ALU / Tensor Core)

    • 内存带宽(从 DRAM 到寄存器的链路)

    • 并行度(SM 内的 warp 调度、线程数)

  • 计算平台的两个指标

    • 算力 $\pi$​

      • 也称为计算平台的性能上限,指的是一个计算平台倾尽全力每秒钟所能完成的浮点运算数,单位是 FLOP/s

        π=Ncore×f×IPC\pi=N_{\text{core}}\times f\times\text{IPC}
      • $N_{\text{core}}$ 是计算单元数量,$f$ 是主频,IPC 是每周期可完成的运算量(Tensor Core 场景下通常远大于 1)

    • 带宽 $\beta$

      • 也即计算平台的带宽上限,指的是一个计算平台倾尽全力每秒所能完成的内存交换量,单位是 Byte/s

        β=Bus Width×Data Rate\beta = \text{Bus Width}\times\text{Data Rate}
    • 计算强度上限

      • 两个指标相除即可得到计算平台的计算强度上限,它描述的是在这个计算平台上,单位内存交换最多用来进行多少次计算,单位是 FLOP/Byte

  • 模型的两个指标

    • 计算量

      • 指的是输入单个样本(对于CNN而言就是一张图像),模型进行一次完整的前向传播所发生的浮点运算个数,也即模型的时间复杂度

      • 单位是 #FLOP or FLOPs,例如卷积层的计算量公式如下:

        Conv Layer Time Complexity:M2K2CinCout(FLOPS)\text{Conv Layer Time Complexity}: M^2\cdot K^2\cdot C_{in}\cdot C_{out}\quad(\text{FLOPS})
      • $M$:每个卷积核输出特征图(Feature Map)的边长

      • $K$:每个卷积核(Kernel)的边长

      • $C_{in}$:每个卷积核的通道数,也即输入通道数,也即上一层的输出通道数

      • $C_{out}$:本卷积层具有的卷积核个数,也即输出通道数

      • 其中,输出特征图尺寸本身又由输入矩阵尺寸 $X$ 、卷积核尺寸 $K$ 、$Padding$、 $Stride$​ 这四个参数所决定,表示如下

        M=(XK+2Padding)/Stride+1M=(X-K+2* Padding)/Stride+1
    • 访存量

      • 指的是输入单个样本,模型完成一次前向传播过程中所发生的内存交换总量,也即模型的空间复杂度

      • 在理想情况下(即不考虑片上缓存),模型的访存量就是模型各层权重参数的内存占用(Kernel Memory)与每层所输出的特征图的内存占用(Output Memory)之和,单位是Byte

      • 由于数据类型通常为 float32 ,因此需要乘以四

        Conv Layer Space Complexity:(K2CinCout+M2Cout)4(Bytes)\text{Conv Layer Space Complexity}:(K^2\cdot C_{in}\cdot C_{out}+M^2\cdot C_{out})\cdot 4\quad(\text{Bytes})
    • 模型的计算强度/算术强度(Arithmetic Intensity) $I$

      • 由计算量除以访存量就可以得到模型的计算强度,它表示此模型在计算过程中,每Byte内存交换到底用于进行多少次浮点运算,单位是 FLOPs/Byte

        I=FLOPsBytes accessed\text{I}=\frac{\text{FLOPs}}{\text{Bytes accessed}}
      • 可以看到,模计算强度越大,其内存使用效率越高

    • 模型的理论性能 $P$:最关心的指标,即模型在计算平台上所能达到的每秒浮点运算次数(理论值),单位是 FLOPSFLOP/s

  • 算力瓶颈

    • 在算力瓶颈中,GPU 的执行单元(CUDA Core、Tensor Core、SFU 等)已经接近满负载运行,而内存系统还有富余,此时程序的性能主要受限于单位时间内能执行多少条算术或矩阵指令

    • GPU 的峰值算力也即上文提到的计算平台的指标之一——算力

    • 如果 kernel 的实际吞吐率已经逼近 $\pi$,那么再去做诸如减少访存、增加 cache 命中率之类的优化,往往几乎看不到收益,因为瓶颈不在内存而在计算单元本身

    • 典型例子包括大规模矩阵乘、卷积在高复用率下的核心计算阶段

  • 内存带宽瓶颈

    • 当 kernel 的执行速度主要受限于数据拷贝速度时,就进入了内存带宽瓶颈区间,这在 embedding lookup、稀疏访问、注意力中的 KV cache 读取等场景尤为常见

    • GPU 的显存系统有一个理论带宽上限,也即上文提到的计算平台的指标之一——带宽

    • 而 kernel 的实际性能往往取决于单位时间内需要从显存中读写多少字节

    • 一个非常核心的衡量标准是模型的算术强度,当 $I$ 很低时,哪怕 GPU 的算力再强,也会因为“算得太少、搬得太多”而被内存拖慢,这类瓶颈的本质不是“GPU 不够快”,而是“数据供不上”

  • Roofline 模型:模型在一个计算平台的限制下,到底能达到多快的浮点计算速度

    • Roofline 认为模型的理论性能上限满足

      P={βI,when I<ImaxMemory Boundπ,when IImaxCompute BoundP = \begin{cases} \beta \cdot I, & \text{when } I < I_{\text{max}} \quad \text{\color{red}Memory Bound} \\ \pi, & \text{when } I \ge I_{\text{max}} \quad \text{\color{green}Compute Bound} \end{cases}
    • 其示意图如下

    img

    • 可以看到,最高算术强度为

      Imax=πβI_{max}=\frac{\pi}{\beta}
    • 当 $I \times \beta < \pi$ 时,性能由内存带宽决定;反之则由算力决定——算力决定屋顶的高度,而带宽决定屋檐的斜率

    • Roof-line Model 解决的,是计算量为 A 且访存量为 B 的模型在算力为 C 且带宽为 D 的计算平台所能达到的理论性能上限 E 是多少这个问题

    • 在大模型推理中,attention 的 score 计算往往接近算力屋顶,而 KV cache 的读取阶段则明显落在带宽斜线区域

  • 延迟与并行度瓶颈(Latency & Occupancy)

    • 即便既不是算力瓶颈,也不是带宽瓶颈,GPU 程序仍然可能很慢,其原因通常是延迟无法被隐藏

    • GPU 的核心思想是用大量线程来掩盖长延迟操作(如 global memory load):如果每个 SM 上活跃的 warps 数量不足,或者寄存器/共享内存占用过高导致 occupancy 下降,那么一旦线程遇到长延迟,就没有其他 warp 可以切换执行,硬件只能空转

    • 这类瓶颈的典型表现是:理论带宽没打满、算力也没打满,但 SM 利用率很低

    • 根源不是硬件能力不足,而是并行度不够或资源分配不合理

    • 可以用 Occupancy 来衡量并行度:

      Occupancy=Active Warps per SMMaximum Warps per SM\text{Occupancy} = \frac{\text{Active\ Warps\ per\ SM}}{\text{Maximum\ Warps\ per\ SM}}
    • SM 的资源是有限的,比如 register 总量、shared memory 总量、最大线程数、最大 block 数等都是有限的

    • 假如一个 SM 有 65536 个 register,每个 thread 用 64 个 register,那么线程数最多为 1024 个

    • Occupancy 越高,说明正在使用的 warp 越多,但每个 warp 算的慢(内存多、复用少),而使用 Register Blocking 后,使用的 warp 变少,但是每个 warp 计算密度高

    • 很多时候,低一点 occupancy + 高算术强度 = 更快,在真实高性能 kernel 中 Occupancy 常常只有 30%~60%,但此时性能是峰值,因为内存访问已经被压到最低

  • 控制流与指令级瓶颈

    • GPU 是 SIMT 架构,warp 内线程执行同一条指令;当存在严重分支发散时,同一个 warp 内的不同路径会被串行执行,等价于降低了有效并行度,而这类瓶颈在复杂条件判断、稀疏结构遍历中很常见

    • 此外,指令混合也可能成为限制因素,例如大量依赖链导致指令级并行度(ILP)不足,或者 SFU 指令(exp、log、sqrt)占比过高,使特定执行单元成为热点

  • 系统层瓶颈(Kernel 粒度与调度)

    • 在真实系统中,GPU 性能还可能被更“外层”的因素限制,例如 kernel 太小导致 launch overhead 占比过高,或者频繁 kernel 切换、PCIe/NUMA 访问、CPU–GPU 同步等

    • 这些瓶颈不体现在单个 kernel 的 Roofline 上,但会显著拉低端到端吞吐

    • 在大模型 serving 场景中,batch size 过小导致 GPU 处于“算力和带宽都没吃饱”的状态,本质上就是一种调度层面的瓶颈

  • GPU kernel 的性能瓶颈,本质上不是由某一个单独因素决定的

    • 而是由 Roofline 模型、并行度(Occupancy)以及内存层级结构(Memory Hierarchy)三者的共同约束所决

    • Roofline 给出了算力与带宽的理论上限,Occupancy 决定这些上限能否在现实执行中被逼近,而内存层级则决定数据在不同速度层次之间流动时的真实成本

    • 这三者叠加在一起,构成了判断一个 kernel “为什么慢、慢在哪里、还能不能再快”的完整框架

  • 因此,从结果形态上看,kernel 的瓶颈通常会落入三种基本类型之一

    • 内存带宽受限:也就是算术强度较低,性能主要由 DRAM 或更低层级内存的吞吐能力决定;这时优化的核心不是“算得更快”,而是“让数据走得更短、更顺”,例如使用 shared memory 做 tiling、通过更好的 memory coalescing、减少不必要的全局访存,或者避免 shared memory 的 bank conflict

    • 算力受限:即 kernel 已经接近 GPU 的峰值 FLOPs,此时性能提升依赖于更高效地利用计算单元,包括更合理的寄存器使用(Register Blocking)、更高的指令级并行度、Tensor Core 的充分利用,以及更好的指令调度、Fused Kernels

    • 延迟受限:这类 kernel 往往既没有打满算力,也没有打满带宽,但仍然很慢,其根本原因在于延迟没有被足够的并行度隐藏,此时需要从 warp 调度、occupancy、Launch Overhead、persistent kernel 或线程粗化等角度入手

  • 在实践中,一条相对稳健的优化路径通常是:先通过算术强度判断算力与带宽的主导关系,再检查寄存器、shared memory 和 warp 分布等资源占用情况,明确是否存在并行度或调度层面的限制,随后选择与瓶颈类型相匹配的优化手段,并在每一次修改后重新测量性能,观察瓶颈是否发生转移

  • Memory-bound 优化路线:降低访存次数,提高数据复用率

    • 第一步:访问模式——如何高效地将数据从 DRAM 拷贝到片上存储

      • Memory Coalescing:warp 内线程访问连续地址以充分利用 cache line / memory segment

      • Layout Transformation(AoS ↔ SoA):调整数据存储布局以优化访问连续性

      • 数据重排 / Padding / Alignment:对齐数据、填充空位以减少非对齐访问和伪共享

    • 第二步:数据复用——通过尽可能复用已经拷贝的数据,来减少内存访问次数

      • Block Tiling:将全局内存数据划分为可重复利用的块,减少 DRAM 访问

      • Warp Tiling:在 block 内进一步划分 warp 工作负载,实现局部复用

      • Thread / Register Tiling:更细粒度的数据复用,充分利用寄存器和线程局部存储

    • 第三步:片上存储优化——优化片上存储的并行访问

      • Bank Conflict:优化 shared memory 访问,避免不同线程访问同一 bank 产生冲突

      • Register Blocking:通过寄存器划分和复用,提高片上并行访问效率

  • Compute-bound 优化路线:提升指令级并行度与算力密度,充分利用专用硬件单元

    • 第一步:指令级并行度与流水线效率——执行单元是否能保持高利用率,减少依赖链阻塞

      • Thread Coarsening:增加每个线程的工作量以提升指令级并行度(ILP)

      • 减少长依赖链:优化指令调度,缩短关键路径

    • 第二步:Kernel 优化——减少不必要的延迟开销和气泡,提升算力密度

      • Fused Kernels:将多个 kernel 融合,减少中间结果读写和依赖延迟

      • Pipeline Kernel:流水线处理多个计算阶段,提高连续计算密度

    • 第三步:专用计算单元与指令调度

      • Tensor Core 利用:调用专用硬件单元加速矩阵运算

      • Instruction Scheduling:优化指令顺序,减少 SFU 或其他执行单元热点

      • 该层通常由框架或编译器完成,但手写 kernel 时仍然关键

  • Latency-bound 优化路线:隐藏硬件不可避免的延迟,保持执行单元持续工作

    • 第一步:Occupancy 与可调度性——是否有足够可调度的 warp 来隐藏延迟

      • Register / Shared Memory 占用控制:限制资源使用以提升 occupancy

      • Block 配置优化:合理设置线程块大小,避免过度占用资源

      • Warp 数量充足:保证 scheduler 有足够 warp 切换执行

    • 第二步:调度结构优化

      • Persistent Kernel:持久化 kernel 以减少重复调度开销

      • Launch Overhead:降低启动 kernel 的延迟

      • Software-managed work queue:软件层面管理任务分发,提升调度效率

    • 第三步:执行路径稳定性——减少分歧和同步

      • 最小化 Control Divergence:减少 warp 内分支差异

      • 减少不必要的同步:降低延迟叠加

      • 分支和同步优化:保持峰值算力不变的情况下降低延迟放大效应

Memory bound 优化之访存模式

Memory Coalescing

  • GPU 访问全局内存(DRAM)相比访问寄存器或共享内存要慢很多,所以需要尽量减少访问全局内存的开销

  • GPU 的内存组织

    • 内存合并访问指的是将一个 warp(32 个线程)的全局内存访问请求合并成一次单一的内存操作

    • 实现这种优化最容易的情况是,当访问的数据连续存储在内存中,并且从一个良好的内存对齐边界开始

    • 当 warp 中的 32 个线程都参与访问时,这类请求相当于从 L1 缓存中加载一个缓存行(cache line),速度非常快

    • NVIDIA GPU 的缓存行大小设计为 32 个单精度浮点数(4 字节/float),总共 128 字节,然而,大多数类型的 RAM 数据通道宽度通常不超过 64 字节,这意味着从全局内存填充一个 GPU 缓存行需要多次事务,为此,NVIDIA 将 GPU 设计为每次可以访问全局内存中的连续 32 字节(8 个单精度浮点数)

    • 因此,一个缓存行被划分为 4 个“扇区”(sector),每个扇区 32 字节,可以独立从全局内存加载。例如,如果 warp 的 32 个线程访问内存中连续的 32 个单精度浮点数,并且数据不在 L1 或 L2 缓存(发生 L2 cache miss),则最少需要 4 次全局内存事务才能加载全部数据;但实际可能需要更多次,因为内存访问必须对齐到特定地址

  • DRAM 的突发传输(burst)特点

    • DRAM 不是一次只返回一个数据,它一次访问一个地址时,会顺便把这个地址附近的一段连续数据也读出来

    • 这些连续数据由 DRAM 内的多个传感器并行读取,然后通过高速通道一次性传给处理器

    • 也就是说,如果程序访问连续地址,DRAM 可以一次性高效地返回一大块数据,而不是每次访问都慢慢取

  • 使用 cudaMalloc() 分配的内存至少按 256 字节对齐,这有助于保证内存访问能够合并,然而,以下情况仍可能导致内存访问无法合并:

    • 数组索引不是连续的

    • 内存访问稀疏,例如通过结构体访问

    • 内存访问未对齐,例如数组索引与线程索引存在偏移

  • 内存访问模式示例

    • 连续且对齐(Sequential and Aligned):如果 32 个连续线程访问连续内存地址,访问是顺序且对齐的,因此可以合并为一次操作,只需一次事务即可加载 32 个单精度浮点数

    Aligned and Consecutive

    • 对齐但不连续(Aligned but Non-sequential):内存地址虽然对齐,但线程访问的不是连续地址。在现代 NVIDIA GPU 上,这种模式仍可能被合并为一次操作,但早期计算能力的 GPU 不支持这种合并

    Not Consecutive

    • 不连续访问(Not Consecutive):内存访问既不连续也可能未对齐,需要多次内存事务才能完成

    Misaligned

    • 未对齐访问(Misaligned Access):如果访问起始地址未对齐,即使访问连续内存,也可能需要两次或更多内存事务才能完成

    Misaligned

  • memory coalescing

    • 内存合并的目标是让同一个 warp 内的线程访问连续的内存地址

    • 如线程 0 访问 M、线程 1 访问 M+1、线程 2 访问 M+2,以此类推

    • 这样 GPU 硬件可以把这些访问请求“合并”为一次大的 DRAM 访问请求(一次突发传输),效率比每个线程单独访问高很多

  • 朴素的矩阵乘法实现

  • Simon Boehm 的 blog postarrow-up-right 可视化说明了这个过程

image.png

  • 上述 kernel 是在做矩阵乘法,而矩阵天然是二维的,行 → x、列 → y

  • 所以常见的映射方式是一个线程 ↔ 矩阵里的一个元素,用 (x, y) 表示这个元素的位置

  • 在这个实现中,每个线程负责计算输出矩阵 $C$ 中的一个元素 $C[x, y]$,计算公式是

    C[x,y]=i=0K1A[x,i]B[i,y]C[x, y] = \sum_{i=0}^{K-1} A[x, i] \cdot B[i, y]
  • 对应到代码里,x:输出矩阵的行号,y:输出矩阵的列号,循环里的 i:沿着 $K$​ 维度做点积

  • 对于一个 warp 中的线程 $(x,y)$ 和 $(x+1,y)$(即 threadIdx.x 连续、threadIdx.y 相同)

  • 在同一次循环 $i$ 中,线程 $(x,y)$ 访问 $A[x, i]$,线程 $(x+1,y)$ 访问 $A[x+1, i]$

  • 由于矩阵 A 采用 row-major 存储,同一行的数据在内存中是连续的,不同行之间相隔 $K$ 个 float

  • 因此 $A[x,i]$ 和 $A[x+1,i]$ 在内存中相隔 $K$ 个 float,不是连续地址,导致 warp 内的全局内存访问无法合并(uncoalesced)

image.png

  • 原有方法使用 2D block,threadIdx.x → 列方向,threadIdx.y → 行方向,使得同一个 warp 内的线程 x 不同,导致访问的是 A 的不同“行”

  • 可以改为使用 1D block,用一个 threadIdx.x 手动“拆”出二维坐标:

    x=blockIdx.xBLOCKSIZE+threadIdx.xBLOCKSIZEy=blockIdx.yBLOCKSIZE+(threadIdx.xmodBLOCKSIZE)x = \text{blockIdx.x} \cdot \text{BLOCKSIZE} + \left\lfloor \frac{\text{threadIdx.x}}{\text{BLOCKSIZE}} \right\rfloor\\ y = \text{blockIdx.y} \cdot \text{BLOCKSIZE} + (\text{threadIdx.x} \bmod \text{BLOCKSIZE})
  • 直观理解上,threadIdx.x / BLOCKSIZE → 行内编号(决定 x),threadIdx.x % BLOCKSIZE → 列内编号(决定 y)

  • 这也使得同一个 warp 中的线程 x 相同(同一行),而 y 连续变化(不同列)

  • 对于矩阵 A,访问模式是 $A[x \cdot K + i]$,由于 warp 内 x 相同,i 相同(warp 锁步执行),所以所有线程访问的是同一个 A 元素,要么命中 L1 / L2 cache,要么由硬件自动广播

  • 对于矩阵 B,访问模式是 $B[i \cdot N + y]$,由于 warp 内 i 相同,y 连续,因此线程访问的是:

    B[iN+y0],;B[iN+y0+1],;B[iN+y0+2],;B[i \cdot N + y_0],; B[i \cdot N + y_0 + 1],; B[i \cdot N + y_0 + 2],;\dots
  • 而矩阵 B 是 row-major 存储的,这正好对应内存中的连续地址,完美符合 DRAM burst 的访问模式,内存合并成功发生

  • 经过改进后,内存带宽提高 10 倍,执行时间下降 10 倍

Layout transformation

数据重排

Memory bound 优化之数据复用

Block Tiling

image.png

  • 问题分析

    • 在 GPU 的内存层次中,全局内存的容量大,但慢,而共享内存的容量小,但非常快,且一个 block 内的线程都能访问

    • 问题在于,在朴素矩阵乘法中,每个线程都会从全局内存中独立读取自己需要的 A 的一行和自己需要的 B 的一列

    • 但同一个 block 内的线程实际上会反复读取大量相同的 A 和 B 元素

    • 结果导致全局内存被重复访问,带宽浪费严重,算力被共享内存拖慢

    • 因此,Tiling 的核心思想是把“大家都会用到的数据”,只从全局内存读一次,然后放进共享内存反复用

  • 以矩阵乘法为例

    • 假设进行运算:$C = A \times B$

    • Tiling 会把大矩阵拆成很多小块:

      • A 的一个 tile:$\text{BLOCK_SIZE_M} \times \text{BLOCK_SIZE_K}$

      • B 的一个 tile:$\text{BLOCK_SIZE_K} \times \text{BLOCK_SIZE_N}$

    • 这些 tile 的大小通常正好适合放进共享内存,并且和 block 中线程的排布对应

  • 一次 Tiling 迭代,本质上就是沿着 $K$ 维“推进一步”

    • 协作加载

      • block 内所有线程一起行动,从全局内存中加载一个 A 的 tile,加载一个 B 的 tile

      • 每个线程只负责加载其中的一小部分,加载完成后,数据放在共享内存里

      • 这一步的特点是:全局内存访问次数少,访问模式可以做到内存合并

    • 同步:使用 __syncthreads(),确保 A、B 的 tile 已经全部加载完成,没有线程读到未初始化的数据

    • 在 shared memory 中计算

      • 现在每个线程需要的数据都已经在共享内存里,访问延迟极低

      • 于是线程就可以计算 $C_{\text{tile}} += A_{\text{tile}} \times B_{\text{tile}}$

      • 并把结果累加到一个寄存器里的 accumulation matrix 中

      • 这个 accumulation 是 跨 tile 的,不会每次清零

    • 沿着 $K$ 维移动到下一个 tile,重复加载、同步、计算、累加,直到所有 tile 都处理完

  • 通过 Tiling,内存吞吐量提高到 410 Gb/s,执行时间下降 ~43%,计算性能达到 ~6.6 TFLOPS

Warp Tiling

Thread / Register Tiling

多级 Tiling

  • 可以看到,前面提到的 Tiling 技术与 Register Blocking 的核心思想一致,Tiling 是把大问题拆成能放进高速存储的小块,将数据从 DRAM 移至 shared memory;而 Register Blocking 继续做同一件事,但更激进,将数据从 shared memory 移至 register 中

  • GPU 的并行层次可以总结如下

    层次
    硬件概念
    数量
    共享存储
    备注

    Grid

    整个 kernel

    1

    Global Memory

    Block

    Thread Block / Workgroup

    数百到数千

    Shared Memory

    前面提到的 Tiling

    Warp/Subgroup

    32 线程一组

    每 Block 若干

    寄存器(shuffle 共享)

    Register Blocking

    Thread

    单个线程

    每 Block 256-1024

    寄存器

    Register Blocking

  • 相关知识参考:

  • Tiling 代码示例

Memory bound 优化之片上存储优化

共享内存优化

  • 在 GPU 内存层级中

    • Register → 线程私有(最快)

    • Shared Mem → Block 共享(很快)

    • Global Mem → 显存(慢)

  • 因此 Kernel 设计中,需要尽量少访问 global memory,尽量把数据搬进 shared 或 register

Bank Conflict

  • shared memory 是在一个 SM 内给一个 block 的所有线程提供的高速 SRAM,它的设计目标是让一个 block 里的线程可以非常快地共享数据

  • 为了并行访问,GPU 会把 shared memory 划分为多个 bank(通常 32 个,恰好等于 warp 的线程数),即一个 shared memory 包含多条独立告诉通道

  • 当内核(kernel)为每个 block 分配共享内存时,地址会根据以下公式映射到 bank:

    bank_index=(address/word_size) mod (number_of_banks)\rm{bank\_index=(address/word\_size)\ mod\ (number\_of\_banks)}
    • NVIDIA GPU 的 bank 数固定为 32(与 warp 大小相同)

    • 默认 word 大小为 4 字节(float 或 int),也可以设置为 8 字节(double)

    • 变量在共享内存中声明的顺序决定了它们的 bank 分配

    • 例如,如果声明了一个 16 个 float 的数组(16 个 word,64 字节):

      • 第 1 个 word 映射到 bank 0,第 2 个到 bank 1,以此类推,第 16 个到 bank 15

      • 下一个声明的变量接着映射到 bank 16,依此类推

      • 动态分配的数组会在所有静态声明之后映射

  • 每个 bank 在一个 cycle 内只能服务于一次访问,理想情况下,一个 warp 有 32 个线程,shared memory 有 32 个 bank,那么 32 个线程同时读,1 个 cycle 即可完成

  • bank conflict 发生在多个线程访问同一个 bank 的不同 word 时,因为一个 bank 一次只能返回一个 word,访问会变为串行

  • 例如,如果在共享内存中分配了 256 字节(64 个 word):

    • 线程访问 bytes 0-3(word 0)和 bytes 128-131(word 32) → 都在 bank 0 → 发生冲突

    • 如果访问 bytes 132-135(word 33) → 在 bank 1 → 不发生冲突

    • 如果多个线程访问同一个地址 → 值会被广播,不产生冲突

  • 解决方法

    • Padding:例如 __shared__ float tile[32][33];,使地址分布改变,冲突消失

    • 改变访问顺序:例如行访问替换为列访问,或者转置 tile

    • Warp-level shuffle:绕过 shared memory

  • Bank Conflict 的示例代码

  • 在上述代码中,OFFSET 决定 bank 冲突的程度

    • OFFSET = 32 → 线程都访问同一个 bank → 最大冲突

    • OFFSET = 1 → 无冲突

  • 双缓冲

  • Async copy

Register 优化

Register Blocking

  • register 是每个 thread 私有,是在 SM 内速度最快的存储,访问 register 的开销基本是 1 cycle

  • 例如对于矩阵乘法,假如一个线程只处理一个输出元素,那么每算一步,都要从共享/全局内存中读取数据,数据复用率很低

  • Register Blocking 的核心思想是,每个 thread 一次计算多个输出元素,并把中间数据存储在 register 中反复使用

  • 这种方法用算力换掉了内存访问,从而减少了 shared memory 访问,增加计算密度,提高算术强度

  • 然而,Register Blocking 会导致每个 thread 用更多 register,SM 中能同时放的 thread 变少,Occupancy 下降

Register 压力控制

Comupute bound 优化

Thread Coarsening(线程粗化)

  • 通过分析 warp states(warp 停顿状态),可以发现大量周期消耗在一个看起来很“神秘”的状态上:smsp__pcsamp_warps_issue_stalled_mio_throttle

  • 通过查询 NVIDIA 的 Kernel Profiling Guide 后,发现该状态的含义是 warp 因为 MIO(memory input/output)指令队列已满而停顿,其中 MIO pipeline 负责共享内存访问、特殊数学指令、动态分支等

  • Tiling 虽然大幅减少了共享内存的访问,把大部分数据访问转移到共享内存中,但共享内存不是“免费”的,它仍然有端口数量限制,有指令发射和 pipeline 吞吐上限

  • 在 tiled GEMM 中,每个线程在内层循环中会频繁地从共享内存读取 A_tile 和 B_tile,这就导致 warp 内 32 个线程每一轮都会发射大量共享内存 load 指令

  • 最终结果是共享内存流水线被占满,warp 只能等待,所以现在的瓶颈已经从全局内存延迟变成了共享内存指令吞吐

  • 对于这种情况,NVIDIA 建议不要让很多线程反复做小而碎的共享内存访问,而是用更少的线程,每个线程一次处理更多数据,这正是 thread coarsening

  • 线程粗化也即让一个线程负责计算多个输出元素,而不是一个线程只算一个,也可以理解为把原来多个“细线程”合并成一个“粗线程”

  • 这样可以减少共享内存访问次数,提高吞吐量

Fused Kernels

  • CPU 和 GPU 可以异步执行

    • CPU(host)可以异步向 GPU 提交工作

    • 提交 kernel 之后,CPU 不需要等 GPU 执行完,可以继续做别的事情

    • GPU 自己按顺序执行队列里的 kernel

    • 这使得通信(memory copy)和计算(kernel 执行)可以重叠,是 GPU 高吞吐的基础之一

  • 虽然 CPU 和 GPU 可以异步,但如果程序结构是:

    • 启动 kernel A

    • 把结果写回 global memory

    • 再启动 kernel B

    • kernel B 再从 global memory 读回刚算完的数据

  • 那么问题是全局内存是慢的,kernel launch 是有开销的,而数据在计算单元和全局内存之间反复“折返跑”

image.png

  • 在这个反模式中,每一步计算写回全局内存,然后再读出来,即使这些中间结果马上就会被下一个 kernel 使用,这就导致了多次 DRAM 访问、多次 kernel 启动、pipeline 被打断

  • 因此引入了核融合(kernel fusion):把多个原本需要分别启动的 kernel,合并成一个更大的 kernel

image.png

  • 数据一旦进了 SM,就放在寄存器 / 共享内存,连续执行多个操作,最后只在必要时写回全局内存

  • 这样做可以减少 kernel launch 次数、减少全局内存读写,提高数据局部性,让 GPU 更“自主”地完成一整段计算

  • 融合内核尤其适合一系列逐点(point-wise)操作,其中 point-wise 操作指的是每个元素独立计算,不依赖其他元素,典型形式如加、乘、exp、sqrt、clamp、where 等运算,这些操作有一个共同点,即对每个 token / 元素独立,不存在跨元素的数据依赖

  • 在 Transformer 里,这种“逐点连续操作”非常多,最典型的就是 LayerNorm

    • LayerNorm 实际包含计算均值、计算方差、减均值、除标准差、乘 $\gamma$、加 $\beta$

    • 如果使用 naive 方式,每一步一个 kernel,每一步都写回共享内存,性能会非常差

    • 真正高性能实现一定是合并均值 + 方差 + 归一化 + affine,尽可能融合进少量 kernel

Pipeline Kernel

  • Producer-consumer pipeline

  • 多阶段流水

Latency-bound 优化

Thread Divergence(线程分歧)

  • Single Instruction, Multiple Data (SIMD) 执行模型

    • 一个 warp = 32 个线程,同一时刻只取一条指令,这条指令会同时在 warp 内所有线程上执行,每个线程处理不同数据

    • 这样设计的好处是指令取值、调度、分发这些控制逻辑只做一份,最大限度地减少了控制功能相关的硬件开销

    • 那么就可以把更多硬件用在提高算术吞吐率上,这是 GPU 高性能的根本原因之一

    • 同一个 warp 内的线程必须同时执行 相同的指令,也就是说,它们必须遵循同一条执行路径,那么当 warp 内出现条件判断(if 语句)时,线程如何执行?

  • if /else 中的线程分歧

    • 假设 warp 内部分线程判断 if 条件为 true,部分为 false,则 “true” 线程希望执行 if 块,“false” 线程希望执行 else 块

    • 这种情况称为线程分歧(thread divergence),通常由分支语句引起

    • CUDA 的处理方式是:

      • 先执行 if 块,即 if 条件为 false 的线程被“停用”,保持空闲

      • 再执行 else 块,之前停用的线程现在执行 else,而 if 条件为 true 的线程停用

    • 可以看到,if 和 else 并不是并行执行的,而是依次执行,在每个块中总有部分线程空闲,这会造成显著的性能损失

    • 因为 warp 是 GPU 的最小调度单位,只要 warp 内有分歧,那么整个 warp 都要付出代价

    • 这和 CPU 不同,CPU 的线程是独立调度的,而 GPU 的线程是“绑在一起”的

    • 示例如下

    • 相反,下面的条件不会导致线程分歧(warp 内不会出现分支):

  • for 循环中的线程分歧

    • 如果 warp 内不同线程的循环次数不一致,也会导致线程分歧

    • 假设某些线程循环 5 次,其他线程循环 10 次,前 5 次迭代可以并行执行,而从第 6 次迭代开始,已完成循环的线程空闲,其他线程继续执行

    • 当循环次数高度依赖线程时,循环可能完全串行化,性能损失严重;嵌套 if-else 的循环会进一步加剧分歧

    • 在循环后调用 __syncthreads() 会放大问题,因为整个线程块必须同步,而不仅仅是 warp

  • 线程分歧与 __syncthreads() 结合可能导致死锁:

    • 示例如下

    • 问题分析:偶数线程执行 if 块,到达 __syncthreads() 等待所有线程,而奇数线程执行 else 块,到达自己的 __syncthreads(),从而偶数和奇数线程等待不同的同步点 → 死锁

    • 解决方法:将 __syncthreads() 放在 if-else 语句之后

    • 这样,所有线程都会执行同一个同步点,避免死锁

  • 相关的处理方法

    • 优化代码,减少或消除分支

    • 如果必须分支,确保 warp 内线程条件结果相同,如按数据块对齐线程,让相邻线程处理“相似数据”

    • 使用 predication(谓词化)

    • 谨慎处理边界条件:在 block 边缘的 warp,很容易出现部分线程越界、部分不越界,从而引入分歧(例如上文中矩阵乘法运算的实现)

Launch Overhead

  • 解决方案:

    • fusion

    • persistent kernel

    • CUDA graph

Persistent Kernel

并行编程原语

什么是原语

  • 在并行编程中,有若干个相互独立的节点,它们可以是 CPU 核心、GPU,或者计算节点

  • 每个节点先各自执行一些计算,然后将计算结果或其中的一部分传递给其他节点,用于下一步计算($t+1$)

image.png

  • 有时,需要把某一个节点的结果发送给所有其他节点

  • 有时,则需要把每个节点的中间结果汇总求和,以得到整体结果并进行汇报

  • 通常,会有一个具有更高地位的节点在这些操作中扮演核心角色,这里称之为 root,它可能是某些操作的源或目标

Broadcast

image.png

  • Broadcast(广播):某个节点(比如节点 1)上已经有了一份数据,我们希望把它共享给所有其他节点,让它们都能基于这份数据进行计

  • 集合通信操作在 PyTorch 中是原生支持的,因此可以很容易地写一个小示例来展示广播是如何工作的

    • 首先,需要使用 dist.init_process_group 来初始化一个进程组,这一步会设置通信后端

    • 它会确定一共有多少个工作进程(也叫节点),并为每一个进程分配一个 rank(可以通过 dist.get_rank 获取)

    • 最后,这一步还会在各个工作进程之间建立通信连接

  • 为了演示 dist.broadcast 操作,在 rank=0 的进程上创建一个包含非零值的张量,而在其他进程上创建全为零的张量

  • 然后,通过 dist.broadcast(tensor, src=0) 将 rank=0 上的张量分发到所有其他 rank 上

  • 可以使用下面的命令来运行这个脚本:

  • 以上命令需要三块 GPU(或者相应地修改 nproc_per_node),运行后输出如下:

  • 需要注意的是,不同 rank 的打印信息可能会乱序出现,因为无法控制各个进程中 print 语句的执行先后顺序

Reduce & AllReduce

image.png

  • Reduce(归约)模式是分布式数据处理中最基础、也最常见的模式之一

  • 其核心思想是:希望通过某个函数 $f()$,将分布在各个节点上的数据组合起来,例如执行求和或求平均等操作

  • 在 Reduce 范式中,最终结果只会被发送到 root 节点;而在 AllReduce 的情况下,最终结果会被广播到所有节点

  • 通常的实现方式是:每个节点只负责一部分计算,节点之间按照环形(ring)或树形(tree)结构组织通信

  • 环形结构:假设需要计算每个节点上的数值之和,并且节点是按环形结构连接的。第一个节点把自己的数值发送给相邻节点,后者在接收到该数值后,加上自己的数值,再转发给下一个节点。经过一整圈之后,第一个节点就会收到所有节点数值的总和

  • 下面是一个运行 Reduce 操作、对张量进行求和的简单示例,通过 op=dist.ReduceOp.SUM 来指定所使用的归约操作

  • 需要注意的是,在 Reduce 操作中,只有 dst 节点上的张量会被更新:

  • 类似地,可以执行 AllReduce 操作:

  • 在这种情况下,归约后的结果会在所有节点上都可用:

Gather & AllGather

image.png

  • Gather 和 AllGather 与 Broadcast 操作非常相似,它们都允许在节点之间分发数据,并且不对数据本身做任何修改

  • 与 Broadcast 的主要区别在于:Broadcast 是将一个节点上的同一份数据发送给所有其他节点;而在 Gather / AllGather 中,每个节点都持有一份不同的数据块,要做的是把这些数据块收集起来

    • Gather 把所有节点的数据收集到一个指定节点上

    • AllGather 把所有节点的数据收集到每一个节点上

  • 上图中的虚线表示有些数据实际上并不会发生移动,因为这些数据本来就已经位于对应的节点上

  • 在 Gather 操作中,需要事先准备一个容器(gather_list),用来存放最终收集到的张量

  • 可以看到,gather_list 中确实包含了来自所有 rank 的张量:

  • 对于 AllGather 示例,唯一需要做的改变是,每一个节点都要准备一个用于存放结果的占位容器:

  • 此时可以看到,每一个节点上都拥有了所有节点的数据:

  • Gather 的“逆操作”:所有数据最初集中在一个节点上,并且需要将它们切分并分发到各个节点上,中间可能还会伴随一些额外的处理,这时可以使用 Scatter 操作;如果在分发之前还需要先对数据做一次归约操作,就可以使用 ReduceScatter 模式

Scatter & ReduceScatter

image.png

  • 顾名思义,Scatter(散射)操作的目标是把集中在某一个节点上的数据分发到所有节点上,它通过将数据切分成若干片段,并把其中一片分配给每个节点来实现这一点

  • 因此,Scatter 不同于 Broadcast 操作:广播是把完整的数据副本发送给每一个节点,而不对数据进行切分;Scatter 可以看作是 Gather 操作在逻辑上的逆过程

  • ReduceScatter 和 AllReduce 类似,都会先对来自所有节点的数据应用某种归约操作;但不同之处在于:在 AllReduce 中,每个节点都会收到完整的输出张量,而在 ReduceScatter 中,每个节点只会收到最终输出张量的一部分切片

  • Scatter 操作在代码层面上可以看作是 Gather 的反向过程:不再是准备一个张量列表作为接收目标,而是把源数据准备成一个张量列表,表示希望分发给各个节点的数据。同时,还需要指定源节点 src:

  • 执行结果表明,原本为空的张量被填充成了 scatter_list 中对应的内容:

  • 下面展示 ReduceScatter 的工作机制,在每个节点上,都会创建一个由二维向量组成的列表,这些向量由节点 rank 相关的倍率以及幂次关系生成:

  • 从打印结果中可以清楚地看到刚刚构造的数据模式,同时也能直观地看到 ReduceScatter 的效果:第一个 rank 收到的是所有节点上第一个张量的求和结果,第二个 rank 收到的是所有节点上第二个张量的求和结果,以此类推:

Ring AllReduce

  • Ring AllReduce 是一种针对可扩展性优化的 AllReduce 实现方式

  • 与所有设备彼此直接通信(这很容易造成通信瓶颈)不同,Ring AllReduce 可以拆解为两个关键步骤:ReduceScatter 和 AllGather,其工作流程如下:

    • ReduceScatter

      • 每个设备将自己的数据(例如梯度)拆分成 $N$ 个块(其中 $N$ 是 GPU 的数量),并将其中一个数据块发送给相邻的设备

      • 同时,每个设备也会从另一个相邻设备那里接收一个数据块

      • 当设备接收到一个数据块时,会将自己对应位置的数据块与接收到的数据块做归约操作(例如求和)

      • 这一过程沿着环不断进行,直到每个设备都持有一个已经完全归约的块,该块表示该分片在所有设备上的梯度之和

    • AllGather

      • 接下来,每个设备需要从其他设备那里收集这些已经完全归约好的数据块

      • 每个设备会将自己持有的归约后数据块发送给相邻设备,并从另一个相邻设备接收数据块

      • 设备会不断转发接收到的数据块,直到每个设备都拥有所有已经归约完成的数据块,从而使每个设备都获得完整的、加总后的梯度

  • 可以通过下面的动画来直观理解这一过程。假设有 5 个 GPU,每个 GPU 上都有一个长度为 5 的张量。第一个动画展示的是 ReduceScatter 阶段,在该阶段结束时,每个 GPU 都得到了某一个特定数据块(橙色矩形)对应的归约结果

image.png

  • 下一个动画展示的是 AllGather 阶段,在该阶段结束时,每个 GPU 都获得了 AllReduce 操作的完整结果:

image.png

  • 在 ReduceScatter 和 AllGather 这两个阶段中,每个 GPU 都会发送和接收 $N-1$ 次数据

  • 每一次传输中,每个 GPU 会发送 $\frac{K}{N}$ 个数值,其中 $K$ 是在所有 GPU 之间参与求和的数值总数

  • 因此,每个 GPU 的总通信数据量为

    2×(N1)×KN2 \times (N - 1) \times \frac{K}{N}
  • 当 $N$(GPU 的数量)很大时,每个 GPU 的通信数据量近似为

    2×K2 \times K
    • 其中 $K$ 是参数(或梯度)总数

  • 关于 AllReduce,有两个关键点需要记住:

    • 当 $N$(GPU 数量)较大时,AllReduce 的通信成本大约是 $2 \times K$

    • 一个 AllReduce 操作可以分解为一次 ReduceScatter 加一次 AllGather,而这两个操作各自的通信成本大约是 AllReduce 的一半,即约为 $K$

    • 可以看到,这种实现方式即使在节点之间带宽有限的情况下,也能够非常高效地利用通信资源

Barrier

image.png

  • Barrier(屏障)是一种非常简单的同步操作,用于让所有节点在某个执行点上对齐

  • 只有当所有节点都到达这个屏障之后,Barrier 才会被解除,随后所有节点才能继续进行后续计算

  • 可以通过在不同节点上设置不同的休眠时间,来模拟节点处理速度不一致的情况,并观察它们通过 Barrier 所需的时间

  • 从输出中可以看到,尽管 rank 0 完全没有休眠,它依然需要等待最慢的节点,最终也花了大约 2 秒才能通过 Barrier:

  • 这说明 Barrier 会强制所有节点以“最慢节点”的速度前进

  • 因此,在使用这类全局同步操作时需要格外谨慎,因为它违背了并行、独立执行的初衷,可能会拖慢整体处理速度

  • 在很多场景下,让更快的节点先一步开始处理下一项任务其实是完全可以接受的——因为在后续迭代中,这些节点可能反而会变慢,从而在整个计算过程中自然地平衡掉这些延迟

NCCL

  • 有多种实现了集合通信(collective communication)的库,并且都得到了 PyTorch 的支持

    • 经典的 MPI(Message Passing Interface,消息传递接口)

    • Meta 开发的 Gloo

    • NCCL(NVIDIA Collective Communications Library,NVIDIA 集体通信库)

  • 它们在提供的集合通信模式上功能相似,但针对不同的硬件配置进行了优化:

    • NCCL 专门为 GPU-GPU 之间的高效通信而设计

    • MPI 和 Gloo 则更适合 CPU-CPU 或 CPU-GPU 之间的通信

  • PyTorch 官方提供了一份很好的指南arrow-up-right,帮助选择合适的通信后端

    • GPU 训练:使用 NCCL

    • CPU 训练:使用 Gloo

Last updated

Was this helpful?