Attention please

@stevhliu|Aug 10, 2025 (5 months ago)36 views

An attention function can be described as mapping a query and a set of key-value pairs to an output, where the query, keys, values, and output are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key.

Attention Is All You Need

It's amazing that an idea so simple has such a tremendous and lasting impact. This idea is described in only two sentences, but it is consumed by millions of people using large language models daily.

This post tries to very simply and very plainly explain how attention, and variants of it, work.

#self-attention

Self-attention (scaled dot-product attention) computes how much attention each word in a sequence assigns to every other word in that same sequence. It's what makes transformers contextually aware.

Q × K^T produces attention scores, scaled and softmaxed into weights, then multiplied by V.

You need 3 matrices to compute self-attention. These matrices are created by multiplying word embeddings by 3 weight matrices (W_k, W_q, W_v).

Query (Q) is the information you're looking for. It is compared to every K (including itself) to calculate how much attention to pay to each word.
Key (K) is the information a word contains. It is multiplied (dot product) by Q to get the attention scores for each word.

The scores are scaled by dividing by the square root of the vector dimension. For example, if the dimension is 64, divide the attention scores by 8. Scaling prevents the attention scores from becoming too large or too small.

Scores are converted into probabilities that add up to 1 with the softmax function. Scaling also smooths out the softmax function by preventing any one value from getting too large and overwhelming the rest.
Value (V) is the information each word contributes. It weights each word with the attention score to determine what information every other word offers.

Try calculating the self-attention score by hand in the following sequence to really get a feel for how it works.

#causal self-attention

Causal self-attention is used in decoder models like GPT. Decoder models predict the next word in a sequence, so it is important to mask the next words to prevent the model from seeing them.

Q, K, V per head. Future positions masked (gray) so tokens only attend to past.

The mask sets words that should be blocked to -∞, which is converted to ~0 by the softmax.

Try calculating the causal self-attention score by hand in the following sequence.

#multi-head attention

Multi-head attention (MHA) adds more heads to learn from different perspectives.

Each head has its own Q, K, and V projections.

The calculations are the same as self-attention, but the embedding size is split by the number of heads. Each head independently computes the scaled dot-product attention on their slice of data. At the end, the outputs are combined and multiplied by a weight matrix to blend all the information together.

BERT uses MHA.

Try calculating the multi-head attention score by hand in the following sequence.

#multi-query attention

Multi-query attention (MQA) is the same as MHA except every Q shares the same K, V.

Each head has its own Q. All heads share a single K and V.

MQA is more memory-efficient and faster at decoding because each head doesn't need to store a separate K, V. This makes an especially big difference for really long sequences.

Gemma 2B uses MQA.

Try calculating the multi-query attention score by hand in the following sequence.

#grouped-query attention

Grouped-query attention (GQA) is similar to MHA and MQA except groups of Q share the same K, V. K, V is different for each group.

Each head has its own Q. K and V are shared within each group.

GQA is the middle ground between MHA and MQA. It's faster than MHA and more expressive than MQA.

gpt-oss-20b uses GQA.

Try calculating the grouped-query attention score by hand in the following sequence.

#flashattention

FlashAttention changes how the attention calculation is executed to minimize memory traffic to and from the GPU's slower high-bandwidth global memory (HBM). The calculation is done on the faster on-chip shared memory (SRAM) and registers, without fully materializing the entire attention score matrix in HBM.

Block 2 uses m₁, l₁ from block 1 to compute new stats m₂, l₂ and rescale l₁ before accumulating.

Q, K, and V are tiled into blocks that fit in SRAM. Tiling lets the calculation be performed in a CUDA kernel.
In the first block, calculate a local max m₁ and a sum l₁ for Q, K.
After processing the second block, update the local max to m₂. Rescale l₁ and the partial output before adding the new updates.
Repeat for each calculation.

Attention scores are streamed over blocks to avoid storing the entire attention score matrix in memory. The softmax calculation is exactly the same. Only Q, K, V, the final output, and the softmax normalization factor from the forward pass are stored in HBM.

The backward pass uses the same idea to avoid storing the attention probabilities P in memory. P is required for computing gradients with respect to Q, K, and V. FlashAttention recomputes P block-by-block, uses the same softmax statistics from the forward pass (m, l), and accumulates the gradients.

Speed ups increase with sequence length. But for short sequences, attention is less memory-bound and there's less benefit.

#flashattention-2

FlashAttention-2 improves FlashAttention by reducing non-matmul FLOPs and scheduling more work on the GPU.

GPUs use specialized compute units, Tensor Cores on Nvidia GPUs, for fast matmuls. In FlashAttention, every block rescales the output with the sum l.

FlashAttention-2 delays rescaling to the end to avoid many scalar divisions. This reduces the number of non-matmul operations.

FlashAttention-2 splits work across sequence blocks, using more SMs when batch size and heads are small.

Each attention head is processed by a single thread block. FlashAttention parallelized by batch size x num heads. For long sequences, when batch sizes are small, there are fewer heads and not all the streaming multiprocessors (SM) are used.

FlashAttention-2 parallelizes over batch size x num heads x seq length. The work within a single attention head is split across more thread blocks to use more SMs.

FlashAttention-2 assigns Q blocks to warps instead of K,V, reducing shared memory writes and synchronization.

Warps are groups of threads in a thread block. FlashAttention split the work on K, V and Q is accessible by all warps. After computing attention, warps write the partial results to shared memory, synchronize, and reduce.

FlashAttention-2 splits the work on Q and K, V is accessible by all warps. Each warp computes a partial result that doesn't require synchronization until the end. Assigning different Q blocks to different warps minimizes shared memory traffic.

#flashattention-3

FlashAttention-3 optimizes for newer hardware, like H100 GPUs, that support asynchrony. Specialized compute units run in parallel. Tensor cores can do fast matmuls, Tensor Memory Accelerator (TMA) loads data from HBM, and CUDA cores can perform slower computations like softmax.

FlashAttention-3 overlaps memory transfers and computation using specialized producer/consumer warps.

Warps performed the same work in FlashAttention-2, loading and computing.

FlashAttention-3 specializes warps to different tasks. Producer warps load K, V blocks from HBM to SRAM with TMA. Consumer warps compute attention using data in SRAM with tensor cores. The producer can prefetch the next block while the consumer is working on the current block. It overlaps loading and computation.

WGMMA issues async matmuls, allowing softmax to run in parallel on the previous block's scores.

The output O depends on P. P is calculated from the softmax(S), and S depends on the attention score calculation.

FlashAttention-3 uses Warp Group Matrix Multiply-Accumulate (WGMMA) to asynchronously perform the computation. It pipelines blocks so while softmax executes on one block of the scores matrix, WGMMA executes to compute the next block. There are two pipelines running in parallel, WGMMA executions and non-WGMMA executions.
FP8 quantization is sensitive to outlier features in large language models (LLMs). A few dimensions in hidden states can have very large values. When quantized to FP8, the issue manifests as precision loss.

Each block computes its own scale so precision is preserved locally. Both use the full FP8 range.

FlashAttention-3 uses a different scale per block. Each block can use the full FP8 range for its values. The outlier block gets a big scale and the normal blocks get a tiny scale. This way, precision is preserved locally.

Multiplying by random orthogonal matrix R redistributes the outlier across all dimensions while preserving dot products.

If an outlier is in a normal block, FlashAttention-3 fixes this by spreading the outlier across all dimensions before quantization. It multiplies Q, K by a random orthogonal matrix to redistribute the outlier. This way, no single dimension dominates a block.

#summary

Self-attention is a weighted sum of contextual information of all words. The weights reflect each surrounding word's relevance to the current word.

But attention is memory-bound. New attention algorithms optimize how the attention calculation is executed, whether it's sharing the same K, V or streaming the attention function.