Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

FlashAttention Algorithm › Worked Example

Michael BrenndoerferUpdated June 29, 20257 min read

Language AI HandbookAll chapters

Part 17 of 67Efficient Attention

Chapter 124FlashAttention Algorithm › Worked Example

Reading section 8 of 12 · Worked Example · 7 minSection 8 of 12 · 7 min

Section 8 of 12 · 7 min read

FlashAttention Algorithm › Worked ExampleOne section · Sections in this chapter8 / 12 ↓

Sections in this chapter12

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

The best way to internalize the FlashAttention algorithm is to trace through it step by step with concrete numbers. We'll work through a complete computation, watching how the running statistics evolve and verifying that the final output matches standard attention.

Setting Up the ExampleLink Copied

Consider a small example designed to make the arithmetic tractable:

Sequence length $n = 8$ (giving us 8 query and 8 key/value positions)
Head dimension $d = 4$ (each Q, K, V vector has 4 components)
Block size $B = 4$ (we process 4 positions at a time)

This configuration divides the computation into 2 query blocks and 2 key/value blocks. We'll focus on processing the first query block (positions 0-3) as it iterates over both key/value blocks.

We'll use carefully constructed Q and K matrices where the dot products create interpretable attention patterns:

In[25]:

Code

# Set up the example
n, d = 8, 4
block_size = 4

# Create simple Q, K, V matrices with interpretable structure
# Each query/key has a distinctive pattern that creates clear attention scores
Q = np.array(
    [
        [1.0, 0.0, 0.0, 0.0],  # Query 0: attends to keys with component 0
        [0.0, 1.0, 0.0, 0.0],  # Query 1: attends to keys with component 1
        [0.0, 0.0, 1.0, 0.0],  # Query 2: attends to keys with component 2
        [0.0, 0.0, 0.0, 1.0],  # Query 3: attends to keys with component 3
        [0.5, 0.5, 0.0, 0.0],  # Query 4: mixed attention
        [0.0, 0.5, 0.5, 0.0],  # Query 5: mixed attention
        [0.0, 0.0, 0.5, 0.5],  # Query 6: mixed attention
        [0.5, 0.0, 0.0, 0.5],  # Query 7: mixed attention
    ]
)

K = np.array(
    [
        [1.0, 0.0, 0.0, 0.0],  # Key 0: matches query 0
        [0.0, 1.0, 0.0, 0.0],  # Key 1: matches query 1
        [0.5, 0.5, 0.0, 0.0],  # Key 2: partial match with queries 0, 1
        [0.0, 0.5, 0.5, 0.0],  # Key 3: partial match with queries 1, 2
        [0.0, 0.0, 1.0, 0.0],  # Key 4: matches query 2
        [0.0, 0.0, 0.0, 1.0],  # Key 5: matches query 3
        [0.0, 0.5, 0.5, 0.0],  # Key 6: partial match with queries 1, 2
        [0.5, 0.0, 0.0, 0.5],  # Key 7: partial match with queries 0, 3
    ]
)

# Simple value matrix: identity-like for easy interpretation
V = np.eye(n)[:, :d]

Step 1: Process Key Block 0Link Copied

We start by processing the first query block against the first key block. The running statistics are initialized to their "empty" states: maximum set to $-\infty$ and sum set to 0.

In[26]:

Code

# Process Query Block 0 (positions 0-3)
Q_block = Q[:4]
scale = 1.0 / np.sqrt(d)

# Initialize running statistics to "empty" state
O_block = np.zeros((4, d))  # Accumulator for weighted values
m_block = np.full(4, -np.inf)  # No maximum seen yet
ell_block = np.zeros(4)  # No exponentials accumulated yet

# Load first key/value block
K_block_0 = K[:4]
V_block_0 = V[:4]

# Compute attention scores: Q_block @ K_block^T * scale
S_block_0 = Q_block @ K_block_0.T * scale

# Update running maximum (any finite value beats -inf)
m_new_0 = np.maximum(m_block, S_block_0.max(axis=1))

# Compute exponentials relative to current maximum
P_block_0 = np.exp(S_block_0 - m_new_0[:, None])

# Accumulate weighted values and exponential sums
O_block = P_block_0 @ V_block_0
ell_block = P_block_0.sum(axis=1)
m_block = m_new_0

Out[27]:

Console

Processing Query Block 0 (positions 0-3):

Scores against Key Block 0:
[[0.5  0.   0.25 0.  ]
 [0.   0.5  0.25 0.25]
 [0.   0.   0.   0.25]
 [0.   0.   0.   0.  ]]

Running max after block 0: [0.5  0.5  0.25 0.  ]
Running sum after block 0: [2.992 3.164 3.336 4.   ]

The $4 \times 4$ score matrix shows how strongly each query attends to each key. Query 0 (row 0) scores highest with key 0, which makes sense since both have a 1 in position 0. The running maximum (0.5 for all queries) establishes our normalization baseline. The running sum (around 3.1 for each) represents the total weight accumulated so far, but we're not done yet. There's another key block to process.

Step 2: Process Key Block 1 with RescalingLink Copied

Now comes the critical step: processing the second key block while maintaining correctness. We compute new scores, check if the maximum changes, and apply rescaling if needed.

In[28]:

Code

# Load second key/value block
K_block_1 = K[4:]
V_block_1 = V[4:]

# Compute attention scores against new keys
S_block_1 = Q_block @ K_block_1.T * scale

# Check if maximum needs to update
m_new_1 = np.maximum(m_block, S_block_1.max(axis=1))

# Compute rescaling factor: e^(old_max - new_max)
# If new_max > old_max, this shrinks previous contributions
scale_factor = np.exp(m_block - m_new_1)

# Apply rescaling to previous accumulator and sum
O_block = O_block * scale_factor[:, None]
ell_block = ell_block * scale_factor

# Compute and add new block's contribution
P_block_1 = np.exp(S_block_1 - m_new_1[:, None])
O_block = O_block + P_block_1 @ V_block_1
ell_block = ell_block + P_block_1.sum(axis=1)
m_block = m_new_1

Out[29]:

Console

Scores against Key Block 1:
[[0.   0.   0.   0.25]
 [0.   0.   0.25 0.  ]
 [0.5  0.   0.25 0.  ]
 [0.   0.5  0.   0.25]]

Scale factor for rescaling: [1.    1.    0.779 0.607]
Final running max: [0.5 0.5 0.5 0.5]
Final running sum: [5.59  5.763 5.59  5.418]

The scale factors are all 1.0, meaning the maximum didn't change. The highest scores in block 1 (0.5) equal the highest scores from block 0. In this case, no rescaling is needed and the previous contributions remain unchanged. If block 1 had contained a score higher than 0.5, the scale factor would be less than 1.0, correctly shrinking the previous contributions.

The final running sum (around 6.3 for each query) represents the total weight across all 8 key positions, which is the denominator for softmax normalization.

Step 3: Final Normalization and VerificationLink Copied

After processing all key blocks, we divide the accumulated output by the running sum to apply softmax normalization. Let's compare this result against standard attention computed in one pass:

In[30]:

Code

# Final step: divide accumulated output by normalization constant
O_final = O_block / ell_block[:, None]

# For comparison: compute standard attention in one pass
full_scores = Q_block @ K.T * scale
full_weights = np.exp(full_scores - full_scores.max(axis=1, keepdims=True))
full_weights = full_weights / full_weights.sum(axis=1, keepdims=True)
O_standard = full_weights @ V

Out[31]:

Console

FlashAttention output (first 4 rows):
[[0.1789 0.1085 0.1393 0.1085]
 [0.1053 0.1735 0.1351 0.1351]
 [0.1085 0.1085 0.1085 0.1393]
 [0.1119 0.1119 0.1119 0.1119]]

Standard attention output (first 4 rows):
[[0.1789 0.1085 0.1393 0.1085]
 [0.1053 0.1735 0.1351 0.1351]
 [0.1085 0.1085 0.1085 0.1393]
 [0.1119 0.1119 0.1119 0.1119]]

Outputs match: True

The outputs match exactly. This worked example demonstrates the complete FlashAttention mechanism:

Initialize running statistics to empty state ( $m = -\infty$ , $L = 0$ , $O = 0$ )
For each key block: compute scores, update maximum, rescale previous contributions, add new contributions
Normalize the final accumulated output by dividing by the running sum

The key insight is that we never stored the full $8 \times 8$ attention matrix. Only the current $4 \times 4$ block resided in memory at any time. Yet the mathematical result is identical to standard attention.

Out[32]:

Visualization

Heatmap showing attention weights from 4 queries to 8 keys, with diagonal-like pattern in the first block. — Attention weights for the first 4 queries across all 8 key positions. Each row shows how a query distributes attention across keys. Query 0 attends most strongly to key 0, query 1 to key 1, and so on, matching the structure we designed in the Q and K matrices.

The heatmap reveals the attention patterns we computed incrementally. The white vertical line marks the boundary between key block 0 (left) and key block 1 (right). During FlashAttention, we never see both halves simultaneously. Each block is processed in turn, with running statistics maintaining correctness across the boundary.