Tiled attention with an online softmax · AI Engineering

Standard attention builds the full matrix of scores, of size T_q \times T_k. FlashAttention never builds that matrix. It reads the keys and values in blocks and keeps a running softmax for each query row. The result is exact, not an approximation. Here you write that running softmax in numpy.

Task

Write two functions.

online_softmax_step(m, l, acc, scores, v_block) adds one block of keys to the running state of each query row.
flash_attention(Q, K, V, block_size, causal=False) runs the step over the key blocks and returns the attention output and the log-sum-exp of each row.

The state of query row i has three parts: the running max m_i, the running sum l_i, and the accumulator a_i, a vector of size d_v. At the start, m_i = -\infty, l_i = 0 and a_i = 0. For a block with scores s_{ij} and value rows v_j, the step computes:

m_i' = \max(m_i, \max_j s_{ij}), \qquad c_i = e^{m_i - m_i'}

l_i' = c_i \, l_i + \sum_j e^{s_{ij} - m_i'}, \qquad a_i' = c_i \, a_i + \sum_j e^{s_{ij} - m_i'} \, v_j

A score of -\infty is a blocked key. When m_i' is -\infty, the row has no allowed key yet. Then use 0 in place of m_i' inside both exponentials, so that c_i = 0 and each term is 0. The step then returns m_i' = -\infty, l_i' = 0 and a_i' = 0.

In flash_attention, the scores of a block are Q K_b^T / \sqrt{d}, where K_b holds the key rows from start to start + block_size and d is the size of the last axis of Q. The blocks start at 0, block_size, 2 * block_size and so on. The last block can be shorter. With causal=True, key j is blocked for query i when j > i. The indices i and j are global row indices, not indices inside a block. After the last block, the output row is a_i / l_i and the log-sum-exp is m_i + \log l_i.

Input and output

Name	Type	Shape	Meaning
`m`	float array	`(Tq,)`	the running max of each query row
`l`	float array	`(Tq,)`	the running sum of each query row
`acc`	float array	`(Tq, d_v)`	the running accumulator of each query row
`scores`	float array	`(Tq, b)`	the scores of one block of `b` keys, `-inf` for a blocked key
`v_block`	float array	`(b, d_v)`	the value rows of the same block
`Q`	float array	`(Tq, d)`	the queries
`K`	float array	`(Tk, d)`	the keys
`V`	float array	`(Tk, d_v)`	the values
`block_size`	int		the number of keys in each block
`causal`	bool		`True` blocks key `j` for query `i` when `j > i`

online_softmax_step returns a tuple (m_new, l_new, acc_new) with the shapes of m, l and acc.
flash_attention returns a tuple (output, lse). output has shape (Tq, d_v). lse has shape (Tq,).

Constraints

1 <= Tq, Tk, d, d_v <= 64 and 1 <= block_size <= 128. The block size can be larger than Tk.
Process the keys block by block. Do not build the full (Tq, Tk) score matrix.
Scores can be larger than 1000 in one block and smaller than minus 1000 in the next. The output must stay finite and exact.
The result must equal standard attention for every block size, to about 1e-9.
Do not change the input arrays.
Use numpy only.

Examples

python

Q = np.array([[1.0, 0.0], [0.0, 1.0]])
K = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])
V = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
output, lse = flash_attention(Q, K, V, 2)

text

output
[[3.         4.        ]
 [3.40667256 4.40667256]]
lse
[1.62062114 1.62062114]

Block sizes 1 and 3 give the same values.

python

output, lse = flash_attention(K, K, V, 2, causal=True)

text

output
[[1.         2.        ]
 [2.3395231  3.3395231 ]
 [3.51046953 4.51046953]]
lse
[0.70710678 1.10794031 2.1004053 ]

python

m = np.array([1.0, -np.inf])
l = np.array([2.0, 0.0])
acc = np.array([[1.0, 1.0], [0.0, 0.0]])
scores = np.array([[3.0, 0.0], [-np.inf, 0.5]])
v_block = np.array([[1.0, 0.0], [0.0, 1.0]])
m_new, l_new, acc_new = online_softmax_step(m, l, acc, scores, v_block)

text

m_new
[3.  0.5]
l_new
[1.32045763 1.        ]
acc_new
[[1.13533528 0.18512235]
 [0.         1.        ]]