Hard Research

Attention with Online Softmax

Attention usually computes a score for every query-key pair before applying softmax. Online softmax lets you build the same weighted average incrementally, without keeping the whole score matrix.

The challenge is numerical stability. If a later key produces a larger score, you must rescale both the accumulated weight and the accumulated weighted values. Otherwise, contributions computed with different offsets will not be comparable.

Your task

Implement flash_attention(Q, K, V) using online softmax, processing one key/value pair at a time for each query. Keep the supplied function name. All three inputs are floating-point arrays of shape (seq_len, d_k), with positive dimensions and finite scaled dot products. They share a dtype and, in PyTorch, a device. Return shape (seq_len, d_k) with the same dtype/device.

For each query, initialize a running maximum m = -infinity, a scalar weight sum l = 0, and a length-d_k weighted-value vector acc = 0. For key j, compute the following updates in order:

s = dot(query, K[j]) / sqrt(d_k)
m_new = max(m, s)
rescale = exp(m - m_new)
weight = exp(s - m_new)
l = rescale * l + weight
acc = rescale * acc + weight * V[j]
m = m_new

After all keys, the query's output is acc / l. Reset all three running statistics for each new query. There is no mask: every query reads every key.

For the first finite score, rescale is zero and weight is one, so the initial negative infinity is safe. Keep the statistics as array scalars rather than converting scores to Python floats.

Why the rescaling works

Both l and acc are expressed relative to the current maximum. Replacing that maximum with a larger one multiplies every previous exponential by exp(m - m_new). Rescaling both accumulators keeps the final ratio unchanged.

For scores [0, log(3)] and scalar values [2, 6], the first step gives l = 1 and acc = 2. The second rescales them by 1/3, then adds weight 1 and weighted value 6. The output is (2/3 + 6) / (1/3 + 1) = 5, the same as weighting the values by [1/4, 3/4].

Connection to FlashAttention

This exercise isolates an online-softmax idea used in FlashAttention. It is an educational scalar-loop implementation, not the tiled GPU algorithm from the paper. It avoids storing the full score matrix during this forward calculation, but still evaluates every query-key pair. Python loops do not imply a speed improvement, and this task makes no claim about backward-pass memory. The numerical tests check attention outputs; they do not certify a particular memory-access pattern or optimized kernel.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Share keys and values across attention heads

Multi-query attention lets several query heads share one key head and one value head. Grouped-query attention shares keys and values within groups of query heads. For example, eight query heads and two key/value heads can form two groups of four query heads each. Queries remain distinct within each group.

Follow the specified mapping between query and key/value heads. Their counts must support that grouping. Sharing keys and values reduces their storage, including the cache used during generation.

Restrict which keys each query can read

Sliding-window attention limits a query to nearby keys. A symmetric window can include positions on either side; a causal window includes only permitted past positions and, where specified, the current position.

Masking a full score matrix changes which positions contribute, but still requires storage for that full matrix. To avoid quadratic score storage, the implementation must also avoid creating the full matrix.

Compute attention a block at a time

For each query, a blockwise calculation tracks three quantities: the largest score seen so far, the sum of exponentials relative to that maximum, and the sum of values weighted by those exponentials.

If a new block raises the maximum from m to m_new, multiply both old sums by exp(m - m_new) before adding terms from the new block. This puts old and new terms on the same scale. Divide the final weighted-value sum by the final exponential sum to get the attention output.

Do not simply average the attention outputs of independently normalized blocks. Their total probability weights generally differ. Handle fully masked blocks without subtracting negative infinity from itself, and follow the task's rule for a query with no allowed keys.

With valid rows, this method matches dense attention up to rounding. FlashAttention also organizes computation and memory transfers to use hardware efficiently; implementing equivalent arithmetic alone does not provide the same speedup. See the FlashAttention paper.

Detecting runtime…