---
kind: concept
slug: warp-specialization
title: Warp specialization
summary: >-
  Warp specialization gives each group of GPU threads one fixed job, such as loading data into shared memory or doing matrix multiplies and softmax, so that the jobs overlap. It needs Hopper-class hardware features. The only evidence so far is FlashAttention-3, where removing the role split or the matrix-softmax overlap each cost about 12-14% of throughput in one benchmark setting.
tags: [gpu-kernels, hopper, asynchrony]
papers: ["2407.08608"]
related: [fused-attention-kernel, arithmetic-intensity, model-flops-utilization]
status: draft
updated_at: 2026-10-18
---

## Overview

A kernel that does a matrix multiply, then a softmax, then another matrix
multiply can be slow even when each step is fast on its own, because the
units that do each step wait for one another. Warp specialisation assigns
each group of threads a single role. A producer group only issues the
asynchronous loads that bring the next blocks into shared memory. Consumer
groups only compute [arxiv:2407.08608 §3.1]. Because the loads run on their
own, the matrix units and the softmax can work while the next data is still
arriving [arxiv:2407.08608 §2.2].

The paper calls the Hopper hardware features it relies on the Tensor Cores
and the Tensor Memory Accelerator (TMA), which run apart from the general
purpose cores [arxiv:2407.08608 §2.2]. The role split is the shared idea of
every variant below. The variants differ in how the consumers are scheduled
and whether the matrix multiply and the softmax inside one consumer overlap.
The one card so far is the Hopper attention kernel:
[FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](/papers/2407.08608).

## Variants

### Producer and consumer roles with register reallocation

Threads are divided into one producer warpgroup and one or more consumer
warpgroups. The producer fills a circular buffer in shared memory with s
stages, and it waits for a slot only after the first s blocks have been
loaded [arxiv:2407.08608 §3.1] [arxiv:2407.08608 alg. 1]. Since the producer
does little work, the paper frees registers on that role and gives them to the
consumers with `setmaxnreg` [arxiv:2407.08608 §2.2] [arxiv:2407.08608 §3.1].

Each consumer issues two asynchronous matrix multiplies per key block. One
computes the scores from operands both in shared memory. The other adds the
probability-weighted values, with the probabilities taken from registers
[arxiv:2407.08608 alg. 1] [arxiv:2407.08608 §3.1]. Without the overlap of
§3.2 below, the consumers still wait for each softmax in turn. On the paper's
ablation setting this variant alone gives 582 TFLOPs/s, against 661 for the
full design [arxiv:2407.08608 tab. 2] [arxiv:2407.08608 §4.2].

### Pingpong scheduling between two consumer warpgroups

With two consumer warpgroups, barriers can hold the matrix multiplies of one
warpgroup until the other has finished its softmax, so the two take turns on
the matrix units [arxiv:2407.08608 §3.1] [arxiv:2407.08608 fig. 1]. The paper
reports that this raises FP16 forward throughput from 570 to 620-640 TFLOPs/s
for head dimension 128 and sequence length 8192. It also says the real
schedule is less tidy than the figure suggests [arxiv:2407.08608 §3.1].

### Overlap of matrix multiplies and softmax inside one warpgroup

This is a separate technique that the paper combines with the role split. It
keeps one extra score block in registers, so the softmax for block j runs while
the score multiply for block j+1 is in flight [arxiv:2407.08608 §3.2]
[arxiv:2407.08608 alg. 2]. The paper gives a three-stage version as well. That
version was slower, because the compiler did not overlap the second matrix
multiply and register pressure forced smaller tiles [arxiv:2407.08608 app. B.3].
On the ablation setting, the overlap alone gives 570 TFLOPs/s and the full
design gives 661 [arxiv:2407.08608 tab. 2].

## When to use which

| Situation | Use | Why |
|---|---|---|
| Hopper GPU, matrix multiplies and softmax in one kernel | Producer and consumer roles, plus the two-stage overlap | The paper measures both features, and the combination is the fastest of the three ablation configurations, at 661 TFLOPs/s [arxiv:2407.08608 tab. 2] |
| Only one of the two features is feasible | Either one; the paper does not compare them against each other on other hardware | Each one alone is 570 or 582 TFLOPs/s on the ablation setting, against 661 combined [arxiv:2407.08608 tab. 2] |
| Two consumer warpgroups available | Add pingpong scheduling | Reported gain from 570 to 620-640 TFLOPs/s, for head dimension 128 and sequence length 8192 only [arxiv:2407.08608 §3.1] |
| Registers are scarce or tiles are large | Keep the two-stage overlap small, or drop it | The extra score block costs registers, and the paper picks tile sizes by profiling [arxiv:2407.08608 §3.2] |
| More than two stages are tempting | Measure before adopting | The three-stage version was slower in the paper's tests [arxiv:2407.08608 app. B.3] |
| Non-Hopper GPU | Not stated | The design relies on TMA, warpgroup matrix multiplies and per-role register allocation [arxiv:2407.08608 §2.2]. The paper gives no results for other GPUs |

## Implementation sketch

The producer and consumer roles of one thread block. The pingpong barriers
and the overlap of §3.2 are left out. They change the schedule of the loop,
not its result.

```pseudocode
# Warp-specialised block, producer and consumer roles  [arxiv:2407.08608 alg. 1]
# Input:  query block; keys and values in global memory
# Output: output block; logsumexp per row

def producer(query_block, keys, values):
    # Give registers to the consumers, then load the query block  [§2.2]
    release_registers()
    load_async(query_block)
    for j in range(T_c):
        # Wait until ring slot j % s has been consumed  [alg. 1]
        wait_slot_free(j % s)
        # Stream key and value block j into that slot  [alg. 1]
        load_async(keys[j], values[j], slot=j % s)

def consumer(query_block):
    # Take the registers, then keep running max and sum  [alg. 1]
    acquire_registers()
    output, row_max, row_sum = zeros(B_r, d), full(B_r, -inf), zeros(B_r)
    for j in range(T_c):
        # Scores for block j, from shared memory operands  [alg. 1]
        scores = query_block @ keys[j].T
        # Update the running max, the rescale factor and the sum  [alg. 1]
        new_max = maximum(row_max, rowmax(scores))
        probs = exp(scores - new_max)
        row_sum = exp(row_max - new_max) * row_sum + rowsum(probs)
        # Rescale the old output and add probs times values  [alg. 1]
        output = exp(row_max - new_max)[:, None] * output
        output = output + probs @ values[j]
        row_max = new_max
        # Free the ring slot for the producer  [alg. 1]
        release_slot(j % s)
    return output / row_sum[:, None], row_max + log(row_sum)
```

The paper's listing writes the output rescale with the inverse factor. The
form above is the one that agrees with the normalising sum
[arxiv:2407.08608 alg. 1]. The card's Pitfalls section gives the detail.

## Contradictions

None known. Only one card cites this concept so far. Two numbers of 570
TFLOPs/s appear in the paper, in §3.1 and in tab. 2, but they come from
different configurations and are not comparable [arxiv:2407.08608 §3.1]
[arxiv:2407.08608 tab. 2].

## Papers

- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](/papers/2407.08608): the warp-specialised Hopper attention kernel, with producer and consumer roles, pingpong scheduling and the two-stage overlap.