FMHA 101: From Transformer Attention to a CUDA Flash Attention Kernel

Flash Multi-Head Attention (FMHA) can look like a specialized GPU trick, but it starts from a simple question: when a model processes one token, which other tokens should influence it most? This tutorial begins with that question, builds the standard attention calculation with a small example, then explains how Flash Attention reorganizes the same mathematics to avoid a major GPU memory bottleneck.

Who this is for. You should be comfortable reading basic matrix notation and CUDA-style C++. No prior experience with Transformer internals or Flash Attention is required.

1. Why Transformers need attention

A language model cannot interpret a token in isolation. In the sentence “The animal didn’t cross the street because it was too tired,” resolving it depends on information elsewhere in the sequence. Attention gives every token a way to gather a weighted mixture of relevant context. The original Transformer introduced this mechanism as a core building block that can be computed in parallel across a sequence. [1]

For each token position, a model derives three learned vectors:

2. Standard scaled dot-product attention

Stacking all token vectors into matrices gives the familiar equation:

O = softmax((QKᵀ) / √d)V

Here, d is the head dimension. The matrix product QKᵀ measures how well each query matches every key. Dividing by √d keeps the logits at a manageable scale before softmax turns each row into non-negative weights that sum to one. The final multiplication by V retrieves a weighted mixture of values.

3. A tiny worked example

Imagine a query token that is comparing itself with four candidate tokens. Suppose the scaled scores are [2.0, 1.0, 0.0, -1.0]. Applying softmax yields approximately [0.64, 0.24, 0.09, 0.03]. The model therefore assigns most of the context to the first token, some to the second, and very little to the others.

scores:              [ 2.00,  1.00,  0.00, -1.00]
softmax(scores):     [ 0.64,  0.24,  0.09,  0.03]
value vectors:       [ V₁,    V₂,    V₃,    V₄   ]
output for this token = 0.64V₁ + 0.24V₂ + 0.09V₃ + 0.03V₄

This is attention in its most concrete form: score potential sources, convert scores into a distribution, then retrieve a weighted average. A real attention head performs this for every query token in parallel.

4. Where “multi-head” enters

One attention pattern is rarely enough. Multi-head attention creates several independent sets of Q, K, and V; each head can specialize in a different relationship, such as nearby syntax, long-range reference, or position. In a CUDA kernel, inputs are commonly stored as [B, H, N, d]: batch size B, number of heads H, sequence length N, and head dimension d.

5. Why ordinary attention becomes expensive

The standard calculation is mathematically straightforward, but a naïve GPU implementation materializes both the score matrix S = QKᵀ and the probability matrix P = softmax(S). Each is N × N per head. At long sequence lengths, repeatedly writing and rereading these intermediates from high-bandwidth memory becomes more costly than the floating-point arithmetic itself. Flash Attention retains the exact result while changing the order of computation to reduce those memory transfers. [2]

6. The Flash Attention idea: work in tiles

Rather than form every score at once, FMHA divides the matrices into tiles that fit in fast on-chip shared memory. The kernel loads a block of queries, then walks through blocks of keys and values. For each key/value tile, it computes a local score tile, contributes to the output tile, and discards the temporary scores. The full N × N matrix never needs to exist in global memory.

7. Shared memory is the workspace

In the educational CUDA kernel, the dynamic shared-memory buffer is partitioned into input tiles, a temporary score tile, an output accumulator, and small per-row statistics. Br is the number of query rows in a tile; Bc is the number of key/value rows.

extern __shared__ float smem[];
float* Qi  = smem;                    // [Br * d]  query tile
float* Kj  = Qi + Br * d;             // [Bc * d]  key tile
float* Vj  = Kj + Bc * d;             // [Bc * d]  value tile
float* Sij = Vj + Bc * d;             // [Br * Bc] score/probability tile
float* Oi  = Sij + Br * Bc;           // [Br * d]  output accumulator
float* li  = Oi + Br * d;             // [Br]      running softmax sum

The thread block cooperatively loads the key and value tile so that the data can be reused many times. A thread then computes a scaled query–key dot product for one element of the score tile:

// One thread computes one element of Sij = scale * Qi * Kj^T.
float score = 0.0f;
for (int k = 0; k < d; ++k) {
    score += Qi[s_row * d + k] * Kj[s_col * d + k];
}
Sij[s_row * Bc + s_col] = score * scale;

8. The one complication: softmax needs the whole row

Softmax is made stable by subtracting the maximum logit in a row. But a tiled kernel sees only part of the row at a time. Flash Attention solves this with an online softmax. For every query row, it keeps:

If a later key tile contains a larger maximum, the kernel rescales the previous l and Oi into the new numerical reference frame, then adds the current tile’s contribution. This is why the algorithm remains both stable and exact:

// Move previous state to the new numerical reference frame.
li_new[row] = expf(mi[row] - mi_new[row]) * li[row]
            + expf(row_m - mi_new[row]) * row_l;

Oi[row * d + col] = expf(mi[row] - mi_new[row]) * Oi[row * d + col]
                   + expf(row_m - mi_new[row]) * pv;

9. Final output and what the kernel does not yet optimize

Only after every key/value tile has contributed does the kernel divide the accumulator by its final normalizer:

// Normalize only after every K/V tile has contributed.
O[qkv_off + global_row * d + col] = Oi[row * d + col] / li[row];

This is intentionally an educational implementation, not a drop-in replacement for production FlashAttention. High-performance kernels add vectorized memory operations, warp-level reductions, Tensor Core matrix-multiply instructions, architecture-specific pipelining, causal and padding masks, dropout, mixed precision, and backward-pass support. FlashAttention-2 further improves work partitioning and reduces non-matrix-multiplication overhead. [3]

10. When FMHA matters

FMHA is particularly valuable for long sequences, many attention heads, and training or inference workloads where attention memory traffic is substantial. In application code, prefer a tested framework implementation when available. Writing a CUDA kernel is most useful when you need to learn the algorithm, inspect a performance bottleneck, support an unusual layout or mask, or experiment with hardware-aware techniques.

Takeaway

Attention is weighted retrieval: queries choose among keys, and the corresponding values are blended into an output. FMHA does not approximate that operation. Instead, it preserves the same result while ensuring that temporary data is consumed in shared memory rather than repeatedly materialized in global memory. That reordering—together with online softmax—is the central idea behind Flash Attention.

References

1. Vaswani et al., Attention Is All You Need.
2. Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.
3. Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning.