A Framework for Fine-Grained Synchronization of Dependent GPU Kernels

作者/机构: Abhinav Jangda (Microsoft Research), Saeed Maleki (Microsoft Research), Maryam Mehri Dehnavi (University of Toronto), Madan Musuvathi (Microsoft Research), Olli Saarikivi (Microsoft Research)

A1 主要贡献

本文旨在解决在GPU上执行机器学习（ML）模型时，由于依赖性计算（如GeMM、Conv2D）导致的GPU资源利用率不足的问题。

核心问题: 传统的GPU计算将任务分解为多个独立的块（tiles），由线程块（thread blocks）在流式多处理器（SMs）上分多波（waves）执行。当线程块数量不是SM数量的整数倍时，最后一波执行将导致部分SM闲置，造成GPU资源浪费。传统的流同步（stream synchronization）机制要求一个核心（kernel）的所有线程块执行完毕后，依赖于它的下一个核心才能开始执行，这加剧了最后一波的资源浪费问题。如图1(b)所示，两个依赖的GeMM操作，每一波的第二个部分都只有一半的SM在工作。Table I的数据也显示，在MegatronLM GPT-3的推理中，两个依赖的GeMM在不同批量大小下的GPU利用率仅为60-80%。

研究目标: 提出一个框架，通过对依赖核心进行细粒度的同步来提高GPU利用率。核心思想是同步计算块（tiles）而非整个核心，从而允许依赖核心中的独立计算块并发执行，填补最后一波的空闲资源，如图1(c)所示。

创新点:
1. cuSync框架: 提出了一个名为cuSync的框架，用于根据用户定义的细粒度同步策略来同步依赖的GPU核心。cuSync通过同步tiles而非kernels，实现了依赖核心中独立tiles的并发执行。其关键机制包括：
* 确保生产者核心的线程块在消费者核心之前执行。
* 允许以特定顺序处理生产者和消费者的tiles，以最小化消费者的同步等待时间。
* 使用信号量（semaphores）和内存栅栏（memory fences）维护生产者和消费者计算tiles之间的依赖关系。
2. cuSyncGen编译器: 开发了一个编译器cuSyncGen，它接收一个描述GPU核心间依赖关系的领域特定语言（DSL）规范，并据此生成多样的细粒度同步策略。
3. 广泛适用性和显著性能提升: 将cuSync应用于GeMM、2-D卷积、Dropout等多种GPU计算，并在四个流行的ML模型上验证了其效果。实验表明，使用cuSync可将以下模型的推理时间显著降低：
* MegatronLM GPT-3：最多降低15%
* LLaMA：最多降低14%
* ResNet-38：最多降低22%
* VGG-19：最多降低16%

本文提出的方法克服了现有技术Stream-K的三个主要缺陷：避免了因划分tile导致的额外全局内存访问、无需为完整波和部分波调用不同的核心，并且能够直接扩展到GeMM之外的其他基于tile的计算。

A3 背景知识

本节提供关于NVIDIA GPU和ML模型的背景知识。

A. NVIDIA图形处理单元和CUDA

CUDA核心与线程组织。在NVIDIA GPU上执行的并行计算被称为CUDA核心（kernel）。一个CUDA核心执行多个并发的线程，这些线程被组织在一个三维网格（grid）中，并被分组为大小相等的线程块（thread blocks）。在CUDA中，dim3结构体用于表示线程和线程块在x、y、z三个维度上的三维网格大小和标识符。一个NVIDIA GPU包含多个流式多处理器（Streaming Multiprocessors, SMs），每个SM执行一个或多个线程块。每个SM能容纳的线程块数量，即占用率（occupancy），取决于CUDA核心的寄存器和共享内存使用量、线程块大小以及线程块的数量。

线程块的波次执行。线程块在所有SM上分波次（waves）执行，波次数为 $\lceil \frac{\text{网格中的线程块总数}}{\text{占用率} \times \text{SM数量}} \rceil$。初始的完整波次会执行占用率 × SM数量个线程块，而最后的非完整波次则执行剩余的线程块。NVIDIA并未公开CUDA和GPU调度线程块到SM的具体机制。

流同步。一个CUDA流（stream）是一系列按发布顺序执行的CUDA操作。当两个有依赖关系的核函数在同一个流上调用时，消费者核函数（consumer-kernel）必须等待生产者核函数（producer-kernel）的所有线程块都执行完毕后才能启动。我们将这种同步方式称为流同步（stream synchronization）。我们可以将独立的CUDA核函数在不同的流上调用，以实现并发执行。流还具有关联的优先级值，高优先级流上的操作会先于低优先级流上的操作被发布。

B. 大型ML模型中的计算

现代ML模型包含大量可并行化的计算，如通用矩阵乘法（GeMM）、二维卷积（Conv2D）、Dropout和Softmax。本文考虑了四种广泛使用的机器学习模型：MegatronLM GPT-3 145B 【索引12，Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2020】，LLaMA 65.2B 【索引15，LLaMA: Open and Efficient Foundation Language Models, 2023】，ResNet-38 【索引6，Deep Residual Learning for Image Recognition, 2015】，和VGG-19 【索引13，Very Deep Convolutional Networks for Large-Scale Image Recognition, 2015】。下面简要解释这些模型中涉及的计算。

1) Transformer模型: Transformer是一种用于自然语言任务的深度学习架构，是两种广泛使用的模型的基础：MegatronLM GPT-3 【索引12，Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2020】 和 LLaMA 【索引15，LLaMA: Open and Efficient Foundation Language Models, 2023】。

1 //X: [B,S,H]; W1: [H,4H/8]; W2: [4H/8,H]
2 //1st GeMM fused with GeLU XW1: [B,S,4*H/8]
3 XW1 = GeLU(X × W1)
4 //2nd GeMM XW2: [B,S,H]
5 XW12 = XW1 × W2

(a) 多层感知机（MLP）包含两个权重矩阵：形状为[H, 4H/8]的W1和形状为[4H/8, H]的W2。

1 //X: [B,S,H]; QKV: [H,3H/8]; W2: [H/8,H]
2 //1st GeMM XQKV: [B,S,3H/8]
3 XQKV = X × QKV
4 //XQ: [B,S,H/8]; XK: [B,S,H/8]; XV: [B,S,H/8]
5 XQ = XQKV[:,:,0:H/8] //1st matrix slice
6 XV = XQKV[:,:,H/8:2*H/8]//2nd matrix slice
7 XK = XQKV[:,:,2*H/8:] //3rd matrix slice
8 //Cached Attention Mechanism
9 //CachedK: [H/8,S’,B ]
10 //CachedV: [B ,S’,H/8]
11 P = XQ × Concat(CachedK, XK.T)
12 R = Softmax(Dropout(P))
13 T = R × Concat(CachedV, XV)
14 CachedV[:S’+S:] = XV
15 CachedK[:S’+S:] = XK.T
16 //2nd GeMM XW2: [B,S,H]
17 XW12 = R × W2

(b) 注意力机制包含两个权重矩阵：形状为[H, 3H/8]的QKV和形状为[H/8, H]的W2。注意力机制会缓存生成的键和值，以避免在推理过程中为每个token重新计算所有先前的token。

Transformer推理过程。对Transformer模型的推理请求包括一个提示（prompt），分两个阶段处理：(i) 提示处理（prompt processing），处理输入的提示；(ii) token生成（token generation），增量地生成一系列代表输出响应文本的token。模型可以将B个请求批处理为单个推理任务。序列长度S表示在提示处理阶段每个请求处理的token数，或在token生成阶段每个请求生成的token数。因此，在提示处理期间B≥1, S>1；在token生成期间B≥1, S=1。一个Transformer由多个多层感知机（MLP）和注意力（Attention）块组成，其设计因模型而异。

GPT-3。在GPT-3中，MLP和Attention都接收一个输入矩阵，使用其两个权重矩阵进行运算，并输出一个矩阵。通过模型并行，这些权重矩阵被分配到所有GPU上【索引12，Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2020】。图2展示了在8个GPU上进行模型并行的GPT-3计算。MLP和Attention首先对输入进行线性变换（即输入与权重矩阵的GeMM），然后执行GeLU和注意力机制等操作，最后将操作结果应用于第二次线性变换。现有的MLP实现将GeLU激活函数与第一个GeMM融合（图2a第4行）。先进的注意力实现【索引4，FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022, Advances in Neural Information Processing Systems】会缓存已处理和生成的token到KV缓存中，提示处理后缓存的token数S'被设为S，生成token时S'递增。这些实现也将注意力机制融合到单个CUDA核心中（图2b第11-13行）。

1 //X: [B, S, H]; W1: [H, H/3];
2 //V: [H, H/3]; W2: [H/3, H]
3 //1st GeMM XW1: [B, S, H/3]
4 XW1 = GeLU(X × W1)
5 //2nd GeMM XV: [B, S, H/3]
6 XV = X × V
7 //SwiGLU fused with 3rd GeMM XW2: [B, S, H]
8 SwiGLU = Swish(XW1) · XV
9 XW12 = SwiGLU × W2

图3：LLaMA的多层感知机（MLP）架构。

LLaMA。LLaMA使用大小为8192的隐藏维度。LLaMA的MLP包含三个GeMM和SwiGLU【索引11，GLU Variants Improve Transformer, 2020】激活函数，如图3所示。先进的实现将前两个GeMM合并为单个GeMM，并将SwiGLU激活函数与第三个GeMM融合。此外，LLaMA使用与GPT-3相同的注意力架构。

2) 计算机视觉模型: ResNet-38【索引6，Deep Residual Learning for Image Recognition, 2015】和VGG-19【索引13，Very Deep Convolutional Networks for Large-Scale Image Recognition, 2015】是两种先进的计算机视觉模型，每层都执行多个Conv2D操作。表II显示了每个卷积层的详细信息。

A2 方法细节

我们对依赖的CUDA核心进行的细粒度同步包含四种新颖的机制。这些机制 (i) 确保依赖核心的同时分配（第III-A节），(ii) 在消费者核心之前执行生产者核心的线程块（第III-B节），(iii) 控制每个核心中tile的处理顺序以最小化同步等待时间（第III-C节），以及 (iv) 对生产者和消费者核心中仅有依赖关系的tiles执行细粒度同步（第III-D节）。我们已将这些机制实现为一个仅需头文件的独立CUDA库，名为cuSync。图4a通过一个使用cuSync同步MLP中两个依赖GeMM核心的例子来解释这些机制。gemm函数是标准的GeMM GPU核心（我们实验中使用NVIDIA CUTLASS【索引1，NVIDIA cuTLASS: CUDA Templates for Linear Algebra Subroutines, https://github.com/NVIDIA/cutlass】），并增加了调用cuSync的代码（以下划线标出）。cuSync为每个核心关联一 个CuStage对象，该对象提供核心间的同步功能。MLP函数创建这些stage对象，声明它们之间的依赖关系，并调用核心。

A. 调用依赖的核心

在不同流上调用以消除流同步。细粒度同步的首要要求是消除核心之间的流同步。cuSync通过在不同的CUDA流上调用所有核心来实现这一点。示例中为两个GeMM核心创建了生产者（prod）和消费者（cons）阶段（图4a第18-20行）。然后，通过指定生产者的输出是消费者的输入来声明两个阶段之间的依赖关系（第23行）。最后，示例在与各自阶段相关联的不同流上调用两个核心（第26和30行）。第III-D节描述了cuSync如何强制执行此依赖关系。

B. 阶段处理顺序

使用等待核心确保执行顺序。细粒度同步的第二个要求是在消费者核心之前执行生产者核心的所有完整波次。然而，CUDA运行时缺乏在属于不同流的核心之间强制执行此顺序的机制。因此，消费者核心有可能在生产者核心之前被调度到GPU上。这可能导致性能不佳，因为消费者核心的线程块会占用SM而不做任何有效工作。在最坏的情况下，如果生产者核心没有可用的SM，这可能导致死锁。

等待核心机制。cuSync通过其等待核心（wait-kernel）机制来确保这一要求，强制规定核心的调度顺序。等待核心由消费者阶段在消费者流上、在消费者核心之前调用（图4a第28行）。等待核心包含一个单线程，它在一个while循环中忙等待一个全局内存信号量。当生产者核心调用stage.start()函数时（第4行），该函数由第一个线程块的第一个线程设置信号量，这反过来会使等待核心退出。等待核心退出后，CUDA运行时可以调用消费者核心。因此，cuSync确保在生产者核心的至少一个线程块被调度之前，消费者核心的任何线程块都不会被调度。该机制假设CUDA按照核心被调用的顺序调度线程块，我们发现在基于Volta和Ampere架构的NVIDIA GPU上运行的CUDA 11和12最新版本遵循此调度方式。

C. 自定义Tile顺序

通过自定义调度顺序最小化等待时间。高效同步的第三个要求是最小化消费者核心的等待时间。然而，CUDA运行时可以以任意顺序将线程块调度到SM上，这可能导致不可预测的等待时间。理想情况下，消费者核心的线程块应该按照生产者核心生成tile的顺序进行调度。

实现自定义顺序。cuSync使得生产者和消费者核心的线程块能够以自定义的调度顺序执行，而与CUDA运行时的调度方式无关。在我们的示例中，每个线程块调用stage.tile()（第4行）以获取它需要计算的tile。参数RowMajor（第18-20行）确保两个核心都以行主序（row major order）生成tile，即首先是x维度的所有线程块，然后是y维度，最后是z维度。图4b将RowMajor顺序定义为一个函数（第29行）。tile顺序函数接收一个3D网格中的tile索引，并为该tile返回一个唯一的1D索引。在内部，cuSync维护一个数组，将线性tile索引映射到3D索引。对于每个线程块，cuSync递增一个原子全局计数器，并返回数组中对应于前一个计数器值的3D索引。总之，cuSync允许用户轻松地试验各种调度顺序以获得最佳性能。

D. 同步依赖的Tiles

使用wait和post函数进行同步。细粒度同步的最终要求是使用同步机制来确保生产者和消费者核心的tile之间的依赖关系得到维护。cuSync提供了两个函数wait和post来强制执行此依赖关系。例如，在我们的示例中，wait函数在加载A和B的tile之前被调用两次（第6和8行），而post函数在计算完tile后为生产者核心调用一次（第12行）。然而，消费者核心只需要等待生产者核心的输出，即消费者核心的输入A。

依赖声明与同步策略。这个依赖关系在第23行中指定。因此，在加载A的tile之前的wait会等待生产者核心相应的post，而在加载B的tile之前的wait则成为一个空操作（no-op）。由于生产者核心没有依赖关系，所以对于生产者核心来说，两个wait都是空操作。cuSync提供了一种基于任意同步策略（简称policy）同步生产者和消费者tile的机制。cuSync使用一个全局内存信号量数组进行同步，其中每个生产者tile只与一个信号量相关联，一个信号量的值代表其生产者tile的状态。因此，一个策略是一个或多个生产者tile到一个信号量的映射。例如，最细粒度的同步策略，我们称之为TileSync，它等待每个生产者tile，并被定义为生产者tile到信号量的一对一映射。一个策略需要实现两个方法：(i) sem，返回给定tile的信号量，以及 (ii) value，返回当tile准备好时期望的信号量值。下面我们详细描述CuStage中实现我们同步机制所需的三个方法（图4b第2-9行）。
* init: init方法根据给定的策略在全局内存中分配和初始化信号量数组。
* post: post方法调用__syncthreads和一个内存屏障，以确保线程块的所有线程都已计算完tile，并且所有全局内存写入对其他核心可见（第5-7行）。最后，该方法使用策略获取tile的信号量并递增该信号量（第9行）。
* wait: wait方法使用策略获取给定tile的信号量，然后仅使用线程块的第一个线程在一个while循环中等待信号量的值（第13行）。当第一个线程在等待时，线程块的所有其他线程都在__syncthreads上被阻塞（第14行）。当信号量变为期望值时，线程块的所有线程从__syncthreads继续执行。

E. 同步策略

多种同步策略的实现。cuSync允许轻松实现多样的同步策略。如前所述，每个策略需要实现sem和value方法。下面我们讨论两种适用于我们工作负载中所有核心的通用策略。

TileSync。TileSync是同步每个生产者tile的最细粒度策略（图4b第16-20行）。为了最小化消费者核心的等待时间，两个核心都以行主序计算它们的tile。sem方法为每个tile返回一个不同的信号量（第17行），value方法返回1以表示tile已计算完成（第20行）。例如，在图4a中，要计算一个tile $E_{xy}$，TileSync策略需要先等待$C_{x0}$，然后等待$C_{x1}$。

RowSync。RowSync在生产者核心的每一行上进行同步，需要的同步次数比TileSync少（图4b第22-27行）。例如，对于图4a的两个GeMM，TileSync总共需要12次同步，而RowSync通过为计算C的同一行的所有tile共享同一个信号量，只需要6次同步。因此，sem方法返回给定tile的行号，value方法返回行准备好时的值，即一行中的tile数量（第23-26行）。为了最小化等待时间，两个核心都以行主序调度它们的tile。RowSync也可用于同步Conv2D核心。第五节显示，对于大型GeMM和Conv2D，大量的同步是一个瓶颈。

IV. 策略和Tile顺序的自动调优

手动调优的挑战。获得最佳性能的过程涉及试验多种同步策略和tile处理顺序。最佳的策略和tile顺序取决于计算、数据大小和GPU架构。然而，手动执行这个过程既繁琐又容易出错。

cuSyncGen工具。因此，我们开发了cuSyncGen，这是一个工具，它接收用户指定的依赖关系，并为cuSync生成最优的tile处理顺序和多种同步策略的CUDA代码。cuSync目前需要用户手动修改GPU核心，以使用生成的策略和tile处理顺序来实例化CuStage，类似于MLP的例子（图4a）。cuSync的模块化设计使用户可以轻松地插入多样的同步策略和tile处理顺序。

A. 工作流程

1 Dim x, y;
2 //Max value of all dimensions of both GeMMs
3 Grid g1(x, y, 4*H/8*TileN, B*S/TileM);
4 Grid g2(x, y, H/TileN, B*S/TileM);
5 //Tile is produced by each thread block
6 Tile prod(x, y), cons(x, y);
7 //All col tiles for a row from 0 to H/2*TileN
8 ForAll prodCols(prod, x, Range(g1.x));
9 //Tile of 2nd GeMM depends on all
10 //col tiles of 1st GeMM
11 Dep dep({g2, cons}, {g1, prodCols});

(a) GPT-3的MLP

1 Dim x, y;
2 //First GeMM Grid
3 Grid g1(x, y, 3*H/8*TileN, B*S/TileM);
4 //P, R, and T Grid
5 Grid gP(x, y, B*(S+S′), B*(S+S′));
6 Grid gR(x, y, B*(S+S′), 1);
7 Grid gT(x, y, H/8*TileN, B*(S+S′)/TileM);
8 //Second GeMM Grid
9 Grid g2(x, y, H/8*TileN, B/TileM);
10 //P to 1st GeMM
11 //Strided Tile Dependencies: stride=H/8*TileN
12 Dep dep1P({gP, Tile(x,y)},
13 {g1, Tile(x,y), Tile(x+H/8*TileN,y)});
14 //R to P
15 Dep depPR({gR, Tile(x,y)},
16 {gP, ForAll(Tile(x,y), y, Range(gP.y))});
17 //T to R and 1st GeMM
18 Dep depTR1({gT, Tile(x,y)},
19 {gR, Tile(x, y)}, {g1, Tile(x+2*H/8*TileN,y)});
20 //2nd GeMM to T
21 dep23({g3, Tile(x,y)}, {gT, Tile(x/TileM,y)});

(b) 注意力机制

1 Dim x, y;
2 //First GeMM Grid
3 Grid g1(x, y, C/TileM, B*P*Q);
4 //Second GeMM Grid
5 Grid g2(x, y, C/TileM, B*P*Q/TileN);
6 //2nd Conv2D to 1st Conv2D
7 Dep dep({g2, Tile(x,y)}, {g1, Tile(x/R*S,y)});

(c) 两个Conv2D

cuSyncGen的工作流程。cuSyncGen的工作流程如下：
1. 用户描述核心tile之间的依赖链以及所有核心的网格值。
2. cuSyncGen根据网格值检查生产者和消费者tile的边界。
3. cuSyncGen生成能够最小化等待时间的tile处理顺序的CUDA代码。
4. cuSyncGen为多种策略生成CUDA代码。
5. 用户修改工作负载以支持cuSync，并将生成的CUDA代码插入到cuSync中。

本节的其余部分将描述这些步骤。

描述依赖关系。用户使用嵌入在C++中的DSL来描述核心tile之间的依赖关系。图5a展示了MLP中两个GeMM之间在DSL中描述的依赖关系。首先，DSL代码必须为每个核心的网格维度定义其最大值。该示例为两个网格定义了x和y维度（第1-4行）。指定网格的精确值可以生成高效的代码并进行边界检查以确保正确性。然后，DSL代码通过为网格的每个维度指定一个仿射函数来构建生产者和消费者的tile。该示例为网格中的每个线程块创建了一个生产者和消费者tile，并使用ForAll创建了一系列列tile（第6-8行）。最后，代码指定了一个消费者tile与一个或多个生产者tile之间的依赖关系（第11行）。

生成Tile处理顺序。cuSyncGen为每个核心生成一个tile处理顺序，以最小化等待时间。为了讨论这个过程，考虑一个依赖关系，其中一个消费者tile $C(x, y)$ 依赖于N个生产者tile，$\{P(x, a_0y + b_0), P(x, a_1y + b_1), \dots, P(x, a_{N-1}y + b_{N-1})\}$。当消费者核心以生产者核心产生tile的相同顺序来消费tile时，我们实现了最小的等待时间。因此，我们使用以下代码为每个消费者tile连续调度所有N个生产者tile：

1 int prodOrder(dim2 tile, dim2 grid) {
2   int linear = bid.y*gDim.x + bid.x, y = 0;
3   if (tile.y%a0 <= b0) y = 0;
4   //Similarly for tiles till N-2
5   else if (tile.y%aN-1 <= bN-1) y = N-1;
6   return linear/N+y;
7 }

此代码获取tile的1-D线性索引，找到N个tile组内的tile索引，并返回新的线性索引。我们还将消费者核心设置为遵循tile的行主序。我们的MLP示例使用行主序，即所有$H/TileN$个连续生产者tile的组被连续调度。通过将依赖关系从最后一个消费者核心扩展到第一个生产者核心，然后为每个核心生成代码，可以很容易地将此方法扩展到依赖核心链。

生成策略。cuSyncGen为每个依赖关系生成多个同步策略。对于接下来的讨论，考虑一个依赖关系，其中一个消费者tile，$C(x, y)$，依赖于N个生产者tile，$\{P(x, a_0y + b_0), P(x, a_1y + b_1), \dots, P(x, a_{N-1}y + b_{N-1})\}$。cuSyncGen为每个维度的依赖关系生成两种策略：(i) 将每个tile映射到一个不同的信号量，或 (ii) 将所有N个tile映射到同一个信号量。对于所考虑的依赖关系和$M \in \{1, N\}$的值，生成的代码是：

1 int sem(dim2 tile, dim2 grid) {
2   int y = 0;
3   if (tile.y%a0 <= b0)
4     y = (tile.y-b0)/a0;
5   //Similarly for tiles till M-2
6   else if (tile.y%aM-1 <= bM-1)
7     y = (tile.y-bM-1)/aM-1;
8   else y = tile.y;
9   return y*grid.x + tile.x;
10 }
11 int value(dim2 tile, dim2 grid) {
12   return M;
13 }

在考虑了最内层维度的两种情况后，该阶段移至外层维度，并遵循相同的方法。在我们的MLP示例中，cuSyncGen生成了两种策略：(i) TileSync，将每个tile映射到一个不同的信号量，以及 (ii) RowSync，将同一行的所有列tile映射到同一个信号量。

运行生成的代码。我们要求用户修改工作负载以支持运行cuSync，方法是在每次加载tile前添加wait调用，在计算tile后添加post调用。例如，在MLP的情况下，我们要求用户进行图4a中的更改。cuSync的模块化使得插入多种策略和tile处理顺序成为可能。因此，用户可以执行所有生成的策略，并获得执行时间最短的策略。

B. ML模型中的计算依赖关系

指定Attention和Conv2D的依赖。我们现在展示如何在cuSyncGen中为Attention和Conv2D情况指定依赖关系。

Attention。Attention在其三个核心之间包含两个依赖关系（图5b）。在第一个依赖关系中，点积的一个元素依赖于第一个GeMM输出中同一行的三个元素，步长为$H/8$（第13行）。除了TileSync和RowSync，cuSyncGen还为此依赖关系生成了一个我们称为StridedSync的策略，该策略将第一个GeMM的所有三个生产者tile映射到同一个信号量。因此，StridedSync会等待第一个GeMM的所有三个tile都计算完毕后，才继续进行tile的点积。此外，cuSyncGen生成了连续调度这三个tile的tile顺序。对于其他依赖关系，cuSyncGen生成TileSync和RowSync，同时以RowMajor顺序处理tile。

Conv2D。使用隐式GeMM算法的Conv2D将一个大小为[P, Q, C]的B个输入图像与一个大小为[R, S]的核心矩阵的卷积，转换为矩阵[B*P*Q, C*R*S]与[C*R*S, C]的GeMM。图5c显示了使用隐式GeMM算法的两个Conv2D之间的依赖关系。因此，该依赖关系描述了第二个隐式GeMM的每个tile都依赖于第一个隐式GeMM输出的所有列tile（第7行）。cuSyncGen为此依赖关系生成了两种策略：(i) RowSync用于同步每一行，和 (ii) Conv2DTileSync策略用于同步每个tile。此外，cuSyncGen为两个Conv2D生成了行主序。

C. 优化

cuSyncGen的自动优化。cuSyncGen自动执行多项优化，以提高cuSync同步工作负载的性能。这些优化取决于GPU的架构细节、CUDA核心的占用率和网格大小。优化如下：
* 避免等待核心：等待核心机制确保生产者核心的所有线程块在消费者核心之前调度到GPU上。但是，如果生产者和消费者核心都可以在不到两个波次内执行完毕，我们就不需要等待核心机制。
* 避免自定义Tile处理顺序：当生产者和消费者核心的所有tile都可以在两个波次内执行时，我们也可以避免自定义tile处理顺序。
* 重排Tile加载和同步：基于tile的CUDA核心的通用工作流程是先将所有输入的tile加载到共享内存或寄存器中，然后对所有tile执行操作。我们可以将等待一个输入的tile与加载另一个tile的操作重排序，以将等待一个tile的时间与加载另一个输入tile的时间重叠。例如，在图4a中，第二个GeMM核心加载输入（A和B）的tile并计算输出矩阵（C）的tile（第6-9行）。我们可以将加载B tile与等待A tile的操作重排序，即交换第6-7行与第8-9行。由于没有等待B的tile，加载B tile可以与等待A tile重叠，从而提高性能。如果用户使用#pragma tile在核心中标注tile加载，cuSyncGen会自动执行此重排序。

D. 在ML模型推理中的适用性

cuSync适用性分析。我们现在从计算类型和GPU平均利用率的角度讨论cuSync在提高ML模型性能方面的适用性。首先，ML模型主要由基于tile的GPU核心组成，如GeMM和Conv2D。由于cuSync支持任何基于tile的核心，我们可以使用cuSync来同步ML模型的核心。其次，由于每个核心的波次数随着批处理大小的增加而增加，所有波次的平均利用率也会增加。然而，每个ML模型在训练和推理阶段都支持最大批处理大小限制。例如，GPT-3和LLaMA支持的最大token长度为2048。我们在实验中表明，即使对于这个最大批处理大小，GPU核心仍然会因为波次数少而导致平均利用率低。总之，cuSync适用于多种ML模型，因为ML模型主要包含基于tile的核心，并且ML模型支持的最大批处理大小仍然存在利用率不足的问题。

A4 实验环境

• 硬件配置:
  • 平台: 一台配备了2.60GHz 12核Intel Xeon E5-2690 CPU和448GB RAM的机器。
  • GPU: 8块NVIDIA Tesla V100 32GB GPU，通过NVLINK连接。
• 软件配置:
  • CUDA: 12.2。
  • 代码实现: GeMM和Conv2D的CUDA核心使用了NVIDIA CUTLASS 3.1。作者开发了一个融合Softmax和Dropout的核心。点态计算与GeMM和Conv2D核心融合。
• 模型与数据集:
  • 模型: MegatronLM GPT-3 145B, LLaMA 65.2B, ResNet-38, VGG-19。
  • 用途: 评估在1到各个模型支持的最大批处理大小范围内，cuSync对推理时间的缩减效果。
• 基线:
  • StreamSync: 标准的CUDA流同步。
  • Stream-K: 一种通过在所有SM中划分GeMM和Conv2D的最后一个线程块波次来提高GPU利用率的技术【索引10，Stream-K: Work-Centric Parallel Decomposition for Dense Matrix-Matrix Multiplication on the GPU, 2023, PPoPP ’23】。
• 实验方法: 报告了5次预热执行后20次执行的平均时间。

A4 实验结果

B. 编程易用性

代码修改量小。如表III所示，为支持使用cuSync对GeMM、Conv2D和Softmax-Dropout核心进行细粒度同步而增加和更改的代码行数与这些核心的总代码行数相比微不足道（0.5%-1%）。这表明cuSync方法只需少量修改即可为基于tile的计算核心实现多样的同步。

D. 最大同步开销

同步开销低。同步机制的开销主要来自全局内存访问和__syncthreads。总开销的百分比取决于GPU核心执行的计算量。在一个设计的极端实验中（最小计算量、最大线程块数），cuSync的同步开销比StreamSync高2-3%。这证明cuSync的同步机制开销很低。

E. 大语言模型推理结果

实验设置。我们评估了在8个GPU上使用模型并行的GPT-3和LLaMA中，cuSync在提示处理和token生成阶段对推理时间的缩减。提示处理阶段，我们考虑总token数（B×S）从512到2048；token生成阶段，我们考虑批处理请求数（B）从1到4，已生成token数（S'）从512到2048。cuSyncGen生成了多种策略，包括RowSync+WRT、TileSync、TileSync+WRT和Strided+TileSync+WRT（其中W、R、T分别代表避免等待核心、重排tile加载和避免自定义tile顺序的优化）。

1) MLP结果:
* 性能提升: 如图6a和6c所示，使用cuSync同步GPT-3 MLP和LLaMA MLP中的依赖GeMM，可将两者的执行时间降低高达20%。
* 性能分析 (基于表IV):
* 对于小批量(B×S=1到256)，TileSync+WRT表现最佳。在B×S=256时，cuSync比StreamSync减少了1个波次，性能提升显著。即使波次数没有减少，通过重叠第二个GeMM的W2 tile加载与第一个GeMM的计算，也能获得7%的性能提升。
* 对于大批量(B×S>512)，RowSync表现最佳。这是因为对整行进行一次同步比对多个tile同步产生的内存访问更少，且更多的行提供了更多重叠机会。
* 在最大批量(2048)时，加速比下降，因为cuSync减少的波次数占总波次数的比例变小。
* 核心调用开销: 核心调用开销约为6µs，远小于StreamSync和cuSync的执行时间差。因此，cuSync的性能提升远不止是重叠核心调用所能达到的。

2) 注意力机制结果:
* 性能提升: 如图6b所示，使用cuSync同步Attention的所有核心，相比StreamSync在GPT-3和LLaMA上均能带来6-16%的性能提升。
* 性能分析:
* 在提示处理阶段(S'=0)，StridedTileSync+WRT策略效果最好，因为它比TileSync同步次数少，比RowSync重叠机会多。
* 在token生成阶段(S'>1, S=1)，所有策略表现相似，因为不同同步策略在不同核心间达到了最佳平衡。

F. 计算机视觉模型推理结果

实验设置: 我们通过使用cuSync同步ResNet-38和VGG-19每层的所有Conv2D核心来评估推理时间的缩短。cuSyncGen生成了RowSync+WRT、Conv2DTileSync和Conv2DTileSync+WRT策略。

性能提升与分析:
* 如图7所示，使用cuSync同步ResNet-38和VGG-19每层的所有Conv2D核心，相比StreamSync可带来高达24%的性能提升。
* 性能提升随批处理大小的增加呈现振荡行为。例如，对于128个通道，性能提升从批大小为1的20%上升到4的24%，然后下降到8的3%，再上升到12的18%，最后又下降到16的3%。这种振荡行为是由于批处理大小的增加导致调用的线程块数量增加，从而使得cuSync减少的波次数比例呈现振荡。

G. 端到端模型推理

整体性能提升: 我们将cuSync同步的CUDA核心集成到所有四个ML模型中，并评估了端到端推理时间的改善。如图8所示，使用cuSync同步的核心将推理时间分别降低了：
* GPT-3: 6–15%
* LLaMA: 9–13%
* ResNet-38: 5–22%
* VGG-19: 6–16%
这表明cuSync能显著减少流行ML模型的推理时间。

H. 与Stream-K的比较

性能优越性: cuSync的最佳策略在GPT-3和LLaMA的GeMM核心上比Stream-K性能高出15%（如图6所示）。
原因分析:
1. 实现机制: Stream-K将GeMM工作负载分为两个核心调用，第二个核心专门处理最后一波的工作负载划分，这需要多次内存访问。而cuSync仅需一次原子加法来发布tile状态，一次读取来等待生产者tile状态。
2. 通用性: Stream-K的思想不易应用于基于tile的所有核心。目前，NVIDIA CUTLASS中的Stream-K仅支持GeMM计算。相比之下，cuSync适用于任何基于tile的核心，因此可以应用于Conv2D。

I. 优化的影响

优化效果显著。如表V所示，对TileSync应用所有优化（避免等待核心+重排加载+避免自定义顺序）可以显著降低执行时间，尤其是在线程块数量较少的 kernels 中。

A7 补充细节 (相关工作)

核心内同步。已有多项工作关注于单个CUDA核心内部线程的高效软件同步，主要用于不规则应用【索引9，Fine-Grained Synchronizations and Dataflow Programming on GPUs, 2015, ICS ’15】，【索引16，Fast Fine-Grained Global Synchronization on GPUs, 2019, ASPLOS ’19】，【索引17，Lock-Based Synchronization for GPU Architectures, 2016, CF ’16】。例如，Li等人【索引9】通过重组共享内存原子操作的微指令来改进线程间同步。Kai等人【索引16】提出了用于不规则应用的分层同步方法。Xu等人【索引17】设计了一种使用锁窃取避免死锁的锁设计。COCONET【索引8，Breaking the Computation and Communication Abstraction Barrier in Distributed Machine Learning Workloads, 2022, ASPLOS ’22】在计算和通信核心之间进行同步，以重叠通信传输和计算。cuSync与这些工作的不同之处在于，它专注于多个CUDA核心之间的线程同步，并提供了一个抽象框架来轻松设计多种同步策略。

硬件支持的同步。一些工作研究了硬件支持的核间线程同步原语。GLocks【索引2，GLocks: Efficient Support for Highly-Contended Locks in Many-Core CMPs, 2011, IPDPS ’11】是首个使用消息传递的硬件支持高竞争锁实现。HQL【索引18，HQL: A Scalable Synchronization Mechanism for GPUs, 2013, IPDPS ’13】是一种硬件加速的细粒度锁方案。ElTantway等人【索引5，Warp Scheduling for FineGrained Synchronization, 2018, HPCA】提出了一种硬件warp调度策略，通过降低等待自旋锁的线程所在warp的优先级来减少锁重试。Dalmia等人【索引3，Improving the Scalability of GPU Synchronization Primitives, 2023, IEEE Transactions on Parallel and Distributed Systems】为基于GPU的同步原语设计了多级屏障和信号量优先级机制。cuSync是一个用于同步多个CUDA核心线程的软件解决方案，这些硬件支持的机制与cuSync是互补的。

同步性能研究。Lingqi等人【索引19，A Study of Single and Multi-device Synchronization Methods in Nvidia GPUs, 2020, IPDPS】研究了用于归约操作的多种CUDA同步方法的性能和陷阱。Sinclair等人【索引14，HeteroSync: A benchmark suite for fine-grained synchronization on tightly coupled GPUs, 2017, IISWC】提出了一个基准套件，用于测量不同一致性协议和一致性模型下的同步原语性能。

Stream-K。Stream-K【索引10，Stream-K: Work-Centric Parallel Decomposition for Dense Matrix-Matrix Multiplication on the GPU, 2023, PPoPP ’23】是一种通过在GPU的所有SMs之间划分工作负载来提高利用率的GeMM实现。然而，Stream-K不易应用于GeMM之外的计算。相比之下，cuSync将多个核心的线程块适配到每个波次中，并适用于任何基于tile的计算。

A5 结论

当前先进的ML模型包含数千个独立的计算，这些计算在一个或多个GPU上执行。然而，这些模型未能充分利用GPU，因为单个计算无法完全占满GPU，并且模型中大量存在依赖性计算。本文提出了cuSync，一个用于对依赖计算的tiles进行细粒度同步的框架。通过仅同步依赖的tiles，我们的框架允许独立tiles的并发执行，从而提高了GPU的利用率。我们的实验表明，使用cuSync同步现有机器学习模型的计算可以显著减少这些模型的推理时间。

A6 附录

项目制品。制品【索引7，(Artifact) A Framework for Fine-Grained Synchronization of Dependent GPU Kernels, 2023】包含了cuSync的CUDA实现和复现我们所有结果的脚本。制品提供了包含所有预装依赖的Dockerfile和源代码。最新的源代码可在 https://github.com/microsoft/cusync 获取。该制品可复现第五节中的图6、7和8。

系统。我们在一个NVIDIA DGX-2系统上执行了实验，该系统包含8个通过NVLINK连接的NVIDIA Tesla V100 GPU。我们的实验可以在任何带GPU的系统上运行，但图8中的端到端推理结果可能无法在其他系统上复现。

提取制品。从【索引7】下载制品并解压zip文件：

unzip cusync-cgo-24.zip cd cusync-cgo-24

A. Docker容器

运行步骤。要在Docker容器内运行制品，请遵循以下步骤：
1. 安装docker: 遵循 https://docs.docker.com/engine/install/ubuntu/ 上的步骤安装docker引擎。
2. 安装NVIDIA容器工具包: 遵循 https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html 上的步骤安装NVIDIA容器工具包。
3. 创建容器: 使用Dockerfile创建docker容器，启动容器，并进入目录：

docker build -t cusync-cgo-24 .
docker run -it --gpus all cusync-cgo-24
cd /cusync

1. 检查PyTorch和CUDA安装: 检查torch是否支持CUDA，如果torch.cuda.is_available()返回True：

python
>>> import torch
>>> torch.cuda.is_available()
True

B. 本地运行

依赖安装。我们也可以本地运行代码，这需要安装所有依赖项。如果使用上述docker步骤，可以忽略这些步骤。
* Linux安装: 我们推荐使用Ubuntu 22.04作为Linux操作系统。我们没有在其他操作系统上测试过我们的制品，但我们相信Ubuntu 20.04和23.04也应该可以工作。
* 安装依赖: 执行以下命令安装依赖：

sudo apt update
sudo apt install gcc linux-headers-$(uname -r) \
make g++ git python3 wget \
unzip python3-pip build-essential cmake

```
*   **安装CUDA**: 我们需要先安装CUDA。在我们的实验中，我们在Ubuntu 22.04上使用了CUDA 12.2。CUDA 12.2工具包可以从 https://developer.nvidia.com/cuda-12-1-0-download-archive 下载。安装CUDA后，设置`nvcc`和CUDA路径：
```bash
export PATH="/usr/local/cuda/bin:$PATH"
export LD_LIBRARY_PATH="/usr/local/cuda/lib64:$LD_LIBRARY_PATH"

• 检查CUDA安装: 要检查CUDA安装，运行nvidia-smi，它应该会打印出系统中的所有GPU。否则，CUDA安装有问题。
• 安装Pytorch: 使用pip安装PyTorch：

sudo pip3 install torch torchvision torchaudio

• 检查Pytorch CUDA安装: 检查torch是否支持CUDA，如果torch.cuda.is_available()返回True。

获取源代码。源代码可以从【索引7】下载。最新的源代码可从cuSync仓库的CGO AE分支获取：

git clone --recurse-submodules \
https://github.com/microsoft/cusync
cd cusync
git checkout cgo-24-ae

C. 功能性和可重用性

运行测试。README.md文件包含了如何将代码编译到其他NVIDIA GPU架构的说明、一个示例和测试用例。可以通过执行这些测试用例来检查功能性。要运行测试，请执行：

make tests -j

如果所有测试都通过，那么我们就可以准备复现结果了。

D. 复现结果

复现主要结果。我们现在将复现图6、7b和8的主要结果。所有命令都应在cusync目录中执行。
* 大语言模型推理结果 [时间 60分钟]:

cd src/ml-bench/volta_transformer
python3 eval_llm.py mlp gpt3
python3 eval_llm.py attention gpt3
python3 eval_llm.py mlp llama
python3 eval_llm.py attention llama
python3 allreduce_times.py

• 计算机视觉推理结果 [时间 60分钟]:

cd src/ml-bench/volta_conv2d
python3 eval_conv.py resnet
python3 eval_conv.py vgg

• 生成图表 [时间 5分钟]:

cd src/ml-bench/plots make -j

当前目录将生成PDF格式的图表，可以与论文中的图表进行核对。