一起聊聊Nvidia Hopper新特性之Pipeline

前言

大家好，我是企鹅。欢迎来到我们关于GEMM(通用矩阵乘法)教程系列的第二部分。在第一部分中，我们讨论了GEMM的计算方面，介绍了WGMMA，这是基于NVIDIA® Hopper™架构在GPU上执行小矩阵块乘法的基本指令。在这一部分，我们将重点关注GEMM的内存方面。具体来说，我们将解释如何有效地将操作数张量的小块从GPU的全局内存传输到片上内存，从而可以传递给WGMMA(或其他基本MMA指令)。

相关阅读：

• 一起聊聊Nvidia Hopper新特性之WGMMA
• 一起聊聊Nvidia Hopper 新特性之TMA

需要解释的主要概念是如何编排数据管道，以有效地为张量核心提供数据。在GEMM内核设计的背景下，流水线指的是通过维护多个数据缓冲区来重叠复制和MMA操作的思想。在本文中，我们将介绍两种在Hopper架构上有效的流水线策略：

• warp specialization。是一种将线程束专门化为生产者(数据传输)和消费者(计算)，并让它们并发运行的方法。
• multi stage。通过使用异步复制(Hopper上的TMA或Ampere上的cp.async)来加载下一组数据，同时对当前数据集进行计算，来掩盖数据传输。线程束兼具生产者和消费者的角色。

为了确保内核的正确性，需要仔细注意当前的数据依赖关系，这决定了缓冲区何时可以被MMA指令读取或由复制操作填充。我们将详细介绍如何使用CUTLASS库中的工具(特别是CUTLASS Pipeline类)编写用于流水线GEMM内核的必要同步逻辑。然后，我们将评估流水线的性能，并展示如何利用这一优化思想就能达到Hopper GEMM内核在半精度下约65%的利用率。

大局观之见”喂养野兽（tensor core）”

GEMM内核中有两个主要操作：将数据复制到正确的内存地址，以及对它们进行乘法累加。前者由访存指令处理：Hopper中的TMA、Ampere中的cp.async，以及早期架构中的普通复制。从2017年的Volta架构开始，就成为了Tensor core的专属操作。

经过多代发展，Tensor core已经成为一头能够高效消耗输入数据的”野兽”。例如，H200 SXM GPU的张量核心可以提供高达3，958 TFLOPS(每秒万亿次浮点运算)的性能。而同样的H200 SXM GPU的内存带宽只有4.8 TB/s。这种数据传输速度远远慢于张量核心的速度，并且通常很难完全利用!因此，CUDA编程的一个共同主题，特别是GEMM内核设计，就是找出如何快速复制数据，以保持张量核心忙碌。我们称这个过程为”喂养野兽”。

总的来说，有两种主要的”喂养野兽”策略，它们是互补的，在不同范围(网格or块)发挥作用。第一个策略是有效的线程块调度，即在CTAs之间分配计算任务，以获得良好的负载均衡和更高的L2缓存命中率。我们将在后续的博客文章中讨论这一点。第二个策略，也是我们在本教程中重点关注的，是重叠复制与数学运算。具体来说，当张量核心忙于乘法累加一批数字时，我们应该告诉复制单元去复制下一批数字。这样，我们就可以有效地隐藏部分复制延迟。这就是流水线的目标。

延迟，warps， and warp-specialization

在讨论流水线的机制之前，让我们回顾一下介绍中提到的两种重叠策略的历史：multi stage和warp specialization。

首先，重叠内存复制和数学运算的想法并不新鲜，也不仅限于GPU。熟悉CPU的读者可能会发现这类似于缓存预取技术，即在需要数据之前异步发出获取请求。事实上，我们在本文中讨论的流水线技术在概念上与CPU缓存预取是相同的!然而，由于在GPU上进行预取代价很高(占用芯片面积)，因此实现方式有所不同。

GPU程序员创建重叠的最基本方法是利用额外的线程束(线程束是由32个连续线程组成的)。Nvidia GPU允许每个SM(流式多处理器)拥有大量的线程束，并且可以在它们之间进行快速切换。特别是，如果一个线程束遇到缓慢的内存获取，调度器可以简单地切换到另一个线程束。为了给调度器更多机会隐藏延迟，约在2011年引入了一种称为线程束专业化的技术。在线程束专业化中，一些线程束专门用于内存获取(生产者)，而其他线程束专门用于计算(消费者)，并使用命名屏障进行同步。这样做的目的是，调度器可以更容易地隐藏复制操作的延迟(反之亦然)。

从Ampere架构开始，Nvidia引入了cp.async，它允许内存复制在执行数学运算的同一个线程束中异步进行。具体来说，异步意味着一个线程束可以发出cp.async来加载下一个缓冲区的数据，然后在不等待异步加载完成的情况下执行当前缓冲区的数学运算。这样就不需要使用线程束专业化来掩盖数据传输和计算了。multi stage内核设计利用了这一思想。最快的Ampere GEMM内核以及著名的FlashAttention-2都使用了multi stage内核设计。

最后，在最新的Hopper GPU架构中，引入了诸如TMA异步复制和跨线程束的寄存器重新分配等新特性，这些特性使得线程束专业化在Hopper上变得非常有效(如下所述)。事实上，最快的CUTLASS Hopper GEMM内核使用了warp specialization。

Pipeline 解析

下图描绘了一个理想的LOAD和MMA的流水线。这里，LOAD指的是将操作数矩阵块从全局内存(GMEM)复制到共享内存(SMEM)的过程，而MMA指的是在SMEM中存储的操作数块上执行的张量核心运算。如图所示，通过重叠两个LOAD和两个MMA，我们节省了2个时间单位。

这种流水线设计的关键思想是：在张量核心忙于执行一批数据的MMA运算时，同时启动下一批数据的LOAD操作。这样可以有效地隐藏LOAD操作的延迟，提高整体的计算效率。

实现这种流水线需要仔细管理数据依赖关系，确保LOAD操作不会覆盖正在被MMA使用的数据。我们需要使用同步机制，如屏障和信号量，来协调LOAD和MMA的执行时序。

这里需要考虑一个问题：LOAD_1和LOAD_2将数据复制到哪里？显然，我们不希望后续的加载操作覆盖MMA还未使用的先前加载的数据。同时也不希望因为等待SMEM可用而造成不必要的停滞。否则，预期的2个时间单位的收益将无法实现。

解决这个问题的简单方法是在SMEM中预留两倍于MMA所需的内存空间，并交替使用它们。这种策略称为双缓冲，如下图所示。当然，我们也可以推广到使用更多的交替缓冲区。这样做可以创造更多的重叠机会，从而更有效地利用可用的硬件，但代价是需要使用更多的SMEM。

通过这种双缓冲或多缓冲的方式，我们可以确保MMA操作始终有可用的数据，同时LOAD操作也可以异步地将下一批数据加载到空闲的缓冲区中。这样就实现了计算和数据传输的重叠，从而提高整体的性能。

实现正确高效的流水线并非易事。程序员必须处理多个缓冲区以及跨多个线程的异步加载调用。在下一节中，我们将展示如何通过CUTLASS的Pipeline类抽象来实现流水线。

这需要仔细的设计和编码工作。我们必须确保数据依赖关系得到正确管理，避免出现竞争条件或数据覆盖等问题。使用CUTLASS提供的Pipeline类可以大大简化这一过程，让我们专注于高层次的流水线逻辑，而不必过多地关注底层的同步细节。

CUTLASS的Pipeline类为我们提供了一个优雅的接口，帮助我们有条不紊地管理多个缓冲区，协调异步的加载和计算操作。通过使用这个抽象，我们可以更轻松地编写出正确且高效的流水线GEMM内核。这对于充分利用GPU的计算能力至关重要。我期待在下一节中深入了解CUTLASS Pipeline类的使用方法。

CUTLASS中流水线抽象

CUTLASS的异步Pipeline类提供了一个有效的抽象，用于管理跨多个数据缓冲区和参与线程的复制和计算。它们包括PipelineAsync、PipelineTmaAsync和PipelineTransactionAsync类，其中”Pipeline”是一个通用的引用。

我们首先以高层次的角度解释CUTLASS Pipeline如何协调数据流水线。假设buffers是一个具有N个阶段的共享内存缓冲区。我们希望在生产者(如TMA)向缓冲区写入数据和消费者(如WGMMA)使用可用数据之间进行同步。

屏障。为了在生产者和消费者之间同步缓冲区阶段，Pipeline遵循标准的获取和释放模型，使用锁来管理对缓冲区的访问。为此，让full_barrier和empty_barrier是两个大小为N的屏障对象数组。这些屏障对象具有一个相位位，初始化为0，并在0和1之间切换。

具体来说，这些屏障对象将是驻留在SMEM中的mbarrier对象。mbarrier对象初始化时包含上述相位位和到达计数。它支持arrive-on和wait操作，并根据达到到达计数阈值而切换相位。重要的是，这些屏障对象的值应该对所有线程可见。

线程本地的流水线状态。接下来，我们有PipelineState类作为线程本地的枚举器，用于跟踪线程的当前索引和相位，其中N个阶段作为模板参数传入。索引取整数值模N，相位为0或1。此外，PipelineState类的++运算符被重载，使得索引在模N的范围内递增，当索引递增到0时相位会翻转。

同步。我们现在解释如何使用屏障对象和线程本地的流水线状态来同步生产者和消费者。为了避免混淆，让我们区分生产者动作和发出该动作的生产者线程，因为它们可能会被解耦(想想TMA)。首先，生产者动作将翻转full_barrier[i]的相位，以表示它已经填充了缓冲区的第i个阶段，这样消费者线程就可以读取它了。同样，消费者线程将翻转empty_barrier[i]的相位，以表示它们已经完成了对缓冲区第i个阶段的消费，这样生产者就可以写入了。

需要注意的是，我们不关心生产者动作或消费者线程到底是如何在SMEM中翻转相位位的，只要它是通过到达计数机制完成的。例如，所有消费者线程可以集体执行到达计数递增，或者每个线程束中选出一个消费者线程来执行同样的操作。

最后，无论是消费者还是生产者，每个线程都会跟踪一个相位，以与屏障对象的相位进行匹配，事实上，同时担任消费者和生产者角色的线程需要跟踪两个相位。这些线程的”内部”相位也需要随着内核主循环的进行而翻转。

四个流水线方法。现在让pipeline是一个Pipeline类的实例，它被初始化为full_barrier和empty_barrier的指针，而pipe_state是一个PipelineState类的实例。那么pipeline可以调用以下四个关键方法：

• pipeline.producer_acquire(pipe_state)。阻塞调用线程，直到empty_barrier[pipe_state.index()]的相位与pipe_state.phase()相反。
• pipeline.producer_commit(pipe_state)。向full_barrier[pipe_state.index()]发信号，使其到达计数递增。
• pipeline.consumer_wait(pipe_state)。阻塞调用线程，直到full_barrier[pipe_state.index()]的相位与pipe_state.phase()相反。
• pipeline.consumer_release(pipe_state)。向empty_barrier[pipe_state.index()]发信号，使其到达计数递增。

在描述阻塞指令producer_acquire和consumer_wait时，我们所说的”翻转相对于pipe_state的相位”的意思是，例如，如果屏障的当前相位是0，那么当pipe_state的相位也是0时该方法会阻塞，而当相位是1时则不会阻塞。

需要注意的是，按照编写方式，(producer_acquire， consumer_release)和(producer_commit， consumer_wait)这两对方法在功能上是完全对称的。然而，如果涉及的Pipeline类是PipelineTmaAsync，那么full_barrier就会被封装为cutlass::arch::ClusterTransactionBarrier类的一个实例，full_barrier的信号机制由TMA加载方法本身通过递增事务计数来处理。在这种情况下，producer_commit方法实际上是一个空操作;我们稍后会回到这一点。但是，在伪代码中，我们仍然会插入producer_commit，就像现在一样，前提是TMA拷贝还没有写出。

综合起来，以下伪代码展示了这四个流水线方法的运作：

using PipelineState = typename cutlass::PipelineState<N>;
// We initialize smem_pipe_write to start with an opposite phase
// (i.e., 1 instead of 0), since the buffers start out as empty.
PipelineState smem_pipe_write = cutlass::make_producer_start_state<Pipeline>();
PipelineState smem_pipe_read;
for (int i = 0; i < total_steps; ++i) {
  pipeline.producer_acquire(smem_pipe_write);
// Acquire data (e.g. TMA, cp.async, etc.)  
  pipeline.producer_commit(smem_pipe_write);
  ++smem_pipe_write;

  pipeline.consumer_wait(smem_pipe_read);
// Compute workload (e.g. WGMMA)
  pipeline.consumer_release(smem_pipe_read);
  ++smem_pipe_read;
}

我们发现上述代码片段有助于说明生产者/消费者获取和释放模式。我们邀请读者跟踪所有相关状态，逐步执行循环的几个步骤，并将这段伪代码与之前给出的详细同步描述联系起来。

然而，这段代码展示的是一种串行执行流程，其中生产者和消费者操作从不并发运行，因此在实践中并不实用。在高效的流水线工作负载中，生产者和消费者必须重叠执行。接下来我们将讨论多阶段内核设计，这是实现这一目标的一种方法。

Multistage kernel 设计

让我们使用Pipeline类的TMA专用版本PipelineTmaAsync来创建一个2阶段流水线，用于Hopper GEMM内核中使其TMA与WGMMA重叠执行。这个内核以128个线程启动（即1个warpgroup）。我们假设读者熟悉CUTLASS中TMA和WGMMA的语法，我们在之前的两篇博客文章中已详细讨论过这些内容。因此，我们省略了用于cute::copy和cute::gemm调用的张量准备工作。

using MainloopPipeline = typename cutlass::PipelineTmaAsync<2>;
using PipelineState = typename cutlass::PipelineState<2>;

typename MainloopPipeline::Params params;
// number of bytes transferred by TMA load per stage (A and B)
params.transaction_bytes = TmaTransactionBytes;
params.role = MainloopPipeline::ThreadCategory::ProducerConsumer;
params.is_leader = threadIdx.x == 0;
params.num_consumers = 128;

// Disregard clusters for this example
auto cluster_shape = Shape<_1,_1,_1>{};

// pipeline_storage is instance of cutlass::PipelineTmaAsync<2>::SharedStorage
// Has full_barrier and empty_barrier as members
// Located in the SharedStorage struct that manages objects in smem
MainloopPipeline pipeline(shared_storage.pipeline_storage, params, cluster_shape);

__syncthreads();

PipelineState smem_pipe_write = 
    cutlass::make_producer_start_state<MainloopPipeline>();
PipelineState smem_pipe_read;

// Prepare tensors for GEMM
// ...

// Issue the first TMA load with leader thread
if(threadIdx.x == 0) {
  pipeline.producer_acquire(smem_pipe_write);
  BarrierType *tmaBar = pipeline.producer_get_barrier(smem_pipe_write);
// smem_pipe_write.index() == 0  
  copy(tma_load_a.with(*tmaBar, 0), tAgA(_,0), tAsA(_,0));
  copy(tma_load_b.with(*tmaBar, 0), tBgB(_,0), tBsB(_,0));
  ++smem_pipe_write;
}

for (int i = 0; i < k_tile_count - 1; ++i) {
// Only leader thread issues TMA load
if(threadIdx.x == 0) {
    pipeline.producer_acquire(smem_pipe_write);
    BarrierType *tmaBar = pipeline.producer_get_barrier(smem_pipe_write);
    auto write_stage = smem_pipe_write.index();
    copy(tma_load_a.with(*tmaBar, 0), tAgA(_,i+1), tAsA(_,write_stage));
    copy(tma_load_b.with(*tmaBar, 0), tBgB(_,i+1), tBsB(_,write_stage));
    ++smem_pipe_write;
  }

// Compute on the completed load from prior iteration
  pipeline.consumer_wait(smem_pipe_read);
auto read_stage = smem_pipe_read.index();
// WGMMA
  warpgroup_arrive();
  gemm(tiled_mma, tCrA(_,_,_,read_stage), tCrB(_,_,_,read_stage), tCrC);
  warpgroup_commit_batch();
  warpgroup_wait<0>();
  pipeline.consumer_release(smem_pipe_read);
  ++smem_pipe_read;
}

// Handle the last compute iteration
pipeline.consumer_wait(smem_pipe_read);
auto read_stage = smem_pipe_read.index();
warpgroup_arrive();
gemm(tiled_mma, tCrA(_,_,_,read_stage), tCrB(_,_,_,read_stage), tCrC);
warpgroup_commit_batch();
warpgroup_wait<0>();
pipeline.consumer_release(smem_pipe_read);

// Epilogue for writing out accumulator
axpby(alpha, tCrC, beta, tCgC);

在这里，主循环的每次迭代中，异步发出第(i+1)个TMA加载指令，并执行第i个WGMMA计算，注意smem_pipe_write和smem_pipe_read之间相差一个偏移量。

在这段伪代码中，请注意我们在TMA博客文章中使用的cute::set_barrier_transaction_bytes方法（或其等效方法cutlass::arch::arrive_and_expect_tx）并未出现。相反，它的功能被PipelineTmaAsync类中的producer_acquire方法所接管。实际上，该方法在内部执行以下操作，其中stage和phase是其PipelineState参数的索引和阶段：

if (barrier_token != BarrierStatus::WaitDone) {
   empty_barrier_ptr_[stage].wait(phase);
}
 
if (params_.is_leader) {
   full_barrier_ptr_[stage].arrive_and_expect_tx(params_.transaction_bytes);
}

此外，我们使用producer_get_barrier方法并传入smem_pipe_write参数来获取指向full_barrier[smem_pipe_write.index()]的指针，这正是cute::copy调用中TMA TiledCopy对象tma_load_a和tma_load_b所需要的。

通过这种方式将cute::copy调用与流水线的mbarrier对象full_barrier关联起来，我们可以使用TMA的基于事务计数的完成机制向消费者发出缓冲区已准备好可以使用的信号，这样就不需要从流水线对象本身调用producer_commit。这就是为什么CUTLASS将PipelineTmaAsync的producer_commit实现为空操作的原因。

这种流水线结构方式允许数据传输和计算重叠执行，充分发挥异步操作隐藏延迟的潜力。虽然我们在此示例中使用了TMA，但在Ampere架构中也可以使用cp.async实现类似的技术。然而，在Hopper架构中，有时使用warp specialization 设计比multi stage设计更为可取，下面我们将对此进行解释。

Warp specialization

在multi stage中，每个线程束同时扮演生产者和消费者角色。使用PipelineState抽象来处理两种角色之间的切换，而TMA加载的异步性使两种类型的操作能够重叠执行。另一种策略是warp specialization设计，它为不同的线程束分配不同的角色，这样我们就有专门负责内存复制的生产者线程束和专门负责计算的消费者线程束。如上所述，线程束调度器可以通过在两种类型的线程束之间切换来隐藏延迟。值得注意的是，与多阶段内核不同，线程束专用设计本身并不依赖于异步执行，但在实践中仍然能从中获得巨大益处。

具体到我们的GEMM，生产者线程束使用TMA将数据从全局内存加载到共享内存，而消费者线程束使用WGMMA计算分块GEMM。值得注意的是，在我们简化的设置中，两种类型线程束中的执行流在内部是串行的，即TMA和WGMMA指令本身在线程束组内并没有被重叠执行。然而，有更复杂的内核调度方式能够利用TMA和WGMMA的异步性，也可以实现与其他指令在线程束组内的重叠执行，例如FlashAttention-3中的实现。

• TMA比早期的复制操作消耗更少的寄存器。
• WGMMA可以直接从共享内存获取操作数，这意味着消费者线程束不需要执行自己的内存加载操作。
• Hopper允许通过setmaxnreg指令手动进行线程束组范围的寄存器(去)分配。因此，更大比例的寄存器可以分配给通常需要更多寄存器的消费者线程束。

进一步说明最后一点，每个SM有限定数量的寄存器，在Hopper之前的架构中，每个线程束在内核启动时被分配固定且相等数量的寄存器。这对于多阶段流水线来说是可以的，因为每个线程束做相同的工作，但对于线程束专用模式来说通常是浪费的：生产者线程束（仅加载数据）通常比消费者线程束（进行计算）需要更少的寄存器，尤其是使用TMA时。对于寄存器密集型工作负载，能够利用这些浪费的寄存器可能意味着允许每个SM有更多的线程束或避免寄存器溢出。

现在让我们展示一段线程束专用代码片段。和之前一样，Pipeline类抽象了设置线程束专用内核的复杂性。

// Create the pipeline and the iterator for the stage
using MainloopPipeline = typename cutlass::PipelineAsync<2>;
using PipelineState = typename cutlass::PipelineState<2>;

// Producer warps
if (isProducerWarp(threadIdx.x)) {
// Only one thread should be calling TMA
if(isTMAThread(threadIdx.x)) { 
    PipelineState smem_pipe_write = 
      cutlass::make_producer_start_state<MainloopPipeline>();
    for (...) {
      pipeline.producer_acquire(smem_pipe_write);
      copy(...); // TMA
      ++smem_pipe_write;
    }
  }
}
// Consumer warps
else {
  PipelineState smem_pipe_read;
for (...) {
    pipeline.consumer_wait(smem_pipe_read);
    // WGMMA
    pipeline.consumer_release(smem_pipe_read);
    ++smem_pipe_read;
  }
// Epilogue
}

格式与我们之前讨论的基本流水线类似，但这次有一个外部条件判断，将工作负载分为生产者线程束和消费者线程束。尾部处理属于消费者线程束，因为它涉及将消费者线程寄存器中保存的累加器写出。

要确定一个线程属于哪个线程束和线程束组，我们可以执行以下操作：

int warp_group_idx = __shfl_sync(0xffffffff, threadIdx.x / 128, 0);
int warp_idx_in_warpgroup = __shfl_sync(0xffffffff, (threadIdx.x / 32) % 4, 0);
int warp_group_thread_idx = threadIdx.x % 128;

上面的代码片段还使用了__shfl_sync操作，这是一个线程束范围内的值广播（更多信息请参阅此处）。这是为了确保线程束中的所有线程获得相同的值。

现在让我们关注这如何应用于GEMM。在本系列的第一部分中，我们讨论了在线程束组级别组织的WGMMA指令。因此，我们也在线程束组级别组织生产者和消费者。我们使用TMA流水线，以便在生产者端使用TMA。

对于2个阶段和2个线程束组，我们首先按如下方式更改线程束专用内核的流水线初始化：

using MainloopPipeline = typename cutlass::PipelineTmaAsync<2>;
using PipelineState = typename cutlass::PipelineState<2>;

typename MainloopPipeline::Params params;
params.transaction_bytes = TmaTransactionBytes; 
constint producerWarpGroupId = 0; 
if (warp_group_idx == producerWarpGroupId)
  params.role = MainloopPipeline::ThreadCategory::Producer;
else
  params.role = MainloopPipeline::ThreadCategory::Consumer;
params.is_leader = warp_group_thread_idx == 0;  
params.num_consumers = 128; 

auto cluster_shape = make_shape(Int<1>{},Int<1>{},Int<1>{});

// Create the pipeline
MainloopPipeline pipeline(shared_storage.pipeline_storage, params, cluster_shape);

我们强调第12行，以突出虽然params.num_consumers仍然等于128，但这现在只计算消费者线程束组的128个线程，而不是所有256个线程。

现在进入主循环。总体结构与初始代码示例相同，但生产者端有一些区别：

// Example values for Hopper GEMM with 1 consumer warpgroup
using LowerRegisterCount = Int<40>;
using HigherRegisterCount = Int<256>;

if (warp_group_idx == producerWarpGroupId) {
  cutlass::arch::warpgroup_reg_dealloc<LowerRegisterCount{}>();
int lane_predicate = cute::elect_one_sync();
if (warp_idx_in_warpgroup == 0 && lane_predicate) {
    PipelineState smem_pipe_write = 
      cutlass::make_producer_start_state<MainloopPipeline>();
    for (...) {
      pipeline.producer_acquire(smem_pipe_write);
      copy(...); // TMA
      ++smem_pipe_write;
    }
  }
} else { // consumer warpgroup
  cutlass::arch::warpgroup_reg_alloc<HigherRegisterCount{}>();
  PipelineState smem_pipe_read;
for (...) {
    pipeline.consumer_wait(smem_pipe_read);
    gemm(...); // WGMMA
    pipeline.consumer_release(smem_pipe_read);
    ++smem_pipe_read;
  }
// Epilogue to write out accumulator
  axpby(...);
}

在第6行和第18行，我们使用CUTLASS调用手动(去)分配额外的寄存器，这反过来调用PTX原语setmaxnreg，调整分配给线程束组中线程的寄存器。正如文档中所解释的，warpgroup_reg_dealloc<M>()释放额外的寄存器以将每线程最大寄存器数量减少到M，而warpgroup_reg_alloc<N>()请求额外的寄存器以将每线程最大寄存器数量提高到N。

这些寄存器数量的确切值取决于算法和硬件施加的约束。在Hopper架构中，一个线程最多可以拥有255个寄存器，setmaxnreg可以设置为24到256之间（包含两端）的8的倍数。一般而言，对于Hopper GEMM线程束专用内核，建议安排一个CTA占用整个SM。因此，我们应该尝试选择寄存器数量，使得(a)分配给发出TMA的生产者线程束组的寄存器数量最少，以及(b)利用每个SM全部64K的寄存器文件大小。例如，对于1个生产者线程束组和2个消费者线程束组，24/240/240的分配通常是有效的（总和为504 < 512，而512128 = 641024），同样，对于1个生产者和3个消费者线程束组，32/160/160/160的分配会被使用。还要注意，如果尝试分配超过寄存器文件大小的总寄存器数，程序将崩溃。

此外，我们必须确保一个线程束组中只有一个线程调用TMA。在我们的代码示例中，我们确保只有第一个线程束参与这一过程，并且由使用elect_one_sync选择的一个线程负责TMA调用。这段代码是为2个线程束组设计的，但对于更多数量的线程束组和阶段，也可以进行最小的修改来使用。

选择使用的线程束组数量和阶段数应通过对内核进行仔细的性能分析来确定。作为一般经验法则，更多的阶段和更多的线程束组意味着更多的并行和重叠机会，但也会使用更多的资源。特别是，使用更多阶段需要为缓冲区提供更多的共享内存(SMEM)，而使用更多线程束组会增加寄存器压力。

性能分析

我们使用CUTLASS Hopper GEMM教程代码作为我们半精度(FP16)数据类型的多阶段和线程束专用GEMM内核的基础。我们还修改了代码以支持FP32累加并使用TMA存储写出输出。然后，我们针对MxNxK = 8192x8192x8192进行了两个版本的调优，为FP16累加和FP32累加选择了不同的块大小。我们选择的块大小和阶段数如下（其中bMxbNxbK可以整除MxNxK）：

• FP16累加：bM = 256，bN = 256，bK = 96，2个阶段，4个MMA线程束组。集群大小(1， 2， 1)。
• FP32累加：bM = 256，bN = 192，bK = 128，2个阶段，2个MMA线程束组。集群大小(1， 2， 1)。

我们使用转换为FP16的随机浮点数初始化矩阵，并记录了以下TFLOP/s（10次迭代，5次测量的平均值）：

• FP16累加：多阶段531，线程束专用536。
• FP32累加：多阶段477，线程束专用485。

请注意，H100 PCIe GPU上半精度MMA的理论峰值性能为750 TFLOP/s，因此在FP32累加的标准设置下，我们达到了理论峰值的约65%。

需要注意的是，CUTLASS Hopper GEMM教程代码使用随机选择的±1初始化矩阵，因此会报告不切实际的高性能；参见这篇文章。例如，使用±1初始化矩阵时，我们的FP16累加多阶段内核的性能会从约530膨胀到约630 TFLOP/s。

现在为了比较，使用CUTLASS分析器测量的最快的CUTLASS FP16 Hopper GEMM内核（进行10次分析迭代）产生了630 TFLOP/s（约84%的利用率）。（注：本文的早期版本报告了较低的约74%利用率数字，因为它使用了过多的分析迭代次数，导致H100 PCIe GPU的350W TDP下出现热降频。）这个数字是通过以下内核获得的：

cutlass3x_sm90_tensorop_s64x256x16gemm_f16_f16_f32_void_f16_128x256x64_2x1x1_0_tnn_align8_warpspecialized_cooperative_epi_tma

请注意，这个CUTLASS内核采用了此处描述的”线程束专用持久合作”设计。我们预计，通过实现线程块光栅化和在CTA之间重叠前序和后序处理的持久内核，我们当前的流水线GEMM内核与最快GEMM内核之间的差距将大大缩小。在更非典型的问题几何形状下，使用Stream-K进行负载平衡也是一个因素。在这个正方形例子中，Stream-K CUTLASS内核表现几乎同样出色（625 TFLOP/s）。

现在我们来评论线程束组范围寄存器重新分配对线程束专用内核的相关性。要查看寄存器使用情况，我们可以使用以下标志编译我们的内核：-Xptxas=–verbose。（注意：此标志与–generate-code不兼容。请使用–gencode代替。）使用寄存器重新分配时，您将看到寄存器使用量是作为所使用的线程束组数量的函数固定的。例如，总共使用3个线程束组时：

0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info    : Used 168 registers

4 warpgroups：

0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info    : Used 128 registers

请注意1683 = 504和1284 = 512，这些数字是生产者和消费者寄存器计数总和必须小于或等于的值（相关地：这就是为什么32/240/240分割在3个线程束组的情况下不起作用的原因）。

另一方面，也可能寄存器使用率本来就很低，以至于寄存器重新分配没有任何实际影响。例如，对于FP16累加，当移除寄存器重新分配时，我们看到：

0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info    : Used 90 registers

此外，重新测量时间表明这种变化没有影响。但对于FP32累加，我们看到：

2784 bytes stack frame, 4764 bytes spill stores, 4760 bytes spill loads
ptxas info    : Used 168 registers

当重新测量时间时，我们现在获得约21 TFLOP/s，性能灾难性下降！然而，我们注意到，调整优化参数为(bM = 128, bN = 256, bK = 128, 2个阶段, 2个MMA线程束组, 集群(2,1,1))可以产生几乎同样良好的性能(460 TFLOP/s)，没有溢出，也不需要寄存器重新分配。

最后，在像FlashAttention-3这样的融合线程束专用内核设计中，寄存器中保存多个累加器，使用寄存器重新分配变得必不可少，以避免过度溢出。

总结

在本文中，我们全面介绍了流水线技术。我们介绍了其通过重叠内存复制和数学运算来隐藏延迟的目标，以及为什么这对良好性能至关重要。然后，我们介绍了两种流水线设计：

Multi stage：通过使用异步复制（Hopper上的TMA或Ampere上的cp.async）加载下一组数据，同时计算当前数据来掩盖数据传输延迟。线程束同时承担生产者和消费者角色。

Warp specialization：将线程束专门化为生产者和消费者，并让它们并发运行。此外，生产者或消费者操作还可以是异步的（例如，Hopper上的TMA和WGMMA）。

我们详细介绍了如何使用CUTLASS Pipeline类来管理在Hopper GEMM内核中实现这两种流水线策略所需的同步逻辑。最后，我们对GEMM示例中的两种流水线类型进行了比较。虽然在我们的简化设置中两者表现几乎相当，但在实践中，性能最佳的Hopper GEMM内核使用线程束专用化（例如，CUTLASS分析器所展示的）。