跳转至

第10章 内核内流水线、线程束专用化和协作线程块集群 (Intra-Kernel Pipelining, Warp Specialization, and Cooperative Thread Block Clusters)

在前几章中,我们介绍了基础优化,如调整内存访问、最大化并行性、重叠计算和数据传输、提高占用率以及最小化线程束停顿。这些优化帮助隐藏了延迟并消除了瓶颈。

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.

现代GPU提供了硬件和软件机制,使线程束和线程块能够以细粒度进行协作。这些机制包括持久内核、线程束专用化、协作组和线程块集群。这些技术是高性能库和框架的基础,如CUTLASS和FlashAttention,它们在许多情况下实现了接近光速的性能。

Modern GPUs, however, offer advanced hardware features and execution models that let us take the fundamental optimization techniques even further.

在本章中,您将学习如何实现这些高级并行模式。我们将从线程束专用化开始,它将线程束分配给生产者-消费者流水线中的特定角色。然后我们将介绍持久内核,它保持GPU运行一个长期存在的内核,从设备端队列中拉取工作。接下来,我们将讨论协作组,它提供了灵活的同步原语。最后,我们将介绍线程块集群和分布式共享内存(DSMEM),它们使多个线程块能够在片上高效共享数据。

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.

在本章结束时,您将能够实现复杂的内核内流水线,减少全局内存流量,并充分利用现代GPU的协作能力。

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.

通过理解这些现代GPU执行模型,您将准备好进入下一章,在那里我们将通过使用CUDA流探索内核间流水线来进一步扩展这些优化。

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

内核内流水线技术 (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.)

核心思想是将内核结构化为并发阶段,以便在加载或存储一块数据时,先前加载的数据正在被处理。这些阶段在不同的块或数据块上并行运行。这提高了吞吐量并有效地隐藏了延迟。

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.

传统上,GPU依赖线程束级多线程来隐藏延迟。当一个线程束在内存加载上停顿时,其他线程束继续进行计算。这是执行模型中单指令多线程(SIMT)延迟隐藏的基础。

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.

内核内流水线通过在同一线程束或内核内重叠内存和计算进一步推进了这一点。它使用细粒度协调来交错内存加载和计算--有时在单个线程束内。

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.

使用CUDA Pipeline API的内核内流水线可以在没有任何__syncthreads()的情况下重叠异步内存传输和计算。内核内流水线的两种常见方法是双缓冲和线程束专用化。

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.

在双缓冲(两阶段)流水线方法中,所有线程统一协作。在线程束专用流水线方法中,线程束被专用化为不同的角色,如内存加载器、计算器和内存存储器。选择取决于您的工作负载和性能要求。表10-1总结了这两种<cuda/pipeline>变体。

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 variants.

表10-1. 使用CUDA Pipeline API在现代GPU上进行内核内流水线的两种方法

API变体 最适用于 主要用途
双缓冲流水线 基于循环的分块和双缓冲 在同一线程束或块中重叠加载和计算
线程束专用流水线(例如,三级内存加载器、计算器、内存存储器) 具有多个不同线程束角色(在我们的例子中为3个)的持久内核 将线程束分配给单独的角色/阶段,如内存加载、计算和内存存储

Table 10-1. Two approaches for intra-kernel pipelining on modern GPUs using the CUDA Pipeline API

使用CUDA Pipeline API进行协作分块和双缓冲 (Cooperative Tiling and Double-Buffering with the CUDA Pipeline API)

您可以使用C++ Pipeline API实现传统的双缓冲分块模式,方法是实例化两阶段流水线来重叠内存加载和计算。具体来说,您可以声明一个cuda::pipeline_shared_state<cuda::thread_scope_block, 2>对象,它使用协作组(稍后讨论)作用于特定线程块。这本质上是一个生产者-消费者模式,如图10-1所示。

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 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.

图10-1 使用CUDA Pipeline API的两阶段生产者-消费者模式

Figure 10-1. Two-stage producer-consumer pattern with the CUDA Pipeline API

关键的CUDA Pipeline API调用如下所示。它们后面是一个使用此API的双缓冲、协作分块内核的实现,以演示与现代GPU硬件功能一致的CUDA技术:

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:

pipe.producer_acquire() - 为写入保留下一个流水线阶段

pipe.producer_acquire() - Reserves the next pipeline stage for writing

pipe.producer_commit() - 发信号表示此阶段先前发出的异步操作已准备好供消费

pipe.producer_commit() - Signals that the previously issued asynchronous operations for this stage are ready for consumption

pipe.consumer_wait() - 等待直到此阶段先前提交的操作完成,以避免循环中的竞争条件

pipe.consumer_wait() - Waits until the previously committed operations for this stage complete to avoid race conditions in loops

pipe.consumer_release() - 释放当前阶段以便可以重用

pipe.consumer_release() - Releases the current stage so it can be reused

流水线中的两个阶段将全局内存加载与计算重叠。在第一阶段(Stage 0),线程块中的一个线程束为下一个块发出异步预取到共享内存。预取发出协作的cuda::memcpy_async副本,这些副本降低为每线程cp.async到共享内存。

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::memcpy_async copies that lower to per thread cp.async into shared memory.

当阶段0在一个线程束中生产(加载)数据时,线程块中的其余线程束在第二阶段(Stage 1)中消费(计算)加载的数据。这种简单的生产者-消费者实现通过正在进行的计算隐藏DRAM延迟,如图10-2所示。

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.

图10-2 使用生产者-消费者流水线隐藏全局DRAM加载(L)延迟与(C)计算

Figure 10-2. Hiding global DRAM load (L) latency with (C) compute using a producer-consumer pipeline

这种模式可以提高SM利用率并减少内核时间,但您是计算受限还是内存受限取决于操作、块大小和重叠效率。即使是高度优化的注意力内核如FlashAttention-3,由于重叠和数据移动的实际限制,也只能达到峰值FP16 FLOPS的约75%。

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.

这是一个使用CUDA C++ Pipeline API的两阶段双缓冲示例。此API启用了这里使用的细粒度生产者-消费者同步代码:

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:

#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) {
    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;
        }
        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;
    }
}

代码首先使用协作组(CG)获取当前线程块的句柄,稍后讨论。然后它立即实例化绑定到该块的两阶段cuda::pipeline对象。通过在任何异步操作之前创建流水线,流水线的内部屏障和内部同步机制在执行数据移动之前就已就位。

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.

内核然后通过定义extern __shared__ float shared_mem[]为A和B块分配一个连续的共享内存区域。它使用指针算术(float* A_buf[2]float* B_buf[2])将此缓冲区分为四个子区域,两个用于A,两个用于B。这允许真正的双缓冲而无需额外的动态分配。

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.

在进入主循环之前,内核使用绑定到以下流水线的异步副本异步预取前STAGES个块:producer_acquire()memcpy_async()producer_commit()。消费者线程束使用consumer_wait()consumer_release(),以便计算恰好在预取块准备好时开始。

Before entering the main loop, the kernel asynchronously prefetches the first STAGES tiles using asynchronous copies bound to the following pipeline: producer_acquire() → memcpy_async() → producer_commit(). Consumer warps use consumer_wait() and consumer_release() so the compute starts exactly when the prefetched tile is ready.

这个初始屏障替代了将来对__syncthreads()的需求,并确保流水线的阶段0和阶段1缓冲区在生产者-消费者序列的第一次迭代中正确填充。

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.

在每次迭代中,内核通过调用pipe.producer_acquire()保留下一个缓冲区以加载后续块。然后它启动两个cuda::memcpy_async操作--一个用于A块,一个用于B块。每个加载都绑定到流水线对象cuda::memcpy_async(..., pipe),它发出从全局内存到共享内存的异步副本。

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.

当访问被正确合并和分块时,这些异步副本可以与计算很好地重叠。这样,流水线保持数据供给以执行最大的有用工作。在排队memcpy_async调用后立即,内核使用pipe.producer_commit()发出完成信号。此提交记录副本的到达,并允许消费者线程束在该特定阶段上等待,而不会阻塞整个线程块。

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.

同时,块中的其他线程束调用pipe.consumer_wait()。这只会有效地停顿那些依赖于当前缓冲区中数据的线程,直到生产者提交它。一旦等待完成,每个线程调用设备函数computeTile(...)执行其TILE_SIZE x TILE_SIZE点积计算。结果在寄存器中增量累积。

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_release()释放该阶段以在后续迭代中重用。这种细粒度释放防止了块范围的停顿,并最大化了计算和内存传输阶段之间的重叠。

After finishing the computation on the current buffer, the warps invoke pipe.consumer_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.

在每次迭代结束时,交换curr和next,这循环两个双缓冲阶段。这样,刚刚计算完的内存缓冲区现在成为下一个异步内存的目标。这些生产者和消费者阶段对所有K维度块重复。

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. These producer and consumer stages are repeated for all K-dimension tiles.

在处理完最后一个块后,每个线程的累加器包含完整的TILE_SIZE x TILE_SIZE点积结果。然后将此结果写回C_global的相应元素。

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.

使用双缓冲时,建议确保每个异步副本(memcpy_async)后面跟着适当的producer_commit()consumer_wait()调用以进行同步。这保证了计算内核只会在数据正确加载后使用它。现代CUDA编译器通常会自动为简单循环执行这些优化。

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.

建议使用Nsight Compute的异步副本指标验证流水线是否按预期执行。特别注意线程束占用率和共享内存bank冲突。目标是提高SM Active百分比并减少停顿周期。完全隐藏DRAM延迟取决于占用率、访问模式和共享内存bank行为。

It's recommended to validate that the pipeline is executing as expected using Nsight Compute's asynchronous copy metrics. In particular, pay attention to warp occupancy and shared-memory bank conflicts. The goal is to increase SM Active percent and reduce stall cycles. Full hiding of DRAM latency depends on occupancy, access patterns, and shared-memory bank behavior.

这种简单的双缓冲方案比朴素分块快约2倍,如表10-2中朴素分块实现与优化的双缓冲流水线实现的性能比较所述。

This simple double-buffering scheme achieves about a 2x speedup over naive tiling, as described in Table 10-2, which compares the performance of the naive tiling implementation versus the optimized double-buffered pipeline implementation.

表10-2. 使用CUDA Pipeline API的朴素分块与双缓冲内核性能比较

指标 朴素分块 两阶段双缓冲流水线
内核执行时间 41.3 ms 20.5 ms(比朴素快约2倍)
SM Active % 68% 92%(比朴素+24%)

Table 10-2. Performance comparison between the naive tiling and double-buffered kernel using the CUDA Pipeline API. Note: The numeric values in all metrics tables are illustrative to explain the concepts. For actual benchmark results on different GPU architectures, see the GitHub repository.

在这个实验中,使用CUDA C++ Pipeline API的gemm_tiled_pipeline内核比朴素分块版本实现了2倍加速。通过使用细粒度的pipe.producer_commit()pipe.consumer_wait()原语,流水线保持填充,SM Active %从68%跃升24%到92%。

In this experiment, the gemm_tiled_pipeline kernel using the CUDA C++ Pipeline API achieves a 2x 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%.

线程束专用化和生产者-消费者模式 (Warp Specialization and the Producer-Consumer Model)

线程束专用化是一种编程模式,其中线程块内的不同线程束被分配不同的角色。例如,一个线程束可能专门负责从全局内存加载数据,而另一个线程束专门负责对这些数据进行计算,第三个线程束专门负责将结果写回全局内存。

Warp specialization is a programming pattern where different warps within a thread block are assigned different roles. For example, one warp may specialize in loading data from global memory, another warp may specialize in computing on that data, and a third warp may specialize in writing results back to global memory.

通过将加载、计算和存储角色分离到不同的线程束,每个线程束可以专注于其特定任务。这实现了这些阶段之间的重叠。当一个线程束在等待内存操作时,其他线程束可以继续进行有用的工作。

By separating load, compute, and store roles into different warps, each warp can focus on its specific task. This enables overlap between these phases. When one warp is waiting for memory operations, other warps can continue doing useful work.

这种专用化通过确保每个线程束始终有工作可做来提高整体SM利用率。它还减少了同步开销,因为线程束可以以细粒度方式协调而无需全块屏障。

This specialization improves overall SM utilization by ensuring each warp always has work to do. It also reduces synchronization overhead because warps can coordinate in a fine-grained manner without full-block barriers.

线程束专用化通过将操作分配给使用不同硬件的线程束(如数据移动(例如TMA)和计算(例如Tensor Core))来扩展双缓冲。这与将同一线程束重用于加载数据和计算形成对比,如图10-3所示。

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.

图10-3 非线程束专用化内核,每个线程束执行数据加载和计算的混合操作

Figure 10-3. Nonwarp specialized kernel with each warp performing a mix of both data loading and compute (source: https://oreil.ly/WZDbM)

这种类型的专用化允许每组线程束拥有自己的指令序列。因此,指令可以连续发出和执行,而不会被其他类型的操作中断。具体来说,线程束专用化允许您分配一组"生产者"或"内存"线程束使用cuda::memcpy_async异步预取块。然后所有其他"消费者"或"计算"线程束执行计算,如图10-4所示。

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 computation, as shown in Figure 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.

因为每个线程束都有自己的调度器,GPU可以在同一周期内从不同线程束使用不同的线程束调度器发出加载指令、数学指令和写入指令。因此,具有多个调度器的SM可以在一个周期内从线程束0发出内存指令,从线程束1发出数学指令,依此类推。

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.

图10-4 线程束专用化内核,一组线程束用于加载数据,所有其他线程束用于计算

Figure 10-4. 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.

线程束专用化的一个有趣模式是使用三种不同类型的线程束,如"加载器"、"计算器"和"存储器"线程束。加载器线程束将块推入流水线队列。计算器线程束在每个块上运行计算内核。存储器线程束写出结果,如图10-5所示。

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.

这种线程束专用流水线挤出单线程束、双缓冲、顺序加载和计算循环无法解决的空闲周期。线程束专用化高效重叠数据传输和计算,提高了GPU利用率--特别是对于长时间运行的循环和持久内核。在这些情况下,角色协调和数据传递的开销在许多迭代中分摊。

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内核具有明显的内存和计算阶段,使得内存和计算轮流成为瓶颈。通过应用线程束专用化,他们的工作负载转变为一种状态,其中SM的内存子系统和计算单元几乎整个时间都同时忙碌。

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.

图10-5 三级线程束专用流水线配置,一组线程束用于加载数据,另一组用于计算,另一组用于数据存储

Figure 10-5. 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)

分析显示,在使用线程束专用化之前,他们的L2带宽利用率和Tensor Core利用率是不同相的。在线程束专用化之后,L2带宽和Tensor Core利用率变得同相。这导致了更高的有效吞吐量,并表明线程束专用化可以挤出最后一点空闲时间--即使对于可能被忽略的精心调优的异步流水线。

Analysis 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.

另一种线程束专用化模式是三级线程束专用流水线的修改版。它像以前一样将一组线程束分配给内存加载器,但随后使用两组消费者线程束在计算和内存存储器的角色之间"乒乓"切换。这种三级线程束专用架构在CUTLASS中作为gemm_tma_warpspecialized_pingpong公开,如图10-6所示。

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 CUTLASS as gemm_tma_warpspecialized_pingpong and shown in Figure 10-6.

图10-6 具有三级线程束专用内核的乒乓架构

Figure 10-6. Ping-pong architecture with three-role warp-specialized kernel (source: https://oreil.ly/xs7YN)

在这里,消费者MMA重叠并包括少量的MMA后收尾或结尾处理。在启动下一个MMA之前需要这种每个MMA的结尾清理。具体来说,结尾可以包括累加、缩放、写回全局内存或将结果洗牌到另一个线程束。此外,结尾可以执行内务处理,如推进块指针、更新循环计数器和发信号表示此块已完成,以便下一个TMA加载或MMA可以启动。

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消费者可以在Tensor Core上开始执行有用工作之前,用一些TMA请求填充流水线,将数据移动到寄存器中。

While not shown here, there is also an equivalent prologue processing step that happens before the MMA operation. In this case, the prologue phase would fill the pipeline with a few TMA requests to move data into registers before the MMA consumers can start doing useful work on the Tensor Cores.

在实践中,线程束专用化在挤出最后一点性能方面非常有效。事实上,FlashAttention v3将其加速部分归因于线程束专用流水线,该流水线重叠GEMM和softmax计算--以及数据传输--以保持所有硬件单元忙碌。这有助于通过计算和TMA驱动的数据移动的激进重叠,为注意力计算实现接近峰值FLOPS。

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.

此外,PyTorch编译器(在第13章和第14章中介绍)生成使用线程束专用化的内核,以调度单独的线程束来加载和计算数据。它还使用低成本屏障来同步线程束,类似于下一节中详述的CUDA Pipeline API实现。PyTorch编译器系统还与CUTLASS的乒乓GEMM集成。torch.compile和Triton都可能为支持的操作生成线程束专用内核。但是,它们根据启发式方法有选择地应用线程束专用化,并不会为每个操作符启用线程束专用化。

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.

对于不平衡或延迟隐藏场景使用线程束专用化--特别是当单个线程束的计算不足以隐藏内存加载延迟时。但是,如果内核很小--或极度受内存限制--坚持使用更简单的双缓冲方案可能会产生类似的好处,而无需额外的代码复杂性。

Use warp specialization for imbalanced or latency-hiding scenarios-especially when a single warp's compute is not enough to hide memory-load latency. However, if a kernel is small-or extremely memory bound-sticking to a simpler double-buffering scheme may produce similar benefits without the extra code complexity.

生产者-消费者模式 (Producer-Consumer Pattern)

线程束专用化的核心是生产者-消费者模式。在这个模式中,一个或多个线程束(生产者)生成数据并将其放入共享缓冲区,而其他线程束(消费者)从该缓冲区读取数据并进行处理。

At the core of warp specialization is the producer-consumer pattern. In this pattern, one or more warps (producers) generate data and place it into a shared buffer, while other warps (consumers) read data from that buffer and process it.

例如,在矩阵乘法内核中,"加载器"线程束可能负责从全局内存获取A和B矩阵的块到共享内存。"计算器"线程束等待这些块可用,然后执行矩阵乘法。"存储器"线程束获取结果并将其写回全局内存。

For example, in a matrix multiplication kernel, a "loader" warp might be responsible for fetching tiles of the A and B matrices from global memory into shared memory. A "compute" warp waits for these tiles to be available, then performs the matrix multiplication. A "storer" warp takes the results and writes them back to global memory.

这三个阶段--加载、计算、存储--可以流水线化,以便它们并发执行。当计算器线程束正在处理块N时,加载器线程束可以获取块N+1,存储器线程束可以写出块N-1的结果。这种重叠有效地隐藏了内存延迟并保持所有线程束忙碌。

These three phases-load, compute, store-can be pipelined so that they execute concurrently. While the compute warp is processing tile N, the loader warp can fetch tile N+1, and the storer warp can write out the results of tile N-1. This overlap effectively hides memory latency and keeps all warps busy.

使用CUDA Pipeline API实现线程束专用化 (Implementing Warp Specialization with the CUDA Pipeline API)

线程束专用化建立在CUDA Pipeline API之上,允许专用线程束使用细粒度的生产者和消费者原语进行通信。这些调用避免了全块屏障,同时与异步副本(如cuda::memcpy_async)自然组合。

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的生产者和消费者调用(例如pipe.producer_acquire()pipe.producer_commit()pipe.consumer_wait()pipe.consumer_release())的关键优势在于,它们只同步实际需要传递数据的特定线程束或阶段。这与强制块中的每个线程等待形成对比。

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.

块范围屏障会停顿每个线程束--即使是那些不参与生产者-消费者流水线的线程束。该块中的所有执行必须暂停,直到每个线程到达屏障,如图10-7所示。

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 块范围屏障阻止线程继续执行,直到它们同步并加载新数据

Figure 10-7. Block-wide barrier prevents threads from proceeding until they synchronize and load the new data

相比之下,Pipeline API在内部维护每阶段状态。当生产者线程束完成其异步副本并调用pipe.producer_commit时,只有调用pipe.consumer_wait的线程束会阻塞直到数据准备好。块中其他不依赖于该阶段的线程束可以继续运行任何工作。

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.

在实践中,CUDA Pipeline API减少了空闲时间并减少了停顿的线程束,因为它消除了使用屏障暂停整个块的需要。使用流水线,您可以以比手工编码的异步副本序列(例如PTX cp.async + __syncthreads())更细的粒度协调生产者和消费者传递。

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()).

您可以使用CUDA Pipeline API以三级模式实现线程束专用化。加载器线程束为计算器线程束生成输入,计算器线程束消费这些输入并生成结果,存储器线程束消费这些结果并将其写出。流水线对象是块范围的,它在内部跟踪阶段顺序。

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 internally tracks stage ordering.

CUDA Pipeline API(在<cuda/pipeline>中)为实现线程束专用化提供了一组原语。该API允许您创建流水线阶段,异步执行内存传输,并在线程束之间进行协调。

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的线程束专用化矩阵乘法的具体示例。在这个示例中,我们将使用三个线程束:一个用于加载,一个用于计算,一个用于存储。

Let's look at a concrete example of warp-specialized matrix multiplication using the CUDA Pipeline API. In this example, we'll use three warps: one for loading, one for computing, and one for storing.

#include <cuda/pipeline>
#include <cuda_runtime.h>

#define TILE_SIZE 128

__global__ void warp_specialized_pipeline(
    const float* __restrict__ A_global,
    const float* __restrict__ B_global,
    float* __restrict__ C_global,
    int numTiles) {

    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, 2> pipe_state;
    auto pipe = cuda::make_pipeline(
        cuda::this_thread_block(), &pipe_state);

    const int lane_id = threadIdx.x & 31;
    const int warp_id = threadIdx.x >> 5;

    for (int tile = blockIdx.x; tile < numTiles; tile += gridDim.x) {
        const size_t offset = static_cast<size_t>(tile) * 
                              TILE_SIZE * TILE_SIZE;

        // Loader warp (warp 0): asynchronously load A and B tiles
        if (warp_id == 0) {
            pipe.producer_acquire();
            cuda::memcpy_async(A_tile, A_global + offset,
                cuda::aligned_size_t<32>{TILE_SIZE * TILE_SIZE * sizeof(float)},
                pipe);
            cuda::memcpy_async(B_tile, B_global + offset,
                cuda::aligned_size_t<32>{TILE_SIZE * TILE_SIZE * sizeof(float)},
                pipe);
            pipe.producer_commit();
        }

        // Compute warp (warp 1): wait for tiles and compute
        if (warp_id == 1) {
            pipe.consumer_wait();
            // Perform matrix multiplication on tiles
            for (int row = lane_id; row < TILE_SIZE; row += 32) {
                for (int col = 0; col < TILE_SIZE; ++col) {
                    float acc = 0.0f;
                    for (int k = 0; k < TILE_SIZE; ++k) {
                        acc += A_tile[row * TILE_SIZE + k] * 
                               B_tile[k * TILE_SIZE + col];
                    }
                    C_tile[row * TILE_SIZE + col] = acc;
                }
            }
            pipe.consumer_release();
            pipe.producer_acquire();
            pipe.producer_commit();
        }

        // Storer warp (warp 2): write C tile back to global memory
        if (warp_id == 2) {
            pipe.consumer_wait();
            for (int idx = lane_id; idx < TILE_SIZE * TILE_SIZE; 
                 idx += 32) {
                C_global[offset + idx] = C_tile[idx];
            }
            pipe.consumer_release();
        }
    }
}

在这个内核中,每个线程束在独特的角色中工作。加载器线程束调用producer_acquire(),使用cuda::memcpy_async执行两次协作复制到A_tile和B_tile,然后调用producer_commit()。计算器线程束调用pipe.consumer_wait()观察新提交的数据,并立即调用consumer_release()释放该阶段以供重用。计算器线程束然后通过调用producer_acquire()、计算C_tile并调用producer_commit()成为下一次传递的生产者。存储器线程束调用consumer_wait()观察计算出的C_tile,并在线程束条带化循环中将其写入全局内存,然后调用consumer_release()。这个序列使用单个块范围流水线,没有显式的阶段编号,并避免任何块范围的__syncthreads

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. 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.

这个内核作为持久内核跨多个块运行以分摊启动开销。稍后我们将详细介绍持久内核。

This kernel runs as a persistent kernel across many tiles to amortize launch overhead. More on persistent kernels in a bit.

简而言之,将CUDA Pipeline API与协作组一起使用,可以实现细粒度、SM范围的生产者-消费者传递,而无需任何显式的__syncthreads()调用。表10-3比较了三种实现:朴素分块内核、使用double_buffered_pipeline的两阶段双缓冲GEMM,以及我们的线程束专用流水线内核warp_specialized_pipeline

In short, using the CUDA Pipeline API together with cooperative groups allows fine-grained, SM-wide producer-consumer handoffs without any explicit __syncthreads() 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.

表10-3. 三种实现的比较:朴素分块内核、两阶段双缓冲GEMM和线程束专用流水线内核

指标 朴素分块 两阶段双缓冲流水线 线程束专用流水线
内核执行时间 41.3 ms 20.5 ms(比朴素快2.01倍) 18.4 ms(比两阶段快10.2%)
线程束执行效率 68% 92%(比朴素+24%) 96%(比两阶段+4%)
共享内存停顿延迟/线程束同步停顿 最小
L2吞吐量 80 GB/s 155 GB/s(比朴素+94%) 165 GB/s(比两阶段+6.45%)
吞吐量可扩展性 每个SM最多只能扩展到2-3个线程束 每个SM可良好扩展到约6个线程束 几乎线性扩展到SM的线程束(例如64个线程束)
DRAM读取与SM周期 重叠差 重叠好 重叠优秀
指令计数 1.7 B 1.05 B(比朴素-38%) ~1.00 B(比两阶段-5%)

Table 10-3. Comparison of three implementations: a naive tiled kernel, a two-stage double-buffered GEMM, and a warp-specialized pipeline kernel

在这里,双缓冲流水线在20.5 ms内完成GEMM,而线程束专用版本仅需18.4 ms完成。这10.2%的改进是因为线程束专用内核只停顿消费者线程束(例如计算)。其他线程束(例如加载器和存储器)独立进行。在之前的两阶段双缓冲内核中,每个线程都参与消费者阶段。因此,consumer_wait()有效地停顿整个线程块。

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.

这种更细粒度的每线程束同步消除了隐式的全块等待,允许所有三个线程束(加载器、计算器和存储器)连续重叠。因此,平均SM利用率从双缓冲设计的大约92%上升到线程束专用版本的大约96%--线程束停顿周期降至接近零。

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.

从可扩展性角度来看,Nsight Compute显示朴素分块内核在每个SM仅有2-3个活跃线程束后就饱和。这是因为每个块加载必须在任何计算开始之前完成。

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.

两阶段双缓冲内核通过重叠加载和计算改进了这一点。此实现在共享内存或寄存器限制成为问题之前,每个SM可扩展到6个线程束。

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.

相比之下,warp_specialized_pipeline几乎线性扩展,只要您可以为加载、计算和存储分配额外的线程束。在Blackwell上,例如,您可以保持每个SM最多64个驻留线程束的架构限制。(实际驻留取决于寄存器、共享内存和块大小。)

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.)

如表10-3所示,双缓冲和线程束专用方法都大大优于朴素分块内核。double_buffer_pipeline通过重叠块加载和计算将运行时间减半,而warp_specialized_pipeline通过避免任何隐式的块范围等待增加了额外的10.2%加速。只有专用的"计算"线程束会停顿。

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_specialized_pipeline adds another 10.2% speedup by avoiding any implicit block-wide waits. Only the dedicated "compute" warp ever stalls.

指令计数从朴素版本的17亿条下降到两阶段流水线的10.5亿条,减少38%,并进一步下降到线程束专用内核的约10亿条,额外减少4.76%。

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.

L2加载吞吐量从朴素分块的80 GB/s上升到两阶段方法的155 GB/s(+94%),然后上升到线程束专用内核的165 GB/s(比两阶段+6.45%)。这是因为线程束专用内核专门使用一个线程束将每个块加载到共享内存一次--然后将该单个副本多播到所有计算通道。这消除了任何剩余的冗余L2读取。因此,在分块和双缓冲之后,流水线中几乎所有冗余都已被消除。

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.

在实践中,两阶段双缓冲流水线非常适合均匀分块的GEMM工作负载。其更简单的生产者/消费者模型将大部分DRAM延迟隐藏在计算之下。同时,线程束专用方法针对不规则或更深的流水线(如融合注意力内核)进行了优化。这是因为每个线程束可以连续执行其分配的角色--加载、计算或存储--而不会强制块的其余部分停顿。

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.

PyTorch、CUDA Pipeline API和线程束专用化 (PyTorch, CUDA Pipeline API, and Warp Specialization)

PyTorch的公共API不直接暴露cuda::pipeline,但当您调用torch.compile时,编译器生成内核(用OpenAI的Triton语言实现),这些内核使用实现与<cuda/pipeline>原语和线程束专用生产者/消费者等效功能的优化。

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 primitives and warp-specialized producers/consumers.

因此,虽然您不会看到显式生成的cuda::pipeline调用,但您会看到底层指令和屏障。这些优化有助于提高占用率并增加数据传输/计算重叠。换句话说,您无需编写任何CUDA代码即可获得手工编写的<cuda/pipeline>实现的相同低延迟、高吞吐量行为。

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 hand-written implementation without writing any CUDA code.

PyTorch的融合注意力内核由TorchInductor生成,使用带有生产者和消费者组的线程束专用化。例如,考虑PyTorch中三个独立的GPU内核:

PyTorch's fused attention kernels, generated by TorchInductor, use warp specialization with producer and consumer groups. For instance, consider three separate GPU kernels in PyTorch:

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

您可以简单地在PyTorch中用一行重写:

You can simply rewrite this in PyTorch in one line:

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

在底层,加载器线程束将块获取到共享内存中。计算器线程束处理块,存储器线程束将结果写回。这消除了matmul和softmax阶段之间昂贵的全局内存往返。

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.

总的来说,torch.compile非常适合性能敏感和不规则的工作负载。像Blackwell这样的现代GPU拥有丰富的SM资源和高内存带宽。因此,线程束专用化等技术可最小化空闲周期,并允许跨线程束的近线性扩展--至少直到SM或内存带宽完全饱和。

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.

即使像FlashAttention-3这样高度优化的内核,使用线程束专用重叠也只能达到峰值FP16 FLOPS的约75%。这表明您不需要达到100%的计算利用率就能实现重要的优化里程碑。

Even highly optimized kernels like FlashAttention-3 reach only about ~75% of peak FP16 FLOPS using warp-specialized overlap. This shows that you don't need to achieve 100% compute utilization to achieve a significant optimization milestone.

通过将加载、计算和存储解耦为独立阶段,流水线模型最大化了吞吐量和资源利用率,使其成为长时间运行内核的首选方法,包括transformer注意力、融合操作符流水线和自定义任务调度器等过程。这些内核使用细粒度的线程束间通信、深度流水线和线性线程束扩展来提供高性能。

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 pipelines, and linear warp scaling to deliver high performance.

持久内核和兆内核 (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.

当持久线程循环时,它们在数据块到达时处理它们--通常使用内存副本或主机信号--而不退出内核。这完全避免了重复的内核启动开销。例如,持久内核可能使用一个线程块的线程束0将数据从全局内存(或CPU主机内存)复制到共享内存。同时,线程束1计算上一批次。这是GPU上软件流水线的一种形式。

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.

例如,考虑有1,000个微小的独立任务。传统上,可能会启动1,000个单独的内核。每个内核只占用几个SM很短的时间然后退出。

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.

在实践中,GPU会为每个小内核反复加速和减速。这将使大多数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在整个工作负载期间保持忙碌。在具有132个SM的GPU上,例如,这可能意味着每个SM启动一个块,每个块256个线程。总共33,792个线程。然后每个线程执行大致如下的代码:

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:

__device__ int g_index;  // next task的全局计数器;在主机上初始化为0

__global__ void persistentKernel(Task* tasks, int totalTasks) {
    // 每个线程循环,原子地获取下一个任务索引直到没有剩余
    while (true) {
        int idx = atomicAdd(&g_index, 1);
        if (idx >= totalTasks) break;
        processTask(tasks[idx]);
    }
}

在启动此内核之前,主机会在设备内存中设置g_index = 0。然后它会调用内核:

Before launching this kernel, the host would set g_index = 0 in device memory. It would then invoke the kernel as shown here:

cudaMemset(&g_index, 0, sizeof(int));

// 在132 SM GPU上每个SM一个块
int blocks         = 132;  
int threadsPerBlock = 256;

persistentKernel<<<blocks, threadsPerBlock>>>(d_tasks, totalTasks);

cudaDeviceSynchronize();

现在,您只需为每个任务支付一次启动开销,而不是为每个单独的任务支付启动开销。假设启动大约需要0.02 ms,1,000个任务中的每个运行0.1 ms,内核将连续运行大约100 ms的总工作量。

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.

相比之下,连续运行1,000个小内核,每个需要0.02 ms启动,将在执行时间中增加20 ms的开销--与所有1,000个任务的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.

使用持久内核,GPU保持高利用率--几乎所有SM同时活跃处理任务--因为有数万个线程可用于处理约1,000个任务。相比之下,连续启动1,000个小内核会使GPU在任何给定时间大部分未被充分利用。在我们的示例中,这相当于平均只使用了GPU容量的约35%。

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.

在这个持久内核场景中,每个SM运行一个256线程的块(最多2,048个线程中的256个),因此每个SM都在工作。因此,即使每个SM自己的占用率相对较低(12%,即256 / 2048个线程),近100%的SM都是活跃的。

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.)

以这种方式使用持久内核,GPU可以保持高测量的SM活跃百分比,具体取决于寄存器和共享内存使用情况。记得始终使用Nsight Compute进行验证。

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上,持久内核特别有效,因为它们更大的共享内存容量和扩展的寄存器文件允许每个线程在片上保存更多中间状态。线程可以使用TMA为即将到来的任务预取张量块,而其他线程束进行计算。因此,当一些线程正在处理一个任务时,其他线程使用TMA为即将到来的任务预取数据--而不会给SM的计算流水线带来内存传输指令的负担。

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.

持久内核的常见工作负载 (Common Workloads for Persistent Kernels)

持久内核在处理许多小型或不均匀大小的任务时表现出色,如果单独处理这些任务会产生高启动开销。它们允许动态负载均衡。这允许更快的线程继续循环并获取更多工作。使用持久内核,没有SM会过早空闲。

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.

这种模式在LLM推理中常见的不规则工作负载中很常见,如图遍历、自定义批处理转换和每token操作。在这些情况下,每个任务的时间可能差异很大。

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.

然而,持久内核也有缺点。首先,您必须使用原子操作显式管理任务队列和同步。如果许多线程同时尝试增加同一个计数器,这可能会引入争用。

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.

调试单个巨大的持久循环比调试多个小内核更复杂。这是因为单个分歧线程或意外分支可能导致整个内核挂起。此外,一个持久内核可以无限期地独占GPU。因此,如果其他工作负载需要并发运行,您必须仔细分配流或分区资源。

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.

简而言之,与朴素内核相比,持久内核可以大幅提高整体吞吐量(例如2-3倍),方法是将GPU转变为持续获取和处理任务的动态"工作线程"池。在现代GPU硬件上,这种方法消除了启动开销,最大化了SM占用率,并且当与协作组或线程块集群(稍后描述)结合时,可以在多阶段流水线中保持数据在片上。

In short, persistent kernels can increase overall throughput substantially (e.g., 2-3x) 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上的设备端任务队列来保持SM完全占用并执行有用的工作。

As of this writing, PyTorch does not automatically fuse an entire multistage workload into one kernel due to scheduling complexity. As such, achieving the full benefits of persistent kernels and megakernels requires custom CUDA code or specialized compilers. Nevertheless, for multiphase algorithms, refactoring into persistent megakernels can produce significant performance gains-as long as you properly handle synchronizations and avoid deadlocks.

用于推理的兆内核 (Megakernels for Inference)

此外,一种源自大规模推理的现代持久内核方法称为兆内核。兆内核将跨层--甚至跨GPU--的整个操作序列融合为单个大型内核。如图10-8所示,持久兆内核已显示出通过消除重复的内核启动开销,将延迟比传统的每层启动降低1.2倍到6.7倍。

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 been shown to reduce latency by 1.2x to 6.7x versus traditional per-layer launches by eliminating repeated kernel launch overhead.

图10-8 使用兆内核相对于vLLM和SGLang的解码吞吐量改进(来源:https://oreil.ly/2aZiF)

Figure 10-8. 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.

持久内核调度的一个限制是为持久内核找到足够的SM来利用。如果太多SM被另一个内核占用,可能没有足够的资源来启动持久内核。这使得在尝试跨SM调度和负载均衡工作时具有挑战性。

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.

为了促进持久内核(以及因此的线程束专用化),现代GPU支持线程块集群--也称为协作线程阵列(CTA)集群,因为线程块也称为协作线程阵列。我们将在接下来的部分讨论线程块(CTA)集群,但简而言之,它们让您将线程块组合成"集群",占据GPU上多个附近的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.

协作组 (Cooperative Groups)

协作组让您以任意粒度定义和同步线程组。例如,您可以创建具有单个线程、线程束、tile、块和集群的组,如图10-9所示。

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.

图10-9 使用协作组跨不同粒度的线程进行同步

Figure 10-9. Synchronizing across different granularities of threads using cooperative groups

协作组提供安全、可重用的集合操作,如同步、广播和归约。这与使用临时同步屏障形成对比。通常,线程只能使用__syncthreads()在自己的块内同步--例如,整个网格没有内置的全局屏障。

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.

协作组为您提供内核内的细粒度同步。该API非常适合跨线程束、块和集群协调多阶段流水线。要在内核中使用协作组API,只需包含<cooperative_groups.h>,获取组对象,然后在这些组之间同步和协调。

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 , obtain group objects, then synchronize and coordinate across these groups.

API包括cg::this_thread_block()cg::tiled_partition()cg::this_cluster()等调用。然后您调用group.sync()--或类似的集合操作--来协调这些组中的线程。

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.

要以协作模式启动内核,您使用cudaLaunchCooperativeKernel()。在协作模式下,CUDA确保您启动的网格可以同时驻留--否则,启动失败。因此,建议始终使用cudaOccupancyMaxActiveBlocksPerMultiprocessor调整协作网格大小,并限制网格大小以避免启动失败。在内核内部,您可以调用以下内容来实现全线程块屏障:

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 limit the grid size to avoid launch failure. Inside the kernel, you can call the following to implement a full thread block barrier:

cooperative_groups::this_grid().sync(); // 或 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.

例如,考虑LLM中常见的softmax算法,它有两个阶段:跨整个数组归约以计算聚合和,然后使用聚合和进行后续的每元素计算。传统上,您会启动一个内核进行归约,将结果复制回主机内存或全局内存,启动第二个内核从主机或全局内存消费结果,然后计算softmax。这需要大量相对较慢的内存移动。

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.

使用协作组,您可以在一个内核中执行两个阶段,使每个块计算其部分和,一个块将这些部分和聚合为最终结果,所有块调用grid.sync()等待聚合完成,然后所有线程进入第二阶段。然后每个线程将从寄存器或共享内存读取聚合和--而不是从全局内存。

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.

CUDA保证协作启动中的每个块同时在GPU上驻留。如果您请求的块超过可以同时容纳的数量,启动只会失败。

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.

因为协作内核必须以GPU可以同时运行的网格大小启动,grid.sync()不会挂起等待不存在的块。换句话说,CUDA运行时保证所有线程块都是活跃的并将到达grid.sync()屏障。如果内核启动太大而无法一次运行,它只会启动失败。因此,检查cudaLaunchCooperativeKernel()的返回状态很重要。

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 cudaLaunchCooperativeKernel().

例如,如果GPU有132个SM,您的内核使用足够多的资源使每个SM可以运行4个块,则网格必须不超过528个块才能成功。如果超过该限制,协作启动只会失败。使用cudaOccupancyMaxActiveBlocksPerMultiprocessor调整协作内核的网格大小,并在假设取得进展之前检查cudaLaunchCooperativeKernel返回状态。

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 cudaOccupancyMaxActiveBlocksPerMultiprocessor to size the grid for a cooperative kernel and check the cudaLaunchCooperativeKernel return status before assuming progress has been made.

在Blackwell之前,开发人员经常使用cudaLaunchCooperativeKernel与全局内存原子操作或标志来协调不共享片上内存的多个块。这种方法有效,但强制中间结果移动到全局内存。这会产生额外的HBM流量。

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上,线程块集群和DSMEM提供了更高效的替代方案。线程块集群可以在片上SRAM中共享数据--并在没有全局内存往返的情况下同步。我们稍后将在本章介绍线程块集群和DSMEM。

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.

当您需要在单个内核启动中使用真正的全线程块屏障时,应该使用协作内核。此外,当您希望将中间结果保持在快速内存(例如寄存器或共享内存)中,而不是反复写入和读取全局内存时,它们也很有用。建议您将网格大小限制为GPU的最大并发块容量--或选择更大的块--以减少总体块数和失败的内核启动。

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.

协作内核的缺点是您的网格大小受容量限制。这可能会迫使您使用更少、更大的块--或依赖线程块集群(稍后介绍)。

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).

即使所有块并发运行,协作屏障仍然需要每个块中的每个线程调用grid.sync()。如果任何单个线程跳过--或从未到达--grid.sync()调用(例如,由于分歧的if语句),那么确实调用了grid.sync()的每个其他线程将永远等待。这会导致死锁。

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.

简而言之,协作组内核让您将整个GPU视为单个协作资源,grid.sync()充当全局屏障。这对于需要全局同步和数据共享的多阶段算法是理想的。只需记住grid.sync()是相对重量级的同步,比块范围或集群范围的屏障开销更高。这是因为它必须协调跨多个SM运行的所有线程块。因此,您应该谨慎使用grid.sync(),并且至少确保您的内核在重量级同步调用之间做大量工作。

In short, cooperative group kernels let you treat the entire GPU as a single collaborative 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.

对于只需要有限跨块协调的情况,使用全局内存原子操作或每块标志可能比依赖全网格屏障更安全、更简单。但是,它们不如使用线程块集群和DSMEM高效,我们稍后将讨论。

For cases where you need only limited cross-block coordination, using global-memory atomics or per-block flags can be safer and simpler than relying on a full-grid barrier. However, they are less efficient than using thread block clusters and DSMEM, as we'll discuss in a bit.

协作网格同步和持久内核 (Cooperative Grid Synchronization and Persistent Kernels)

如果您的工作负载偶尔需要在持久循环内进行跨线程块屏障来聚合部分结果,您可以通过调用grid.sync()将持久内核与协作组结合。这将提供网格范围的屏障,以避免结束内核并重新启动它。这样,归约和其他全局步骤的多阶段流水线完全保留在设备上。

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.

例如,考虑一个工作负载在1,000次迭代中重复执行两种不同的计算,使得每次计算由于跨块数据依赖性而需要全局屏障。朴素实现可能每次迭代启动两个单独的内核,导致2,000次内核启动。在单个流中,第二个内核自动等待第一个--无需显式的主机端同步--但您仍然要支付启动开销。相反,您可以将所有内容融合到一个协作、持久的内核中,如下所示:

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:

#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(再次网格步长)
        for (int i = tid; i < N; i += stride) {
            const float mid = dataA[i];
            dataB[i] = someComputationB(mid, dataB[i]);
        }

        grid.sync();  // 下一次迭代前的屏障
    }
}

在主机端,您只需调用此组合的协作和持久内核:

On the host side, you would simply invoke this combined cooperative and persistent kernel, as shown here:

// ... 分配/复制 dA, dB, 设置 N, iterations, 选择 blockSize
cudaDeviceProp prop{};
cudaGetDeviceProperties(&prop, 0);
if (!prop.cooperativeLaunch) {
    // 回退:每次迭代运行两个内核而不是 grid.sync()
    // (见下面的替代方案)
}

// 计算协作内核允许的最大网格大小
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();

这个单一内核取代了原本需要2 x 1,000 = 2,000次内核启动的操作,并避免了重复的启动开销。在循环内部,grid.sync()确保正确的顺序(每个块在任何块开始阶段2之前完成阶段1--以及在下一次迭代之前),无需任何主机端同步。

This single kernel replaces what would otherwise be 2 x 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.

每线程和每块数据可以跨grid.sync()保留在寄存器/共享内存中。对于跨块数据交换,使用全局内存(或线程块集群和分布式共享内存,稍后介绍)。同时,因为工作在GPU上循环,第一次调用后没有启动开销。

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.

协作内核需要设备支持(cooperativeLaunch)并且必须使用cudaLaunchCooperativeKernel启动。网格中的所有块必须同时驻留。使用CUDA占用率API调整和限制网格大小,以便所有CTA可以共存。否则启动将失败。一个例子是当((N + threads - 1) / threads)超过协作容量时。

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.

使用CUDA占用率API调整协作网格大小。

Use the CUDA Occupancy APIs to size cooperative grids.

现代GPU可以同时运行大约几千个线程块(假设每个使用适量的资源,包括寄存器和共享内存),因此有相当大的空间。但是,您仍应在继续此实现之前验证块大小和资源消耗。

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.

在这个内核示例中,内核在迭代之间从不将控制权返回给主机。因此,它在外部循环中表现得像持久内核,但在内部阶段强制执行全局同步点。

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.

这种组合模式对于LLM推理中的多步算法特别强大,其中每层可能需要归约或归一化(阶段1),然后是每元素转换(阶段2)。通过将所有内容捕获在协作持久内核中,您消除了所有内核间往返并最大化片上数据局部性。

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.

何时结合持久内核和协作组 (When to Combine Persistent Kernels and Cooperative Groups)

一个常见的最佳实践是如果资源允许,每个SM启动一个带有单个线程块的持久内核。这样,每个SM都有一个驻留的线程块,迭代地从工作队列处理任务。这最大化了占用率并保持SM不会空闲。然后线程块将使用grid.sync()在它们的阶段之间协调。但是,在决定协作和持久策略--以及是否结合它们时--请考虑以下问题:

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:

您是否有需要全局同步的多个顺序阶段? 如果是,使用协作内核(或协作持久内核),以便您可以在阶段之间调用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.

您是否有许多受启动开销影响的小型或不规则任务? 如果是,使用持久内核,以便您的线程在共享队列上循环而不返回主机。

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.

您是否需要与其他工作共享SM? 如果是,考虑使用线程块集群(下一节),以便持久或协作内核不会服务其他工作负载。

Do you need to share SMs with other work? If yes, consider using thread block clusters (next section) so the persistent or cooperative kernel does not serve other workloads.

简而言之,当您使用单个内核既持久地在任务上循环又包含grid.sync()调用以在阶段之间同步所有线程块时,您消除了通常在单独的内核启动之间发生的暂停和额外内存传输。

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.

使用现代GPU,这意味着数据在整个计算过程中保持在共享内存或寄存器中。这与每个阶段后将数据写回全局内存形成对比。因此,GPU几乎一直都在做有用的工作--实现接近其峰值硬件限制的性能。

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.

一个重要的注意事项:协作内核保留网格中的所有SM,因此没有其他内核可以在这些SM上并发运行。如果您需要在单独的流上共同调度其他工作,如异步数据预取或低优先级推理内核,您可能需要将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.

线程块集群和分布式共享内存 (Thread Block Clusters and Distributed Shared Memory)

协作组是一个软件级抽象,提供API让您将内核线程划分为任意集合用于同步和数据移动。这包括线程束、tile、线程块--甚至整个网格或多设备网格。

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.

相比之下,线程块集群(或协作线程阵列[CTA]集群)是一个硬件级层次结构。它为您的协作网格授予一部分SM--并将剩余部分留给其他内核使用。这减轻了一个内核独占GPU的风险。GPU保证线程块将在同一个GPU处理集群(GPC)上共同调度,如图10-10所示。

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.

图10-10 多个线程块集群保证在同一个GPC或GPC分区上共同调度

Figure 10-10. Multiple thread block clusters are guaranteed to be coscheduled on the same GPC or GPC partition

GPC是附近SM的集合。GPU将线程块调度到GPC上,类似于它将线程块的线程调度到同一个SM上。

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.

实际上,在多芯片模块如NVIDIA Blackwell B200/300和GB200/300上存在多个GPC"分区"--每个芯片一个GPC分区。由于Blackwell是双芯片GPU,它有两个GPC分区。请记住,Blackwell的两个GPU芯片通过NV-HBI连接,并作为单个CUDA设备呈现,具有跨芯片的完整缓存一致性。L2缓存在芯片间也是一致的。因此,这些芯片形成一个组合的逻辑GPC,架构会为您处理单独的GPC分区。

There are actually multiple GPC "partitions" on multidie modules like the NVIDIA Blackwell B200/300 and GB200/300-one GPC partition per die. Since Blackwell is a two-die GPU, it has two GPC partitions. Remember that Blackwell's two GPU dies are linked with NV-HBI and present as a single CUDA device with full cache coherence across dies. The L2 caches are coherent across dies, as well. As such, the dies form a combined logical GPC so the architecture handles the separate GPC partitions for you.

GPU为线程块集群提供分布式共享内存(DSMEM),用于跨这些块使用。它还支持使用集群组API(cluster.sync())的集群级屏障。

The GPU provides distributed shared memory (DSMEM) for use across these blocks for a thread block cluster. It also supports cluster-level barriers using the cluster group API (cluster.sync()).

此集群级屏障让您只同步一部分线程块而不阻塞整个GPU。线程块集群让您启动一个协作内核,将网格细分为更小的块组,如图10-11所示。

This cluster-level barrier lets you synchronize just 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.

图10-11 对于这四个(2x2)线程集群,A和B的每个块同时加载到两个线程块中(来源:https://oreil.ly/kEZsv)

Figure 10-11. For these four (2 x 2) thread clusters, each tile of A and B is loaded into two thread blocks simultaneously (source: https://oreil.ly/kEZsv)

在每个组内,调用cluster.sync()提供本地屏障。这让集群内的块通过专用片上资源共享数据而不独占每个SM。在现代GPU上,您可以使用DSMEM,它允许线程块集群共享片上SRAM的连续区域。这实现了同一集群中块之间的低延迟通信,并有原生硬件支持。

Within each group, calling cluster.sync() provides a local barrier. This lets blocks in the cluster share data via dedicated on-chip resources without monopolizing every SM. On modern GPUs, you can use DSMEM, which allows a thread block cluster 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()来"本地"同步这些线程块。虽然线程块内的线程传统上可以使用共享内存进行协作,但在现代GPU上,它们现在可以使用线程块集群和DSMEM相互协作。

Thread block clusters group a subset of thread blocks and synchronize them "locally" by calling cluster.sync() within each cluster. While threads within a thread block have traditionally been able to cooperate using shared memory, on modern GPUs they can now cooperate across thread blocks using thread block clusters and DSMEM.

换句话说,没有线程块集群,只有同一线程块内的线程可以在共享内存中共享数据。不同的线程块必须使用全局内存和全局屏障(例如grid.sync())进行协调,这会停顿所有线程并限制可扩展性。

In other words, without thread block clusters, only threads within the same thread block can share data in shared memory. Different thread blocks must coordinate using global memory and global barriers (e.g., grid.sync()), which stalls all threads and limits scalability.

此外,以前不同线程块中的线程无法高效共享状态和同步,除非使用全局内存或grid.sync()进行粗粒度的网格范围同步,如前一节所述。不幸的是,全局内存和grid.sync()都相对较慢,因此可能成为瓶颈。

Furthermore, previously threads in different thread blocks could not efficiently share state and synchronize without using 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 thus can be a bottleneck.

线程块集群由GPU硬件原生支持,包括集群启动控制。集群启动控制是一种硬件级机制,用于启动和调度持久线程块集群。具体来说,它允许持久内核(及其线程块集群)保持平衡的工作负载--即使一些SM被占用。这为高效的线程束专用化实现提供了基础。

Thread block clusters are natively supported by GPU hardware, including cluster launch control. Cluster launch control is a hardware-level mechanism for launching and scheduling persistent thread block clusters. Specifically, it allows persistent kernels (and their thread block clusters) to maintain a balanced workload-even if some SMs are occupied. This provides the basis for efficient warp specialization implementations.

使用硬件支持的通信和同步构造,线程块集群可以使用低级PTX指令和CUDA内建函数的形式进行集群范围屏障同步。因此,集群中的块可以比使用grid.sync()的全网格同步快得多地执行屏障同步,这是由于线程块之间的硬件支持屏障同步。

Using hardware-supported communication and synchronization constructs, thread block clusters can perform cluster-wide barrier synchronization in the form of low-level PTX instructions and CUDA intrinsics. As a result, blocks in a cluster can perform barrier synchronization much faster than a full-grid synchronization using grid.sync(), due to hardware-supported barrier synchronization between thread blocks.

线程块搅动 (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.

相反,您希望单个波中的块A和B可以从L2缓存中读出。

Instead, you want tiles A and B within a single wave, which can be read out of L2 cache.

为了解决这种低效问题,您可以使用线程块搅动。与使用搅动优化内存访问和避免共享内存bank冲突类似,您可以使用线程块搅动避免以低效的行优先和列优先顺序分配块,如图10-12所示。

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.

图10-12 线程块搅动以在单个波中从L2缓存读取块A和B

Figure 10-12. Thread block swizzling to read tiles A and B in a single wave out of L2 cache

线程块搅动让同一波所需的A和B矩阵的块保持在L2中以实现最大重用。当应用于持久和分块GEMM工作负载时,这种类型的搅动可以通过减少内存未命中和带宽压力产生两位数的性能提升。

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.

有了对线程块集群的硬件支持、相对大量的SM和改进的流水线调度,现代GPU可以托管具有复杂线程束专用化流水线的大型持久内核来优化内核性能。

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 complex warp-specialized pipelines to optimize kernel performance.

简而言之,线程块集群建立在协作组模型之上,允许一部分块同步和共享状态而不锁定整个GPU。这让您在一次启动内构建多阶段、细粒度流水线而不锁定设备的其余部分。这使剩余的SM可以自由执行独立工作。

In short, thread block clusters build on the cooperative groups model by allowing a subset of blocks to synchronize and share state without locking the entire GPU. This lets you build multistage, fine-grained pipelines within one launch without locking the rest of the device. This leaves remaining SMs free to perform independent work.

在线程块集群内,有两种在线程块之间共享状态的关键机制:DSMEM和暂存内存。让我们接下来描述这些。

Within a thread block cluster, there are two key mechanisms for sharing state between thread blocks: DSMEM and scratch memory. Let's describe these next.

分布式共享内存 (Distributed Shared Memory)

分布式共享内存(DSMEM)将__shared__内存的概念扩展到单个线程块之外,跨越整个线程块集群。在传统内核中,每个线程块都有自己的私有__shared__区域,其他线程块无法访问。然而,使用DSMEM,同一集群中的多个块可以读/写到逻辑上组合了它们所有本地__shared__区域的共享内存空间。实际上,集群的共享内存被拼接成一个分布式片上缓冲区。

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. 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显著提高了有效算术强度,因为块可以交换中间结果或执行归约而无需往返全局内存。这对于数据集对单个块的共享内存来说太大的情况特别有价值。

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同步和共享。所有通信都以SM到SM的速度发生,避免了昂贵的HBM流量。

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是使用全局内存进行块到块通信的更快、更有结构的替代方案。但其范围仅限于同一集群中的块。DSMEM非常适合需要频繁、低延迟块间协调的持久内核、注意力机制或LLM中其他多阶段算法,其中每个块在继续之前必须交换中间状态--以及任何受益于将临时结果保持在片上而不是写回和重新加载全局内存的工作负载。

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.

可移植的最大集群大小是8个块。一些GPU支持更大的不可移植集群大小(例如,最多16个CTA),当使用cudaFuncAttributeNonPortableClusterSizeAllowed属性显式启用时。这会增加您的DSMEM占用空间,但代价是占用集群中更多的SM。

The portable maximum cluster size is 8 blocks. Some GPUs support larger non-portable cluster sizes (e.g., up to 16 CTAs) when explicitly enabled with the cudaFuncAttributeNonPortableClusterSizeAllowed attribute. This increases your DSMEM footprint at the cost of occupying more SMs in the cluster.

暂存内存 (Scratch Memory)

暂存内存是支持DSMEM和线程块集群同步的底层硬件基础设施,而DSMEM是对CUDA代码隐式可见的共享数据缓冲区,暂存内存是一个单独的片上SRAM区域,GPU用它来跟踪协调元数据,如屏障状态、组进度和DSMEM的访问标志。

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.

您不从内核直接访问暂存内存。相反,GPU自动管理它。当您启动线程块集群时,硬件分配一部分暂存内存来维护集群范围的屏障计数器(cluster.sync()状态)、跟踪哪些块已到达DSMEM操作或同步点,并协调跨SM对分布式共享区域的安全访问。

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.

因为暂存内存针对这些元数据操作进行了优化,它使集群级屏障和DSMEM访问能够非常快速地完成。如果元数据大小超过了非常大集群或复杂同步模式的可用暂存SRAM,GPU会透明地将一些状态溢出到本地内存,后者由L1/L2缓存支持。这种溢出确保正确性,同时仍尽可能保持片上效率。

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是您用来在块之间共享数据的抽象,暂存内存是使这些DSMEM操作快速和可扩展的幕后设施。它们共同使线程块集群表现得像一个单一逻辑单元,打破了线程块之间的传统屏障。这提高了需要跨多个线程块紧密协调并行性的工作负载的性能。

In short, DSMEM is the abstraction you use to share data between blocks and scratch memory is 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.

启动线程块集群 (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.

在CUDA中,这通过扩展启动API完成,允许指定集群大小。例如,如果我们想要两个块的集群,也称为线程块对或CTA对,我们可以配置特殊的启动属性:

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:

// 主机代码:启动具有大小为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);

在这个主机代码中,我们使用cudaLaunchKernelEx启动MyClusterKernel,集群大小为2个线程块(总共64个集群,因为128个块 / 每集群2个)。clusterDim设置为2,意味着每个集群将包含2个线程块。这取代了传统的<<<gridDim, blockDim>>>语法,用于我们想要集群协作内核的情况。在底层,CUDA运行时确保这些配对的块被调度,使它们可以使用DSMEM通信。

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 <<>> 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.

请记住,每个线程块集群需要物理上适合您指定的SM和资源。可移植的最大集群维度是8个线程块。要在支持的Blackwell部件上启动16个线程块,您必须启用不可移植属性并调整块大小,以便集群保持在SM上驻留。在调整线程块集群大小时请记住这一点。

Remember that each thread block cluster needs to physically fit into the SMs and resources that you specify. The portable maximum cluster dimension is 8 thread blocks. To launch 16 on supported Blackwell parts, you must enable the non-portable attribute and size blocks so the cluster stays resident on the SMs. Keep this in mind when sizing your thread block clusters.

使用协作组API协调线程块集群 (Coordinating Thread Block Clusters with Cooperative Groups API)

要协调集群中多个线程块的工作,您首先使用cooperative_groups::this_cluster()获取该块组的句柄。此集群句柄让您执行硬件加速的集群范围屏障--并直接访问另一个块的共享内存。

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.

这一切都在不离开内核或诉诸全局内存标志的情况下发生。以下是一个示例内核,它将每个线程块的本地值求和到线程块0的共享内存中:

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:

#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 __shared__ int shared_buffer[];

    // 3. 每个块了解其排名和总集群大小
    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);

    if (clusterRank != 0) {
        // 7. 非零块将其本地shared_buffer[0]原子添加到
        //    排名0的缓冲区。此atomicAdd通过DSMEM在片上路由
        //    (不通过DRAM)
        atomicAdd(&remote_buffer[0], shared_buffer[0]);
    }

    // 8. 另一个集群范围屏障确保所有原子添加已完成
    cluster.sync();

    // 9. 最后,块0(排名0,线程0)可以读取组合结果
    if (clusterRank == 0 && threadIdx.x == 0) {
        printf("Combined sum in cluster[0]: %d\n", shared_buffer[0]);
    }
}

在这里,我们通过调用cg::this_cluster()获取集群句柄。这返回一个cluster_group对象,表示作为线程块集群一起启动的那些块。运行时保证此集群组中的每个线程块同时驻留--否则,启动将失败。

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.

提供了一个统一的共享内存分配。集群中的每个块必须为extern __shared__ int shared_buffer[]声明相同的大小。然后DSMEM硬件将逻辑上把每个块的共享内存组合成一个虚拟地址空间。因为所有块保留相同的shared_buffer大小,指向线程块0(或排名0)缓冲区的指针将跨SM正确协调。

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.

使用线程块集群时,请确保像全局内存合并一样组织每个线程块的DSMEM访问。换句话说,让每个线程束读取或写入连续的、32字节对齐的扇区。这样,硬件可以在没有存储体冲突或意外序列化的情况下路由集群范围的DSMEM传输。在实践中,布置您的块数据,使线程块j中的线程束i始终触及唯一的、对齐的范围。这将避免内存存储体争用,并使DSMEM数据传输保持全速执行。

When using thread block clusters, make sure to organize each thread block's DSMEM accesses just like global-memory coalescing. In other words, have each warp read or write contiguous, 32-byte-aligned sectors. That way, the hardware can route cluster-wide DSMEM transfers without bank conflicts or unexpected serialization. In practice, lay out your tile data so that the warp i in thread block j always touches a unique, aligned range. This will avoid memory bank contention and keep DSMEM data transfers performing at full speed.

在前面的示例代码中,我们还看到cluster.sync()被用作集群中线程块之间的集群范围屏障。与只同步单个块内线程的__syncthreads()不同,cluster.sync()同步此集群中每个块的所有线程。

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.

此屏障在硬件中实现,通常比网格级同步延迟更低,因为它只协调集群中的块。这意味着您可以频繁同步块而影响最小--只要它们在集群内。

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.

没有因缺少块而导致死锁的风险,因为CUDA强制网格正确启动并适合GPU。因此,不存在某些块到达屏障然后永远等待从未启动的"缺失"块的情况。所有块都已经存在,所以一旦每个块调用它,屏障就会干净地完成。

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.

在前面的代码块中,您看到使用map_shared_rank()进行跨块共享内存访问。在第一个屏障之后,每个块的shared_buffer[]都已初始化。要获取指向块0共享内存的指针,其他块调用int* remote_buffer = cluster.map_shared_rank(shared_buffer, 0),这指定了集群中的线程块ID(0)。

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.

GPU硬件自动翻译该指针,以便任何加载或存储都通过片上DSMEM网络而不是全局DRAM。因此,任何写入remote_buffer的操作--无论是原子操作还是常规写入--都将直接在片上更新块0的共享内存。值得强调的是,远程DSMEM访问的性能特征与本地共享内存不同。远程加载和存储受益于合并的、对齐的32字节段。

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.

例如,在前面的代码块中,当块1...n需要将其本地值添加到块ID 0的共享缓冲区时,它们调用atomicAdd(&remote_buffer[0], shared_buffer[0])。因为remote_buffer是通过cluster.map_shared_rank()获得的,atomicAdd通过片上DSMEM网络直接进入块0的SMEM。换句话说,没有到DRAM或L2缓存的往返。每次写入都以片上速度发生。

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 through cluster.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 speed.

在前面的代码示例中,您看到一旦所有块执行了它们的atomicAdd,第二个cluster.sync()确保块0在继续之前在其自己的shared_buffer[0]中看到完整的和。

In the previous code example, you see that once all blocks have performed their atomicAdd, a second cluster.sync() ensures that block 0 sees the full sum in its own shared_buffer[0] before proceeding.

简而言之,cooperative_groups::this_cluster()加上cluster.sync()cluster.map_shared_rank()为您提供了一种简单、高效的方式来同步和共享线程块集群中多个线程块之间的数据。所有这些操作都在片上发生,避免了全局内存往返,并实现了块之间的细粒度协作。这种线程块集群和DSMEM的组合提供了比任何全局内存回退或手动原子标志方法更高的性能。

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 this happens on-chip, avoiding global memory round trips and enabling fine-grained cooperation between blocks. This combination of thread block clusters and DSMEM offers higher performance than any global-memory fallback or manual atomic flag approach.

线程块对 (Thread Block Pair)

现代NVIDIA GPU让您在单个GPC内的SM之间共同调度恰好两个线程块在一个集群中,或线程块对(又名CTA对)。通过将线程块分组在一个集群中(例如,如图10-13所示的2块集群),共享数据的内核可以使用TMA将块移动到每个块的共享内存中。

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.

图10-13 线程块对组合两个线程块

Figure 10-13. Thread block pair combines two thread blocks

单个线程块可能缺乏寄存器或共享内存容量来独自处理非常大的块(例如,256 x 256矩阵子块)。通过在GPC内的附近SM上配对两个线程块并使用DSMEM,这两个块可以在一个大型块上分割工作,但通过统一的共享内存区域共享数据,如图10-14所示。

A single thread block might lack the registers or shared-memory capacity to process a very large tile (for example, a 256 x 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.

图10-14 线程块对(又名CTA对)加载块作为A * B矩阵乘法的操作数(来源:https://oreil.ly/kEZsv)

Figure 10-14. Thread block pair (aka CTA pair) loading tiles as operands for an A * B matrix multiply (source: https://oreil.ly/kEZsv)

在这里,对中的每个线程块可以将操作数块的一部分(例如128 x 16)加载到其片上SMEM中用于矩阵乘法。此外,每个线程块在张量内存(TMEM)中保存累加器的一部分(例如128 x 256)。这允许CTA对中的两个线程块在单个块上协作,如图10-15所示。

Here, each thread block in the pair can load a fraction of the operand tile (e.g., 128 x 16) into its on-chip SMEM for the matrix multiply. In addition, each thread block holds part of the accumulator (e.g., 128 x 256) in Tensor Memory (TMEM). This allows the two thread blocks in the CTA pair to collaborate on a single tile, as shown in Figure 10-15.

图10-15 带有Tensor Core和TMEM的线程块对

Figure 10-15. Thread block pair with Tensor Cores and TMEM

线程块对允许跨越两个物理SM的更大矩阵乘法。使用一对线程块有效地使块大小加倍,因为每个SM处理块数据的一半。

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硬件使用DSMEM在SM之间共享操作数数据。DSMEM减少重复加载,改善数据重用,并增加算术强度。图10-16显示了使用DSMEM在线程块集群中的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.

图10-16 使用DSMEM在线程块集群中的SM之间共享数据

Figure 10-16. Sharing data between SMs in a thread block cluster using DSMEM

CUDA为这些多SM操作提供增强的多播屏障和同步原语。这对使用轻量级cluster.sync()调用相互同步。这样,当一个SM完成使用块时,CTA对中的相邻SM可以消费它。

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.

每个块被送入线程块对而没有冗余的全局内存事务。这更好地占用了否则空闲的寄存器、共享内存bank和Tensor Core。线程块对通过使可用于一个块的线程和共享内存加倍来实现更高的Tensor Core利用率。

Each tile is fed to the thread block pair without redundant global-memory transactions. This better occupies otherwise idle registers, shared-memory banks, and Tensor Cores. Thread block pairs achieve higher Tensor Core utilization by doubling the threads and shared memory available to one tile.

任何性能提升都会透明发生。从程序员的角度来看,您只需启动一个包含两个块的集群,并将其视为跨两个线程块分割的一个大型线程块。例如,在CUTLASS中,这被公开为使用DSMEM和集群屏障进行块间数据共享的2-SM UMMA(统一MMA)GEMM操作。

Any performance gains happen transparently. From the programmer's perspective, you simply launch a cluster of two blocks and treat it as one large thread block split across two thread blocks. For example, in CUTLASS, this is exposed as a 2-SM UMMA (Unified MMA) GEMM operation that uses DSMEM and cluster barriers for inter-block data sharing.

您只需为单个GEMM块请求两个SM,CUTLASS会自动处理集群启动、DSMEM设置和线程块间同步。

You simply request two SMs for a single GEMM tile, and CUTLASS automatically handles the cluster launch, DSMEM setup, and inter-thread-block synchronization.

如果您请求单个线程块无法处理的块大小(例如,FP16 Tensor Core操作的128 x 128),CUTLASS将自动分配两个线程块作为一对。SM分割工作并在底层依赖DSMEM + cluster.sync()来共享块。这样,您可以成对重叠UMMA操作与DSMEM。

If you request a tile size that a single thread block cannot handle (e.g., 128 x 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 tiles. This way, you can overlap UMMA operations with DSMEM in pairs.

简而言之,线程块对和DSMEM以多SM协作的形式提供并行性。它们提供跨线程块的快速、片上数据共享和同步。这消除了许多以前需要全局内存传递或额外内核启动的场景。这有利于许多多线程块算法并简化了它们的实现。

In short, thread block pairs and DSMEM provide parallelism in the form of multi-SM collaboration. They provide fast, on-chip data sharing and synchronization across thread blocks. This eliminates many scenarios that previously required global-memory passes or extra kernel launches. This benefits many multi-thread-block algorithms and simplifies their implementation.

使用线程块集群减少全局内存流量 (Reducing Global Memory Traffic with Thread Block Clusters)

从性能角度来看,DSMEM可以显著减少冗余的全局内存流量并实现更高的有效带宽。例如,在分块GEMM中,集群中的多个线程块共享A或B矩阵的块,一个块可以从全局内存加载块并使用TMA将块多播到集群中的其他块。

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.

TMA引擎支持多播复制模式,当线程块属于同一集群时,直接送入DSMEM。单个TMA从全局内存传输可以同时将数据放入每个参与块的共享内存中。这避免了冗余的DRAM获取。

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多播,GPU确保L2缓存数据在一个过程中广播到每个集群成员(线程块)的共享内存中。因此,您避免了同一块的重复全局内存加载。这改善了带宽利用率并减少了DRAM流量--尤其是当许多块需要相同输入时,如图10-17所示。

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.

图10-17 对于这四个(2x2)线程集群,每个块加载一次并多播到每个集群中所有CTA的共享内存中(来源:https://oreil.ly/kEZsv)

Figure 10-17. For these four (2 x 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)

在这里,TMA引擎执行从全局内存到DSMEM的单个多播,将块广播到集群中每个线程块的SMEM,并消除冗余的DRAM读取。TMA多播通过张量映射描述符配置,并作为针对shared::clustercp.async.bulk.tensor发出。幸运的是,像CUTLASS/cuTe和Triton这样的高级库为您生成这些多播操作、张量映射描述符和批量张量副本。具体来说,这些库可以发出相关的PTX指令,包括带有张量映射操作数的cp.async.bulk.tensor。它们可以发出针对shared::cluster进行多播的PTX指令。在这些情况下,硬件将块传送到每个线程块的SMEM。

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.

另一个常见用例是多块归约或扫描。不是让每个块将其部分和或扫描结果写入全局内存然后启动单独的内核来组合它们,线程块集群中的线程块可以将其部分结果写入DSMEM。

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.

在该集群内,您可以使用共享内存和集群级屏障cluster.sync()对这些部分结果执行快速片上归约或前缀和。只有最终结果需要进入全局内存。这大大减少了全局内存读取和写入。

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, cluster.sync(). Only the final result needs to go to global memory. This greatly reduces global memory reads and writes.

通过结合线程块集群、DSMEM和TMA多播,您可以构建多阶段、细粒度流水线,保持大多数数据交换在片上。无论您是为GEMM共享块、为归约累加部分和,还是执行多块扫描,这些机制都让您最小化HBM往返并最大化算术强度。

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.

块范围的cuda::pipeline不会同步其他块;集群范围分发使用TMA多播或带cluster.sync的DSMEM。

A block scoped cuda::pipeline does not synchronize other blocks; cluster wide distribution uses TMA multicast or DSMEM with cluster.sync.

考虑两个线程块CTA 0和CTA 1,它们都需要矩阵A的同一块。没有DSMEM,每个CTA发出自己的全局内存加载,这通过获取相同的数据两次浪费了DRAM带宽。然而,使用DSMEM,CTA 0将块加载到其共享内存一次,然后通过片上DSMEM网络多播,以便CTA 1可以直接从共享内存读取它。如果超过两个CTA形成一个集群,同一块将在集群中的所有线程块之间共享,但这仍然只需要一次全局HBM加载。表10-4显示了使用两个线程块的集群与两个独立线程块的比较。

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, wasting DRAM bandwidth by fetching the same data twice. With DSMEM, however, CTA 0 loads the tile into its shared memory once and then multicasts it over the on-chip DSMEM network so CTA 1 can read it directly from shared memory. If more than two CTAs form a cluster, the same tile is shared among all thread blocks in the cluster-but still only one global HBM load is required. Table 10-4 shows a comparison of using a cluster of two thread blocks versus two independent thread blocks.

表10-4. DSMEM对两块工作负载的性能影响

指标 两个独立CTA(无DSMEM) 带DSMEM的CTA对(2个集群)
全局加载事务 2x(每个线程块加载块) 1x(块加载一次)
L2缓存命中率 50% 85%
块间数据重用 无(无重用) 显著(CTA 1重用块)
每CTA有效DRAM带宽 300 GB/s 150 GB/s(减少50%)
内核时间(相对) 1.0x 0.6x(加速40%)

Table 10-4. Performance impact of DSMEM on a two-block workload

在这里,我们看到使用带有DSMEM的线程块对(又名CTA对)时,内核执行有40%的加速。这与带宽节省一致。每个线程块的DRAM带宽从300 GB/s下降50%到150 GB/s,因为每个块只获取一次。L2命中率从50%跃升至85%,因为TMA硬件正在执行片上多播复制--因此避免了从全局内存的多次加载。

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. DRAM bandwidth per thread block drops 50% from 300 GB/s to 150 GB/s because each tile is fetched only once. L2 hit rate jumps from 50% to 85% because the TMA hardware is performing on-chip multicast copies-so multiple loads from global memory are avoided.

简而言之,通过多播共享数据,线程块集群允许多个块以片上速度重用数据。这为内存受限的工作负载带来了显著的加速。

In short, by multicast sharing data, thread block clusters allow multiple blocks to reuse data at on-chip speeds. This yields significant speedups for memory-bound workloads.

重要的是要注意,DSMEM和L2并行运行。这为块间数据共享提供了两条"通道"。这种双路径设计将DSMEM的超快集群范围通信与L2更广泛的缓存覆盖相结合。换句话说,DSMEM访问绕过集群本地地址的L2缓存。相反,DSMEM使用专用的SM到SM网络,如图10-18所示。

Importantly, DSMEM and L2 operate in parallel. This provides two "channels" for interblock data sharing. This dual-path design combines DSMEM's ultra-fast cluster-wide communication with L2's broader cache coverage. In other words, DSMEM accesses bypass the L2 cache for cluster-local addresses. Instead, DSMEM uses a dedicated SM-to-SM network, as shown in Figure 10-18.

图10-18 DSMEM使用两个线程块集群之间的SM到SM线程块集群本地网络进行数据交换

Figure 10-18. DSMEM uses an SM-to-SM thread block cluster-local network between two thread block clusters for its data exchange

远程DSMEM访问使用线程块集群互连路由,与全局内存流量不同。例如,当线程块使用DSMEM拉取块时,它使用低延迟共享内存传输。这确保线程块集群中的线程块间数据共享与片上共享内存一样快。

Remote DSMEM accesses are routed using the thread block cluster interconnect, separate from global memory traffic. For example, when a thread block pulls a tile using DSMEM, it uses a low-latency shared-memory transfer. This ensures that inter-thread-block data sharing within a thread block cluster is as fast as on-chip shared memory.

如果块不存在于DSMEM中,它必须从全局内存获取(遵循可能命中L2的正常缓存层次结构)。这样,如果块已经从DSMEM中驱逐,例如,线程块仍然可以从L2缓存检索块数据,而不是一直返回到全局DRAM。

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). This way, if a tile has already been evicted from DSMEM, for example, the thread block can still retrieve tile data from the L2 cache rather than going all the way back to global DRAM.

因此,内存停顿很少,占用率保持高水平,整体执行比未优化、两个独立线程块、非DSMEM实现快得多。

As a result, memory stalls are rare, occupancy stays high, and overall execution is much faster than an unoptimized, two-independent-thread-block, non-DSMEM implementation.

使用线程块集群设计高效算法 (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.

在过去,您可能将工作分割到两个线程块中,但然后每个块必须与全局内存交换部分结果,这相对较慢且低效。否则,您必须启动单独的归约内核来组合两个线程块的输出。

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.

使用线程块集群和DSMEM,这两个块可以形成一个集群,并直接共享片上共享内存的联合区域,使用硬件支持的原语无缝组合它们的结果。DSMEM硬件允许SM通过快速网络对另一个SM的共享内存执行加载/存储/原子操作。

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.

仔细设计算法以有效使用线程块集群很重要。同步开销低,但仍非零。执行非常细粒度的数据共享可能不值得。

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.

线程块集群屏障的工作方式类似于第7章介绍的线程束级内建函数(例如__shfl_sync()),即每个参与者必须一起到达同步点。当您调用cluster.sync()时,集群中的所有线程块必须到达该行,然后任何块才能继续。

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.

如果一个块提前完成工作,它只是等待。如果一个块从未到达屏障,因为其线程采取了不同的分支,例如,整个集群将死锁。

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.

换句话说,正如线程束内建函数要求线程束中的所有线程执行相同的指令路径以避免分歧,线程块集群要求所有块遵循相同的控制流直到每个cluster.sync()调用。

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 inexpensive when done correctly, but if even one block bypasses or delays reaching a cluster.sync(), performance may be impacted-or, even worse, the kernel might hang.

您应该构建代码,使每个块一致地到达每个屏障--就像线程束中的每个线程必须使用线程束级内建函数到达屏障一样。这对于充分利用线程块集群的低延迟、片上通信而不陷入死锁至关重要。

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.

线程块集群对于块稀疏矩阵操作特别有效--这在LLM的稀疏注意力、模型剪枝和压缩中很常见。在这些情况下,块处理不同的非零区域并共享其边界数据。

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.

线程块集群对于多阶段归约也很有用,例如transformer层中的softmax和归一化步骤。这些需要使用线程块集群中每个线程块的部分结果进行最终组合步骤。

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.

更一般地说,当大型GEMM超过单个块的资源时,它们受益于线程块集群。当然,GEMM是现代LLM中常见的transformer注意力和多层感知器(MLP)层以及嵌入查找的核心。

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.

然而,较大的集群大小可能会降低总体占用率。例如,16个块的集群可能会为一个任务独占16个SM。这可能会为该内核启动所需的其他任务留下更少的SM。建议从小型集群开始,两个或四个线程块集群,除非特殊情况需要更大的集群。一如既往,您应该使用特定工作负载进行分析,以确认跨线程块共享片上资源的好处超过潜在的并行性损失。

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.

在使用集群启动时,请使用Nsight Compute启动统计信息和标准占用率API(如cudaOccupancyMaxActiveBlocksPerMultiprocessor)验证活跃块和集群驻留情况。对于不可移植的集群大小,请设置cudaFuncAttributeNonPortableClusterSizeAllowed(或使用cudaLaunchKernelEx属性传递cudaLaunchAttributeNonPortableClusterSizeAllowed)。否则启动可能会失败或测量到低占用率。

When using cluster launches, verify active blocks and cluster residency with Nsight Compute launch statistics and the standard occupancy APIs, such as cudaOccupancyMaxActiveBlocksPerMultiprocessor. For nonportable cluster sizes, set cudaFuncAttributeNonPortableClusterSizeAllowed (or pass cudaLaunchAttributeNonPortableClusterSizeAllowed using cudaLaunchKernelEx attributes). Otherwise the launch may fail or measure low occupancy.

使用线程块集群的线程束专用化 (Warp Specialization with Thread Block Clusters)

现在让我们重新审视使用带有CUDA Pipeline API的线程块集群进行线程束专用化。我们使用块范围的流水线作为集群领导者来暂存副本并使用DSMEM执行集群范围屏障。这样,集群中的每个块消费相同的输入块而不从全局内存重新加载它们。

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.

角色在每个块内按线程束分配。领导者块的加载器线程束对每个块执行一次协作复制到其共享内存中。在集群范围屏障发布这些块后,每个块使用计算器线程束通过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和块范围流水线跨线程块集群的线程束专用化

#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_SIZExTILE_SIZE乘积的一个行带。
// 每个通道以32路条带化循环处理行[row_begin, row_end)。
__device__ void compute_rows_from_ds(const float* __restrict__ A_src,
                                     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;

    // 块范围流水线仅由集群领导者用于暂存异步复制
    __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的1D集群排列;
    // 每次迭代每个集群处理一个块
    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;

        // 领导者块的加载器线程束为整个集群暂存A和B一次
        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();
            // 在集群范围发布之前使加载对领导者CTA可见
            // 在发布之前等待已提交的阶段
            pipe.consumer_wait();
            pipe.consumer_release();
        }

        // 通过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);

        // 在集群中的块之间分配行
        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();
        }

        // 所有块在领导者重用其缓冲区之前完成此块
        cluster.sync();
    }
    // 动态共享内存大小:3 * TILE_ELEMS * sizeof(float)
}

在这里,领导者使用块范围流水线对每个块执行一对协作复制到其共享内存中。集群范围屏障通过DSMEM使数据对所有集群成员可见。

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共享"模式在集群屏障后使用DSMEM和map_shared_rank共享领导者的块。它不执行我们之前用张量映射描述符、cp.async.bulk.tensorshared::cluster介绍的TMA多播。

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.

这是一个需要理解的重要区别,因为没有DSMEM模式中的cluster.sync(),跟随者可能会读取陈旧的领导者SMEM数据。因此,在map_shared_rank()之前需要cluster.sync()

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().

以下内容应该帮助您选择TMA多播还是DSMEM共享:如果每块块重用很高(例如,每个块多次触摸同一块)或集群大小(C)很大(8-16个CTA),TMA多播通常获胜,因为此模式写入一次,然后本地读取多次。如果SMEM紧张(例如,大块)或集群大小(C)很小(2-4个CTA)且每个跟随者触摸块一次,DSMEM共享通常是更好的选择,因为它使用更小的占用空间。

The following should help you choose TMA multicast versus DSMEM sharing: If tile reuse per block is high (e.g., each block touches the same tile many times) or the cluster size (C) is large (8-16 CTAs), TMA multicast usually wins since this pattern writes once, then reads locally many times. If SMEM is tight (e.g., large tiles) or cluster size (C) is small (2-4 CTAs) and each follower touches the tile once, DSMEM sharing is usually the better option since it uses a smaller footprint.

每个块通过使用map_shared_rank直接从领导者的共享内存计算不同的行带,并将其结果写入全局内存。这消除了集群中的重复全局加载,并在重要的点上保持流水线的重叠优势。

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.

领导者块的加载器线程束调用producer_commit()将其本地阶段发布到该块。其他块不等待领导者的流水线。相反,它们在集群范围屏障后观察领导者的块,然后从分布式共享内存读取。

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.

这允许块加载、计算和存储跨多个线程块交错,这将SM驱动到更高的利用率。表10-5比较了单块线程束专用内核与本节介绍的线程块集群(多块)线程束专用内核。

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.

表10-5. 朴素分块、两阶段双缓冲、线程束专用和线程块集群流水线内核的比较

指标 朴素分块 两阶段双缓冲 线程束专用 线程块集群流水线
内核执行时间 41.3 ms 20.5 ms(比朴素快2.01倍) 18.4 ms(比两阶段快10.2%) 17.2 ms(比线程束专用快6.5%)
线程束执行效率 68% 92%(比朴素+24%) 96%(比两阶段+4%) 97%(比线程束专用+1%)
线程束状态停顿%(共享内存和屏障等待) 最小 最小(进一步减少)
L2吞吐量 80 GB/s 155 GB/s(比朴素+94%) 165 GB/s(比两阶段+6.45%) 170 GB/s(比线程束专用+3%)
吞吐量可扩展性 每个SM最多扩展到2-3个线程束 每个SM扩展到约6个线程束 几乎线性扩展到SM线程束限制(Blackwell上每个SM 64个驻留线程束) 完全跨线程块扩展,直到受SM线程束限制(Blackwell上每个SM 64个驻留线程束)
DRAM读取吞吐量与内核持续时间 重叠差 重叠好 重叠优秀 重叠优秀(即使在线程块传递下)
指令计数 1.7 B 1.05 B(比朴素-38%) ~1.00 B(比两阶段-4.76%) ~0.98 B(比线程束专用-2%)

Table 10-5. Comparison of naive tiling, two-stage double buffering, warp-specialized, and thread block cluster pipeline kernels.

线程块集群流水线内核进一步减少了17.2 ms的执行时间,比单块线程束专用内核提高了6.5%。这种改进来自跨多个SM共享数据而没有冗余的全局内存加载。线程束执行效率从96%上升到98%,因为更多的线程束可以并行工作,每个线程束等待的时间更少。

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. Specifically, the thread block cluster pipeline kernel implementation (warp_specialized_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.

在手动调优之前验证编译器生成的流水线

Verify compiler-generated pipelines before hand-tuning

使用性能分析工具检查框架生成的代码,例如PyTorch的编译器和NVCC。如果编译后的代码使用了使用cuda::memcpy_asyncproducer_commit()consumer_wait()的异步流水线,手动调优可能不会产生太大的加速效果。

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.

关键要点 (Key Takeaways)

本章的关键要点如下,专注于在现代GPU上提取峰值性能。这些将使内核不会在DRAM、线程束调度器和全局屏障上空闲:

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:

使用流水线深度隐藏延迟

使用两阶段(cuda::pipeline_shared_state<cuda::thread_scope_block, 2>)分块来重叠异步加载和计算,当计算超过内存时(例如,像Blackwell这样的现代GPU)添加另一个阶段(cuda::pipeline_shared_state<cuda::thread_scope_block, 3>)。这将有助于消除空闲线程束。

Hide latency with pipeline depth - Use two-stage (cuda::pipeline_shared_state) tiling to overlap asynchronous loads with compute and add another stage (cuda::pipeline_shared_state) when compute outweighs memory (e.g., modern GPUs like Blackwell.) This will help to eliminate idle warps.

使用线程束专用化平衡工作负载

当计算阶段占主导时,将单独的线程束分配给加载、计算和存储,确保现代GPU硬件上接近峰值的线程束效率。

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.

使用持久内核消除启动开销

在设备端工作队列上运行单个长期存在的内核,并使用grid.sync()进行多阶段算法。这将减少主机-设备往返和总体启动成本。

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.

使用线程块集群和DSMEM实现片上共享

将线程块(CTA)分组到集群中,使它们共享连续的片上缓冲区。对于集群范围广播,使用带有TMA多播描述符的cp.async.bulk.tensor。使用TMA多播将块一次广播到每个线程块。这将提高L2命中率并减少DRAM带宽。

Enable on-chip sharing with thread block clusters and 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.

特别注意屏障语义

cluster.sync()grid.sync()都要求每个参与的线程块到达同一同步点。不匹配的控制流--或过大的集群大小--可能导致死锁或启动失败。

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.

调优前进行性能分析

使用Nsight Compute识别您的内核是内存受限还是计算受限。如果是内存受限,从两阶段流水线开始。如果是计算受限,考虑线程束专用化或线程块集群。

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.

在手动调优之前验证编译器生成的流水线

分析并检查来自PyTorch编译器和NVCC等编译器的框架生成代码。如果编译后的代码使用cuda::memcpy_asyncproducer_commit()consumer_wait()的异步流水线,手动调优可能不会产生太大的加速。

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.

结论 (Conclusion)

持久内核、线程束专用化、协作组和线程块集群是现代GPU编程的高级技术。它们使复杂的内核内流水线、减少的全局内存流量和更高的SM利用率成为可能。

Persistent kernels, warp specialization, cooperative groups, and thread block clusters are advanced techniques in modern GPU programming. They enable complex intra-kernel pipelines, reduced global memory traffic, and higher SM utilization.

这些技术是高性能库和框架的基础,如CUTLASS、FlashAttention和TorchInductor。通过掌握这些技术,您可以实现接近硬件峰值的性能。

These techniques will move performance closer to the GPU's peak theoretical limits.

在下一章中,我们将介绍CUDA流和CUDA Graphs,它们提供了额外的并发和优化机会。这些技术允许您重叠内核执行与数据传输,并高效启动复杂的内核工作流。

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.