Warp specialization

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

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 2407.08608 §3.1 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 2407.08608 §2.2 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 2407.08608 alg. 1 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 2407.08608 tab. 2 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 2407.08608 §3.1 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 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 2407.08608 §3.2 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 2407.08608 app. B.3. On the ablation setting, the overlap alone gives 570 TFLOPs/s and the full design gives 661 2407.08608 tab. 2.

When to use which

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.

# 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 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 2407.08608 §3.1 2407.08608 tab. 2.

Papers

• FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision: the warp-specialised Hopper attention kernel, with producer and consumer roles, pingpong scheduling and the two-stage overlap.