Fused attention kernel

Overview

Standard attention writes an N x N score matrix and an N x N probability matrix to high-bandwidth GPU memory, so its cost grows quadratically with sequence length 2407.08608 §2.1 2407.08608 §2.3. A fused kernel does the matrix products, the softmax and the output sum in one pass. It loads small blocks of queries, keys and values into on-chip memory, and keeps a running row maximum and row sum so that the softmax can be finished block by block 2407.08608 §3.1. The score and probability matrices never leave the chip 2407.08608 §2.3.

The shared idea behind every variant is this running-softmax loop. The variants differ in three things: how the warps and warpgroups divide the work so that loads, matrix multiplies and softmax overlap; whether the matrix multiplies use 16-bit or 8-bit numbers; and which GPU generation they are written for. The one card so far targets NVIDIA Hopper (H100). Card: FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision.

The earlier kernels in this line are FlashAttention (Dao et al., 2022) and FlashAttention-2. Neither has a card here, so they are named but not described.

Variants

Warp-specialised asynchronous forward pass (16-bit, Hopper)

Warps are split into a producer warpgroup and consumer warpgroups. The producer only issues TMA loads of the query block and of the key and value blocks into a circular buffer in shared memory. The consumers only compute. Registers are reallocated per role 2407.08608 §3.1 2407.08608 alg. 1. Each consumer issues two asynchronous matrix multiplies per key block: scores from operands in shared memory, and the output update with the probabilities taken from registers [arxiv:2407.08608 §3.1] 2407.08608 alg. 1.

The paper adds a pingpong schedule for two consumer warpgroups. Barriers make one warpgroup's matrix multiplies run while the other does its softmax. The paper reports 570 TFLOPs/s without it and 620-640 TFLOPs/s with it, for head dimension 128 and sequence length 8192. It also says the real schedule is less clean than its figure shows 2407.08608 §3.1 2407.08608 fig. 1. Card: FlashAttention-3.

Two-stage overlap of matrix multiplies and softmax

Inside one warpgroup, the chain of score multiply, softmax and output multiply serialises the units. The fix keeps one extra score block in registers, so the softmax of block j runs while the score multiply for block j+1 is in flight 2407.08608 §3.2 2407.08608 alg. 2 2407.08608 fig. 2. The extra block costs B_r x B_c x 4 bytes of registers per thread block, which is the paper's stated conflict with large tiles 2407.08608 §3.2.

A three-stage version was tried and was slower. The compiler did not overlap the second matrix multiply, and register pressure forced smaller tiles 2407.08608 app. B.3. On the paper's ablation setting (batch 4, sequence length 8448, 16 heads, head dimension 128, non-causal, 16-bit), the full design reaches 661 TFLOPs/s. Removing warp specialisation gives 582, and keeping only the overlap gives 570 2407.08608 tab. 2 [arxiv:2407.08608 §4.2]. The two features help separately, and the combination is fastest.

8-bit forward pass (FP8, Hopper)

The 8-bit path uses the e4m3 format 2407.08608 §3.3. It needs three extra steps that the 16-bit path does not:

• Transposing V. FP8 matrix multiplies take both operands along the sequence dimension, so V tiles are transposed inside shared memory during the kernel. The paper rejects a global transpose before the call, because it costs memory bandwidth 2407.08608 §3.3 2407.08608 fig. 3 2407.08608 fig. 4.
• Block quantisation. Each block of Q, K and V gets its own scale, rather than one scale per tensor 2407.08608 §3.3.
• Incoherent processing. Q and K are multiplied by an orthogonal matrix before quantisation. Because the matrix cancels in Q K^T, the output is unchanged while large outlier values are spread across entries. The paper builds the matrix from a random sign diagonal and a Hadamard matrix 2407.08608 §3.3.

The paper reports FP8 forward throughput close to 1.2 PFLOPs/s [arxiv:2407.08608 §1] 2407.08608 fig. 7. On its synthetic test, with 0.1% outlier entries, the RMSE against an FP64 reference is 2.4e-2 for the per-tensor baseline and 9.1e-3 for this kernel, a ratio of about 2.6 [arxiv:2407.08608 tab. 3] 2407.08608 §4.3. The test uses synthetic inputs, not real model activations 2407.08608 §4.3.

When to use which

Implementation sketch

The core loop for one query block, with a running max and sum. It is the same for every variant; the overlap and the 8-bit steps change how the loop body is scheduled, not what it computes.

# Tiled attention for one query block, running softmax  [arxiv:2407.08608 §3.1]
# Input:  q (B_r, d) in fast memory; keys, values (N, d) streamed in blocks
# Output: o (B_r, d); logsumexp lse (B_r,)

def fused_attention_block(q, keys, values, block_c):
    # Start with empty running state: output, row max and row sum  [§3.1]
    o = zeros(B_r, d)
    m = full(B_r, -inf)
    l = zeros(B_r)
    for j in range(ceil(N / block_c)):
        # Scores for key block j; they stay on chip  [arxiv:2407.08608 §2.3]
        s = q @ keys[j].T                      # (B_r, block_c)
        # New row max, and the rescale factor for the old state  [§3.1]
        m_new = maximum(m, rowmax(s))
        alpha = exp(m - m_new)                 # (B_r,)
        p = exp(s - m_new[:, None])            # (B_r, block_c)
        # Update the row sum and the output with the same factor  [§3.1]
        l = alpha * l + rowsum(p)
        o = alpha[:, None] * o + p @ values[j]
        m = m_new
    # Normalise once at the end and return the logsumexp  [§3.1]
    o = o / l[:, None]
    lse = m + log(l)
    return o, lse

Two notes from the card's checks against the paper:

• The rescale factor here is exp(old max - new max) on the output. The paper's Algorithm 1 prints the inverse for the output update 2407.08608 alg. 1. The card follows the form used in the released code; see the card.
• The paper's Algorithm 2 loop bound leaves the last key block unscored in the form printed. This sketch scores every block 2407.08608 alg. 2.

Contradictions

No disagreement between papers: there is one card so far.

One gap between the paper and its released code is worth knowing before you rely on the 8-bit accuracy result. The paper's incoherent processing step 2407.08608 §3.3 was not found in the released repository at the pinned commit, according to the card's code review. Reproduce the accuracy claim before relying on it. Card: FlashAttention-3.

Papers

• FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision: warp-specialised, overlapped Hopper kernel for 16-bit and 8-bit attention, with author-reported speed-ups over FlashAttention-2.