Compute single-head scaled dot-product attention by scanning over blocks of keys and values with lax.scan, keeping a running maximum, a running sum and a running output for every query (the online softmax used by flash attention).
The result must equal the full attention
O=softmax(dqk⊤+M)v,where q has shape (Tq, d), k and v have shape (Tk, d), and the softmax runs over the key axis. M is 0 everywhere when causal is false. When causal is true, Tq equals Tk and M is minus infinity wherever the key position is greater than the query position (positions start at 0), so each query sees only itself and earlier keys.
Split k and v into Tk / block_size consecutive blocks of block_size keys each.
Return a float32 JAX array of shape (Tq, d).
Input: q = [[1.0,0.0],[0.0,1.0]], k = [[1.0,0.0],[0.0,1.0],[1.0,1.0],[-1.0,0.0]], v = [[1.0,2.0],[3.0,4.0],[5.0,6.0],[7.0,8.0]], block_size = 2, causal = false
Output: [[3.35541,4.35541],[4.0,5.0]]
Explanation: With block_size = 2 the scan runs over two blocks of two keys. The second query has dot product 1 with the second and third keys and 0 with the others, so its output leans towards their values, and the online result equals the full softmax average.
Input: q = [[0.5,-1.0],[1.0,0.5],[-0.5,0.5],[2.0,0.0],[0.0,-2.0],[1.0,1.0]], k = [[1.0,0.0],[0.0,1.0],[-1.0,1.0],[0.5,0.5],[2.0,-1.0],[0.0,0.0]], v = [[1.0,0.0],[0.0,1.0],[1.0,1.0],[-1.0,2.0],[3.0,-1.0],[0.5,0.5]], block_size = 2, causal = true
Output: [[1.0,0.0],[0.587479,0.412521],[0.657204,0.830978],[0.31529,0.717633],[2.148492,-0.43346],[0.75,0.549444]]
Input: q = [[40.0,40.0,0.0]], k = [[40.0,40.0,0.0],[-40.0,40.0,0.0],[40.0,39.9,0.0],[0.0,-40.0,40.0]], v = [[1.0,0.0,0.0],[0.0,1.0,0.0],[0.0,0.0,1.0],[1.0,1.0,1.0]], block_size = 2, causal = false
Output: [[0.90965,0.0,0.09035]]
Explanation: After dividing by the square root of 3, the scores of the first and third keys are about 1847.520861 and 1845.211496, so a plain exp overflows. Subtracting the running maximum leaves a gap of about 2.309366 between them, which gives the first key a weight of 0.90965 and the third key a weight of 0.09035. The other two keys score so much lower that their weights round to 0. The input 39.9 is stored in float32 as 39.9000015, which is why these weights differ from exact decimal arithmetic in the sixth decimal place.
k.reshape(num_blocks, block_size, d) splits keys into blocks for jax.lax.scan.
jax.lax.scan(step, init, (k_blocks, v_blocks)) scans several arrays together.
jnp.max(s, axis=-1, keepdims=True) keeps running stats shaped (Tq, 1).
Accepts: array
Accepts: array
Accepts: array
Accepts: integer
Compute single-head scaled dot-product attention by scanning over blocks of keys and values with lax.scan, keeping a running maximum, a running sum and a running output for every query (the online softmax used by flash attention).
The result must equal the full attention
O=softmax(dqk⊤+M)v,where q has shape (Tq, d), k and v have shape (Tk, d), and the softmax runs over the key axis. M is 0 everywhere when causal is false. When causal is true, Tq equals Tk and M is minus infinity wherever the key position is greater than the query position (positions start at 0), so each query sees only itself and earlier keys.
Split k and v into Tk / block_size consecutive blocks of block_size keys each.
Return a float32 JAX array of shape (Tq, d).
Input: q = [[1.0,0.0],[0.0,1.0]], k = [[1.0,0.0],[0.0,1.0],[1.0,1.0],[-1.0,0.0]], v = [[1.0,2.0],[3.0,4.0],[5.0,6.0],[7.0,8.0]], block_size = 2, causal = false
Output: [[3.35541,4.35541],[4.0,5.0]]
Explanation: With block_size = 2 the scan runs over two blocks of two keys. The second query has dot product 1 with the second and third keys and 0 with the others, so its output leans towards their values, and the online result equals the full softmax average.
Input: q = [[0.5,-1.0],[1.0,0.5],[-0.5,0.5],[2.0,0.0],[0.0,-2.0],[1.0,1.0]], k = [[1.0,0.0],[0.0,1.0],[-1.0,1.0],[0.5,0.5],[2.0,-1.0],[0.0,0.0]], v = [[1.0,0.0],[0.0,1.0],[1.0,1.0],[-1.0,2.0],[3.0,-1.0],[0.5,0.5]], block_size = 2, causal = true
Output: [[1.0,0.0],[0.587479,0.412521],[0.657204,0.830978],[0.31529,0.717633],[2.148492,-0.43346],[0.75,0.549444]]
Input: q = [[40.0,40.0,0.0]], k = [[40.0,40.0,0.0],[-40.0,40.0,0.0],[40.0,39.9,0.0],[0.0,-40.0,40.0]], v = [[1.0,0.0,0.0],[0.0,1.0,0.0],[0.0,0.0,1.0],[1.0,1.0,1.0]], block_size = 2, causal = false
Output: [[0.90965,0.0,0.09035]]
Explanation: After dividing by the square root of 3, the scores of the first and third keys are about 1847.520861 and 1845.211496, so a plain exp overflows. Subtracting the running maximum leaves a gap of about 2.309366 between them, which gives the first key a weight of 0.90965 and the third key a weight of 0.09035. The other two keys score so much lower that their weights round to 0. The input 39.9 is stored in float32 as 39.9000015, which is why these weights differ from exact decimal arithmetic in the sixth decimal place.
k.reshape(num_blocks, block_size, d) splits keys into blocks for jax.lax.scan.
jax.lax.scan(step, init, (k_blocks, v_blocks)) scans several arrays together.
jnp.max(s, axis=-1, keepdims=True) keeps running stats shaped (Tq, 1).
Accepts: array
Accepts: array
Accepts: array
Accepts: integer
