第 10 章 内核内流水线、线程束 专用化与协作线程块集群
第 10 章

内核内流水线、线程束 专用化与协作线程块集群

Chapter 10. Intra-Kernel Pipelining, Warp Specialization, and Cooperative Thread Block Clusters

原书页码 401–456
本章目录
  1. 内核内流水线技术
  2. 使用 CUDA Pipeline API 的协作式分块与双缓冲
  3. 线程束专用化与生产者-消费者模型
  4. 使用 CUDA Pipeline API 实现线程束专用化
  5. PyTorch、CUDA Pipeline API 与线程束专用化
  6. 持久化内核与巨型内核
  7. 持久化内核的常见工作负载
  8. 用于推理的巨型内核
  9. 持久化内核与线程束专用化
  10. 协作组
  11. 协作网格同步与持久化内核
  12. 何时结合持久化内核与协作组
  13. 线程块集群与分布式共享内存
  14. 线程块重排(Swizzling)
  15. 分布式共享内存
  16. 暂存内存(Scratch Memory)
  17. 启动线程块集群
  18. 使用协作组 API 协调线程块集群
  19. 线程块对
  20. 使用线程块集群减少全局内存流量
  21. 使用线程块集群设计高效算法
  22. 使用线程块集群的线程束专用化
  23. 关键要点
  24. 结语

In the previous chapters, we covered fundamental optimizations such as tuning memory access, maximizing parallelism, overlapping computation and data transfer, boosting occupancy, and minimizing warp stalls. These helped hide latency and eliminate bottlenecks. Modern GPUs, however, offer advanced hardware features and execution models that let us take the fundamental optimization techniques even further.在前面的章节中,我们介绍了若干基础优化技术,例如调整显存访问、最大化并行度、重叠计算与数据传输、提升占用率以及减少线程束停顿。这些技术有助于隐藏时延并消除瓶颈。然而,现代 GPU 提供了先进的硬件特性与执行模型,使我们能够把基础优化技术进一步推向深入。

In this chapter, we introduce some more advanced CUDA techniques such as warp-specialized pipelines, cooperative groups with grid-level and cluster-level synchronization, persistent kernels that loop over dynamic work queues, and thread block clusters (aka cooperative thread array cluster [CTA]) that use distributed shared memory (DSMEM or DSM) and Tensor Memory Accelerator (TMA) multicast. At a high level, a thread block cluster is a group of thread blocks that are guaranteed to run concurrently. They can read, write, and perform atomics to each other’s shared memory using DSMEM.本章将介绍一些更高级的 CUDA 技术,例如线程束专用化流水线、支持网格级与集群级同步的协作组、在动态工作队列上循环的持久化内核,以及使用分布式共享内存(DSMEM 或 DSM)和 Tensor Memory Accelerator(TMA)多播的线程块集群(又称*协作线程数组集群*[CTA])。从高层来看,线程块集群是一组保证并发运行的线程块。它们可以使用 DSMEM 读取、写入彼此的共享内存,并对其执行原子操作。

These methods let us overlap memory accesses and compute operations without host intervention. We can also share data on-chip across thread blocks—and keep every SM fully utilized.这些方法使我们无需主机干预即可重叠显存访问与计算操作。我们还可以跨线程块在片上共享数据,并让每个 SM 都得到充分利用。

By understanding these modern GPU execution models, you’ll be ready to progress to the next chapter where we extend these optimizations even further by exploring inter-kernel pipelines with CUDA streams. The next chapter builds inter-kernel pipelines on the foundation of intra-kernel optimizations discussed throughout this chapter.通过理解这些现代 GPU 执行模型,你将为进入下一章做好准备——在下一章中,我们将通过探索基于 CUDA 流的核间流水线,把这些优化进一步扩展。下一章将在本章通篇讨论的核内优化基础上构建核间流水线。

表 10-1使用 CUDA Pipeline API 在现代 GPU 上实现核内流水线的两种方法Two approaches for intra-kernel pipelining on modern GPUs using the CUDA Pipeline API
API variantAPI 变体Best for最适合Main use主要用途
Double-buffered pipeline双缓冲流水线Loop-based tiling and double buffering基于循环的分块与双缓冲Overlapping loads and compute in the same warp or block在同一线程束或线程块内重叠加载与计算
Warp-specialized pipeline (e.g.,线程束专用化流水线(例如,Persistent kernels with multiple具有多个Assigning warps to separate roles/stages such as将线程束分配到不同角色/阶段,例如
three-stage memory loader,三阶段内存加载器,distinct warp roles (3 in our case)不同的线程束角色(本例中为 3 个)memory load, compute, and memory store内存加载、计算与内存存储
compute, memory storer)计算、内存存储)

内核内流水线技术

Intra-Kernel Pipelining Techniques

Intra-kernel pipelining refers to a set of techniques that overlap memory operations and computations within a single kernel execution. (In the next chapter, we’ll explore inter-kernel pipelining, which overlaps work across multiple kernels running in different streams.)内核内流水线(intra-kernel pipelining)是指在单次内核执行内重叠内存操作与计算的一组技术。(下一章我们将探讨内核间流水线(inter-kernel pipelining),它在不同流中运行的多个内核之间重叠工作。)

The core idea is to structure a kernel into concurrent stages such that while one piece of data is being loaded or stored, previously loaded data is being processed. These stages operate in parallel over different tiles or data chunks. This improves throughput and efficiently hides latency.核心思想是将内核组织为多个并发阶段,使得在加载或存储某一块数据的同时,对先前已加载的数据进行处理。这些阶段在不同的瓦片(tile)或数据块上并行运行。这提升了吞吐量并高效地隐藏了时延。

Traditionally, GPUs rely on warp-level multithreading to hide latency. While one warp stalls on a memory load, other warps proceed with computation. This is the foundation of single instruction, multiple threads (SIMT) latency hiding in the execution model.传统上,GPU 依赖线程束级的多线程来隐藏时延。当一个线程束因内存加载而停顿时,其他线程束继续执行计算。这是执行模型中单指令多线程(SIMT)时延隐藏的基础。

Intra-kernel pipelining pushes this further by overlapping memory and compute within the same warp or kernel. It uses fine-grained coordination to stagger memory loads and compute—sometimes within a single warp.内核内流水线进一步推进了这一点,它在同一个线程束或内核内重叠内存操作与计算。它使用细粒度协调来错开内存加载与计算——有时是在单个线程束内。

Intra-kernel pipelining with the CUDA Pipeline API overlaps asynchronous memory transfers and computation without any __syncthreads(). The two common approaches to intra-kernel pipelining are double buffering and warp specialization.使用 CUDA Pipeline API 的内核内流水线无需任何 __syncthreads() 即可重叠异步内存传输与计算。内核内流水线的两种常见方法是双缓冲与线程束专用化。

In the double-buffered (two stages) pipeline approach, all threads cooperate uniformly. In the warp-specialized pipeline approach, warps are specialized into distinct roles like memory loader, compute, and memory storer. The choice depends on your workload and performance requirements. Table 10-1 summarizes these two <cuda/ pipeline> variants.在双缓冲(两阶段)流水线方法中,所有线程以统一方式协作。在线程束专用化流水线方法中,线程束被专用化为不同角色,如内存加载器、计算和内存存储。选择取决于你的工作负载与性能需求。表 10-1 总结了这两种 <cuda/ pipeline> 变体。

图 10-1
图 10-1使用 CUDA Pipeline API 的两阶段生产者-消费者模式Two-stage producer-consumer pattern with the CUDA Pipeline API

使用 CUDA Pipeline API 的协作式分块与双缓冲

Cooperative Tiling and Double-Buffering with the CUDA Pipeline API

You can implement the traditional double-buffered tiling pattern using the C++ Pipeline API by instantiating a two-stage pipeline to overlap memory loads and computations. Specifically, you can declare a two-stage cuda::pipeline_shared_state <cuda::thread_scope_block, 2> object, which is scoped to a specific thread block using cooperative groups (discussed in a bit). This is essentially a producer-consumer pattern, as shown in Figure 10-1.你可以使用 C++ Pipeline API 实现传统的双缓冲分块模式,方法是实例化一个两阶段流水线以重叠内存加载与计算。具体而言,你可以声明一个两阶段的 cuda::pipeline_shared_state <cuda::thread_scope_block, 2> 对象,它通过协作组(稍后讨论)限定到特定的线程块。这本质上是一种生产者-消费者模式,如图 10-1 所示。

The key CUDA Pipeline API calls are shown next. They are followed by an implementation of a double-buffered, cooperative tiling kernel using this API to demonstrate modern CUDA techniques that align with hardware features:接下来展示关键的 CUDA Pipeline API 调用。随后是使用该 API 实现的双缓冲协作分块内核,用以演示与硬件特性相契合的现代 CUDA 技术:

Text
pipe.producer_acquire()

Reserves the next pipeline stage for writing为写入预留下一个流水线阶段

Text
pipe.producer_commit()

Signals that the previously issued asynchronous operations for this stage are ready for consumption发出信号,表示此前为该阶段发出的异步操作已可供消费

Text
pipe.consumer_wait()

Waits until the previously committed operations for this stage complete to avoid race conditions in loops等待此前为该阶段提交的操作完成,以避免循环中的竞态条件

Text
pipe.consumer_release()

Releases the current stage so it can be reused释放当前阶段,使其可被复用

The two stages in the pipeline overlap global-memory loads with computations. In the first, Stage 0, one warp in the thread block issues an asynchronous prefetch for the next tile into shared memory. The prefetch issues cooperative cuda::mem cpy_async copies that lower to per thread cp.async into shared memory.流水线中的两个阶段将全局内存加载与计算重叠起来。在第一个阶段 Stage 0 中,线程块内的一个线程束为下一个分块(tile)向共享内存发起异步预取。该预取发起协作式的 cuda::mem cpy_async 拷贝,这些拷贝会下沉为每个线程的 cp.async 写入共享内存。

图 10-2
图 10-2使用生产者—消费者流水线,以(C)计算隐藏全局 DRAM 加载(L)的时延Hiding global DRAM load (L) latency with (C) compute using a producer– consumer pipeline

While Stage 0 is producing (loading) the data in one warp, the remaining warps in the thread block are consuming (computing) the loaded data in the second stage, 1. This simple producer-consumer implementation hides DRAM latency with ongoing computations, as shown in Figure 10-2.当 Stage 0 在一个线程束中生产(加载)数据时,线程块中其余的线程束在第二个阶段 Stage 1 中消费(计算)已加载的数据。这一简单的生产者—消费者实现用持续进行的计算隐藏 DRAM 时延,如图 10-2 所示。

This pattern can raise SM utilization and reduce kernel time, but whether you are compute bound or memory bound depends on the operation, tile sizes, and overlap efficiency. Even highly optimized attention kernels like FlashAttention-3 report around 75% percent of peak FP16 FLOPs due to practical limits in overlap and data movement.这种模式可以提升 SM 利用率并缩短内核时间,但究竟处于计算受限还是显存受限,取决于具体操作、分块大小和重叠效率。即便是 FlashAttention-3 这类高度优化的注意力内核,由于重叠与数据搬运的实际限制,报告的 FP16 FLOPs 也仅达到峰值的约 75%。

This is a two-stage, double-buffering example using the CUDA C++ Pipeline API. This API enables the fine-grained producer-consumer synchronization code used here:这是一个使用 CUDA C++ Pipeline API 的两阶段双缓冲示例。该 API 支持此处所用的细粒度生产者—消费者同步代码:

Python
#include <cuda/pipeline>
#include <cooperative_groups.h>
#include <algorithm>
namespace cg = cooperative_groups;
#ifndef TILE_SIZE
#define TILE_SIZE 32
#endif
#ifndef STAGES
#define STAGES 2          // 2 = double buffer, 3 = triple buffer, etc.
#endif
__device__  float computeTile(const float* __restrict__ A_sub,
                              const float* __restrict__ B_sub,
                              int tx, int ty) {
Text
    float s = 0.0f;
    #pragma unroll
    for (int k = 0; k < TILE_SIZE; ++k) {
        s += A_sub[ty * TILE_SIZE + k] * B_sub[k * TILE_SIZE + tx];
    }
    return s;
}
extern "C" __global__
void gemm_tiled_pipeline(const float* __restrict__ A_global, // [M x K]
                         const float* __restrict__ B_global, // [K x N]
                         float* __restrict__ C_global,       // [M x N]
                         int M, int N, int K) {
    cg::thread_block cta = cg::this_thread_block();
    // 共享内存布局:先 A[STAGES],再 B[STAGES]
    extern __shared__ float shared_mem[];
    float* A_buf[STAGES];
    float* B_buf[STAGES];
    {
        float* p = shared_mem;
        for (int s = 0; s < STAGES; ++s) {
            A_buf[s] = p;
            p += TILE_SIZE * TILE_SIZE;
       }
        for (int s = 0; s < STAGES; ++s) {
            B_buf[s] = p;
            p += TILE_SIZE * TILE_SIZE;
       }
    }
    __shared__ cuda::pipeline_shared_state<cuda::thread_scope_block, STAGES> state;
    auto pipe = cuda::make_pipeline(cta, &state);
    int tx = threadIdx.x, ty = threadIdx.y;
    int block_row = blockIdx.y * TILE_SIZE;
    int block_col = blockIdx.x * TILE_SIZE;
    float accum = 0.0f;
    int numTiles = (K + TILE_SIZE - 1) / TILE_SIZE;
    // 序言(Prologue):加载前 STAGES 个分块(若数量不足则更少)
    for (int s = 0; s < std::min(STAGES, numTiles); ++s) {
        int aRow = block_row + ty;
        int aCol = s * TILE_SIZE + tx;
        int bRow = s * TILE_SIZE + ty;
        int bCol = block_col + tx;
        pipe.producer_acquire();
        if (aRow < M && aCol < K) {
            cuda::memcpy_async(cta, A_buf[s] + ty*TILE_SIZE + tx,
                               &A_global[aRow*K + aCol],
                               cuda::aligned_size_t<32>(sizeof(float)), pipe);
        } else {
            A_buf[s][ty*TILE_SIZE + tx] = 0.0f;
        }
    float s = 0.0f;
    #pragma unroll
    for (int k = 0; k < TILE_SIZE; ++k) {
        s += A_sub[ty * TILE_SIZE + k] * B_sub[k * TILE_SIZE + tx];
    }
    return s;
}
extern "C" __global__
void gemm_tiled_pipeline(const float* __restrict__ A_global, // [M x K]
                         const float* __restrict__ B_global, // [K x N]
                         float* __restrict__ C_global,       // [M x N]
                         int M, int N, int K) {
    cg::thread_block cta = cg::this_thread_block();
    // Shared memory layout: A[STAGES] then B[STAGES]
    extern __shared__ float shared_mem[];
    float* A_buf[STAGES];
    float* B_buf[STAGES];
    {
        float* p = shared_mem;
        for (int s = 0; s < STAGES; ++s) {
            A_buf[s] = p;
            p += TILE_SIZE * TILE_SIZE;
       }
        for (int s = 0; s < STAGES; ++s) {
            B_buf[s] = p;
            p += TILE_SIZE * TILE_SIZE;
       }
    }
    __shared__ cuda::pipeline_shared_state<cuda::thread_scope_block, STAGES> state;
    auto pipe = cuda::make_pipeline(cta, &state);
    int tx = threadIdx.x, ty = threadIdx.y;
    int block_row = blockIdx.y * TILE_SIZE;
    int block_col = blockIdx.x * TILE_SIZE;
    float accum = 0.0f;
    int numTiles = (K + TILE_SIZE - 1) / TILE_SIZE;
    // Prologue: load first STAGES tiles (or fewer if short)
    for (int s = 0; s < std::min(STAGES, numTiles); ++s) {
        int aRow = block_row + ty;
        int aCol = s * TILE_SIZE + tx;
        int bRow = s * TILE_SIZE + ty;
        int bCol = block_col + tx;
        pipe.producer_acquire();
        if (aRow < M && aCol < K) {
            cuda::memcpy_async(cta, A_buf[s] + ty*TILE_SIZE + tx,
                               &A_global[aRow*K + aCol],
                               cuda::aligned_size_t<32>(sizeof(float)), pipe);
        } else {
            A_buf[s][ty*TILE_SIZE + tx] = 0.0f;
        }
Text
        if (bRow < K && bCol < N) {
            cuda::memcpy_async(cta, B_buf[s] + ty*TILE_SIZE + tx,
                               &B_global[bRow*N + bCol],
                               cuda::aligned_size_t<32>(sizeof(float)), pipe);
        } else {
            B_buf[s][ty*TILE_SIZE + tx] = 0.0f;
        }
        pipe.producer_commit();
    }
    // 稳态
    for (int tile = 0; tile < numTiles; ++tile) {
        int s = tile % STAGES;
        // 块作用域等待
        pipe.consumer_wait();
        accum += computeTile(A_buf[s], B_buf[s], tx, ty);
        pipe.consumer_release();
        // 将下一个分块预取到同一槽位 s(环形缓冲区)
        int nextTile = tile + STAGES;
        if (nextTile < numTiles) {
            int aRow = block_row + ty;
            int aCol = nextTile * TILE_SIZE + tx;
            int bRow = nextTile * TILE_SIZE + ty;
            int bCol = block_col + tx;
            pipe.producer_acquire();
            if (aRow < M && aCol < K) {
                cuda::memcpy_async(cta, A_buf[s] + ty*TILE_SIZE + tx,
                                   &A_global[aRow*K + aCol],
                                   cuda::aligned_size_t<32>{sizeof(float)}, pipe);
            } else {
                A_buf[s][ty*TILE_SIZE + tx] = 0.0f;
            }
            if (bRow < K && bCol < N) {
                cuda::memcpy_async(cta, B_buf[s] + ty*TILE_SIZE + tx,
                                   &B_global[bRow*N + bCol],
                                   cuda::aligned_size_t<32>{sizeof(float)}, pipe);
            } else {
                B_buf[s][ty*TILE_SIZE + tx] = 0.0f;
            }
            pipe.producer_commit();
        }
    }
    // 收尾(Epilogue):最终写回(保护尾部)
    int cRow = block_row + ty;
    int cCol = block_col + tx;
    if (cRow < M && cCol < N) {
        C_global[cRow * N + cCol] = accum;
    }
}
        if (bRow < K && bCol < N) {
            cuda::memcpy_async(cta, B_buf[s] + ty*TILE_SIZE + tx,
                               &B_global[bRow*N + bCol],
                               cuda::aligned_size_t<32>(sizeof(float)), pipe);
        } else {
            B_buf[s][ty*TILE_SIZE + tx] = 0.0f;
        }
        pipe.producer_commit();
    }
    // Steady state
    for (int tile = 0; tile < numTiles; ++tile) {
        int s = tile % STAGES;
        // Block-scope wait
        pipe.consumer_wait();
        accum += computeTile(A_buf[s], B_buf[s], tx, ty);
        pipe.consumer_release();
        // Prefetch next tile into the same slot s (ring buffer)
        int nextTile = tile + STAGES;
        if (nextTile < numTiles) {
            int aRow = block_row + ty;
            int aCol = nextTile * TILE_SIZE + tx;
            int bRow = nextTile * TILE_SIZE + ty;
            int bCol = block_col + tx;
            pipe.producer_acquire();
            if (aRow < M && aCol < K) {
                cuda::memcpy_async(cta, A_buf[s] + ty*TILE_SIZE + tx,
                                   &A_global[aRow*K + aCol],
                                   cuda::aligned_size_t<32>{sizeof(float)}, pipe);
            } else {
                A_buf[s][ty*TILE_SIZE + tx] = 0.0f;
            }
            if (bRow < K && bCol < N) {
                cuda::memcpy_async(cta, B_buf[s] + ty*TILE_SIZE + tx,
                                   &B_global[bRow*N + bCol],
                                   cuda::aligned_size_t<32>{sizeof(float)}, pipe);
            } else {
                B_buf[s][ty*TILE_SIZE + tx] = 0.0f;
            }
            pipe.producer_commit();
        }
    }
    // Epilogue: final store (guard tails)
    int cRow = block_row + ty;
    int cCol = block_col + tx;
    if (cRow < M && cCol < N) {
        C_global[cRow * N + cCol] = accum;
    }
}

The code first retrieves a handle to the current thread block using cooperative groups (CG), discussed in a bit. It then immediately instantiates a two-stage cuda::pipeline object bound to that block. By creating the pipeline before any asynchronous operations, the pipelines’ internal barriers and internal synchronization mechanisms are in place prior to performing data movement.代码首先使用协作组(CG)获取当前线程块的句柄,这一点稍后会讨论。随后它立即实例化一个绑定到该线程块的两阶段 cuda::pipeline 对象。由于流水线在任何异步操作之前就已创建,其内部屏障与内部同步机制在执行数据搬运之前便已就位。

The kernel then allocates one contiguous shared-memory region for both A and B tiles by defining an extern __shared__ float shared_mem[]. It divides this buffer into four subregions, two for A and two for B, using pointer arithmetic (float* A_buf[2] and float* B_buf[2]). This allows true double-buffering without extra dynamic allocations.随后,内核通过定义 extern __shared__ float shared_mem[],为 A 和 B 两个分块分配一块连续的共享内存区域。它使用指针运算(float* A_buf[2] 和 float* B_buf[2])将该缓冲区划分为四个子区域,两个用于 A,两个用于 B。这样无需额外的动态分配即可实现真正的双缓冲。

Before entering the main loop, the kernel asynchronously prefetches the first STAGES tiles using asynchronous copies bound to the following pipeline: pro ducer_acquire() → memcpy_async() → producer_commit(). Consumer warps use consumer_wait() and consumer_release() so the compute starts exactly when the prefetched tile is ready.在进入主循环之前,内核使用绑定到下述流水线的异步拷贝,异步预取最初的 STAGES 个分块:pro ducer_acquire() → memcpy_async() → producer_commit()。消费者线程束使用 consumer_wait() 和 consumer_release(),使计算恰好在预取的分块就绪时开始。

This initial barrier replaces a future need for __syncthreads() and ensures the pipeline’s stage 0 and stage 1 buffers are correctly populated for the first iteration of the producer-consumer sequence.这一初始屏障替代了后续对 __syncthreads() 的需求,并确保流水线的 stage 0 与 stage 1 缓冲区在生产—消费者序列的第一次迭代时已正确填充。

Within each iteration, the kernel reserves the next buffer to load subsequent tiles by calling pipe.producer_acquire(). It then launches two cuda::memcpy_async operations—one for the A tile and one for the B tile. Each load is bound to the pipeline object with cuda::memcpy_async(..., pipe), which issues asynchronous copies from global memory into shared memory.在每次迭代中,内核通过调用 pipe.producer_acquire() 预留下一个缓冲区,用于加载后续分块。随后它发起两个 cuda::memcpy_async 操作——一个用于 A 分块,一个用于 B 分块。每次加载都通过 cuda::memcpy_async(..., pipe) 绑定到流水线对象,该调用会发起从全局内存到共享内存的异步拷贝。

These asynchronous copies can overlap well with compute when accesses are coalesced and tiled correctly. This way, the pipeline stays fed with data to perform maximum useful work. Immediately after queuing the memcpy_async calls, the kernel signals completion with pipe.producer_commit(). This commit records the copy’s arrival and allows consumer warps to wait on that specific stage without blocking the entire thread block.当访问被正确合并与分块时,这些异步拷贝可以与计算良好重叠。这样一来,流水线始终有数据供给,从而完成尽可能多的有效工作。在将 memcpy_async 调用入队之后,内核立即通过 pipe.producer_commit() 发出完成信号。这次提交记录了拷贝的到达,并允许消费者线程束等待该特定阶段,而不会阻塞整个线程块。

Concurrently, other warps in the block invoke pipe.consumer_wait(). This efficiently stalls only those threads dependent on the data in the curr buffer until the producer has committed it. Once the wait completes, each thread calls the device function computeTile(...) to perform its TILE_SIZE x TILE_SIZE dot-product computation. The results are incrementally accumulated in registers.与此同时,块中的其他线程束调用 pipe.consumer_wait()。这只会高效地阻塞那些依赖 curr 缓冲区数据的线程,直到生产者完成提交。等待完成后,每个线程调用设备函数 computeTile(...),执行其 TILE_SIZE x TILE_SIZE 的点积计算。结果在寄存器中逐步累加。

After finishing the computation on the current buffer, the warps invoke pipe.con sumer_release() to free that stage for reuse in subsequent iterations. This fine-grained release prevents block-wide stalls and maximizes overlap between compute and memory transfer phases.完成对当前缓冲区的计算后,线程束调用 pipe.con sumer_release() 释放该阶段,以便在后续迭代中复用。这种细粒度的释放避免了整个块的停顿,并最大化计算与内存传输阶段之间的重叠。

表 10-2使用 CUDA Pipeline API 的朴素分块与双缓冲内核的性能对比Performance comparison between the naive tiling and double-buffered kernel using the CUDA Pipeline API
Metric指标Naive tiling朴素分块Two-stage, double-buffered pipeline (double_buffered_pipeline)两阶段双缓冲流水线(double_buffered_pipeline)
Kernel execution time内核执行时间41.3 ms41.3 ms20.5 ms (~2× faster versus naive)20.5 ms(相比朴素实现约快 2×)
SM Active %SM 活跃占比 %68%68%92% (+24% versus naive)92%(比朴素实现 +24%)

At the end of each iteration, swap the curr and next, which cycles the two double-buffer stages. This way, the memory buffer that was just computed now becomes the target for the next asynchronous memory.在每次迭代结束时,交换 curr 与 next,从而在两个双缓冲阶段之间循环。这样,刚刚完成计算的内存缓冲区现在成为下一次异步内存传输的目标。

These producer and consumer stages are repeated for all K-dimension tiles. After processing the final tile, each thread’s accumulator contains the full TILE_SIZE x TILE_SIZE dot-product result. This result is then written back to the appropriate element of C_global.这些生产者阶段和消费者阶段会对所有 K 维度分块重复执行。处理完最后一个分块后,每个线程的累加器中保存了完整的 TILE_SIZE x TILE_SIZE 点积结果。随后该结果被写回 C_global 中对应的元素。

When using double buffering, it’s recommended to make sure that each asynchronous copy (memcpy_async) is followed by the appropriate producer_commit() and consumer_wait() calls for synchronization. This guarantees that the compute kernels will use the data only after it’s properly loaded. Modern CUDA compilers will often perform these optimizations automatically for simple loops.使用双缓冲时,建议确保每次异步拷贝(memcpy_async)之后都跟随相应的 producer_commit() 和 consumer_wait() 调用以完成同步。这保证了计算内核只会在数据被正确加载之后才使用这些数据。现代 CUDA 编译器通常会对简单循环自动执行这些优化。

This simple double-buffered scheme is roughly 2× faster than naive tiling, as described in Table 10-2’s performance comparison between the naive tiling implementation and the optimized, double-buffered, pipeline implementation.这种简单的双缓冲方案比朴素分块大约快 2 倍,如表 10-2 中朴素分块实现与优化的双缓冲流水线实现之间的性能对比所示。

In this experiment, the gemm_tiled_pipeline kernel using the CUDA C++ Pipeline API achieves a 2× speedup over the naive tiling version. By using the fine-grained pipe.producer_commit() and pipe.consumer_wait() primitives, the pipeline remains filled, and SM Active % jumps 24% from 68% to 92%.在该实验中,使用 CUDA C++ Pipeline API 的 gemm_tiled_pipeline 内核相比朴素分块版本实现了 2 倍加速。通过使用细粒度的 pipe.producer_commit() 和 pipe.consumer_wait() 原语,流水线始终保持填满,SM Active % 提升了 24%,从 68% 升至 92%。

图 10-3
图 10-3非线程束专用化内核,每个线程束同时执行数据加载与计算的混合任务(来源:https://oreil.ly/WZDbM)Nonwarp specialized kernel with each warp performing a mix of both data loading and compute (source: https://oreil.ly/WZDbM)

线程束专用化与生产者-消费者模型

Warp Specialization and the Producer-Consumer Model

Warp specialization extends double buffering by assigning operations to warps that use different hardware, such as data movement (e.g., TMA) and compute (e.g., Tensor Cores). This is in contrast to reusing the same warps for both loading data and computing, as shown in Figure 10-3.线程束专用化扩展了双缓冲,它把操作分配给使用不同硬件的线程束,例如数据搬运(如 TMA)和计算(如张量核心)。这与图 10-3 所示的复用同一批线程束同时进行数据加载和计算的做法形成对比。

This type of specialization allows each set of warps to have their own instruction sequences. As such, instructions are issued and executed continuously without being interrupted by other types of operations. Specifically, warp specialization lets you assign one set of “producer” or “memory” warps to prefetch tiles asynchronously using cuda::memcpy_async. Then all other “consumer” or “compute” warps perform the computations, as shown in Figure 10-4.这种专用化允许每组线程束拥有各自的指令序列。因此,指令可以连续发射和执行,不会被其他类型的操作打断。具体而言,线程束专用化让你可以把一组“生产者”或“内存”线程束指定为使用 cuda::memcpy_async 异步预取分块,然后由所有其他“消费者”或“计算”线程束执行计算,如图 10-4 所示。

Here, four warps are assigned the producer role, while the remaining eight warps are assigned the consumer role. Like most producer-consumer patterns, you can assign a different number of warps for the producer and consumer.这里,四个线程束被分配生产者角色,其余八个线程束被分配消费者角色。与大多数生产者-消费者模式一样,你可以为生产者和消费者分配不同数量的线程束。

Because each warp has its own scheduler, the GPU can issue a load instruction, a math instruction, and a write instruction—all in the same cycle from different warps using different warp schedulers. So an SM with multiple schedulers can issue a memory instruction from Warp 0, a math instruction from Warp 1, and so on, in one cycle.由于每个线程束都有自己的调度器,GPU 可以在同一周期内从使用不同线程束调度器的不同线程束发射一条加载指令、一条数学指令和一条写指令。因此,拥有多个调度器的 SM 可以在一个周期内发射来自 Warp 0 的内存指令、来自 Warp 1 的数学指令,依此类推。

图 10-4
图 10-4线程束专用化内核,其中一组线程束负责加载数据,所有其他线程束负责计算(来源:https://oreil.ly/WZDbM)Warp-specialized kernel with one set of warps for loading data and all other warps for computations (source: https://oreil.ly/WZDbM)

This effectively creates a thread-block-level multi-issue scenario across warps. This is not possible with single-warp double buffering because a single warp’s scheduler can issue only one instruction per cycle.这在跨线程束的层面有效地创造了线程块级的多发射场景。单线程束双缓冲无法做到这一点,因为单个线程束的调度器每周期只能发射一条指令。

An interesting pattern for warp specialization is using three different types of warps, such as “loader,” “compute,” and “storer” warps. The loader warp pushes tiles into the pipeline’s queue. The compute warp runs the compute kernel on each tile. And the storer warp writes out the results, as shown in Figure 10-5.线程束专用化的一种有趣模式是使用三种不同类型的线程束,例如“加载器”(loader)、“计算”(compute)和“存储”(storer)线程束。加载器线程束将瓦片推入流水线的队列。计算线程束对每个瓦片运行计算内核。存储线程束则写出结果,如图 10-5 所示。

This warp-specialized pipeline squeezes out the idle cycles that a single-warp, double-buffered, sequential load-and-compute loop cannot address. Warp specialization’s efficient overlap of data transfer and computation increases GPU utilization—especially for long-running loops and persistent kernels. In these cases, the overhead of role coordination and data handoff is amortized over many iterations.这种线程束专用化的流水线能够挤出单线程束、双缓冲、顺序“加载—计算”循环无法消除的空闲周期。线程束专用化对数据传输与计算的高效重叠提升了 GPU 利用率——尤其对于长时间运行的循环和持久化内核。在这些情况下,角色协调与数据交接的开销会在多次迭代中被摊薄。

A paper on warp scheduling demonstrated that warp specialization can achieve nearly perfect overlap of memory and compute. In this case, the GPU kernel had distinct memory and compute phases such that memory and compute took turns being the bottleneck. By applying warp specialization, their workload transformed into a state in which both the SM’s memory subsystem and compute units were simultaneously busy almost the entire time.一篇关于线程束调度的论文表明,线程束专用化可以实现内存与计算近乎完美的重叠。在该案例中,GPU 内核具有明显的内存阶段和计算阶段,内存与计算轮流成为瓶颈。通过应用线程束专用化,他们的工作负载转变为 SM 的内存子系统和计算单元几乎全程同时忙碌的状态。

图 10-5
图 10-5三角色线程束专用化流水线配置:一组线程束负责加载数据,另一组负责计算,还有一组负责数据存储(来源:https://oreil.ly/xs7YN)Three-role warp-specialized pipeline configuration with one set of warps for loading data, another set for compute, and another set for data storing (source: https:// oreil.ly/xs7YN)

The profiling showed that before using warp specialization, their L2 bandwidth utilization and Tensor Core utilization were out of phase. After warp specialization, L2 bandwidth and Tensor Core utilization became in-phase. This resulted in much higher effective throughput and showed that warp specialization can squeeze out the last bits of idle time—even for a well-tuned asynchronous pipeline that might be left on the table.性能剖析显示,在使用线程束专用化之前,他们的 L2 带宽利用率与张量核心利用率是不同相的。使用线程束专用化之后,L2 带宽与张量核心利用率变为同相。这带来了高得多的有效吞吐,并表明线程束专用化能够挤出最后一点空闲时间——即便是对于一条调优良好、但可能仍有性能留在桌面上的异步流水线也是如此。

Another warp specialization pattern is a modification of the three-role warp-specialized pipeline. It assigns a set of warps to the memory loader as before, but then uses two sets of consumer warps that “ping-pong” between the roles of compute and memory-storer. This three-role warp-specialization architecture is exposed in CUT-LASS as gemm_tma_warpspecialized_pingpong and shown in Figure 10-6.另一种线程束专用化模式是对三角色线程束专用化流水线的改造。它和之前一样为一组线程束分配内存加载任务,但随后使用两组消费者线程束,在计算与内存存储两种角色之间“乒乓”切换。这种三角色线程束专用化架构在 CUTLASS 中以 gemm_tma_warpspecialized_pingpong 的形式提供,如图 10-6 所示。

图 10-6
图 10-6采用三角色线程束专用化内核的乒乓架构(来源:https://oreil.ly/xs7YN)Ping-pong architecture with three-role warp-specialized kernel (source: https://oreil.ly/xs7YN)

Here, the consumer MMAs are overlapping and include a small amount of post-MMA wrap-up, or epilogue, processing. This per-MMA epilogue cleanup is required before launching the next MMA. Specifically, the epilogue can include accumulating, scaling, writing back to global memory, or shuffling results to another warp. Additionally, the epilogue can perform housekeeping like advancing tile pointers, updating loop counters, and signaling that this tile is done so the next TMA load or MMA can kick off.在这里,消费者 MMA 相互重叠,并包含少量 MMA 之后的收尾处理,即*尾声*(epilogue)。每次 MMA 之后的尾声清理是启动下一次 MMA 之前所必需的。具体而言,尾声可以包括累加、缩放、写回全局内存,或将结果混洗给另一个线程束。此外,尾声还可以执行一些例行维护工作,例如推进瓦片指针、更新循环计数器,以及发出该瓦片已完成的信号,以便下一次 TMA 加载或 MMA 可以启动。

In practice, warp specialization is extremely effective at squeezing out the last bits of performance. In fact, FlashAttention v3 attributes its speedups partially to warp-specialized pipelines that overlap GEMM and softmax computations—along with data transfers—to keep all hardware units busy. This helps achieve near-peak FLOPS for attention computations due to aggressive overlap of compute and TMA-driven data movement.在实践中,线程束专用化在挤出最后一点性能方面极其有效。事实上,FlashAttention v3 将其加速部分归功于线程束专用化的流水线——它重叠 GEMM 与 softmax 计算以及数据传输,使所有硬件单元保持忙碌。由于计算与 TMA 驱动的数据移动实现了激进的重叠,这有助于让注意力计算达到接近峰值的 FLOPS。

In addition, the PyTorch compiler (covered in Chapters 13 and 14) generates kernels that use warp specialization to schedule separate warps for loading and computing data. It also uses low-cost barriers to synchronize the warps, similar to the CUDA Pipeline API implementation detailed in the next section. The PyTorch compiler system also integrates with CUTLASS’s ping-pong GEMM. Both torch.compile and Triton may generate warp specialized kernels for supported operations. However, they apply warp specialization selectively based on heuristics and do not enable warp specialization for every operator.此外,PyTorch 编译器(第 13 章和第 14 章介绍)生成的内核会使用线程束专用化,为加载数据和计算数据调度各自独立的线程束。它还使用低开销的屏障来同步这些线程束,类似于下一节详述的 CUDA Pipeline API 实现。PyTorch 编译器系统还与 CUTLASS 的 ping-pong GEMM 集成。torch.compile 和 Triton 都可能为受支持的操作生成线程束专用化的内核。不过,它们会基于启发式规则有选择地应用线程束专用化,并不会为每个算子都启用线程束专用化。

使用 CUDA Pipeline API 实现线程束专用化

Using CUDA Pipeline API for Warp Specialization

Warp specialization builds on the CUDA Pipeline API by allowing specialized warps to communicate using fine-grained producer and consumer primitives. These calls avoid full block barriers while composing naturally with asynchronous copies such as cuda::memcpy_async.线程束专用化构建在 CUDA Pipeline API 之上,允许专用化的线程束使用细粒度的生产者与消费者原语进行通信。这些调用避免了全块屏障,同时能与 cuda::memcpy_async 等异步拷贝自然地组合使用。

The key advantage of using the CUDA Pipeline API’s producer and consumer calls (e.g., pipe.producer_acquire(), pipe.producer_commit(), pipe.consumer_wait(), and pipe.consumer_release()) is that they synchronize only the specific warps or stages that actually need to hand off data. This is in contrast to forcing every thread in a block to wait.使用 CUDA Pipeline API 的生产者与消费者调用(例如 pipe.producer_acquire()、pipe.producer_commit()、pipe.consumer_wait() 和 pipe.consumer_release())的关键优势在于,它们只同步那些真正需要交接数据的特定线程束或阶段。这与强制块内每个线程都等待形成对比。

A block-wide barrier would stall every warp—even those that are not involved with the producer-consumer pipeline. All execution in that block must pause until every thread reaches the barrier, as shown in Figure 10-7.全块屏障会让每个线程束都停顿——即使是那些未参与生产者-消费者流水线的线程束。该块内的所有执行都必须暂停,直到每个线程都到达屏障,如图 10-7 所示。

By comparison, the Pipeline API maintains per-stage state internally. When a producer warp finishes its asynchronous copy and calls pipe.producer_commit, only the warps that call pipe.consumer_wait will block until the data is ready. Other warps in the block can continue running any work that does not depend on that stage. In practice, the CUDA Pipeline API reduces idle time and decreases stalled warps because it eliminates the need to pause the entire block with a barrier. With pipelines you coordinate producer and consumer handoffs at a finer granularity than a hand-coded async-copy sequence (e.g., PTX cp.async + __syncthreads()).相比之下,Pipeline API 在内部维护每个阶段的状态。当生产者线程束完成其异步拷贝并调用 pipe.producer_commit 时,只有调用 pipe.consumer_wait 的线程束会阻塞,直到数据就绪。块内的其他线程束可以继续执行任何不依赖该阶段的工作。在实践中,CUDA Pipeline API 减少了空闲时间并减少了停顿的线程束,因为它消除了用屏障暂停整个块的需要。借助流水线,你可以以比手写 async-copy 序列(例如 PTX cp.async + __syncthreads())更细的粒度协调生产者与消费者的交接。

You can implement warp specialization with the CUDA Pipeline API in a three-role pattern. A loader warp produces inputs for a compute warp, the compute warp consumes those inputs and produces results, and a storer warp consumes those results and writes them out. The pipeline object is block scoped, and it tracks the stage order internally.你可以用 CUDA Pipeline API 以三种角色的模式实现线程束专用化。加载线程束为计算线程束生产输入,计算线程束消费这些输入并生产结果,存储线程束消费这些结果并将其写出。流水线对象的作用域为块级,并在内部跟踪阶段顺序。

图 10-7
图 10-7全块屏障会阻止线程继续执行,直到它们完成同步并加载新数据Block-wide barrier prevents threads from proceeding until they synchronize and load the new data

Here is an example of a warp-specialized, three-role kernel that computes tiles of data using a loader warp (warp_id 0), compute warp (warp_id 1), and storer warp (warp_id 2):下面是一个线程束专用化的三角色内核示例,它使用加载线程束(warp_id 0)、计算线程束(warp_id 1)和存储线程束(warp_id 2)来计算数据分块:

CUDA / C
// 使用 CUDA Pipeline API 的线程束专用化角色
#include <cuda/pipeline>
#include <cooperative_groups.h>
namespace cg = cooperative_groups;
// smem_bytes = 3×TILE_SIZE^2×sizeof(float)(三个角色,单流水线)
// 3 个分块 * 112 * 112 * 每个 float 4 字节 =
// 150,528 字节 < 227,328 字节(约 227 KB)
// Blackwell 上每线程块的动态 SMEM 上限
// 选择 TILE_*,使每线程块的总共享内存 ≤ 设备的
// 所报告的每线程块共享内存上限。
// 查询 cudaDevAttrMaxSharedMemoryPerBlockOption 并设置
// 相应地设置 cudaFuncAttributeMaxDynamicSharedMemorySize。
#define TILE_SIZE 112
// 示例分块计算:对一个 TILE_SIZE × TILE_SIZE 的分块执行 C = A x B
__device__ void compute_full_tile(const float* __restrict__ A_tile,
                                  const float* __restrict__ B_tile,
                                  float* __restrict__ C_tile,
                                  int lane_id) {
    for (int idx = lane_id; idx < TILE_SIZE * TILE_SIZE; idx += warpSize) {
// warp specialized roles using the CUDA Pipeline API
#include <cuda/pipeline>
#include <cooperative_groups.h>
namespace cg = cooperative_groups;
// smem_bytes = 3×TILE_SIZE^2×sizeof(float) (three roles, single-pipeline)
// 3 tiles * 112 * 112 * 4 byte per float =
// 150,528 bytes < 227,328 bytes (~227 KB)
// per-block dynamic SMEM limit on Blackwell
// Choose TILE_* so total shared memory per block is ≤ the device’s
// reported per-block shared memory limit.
// Query cudaDevAttrMaxSharedMemoryPerBlockOption and set
// cudaFuncAttributeMaxDynamicSharedMemorySize accordingly.
#define TILE_SIZE 112
// Example tile compute: C = A x B for one TILE_SIZE by TILE_SIZE tile
__device__ void compute_full_tile(const float* __restrict__ A_tile,
                                  const float* __restrict__ B_tile,
                                  float* __restrict__ C_tile,
                                  int lane_id) {
    for (int idx = lane_id; idx < TILE_SIZE * TILE_SIZE; idx += warpSize) {
Text
        int row = idx / TILE_SIZE;
        int col = idx % TILE_SIZE;
        float acc = 0.0f;
        #pragma unroll
        for (int k = 0; k < TILE_SIZE; ++k) {
            acc += A_tile[row * TILE_SIZE + k] * B_tile[k * TILE_SIZE + col];
        }
        C_tile[idx] = acc;
    }
}
extern "C"
__global__ void warp_specialized_pipeline_kernel(
        const float* __restrict__ A_global,
        const float* __restrict__ B_global,
        float* __restrict__ C_global,
        int numTiles) {
    thread_block cta = this_thread_block();
    // 动态共享内存中的三个方形分块:A、B 和 C
    extern __shared__ float shared_mem[];
    float* A_tile = shared_mem;
    float* B_tile = A_tile + TILE_SIZE * TILE_SIZE;
    float* C_tile = B_tile + TILE_SIZE * TILE_SIZE;
    // 由该线程块共享的三级流水线
    __shared__ cuda::pipeline_shared_state<cuda::thread_scope_block, 3>
        pipe_state;
    auto pipe = cuda::make_pipeline(cta, &pipe_state);
    int warp_id  = threadIdx.x >> 5;
    int lane_id  = threadIdx.x & 31;
    int warps_per_block = blockDim.x >> 5;
    // 用于持久化分块的网格级线程束索引
    int totalWarps = gridDim.x * warps_per_block;
    int global_warp = warp_id + blockIdx.x * warps_per_block;
    for (int tile = global_warp; tile < numTiles; tile += totalWarps) {
        size_t offset = static_cast<size_t>(tile) * TILE_SIZE * TILE_SIZE;
        if (warp_id == 0) {
            // loader 生产 A_tile 和 B_tile
            pipe.producer_acquire();

            auto bytes = TILE_SIZE * TILE_SIZE * sizeof(float);
            cuda::memcpy_async(cta, A_tile, A_global + offset,
                               cuda::aligned_size_t<32>{bytes}, pipe);
            cuda::memcpy_async(cta, B_tile, B_global + offset,
                               cuda::aligned_size_t<32>{bytes}, pipe);
            pipe.producer_commit();
        }
        int row = idx / TILE_SIZE;
        int col = idx % TILE_SIZE;
        float acc = 0.0f;
        #pragma unroll
        for (int k = 0; k < TILE_SIZE; ++k) {
            acc += A_tile[row * TILE_SIZE + k] * B_tile[k * TILE_SIZE + col];
        }
        C_tile[idx] = acc;
    }
}
extern "C"
__global__ void warp_specialized_pipeline_kernel(
        const float* __restrict__ A_global,
        const float* __restrict__ B_global,
        float* __restrict__ C_global,
        int numTiles) {
    thread_block cta = this_thread_block();
    // three square tiles in dynamic shared memory: A, B, and C
    extern __shared__ float shared_mem[];
    float* A_tile = shared_mem;
    float* B_tile = A_tile + TILE_SIZE * TILE_SIZE;
    float* C_tile = B_tile + TILE_SIZE * TILE_SIZE;
    // three stage pipeline shared by the block
    __shared__ cuda::pipeline_shared_state<cuda::thread_scope_block, 3>
        pipe_state;
    auto pipe = cuda::make_pipeline(cta, &pipe_state);
    int warp_id  = threadIdx.x >> 5;
    int lane_id  = threadIdx.x & 31;
    int warps_per_block = blockDim.x >> 5;
    // grid wide warp indexing for persistent tiling
    int totalWarps = gridDim.x * warps_per_block;
    int global_warp = warp_id + blockIdx.x * warps_per_block;
    for (int tile = global_warp; tile < numTiles; tile += totalWarps) {
        size_t offset = static_cast<size_t>(tile) * TILE_SIZE * TILE_SIZE;
        if (warp_id == 0) {
            // loader produces A_tile and B_tile
            pipe.producer_acquire();

            auto bytes = TILE_SIZE * TILE_SIZE * sizeof(float);
            cuda::memcpy_async(cta, A_tile, A_global + offset,
                               cuda::aligned_size_t<32>{bytes}, pipe);
            cuda::memcpy_async(cta, B_tile, B_global + offset,
                               cuda::aligned_size_t<32>{bytes}, pipe);
            pipe.producer_commit();
        }
Text
        if (warp_id == 1) {
            // 等待 A_tile/B_tile
            pipe.consumer_wait();
            // 由 A_tile、B_tile 计算 C_tile
            compute_full_tile(A_tile, B_tile, C_tile, lane_id);
            // 已消费完 A/B;释放该流水级
            pipe.consumer_release();
            // 将 C_tile 发布给 storer 线程束
            pipe.producer_acquire();
            pipe.producer_commit();
        }
        if (warp_id == 2) {
            // store 消费已计算出的 C_tile
            // 消费前先等待已提交的流水级
            pipe.consumer_wait();
            for (int idx=lane_id; idx<TILE_SIZE*TILE_SIZE; idx+=warpSize) {
                C_global[offset + idx] = C_tile[idx];
            }
            pipe.consumer_release();
        }
    }
    // 启动时动态共享内存大小设为
    // 3 * TILE_SIZE * TILE_SIZE * sizeof(float)
    // 当以大于 48 KB 的动态共享内存启动时
    // 需要设置 cudaFuncAttributeMaxDynamicSharedMemorySize
    // 针对该内核
}
        if (warp_id == 1) {
            // wait for A_tile/B_tile
            pipe.consumer_wait();
            // compute C_tile from A_tile, B_tile
            compute_full_tile(A_tile, B_tile, C_tile, lane_id);
            // finished consuming A/B; free that stage
            pipe.consumer_release();
            // publish C_tile for the storer warp
            pipe.producer_acquire();
            pipe.producer_commit();
        }
        if (warp_id == 2) {
            // store consumes computed C_tile
            // Wait for the committed stage before consuming
            pipe.consumer_wait();
            for (int idx=lane_id; idx<TILE_SIZE*TILE_SIZE; idx+=warpSize) {
                C_global[offset + idx] = C_tile[idx];
            }
            pipe.consumer_release();
        }
    }
    // launch with dynamic shared memory size equal to
    // 3 * TILE_SIZE * TILE_SIZE * sizeof(float)
    // When launching with dynamic shared memory >48 KB
    // You need to set cudaFuncAttributeMaxDynamicSharedMemorySize
    // for the kernel
}

Here each warp works in a distinct role. The loader warp calls producer_acquire(), performs two cooperative copies with cuda::memcpy_async into A_tile and B_tile, then calls producer_commit(). The compute warp calls pipe.consumer_wait() to observe the newly committed data and immediately calls consumer_release() to free that stage for reuse.这里每个线程束承担不同的角色。加载线程束调用 producer_acquire(),用 cuda::memcpy_async 执行两次协作拷贝,分别写入 A_tile 和 B_tile,然后调用 producer_commit()。计算线程束调用 pipe.consumer_wait() 来观察新提交的数据,并立即调用 consumer_release() 以释放该阶段供复用。

The compute warp then becomes a producer for the next handoff by calling producer_acquire(), computing C_tile, and calling producer_commit(). The storer warp calls consumer_wait() to observe the computed C_tile and writes it to global memory in a warp-striped loop, then calls consumer_release(). This sequence uses a single block-scoped pipeline with no explicit stage numbers, and it avoids any block-wide __syncthreads.随后,计算线程束通过调用 producer_acquire()、计算 C_tile 并调用 producer_commit(),成为下一次交接的生产者。存储线程束调用 consumer_wait() 观察到已计算出的 C_tile,并在线程束条带化(warp-striped)循环中将其写入全局内存,然后调用 consumer_release()。该序列使用单个块作用域的流水线,不含显式的阶段编号,并且避免了任何块范围的 __syncthreads。

表 10-3三种实现的对比:朴素分块内核、两阶段双缓冲 GEMM 与线程束专用化流水线内核Comparison of three implementations: a naive tiled kernel, a two-stage double-buffered GEMM, and a warp-specialized pipeline kernel
Metric指标Naive tiling朴素分块Two‐stage, double‐buffered pipeline (double_buffered_pipeline)两阶段双缓冲流水线(double_buffered_pipeline)Warp‐specialized pipeline (warp_specialized_pipeline)线程束专用化流水线(warp_specialized_pipeline)
Kernel内核41.3 ms41.3 ms20.5 ms (2.01× faster versus naive)20.5 ms(比朴素实现快 2.01×)18.4 ms (10.2% speedup versus two‐stage)18.4 ms(比两阶段实现快 10.2%)
execution time执行时间
Warp execution线程束执行68%68%92% (+24% versus naive)92%(比朴素实现 +24%)96% (+4% versus two‐stage)96%(较两级提升 4%)
efficiency效率
Shared memory共享内存High高Low低Minimal极小
stall latency/停顿延迟/
warp sync stalls线程束同步停顿
L2 throughputL2 吞吐量80 GB/s80 GB/s155 GB/s (+94% versus naive)155 GB/s(较朴素实现提升 94%)165 GB/s (+6.45% versus two‐stage)165 GB/s(较两级提升 6.45%)
Throughput吞吐量Scales up可扩展Scales well up to ~6 warps per SM每个 SM 最多约 6 个线程束时扩展良好Scales nearly linearly to the SM’s warps (e.g.,几乎线性扩展至 SM 的线程束数(例如,
scalability可扩展性to only 2– 3 warps per SM仅能扩展至每个 SM 2–3 个线程束64 warps)64 个线程束)
DRAM bytesDRAM 字节数Poor差Great overlap重叠良好Excellent overlap重叠极佳
read versus SM读取与 SMoverlap重叠
cycles周期
Instruction指令1.7 B1.7 B1.05 B (–38% versus naive)1.05 B(较朴素实现 –38%)~1.00 B (–5% versus two‐stage)~1.00 B(较两阶段 –5%)
count计数

In short, using the CUDA Pipeline API together with cooperative groups allows fine-grained, SM-wide producer-consumer handoffs without any explicit __sync threads() calls. Table 10-3 compares three implementations: a naive tiled kernel, a two-stage double-buffered GEMM using double_buffered_pipeline, and our warp-specialized pipeline kernel warp_specialized_pipeline.简而言之,将 CUDA Pipeline API 与协作组结合使用,可以在无需任何显式 __sync threads() 调用的情况下,实现细粒度的、SM 范围的生产者—消费者交接。表 10-3 比较了三种实现:朴素分块内核、使用 double_buffered_pipeline 的两阶段双缓冲 GEMM,以及我们的线程束专用化流水线内核 warp_specialized_pipeline。

Here, the double-buffered pipeline finishes GEMM in 20.5 ms, whereas the warp-specialized version completes in just 18.4 ms. This 10.2% improvement is a result of the warp-specialized kernel only stalling the consumer warp (e.g., compute). The other warps (e.g., loader and storer) progress independently. In the previous two-stage, double-buffered kernel, every thread participates in the consumer phase. As such, consumer_wait() effectively stalls the entire thread block.在这里,双缓冲流水线在 20.5 ms 内完成 GEMM,而线程束专用化版本仅需 18.4 ms 即可完成。这 10.2% 的提升源于线程束专用化内核只会让消费者线程束(例如计算)停顿。其他线程束(例如加载器和存储器)则独立推进。在此前的两阶段双缓冲内核中,每个线程都参与消费者阶段。因此,consumer_wait() 实际上会让整个线程块停顿。

This finer-grained, per-warp synchronization eliminates the implicit full-block wait and allows all three warps (loader, compute, and storer) to overlap continuously. As a result, average SM utilization rises from roughly 92% in the double-buffered design to about 96% in the warp-specialized version—and warp-stall cycles drop to near-zero.这种更细粒度的、按线程束进行的同步消除了隐式的全块等待,使全部三个线程束(加载、计算和存储)能够持续重叠。因此,平均 SM 利用率从双缓冲设计中的约 92% 提升到线程束专用化版本中的约 96%——并且线程束停顿周期降至接近零。

From a scalability standpoint, Nsight Compute shows that the naive tiling kernel saturates after just two to three active warps per SM. This is because each tile load must complete before any computation can start.从可扩展性角度看,Nsight Compute 显示,朴素分块内核在每个 SM 仅有 2 到 3 个活跃线程束时就会饱和。这是因为每次分块加载都必须完成后才能开始任何计算。

The two-stage, double-buffered kernel improves on this by overlapping loads and compute. This implementation scales up to 6 warps per SM before shared-memory or register limits become an issue.两阶段双缓冲内核通过重叠加载与计算对此加以改进。该实现可扩展到每个 SM 6 个线程束,之后共享内存或寄存器限制才会成为问题。

In contrast, the warp_specialized_pipeline scales almost linearly as long as you can assign additional warps for load, compute, and store. On Blackwell, for instance, you can keep up to the architectural limit of 64 resident warps per SM. (Actual residency depends on registers, shared memory, and block size.)相比之下,只要能为加载、计算和存储分配更多线程束,warp_specialized_pipeline 几乎可以线性扩展。例如在 Blackwell 上,最多可维持每个 SM 64 个常驻线程束的架构上限。(实际常驻数量取决于寄存器、共享内存和线程块大小。)

As Table 10-3 shows, both the double-buffered and warp-specialized approaches substantially outperform the naive tiled kernel. The double_buffer_pipeline halves the runtime by overlapping tile loads and computation, while the warp_special ized_pipeline adds another 10.2% speedup by avoiding any implicit block-wide waits. Only the dedicated “compute” warp ever stalls.如表 10-3 所示,双缓冲与线程束专用化两种方法都显著优于朴素分块内核。double_buffer_pipeline 通过重叠分块加载与计算将运行时间减半,而 warp_special ized_pipeline 通过避免任何隐式的全块等待又带来 10.2% 的加速。只有专用的“计算”线程束才会停顿。

Instruction counts drop from 1.7 billion in the naive version to 1.05 billion in the two-stage pipeline, a 38% reduction, and further to ~1.00 billion in the warp-specialized kernel for an additional 4.76% reduction.指令数从朴素版本的 17 亿条降至两阶段流水线的 10.5 亿条,减少了 38%,并进一步降至线程束专用化内核的约 10 亿条,再减少 4.76%。

L2 load throughput climbs from 80 GB/s in naive tiling to 155 GB/s in the two-stage approach (+94%) and then to 165 GB/s in the warp-specialized kernel (+6.45% versus two-stage). This is because this warp-specialized kernel dedicates one warp to loading each tile into shared memory once—and then multicasts that single copy to all compute lanes. This eliminates any remaining redundant L2 reads. As such, after tiling and double-buffering, nearly all redundancy is already removed from the pipeline.L2 加载吞吐量从朴素分块的 80 GB/s 攀升至两阶段方法的 155 GB/s(+94%),再升至线程束专用化内核的 165 GB/s(较两阶段 +6.45%)。这是因为该线程束专用化内核专门用一个线程束将每个分块一次性加载到共享内存——然后将这一份副本多播到所有计算通道。这消除了任何残留的冗余 L2 读取。因此,在分块与双缓冲之后,流水线中几乎所有的冗余都已被消除。

In practice, the two-stage double-buffered pipeline is ideal for uniformly tiled GEMM workloads. Its simpler producer/consumer model hides most of the DRAM latency under compute. Meanwhile, the warp-specialized approach is optimized for irregular or deeper pipelines such as fused attention kernels. This is because each warp can continuously perform its assigned role—loading, computing, or storing—without ever forcing the rest of the block to stall.在实践中,两阶段双缓冲流水线非常适合均匀分块的 GEMM 工作负载。其更简单的生产者/消费者模型可在计算之下隐藏大部分 DRAM 时延。与此同时,线程束专用化方法针对不规则或更深的流水线(例如融合注意力内核)进行了优化。这是因为每个线程束都能持续执行其被分配的角色——加载、计算或存储——而绝不会迫使线程块的其余部分停顿。

PyTorch、CUDA Pipeline API 与线程束专用化

PyTorch, CUDA Pipeline API, and Warp Specialization

PyTorch’s public API doesn’t expose cuda::pipeline directly, but when you invoke torch.compile, the compiler generates kernels (implemented in OpenAI’s Triton language) that use optimizations that implement functionality equivalent to <cuda/pipe line> primitives and warp-specialized producers/consumers.PyTorch 的公开 API 并不直接暴露 cuda::pipeline,但当你调用 torch.compile 时,编译器会生成内核(以 OpenAI 的 Triton 语言实现),这些内核所采用的优化实现了与 <cuda/pipe line> 原语以及线程束专用化的生产者/消费者等价的功能。

So while you won’t see cuda::pipeline calls explicitly generated, you will see the underlying instructions and barriers. These optimizations help to improve occupancy and increase data-transfer/computation overlap. In other words, you get the same low-latency, high-throughput behavior of a handwritten <cuda/pipeline> implementation without writing any CUDA code yourself.因此,虽然你不会看到显式生成的 cuda::pipeline 调用,但你会看到底层的指令和屏障。这些优化有助于提升占用率,并增加数据传输与计算的重叠。换言之,你无需自己编写任何 CUDA 代码,就能获得手写 <cuda/pipeline> 实现所具有的低时延、高吞吐行为。

PyTorch’s fused attention kernels, produced by TorchInductor, use warp specialization with producer and consumer groups. For instance, consider three separate GPU kernels in PyTorch shown here:由 TorchInductor 生成的 PyTorch 融合注意力内核使用线程束专用化,并划分出生产者组与消费者组。例如,考虑此处展示的 PyTorch 中三个独立的 GPU 内核:

Python
scores = torch.matmul(queries, keys.transpose(-2, -1))
probabilities = F.softmax(scores, dim=-1)
context = torch.matmul(probabilities, values)

You can simply rewrite this with one line in PyTorch as shown here:你可以像这里展示的那样,用 PyTorch 中的一行代码重写它:

Python
context = torch.nn.functional.scaled_dot_product_attention(queries,
   keys, values)

Under the hood, the loader warps fetch tiles into shared memory. the compute warps process the tiles, and the store warps write results back. This eliminates expensive round trips to global memory between matmul and softmax stages.在底层,加载线程束将数据块取入共享内存,计算线程束处理这些数据块,存储线程束将结果写回。这消除了 matmul 与 softmax 阶段之间往返全局内存的高昂开销。

Overall, torch.compile is ideal for performance-sensitive and irregular workloads. Modern GPUs like Blackwell have abundant SM resources and high memory bandwidth. As such, techniques like warp specialization minimize idle cycles and allow near-linear scaling across warps—at least until the SMs—or memory bandwidth—are fully saturated.总体而言,torch.compile 非常适合对性能敏感且不规则的工作负载。像 Blackwell 这样的现代 GPU 拥有充裕的 SM 资源和高显存带宽。因此,线程束专用化之类的技术能够最小化空闲周期,并使跨线程束接近线性扩展——至少在 SM 或显存带宽被完全饱和之前如此。

By decoupling load, compute, and store into independent stages, the pipeline model maximizes throughput and resource utilization, making it the preferred approach for long-running kernels, including such processes as transformer attention, fused operator pipelines, and custom task schedulers. These kernels use the fine-grained inter-warp communication, deep pipelining, and linear warp scaling to provide high performance.通过将加载、计算和存储解耦为相互独立的阶段,流水线模型最大化了吞吐量和资源利用率,使其成为长时间运行内核的首选方案,包括 Transformer 注意力、融合算子流水线以及自定义任务调度器等过程。这些内核利用细粒度的线程束间通信、深度流水线和线性线程束扩展来提供高性能。

持久化内核与巨型内核

Persistent Kernels and Megakernels

Persistent kernels, also called persistent threads, invert the usual one-kernel-per-task approach. Instead of launching many small kernels in which each incurs significant overhead, you can launch a single, long-running kernel whose threads continually pull work from a shared producer-consumer queue in global or shared memory.持久化内核(persistent kernels),也称为*持久化线程*,颠覆了通常的“一个任务一个内核”的做法。无需启动许多小内核、每个都带来可观开销,你可以启动单个长时间运行的内核,其线程持续从全局内存或共享内存中的共享生产者—消费者队列中拉取工作。

When persistent threads loop, they handle data chunks as they arrive—often using memory copies or host signals—without exiting the kernel. This avoids repeated kernel launch overhead entirely. For instance, a persistent kernel might use one thread block’s Warp 0 to copy data from global memory (or CPU host memory) to shared memory. In the meantime, Warp 1 computes the previous batch. This is a form of software pipelining on GPU.持久化线程在循环时,会在数据块到达时处理它们——通常借助显存拷贝或主机信号——而不退出内核。这完全避免了重复的内核启动开销。例如,一个持久化内核可能使用某个线程块的 Warp 0 将数据从全局内存(或 CPU 主机内存)拷贝到共享内存。与此同时,Warp 1 计算上一批数据。这是 GPU 上软件流水线的一种形式。

For instance, consider having 1,000 tiny, independent tasks. Traditionally, one might launch 1,000 separate kernels. Each kernel occupies only a few SMs for a brief moment before exiting.例如,假设有 1,000 个微小的独立任务。传统做法可能会启动 1,000 个独立的内核。每个内核只在短暂瞬间占用少数几个 SM 便退出。

In practice, the GPU would repeatedly ramp up for each tiny kernel and ramp down afterward. This would leave most SMs idle between launches—and would fail to utilize the hardware fully.在实践中,GPU 会为每个微小内核反复地加速启动、随后减速退出。这会使大多数 SM 在两次启动之间处于空闲状态——无法充分利用硬件。

With a persistent kernel, you instead launch one large grid designed to keep the GPU busy for the entire workload. On a GPU with 132 SMs, for instance, this might mean launching one block per SM with 256 threads per block. That’s 33,792 threads in total. Each thread then executes code that looks roughly like the following:而使用持久化内核时,你改为启动一个大型网格,旨在让 GPU 在整个工作负载期间保持忙碌。例如,在拥有 132 个 SM 的 GPU 上,这可能意味着每个 SM 启动一个线程块、每个线程块 256 个线程。总计 33,792 个线程。随后每个线程执行大致如下所示的代码:

Text
__device__ int g_index;  // 下一个任务的全局计数器;在主机端初始化为 0
__global__ void persistentKernel(Task* tasks, int totalTasks) {
    // 每个线程循环,原子地获取下一个任务索引,直到没有剩余任务
    while (true) {
        int idx = atomicAdd(&g_index, 1);
        if (idx >= totalTasks) break;
        processTask(tasks[idx]);
    }
}
__device__ int g_index;  // global counter for next task; initialize to 0 on host
__global__ void persistentKernel(Task* tasks, int totalTasks) {
    // Every thread loops, atomically grabbing next task index until none remain
    while (true) {
        int idx = atomicAdd(&g_index, 1);
        if (idx >= totalTasks) break;
        processTask(tasks[idx]);
    }
}

Before launching this kernel, the host would set g_index = 0 in device memory. It would then invoke the kernel as shown here:在启动该内核之前,主机需要在设备显存中将 g_index = 0 设置好。然后按此处所示的方式调用该内核:

Text
cudaMemset(&g_index, 0, sizeof(int));
// 在 132 个 SM 的 GPU 上每个 SM 一个线程块
int blocks         = 132;
int threadsPerBlock = 256;
persistentKernel<<<blocks, threadsPerBlock>>>(d_tasks,
    totalTasks);
cudaDeviceSynchronize();
cudaMemset(&g_index, 0, sizeof(int));
// one block per SM on a 132 SM GPU
int blocks         = 132;
int threadsPerBlock = 256;
persistentKernel<<<blocks, threadsPerBlock>>>(d_tasks,
    totalTasks);
cudaDeviceSynchronize();

Now, instead of paying launch overhead for every single task, you pay only once to launch the persistentKernel. Assuming the launch takes roughly 0.02 ms and each of the 1,000 tasks runs in 0.1 ms, the kernels would run back to back for approximately 100 ms of total work.现在,不必为每一个任务都付出启动开销,只需为启动 persistentKernel 付出一次开销。假设启动耗时约 0.02 ms,而 1,000 个任务每个运行 0.1 ms,那么这些内核将背靠背运行,总工作量约为 100 ms。

By comparison, running 1,000 tiny kernels back to back, each taking 0.02 ms to launch, would add 20 ms in overhead to the execution time—separate from the 100 ms of total runtime for all 1,000 tasks. In short, consolidating small tasks into a persistent loop can cut tens of milliseconds of launch overhead.相比之下,背靠背运行 1,000 个微型内核,每个启动耗时 0.02 ms,会在执行时间上额外增加 20 ms 的开销——这还不包括全部 1,000 个任务共 100 ms 的总运行时间。简而言之,将小任务合并进一个持久化循环,可以削减数十毫秒的启动开销。

Using persistent kernels, the GPU stays highly utilized—with nearly all SMs actively working on tasks concurrently—because tens of thousands of threads are available to process the ~1,000 tasks. In contrast, launching 1,000 tiny kernels sequentially leaves much of the GPU underutilized at any given time. In our example this equates to only ~35% of the GPU’s capacity being used on average.使用持久化内核时,GPU 能保持很高的利用率——几乎所有 SM 都在并发地处理任务——因为有数万个线程可用于处理这约 1,000 个任务。相比之下,顺序启动 1,000 个微型内核会使 GPU 的大部分在任何时刻都处于未充分利用状态。在我们的示例中,这相当于平均只使用了 GPU 约 35% 的算力。

In this persistent kernel scenario, each SM is running one block of 256 threads (out of 2,048 total threads max), so every SM is doing work. As such, nearly 100% of SMs are active even though each SM’s own occupancy is relatively low at 12% (256 ÷ 2048 threads.)在这个持久化内核场景中,每个 SM 运行一个 256 线程的线程块(最大总线程数为 2,048),因此每个 SM 都在工作。这样一来,尽管每个 SM 自身的占用率相对较低,仅为 12%(256 ÷ 2048 线程),但仍有接近 100% 的 SM 处于活跃状态。

Using persistent kernels in this manner, the GPU can maintain high measured SM Active percent, subject to register and shared memory usage. Remember to always verify with Nsight Compute.以这种方式使用持久化内核,GPU 可以维持较高的实测 SM 活跃百分比,但这受限于寄存器和共享内存的使用量。请记住始终用 Nsight Compute 进行验证。

On modern GPUs, persistent kernels are particularly effective because their larger shared-memory capacity and expanded register file allow each thread to hold more intermediate state on-chip. Threads can use TMA to prefetch tensor tiles for upcoming tasks while other warps compute. So while some threads are processing one task, other threads use TMA to prefetch data for upcoming tasks—without burdening the SM’s compute pipelines with memory transfer instructions.在现代 GPU 上,持久化内核尤其有效,因为其更大的共享内存容量和扩展的寄存器堆让每个线程能在片上保存更多中间状态。线程可以使用 TMA 为即将到来的任务预取张量分块,同时其他线程束进行计算。因此,当一些线程正在处理某个任务时,另一些线程使用 TMA 为后续任务预取数据——不会用内存传输指令加重 SM 计算流水线的负担。

持久化内核的常见工作负载

Common Workloads for Persistent Kernels

Persistent kernels shine when you have many small or unevenly sized tasks that would incur high launch overhead if handled separately. They allow dynamic load balancing. This allows faster threads to continue looping and grab more work. Using persistent kernels, no SM ever goes idle prematurely.当你拥有大量小型或大小不均的任务、若分别处理会产生高昂启动开销时,持久化内核便能大显身手。它们支持动态负载均衡。这让更快的线程可以继续循环并领取更多工作。使用持久化内核,任何 SM 都不会过早空闲。

This pattern is common in irregular workloads such as graph traversals, custom batched transformations, and per-token operations common in LLM inference. In these cases, each task’s time can vary significantly.这种模式在不规则工作负载中很常见,例如图遍历、自定义的批量变换,以及 LLM 推理中常见的逐 token 操作。在这些情况下,每个任务的耗时可能差异很大。

There are downsides to persistent kernels, however. First, you must explicitly manage your task queue and synchronization using atomics. This can introduce contention if many threads attempt to increment the same counter simultaneously.然而,持久化内核也有缺点。首先,你必须使用原子操作显式管理任务队列和同步。如果许多线程同时尝试递增同一个计数器,就可能引入争用。

Debugging a single, giant persistent loop is more complex than debugging multiple small kernels. This is because a single divergent thread or an unexpected branch can cause the entire kernel to hang. Furthermore, one persistent kernel can monopolize the GPU indefinitely. So if other workloads need to run concurrently, you must carefully assign streams or partition resources.调试一个单一的巨大持久化循环比调试多个小内核更复杂。这是因为单个发生分歧的线程或一个意外的分支就可能导致整个内核挂起。此外,一个持久化内核可能无限期地独占 GPU。因此,如果其他工作负载需要并发运行,你必须谨慎地分配流或划分资源。

In short, persistent kernels can increase overall throughput substantially (e.g., 2–3×) compared to naive kernels by turning the GPU into a dynamic “worker-thread” pool that continuously fetches and processes tasks. On modern GPU hardware, this approach eliminates launch overhead, maximizes SM occupancy, and—when combined with cooperative groups or thread block clusters (described in a bit)—keeps data on-chip throughout multistage pipelines.简而言之,与朴素内核相比,持久化内核能将 GPU 变成一个持续获取并处理任务的动态“工作线程”池,从而大幅提升整体吞吐量(例如 2–3 倍)。在现代 GPU 硬件上,这种方法消除了启动开销,最大化了 SM 占用率,并且——当与协作组或线程块集群(稍后介绍)结合使用时——能在多阶段流水线中始终将数据保留在片上。

More and more frameworks and libraries are using persistent kernels and megakernels to avoid wasted capacity and improve performance of latency-sensitive workloads like inference. The key is to eliminate repeated launches and use device-side task queues on the GPU to keep the SMs fully occupied and performing useful work.越来越多的框架和库正在使用持久化内核(persistent kernel)和巨型内核(megakernel),以避免容量浪费,并提升推理等时延敏感型工作负载的性能。关键在于消除重复的启动开销,并利用 GPU 上的设备端任务队列,使 SM 保持满载并执行有用的工作。

用于推理的巨型内核

Megakernels for Inference

Additionally, a modern approach to persistent kernels originating from large-scale inference is called a megakernel. A megakernel fuses entire sequences of operations across layers—and even across GPUs—into a single large kernel. As shown in Figure 10-8, persistent megakernels have shown to reduce latency by 1.2× to 6.7× versus traditional per-layer launches by eliminating repeated kernel launch overhead.此外,源自大规模推理的一种现代持久化内核方法被称为*巨型内核*。巨型内核将跨层——甚至跨 GPU——的整个操作序列融合为单个大型内核。如图 10-8 所示,通过消除重复的内核启动开销,持久化巨型内核相比传统的逐层启动方式可将时延降低 1.2× 至 6.7×。

图 10-8
图 10-8巨型内核相对于 vLLM 和 SGLang 的解码吞吐量提升(来源:https://oreil.ly/2aZiF)Decode throughput improvement with megakernels relative to vLLM and SGLang (source: https://oreil.ly/2aZiF)

持久化内核与线程束专用化

Persistent Kernels and Warp Specialization

Warp specialization is typically used with persistent kernels in which threads perform many iterations over a relatively long period of time. This allows for deeper pipelines, better overlap, and efficient utilization of long-lived resources. For shorter-running kernels, the added code complexity of persistent kernels and warp specialization might not pay off.线程束专用化通常与持久化内核配合使用,此时线程会在相对较长的时间内执行多次迭代。这样可以实现更深的流水线、更好的重叠,以及长生命周期资源的高效利用。对于运行时间较短的内核,持久化内核和线程束专用化所带来的额外代码复杂性可能并不划算。

And a limitation of persistent-kernel scheduling is finding enough SMs for the persistent kernel to utilize. If too many SMs are occupied by another kernel, there might not be enough resources for the persistent kernel to launch. This makes it challenging when trying to schedule and load-balance work across SMs.持久化内核调度的一个局限在于,需要找到足够多的 SM 供该持久化内核使用。如果过多 SM 被另一个内核占用,持久化内核可能没有足够资源启动。这使得在 SM 之间调度和负载均衡工作变得颇具挑战。

To facilitate persistent kernels (and therefore warp specialization), modern GPUs support thread block clusters—also called cooperative thread array (CTA) clusters since a thread block is also called a cooperative thread array. We will discuss thread block (CTA) clusters in an upcoming section, but, in short, they let you combine thread blocks into “clusters” that occupy multiple nearby SMs on the GPU.为了支持持久化内核(进而支持线程束专用化),现代 GPU 支持*线程块集群*(thread block cluster)——由于线程块也称为协作线程数组,因此它也被称为*协作线程数组(CTA)集群*。我们将在后续小节讨论线程块(CTA)集群,简而言之,它们让你可以把线程块组合成“集群”,占用 GPU 上多个相邻的 SM。

协作组

Cooperative Groups

Cooperative groups let you define and synchronize groups of threads at arbitrary granularities. For example, you can create groups with individual threads, warps, tiles, blocks, and clusters, as shown in Figure 10-9.协作组让你可以按任意粒度定义并同步线程组。例如,你可以创建由单个线程、线程束、瓦片(tile)、线程块和集群组成的组,如图 10-9 所示。

图 10-9
图 10-9使用协作组在不同线程粒度上进行同步Synchronizing across different granularities of threads using cooperative groups

Cooperative groups provide safe, reusable collectives like sync, broadcast, and reduce. This is in contrast to using ad hoc synchronization barriers. Normally, threads can only synchronize within their own block using __syncthreads()—and there is no built-in global barrier for the entire grid, for example.协作组提供安全、可复用的集合操作,如 sync、broadcast 和 reduce。这与使用临时的同步屏障形成对比。通常,线程只能使用 __syncthreads() 在自己的线程块内同步——例如,并不存在针对整个网格的内置全局屏障。

Cooperative groups give you fine-grained synchronization inside the kernel. The API is ideal for coordinating multistage pipelines across warps, blocks, and clusters. To use the Cooperative Groups API with your kernels, you simply include <coopera tive_groups.h>, obtain group objects, then synchronize and coordinate across these groups.协作组让你可以在内核内部进行细粒度同步。该 API 非常适合协调跨线程束、线程块和集群的多阶段流水线。要在内核中使用 Cooperative Groups API,只需包含 <coopera tive_groups.h>,获取组对象,然后在这些组之间进行同步与协调。

The API includes calls like cg::this_thread_block(), cg::tiled_partition(), and cg::this_cluster(). You then call group.sync()—or similar collectives—to coordinate the threads in these groups.该 API 包含 cg::this_thread_block()、cg::tiled_partition() 和 cg::this_cluster() 等调用。随后调用 group.sync()——或类似的集合操作——来协调这些组内的线程。

To launch a kernel in cooperative mode, you use cudaLaunchCooperativeKernel(). In cooperative mode, CUDA ensures the grid you launch can be resident concurrently—otherwise, the launch fails. As such, it’s recommended to always size the cooperative grid using cudaOccupancyMaxActiveBlocksPerMultiprocessor and clamp grid size to avoid launch failure. Inside the kernel, you can then call the following to implement an all-thread-block barrier:要以协作模式启动内核,需使用 cudaLaunchCooperativeKernel()。在协作模式下,CUDA 确保所启动的网格能够同时驻留——否则启动就会失败。因此,建议始终使用 cudaOccupancyMaxActiveBlocksPerMultiprocessor 来确定协作网格的规模,并对网格大小进行钳制,以避免启动失败。在内核内部,随后可以调用以下内容来实现全线程块屏障:

Text
cooperative_groups::this_grid().sync(); // 或 grid.sync()
cooperative_groups::this_grid().sync(); // or grid.sync()

In this case, no thread in any block can proceed past that point until every thread in every block has reached it. This barrier allows you to split a kernel into sequential phases without ending the launch or returning control to the host.在这种情况下,任何块中的任何线程都无法越过该点继续执行,直到每个块中的每个线程都到达该点。该屏障允许你将一个内核拆分为多个顺序阶段,而无需结束本次启动或将控制权交还给主机。

For instance, consider the softmax algorithm, common in LLMs, that has two stages: a reduction across the entire array to compute an aggregate sum, and then a subsequent per-element computation that uses the aggregate sum. Traditionally, you would launch one kernel to do the reduction, copy the result back to host memory or global memory, launch a second kernel to consume the result from host or global memory, then calculate the softmax. This requires a lot of relatively slow memory movement.例如,考虑 LLM 中常见的 softmax 算法,它包含两个阶段:先对整个数组做归约以计算总和,然后使用该总和进行逐元素计算。传统做法是启动一个内核完成归约,将结果复制回主机内存或全局内存,再启动第二个内核从主机内存或全局内存读取该结果,然后计算 softmax。这需要大量相对缓慢的内存搬运。

With cooperative groups, you can perform both stages in one kernel such that each block computes its partial sum, one block aggregates those partial sums into a final result, all blocks call grid.sync() to wait until aggregation is complete, then all threads proceed to the second stage. Each thread would then read the aggregate sum from a register or shared memory—rather than from global memory.借助协作组,你可以在一个内核中完成这两个阶段:每个块计算自己的部分和,由一个块将这些部分和汇总为最终结果,所有块调用 grid.sync() 等待汇总完成,随后所有线程进入第二阶段。此时每个线程从寄存器或共享内存读取总和,而不是从全局内存读取。

CUDA guarantees that every block in a cooperative launch is resident on the GPU at the same time. If you request more blocks than can fit concurrently, the launch simply fails.CUDA 保证协作式启动中的每个块同时驻留在 GPU 上。如果请求的块数超过可同时容纳的数量,启动就会直接失败。

Because cooperative kernels must be launched with a grid size that the GPU can run concurrently, grid.sync() will not hang waiting for nonexistent blocks. In other words, the CUDA runtime guarantees that all thread blocks are active and will reach the grid.sync() barrier. If the kernel launch were too large to run at once, it would simply fail to launch. As such, it’s important to check the return status of cudaLaunch CooperativeKernel().由于协作式内核必须以 GPU 能够并发运行的网格大小启动,grid.sync() 不会因等待不存在的块而挂起。换言之,CUDA 运行时保证所有线程块都处于活动状态,并且都会到达 grid.sync() 屏障。如果内核启动规模过大而无法一次运行完,它就会直接启动失败。因此,检查 cudaLaunch CooperativeKernel() 的返回状态非常重要。

For instance, if a GPU has 132 SMs and your kernel uses enough resources that each SM can run four blocks, the grid must be no larger than 528 blocks to succeed. If you exceed that limit, the cooperative launch will simply fail. Use cudaOccupancyMax ActiveBlocksPerMultiprocessor to size the grid for a cooperative kernel and check the cudaLaunchCooperativeKernel return status before assuming progress has been made.例如,如果一块 GPU 有 132 个 SM,而你的内核使用的资源量使得每个 SM 能运行四个块,那么网格不得超过 528 个块才能成功。一旦超过该上限,协作式启动就会直接失败。应使用 cudaOccupancyMax ActiveBlocksPerMultiprocessor 来确定协作式内核的网格规模,并在假定已有进展之前检查 cudaLaunchCooperativeKernel 的返回状态。

Prior to Blackwell, developers often used cudaLaunchCooperativeKernel together with global-memory atomics or flags to coordinate multiple blocks that did not share on-chip memory. This approach worked but forced intermediate results to be moved to global memory. This incurred extra HBM traffic.在 Blackwell 之前,开发者常将 cudaLaunchCooperativeKernel 与全局内存原子操作或标志位结合使用,以协调多个不共享片上内存的块。这种做法可行,但会迫使中间结果被搬到全局内存,从而带来额外的 HBM 流量。

On Blackwell, thread block clusters and DSMEM provide a far more efficient alternative. A thread block cluster can share data in on-chip SRAM—and synchronize without global memory round trips. We will cover thread block clusters and DSMEM later in this chapter.在 Blackwell 上,线程块集群与 DSMEM 提供了效率高得多的替代方案。线程块集群可以在片上 SRAM 中共享数据,并在无需往返全局内存的情况下完成同步。本章后文将介绍线程块集群与 DSMEM。

You should use cooperative kernels when you need a true, all-thread-block barrier within a single kernel launch. Also, they’re useful when you want to keep intermediate results in fast memory (e.g., registers or shared memory) rather than repeatedly writing to and reading from global memory. It’s recommended that you constrain your grid size to the GPU’s maximum concurrent block capacity—or choose larger blocks—to reduce overall block count and a failed kernel launch.当你需要在单次内核启动内实现真正的、覆盖所有线程块的屏障时,应当使用协作内核。此外,当你希望把中间结果保留在快速存储器(例如寄存器或共享内存)中,而不是反复写入全局内存再从中读回时,协作内核同样有用。建议将网格大小限制在 GPU 的最大并发块容量之内——或者选择更大的线程块——以减少总块数,并避免内核启动失败。

The downsides of cooperative kernels are that your grid size is limited by capacity. This may force you to use fewer, larger blocks—or rely on thread block clusters (later in this chapter).协作内核的缺点是网格大小受容量限制。这可能迫使你使用数量更少、尺寸更大的线程块——或者依赖线程块集群(本章稍后介绍)。

Even though all blocks run concurrently, a cooperative barrier still requires every thread in every block to call grid.sync(). If any single thread skips—or never reaches—the grid.sync() call (e.g., because of a divergent if statement), then every other thread that did call grid.sync() will wait forever. This results in a deadlock.尽管所有线程块都并发运行,协作屏障仍要求每个线程块中的每个线程都调用 grid.sync()。如果任何一个线程跳过——或始终未到达——grid.sync() 调用(例如由于分支发散的 if 语句),那么所有其他确实调用了 grid.sync() 的线程都将永远等待。这会导致死锁。

In short, cooperative group kernels let you treat the entire GPU as a single collabora-tive resource, with grid.sync() acting as a global barrier. This is ideal for multistage algorithms requiring global synchronization and data sharing. Just remember that grid.sync() is a relatively heavyweight synchronization with higher overhead than block-scoped or cluster-scoped barriers. This is because it must coordinate across all thread blocks running on multiple SMs. As such, you should use grid.sync() sparingly and, at the very least, make sure that your kernel does a significant amount of work between heavyweight synchronization calls.简而言之,协作组内核让你可以把整个 GPU 视为单一协作资源,由 grid.sync() 充当全局屏障。这对于需要全局同步与数据共享的多阶段算法非常理想。只需记住,grid.sync() 是一种相对重量级的同步,其开销高于块作用域或集群作用域的屏障。这是因为它必须协调运行在多个 SM 上的所有线程块。因此,你应当谨慎使用 grid.sync(),至少应确保内核在重量级同步调用之间完成大量工作。

协作网格同步与持久化内核

Cooperative Grid Synchronization and Persistent Kernels

If your workload occasionally needs cross-thread-block barriers inside a persistent loop to aggregate partial results, you can combine persistent kernels with cooperative groups by calling grid.sync(). This will provide a grid-wide barrier to avoid ending the kernel and having to relaunch it. In this way, a multistage pipeline of reductions and other global steps remains entirely on-device.如果你的工作负载偶尔需要在持久化循环内部使用跨线程块屏障来聚合部分结果,你可以将持久化内核与协作组结合,调用 grid.sync()。这将提供网格范围的屏障,避免结束内核后不得不重新启动它。这样,由归约和其他全局步骤组成的多级流水线就完全留在设备上执行。

For instance, consider a workload that performs two different computations repeatedly across 1,000 iterations such that each computation requires a global barrier because of cross-block data dependencies. A naive implementation might launch two separate kernels per iteration, resulting in 2,000 kernel launches. In a single stream, the second kernel waits for the first automatically—no explicit host-side sync—but you still pay launch overhead. Instead, you can fuse everything into one cooperative, persistent kernel, as shown here:例如,考虑一个在 1,000 次迭代中重复执行两种不同计算的工作负载,其中每种计算都因跨块数据依赖而需要全局屏障。朴素的实现可能在每次迭代中启动两个独立的内核,从而产生 2,000 次内核启动。在单个流中,第二个内核会自动等待第一个内核——无需显式的主机端同步——但你仍然要付出启动开销。相反,你可以把所有内容融合进一个协作的持久化内核,如下所示:

CUDA / C
#include <cuda_runtime.h>
#include <cooperative_groups.h>
namespace cg = cooperative_groups;
__device__ inline float someComputationA(float x) { return x * 2.0f; }
__device__ inline float someComputationB(float a, float b) { return a + b; }
__global__ void combinedKernel(float* __restrict__ dataA,
                               float* __restrict__ dataB,
                               int N,
                               int iterations) {
    cg::grid_group grid = cg::this_grid();                 // 全网格组
    const int tid = blockIdx.x * blockDim.x + threadIdx.x; // 线程的线性 id
    const int stride = gridDim.x * blockDim.x;             // 网格线程总数
    for (int it = 0; it < iterations; ++it) {
        // 阶段 1:使用网格跨步循环更新 A 的所有元素
        for (int i = tid; i < N; i += stride) {
            dataA[i] = someComputationA(dataA[i]);
        }
        grid.sync();  // 全局屏障;同时提供内存可见性
        // 阶段 2:用已完成的 A 更新 B(同样采用 grid-stride 方式)
        for (int i = tid; i < N; i += stride) {
            const float mid = dataA[i];
            dataB[i] = someComputationB(mid, dataB[i]);
        }
        grid.sync();  // 下一次迭代前的屏障
    }
}
#include <cuda_runtime.h>
#include <cooperative_groups.h>
namespace cg = cooperative_groups;
__device__ inline float someComputationA(float x) { return x * 2.0f; }
__device__ inline float someComputationB(float a, float b) { return a + b; }
__global__ void combinedKernel(float* __restrict__ dataA,
                               float* __restrict__ dataB,
                               int N,
                               int iterations) {
    cg::grid_group grid = cg::this_grid();                 // whole-grid group
    const int tid = blockIdx.x * blockDim.x + threadIdx.x; // thread's linear id
    const int stride = gridDim.x * blockDim.x;             // total grid threads
    for (int it = 0; it < iterations; ++it) {
        // Stage 1: update all elements of A using a grid-stride loop
        for (int i = tid; i < N; i += stride) {
            dataA[i] = someComputationA(dataA[i]);
        }
        grid.sync();  // global barrier; also provides memory visibility
        // Stage 2: use finished A to update B (again grid-stride)
        for (int i = tid; i < N; i += stride) {
            const float mid = dataA[i];
            dataB[i] = someComputationB(mid, dataB[i]);
        }
        grid.sync();  // barrier before next iteration
    }
}

On the host side, you would simply invoke this combined cooperative and persistent kernel, as shown here:在主机端,你只需调用这个协作与持久化相结合的内核,如下所示:

Text
// ……分配/拷贝 dA、dB,设置 N、iterations,选择 blockSize
cudaDeviceProp prop{};
cudaGetDeviceProperties(&prop, 0);
if (!prop.cooperativeLaunch) {
    // 回退方案:每次迭代运行两个内核,而不是使用 grid.sync()
    // (参见下面的替代方案)
}
// ... allocate/copy dA, dB, set N, iterations, choose blockSize
cudaDeviceProp prop{};
cudaGetDeviceProperties(&prop, 0);
if (!prop.cooperativeLaunch) {
    // Fallback: run two kernels per iteration instead of grid.sync()
    // (see alternative below)
}
Text
// 计算协作内核允许的最大网格大小
int maxBlocksPerSm = 0;
const int blockSize = 256;
cudaOccupancyMaxActiveBlocksPerMultiprocessor(&maxBlocksPerSm,
                                              combinedKernel,
                                              blockSize, 0);
const int maxCoopGrid = maxBlocksPerSm * prop.multiProcessorCount;
// 选择网格大小,但将其限制在协作上限之内
int gridSize = (N + blockSize - 1) / blockSize;
if (gridSize > maxCoopGrid) gridSize = maxCoopGrid;
// 以协作方式启动
void* args[] = { &dA, &dB, &N, &iterations };
cudaLaunchCooperativeKernel((void*)combinedKernel, gridSize, blockSize, args);
cudaDeviceSynchronize();
// Compute the maximum grid size allowed for a cooperative kernel
int maxBlocksPerSm = 0;
const int blockSize = 256;
cudaOccupancyMaxActiveBlocksPerMultiprocessor(&maxBlocksPerSm,
                                              combinedKernel,
                                              blockSize, 0);
const int maxCoopGrid = maxBlocksPerSm * prop.multiProcessorCount;
// Choose grid size, but clamp to cooperative limit
int gridSize = (N + blockSize - 1) / blockSize;
if (gridSize > maxCoopGrid) gridSize = maxCoopGrid;
// Launch cooperatively
void* args[] = { &dA, &dB, &N, &iterations };
cudaLaunchCooperativeKernel((void*)combinedKernel, gridSize, blockSize, args);
cudaDeviceSynchronize();

This single kernel replaces what would otherwise be 2 × 1,000 = 2,000 kernel launches and avoids repeated launch overhead. Inside the loop, grid.sync() ensures correct ordering (every block finishes Stage 1 before any block begins Stage 2—and before the next iteration) without any host-side synchronization.这单个内核取代了原本需要 2 × 1,000 = 2,000 次内核启动的做法,避免了重复的启动开销。在循环内部,grid.sync() 确保正确的顺序(每个线程块在任何一个线程块开始阶段 2 之前完成阶段 1——并且在进入下一次迭代之前也是如此),而无需任何主机端同步。

Per-thread and per-block data can remain in registers/shared memory across grid.sync(). For cross-block data exchange, use global memory (or thread block clusters and distributed shared memory, as we’ll cover in a bit). Meanwhile, because the work is looped on the GPU, there is no launch overhead after the first invocation.每个线程和每个线程块的数据可以在 grid.sync() 之间保留在寄存器/共享内存中。对于跨块数据交换,请使用全局内存(或线程块集群与分布式共享内存,我们稍后将介绍)。同时,由于工作是在 GPU 上循环执行的,首次调用之后不再有启动开销。

Cooperative kernels require device support (cooperativeLaunch) and must be launched with cudaLaunchCooperativeKernel. All blocks in the grid must be resident concurrently. Size and clamp the grid using the CUDA Occupancy APIs so all CTAs can co-reside. Otherwise the launch will fail. An example is when ((N + threads - 1) ÷ threads) exceeds the cooperative capacity.协作内核需要设备支持(cooperativeLaunch),并且必须通过 cudaLaunchCooperativeKernel 启动。网格中的所有线程块必须同时驻留。使用 CUDA 占用率 API 来确定网格大小并对其进行钳制,使所有 CTA 都能共同驻留。否则启动将失败。例如,当 ((N + threads - 1) ÷ threads) 超过协作容量时就会发生这种情况。

Modern GPUs can run on the order of a few thousand thread blocks concurrently (assuming each uses a modest amount of resources, including registers and shared memory), so there is considerable headroom. However, you should still verify block size and resource consumption before proceeding with this implementation.现代 GPU 可以同时运行数千个线程块(假设每个线程块使用的资源量适中,包括寄存器和共享内存),因此仍有相当大的余量。不过,在着手此实现之前,你仍应核实线程块大小和资源消耗情况。

In this kernel example, the kernel never returns control to the host between iterations. As such, it behaves like a persistent kernel for the outer loop but enforces global synchronization points in the inner stages.在这个内核示例中,内核在各次迭代之间从不将控制权交还给主机。因此,它对外层循环而言表现得像一个持久化内核,但在内层各阶段中强制执行全局同步点。

This combined pattern is especially powerful for multistep algorithms in LLM inference in which each layer may require a reduction or normalization (Stage 1) followed by a per-element transformation (Stage 2). By capturing everything in a cooperative persistent kernel, you eliminate all inter-kernel round trips and maximize on-chip data locality.这种组合模式对于 LLM 推理中的多步算法尤其强大,其中每一层可能都需要先做一次归约或归一化(阶段 1),再做逐元素变换(阶段 2)。通过将一切捕获在一个协作持久化内核中,你可以消除所有内核间的往返,并最大化片上数据局部性。

何时结合持久化内核与协作组

When to Combine Persistent Kernels and Cooperative Groups

A common best practice is to launch a persistent kernel with one thread block per SM if resources allow. This way, each SM has a resident thread block that iteratively processes tasks from a work queue. This maximizes occupancy and keeps SMs from becoming idle. The thread blocks would then use grid.sync() to coordinate between their phases. However, when deciding between cooperative and persistent strategies— and whether to combine them—ask the following questions:一条常见的最佳实践是:在资源允许的情况下,以每个 SM 一个线程块的方式启动持久化内核。这样,每个 SM 都有一个常驻线程块,从工作队列中迭代地处理任务。这可以最大化占用率,并避免 SM 空闲。随后,各线程块会使用 grid.sync() 在各自的阶段之间进行协调。不过,在协作式与持久化策略之间做选择——以及是否将二者结合——时,请问以下问题:

Do you have multiple sequential phases that need global synchronization?你是否有多个需要全局同步的顺序阶段?
If yes, use cooperative kernels (or cooperative persistent kernels) so you can call grid.sync() between phases.如果是,请使用协作内核(或协作持久化内核),这样你就可以在各阶段之间调用 grid.sync()。
Do you have many small or irregular tasks that suffer from launch overhead?你是否有许多因启动开销而受损的小型或不规则任务?
If yes, use persistent kernels so your threads loop over a shared queue without returning to the host.如果是,请使用持久化内核,让你的线程在共享队列上循环,而不返回主机。
Can you afford to reserve the entire grid (or a thread block cluster) exclusively for this* *workload?你能否承担将整个网格(或一个线程块集群)专门保留给此* *工作负载?
If yes, a cooperative persistent kernel may yield the best performance.如果是,协作持久化内核可能会带来最佳性能。
Do you need to share SMs with other work?你需要与其他工作共享 SM 吗?
If yes, consider using thread block clusters (next section) so the persistent or cooperative kernel does not serve other workloads.如果答案是肯定的,考虑使用线程块集群(见下一节),这样持久化内核或协作内核就不会再服务于其他工作负载。

In short, when you use a single kernel that both loops persistently over tasks and includes grid.sync() calls to synchronize all thread blocks between stages, you eliminate the pauses and extra memory transfers that normally occur between separate kernel launches.简而言之,当你使用单个内核,既持久地循环处理任务,又包含 grid.sync() 调用来在各阶段之间同步所有线程块时,你就消除了通常发生在独立内核启动之间的停顿和额外内存传输。

With modern GPUs, this means data stays in shared memory or registers throughout the entire computation. This is in contrast to writing the data back to global memory after each phase. As a result, the GPU stays busy doing useful work almost all the time—achieving performance close to its peak hardware limits.在现代 GPU 上,这意味着数据在整个计算过程中始终留在共享内存或寄存器中。这与在每个阶段结束后把数据写回全局内存的做法形成对比。因此,GPU 几乎始终忙于做有用的工作——性能接近其硬件峰值极限。

One important caveat: cooperative kernels reserve all SMs in the grid, so no other kernels can run concurrently on those SMs. If you need to coschedule other work such as asynchronous data prefetch on a separate stream or lower-priority inference kernels, you may need to partition the GPU into thread block clusters, covered in the next section.一个重要提醒:协作内核会预留网格中的所有 SM,因此其他内核无法在这些 SM 上并发运行。如果你需要协同调度其他工作,例如在独立流上进行异步数据预取,或运行低优先级推理内核,你可能需要把 GPU 划分为线程块集群,这将在下一节介绍。

图 10-10
图 10-10多个线程块集群保证被协同调度到同一个 GPC 或 GPC 分区上Multiple thread block clusters are guaranteed to be coscheduled on the same GPC or GPC partition

线程块集群与分布式共享内存

Thread Block Clusters and Distributed Shared Memory

A cooperative group is a software-level abstraction that provides an API that lets you carve up kernel threads into arbitrary collectives for synchronization and data movement. This includes warps, tiles, thread blocks—even entire grids or multidevice grids.*协作组*是一种软件层面的抽象,它提供的 API 让你可以把内核线程划分成任意集合,用于同步和数据移动。这包括线程束、瓦片、线程块——甚至整个网格或多设备网格。

In contrast, a thread block cluster (or cooperative thread array [CTA] cluster), is a hardware-level hierarchy. It grants a subset of SMs to your cooperative grid—and leaves the remainder free for other kernels to use. This mitigates the risk of one kernel monopolizing the GPU. The GPU guarantees the thread blocks will be coscheduled on the same GPU processing cluster (GPC), as shown in Figure 10-10.相比之下,线程块集群(或称协作线程数组(CTA)集群)是一种硬件层面的层级结构。它把一部分 SM 分配给你的协作网格,而把其余 SM 留给其他内核使用。这降低了单个内核独占 GPU 的风险。GPU 保证这些线程块被协同调度到同一个 GPU 处理集群(GPC)上,如图 10-10 所示。

图 10-11
图 10-11对于这四个(2 × 2)线程集群,A 和 B 的每个瓦片被同时加载到两个线程块中(来源:https://oreil.ly/kEZsv)For these four (2 × 2) thread clusters, each tile of A and B is loaded into two thread blocks simultaneously (source: https://oreil.ly/kEZsv)

A GPC is a collection of nearby SMs. The GPU schedules thread blocks onto GPCs similar to how it schedules threads of a thread block onto the same SM.GPC 是一组位置相邻的 SM 的集合。GPU 把线程块调度到 GPC 上的方式,类似于它把线程块中的线程调度到同一个 SM 上。

The GPU provides distributed shared memory (DSMEM), discussed in the next section, for the thread block cluster to use across those blocks. It also supports a cluster-level barrier using the Cluster Group API (cluster.sync()).GPU 提供分布式共享内存(DSMEM)(将在下一节讨论),供线程块集群在这些块之间使用。它还通过集群组 API(cluster.sync())支持集群级屏障。

This cluster-level barrier lets you synchronize only a subset of thread blocks without blocking the entire GPU. Thread block clusters let you launch a cooperative kernel that subdivides the grid into smaller groups of blocks, as shown in Figure 10-11.这个集群级屏障让你只同步一部分线程块,而不会阻塞整个 GPU。线程块集群让你可以启动一个协作内核,把网格细分为更小的块组,如图 10-11 所示。

Within each group, calling cluster.sync() provides a local barrier. This lets blocks inside a cluster share data through dedicated on-chip resources without monopolizing every SM. On modern GPUs, you can use DSMEM, which allows thread block clusters to share a contiguous region of on-chip SRAM. This enables low-latency communication between blocks in the same cluster with native hardware support.在每个组内,调用 cluster.sync() 提供一个局部屏障。这让集群内的块可以通过专用的片上资源共享数据,而不会独占每一个 SM。在现代 GPU 上,你可以使用 DSMEM,它允许线程块集群共享一块连续的片上 SRAM。这借助原生硬件支持,实现了同一集群内块之间的低时延通信。

Thread block clusters group a subset of thread blocks and call cluster.sync() within each cluster to synchronize “locally” within just those thread blocks. While threads within a thread block could traditionally cooperate using shared memory, on modern GPUs, they can now collaborate with each other using thread block clusters and DSMEM.线程块集群将一部分线程块组合在一起,并在每个集群内调用 cluster.sync(),从而仅在这些线程块内部进行“本地”同步。传统上,一个线程块内的线程可以借助共享内存协作;而在现代 GPU 上,线程现在可以通过线程块集群和 DSMEM 相互协作。

Said differently, without thread block clusters, only threads within the same thread block could share data in shared memory. Different thread blocks would have to coordinate using global memory and global barriers (e.g., a grid.sync()), which stalls all threads and limits scalability.换言之,如果没有线程块集群,只有同一线程块内的线程才能在共享内存中共享数据。不同的线程块则必须借助全局内存和全局屏障(例如 grid.sync())进行协调,而这会阻塞所有线程并限制可扩展性。

And, furthermore, threads in different thread blocks previously could not efficiently share state and synchronize except using either global memory or grid.sync() for coarse-grained, grid-wide synchronization, as described in the previous section. Unfortunately, both global memory and grid.sync() are relatively slow and, as such, can become bottlenecks.此外,不同线程块中的线程此前无法高效地共享状态并同步,除非使用全局内存或 grid.sync() 进行粗粒度的全网格同步,如上一节所述。遗憾的是,全局内存和 grid.sync() 都相对较慢,因此可能成为瓶颈。

Thread block clusters are supported natively by the GPU hardware, including cluster launch control. Cluster launch control is a hardware-level mechanism that launches and schedules persistent thread block clusters. Specifically, it allows a persistent kernel (and its thread block cluster) to maintain a balanced load of work—even when some of the SMs are occupied. This provides the basis for an efficient warp specialization implementation.线程块集群由 GPU 硬件原生支持,包括*集群* 启动控制(cluster launch control)。集群启动控制是一种硬件级机制,用于启动和调度持久化线程块集群。具体而言,它允许持久化内核(及其线程块集群)维持均衡的工作负载——即使部分 SM 已被占用。这为高效的线程束专用化实现提供了基础。

Using hardware-supported communication and synchronization constructs, thread block clusters can synchronize using cluster-wide barriers in the form of low-level PTX instructions and CUDA intrinsics. As such, blocks in a cluster can perform barrier synchronization much faster than a full grid synchronization using grid.sync() due to the hardware-supported barrier synchronization between the thread blocks.借助硬件支持的通信与同步原语,线程块集群可以通过集群范围的屏障进行同步,其形式为底层 PTX 指令和 CUDA 内建函数(intrinsics)。因此,由于线程块之间存在硬件支持的屏障同步,集群中的线程块执行屏障同步的速度远快于使用 grid.sync() 的完整网格同步。

线程块重排(Swizzling)

Thread Block Swizzling

In a straightforward grid launch, thread blocks process tiles in strict row-major or column-major order. This can cause early blocks to evict data that later blocks will need—resulting in poor reuse and extra memory traffic. Instead, you want tiles A and B within a single wave, which can be read out of L2 cache.在直接的网格启动中,线程块严格按照行主序或列主序处理瓦片。这可能导致较早的线程块驱逐后续线程块所需的数据——造成复用率低下和额外的内存流量。相反,你希望瓦片 A 和 B 位于同一个波次内,从而可以从 L2 缓存中读出。

To work around this inefficiency, you can use thread block swizzling. Similar to using swizzling to optimize memory access and avoid shared-memory bank conflicts, you can use thread block swizzling to avoid assigning tiles in the inefficient row-major and column-major order, as shown in Figure 10-12.为规避这一低效问题,可以使用线程块重排(thread block swizzling)。与使用重排来优化显存访问、避免共享内存存储体冲突类似,你可以用线程块重排来避免以低效的行主序和列主序分配瓦片,如图 10-12 所示。

图 10-12
图 10-12通过线程块重排,在单个波次内从 L2 缓存读取瓦片 A 和 BThread block swizzling to read tiles A and B in a single wave out of L2 cache

Thread block swizzling lets tiles of both A and B matrices, needed by the same wave, stay in L2 for maximum reuse. When applied to persistent and tiled GEMM workloads, this type of swizzling can produce double-digit performance gains by reducing memory misses and bandwidth pressure. Thread-block swizzling is a simple yet powerful pattern-reordering technique that aligns kernel launch order with cache locality.线程块重排使同一波次所需的 A、B 两个矩阵的瓦片都能留在 L2 中,从而实现最大化复用。当应用于持久化和分块 GEMM 工作负载时,这类重排可通过减少内存未命中和带宽压力带来两位数的性能提升。线程块重排是一种简单而强大的模式重排技术,它使内核启动顺序与缓存局部性保持一致。

With hardware support for thread block clusters, a relatively large number of SMs, and improved pipeline scheduling, modern GPUs can host large persistent kernels with a sophisticated warp specialization pipeline to optimize kernel performance.凭借对线程块集群的硬件支持、相对较多的 SM 数量以及改进的流水线调度,现代 GPU 能够承载带有复杂线程束专用化流水线的大型持久化内核,以优化内核性能。

In short, thread block clusters build on the cooperative groups model by allowing a subset of blocks to synchronize and share state without locking out the entire GPU. This lets you build multistage, fine-grained pipelines inside one launch without locking out the rest of the device. This leaves the remaining SMs free to perform independent work.简而言之,线程块集群建立在协作组模型之上,允许一部分线程块进行同步并共享状态,而不会锁住整个 GPU。这使你能够在一次启动内构建多级、细粒度的流水线,而不会锁住设备的其余部分。剩余的 SM 则可以自由执行独立的工作。

Within thread block clusters, there are two key mechanisms to share state between thread blocks: DSMEM and scratch memory. Let’s describe these next.在线程块集群内,有两种在线程块之间共享状态的关键机制:DSMEM 和暂存内存(scratch memory)。接下来我们逐一介绍。

分布式共享内存

Distributed Shared Memory

Distributed shared memory (DSMEM) extends the concept of __shared__ memory beyond a single thread block to span an entire thread block cluster. In a traditional kernel, each thread block has its own private __shared__ region and is inaccessible to other thread blocks.分布式共享内存(DSMEM)将 __shared__ 内存的概念从单个线程块扩展到整个线程块集群。在传统内核中,每个线程块都有自己私有的 __shared__ 区域,其他线程块无法访问。

With DSMEM, however, multiple blocks in the same cluster can read/write into a shared memory space that logically combines all of their local __shared__ regions. In effect, the cluster’s shared memory is stitched together into one distributed on-chip buffer.而借助 DSMEM,同一集群中的多个线程块可以读写一个共享内存空间,该空间在逻辑上合并了它们所有的本地 __shared__ 区域。实际上,集群的共享内存被拼接成一个分布式的片上缓冲区。

By keeping interblock data inside on-chip memory, DSMEM significantly increases effective arithmetic intensity since blocks can exchange intermediate results or perform reductions without round-tripping to global memory. This is especially valuable when a dataset is too large for a single block’s shared memory.通过将块间数据保留在片上内存中,DSMEM 显著提高了有效算术强度,因为线程块可以交换中间结果或执行归约,而无需往返全局内存。当数据集大到单个线程块的共享内存无法容纳时,这一点尤其有价值。

Instead of falling back to global loads and stores, you launch a thread block cluster whose blocks each work on a portion of the data, then synchronize and share using DSMEM. All communication happens at SM-to-SM speed, which avoids expensive HBM traffic.无需退回到全局加载与存储,而是启动一个线程块集群,其中每个线程块处理数据的一部分,然后使用 DSMEM 进行同步与共享。所有通信都以 SM 到 SM 的速度进行,从而避免了昂贵的 HBM 流量。

In short, DSMEM is a faster, more structured alternative to using global memory for block-to-block communication. But its scope is limited to the blocks in the same cluster. DSMEM is ideal for persistent kernels that require frequent, low-latency interblock coordination, attention mechanisms, or other multistage algorithms in LLMs in which every block must exchange intermediate state before proceeding—as well as any workload that benefits from keeping interim results on-chip rather than writing back to and reloading from global memory.简而言之,与使用全局内存进行块间通信相比,DSMEM 是一种更快、更有结构的替代方案。但其作用范围仅限于同一集群内的线程块。DSMEM 非常适合需要频繁、低时延块间协调的持久化内核、注意力机制,或 LLM 中每个线程块在继续执行前都必须交换中间状态的其他多阶段算法——也适合任何受益于将中间结果保留在片上、而非写回全局内存再重新加载的工作负载。

暂存内存(Scratch Memory)

Scratch Memory

Scratch memory is the low-level hardware infrastructure that underpins DSMEM and thread block cluster synchronization, whereas DSMEM is the shared data buffer that is implicitly visible to your CUDA code, scratch memory is a separate on-chip SRAM region used by the GPU to track coordination metadata such as barrier state, group progress, and access flags for DSMEM.暂存内存是支撑 DSMEM 与线程块集群同步的底层硬件基础设施;DSMEM 是 CUDA 代码可隐式看到的共享数据缓冲区,而暂存内存则是 GPU 用于跟踪协调元数据的独立片上 SRAM 区域,这些元数据包括屏障状态、组进度以及 DSMEM 的访问标志。

You do not access scratch memory directly from your kernels. Instead, the GPU manages it automatically. When you launch a thread block cluster, the hardware allocates a portion of scratch memory to maintain cluster-wide barrier counters (cluster .sync() state), track which blocks have arrived at DSMEM operations or synchronization points, and coordinate safe access to the distributed shared regions across SMs.内核不会直接访问暂存内存,而是由 GPU 自动管理。启动线程块集群时,硬件会分配一部分暂存内存,用于维护集群范围的屏障计数器(cluster .sync() 状态)、跟踪哪些线程块已到达 DSMEM 操作或同步点,并协调跨 SM 对分布式 shared 区域的安全访问。

Because scratch memory is optimized for these metadata operations, it enables cluster-level barriers and DSMEM accesses to complete very quickly. If the metadata size exceeds available scratch SRAM for very large clusters or complex synchronization patterns, the GPU transparently spills some state into local memory, which is backed by L1/L2 caches. This spill ensures correctness while still preserving as much on-chip efficiency as possible.由于暂存内存针对这些元数据操作进行了优化,它使集群级屏障和 DSMEM 访问能够非常快速地完成。对于非常大的集群或复杂的同步模式,如果元数据大小超出可用的暂存 SRAM,GPU 会透明地将部分状态溢出到本地内存,后者由 L1/L2 缓存支撑。这种溢出确保了正确性,同时尽可能保留片上效率。

In short, DSMEM is the abstraction you use to share data between blocks and scratch memory in the behind-the-scenes facility that makes those DSMEM operations fast and scalable. Together, they allow a thread block cluster to behave like a single logical unit that breaks the traditional barrier between thread blocks. This improves performance for workloads that need tightly coordinated parallelism across multiple thread blocks.简而言之,DSMEM 是你用来在线程块之间共享数据的抽象,而暂存内存则是幕后设施,使这些 DSMEM 操作快速且可扩展。二者结合,使线程块集群能够像单个逻辑单元一样运作,打破了线程块之间的传统壁垒。这为需要在多个线程块之间进行紧密协调并行的工作负载提升了性能。

启动线程块集群

Launching a Thread Block Cluster

Let’s illustrate how to use thread block clusters in practice. First, we need to launch a kernel with a cluster dimension that we haven’t used before.下面通过实践来说明如何使用线程块集群。首先,我们需要在启动内核时指定一个此前未曾使用过的 cluster 维度。

In CUDA, this is done with an extended launch API that allows specifying the cluster size. For example, if we want clusters of two blocks, also called a thread block pair, or CTA pair, we can configure a special launch attribute, as shown here:在 CUDA 中,这通过扩展的启动 API 完成,该 API 允许指定集群大小。例如,如果我们想要由两个块组成的集群,也称为*线程块对*(thread block pair)或 CTA 对,就可以配置一个特殊的启动属性,如下所示:

Text
// 主机端代码:启动一个线程块集群大小为 2 的内核
cudaLaunchConfig_t config{};
config.gridDim = dim3(128, 1, 1);
config.blockDim = dim3(256, 1, 1);
config.dynamicSmemBytes = 0;
cudaLaunchAttribute attr{};
attr.id = cudaLaunchAttributeClusterDimension;
attr.val.clusterDim.x = 2;
attr.val.clusterDim.y = 1;
attr.val.clusterDim.z = 1;
config.attrs = &attr;
config.numAttrs = 1;
// 如果之后打算使用大于 8 的集群大小,则允许非可移植的集群大小
cudaFuncSetAttribute(MyClusterKernel,
                     cudaFuncAttributeNonPortableClusterSizeAllowed,
                     1);
cudaLaunchKernelEx(&config, MyClusterKernel, args, nullptr);
// Host code: launch a kernel with thread block clusters of size 2
cudaLaunchConfig_t config{};
config.gridDim = dim3(128, 1, 1);
config.blockDim = dim3(256, 1, 1);
config.dynamicSmemBytes = 0;
cudaLaunchAttribute attr{};
attr.id = cudaLaunchAttributeClusterDimension;
attr.val.clusterDim.x = 2;
attr.val.clusterDim.y = 1;
attr.val.clusterDim.z = 1;
config.attrs = &attr;
config.numAttrs = 1;
// Allow non-portable cluster sizes if you intend to use > 8 later
cudaFuncSetAttribute(MyClusterKernel,
                     cudaFuncAttributeNonPortableClusterSizeAllowed,
                     1);
cudaLaunchKernelEx(&config, MyClusterKernel, args, nullptr);

In this host code, we use cudaLaunchKernelEx to launch MyClusterKernel with clusters of 2 thread blocks (so 64 clusters total since 128 blocks ÷ 2 per cluster). The clusterDim is set to 2, meaning each cluster will contain 2 thread blocks. This replaces the traditional <<<gridDim, blockDim>>> syntax for cases where we want cluster-cooperative kernels. Under the hood, the CUDA runtime ensures those paired blocks are scheduled such that they can communicate with DSMEM.在这段主机端代码中,我们使用 cudaLaunchKernelEx 启动 MyClusterKernel,集群大小为 2 个线程块(因此总共 64 个集群,因为 128 个块 ÷ 每个集群 2 个块)。clusterDim 设为 2,表示每个集群包含 2 个线程块。在需要集群协作内核的场景中,这取代了传统的 <<<gridDim, blockDim>>> 语法。在底层,CUDA 运行时确保这些成对的块被调度到能够借助 DSMEM 相互通信的位置。

使用协作组 API 协调线程块集群

Coordinating Thread Block Clusters with Cooperative Groups API

To coordinate work across multiple thread blocks in a cluster, you first obtain a handle to that group of blocks using cooperative_groups::this_cluster(). This cluster handle lets you perform hardware-accelerated cluster-wide barriers—and directly access another block’s shared memory.要在集群内的多个线程块之间协调工作,首先使用 cooperative_groups::this_cluster() 获取该块组的句柄。这个集群句柄让你能够执行由硬件加速的集群级屏障,并直接访问另一个块的共享内存。

This all happens without leaving the kernel or resorting to global-memory flags. Here is an example kernel that sums a local value from each thread block into thread block 0’s shared memory:这一切都无需离开内核,也无需借助全局内存标志位。下面这个示例内核将每个线程块的局部值累加到线程块 0 的共享内存中:

CUDA / C
#include <cuda_runtime.h>
#include <cooperative_groups.h>
namespace cg = cooperative_groups;
__global__ void MyClusterKernel(/* args */) {
    // 1. 为此线程块集群中的所有线程块组建一个集群组
    cg::cluster_group cluster = cg::this_cluster();
    // 2. 在每个块中分配相同的 extern 共享缓冲区
    extern __shared__ int shared_buffer[];
    // 3. 每个块获取自己的 rank 和集群总大小
    int clusterRank = cluster.block_rank();
    int clusterSize = cluster.num_blocks();
    // 4. 初始化本块对应的那部分共享内存(一个简单的局部求和)
    int localSum = threadIdx.x;
    shared_buffer[threadIdx.x] = localSum;
    // 5. 在集群内所有块之间设置屏障;任何块都不会继续执行,直到
    //    每个块都到达此点并已写入其 shared_buffer。
    cluster.sync();
    // 6. 将指针映射到块 0 的 shared_buffer,以便块可以向其中写入
    // 由 cluster.map_shared_rank 返回的指针
    // 指向远程 CTA 的共享内存
    // 并支持远程原子操作
    // 以及其中的内存操作
    // 该集群
    // 在明确定义的位置将更新与 cluster.sync 配对
    int* remote_buffer = cluster.map_shared_rank(shared_buffer, 0);
#include <cuda_runtime.h>
#include <cooperative_groups.h>
namespace cg = cooperative_groups;
__global__ void MyClusterKernel(/* args */) {
    // 1. Form a cluster group for all thread blocks in this thread block cluster
    cg::cluster_group cluster = cg::this_cluster();
    // 2. Allocate the same extern shared buffer in each block
    extern __shared__ int shared_buffer[];
    // 3. Each block learns its rank and the total cluster size
    int clusterRank = cluster.block_rank();
    int clusterSize = cluster.num_blocks();
    // 4. Initialize this block’s portion of shared memory (a simple local sum)
    int localSum = threadIdx.x;
    shared_buffer[threadIdx.x] = localSum;
    // 5. Barrier across all blocks in the cluster; no block proceeds until
    //    every block reaches this point and has written its shared_buffer.
    cluster.sync();
    // 6. Map a pointer to block 0’s shared_buffer so block can write into it
    // Pointers returned by cluster.map_shared_rank
    // refer to the remote CTA’s shared memory
    // and support remote atomics
    // and memory operations within
    // the cluster
    // Pair updates with cluster.sync at well defined point
    int* remote_buffer = cluster.map_shared_rank(shared_buffer, 0);
Text
    if (clusterRank != 0) {
        // 7. 非零线程块以原子方式将其本地 shared_buffer[0] 累加到
        //    rank 0 的缓冲区。该 atomicAdd 在片上经由 DSMEM 路由
        //    (不经过 DRAM。)
        atomicAdd(&remote_buffer[0], shared_buffer[0]);
    }
    // 8. 另一个集群范围的屏障确保所有原子加均已完成
    cluster.sync();
    // 9. 最后,线程块 0(rank 0,thread 0)可以读取合并后的结果
    if (clusterRank == 0 && threadIdx.x == 0) {
        printf("Combined sum in cluster[0]: %d\n", shared_buffer[0]);
    }
}
    if (clusterRank != 0) {
        // 7. Nonzero blocks atomically add their local shared_buffer[0] into
        //    rank 0’s buffer. This atomicAdd is routed on‐chip over DSMEM
        //    (not through DRAM.)
        atomicAdd(&remote_buffer[0], shared_buffer[0]);
    }
    // 8. Another cluster‐wide barrier ensures all atomic adds have completed
    cluster.sync();
    // 9. Finally, block 0 (rank 0, thread 0) can read the combined result
    if (clusterRank == 0 && threadIdx.x == 0) {
        printf("Combined sum in cluster[0]: %d\n", shared_buffer[0]);
    }
}

Here, we obtain a cluster handle by calling cg::this_cluster(). This returns a cluster_group object that represents exactly those blocks launched together as one thread block cluster. The runtime guarantees that every thread block in this cluster group is resident simultaneously—otherwise, the launch will fail.这里,我们通过调用 cg::this_cluster() 获得集群句柄。它返回一个 cluster_group 对象,恰好代表那些作为一个线程块集群一起启动的块。运行时保证该集群组中的每个线程块同时驻留——否则启动就会失败。

A uniform shared memory allocation is implicitly provided. Each block in the cluster must declare the same size for extern __shared__ int shared_buffer[]. The DSMEM hardware will then logically combine each block’s shared memory into one virtual address space. And because all blocks reserve the same shared_buffer size, a pointer to the buffer of thread block 0, or rank 0, will coordinate properly across SMs.系统会隐式提供统一的共享内存分配。集群中的每个块都必须为 extern __shared__ int shared_buffer[] 声明相同的大小。随后 DSMEM 硬件会将各块的共享内存逻辑上合并为一个虚拟地址空间。并且由于所有块预留的 shared_buffer 大小相同,指向线程块 0(即 rank 0)缓冲区的指针就能在多个 SM 之间正确协调。

In the previous example code, we also see that cluster.sync() is used as a cluster-wide barrier across thread blocks in the cluster. Unlike __syncthreads(), which only synchronizes threads within a single block, cluster.sync() synchronizes all threads in every block of this cluster.在前面的示例代码中,我们还看到 cluster.sync() 被用作集群内跨线程块的集群级屏障。与仅同步单个块内线程的 __syncthreads() 不同,cluster.sync() 同步该集群中每个块内的所有线程。

This barrier is implemented in hardware and generally has lower latency than a grid-level synchronization, because it coordinates only the blocks in the cluster. This means you can synchronize blocks frequently with minimal impact—as long as they are within a cluster.该屏障由硬件实现,通常比网格级同步具有更低的时延,因为它只协调集群内的块。这意味着只要这些块位于同一集群内,你就可以频繁地同步它们,而影响极小。

And there is no risk of deadlock caused by missing blocks since CUDA enforces that the grid launches properly and fits on the GPU. As such, there is no scenario in which some blocks reach the barrier and then hang forever waiting for “missing” blocks that never started. All blocks are already present, so the barrier completes cleanly once every block calls it.而且不存在因缺少线程块而导致死锁的风险,因为 CUDA 会强制要求网格正确启动并完整放入 GPU。因此,不会出现某些线程块到达屏障后永远挂起、等待那些从未启动的“缺失”线程块的情形。所有线程块都已就位,所以只要每个线程块都调用该屏障,屏障就会顺利完成。

In the previous code block, you see that cross-block shared memory access happens using map_shared_rank(). After the first barrier, every block’s shared_buffer[] is initialized. To get a pointer into block 0’s shared memory, the other blocks call int* remote_buffer = cluster.map_shared_rank(shared_buffer, 0), which specifies the thread block id (0) in the cluster.在上一个代码块中可以看到,跨线程块的共享内存访问通过 map_shared_rank() 完成。在第一个屏障之后,每个线程块的 shared_buffer[] 都已初始化。为了获得指向线程块 0 共享内存的指针,其他线程块调用 int* remote_buffer = cluster.map_shared_rank(shared_buffer, 0),其中指定了集群内的线程块 id (0)。

The GPU hardware automatically translates that pointer so any load or store goes over the on-chip DSMEM network rather than out to global DRAM. As such, any write operation into remote_buffer—either an atomic or regular write—will update block 0’s shared memory directly on-chip. It’s worth highlighting that the performance characteristics of remote DSMEM accesses differ from local shared memory. Remote loads and stores benefit from coalesced, aligned 32-byte segments.GPU 硬件会自动转换该指针,使任何加载或存储都经由片上 DSMEM 网络,而不是访问全局 DRAM。因此,对 remote_buffer 的任何写操作——无论是原子写还是普通写——都会直接在片上更新线程块 0 的共享内存。值得强调的是,远程 DSMEM 访问的性能特征与本地共享内存不同。远程加载和存储受益于合并的、对齐的 32 字节段。

For instance, in the previous code block, when blocks 1...n need to add their local values into block ID 0’s shared buffer, they call atomicAdd(&remote_buffer[0], shared_buffer[0]). And because remote_buffer was obtained with clus ter.map_shared_rank(), the atomicAdd goes directly into block 0’s SMEM over the on-chip DSMEM network. In other words, there is no round trip to DRAM or L2 cache. Every write happens at on-chip speeds.例如,在前面的代码块中,当块 1...n 需要把各自的本地值累加到块 ID 0 的共享缓冲区时,它们调用 atomicAdd(&remote_buffer[0], shared_buffer[0])。而由于 remote_buffer 是通过 clus ter.map_shared_rank() 获得的,该 atomicAdd 会经由片上 DSMEM 网络直接写入块 0 的 SMEM。换言之,不存在到 DRAM 或 L2 缓存的往返。每次写入都以片上速度完成。

In the previous code example, you see that once all blocks have performed their atom icAdd, a second cluster.sync() ensures that block 0 sees the complete sum in its own shared_buffer[0] before proceeding.在上一个代码示例中可以看到,一旦所有块都完成了各自的 atom icAdd,第二次 cluster.sync() 便会确保块 0 在继续执行之前,能在自己的 shared_buffer[0] 中看到完整的求和结果。

In short, cooperative_groups::this_cluster() plus cluster.sync() and cluster .map_shared_rank() give you a simple, efficient way to synchronize and share data across multiple thread blocks in a thread block cluster. All of these operations occur on-chip, avoid global-memory round trips, and enable fine-grained cooperation among blocks. This combination of thread block clusters and DSMEM provides much higher performance than any global-memory fallback or manual atomic-flag approach.简而言之,cooperative_groups::this_cluster() 加上 cluster.sync() 与 cluster .map_shared_rank(),为你提供了一种简单高效的方式,可在线程块集群内跨多个线程块进行同步与数据共享。所有这些操作都在片上完成,避免了全局内存往返,并实现了块间的细粒度协作。线程块集群与 DSMEM 的这种组合,其性能远高于任何回退到全局内存的方案或手工原子标志方案。

线程块对

Thread Block Pair

With modern NVIDIA, GPUs let you coschedule exactly two thread blocks in a cluster, or a thread block pair (aka CTA pair), across SMs within a single GPC. By grouping thread blocks in a cluster (e.g., a 2-block cluster, as shown in Figure 10-13), kernels that share data can use TMA to move tiles into each block’s shared memory.现代 NVIDIA GPU 允许在一个集群中精确协同调度两个线程块,即线程块对(又称 CTA 对),它们分布在单个 GPC 内的多个 SM 上。通过将线程块组织成集群(例如图 10-13 所示的 2 块集群),需要共享数据的内核可以使用 TMA 将瓦片搬入每个线程块的共享内存。

图 10-13
图 10-13线程块对组合两个线程块Thread block pair combines two thread blocks
图 10-14
图 10-14线程块对(又称 CTA 对)加载瓦片作为 A * B 矩阵乘法的操作数(来源:https://oreil.ly/kEZsv)Thread block pair (aka CTA pair) loading tiles as operands for an A * B matrix multiply (source: https://oreil.ly/kEZsv)

A single thread block might lack the registers or shared-memory capacity to process a very large tile (for example, a 256 × 256 matrix subtile) by itself. By pairing two thread blocks on nearby SMs within a GPC and using DSMEM, those two blocks can split the work on one large tile yet share data through a unified shared-memory region, as shown in Figure 10-14.单个线程块可能缺少寄存器或共享内存容量,无法独自处理非常大的瓦片(例如 256 × 256 的矩阵子瓦片)。通过在一个 GPC 内相邻的 SM 上配对两个线程块并使用 DSMEM,这两个线程块可以拆分同一个大瓦片上的工作,同时通过统一的共享内存区域共享数据,如图 10-14 所示。

Here, each thread block in the pair can load a fraction of the operand tile (e.g., 128 × 16) into its on-chip SMEM for the matrix multiply. In addition, each thread block holds part of the accumulator (e.g., 128 × 256) in Tensor Memory (TMEM). This在这里,配对中的每个线程块都可以将操作数瓦片的一部分(例如 128 × 16)加载到其片上 SMEM 中,用于矩阵乘法。此外,每个线程块在张量内存(TMEM)中保存累加器的一部分(例如 128 × 256)。这

图 10-15
图 10-15使用张量核心和 TMEM 的线程块对Thread block pair with Tensor Cores and TMEM
图 10-16
图 10-16在线程块集群中使用 DSMEM 在 SM 之间共享数据Sharing data between SMs in a thread block cluster using DSMEM

allows the two thread blocks in the CTA pair to collaborate on a single tile, as shown in Figure 10-15.使 CTA 对中的两个线程块能够协作处理单个瓦片(tile),如图 10-15 所示。

Thread block pairs allow larger matrix multiplies that span two physical SMs. Using a pair of thread blocks effectively doubles the tile size since each SM handles half of the tile’s data.线程块对允许跨两个物理 SM 执行更大规模的矩阵乘法。使用一对线程块实际上使瓦片大小翻倍,因为每个 SM 处理该瓦片一半的数据。

The SM hardware shares operand data between the SMs using DSMEM. DSMEM reduces duplicate loads, improves data reuse, and increases arithmetic intensity. Figure 10-16 shows this data sharing between SMs in a thread block cluster using DSMEM.SM 硬件使用 DSMEM 在各 SM 之间共享操作数数据。DSMEM 减少重复加载、提升数据复用并提高算术强度。图 10-16 展示了线程块集群中各 SM 之间借助 DSMEM 进行的数据共享。

CUDA provides enhanced multicast barrier and synchronization primitives for these multi-SM operations. The pair synchronizes with each other using the lightweight cluster.sync() call. This way, when one SM finishes using a tile, the adjacent SM in the CTA pair can consume it.CUDA 为这些多 SM 操作提供了增强的多播屏障与同步原语。该线程块对通过轻量级的 cluster.sync() 调用相互同步。这样,当一个 SM 用完某个瓦片后,CTA 对中相邻的 SM 便可消费它。

Each tile is fed into the pair of thread blocks without redundant global-memory transactions. This better occupies otherwise idle registers, shared-memory banks, and Tensor Cores. And a thread block pair achieves higher Tensor Core utilization by doubling the threads and shared memory available for one tile.每个瓦片被送入这对线程块时不会产生冗余的全局内存事务。这更好地占用了原本空闲的寄存器、共享内存存储体和张量核心。并且,线程块对通过为一个瓦片投入双倍的线程与共享内存,实现了更高的张量核心利用率。

Any boost in performance will happen transparently. From a programmer’s perspective, you simply launch a cluster of two blocks and treat it as one large thread block split across two thread blocks. In CUTLASS, for instance, this is exposed as a 2-SM UMMA (Unified MMA) GEMM operation that uses DSMEM and cluster barriers for inter-CTA data sharing. You simply request two SMs for a single GEMM tile, and CUTLASS handles the cluster launch, DSMEM setup, and interthread block synchronization automatically.任何性能提升都是透明发生的。从程序员的角度看,你只需启动一个包含两个线程块的集群,并将其视为拆分到两个线程块上的一个大线程块。例如在 CUTLASS 中,这以 2-SM UMMA(Unified MMA)GEMM 操作的形式暴露出来,该操作使用 DSMEM 和集群屏障实现 CTA 间的数据共享。你只需为单个 GEMM 瓦片申请两个 SM,CUTLASS 便会自动处理集群启动、DSMEM 设置以及线程块间的同步。

If you request a tile size that a single thread block can’t handle (say, 128 × 128 for FP16 Tensor Core operations), CUTLASS will automatically allocate two thread blocks as a pair. The SMs split the work and rely on DSMEM + cluster.sync() under the hood to share the tile. This way, you can do pairwise overlapping of UMMA operations with DSMEM.如果你请求的瓦片大小是单个线程块无法处理的(例如 FP16 张量核心操作下的 128 × 128),CUTLASS 会自动分配两个线程块组成一对。两个 SM 分担工作,并在底层依赖 DSMEM + cluster.sync() 来共享该瓦片。这样,你就可以借助 DSMEM 实现 UMMA 操作的成对重叠。

In short, thread block pairs and DSMEM provide parallelism in the form of multi-SM cooperation. They provide fast, on-chip data sharing and synchronization across thread blocks. This eliminates many scenarios that previously required global-memory handoffs or additional kernel launches. This benefits many multithread block algorithms and simplifies their implementation.简而言之,线程块对与 DSMEM 以多 SM 协作的形式提供并行性。它们提供跨线程块的快速片上数据共享与同步。这消除了许多以往需要全局内存交接或额外内核启动的场景。这有利于许多多线程块算法,并简化了它们的实现。

使用线程块集群减少全局内存流量

Reducing Global Memory Traffic with Thread Block Clusters

From a performance perspective, DSMEM can significantly cut down redundant global-memory traffic and enable higher effective bandwidth. For instance, in a tiled GEMM where multiple thread blocks within a cluster share chunks of the A or B matrix, one block can load a tile from global memory and multicast the tile to other blocks in the cluster using the TMA.从性能角度看,DSMEM 能显著削减冗余的全局内存流量,并带来更高的有效带宽。例如,在分块 GEMM 中,集群内的多个线程块共享 A 矩阵或 B 矩阵的数据块,此时一个线程块可以从全局内存加载一个瓦片,并使用 TMA 将该瓦片多播到集群中的其他线程块。

The TMA engine supports a multicast copy mode that feeds directly into DSMEM when thread blocks belong to the same cluster. A single TMA transfer from global memory can place data into each participating block’s shared memory simultaneously. This avoids redundant DRAM fetches.当线程块属于同一集群时,TMA 引擎支持一种直接写入 DSMEM 的多播拷贝模式。一次从全局内存发起的 TMA 传输即可同时将数据放入每个参与线程块的共享内存。这避免了冗余的 DRAM 读取。

图 10-17
图 10-17对于这四个(2 × 2)线程集群,每个瓦片仅加载一次,并被多播到每个集群中所有 CTA 的共享内存(来源:https://oreil.ly/kEZsv)For these four (2 × 2) thread clusters, each tile is loaded once and multicast into the shared memory of all CTAs in each cluster (source: https://oreil.ly/kEZsv)

With TMA multicast, the GPU ensures that L2-cached data is broadcast into the shared memory of each cluster member (thread block) in one pass. As a result, you avoid repeated global-memory loads of the same tile. This improves bandwidth utilization and shrinks DRAM traffic—especially when many blocks need the same input, as shown in Figure 10-17.借助 TMA 多播,GPU 确保 L2 缓存中的数据在一次传输中广播到每个集群成员(线程块)的共享内存。因此,你避免了同一瓦片的重复全局内存加载。这提升了带宽利用率并减少了 DRAM 流量——尤其是当许多线程块需要相同输入时,如图 10-17 所示。

Here, the TMA engine performs a single multicast from global memory into DSMEM, broadcasting the tile to every thread block’s SMEM in the cluster and eliminating redundant DRAM reads. TMA multicast is configured through a tensor-map descriptor and issued as cp.async.bulk.tensor targeting shared::cluster. Fortunately, higher-level libraries like CUTLASS/cuTe and Triton generate these multicast operations, tensor-map descriptors, and bulk tensor copies for you. Specifically, these libraries can issue the relevant PTX instructions including cp.async.bulk.tensor with a tensor-map operand. They can issue PTX instructions that target shared::cluster for multicast. In these cases, the hardware delivers the tile to each thread block’s SMEM.在这里,TMA 引擎执行一次从全局内存到 DSMEM 的多播(multicast),将该数据块(tile)广播到集群中每个线程块的 SMEM,从而消除冗余的 DRAM 读取。TMA 多播通过 tensor-map 描述符进行配置,并以面向 shared::cluster 的 cp.async.bulk.tensor 形式发出。幸运的是,CUTLASS/cuTe 和 Triton 等高层库会为你生成这些多播操作、tensor-map 描述符和批量张量拷贝。具体而言,这些库可以发出相关的 PTX 指令,包括带 tensor-map 操作数的 cp.async.bulk.tensor。它们可以发出面向 shared::cluster 的 PTX 指令以实现多播。在这些情况下,硬件会将数据块投递到每个线程块的 SMEM。

Another common use case is multiblock reductions or scans. Instead of having each block write out its partial sum or scan result to global memory and then launching a separate kernel to combine them, thread blocks in a thread block cluster can write their partial results into DSMEM.另一个常见用例是多块归约或扫描。与其让每个块将自己的部分和或扫描结果写入全局内存、再启动一个单独的内核来合并它们,线程块集群中的线程块可以将各自的部分结果写入 DSMEM。

Within that same cluster, you can perform a fast on-chip reduction or prefix-sum across these partial results using shared memory and a cluster-level barrier, clus ter.sync(). Only the final result needs to go to global memory. This greatly reduces global memory reads and writes.在同一集群内,你可以使用共享内存和集群级屏障 clus ter.sync(),对这些部分结果执行快速的片上归约或前缀和。只有最终结果需要写入全局内存。这大幅减少了全局内存的读取和写入。

表 10-4DSMEM 对双块工作负载的性能影响Performance impact of DSMEM on a two-block workload
Metric指标Two independent CTAs (no DSMEM)两个独立的 CTA(无 DSMEM)CTA pair with DSMEM (cluster of 2)使用 DSMEM 的 CTA 对(规模为 2 的集群)
Global load transactions全局加载事务数2× (each thread block loads tile)2×(每个线程块各自加载数据块)1× (tile loaded once)1×(数据块仅加载一次)
L2 cache hit rateL2 缓存命中率50%50%85%85%
Inter-CTA data reuseCTA 间数据复用N/A (no reuse)不适用(无复用)Significant (tile reused by CTA 1)显著(数据块被 CTA 1 复用)
Effective DRAM BW per CTA每个 CTA 的有效 DRAM 带宽300 GB/s300 GB/s50% less 150 GB/s减少 50%,150 GB/s
Kernel time (relative)内核时间(相对值)1.0×1.0×0.6× (40% speedup)0.6×(加速 40%)

By combining thread block clusters, DSMEM, and TMA multicast, you can build multistage, fine-grained pipelines that keep most data exchanges on-chip. Whether you are sharing tiles for GEMM, accumulating partial sums for a reduction, or performing a multiblock scan, these mechanisms let you minimize HBM round trips and maximize arithmetic intensity.通过组合线程块集群、DSMEM 和 TMA 多播,可以构建多级、细粒度的流水线,将大部分数据交换保留在片上。无论你是为 GEMM 共享数据块(tile)、为归约累加部分和,还是执行多块扫描,这些机制都能让你最小化 HBM 往返并最大化算术强度。

Consider two thread blocks, CTA 0 and CTA 1, that both need the same tile of matrix A. Without DSMEM, each CTA issues its own global-memory load, which wastes DRAM bandwidth by fetching the identical data twice.考虑两个线程块 CTA 0 和 CTA 1,它们都需要矩阵 A 的同一个数据块。在没有 DSMEM 的情况下,每个 CTA 各自发起全局内存加载,由于重复取回完全相同的数据而浪费 DRAM 带宽。

With DSMEM, however, CTA 0 loads the tile once into its shared memory and then multicasts it over the on-chip DSMEM network so that CTA 1 can read it directly from shared memory. If more than two CTAs form a cluster, the same tile will be shared across all thread blocks in the cluster, but this still requires only one global HBM load. Table 10-4 shows a comparison using a cluster of two thread blocks versus two independent thread blocks.然而使用 DSMEM 时,CTA 0 将该数据块一次性加载到自己的共享内存,然后通过片上 DSMEM 网络对其进行多播,使 CTA 1 可以直接从共享内存读取。如果两个以上的 CTA 组成集群,同一个数据块将在集群内的所有线程块之间共享,但这仍然只需要一次全局 HBM 加载。表 10-4 展示了使用由两个线程块组成的集群与使用两个独立线程块的对比。

Here, we see a 40% speedup in kernel execution when using a thread block pair (aka CTA pair) with DSMEM. This is consistent with the bandwidth savings as well. The DRAM bandwidth per thread block drops by 50% from 300 GB/s to 150 GB/s since each tile is fetched only once. The L2 hit rate jumps from 50% to 85% since the TMA hardware is performing a multicast copy on-chip—and therefore avoiding multiple loads from global memory.在这里,我们看到使用线程块对(又称 CTA 对)配合 DSMEM 时,内核执行速度提升了 40%。这也与带宽的节省相一致。由于每个数据块只取一次,每个线程块的 DRAM 带宽从 300 GB/s 降至 150 GB/s,下降了 50%。L2 命中率从 50% 跃升至 85%,因为 TMA 硬件在片上执行多播复制——从而避免了对全局内存的多次加载。

In short, by multicasting shared data, thread block clusters allow multiple blocks to reuse data at on-chip speeds. This leads to substantial speedups for memory-bound workloads.简而言之,通过对共享数据进行多播,线程块集群允许多个块以片上速度复用数据。这为显存受限型工作负载带来显著的加速。

图 10-18
图 10-18DSMEM 在两个线程块集群之间使用 SM 到 SM 的线程块集群本地网络进行数据交换DSMEM uses an SM-to-SM thread block cluster-local network between two thread block clusters for its data exchange

It’s important to note that both DSMEM and L2 operate in parallel. This provides two “lanes” for interblock data sharing. This dual-path design combines DSMEM’s ultrafast cluster-wide communication with L2’s broader caching coverage. In other words, DSMEM access is bypassing the L2 cache for cluster-local addresses. Instead, DSMEM uses the dedicated SM-to-SM network, as shown in Figure 10-18.需要注意的是,DSMEM 与 L2 是并行工作的。这为块间数据共享提供了两条“通道”。这种双路径设计将 DSMEM 超快的集群范围通信与 L2 更广泛的缓存覆盖结合起来。换言之,对于集群本地地址,DSMEM 访问会绕过 L2 缓存,转而使用专用的 SM 到 SM 网络,如图 10-18 所示。

Remote DSMEM accesses are routed using the thread block cluster interconnect and are distinct from global-memory traffic. For example, when a thread block pulls a tile using DSMEM, it uses low-latency shared-memory transfers. This ensures that interthread block data sharing in a thread block cluster is as fast as on-chip shared memory.远程 DSMEM 访问通过线程块集群互连进行路由,与全局内存流量相互独立。例如,当某个线程块使用 DSMEM 拉取一个数据块时,它使用的是低时延的共享内存传输。这确保了线程块集群内的线程块间数据共享与片上共享内存一样快。

If a tile is not present in DSMEM, it must be fetched from global memory (following the normal cache hierarchy that may hit in L2, for example). This way, if a tile has already been evicted from DSMEM, for instance, the thread block can still retrieve the tile data from the L2 cache instead of going all the way back to global DRAM.如果某个数据块(tile)不在 DSMEM 中,就必须从全局内存中获取(遵循常规的缓存层级,例如可能命中 L2)。这样一来,例如某个 tile 已被从 DSMEM 中逐出时,线程块仍可从 L2 缓存中取回该 tile 的数据,而不必一路回退到全局 DRAM。

As a result, memory stalls are rare, occupancy remains high, and overall execution runs significantly faster than the unoptimized, two-independent-thread-blocks, non-DSMEM implementation.因此,内存停顿很少发生,占用率保持较高,整体执行速度显著快于未优化、使用两个独立线程块且不使用 DSMEM 的实现。

使用线程块集群设计高效算法

Designing Efficient Algorithms with Thread Block Clusters

Thread block clusters enable new strategies for parallelizing workloads that previously required global memory communication or multiple kernel launches. For example, imagine a large matrix multiplication in which the output tile is too large to be handled by a single thread block due to shared memory limits.线程块集群为并行化工作负载带来了新的策略,而这些工作负载此前需要借助全局内存通信或多次内核启动。例如,设想一个大型矩阵乘法,其输出 tile 因共享内存限制而过大,无法由单个线程块处理。

In the past, you might split the work across two thread blocks, but then each block would have to exchange partial results with global memory, which is relatively slow and inefficient. Otherwise, you’d have to launch a separate reduction kernel to combine the outputs of the two thread blocks.过去,你可能会把工作拆分到两个线程块上,但这样每个块都必须通过全局内存交换部分结果,相对缓慢且低效。否则,你就必须启动一个单独的归约内核来合并两个线程块的输出。

With thread block clusters and DSMEM, those two blocks can form a cluster and directly share a joint region of on-chip shared memory to seamlessly combine their results using hardware-supported primitives. The DSMEM hardware allows an SM to perform loads/stores/atomics to another SM’s shared memory through a fast network.借助线程块集群与 DSMEM,这两个块可以组成一个集群,直接共享一块联合的片上共享内存区域,利用硬件支持的原语无缝合并各自的结果。DSMEM 硬件允许一个 SM 通过高速网络对另一个 SM 的共享内存执行加载/存储/原子操作。

It’s important to carefully design algorithms to use thread block clusters effectively. Synchronization overhead is low, but it is still nonzero. Performing very fine-grained data sharing might not pay off.要有效利用线程块集群,仔细设计算法十分重要。同步开销很低,但并非为零。进行过于细粒度的数据共享可能得不偿失。

Thread block cluster barriers work like warp-level intrinsics (e.g., __shfl_sync()), introduced in Chapter 7, in that every participant must arrive at the synchronization point together. When you call cluster.sync(), all thread blocks in the cluster must reach that line before any block can continue.线程块集群屏障的工作方式类似于第 7 章介绍的线程束级内建函数(例如 __shfl_sync()):每个参与者都必须一起到达同步点。当你调用 cluster.sync() 时,集群中的所有线程块都必须执行到该行,任何块才能继续。

If one block finishes its work early, it simply waits. If a block never reaches the barrier, because its threads took a different branch, for instance, the entire cluster will deadlock.如果一个块提前完成工作,它只需等待。如果某个块始终未到达屏障——例如其线程走了不同的分支——整个集群就会死锁。

In other words, just as warp intrinsics demand that all threads in a warp execute the same instruction path to avoid divergence, thread block clusters demand that all blocks follow the same control flow up to each cluster.sync() call.换言之,正如线程束内建函数要求一个线程束中的所有线程执行相同的指令路径以避免发散,线程块集群要求所有块在到达每个 cluster.sync() 调用之前遵循相同的控制流。

Because of this lockstep requirement, fine-grained data sharing between blocks must be balanced against the overhead and risk of deadlock. Synchronization itself is inex-pensive when done correctly, but if even one block bypasses or delays reaching a clus ter.sync(), performance may be impacted—or, even worse, the kernel might hang.由于这种锁步要求,线程块之间的细粒度数据共享必须在开销与死锁风险之间进行权衡。同步本身若执行正确,开销并不高,但只要有一个线程块绕过或延迟到达 clus ter.sync(),性能就可能受到影响——更糟的是,内核可能会挂起。

You should structure code such that every block arrives at each barrier in unison— just as every thread in a warp must arrive at the barrier with warp-level intrinsics. This is essential to take full advantage of the thread block clusters’ low-latency, on-chip communication without falling into deadlocks.应当这样组织代码:每个块都齐步到达每个屏障——正如使用线程束级内建函数时,线程束中的每个线程都必须到达该屏障。要充分利用线程块集群的低时延片上通信而不陷入死锁,这一点至关重要。

Typically, thread block clusters perform the best when each block has a sizable amount of work that can run independently until a synchronization or data exchange point is needed. Once the synchronization happens and the data is transferred on-chip, the thread blocks can continue processing.通常,当每个块都有相当数量的工作、且这些工作可以独立运行直到需要同步或数据交换点时,线程块集群的性能最佳。一旦同步完成、数据在片上完成传输,各线程块便可继续处理。

Thread block clusters are especially effective for block-sparse matrix operations— common in sparse attention, model pruning, and compression in LLMs. In these cases, the blocks process different nonzero regions and share their boundary data.线程块集群对块稀疏矩阵运算尤为有效——这类运算常见于稀疏注意力、模型剪枝以及 LLM 中的压缩。在这些场景中,各块处理不同的非零区域并共享其边界数据。

Thread block clusters are also useful for multiphase reductions such as the softmax and normalization steps in transformer layers. These require a final combine step using the partial results of each thread block in the thread block cluster.线程块集群对于多阶段归约(例如 Transformer 层中的 softmax 与归一化步骤)同样有用。这些操作需要一个最终的合并步骤,使用线程块集群中每个线程块的部分结果。

More generally, large GEMMs benefit from thread block clusters when they exceed the resources of a single block. GEMMs, of course, are central to the transformer attention and multilayer perceptron (MLP) layers and embedding lookups common in modern LLMs.更一般地说,当大型 GEMM 超出单个块的资源时,线程块集群能为其带来收益。当然,GEMM 是现代 LLM 中常见的 Transformer 注意力层、多层感知机(MLP)层以及嵌入查找的核心。

Larger cluster sizes can reduce overall occupancy, however. For instance, a cluster of 16 blocks might monopolize 16 SMs for one task. This could leave fewer SMs for other tasks needed by that kernel launch. It’s recommended to start with small clusters, two- or four-thread block clusters, unless a bigger cluster is needed for special cases. As always, you should profile with your specific workload to confirm that sharing on-chip resources across thread block clusters outweighs the potential loss of parallelism.然而,更大的集群规模会降低整体占用率。例如,一个由 16 个块组成的集群可能会为单个任务独占 16 个 SM。这会为该内核启动所需的其他任务留下更少的 SM。建议从较小的集群开始,即 2 个或 4 个线程块的集群,除非特殊情况需要更大的集群。一如既往,你应针对自己的具体工作负载进行性能剖析,以确认跨线程块集群共享片上资源所带来的收益超过潜在的并行度损失。

使用线程块集群的线程束专用化

Warp Specialization with Thread Block Clusters

Let’s now revisit warp specialization using a thread block cluster with the CUDA Pipeline API. We use a thread-block scoped pipeline as the cluster leader to stage the copies and perform cluster-wide barriers with DSMEM. This way, every block in the cluster consumes the same input tiles without reloading them from global memory.现在我们重新审视结合线程块集群与 CUDA Pipeline API 的线程束专用化。我们使用线程块作用域的流水线作为集群领导者,来暂存拷贝并通过 DSMEM 执行集群范围的屏障。这样,集群中的每个块都能消费相同的输入分块,而无需从全局内存重新加载它们。

Roles are assigned per warp inside each block. The leader block’s loader warp performs the cooperative copies into its shared memory once per tile. After a cluster-wide barrier publishes those tiles, every block uses a compute warp to read the leader’s tiles through DSMEM and compute a disjoint band of rows. A storer warp in each block writes that band back to global memory. The block-scoped pipeline is used only by the leader for the asynchronous copies. Other blocks do not wait on that pipeline. Here is the code:角色在每个块内部按线程束分配。领导者块的加载线程束对每个分块执行一次协作拷贝,将其写入自己的共享内存。在集群范围的屏障发布这些分块之后,每个块使用一个计算线程束通过 DSMEM 读取领导者的分块,并计算一段互不重叠的行。每个块中的存储线程束将该行段写回全局内存。块作用域的流水线仅由领导者用于异步拷贝。其他块不会等待该流水线。代码如下:

CUDA / C
// 跨线程块集群的线程束专用化
// 使用 DSMEM 和线程块作用域的流水线
#include <cuda/pipeline>
#include <cooperative_groups.h>
#include <algorithm>
namespace cg = cooperative_groups;
#define TILE_SIZE 128
#define TILE_ELEMS (TILE_SIZE * TILE_SIZE)
// 从 DSMEM 源计算 TILE_SIZE×TILE_SIZE 乘积的一个行带。
// 每个 lane 在 32 路条带化循环中处理行 [row_begin, row_end)。
__device__ void compute_rows_from_ds(const float* __restrict__ A_src,
// Warp specialization across a thread-block cluster
// using DSMEM and a block-scoped pipeline
#include <cuda/pipeline>
#include <cooperative_groups.h>
#include <algorithm>
namespace cg = cooperative_groups;
#define TILE_SIZE 128
#define TILE_ELEMS (TILE_SIZE * TILE_SIZE)
// Compute a band of rows of the TILE_SIZE×TILE_SIZE product from DSMEM sources.
// Each lane processes rows [row_begin, row_end) in a 32-way striped loop.
__device__ void compute_rows_from_ds(const float* __restrict__ A_src,
Text
                                     const float* __restrict__ B_src,
                                     float* __restrict__ C_dst,
                                     int row_begin, int row_end,
                                     int lane_id) {
    for (int row = row_begin + lane_id; row < row_end; row += warpSize) {
        for (int col = 0; col < TILE_SIZE; ++col) {
            float acc = 0.0f;
            #pragma unroll
            for (int k = 0; k < TILE_SIZE; ++k) {
                acc += A_src[row * TILE_SIZE + k] * B_src[k * TILE_SIZE + col];
            }
            C_dst[row * TILE_SIZE + col] = acc;
        }
    }
}
extern "C"
__global__ void warp_specialized_cluster_pipeline(
    const float* __restrict__ A_global,
    const float* __restrict__ B_global,
    float* __restrict__ C_global,
    int numTiles) {
    thread_block cta = this_thread_block();
    cluster_group  cluster = this_cluster();
    extern __shared__ float shared_mem[];
    float* A_tile_local = shared_mem;
    float* B_tile_local = A_tile_local + TILE_ELEMS;
    float* C_tile_local = B_tile_local + TILE_ELEMS;
    // 仅由集群 leader 使用的线程块作用域流水线
    // 用于暂存异步拷贝
    __shared__
    cuda::pipeline_shared_state<cuda::thread_scope_block, 2> pipe_state;
    auto pipe = cuda::make_pipeline(cta, &pipe_state);
    const int lane_id = threadIdx.x & 31;
    const int warp_id = threadIdx.x >> 5;
    const int cluster_rank      = cluster.block_rank();
    const dim3 cluster_dims     = cluster.dim_blocks();
    const int  blocks_in_cluster = cluster_dims.x * cluster_dims.y *
                                   cluster_dims.z;
    // 沿 x 方向的一维集群排布;
    // 每次迭代每个集群处理一个 tile
    auto loader = cooperative_groups::tiled_partition<32>(cta);
    for (int tile = blockIdx.x / cluster_dims.x; tile < numTiles;
         tile += gridDim.x / cluster_dims.x) {
        const size_t offset = static_cast<size_t>(tile) * TILE_ELEMS;
        // leader 线程块的加载线程束为整个集群暂存一次 A 和 B
                                     const float* __restrict__ B_src,
                                     float* __restrict__ C_dst,
                                     int row_begin, int row_end,
                                     int lane_id) {
    for (int row = row_begin + lane_id; row < row_end; row += warpSize) {
        for (int col = 0; col < TILE_SIZE; ++col) {
            float acc = 0.0f;
            #pragma unroll
            for (int k = 0; k < TILE_SIZE; ++k) {
                acc += A_src[row * TILE_SIZE + k] * B_src[k * TILE_SIZE + col];
            }
            C_dst[row * TILE_SIZE + col] = acc;
        }
    }
}
extern "C"
__global__ void warp_specialized_cluster_pipeline(
    const float* __restrict__ A_global,
    const float* __restrict__ B_global,
    float* __restrict__ C_global,
    int numTiles) {
    thread_block cta = this_thread_block();
    cluster_group  cluster = this_cluster();
    extern __shared__ float shared_mem[];
    float* A_tile_local = shared_mem;
    float* B_tile_local = A_tile_local + TILE_ELEMS;
    float* C_tile_local = B_tile_local + TILE_ELEMS;
    // Block-scoped pipeline used only by the cluster leader
    // to stage asynchronous copies
    __shared__
    cuda::pipeline_shared_state<cuda::thread_scope_block, 2> pipe_state;
    auto pipe = cuda::make_pipeline(cta, &pipe_state);
    const int lane_id = threadIdx.x & 31;
    const int warp_id = threadIdx.x >> 5;
    const int cluster_rank      = cluster.block_rank();
    const dim3 cluster_dims     = cluster.dim_blocks();
    const int  blocks_in_cluster = cluster_dims.x * cluster_dims.y *
                                   cluster_dims.z;
    // 1D cluster arrangement along x;
    // each iteration processes one tile per cluster
    auto loader = cooperative_groups::tiled_partition<32>(cta);
    for (int tile = blockIdx.x / cluster_dims.x; tile < numTiles;
         tile += gridDim.x / cluster_dims.x) {
        const size_t offset = static_cast<size_t>(tile) * TILE_ELEMS;
        // Leader block’s loader warp stages A and B once for the entire cluster
Text
        if (cluster_rank == 0 && warp_id == 0) {
            pipe.producer_acquire();
            cuda::memcpy_async(loader, A_tile, A_global + offset,
                               cuda::aligned_size_t<32>{TILE_ELEMS*sizeof(float)},
                               pipe);
            cuda::memcpy_async(loader, B_tile, B_global + offset,
                               cuda::aligned_size_t<32>{TILE_ELEMS*sizeof(float)},
                               pipe);
            pipe.producer_commit();
            // 在向集群范围发布之前,使加载对 leader CTA 可见
            // 在发布之前等待已提交的 stage
pipe.consumer_wait();
            pipe.consumer_release();
        }
        // 通过 DSMEM 将 leader 的数据块(tile)发布到每个线程块
        cluster.sync();
        const float* A_src = cluster.map_shared_rank(A_tile_local, 0);
        const float* B_src = cluster.map_shared_rank(B_tile_local, 0);
        // 在集群内将行划分给各线程块
        const int rows_per_block = (TILE_SIZE + blocks_in_cluster - 1)
                                   / blocks_in_cluster;
        const int row_begin = std::min(cluster_rank * rows_per_block, TILE_SIZE);
        const int row_end   = std::min(row_begin + rows_per_block, TILE_SIZE);
        // 计算线程束将本线程块负责的行带写入本地共享内存
        if (warp_id == 1) {
            pipe.producer_acquire();
            compute_rows_from_ds(A_src, B_src, C_tile_local, row_begin, row_end,
                                 lane_id);
            pipe.producer_commit(); // 发布 C_tile_local
        }
        // 存储线程束将本线程块的行写回全局内存
        if (warp_id == 2) {
            pipe.consumer_wait(); // 观察 C_tile_local
            for (int row = row_begin + lane_id; row < row_end; row += warpSize) {
                for (int col = 0; col < TILE_SIZE; ++col) {
                    C_global[offset + row * TILE_SIZE + col] =
                        C_tile_local[row * TILE_SIZE + col];
                }
            }
            pipe.consumer_release();
        }
        // 所有线程块完成此数据块后,leader 才能复用其缓冲区
        cluster.sync();
    }
    // 动态共享内存大小:3 * TILE_ELEMS * sizeof(float)
}
        if (cluster_rank == 0 && warp_id == 0) {
            pipe.producer_acquire();
            cuda::memcpy_async(loader, A_tile, A_global + offset,
                               cuda::aligned_size_t<32>{TILE_ELEMS*sizeof(float)},
                               pipe);
            cuda::memcpy_async(loader, B_tile, B_global + offset,
                               cuda::aligned_size_t<32>{TILE_ELEMS*sizeof(float)},
                               pipe);
            pipe.producer_commit();
            // Make loads visible to leader CTA before publishing cluster-wide
            // wait for committed stage before publishin
pipe.consumer_wait();
            pipe.consumer_release();
        }
        // Publish the leader’s tiles to every block via DSMEM
        cluster.sync();
        const float* A_src = cluster.map_shared_rank(A_tile_local, 0);
        const float* B_src = cluster.map_shared_rank(B_tile_local, 0);
        // Divide rows among blocks in the cluster
        const int rows_per_block = (TILE_SIZE + blocks_in_cluster - 1)
                                   / blocks_in_cluster;
        const int row_begin = std::min(cluster_rank * rows_per_block, TILE_SIZE);
        const int row_end   = std::min(row_begin + rows_per_block, TILE_SIZE);
        // Compute warp produces block’s band of rows into local shared memory
        if (warp_id == 1) {
            pipe.producer_acquire();
            compute_rows_from_ds(A_src, B_src, C_tile_local, row_begin, row_end,
                                 lane_id);
            pipe.producer_commit(); // publish C_tile_local
        }
        // Storer warp writes this block’s rows back to global memory
        if (warp_id == 2) {
            pipe.consumer_wait(); // observe C_tile_local
            for (int row = row_begin + lane_id; row < row_end; row += warpSize) {
                for (int col = 0; col < TILE_SIZE; ++col) {
                    C_global[offset + row * TILE_SIZE + col] =
                        C_tile_local[row * TILE_SIZE + col];
                }
            }
            pipe.consumer_release();
        }
        // All blocks finish this tile before the leader reuses its buffers
        cluster.sync();
    }
    // dynamic shared memory size: 3 * TILE_ELEMS * sizeof(float)
}

Here, the leader performs one pair of cooperative copies of each tile into its shared memory using a block-scoped pipeline. The cluster-wide barrier makes the data visible to all cluster members through DSMEM.在这里,领导者使用块作用域的流水线,对每个分块执行一对协作拷贝写入其共享内存。集群范围的屏障使这些数据通过 DSMEM 对集群所有成员可见。

This “copy once, share through DSMEM” pattern shares the leader’s tiles using DSMEM and map_shared_rank following a cluster barrier. It does not perform a TMA multicast as we covered previously with tensor-map descriptors, cp.async.bulk.tensor, and shared::cluster. This is an important distinction to understand as, without the cluster.sync() in the DSMEM pattern, followers can read stale leader SMEM data. As such, cluster.sync() is required before map_shared_rank().这种“拷贝一次、通过 DSMEM 共享”的模式在集群屏障之后使用 DSMEM 与 map_shared_rank 共享领导者的分块。它并不像我们之前用张量映射描述符、cp.async.bulk.tensor 和 shared::cluster 所介绍的那样执行 TMA 多播。这是一个需要理解的重要区别:在 DSMEM 模式中若没有 cluster.sync(),跟随者可能读到领导者 SMEM 的陈旧数据。因此,在 map_shared_rank() 之前必须调用 cluster.sync()。

Every block computes a distinct band of rows directly from the leader’s shared memory by using map_shared_rank and writes its results to global memory. This removes duplicate global loads across the cluster and keeps the overlap advantages of a pipeline at the points where it matters.每个线程块都通过 map_shared_rank 直接从 leader 块的共享内存中计算各自不同的一段行,并将结果写入全局内存。这消除了集群内重复的全局加载,并在真正重要的环节保留了流水线的重叠优势。

The leader block’s loader warp calls producer_commit() to publish its local stage to that block. Other blocks do not wait on the leader’s pipeline. Instead, they observe the leader’s tiles after the cluster wide barrier and then read from distributed shared memory.leader 块的加载线程束调用 producer_commit(),将其本地阶段发布给该块。其他块并不等待 leader 的流水线,而是在集群级屏障之后观测 leader 的分块(tile),再从分布式共享内存中读取。

This allows tile loading, computing, and storing to interleave across multiple thread blocks, which drives the SMs to even higher utilization. Table 10-5 compares the single-block warp-specialized kernel with the thread-block-clustered (multiblock) warp-specialized kernel presented in this section.这使得分块加载、计算与存储能够在多个线程块之间交错进行,从而将流多处理器(SM)的利用率推向更高水平。表 10-5 对比了单块线程束专用化内核与本节介绍的线程块集群(多块)线程束专用化内核。

Metric Naive tiling指标 朴素分块

Two-stage, double-buffered (double_ buffer_pipeline)两阶段、双缓冲(double_ buffer_pipeline)

Text
Warp-specialized
(warp_special
ized_pipeline)

Thread block cluster pipeline (warp_special线程块集群 流水线(warp_special

Text
ized_cluster_

pipeline) Kernel execution timepipeline) 内核执行时间

17.2 ms (+6.5% speedup versus warp-specialized) Warp execution efficiency17.2 ms(相比线程束专用化提升 6.5%) 线程束执行效率

41.3 ms 20.5 ms (+2.01× faster versus naive)41.3 ms 20.5 ms(相比朴素分块快 2.01 倍)

18.4 ms (+10.2% speedup versus two-stage)18.4 ms(相比两阶段提升 10.2%)

96% (+4% versus two-stage) 97% (+1% versus warp-specialized) Warp state stall % (for shared memory and barrier waits)96%(相比两阶段 +4%) 97%(相比线程束专用化 +1%) 线程束状态停顿占比(针对共享内存与屏障等待)

68% 92% (+24% versus naive)68% 92%(相比朴素分块 +24%)

High Low Minimal Minimal (further reduced)高 低 极小 极小(进一步降低)

L2 throughput 80 GB/s 155 GB/s (+94% versus naive)L2 吞吐量 80 GB/s 155 GB/s(相比朴素分块 +94%)

170 GB/s (+3% versus warp-specialized) Throughput scalability170 GB/s(相比线程束专用化提升 3%)吞吐量可扩展性

165 GB/s (+6.45% versus two-stage)165 GB/s(相比两阶段提升 6.45%)

Scales fully across thread blocks until they are constrained by SM warp limit of 64 resident warps per SM on Blackwell DRAM read throughput versus kernel duration可跨线程块完全扩展,直至受限于 Blackwell 上每 SM 64 个常驻线程束的 SM 线程束上限;Blackwell 上 DRAM 读吞吐量与内核持续时间的关系

Scales up to 2–3 warps/SM可扩展至 2–3 个线程束/SM

Scales to ~6 warps/SM Scales nearly linearly to SM warp limit (64 resident warps per SM on Blackwell)可扩展至约 6 个线程束/SM 近乎线性地扩展至 SM 线程束上限(Blackwell 上每 SM 64 个常驻线程束)

Poor overlap重叠效果差

Great overlap Excellent overlap Excellent overlap (even under thread block handoff)重叠效果良好 重叠效果极佳 重叠效果极佳(即使在线程块交接期间)

Instruction count 1.7 B 1.05 B (–38% versus naive)指令数 1.7 B 1.05 B(相比朴素实现减少 38%)

~1.00 B (–4.76% versus two-stage)约 1.00 B(相比两阶段减少 4.76%)

~0.98 B (–2% versus warp-specialized)约 0.98 B(相比线程束专用化减少 2%)

Here, we see that the thread block cluster implementation further improves on warp-specialized by distributing loader, compute, and storer roles across multiple thread blocks in one cooperative launch. This increases overall SM utilization and reduces both execution time and redundant memory traffic.在这里可以看到,线程块集群实现在线程束专用化的基础上进一步提升:它在一次协作启动中把加载器、计算和存储器的角色分配到多个线程块上。这提高了整体 SM 利用率,并同时缩短了执行时间、减少了冗余内存流量。

Specifically, the thread block cluster pipeline kernel implementation (warp_special ized_cluster_pipeline) squeezes out another ~1.2 ms (6.5%) over the single-thread-block, warp-specialized kernel. This is because the thread block cluster version interleaves tile loads, computes, and stores across all thread blocks.具体而言,线程块集群流水线内核实现(warp_special ized_cluster_pipeline)相比单线程块的线程束专用化内核又省出约 1.2 ms(6.5%)。这是因为线程块集群版本在所有线程块之间交错执行分块加载、计算和存储。

SM utilization reaches ~97%, since any idle SM in one thread block can be kept busy by another thread block performing a different pipeline stage. L2 load throughput peaks around 170 GB/s, thanks to better shared-memory reuse and fewer redundant loads. As such, the cluster can avoid duplicate global reads. With DSMEM, the leader block loads a tile once and other blocks read it using map_shared_rank() after cluster.sync() is called. With TMA multicast, a single tensor copy from global can be broadcast directly into each block’s shared memory.SM 利用率达到约 97%,因为某个线程块中空闲的 SM 可以被执行不同流水线阶段的另一个线程块保持忙碌。得益于更好的共享内存复用和更少的冗余加载,L2 加载吞吐量峰值约为 170 GB/s。因此,集群可以避免重复的全局读取。使用 DSMEM 时,主线程块只加载一次分块,其他线程块在调用 cluster.sync() 之后通过 map_shared_rank() 读取它。使用 TMA 多播时,一次来自全局内存的张量拷贝可以直接广播到每个线程块的共享内存中。

Instruction count drops to ~0.98 B because thread-block-wide pipelining further reduces wasted cycles on redundant global reads. Throughput scalability is also maintained since, as long as you have enough thread blocks and warps, the pipeline can keep them all saturated up to the warps-per-SM hardware limit of your GPU, or 64 warps per SM on Blackwell.指令数降至约 0.98 B,因为线程块级流水线进一步减少了冗余全局读取所浪费的周期。吞吐量可扩展性同样得以保持:只要有足够的线程块和线程束,流水线就能让它们全部保持饱和,直至达到 GPU 的每 SM 线程束硬件上限,即 Blackwell 上每 SM 64 个线程束。

Note that both the single thread block and thread block cluster versions of warp specialization will eventually hit that same “warps-per-SM” ceiling, but they differ in how those warps are fed data and synchronized. This difference translates into a measurable performance advantage for the thread block cluster implementation, as we saw in Table 10-5.注意,线程束专用化的单线程块版本与线程块集群版本最终都会触及相同的“每 SM 线程束数”上限,但二者在为这些线程束供给数据和同步的方式上有所不同。这一差异转化为线程块集群实现可测量的性能优势,正如我们在表 10-5 中所见。

Specifically, in the single thread block, warp-specialized kernel, each thread block allocates exactly three warps: one to load a tile-sized chunk of A and B into shared memory, one to compute on that chunk, and one to store the results back to global memory. Once those three warps complete their roles, the thread block moves on to its next tile, and the cycle repeats.具体而言,在单线程块、线程束专用化的内核中,每个线程块恰好分配三个线程束:一个负责将 tile(数据块)大小的 A 和 B 数据加载到共享内存,一个负责对该数据块进行计算,一个负责将结果写回全局内存。一旦这三个线程束完成各自的职责,线程块便转向下一个 tile,如此循环往复。

Even if every SM has three active warps for load, compute, and store, you are not saturating the warp limit of the SM. Modern GPUs allow up to 64 concurrent warps per SM, and actual residency is limited by registers, shared memory, and blocks per SM.即使每个 SM 都有三个活跃线程束分别负责加载、计算和存储,你也并未饱和 SM 的线程束上限。现代 GPU 每个 SM 最多允许 64 个并发线程束,而实际驻留数受寄存器、共享内存和每 SM 线程块数的限制。

At similar total live warp counts, a single-block warp specialized design and a cooperative grid can show comparable occupancy, but the memory behavior differs. In a single-block design, each block fetches its own copy of any input tile that it touches. When multiple blocks reuse the same tile at about the same time, the device performs duplicate reads that consume L2 capacity and HBM bandwidth.在相近的总活跃线程束数量下,单块线程束专用化设计与协作网格可以表现出相当的占用率,但内存行为不同。在单块设计中,每个块都会获取它所触及的任何输入 tile 的自己的一份副本。当多个块在大致相同的时间复用同一个 tile 时,设备会执行重复读取,消耗 L2 容量和 HBM 带宽。

A thread block cluster changes that pattern. Blocks in the cluster can share data through distributed shared memory and can receive one multicast of a global tile into the shared memory of all blocks in the cluster. Use a block-scoped pipeline in the leader block together with cluster synchronization and map_shared_rank, or configure for multicast when broadcast from global memory is preferred. This removes duplicate loads within the cluster and reduces pressure on L2 and HBM.线程块集群改变了这一模式。集群中的块可以通过分布式共享内存共享数据,并且可以将一个全局 tile 的一次多播接收到集群中所有块的共享内存中。在领导者块中使用块作用域的流水线,配合集群同步和 map_shared_rank,或者在更倾向于从全局内存广播时配置多播。这消除了集群内的重复加载,并减轻了对 L2 和 HBM 的压力。

A cuda::pipeline is scoped to the creating block. A producer_commit() in one block does not release consumer_wait() in other blocks. To coordinate across blocks in a cluster, first publish data (e.g., into the leader’s shared memory), then use clus ter.sync(). Followers then access the leader’s shared memory with map_shared_rank(). Optional TMA multicast can broadcast from global memory directly into each block’s shared memory.cuda::pipeline 的作用域限于创建它的块。一个块中的 producer_commit() 不会解除其他块中的 consumer_wait()。要在集群中的多个块之间协调,首先发布数据(例如发布到领导者的共享内存),然后使用 clus ter.sync()。随后,跟随者使用 map_shared_rank() 访问领导者的共享内存。可选的 TMA 多播可以从全局内存直接广播到每个块的共享内存。

At this point, every block observes the leader’s tiles through DSMEM (or TMA multicast) and proceeds. The block-scoped pipeline does not explicitly synchronize other blocks. cluster.sync() and DSMEM semantics provide the cluster-wide visibility.此时,每个块都通过 DSMEM(或 TMA 多播)观察到领导者的 tile 并继续执行。块作用域的流水线并不显式同步其他块。cluster.sync() 和 DSMEM 语义提供了集群范围的可见性。

When designed to use TMA multicast and DSMEM efficiently, the thread block cluster pipeline will fetch each tile exactly once and multicast it into every thread block’s shared memory. This is in contrast to redundantly reloading the data from each thread block. As such, although every SM still cannot exceed its 64-warp limit, the global thread block cluster coordination ensures that those warps spend far less time on duplicate loads and far more time on actual computation.当设计为高效使用 TMA 多播和 DSMEM 时,线程块集群流水线会恰好获取每个 tile 一次,并将其多播到每个线程块的共享内存中。这与从每个线程块冗余重新加载数据形成对比。因此,尽管每个 SM 仍然无法超过其 64 线程束上限,全局的线程块集群协调确保这些线程束在重复加载上花费的时间少得多,而在实际计算上花费的时间多得多。

Moreover, in the single thread block version, each block loads its next tile as soon as it finishes computing the current one. If one thread block completes its computation slightly earlier than another, it immediately issues its own load for the next tile—even if that tile is already being loaded by another block. This results in redundant memory traffic.此外,在单线程块版本中,每个块在完成当前 tile 的计算后立即加载下一个 tile。如果一个线程块比另一个稍早完成计算,它会立即发出自己下一个 tile 的加载——即使该 tile 已由另一个块正在加载。这导致冗余的内存流量。

However, with the thread block cluster pipeline version, if a compute warp on SM 0 finishes early, it does not fetch its next tile immediately. Instead, it stalls at con sumer_wait() until the global loader—possibly running on SM 7—finishes bringing that tile into shared memory for all thread blocks. In other words, the compute warp waits for the cluster’s single, shared load rather than issuing its own redundant copy.然而,在使用线程块集群流水线版本时,如果 SM 0 上的某个计算线程束提前完成,它不会立即获取下一个分块(tile)。相反,它会阻塞在 con sumer_wait() 处,直到全局加载器(可能运行在 SM 7 上)完成为所有线程块把该分块送入共享内存。换言之,计算线程束等待的是集群中唯一的一次共享加载,而不是自己发起一次冗余的拷贝。

By pausing until the tile is guaranteed available, each SM participating in a thread block cluster avoids spinning idle or performing duplicate loads. This alignment across thread block clusters smooths out variations in load times and keeps every SM’s compute warps busy with data loaded only once per cluster. This improves overall throughput.通过暂停等待,直到该分块确保可用,参与线程块集群的每个 SM 都能避免空转等待或执行重复加载。跨线程块集群的这种对齐可以平滑加载时间的波动,并使每个 SM 的计算线程束始终忙于处理每个集群仅加载一次的数据。这提升了整体吞吐量。

Consider a problem so large and a grid so full of work that every SM already has loader, compute, and storer warps active under the single-block approach. In that case the raw warp counts and per-SM occupancy can look similar between the two kernels. The cluster version helps when multiple blocks reuse the same input tiles by eliminating duplicate global reads within a cluster and by using cluster-wide synchronization to align stages.考虑这样一种情况:问题规模如此之大、网格中的工作如此之满,以至于在单块方案下每个 SM 都已经有加载、计算和存储线程束处于活动状态。在这种情况下,两个内核的原始线程束数量和每 SM 占用率看起来可能相似。当多个块复用相同的输入分块时,集群版本通过消除集群内的重复全局读取,并利用集群范围的同步来对齐各阶段,从而发挥作用。

This often produces a modest throughput improvement on reuse-heavy workloads, but the gain depends on tile reuse, alignment, and resource limits. Both designs remain bounded by the same architectural limits, such as resident warps, registers, and shared memory.这在复用密集的工作负载上通常带来适度的吞吐量提升,但收益取决于分块复用、对齐情况和资源限制。两种设计仍然受限于相同的架构限制,例如常驻线程束、寄存器和共享内存。

All of the kernels here use the CUDA pipeline for producer and consumer handoffs without calling a block-wide __syncthreads. The naive tiled kernel does not overlap memory and compute. The two-stage double-buffered pipeline overlaps copies with compute and can significantly reduce time to solution on memory-bound cases.这里的全部内核都使用 CUDA 流水线来完成生产者与消费者之间的交接,而不调用块范围的 __syncthreads。朴素的分块内核不会重叠显存访问与计算。两阶段双缓冲流水线将拷贝与计算重叠,在受内存限制的场景下可以显著缩短求解时间。

The warp-specialized pipeline dedicates separate warps to load, compute, and store to avoid implicit full-block waits. The thread block cluster variant shares tiles across blocks in a cluster through distributed shared memory or multicast so that a tile is fetched once per cluster rather than once per block. These patterns improve utilization and reduce redundant memory traffic.线程束专用化流水线将加载、计算和存储分配给不同的线程束,以避免隐式的全块等待。线程块集群变体通过分布式共享内存或多播在集群内的块之间共享分块,使得每个分块在每个集群中只获取一次,而不是每个块获取一次。这些模式提升了利用率并减少了冗余的内存流量。

关键要点

Key Takeaways

The following are key takeaways from this chapter, which is focused on extracting peak performance on modern GPUs. These will keep kernels from idling on DRAM, warp schedulers, and global barriers:以下是本章的关键要点,本章聚焦于在现代 GPU 上榨取峰值性能。这些要点可避免内核在 DRAM、线程束调度器和全局屏障上空闲等待:

Hide latency with pipeline depth用流水线深度隐藏时延
Use two-stage (cuda::pipeline_shared_state<cuda::thread_scope_block, 2>) tiling to overlap asynchronous loads with compute and add another stage (cuda::pipeline_shared_state<cuda::thread_scope_block, 3>) when compute outweighs memory (e.g., modern GPUs like Blackwell.) This will help to eliminate idle warps.使用两阶段(cuda::pipeline_shared_state<cuda::thread_scope_block, 2>)分块来将异步加载与计算重叠,并在计算量超过显存访问量时(例如 Blackwell 等现代 GPU)再增加一个阶段(cuda::pipeline_shared_state<cuda::thread_scope_block, 3>)。这有助于消除空闲线程束。
Balance workloads with warp specialization用线程束专用化平衡工作负载
Assign separate warps to loading, computing, and storing when compute phases dominate, ensuring near-peak warp efficiency on modern GPU hardware.当计算阶段占主导时,将加载、计算和存储分配给不同的线程束,确保在现代 GPU 硬件上达到接近峰值的线程束效率。
Remove launch overhead with persistent kernels用持久化内核消除启动开销
Run a single long-lived kernel over a device-side work queue and use grid.sync() for multiphase algorithms. This will reduce host-device round trips and overall launch costs.在设备端工作队列上运行单个长期存活的内核,并对多阶段算法使用 grid.sync()。这将减少主机与设备之间的往返以及整体启动开销。
Enable on-chip sharing with thread block clusters and DSMEM使用线程块集群和 DSMEM 实现片上共享
Group thread blocks (CTAs) into clusters so they share a contiguous on-chip buffer. For a cluster-wide broadcast, use cp.async.bulk.tensor with a TMA multicast descriptor. Use TMA multicast to broadcast tiles once to every thread block. This will boost L2 hit rates and trim DRAM bandwidth.将线程块(CTA)组织成集群,使它们共享一块连续的片上缓冲区。要在整个集群范围内广播,可使用带 TMA 多播描述符的 cp.async.bulk.tensor。使用 TMA 多播将数据块一次性广播到每个线程块。这会提升 L2 命中率并减少 DRAM 带宽占用。
Pay special attention to barrier semantics特别注意屏障语义
Both cluster.sync() and grid.sync() require every participating thread block to reach the same synchronization point. Mismatched control flow—or an excessive cluster size—may lead to deadlock or launch failure.cluster.sync() 和 grid.sync() 都要求每个参与的线程块到达同一个同步点。控制流不一致——或集群规模过大——可能导致死锁或启动失败。
Profile before tuning先剖析,再调优
Use Nsight Compute to identify whether your kernel is memory bound or compute bound. If memory bound, start with a two-stage pipeline. If compute bound, consider warp specialization or thread block clustering.使用 Nsight Compute 判断你的内核是受限于内存还是受限于计算。若受限于内存,先从两阶段流水线入手;若受限于计算,可考虑线程束专用化或线程块集群。
Verify compiler-generated pipelines before hand-tuning手工调优前先验证编译器生成的流水线
Profile and inspect the framework’s generated code from compilers like PyTorch’s compiler and NVCC. If the compiled code uses an asynchronous pipeline using cuda::memcpy_async, producer_commit(), and consumer_wait(), manual tuning likely won’t produce much of a speedup.剖析并检查框架由 PyTorch 编译器、NVCC 等编译器生成的代码。如果编译后的代码已经使用 cuda::memcpy_async、producer_commit() 和 consumer_wait() 构成异步流水线,那么手工调优大概率不会带来多少加速。

结语

Conclusion

The techniques discussed in this chapter help to systematically hide latency and remove redundant loads. This keeps the GPU at near-peak utilization for the duration of the kernel’s execution.本章讨论的技术有助于系统性地隐藏时延并消除冗余加载。这能让 GPU 在内核执行的整个过程中保持接近峰值的利用率。

Warp-specialized pipelines overlap load, compute, and store operations. Cooperative-group barriers (grid.sync() and cluster.sync()) help coordinate multiphase work without host round trips. And persistent kernels loop over a device-side queue to eliminate launch overhead.线程束专用化流水线将加载、计算和存储操作重叠起来。协作组屏障(grid.sync() 和 cluster.sync())有助于协调多阶段工作,而无需往返主机。持久化内核则在设备端队列上循环,以消除启动开销。

As always, you should start by profiling. If global-memory stalls dominate, a two-stage asynchronous-copy pipeline like double buffer usually suffices. If compute warps still stall, switch to a multistage warp-specialized pipeline such that loader, compute, and storer warps can operate with minimal contention. The only contention would be at the memory subsystem.一如既往,你应当从剖析开始。如果全局内存停顿占主导,像双缓冲这样的两阶段异步拷贝流水线通常就足够了。如果计算线程束仍然停顿,就改用多级线程束专用化流水线,使加载、计算和存储线程束能在竞争最小的情况下运行。唯一的竞争会出现在内存子系统。

For multiphase reductions or irregular tasks, replace multiple launches with a single persistent kernel plus grid.sync() to preserve occupancy. And when thread blocks need the same data (e.g., multiheaded attention), you can form a thread block cluster so that DSMEM and TMA load each tile only once—and multicast the data to the other thread blocks—without repeatedly accessing global memory. These techniques will move performance closer to the GPU’s peak theoretical limits.对于多阶段归约或不规则任务,可用单个持久化内核加上 grid.sync() 替代多次启动,以保持占用率。而当多个线程块需要同一份数据时(例如多头注意力),可以组成线程块集群,让 DSMEM 和 TMA 只加载每个数据块一次——并将数据多播给其他线程块——从而避免反复访问全局内存。这些技术将使性能更接近 GPU 的峰值理论极限。

As you continue through the next chapters, it’s important to remember these principles because they apply to many more optimizations. Specifically, you should overlap work at the warp and thread block levels, synchronize directly on-device (versus on the host), and share data on chip. These are essential mechanisms for tuning ultra-high-performance GPU workloads.在继续阅读后续章节时,务必牢记这些原则,因为它们适用于更多优化场景。具体而言,你应当在 warp 和线程块层级重叠工作、直接在设备上(而非主机上)进行同步,并在芯片内共享数据。这些是调优超高性能 GPU 工作负载的关键机制。

In the next chapter, we keep these intra-kernel building blocks—cuda::pipeline double-buffering, warp-specialized roles, and thread-block clusters—and show how to drive them through CUDA streams. The goal is to hide latency between kernels and between host ↔ device communication and not just inside a single kernel. Concretely, we will reuse the kernels from this chapter and run them in multistream pipelines with cudaMemcpyAsync, cudaMallocAsync/cudaFreeAsync, and event-based synchronization. This will help push the entire system toward achieving peak performance across many GPUs in your AI system.在下一章中,我们将保留这些内核内构建模块——cuda::pipeline 双缓冲、线程束专用化角色以及线程块集群——并展示如何通过 CUDA 流来驱动它们。目标是在内核之间以及主机 ↔ 设备通信之间隐藏时延,而不仅仅是在单个内核内部。具体来说,我们将复用本章的内核,并结合 cudaMemcpyAsync、cudaMallocAsync/cudaFreeAsync 以及基于事件的同步,将它们放入多流流水线中运行。这将有助于推动整个系统在 AI 系统中的多块 GPU 上达到峰值性能。

《AI Systems Performance Engineering》中文译本 · 原书作者 Chris Fregly · O'Reilly Media 出版 · 本译本仅供个人学习研究