Skip to content
Markdown

Warp specialization and intra-kernel pipelining

Scope: splitting a thread block into producer (load) and consumer (compute) warps and software-pipelining the stages with the CUDA Pipeline API and async copies plus double buffering, to hide memory latency inside a single kernel (the FlashAttention/GEMM pattern).

What it is

A GPU SM hides latency by switching between resident warps. When a kernel is memory-bound and occupancy is already capped (by registers, shared memory, or block count), there are no more warps to hide a load's latency behind, and the math units stall waiting on HBM. Warp specialization plus intra-kernel pipelining is the response: instead of every warp doing the same load-then-compute, the block is partitioned by role, and the loads for iteration i+1 are issued (asynchronously) while iteration i computes.

Two mechanisms compose here:

  • Warp specialization: partition a thread block's warps into a producer role (warps dedicated to moving tiles from global memory into shared memory) and a consumer role (warps dedicated to the math). The roles run concurrently and communicate through shared-memory buffers. This is the structure FlashAttention-3 uses on Hopper: producer warps issue TMA loads of K/V tiles into shared memory while consumer warps run the WGMMA matmuls, with ping-pong scheduling overlapping the softmax of one block against the GEMM of the next (Shah et al., "FlashAttention-3", 2024).
  • Intra-kernel software pipelining (double / multi-buffering): allocate N shared-memory buffers (stages) and run them as a FIFO. While consumers process the buffer for stage k, the producer fills the buffer for stage k+1 using asynchronous copies that do not block the issuing warp. N=2 is double buffering; deeper pipelines (N=3,4) prefetch further ahead to absorb longer latency. NVIDIA's cuda::pipeline library provides the synchronization primitives for this (CUDA Programming Guide, "Pipelines").

The hardware enabler is the asynchronous global-to-shared copy (cp.async / LDGSTS), introduced with compute capability 8.0 (Ampere). It streams data from global memory directly into shared memory without staging through registers or the L1 data path, freeing the warp to keep issuing other instructions while the copy is in flight (CUDA Programming Guide, "Asynchronous Data Copies").

Why use it

A classic tiled GEMM (see Shared Memory, Bank Conflicts, and Tiling) loads a tile into shared memory, calls __syncthreads(), computes, syncs again, then loads the next tile. The load and the compute are serialized: every iteration pays the full HBM latency before any math starts. On a deep K loop, that latency is exposed on every step.

Pipelining hides it. With double buffering, the producer's async load for the next tile is issued before the consumer touches the current tile, so by the time the consumer finishes, the next tile is already (or nearly) resident. The exposed cost per iteration collapses from latency + compute toward max(latency, compute). When compute dominates, the loads become free; when loads dominate, deeper pipelines and more producer bandwidth help.

This is the core latency-hiding lever once occupancy can no longer be raised. It is also what lets attention and GEMM kernels stay near the Tensor Core roofline (see Roofline Model and Arithmetic Intensity, Tensor Cores and Mixed Precision): the asynchronous TMA + WGMMA pipeline on Hopper is precisely why FlashAttention-3 reaches a reported 1.5-2.0x speedup over FlashAttention-2 on H100 (Shah et al., "FlashAttention-3", 2024).

The book frames this as the natural progression past tiling: once the data is in shared memory cheaply, the remaining stall is the load latency itself, and specialization plus pipelining is how you overlap it away (Fregly, AI Systems Performance Engineering, O'Reilly, Ch. 9). Treat speedup figures as workload- and shape-dependent; confirm yours with Nsight Compute.

That per-iteration claim (from latency + compute toward max(latency, compute)) is a first-order scheduling model, and it is worth making concrete and checkable. The block below simulates the bounded-buffer producer/consumer pipeline with a max-plus recurrence and asserts the collapse against an independently derived closed form, plus the serialized baseline, monotonicity in pipeline depth, and boundary cases:

import numpy as np

def pipeline_cost(num_iters, latency, compute, num_stages):
    # Max-plus recurrence for a bounded-buffer producer/consumer pipeline.
    # Producer fills one shared-memory stage per tile (cost `latency`);
    # consumer drains it (cost `compute`); `num_stages` = in-flight buffers.
    # A stage frees only when the consumer finishes the tile num_stages back,
    # which is exactly the shared-memory back-pressure of double/multi-buffering.
    assert num_iters >= 1 and num_stages >= 1
    assert latency >= 0 and compute >= 0
    prod = np.zeros(num_iters, dtype=np.int64)  # producer finish time per tile
    cons = np.zeros(num_iters, dtype=np.int64)  # consumer finish time per tile
    for i in range(num_iters):
        prev_prod = prod[i - 1] if i >= 1 else 0
        slot_free = cons[i - num_stages] if i >= num_stages else 0
        prod[i] = max(prev_prod, slot_free) + latency
        prev_cons = cons[i - 1] if i >= 1 else 0
        cons[i] = max(prev_cons, prod[i]) + compute
    return prod, cons

def total_time(num_iters, latency, compute, num_stages):
    return int(pipeline_cost(num_iters, latency, compute, num_stages)[1][-1])

rng = np.random.default_rng(0)

# 1) Baseline: with a single buffer the load and compute serialize, so the
#    exposed cost per tile is exactly latency + compute (the un-pipelined loop).
for _ in range(2000):
    n = int(rng.integers(1, 64))
    lat, cmp = int(rng.integers(0, 16)), int(rng.integers(0, 16))
    assert total_time(n, lat, cmp, 1) == n * (lat + cmp)

# 2) Double buffering (num_stages >= 2) fully overlaps a two-stage pipeline:
#    total collapses to n*max(latency,compute) + min(latency,compute), i.e. the
#    per-tile cost falls from latency+compute toward max(latency,compute). This
#    is asserted against an independently derived closed form, not the sim.
for _ in range(2000):
    n = int(rng.integers(1, 64))
    lat, cmp = int(rng.integers(0, 16)), int(rng.integers(0, 16))
    for stages in (2, 3, 4):
        closed = n * max(lat, cmp) + min(lat, cmp)
        assert total_time(n, lat, cmp, stages) == closed

# 3) Steady-state per-tile cost is exactly max(latency, compute).
_, cons = pipeline_cost(200, latency=7, compute=3, num_stages=4)
steady = np.diff(cons[50:])
assert np.all(steady == max(7, 3))               # producer-bound tail -> latency
_, cons2 = pipeline_cost(200, latency=3, compute=7, num_stages=4)
assert np.all(np.diff(cons2[50:]) == max(3, 7))  # compute-bound tail -> compute

# 4) Monotonic: deeper pipelines are never slower than shallower ones.
for _ in range(2000):
    n = int(rng.integers(1, 64))
    lat, cmp = int(rng.integers(0, 16)), int(rng.integers(0, 16))
    assert total_time(n, lat, cmp, 3) <= total_time(n, lat, cmp, 2) <= total_time(n, lat, cmp, 1)

# 5) Adversarial boundaries.
#    a single tile cannot overlap with anything: no speedup regardless of depth.
assert total_time(1, 5, 4, 4) == total_time(1, 5, 4, 1) == 5 + 4
#    a free load (latency 0) leaves only compute; asymptotic speedup is bounded
#    by (latency+compute)/max(latency,compute) <= 2 and hits 2x when balanced.
assert total_time(1000, 0, 4, 2) == 1000 * 4 + 0
bal = (5 + 5) / max(5, 5)
assert abs(bal - 2.0) < 1e-12
speedup = total_time(1000, 5, 5, 1) / total_time(1000, 5, 5, 2)
assert 1.9 < speedup <= 2.0

print("pipeline cost model: serialized=lat+cmp, double-buffered->max(lat,cmp),"
      " steady-state + monotonic + boundaries OK")

When to use it (and when not)

Reach for it when:

  • You are hand-writing a CUDA or Triton kernel (GEMM, attention, convolution, long reduction) that Nsight Compute shows memory-latency-bound: high "Long Scoreboard" / "LG Throttle" stall reasons, low achieved occupancy that you cannot raise, and a Warp State view dominated by waiting on global loads (see Profiling GPUs: Nsight Systems and Nsight Compute).
  • The kernel has a deep, regular inner loop over tiles where the next iteration's input addresses are known ahead of time (so prefetch is possible).
  • You target Ampere or newer (sm_80+), where cp.async exists; on Hopper (sm_90) TMA and WGMMA make the producer/consumer split substantially more effective.

Do not reach for it when:

  • A library already covers your shape and dtype. cuBLAS, cuDNN, CUTLASS, FlashAttention, and the Triton-backed TorchInductor kernels already implement multi-stage pipelined, warp-specialized inner loops. Reimplementing them by hand rarely wins and is easy to get wrong.
  • The kernel is compute-bound at the math roofline: there is no exposed load latency left to hide, so pipelining adds shared-memory pressure for no gain.
  • Occupancy is the real limiter and can still be raised by cutting registers or shared memory. Multi-stage buffering increases shared-memory use (N copies of each tile), which can lower occupancy; the win must outweigh that.
  • You are pre-Ampere (no cp.async); the manual copy-through-registers fallback gives much less overlap.

Architecture

The kernel splits one thread block into two roles that run concurrently and hand tiles across a ring of shared-memory buffers. The producer streams the next tiles from HBM with asynchronous copies while the consumer does the math on tiles already resident, so the load of stage k+1 overlaps the compute of stage k.

flowchart LR
    HBM["Global memory (HBM)"] -->|"cp.async / TMA"| PROD["Producer warps"]
    PROD -->|"producer_acquire / commit"| BUFS["Shared-memory stages 0..N-1"]
    BUFS -->|"consumer_wait / release"| CONS["Consumer warps (FMA / WGMMA)"]
    CONS -->|"results"| OUT["Global memory C"]
    PROD -.->|"prefetch stage k+1 while consumers run stage k"| CONS

The structural pieces the diagram rests on:

  • Producer warps own the load side. On Ampere they issue cp.async; on Hopper they drive the TMA engine to move whole K/V tiles into shared memory. Because the copy is asynchronous, a producer warp keeps issuing the next tile's copy without blocking on HBM latency.
  • Consumer warps own the math (FMA on CUDA cores, or WGMMA on Hopper Tensor Cores) over tiles already resident in shared memory. They never touch HBM on the critical path.
  • The N-stage shared-memory ring is the bounded buffer between the roles. It decouples them: the producer may run up to N tiles ahead, and a stage frees only when the consumer releases it. That back-pressure is why N also sets the shared-memory footprint (N copies of each tile) and therefore the occupancy cost.
  • The pipeline barrier (cuda::pipeline with producer_acquire/producer_commit and consumer_wait/consumer_release) is the handshake that makes stage k visible to the consumer only after its load lands, and frees the slot only after the consumer is done.
  • The async copy engine (cp.async/LDGSTS on sm_80+, TMA on sm_90) is the hardware that lets a single warp launch a global-to-shared transfer and keep issuing, which is what makes the overlap real rather than nominal.

How to use it

Producer/consumer with the CUDA Pipeline API

The cuda::pipeline primitive coordinates producer and consumer threads through an N-stage FIFO. Producers call producer_acquire() / producer_commit() around their async copies; consumers call consumer_wait() / consumer_release() around the math. Shared state lives in a cuda::pipeline_shared_state<scope, N>. All names below are from the official API (CUDA Programming Guide, "Pipelines"; libcu++ pipeline_shared_state).

The example below is a producer/consumer GEMM skeleton built on the CUDA Pipeline API and asynchronous memory copies. Warp 0 is the producer; the remaining warps are consumers. It is a CUDA C++ reference template (it requires nvcc and an NVIDIA GPU and has not been hardware-tested in this knowledge base); its scheduling core math, the load/compute overlap, is the cost model validated in numpy under "Why use it". Compile and profile before trusting any number.

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

namespace cg = cooperative_groups;

template <int TILE_M, int TILE_N, int TILE_K, int NUM_STAGES>
__global__ void warp_specialized_gemm(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int M, int N, int K)
{
    // One shared-memory buffer per pipeline stage (multi-buffering).
    __shared__ float smem_A[NUM_STAGES][TILE_M][TILE_K];
    __shared__ float smem_B[NUM_STAGES][TILE_K][TILE_N];

    // Block-scoped pipeline state sized for NUM_STAGES in flight.
    __shared__ cuda::pipeline_shared_state<
        cuda::thread_scope_block, NUM_STAGES> pipe_state;

    auto block = cg::this_thread_block();
    auto pipe  = cuda::make_pipeline(block, &pipe_state);

    const int warp_id    = threadIdx.x / 32;
    const bool is_producer = (warp_id == 0);
    const int num_k_tiles  = (K + TILE_K - 1) / TILE_K;

    if (is_producer) {
        // Producer: stage async loads ahead of the consumers.
        for (int stage = 0; stage < min(NUM_STAGES, num_k_tiles); ++stage) {
            pipe.producer_acquire();
            // Issue cuda::memcpy_async copies of the A/B tiles for this
            // stage into smem_A[stage] / smem_B[stage] here.
            pipe.producer_commit();   // advance the pipeline head
        }
    } else {
        // Consumer: compute on each stage as it becomes ready.
        for (int kt = 0; kt < num_k_tiles; ++kt) {
            pipe.consumer_wait();     // block until this stage's loads land
            const int stage = kt % NUM_STAGES;
            // Fused multiply-accumulate over smem_A[stage] x smem_B[stage].
            pipe.consumer_release();  // free the buffer for the producer to refill
        }
    }
}

Launch with at least two warps so the producer/consumer split exists; the repo's harness uses a 128-thread block (4 warps: 1 producer, 3 consumers) with NUM_STAGES = 4:

dim3 block(128);                              // 4 warps
dim3 grid((N + TILE_N - 1) / TILE_N,
          (M + TILE_M - 1) / TILE_M);
warp_specialized_gemm<32, 32, 32, 4><<<grid, block>>>(d_A, d_B, d_C, M, N, K);

The asynchronous copy itself

The producer's load uses an async copy so the warp does not block on HBM latency. The collective form is cooperative_groups::memcpy_async, paired with cg::wait; the library notes it is less efficient than the pipeline form because it commits each copy immediately (CUDA Programming Guide, "Asynchronous Data Copies").

#include <cooperative_groups.h>
#include <cooperative_groups/memcpy_async.h>

namespace cg = cooperative_groups;

cg::thread_block block = cg::this_thread_block();
__shared__ float smem[TILE_SIZE];

// Async global -> shared; the warp keeps issuing while the copy is in flight.
cg::memcpy_async(block, &smem[i], &src[global_idx], sizeof(float));
cg::wait(block);   // barrier on outstanding copies before reading smem

Under the hood this lowers to the cp.async / LDGSTS instruction available on compute capability 8.0+ (Ampere and newer). LDGSTS supports 4-, 8-, or 16-byte transfers; the 16-byte path uses L1 BYPASS mode (data is not cached in L1, avoiding pollution), while 4/8-byte transfers use L1 ACCESS mode (CUDA Programming Guide, "Asynchronous Data Copies"). For the lowest-level control, __pipeline_memcpy_async() / __pipeline_commit() / __pipeline_wait_prior() from <cuda_pipeline.h> map almost directly onto the PTX.

How to integrate it

Higher-level tools express this pattern declaratively, so you rarely hand-roll the CUDA form. Triton is the usual entry point: it software-pipelines a loop for you and, on recent versions and hardware, can specialize warps automatically.

The same pattern in Triton

Triton expresses pipelining declaratively on a loop. tl.range(..., num_stages=K) software-pipelines the loop body into K in-flight iterations; this is broader than the num_stages kernel argument, which only pipelines loads feeding dot (Triton docs, triton.language.range).

This is a Triton reference template (it requires Triton and a GPU and is not executed in this page); its elementwise core math, a masked, tiled relu((q*k)*scale) * v, is validated in numpy directly beneath it.

import triton
import triton.language as tl

@triton.jit
def staged_kernel(q_ptr, k_ptr, v_ptr, out_ptr, n_elements,
                  BLOCK_SIZE: tl.constexpr, TILE_SIZE: tl.constexpr,
                  PIPELINE_STAGES: tl.constexpr, SCALE: tl.constexpr):
    pid = tl.program_id(axis=0)
    tile_start = pid * TILE_SIZE
    # Software-pipeline the tile loop into PIPELINE_STAGES in-flight iterations.
    for chunk in tl.range(0, TILE_SIZE, BLOCK_SIZE, num_stages=PIPELINE_STAGES):
        offs = tile_start + chunk + tl.arange(0, BLOCK_SIZE)
        mask = offs < n_elements
        q = tl.load(q_ptr + offs, mask=mask, other=0.0)
        k = tl.load(k_ptr + offs, mask=mask, other=0.0)
        v = tl.load(v_ptr + offs, mask=mask, other=0.0)
        s = tl.maximum((q * k) * SCALE, 0.0)
        tl.store(out_ptr + offs, s * v, mask=mask)

The numpy block below mirrors that kernel's core math: the tiled loop, the mask / other=0.0 boundary handling, and the tl.maximum(..., 0.0) relu, cross-checked against an untiled vectorized reference and exercising the ragged-boundary and negative-input paths:

import numpy as np

def staged_scores(q, k, v, scale, block_size):
    # numpy mirror of the Triton staged_kernel core math: a masked, tiled loop
    # computing s = relu((q*k)*scale) elementwise, then out = s*v. The tiling
    # and mask (other=0.0 for out-of-range lanes) mirror tl.load(..., mask).
    n = q.shape[0]
    assert q.shape == k.shape == v.shape == (n,)
    out = np.zeros(n, dtype=np.float64)
    for start in range(0, n, block_size):
        offs = start + np.arange(block_size)
        mask = offs < n                       # boundary lanes are masked off
        idx = np.where(mask, offs, 0)
        qb = np.where(mask, q[idx], 0.0)      # tl.load(..., other=0.0)
        kb = np.where(mask, k[idx], 0.0)
        vb = np.where(mask, v[idx], 0.0)
        s = np.maximum((qb * kb) * scale, 0.0)
        contrib = s * vb
        out[idx[mask]] = contrib[mask]        # tl.store(..., mask=mask)
    return out

def reference(q, k, v, scale):
    # Independent, untiled, unmasked vectorized reference.
    return np.maximum((q * k) * scale, 0.0) * v

rng = np.random.default_rng(1)
scale = 0.7

# 1) Equivalence to the vectorized reference when n is a multiple of the block.
q = rng.standard_normal(256); k = rng.standard_normal(256); v = rng.standard_normal(256)
assert np.allclose(staged_scores(q, k, v, scale, 64), reference(q, k, v, scale))

# 2) Adversarial boundary: n NOT a multiple of the block. Masked lanes must not
#    read past the end or write garbage; the guarded result stays exact.
q2 = rng.standard_normal(300); k2 = rng.standard_normal(300); v2 = rng.standard_normal(300)
got = staged_scores(q2, k2, v2, scale, 64)             # 300 = 4*64 + 44
assert got.shape == (300,)
assert np.allclose(got, reference(q2, k2, v2, scale))

# 3) The relu is load-bearing: wherever q*k < 0 the output is exactly zero,
#    independent of v (a plain q*k*v kernel would fail this).
neg = np.array([-1.0, 2.0, -3.0, 4.0]); kk = np.array([1.0, 1.0, 1.0, 1.0])
vv = np.array([9.0, 9.0, 9.0, 9.0])
r = staged_scores(neg, kk, vv, 1.0, 2)
assert np.array_equal(r == 0.0, (neg * kk) <= 0.0)
assert r[0] == 0.0 and r[2] == 0.0 and r[1] == 2.0 * 9.0 and r[3] == 4.0 * 9.0

# 4) Block size must never change the result (tiling invariance).
for b in (1, 2, 3, 7, 64, 512):
    assert np.allclose(staged_scores(q, k, v, scale, b), reference(q, k, v, scale))

print("staged kernel math: reference-equiv + boundary(300) + relu-zeroing"
      " + block-invariance OK")

(See the Triton docs for the automated warp_specialize path.)

For true warp specialization, Triton's automated path is version- and hardware-sensitive: verify against the Triton version you ship.

  • The warp_specialize=True flag on tl.range: per current Triton docs it "is only supported on Blackwell GPUs and only works on simple matmul loops" (Triton docs, triton.language.range). The book repo's warp_specialize=True example reflects an earlier/experimental API surface; do not assume it compiles or specializes on your target without checking.
  • The autotuner-config path uses num_consumer_groups (non-zero enables specialization; only one producer group is supported) and num_buffers_warp_spec (producer->consumer buffer count). This automated specialization landed in Triton 3.2 (PyTorch 2.6) targeting Hopper H100, with a reported 10-15% improvement on FlashAttention and FP8 row-wise GEMM (PyTorch blog, "Enabling advanced GPU features: Warp Specialization").

The triton.Config below is a reference template; its num_stages and num_buffers_warp_spec are the pipeline depth N from the cost model in "Why use it", and its num_consumer_groups is the producer/consumer split from the Architecture section.

triton.Config(
    {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8},
    num_stages=2, num_warps=4,
    num_consumer_groups=2, num_buffers_warp_spec=3,
)

How to run it in production

Prefer the shipped kernels. cuBLAS, cuDNN, CUTLASS, FlashAttention, and the Triton-backed TorchInductor kernels already implement multi-stage, warp-specialized inner loops tuned across GPU generations, so a torch.matmul, an attention call, or a torch.compiled region gets the pipeline for free. This is why FlashAttention-3's TMA + WGMMA pipeline reaches a reported 1.5-2.0x over FlashAttention-2 on H100 (Shah et al., "FlashAttention-3", 2024) with no hand-written pipeline on your side. Hand-roll the producer/consumer form only when no library covers your shape, dtype, or fused pattern, and when you do, gate it behind the correctness and profiling checks below on the exact GPU you will ship on: the skeletons here are reference templates and have not been hardware-tested in this knowledge base. Treat any reported speedup (the 1.5-2.0x here, the Triton 10-15% above) as workload- and shape-dependent until Nsight Compute confirms it on your kernel.

How to maintain it

A hand-rolled pipeline is only correct once profiled, and it must be re-checked after every change and on every new GPU generation. Correctness comes first, because a pipelining bug is silent: a missing consumer_wait() / producer_commit() pairing is a read-before-write race that may pass intermittently. Run under CUDA Compute Sanitizer: Correctness Debugging (racecheck, synccheck) before trusting any result.

Then profile with Nsight Compute (see Profiling GPUs: Nsight Systems and Nsight Compute):

# Warp-state (stall reasons) + occupancy sections for the pipelined kernel
ncu --set full \
    --section WarpStateStats \
    --section Occupancy \
    ./warp_specialized_gemm

Success looks like falling memory-latency stalls ("Long Scoreboard", "LG Throttle") and rising issue-slot utilization; the Warp State view should no longer be dominated by waiting on global loads, and the arithmetic-intensity point should have moved toward the compute roofline. Watch occupancy at the same time: each added stage multiplies the shared-memory footprint, so check that launch__occupancy_limit_shared_mem did not crater achieved occupancy (see CUDA Occupancy Tuning).

How to scale it

The two knobs are pipeline depth and the producer/consumer split.

  • Sweep NUM_STAGES (2, 3, 4) and keep the best. N=2 (double buffering) already reaches full two-stage overlap in the cost model above; deeper pipelines prefetch further ahead to absorb longer latency, at the cost of N copies of each tile in shared memory. The gain must outweigh the occupancy hit from that footprint, which is exactly the tradeoff the cost model and the launch__occupancy_limit_shared_mem check quantify.
  • Balance producer against consumer warps to the load/compute ratio. One producer warp starving many consumers (or many producers feeding one consumer) leaves the pipeline unbalanced; the repo harness runs 1 producer to 3 consumers, but the right ratio tracks max(latency, compute) from the cost model.
  • On Hopper, let the hardware specialize further. TMA producers plus WGMMA consumers, exposed through Triton's num_consumer_groups and num_buffers_warp_spec (Triton 3.2 / PyTorch 2.6), extend the same producer/consumer decomposition to warp groups, and are where the reported 10-15% FlashAttention and FP8 row-wise GEMM gains come from (PyTorch blog, "Enabling advanced GPU features: Warp Specialization").

Failure modes

  • Read-before-write race from a dropped handshake. A missing consumer_wait() / producer_commit() pairing lets a consumer read a stage before its load lands. It yields wrong, nondeterministic results, not a crash, and may pass intermittently. Catch it with Compute Sanitizer racecheck / synccheck, not by eyeballing timing.
  • Occupancy collapse from too many stages. Each added stage multiplies the shared-memory footprint (N copies of each tile), which can lower achieved occupancy below the point where the overlap pays for itself. Check launch__occupancy_limit_shared_mem when sweeping NUM_STAGES.
  • Pipelining a compute-bound kernel. If the kernel already sits at the math roofline there is no exposed load latency to hide, so multi-stage buffering only adds shared-memory pressure for no gain. Confirm the kernel is memory-latency-bound in Nsight first.
  • Pre-Ampere target with no cp.async. Below compute capability 8.0 there is no asynchronous global-to-shared copy; the manual copy-through-registers fallback gives much less overlap, so the pattern buys far less.
  • Non-portable Triton warp specialization. The automated path is version- and hardware-sensitive: warp_specialize=True on tl.range "is only supported on Blackwell GPUs and only works on simple matmul loops", and the book repo's example reflects an earlier/experimental API surface. Do not assume it compiles or specializes on your target without checking the Triton version you ship.
  • Unbalanced producer/consumer split. One producer warp starving many consumers, or the reverse, leaves the pipeline waiting on the underprovisioned role; match the split to the load/compute ratio.
  • Trusting a speedup figure blind. The 1.5-2.0x (FlashAttention-3) and 10-15% (Triton warp specialization) numbers are workload- and shape-dependent; confirm yours with Nsight Compute before quoting one.

References

Related: Shared Memory, Bank Conflicts, and Tiling · Memory Coalescing and Vectorized Access · Kernel Fusion · Tensor Cores and Mixed Precision · CUDA Streams and Concurrency · FlashAttention and Multi-Head Latent Attention · CUDA Occupancy Tuning · Profiling GPUs: Nsight Systems and Nsight Compute · CUDA Compute Sanitizer: Correctness Debugging · CUTLASS: Templated GEMM and Kernel Building Blocks · OpenAI Triton: Authoring GPU Kernels in Python · Roofline Model and Arithmetic Intensity · Glossary