Deep Learning - Flash Attention

What FlashAttention Does

FlashAttention computes exact attention without storing the huge \(N\times N\) attention matrix:

Because GPU performance often depends heavily on memory traffic. A GPU has two kinds of memory that matter here:

• HBM (high-bandwidth memory, the GPU’s global memory): large (tens of GB) but relatively slow to read and write.
• SRAM (on-chip shared memory): tiny (around 100–200 KB per streaming multiprocessor) but much faster.

1. Attention is memory-bound. The GPU can multiply much faster than it can move the \(N\times N\) matrix in and out of HBM, so cutting HBM traffic cuts the runtime.
2. Kernel fusion. Matmul, mask, softmax, and the multiply by V happen in one GPU kernel, so the intermediate tiles stay in SRAM instead of making round trips through HBM between kernels.
3. Memory drops from \(O(N^2)\) to \(O(N)\). Only Q, K, V, the output, and two small per-row statistics (m and l, below) live in HBM. This allows longer sequences or larger batches on the same GPU.

So it reduces expensive HBM reads and writes. Suppose the sequence length is N. Then:

Instead of computing all of \(QK^T\) at once, it divides Q, K, V into small blocks along the sequence dimension:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
Q:
┌─────┐
│ Q1  │
├─────┤
│ Q2  │
├─────┤
│ Q3  │
└─────┘

K:
┌─────┐
│ K1  │
├─────┤
│ K2  │
├─────┤
│ K3  │
└─────┘

The full \(N\times N\) attention matrix is never materialized in GPU global memory.

Refresher - Softmax of a Matrix Normalizes Each Row

In attention, softmax(A) is applied row by row: each row is one query’s scores over all keys, so each row of the result sums to 1:

An Example

Take 4×4 matrices Q and K. To get 2×2 attention-score tiles, split Q and K into 2×4 row blocks, since (2×4)(4×2) = 2×2. I’ll use deliberately simple matrices.

1. Define Q, K, V

With \(d_k=4\) (so \(\sqrt{d_k}=2\)), let

just makes the final output very easy to interpret.

2. Split into blocks

Split each matrix into two row blocks. V is split the same way as K, since key j goes with value j:

3. Compute the 2x2 tiles In Parallel

FlashAttention computes one tile \(S_{ij} = Q_iK_j^T/2\) at a time. For reference, the full matrix it never materializes is on the right:

Now comes the FlashAttention part.

4. Q_1, first block K_1, V_1

Each query row has its own running max. In S_11, row 1 is \([1,2]\) and row 2 is \([2,1]\), so both maxima are 2. Exponentiating the shifted scores gives unnormalized softmax weights (not probabilities yet); their row sums are the running denominators:

Q_2 starts the same way with S_21, independently of Q_1 (see step 7).

5. Same Q_1, next row block K_2, V_2

The new tile is \(S_{12} = \begin{bmatrix} 3&4\\ 0&-1 \end{bmatrix}\). Update the row-wise max and compute the rescale factor:

Row 1’s max changed \(2\rightarrow4\), so row 1 of the old O and l shrinks by \(0.1353\); row 2’s max stayed 2, so row 2 is unchanged. Next, exponentiate the new tile relative to m_new (row 1: \([3,4]-4=[-1,0]\); row 2: \([0,-1]-2=[-2,-3]\)) and multiply by V_2:

Add that to the rescaled accumulator \(\alpha O_{\text{old}}\), and update the denominators with \(l_{\text{new}} = \alpha\, l_{\text{old}} + \operatorname{rowsum}(P_{12})\):

6. Finish Q_1: normalize

Divide each row of the accumulator by its denominator:

Because \(V=I\), these rows are the attention probabilities themselves: \(\operatorname{softmax}([1,2,3,4]) = [0.0321,0.0871,0.2369,0.6439]\) for query 1, and \(\operatorname{softmax}([2,1,0,-1]) = [0.6439,0.2369,0.0871,0.0321]\) for query 2.

So FlashAttention got exactly the result we’d get from the full matrix, without ever needing to keep \([1,2,3,4]\) as a complete score row in global memory.

7. Meanwhile Q_2 can do the same thing

Independently:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
Q1 worker:

S11 → update state
       ↓
S12 → update state
       ↓
output rows 1,2


Q2 worker:

S21 → update state
       ↓
S22 → update state
       ↓
output rows 3,4

For Q_2, the tiles S_21 and S_22 go through the same updates and give rows 3 and 4. Stacking both workers’ outputs gives the complete result:

The Entire FlashAttention Flow

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
Q = 4x4                     K = 4x4
   │                           │
   ▼                           ▼
split rows                  split rows

Q1 = 2x4                   K1 = 2x4
Q2 = 2x4                   K2 = 2x4


       Q1                       Q2
        │                        │
        ▼                        ▼
   Q1 K1^T                 Q2 K1^T
      2x2                     2x2
        │                        │
        ▼                        ▼
 update m,l,O              update m,l,O
        │                        │
        ▼                        ▼
   Q1 K2^T                 Q2 K2^T
      2x2                     2x2
        │                        │
        ▼                        ▼
 update m,l,O              update m,l,O
        │                        │
        ▼                        ▼
 final rows 1,2            final rows 3,4

The important pattern is:

The two Q blocks can be handled in parallel, while the K/V blocks are streamed through each Q block’s running state.

• Previous
  Deep Learning - Transformer Series 6 - Decoder-Only LLMs and KV Cache
• Next
  Deep Learning - Contrastive Learning