When you need this. A document assistant needs longer context, and the attention score matrix grows much faster than its input. Implement an online softmax that processes tiles without retaining the full matrix, then compare its outputs with a direct reference.
What you run. CPU attention implementation and logical-access accounting, compared with a direct numerical reference. No FlashAttention GPU benchmark or measured HBM traffic.
This lesson builds on implement paged cache allocation and copy-on-write.
Keep the softmax state between tiles
Attention weights come from scaled query-key scores followed by softmax. The direct reference retains the scores and probabilities before multiplying by values. The tiled algorithm keeps a running maximum, normalization sum, and output accumulator for each query row.
When a later tile has a larger maximum, rescale both the old normalization sum and the old output accumulator. Otherwise the tiles use incompatible exponential scales. This preserves the same attention expression up to floating-point rounding. It does not remove the attention computation’s quadratic work.
tiled_flash_attention() is a readable CPU implementation of the online-softmax mechanism. It checks causal masking and irregular tile sizes against standard_attention(). The access calculator counts repeated K/V reads in its chosen loop order; it no longer claims that every tensor is read once from device memory.
Compare the attention outputs
Download the lab from this page, extract it, and run these commands from its directory. The default command reads saved run records or runs the explicitly described local algorithm; it never downloads a model or calls a paid API.
python3 lab.py
python3 -m unittest -v test_lab.pyRead lab.py
"""Compare dense and tiled CPU attention with online softmax; device traffic is not measured."""
from __future__ import annotations
import json
import math
import random
def make_qkv(
seq_len: int,
head_dim: int,
seed: int = 42,
) -> tuple[list[list[float]], list[list[float]], list[list[float]]]:
rng = random.Random(seed)
q = [[rng.uniform(-0.8, 0.8) for _ in range(head_dim)] for _ in range(seq_len)]
k = [[rng.uniform(-0.8, 0.8) for _ in range(head_dim)] for _ in range(seq_len)]
v = [[rng.uniform(-0.8, 0.8) for _ in range(head_dim)] for _ in range(seq_len)]
return q, k, v
def dot(a: list[float], b: list[float]) -> float:
return sum(x * y for x, y in zip(a, b))
def standard_attention(
q: list[list[float]],
k: list[list[float]],
v: list[list[float]],
causal: bool = True,
subtract_max: bool = True,
) -> tuple[list[list[float]], dict[str, int]]:
"""Materializes the full N x N score and probability matrices before multiplying by V."""
n = len(q)
d = len(q[0])
scale = 1.0 / math.sqrt(d)
# 1. Compute S = Q K^T * scale (materialized N x N in HBM)
scores: list[list[float]] = []
for i in range(n):
limit = (i + 1) if causal else n
row_scores = [dot(q[i], k[j]) * scale for j in range(limit)]
scores.append(row_scores)
# 2. Row-wise softmax P = softmax(S) (materialized N x N in HBM)
probs: list[list[float]] = []
for row_scores in scores:
m = max(row_scores) if subtract_max else 0.0
exps = [math.exp(s - m) for s in row_scores]
denom = sum(exps)
probs.append([e / denom for e in exps])
# 3. Output O = P V
out: list[list[float]] = []
for i in range(n):
row_out = [0.0] * d
for j, p_ij in enumerate(probs[i]):
for col in range(d):
row_out[col] += p_ij * v[j][col]
out.append(row_out)
materialized_nxn_elements = sum(len(r) for r in scores) + sum(len(r) for r in probs)
return out, {"materialized_nxn_elements": materialized_nxn_elements}
def tiled_flash_attention(
q: list[list[float]],
k: list[list[float]],
v: list[list[float]],
block_q: int = 16,
block_kv: int = 16,
causal: bool = True,
rescale_previous_acc: bool = True,
) -> tuple[list[list[float]], dict[str, int]]:
"""Tiles Q, K, and V in SRAM using online softmax without materializing N x N in HBM."""
n = len(q)
d = len(q[0])
scale = 1.0 / math.sqrt(d)
out: list[list[float]] = [[0.0] * d for _ in range(n)]
sram_tile_score_peak = block_q * block_kv
tiles_visited = 0
for q_start in range(0, n, block_q):
q_end = min(q_start + block_q, n)
q_rows = q_end - q_start
running_max = [-math.inf] * q_rows
running_sum = [0.0] * q_rows
output_acc = [[0.0] * d for _ in range(q_rows)]
for kv_start in range(0, n, block_kv):
if causal and kv_start >= q_end:
break
kv_end = min(kv_start + block_kv, n)
tiles_visited += 1
for local_i in range(q_rows):
global_i = q_start + local_i
valid_kv_indices = [
j for j in range(kv_start, kv_end) if (not causal or j <= global_i)
]
if not valid_kv_indices:
continue
tile_scores = [dot(q[global_i], k[j]) * scale for j in valid_kv_indices]
tile_max = max(tile_scores)
prev_max = running_max[local_i]
new_max = max(prev_max, tile_max)
# Rescale factor for previous running_sum and output_acc
alpha = math.exp(prev_max - new_max) if prev_max > -math.inf else 0.0
tile_exps = [math.exp(s - new_max) for s in tile_scores]
tile_sum = sum(tile_exps)
running_sum[local_i] = running_sum[local_i] * alpha + tile_sum
acc_scale = alpha if rescale_previous_acc else 1.0
for col in range(d):
acc_val = output_acc[local_i][col] * acc_scale
for exp_val, j in zip(tile_exps, valid_kv_indices):
acc_val += exp_val * v[j][col]
output_acc[local_i][col] = acc_val
running_max[local_i] = new_max
for local_i in range(q_rows):
global_i = q_start + local_i
inv_sum = 1.0 / running_sum[local_i]
out[global_i] = [val * inv_sum for val in output_acc[local_i]]
return out, {
"materialized_nxn_elements": 0,
"max_sram_score_tile_elements": sram_tile_score_peak,
"tiles_visited": tiles_visited,
}
def max_abs_diff(a: list[list[float]], b: list[list[float]]) -> float:
return max(abs(x - y) for row_a, row_b in zip(a, b) for x, y in zip(row_a, row_b))
def estimate_hbm_bytes(
seq_len: int,
num_heads: int = 32,
head_dim: int = 128,
bytes_per_elem: int = 2,
block_q: int = 128,
block_kv: int = 128,
) -> dict[str, object]:
"""Compare single-layer multi-head attention HBM traffic for standard vs FlashAttention."""
qkv_o_bytes = 4 * num_heads * seq_len * head_dim * bytes_per_elem
# Standard attention writes S (N x N), reads S, writes P (N x N), reads P across all heads
attention_matrix_rw_bytes = 4 * num_heads * (seq_len * seq_len) * bytes_per_elem
standard_total_bytes = qkv_o_bytes + attention_matrix_rw_bytes
# Logical tensor accesses for causal outer-Q tiling, assuming no inter-block cache reuse.
# Hardware HBM traffic requires profiling; this is an access-count model.
num_q_blocks = math.ceil(seq_len / block_q)
kv_rows_read = sum(
min(seq_len, start + block_q) for start in range(0, seq_len, block_q)
)
tiled_bytes = (
2 * num_heads * seq_len * head_dim + 2 * num_heads * kv_rows_read * head_dim
) * bytes_per_elem
return {
"seq_len": seq_len,
"num_heads": num_heads,
"head_dim": head_dim,
"standard_qkv_o_mb": round(qkv_o_bytes / 1e6, 2),
"standard_nxn_matrix_hbm_mb": round(attention_matrix_rw_bytes / 1e6, 2),
"standard_total_hbm_mb": round(standard_total_bytes / 1e6, 2),
"tiled_logical_access_mb_no_cache_reuse": round(tiled_bytes / 1e6, 2),
"qkvo_compulsory_bytes_mb": round(qkv_o_bytes / 1e6, 2),
"num_q_blocks": num_q_blocks,
"accounting": "logical accesses, not measured HBM traffic; dense baseline and causal tiled loop differ in work",
}
def build_report() -> dict[str, object]:
q, k, v = make_qkv(seq_len=64, head_dim=16, seed=42)
std_out, std_meta = standard_attention(q, k, v, causal=True)
flash_out, flash_meta = tiled_flash_attention(
q, k, v, block_q=16, block_kv=16, causal=True
)
diff = max_abs_diff(std_out, flash_out)
hbm_sweep = [
estimate_hbm_bytes(seq_len=n, num_heads=32, head_dim=128, bytes_per_elem=2)
for n in (1024, 4096, 16384, 32768)
]
return {
"equivalence_check": {
"seq_len": 64,
"head_dim": 16,
"block_q": 16,
"block_kv": 16,
"max_abs_diff": float(f"{diff:.3e}"),
"standard_materialized_elements": std_meta["materialized_nxn_elements"],
"flash_materialized_hbm_elements": flash_meta["materialized_nxn_elements"],
"flash_sram_tile_elements": flash_meta["max_sram_score_tile_elements"],
"flash_tiles_visited": flash_meta["tiles_visited"],
"sample_row0_first4": [round(x, 6) for x in flash_out[0][:4]],
"sample_last_row_first4": [round(x, 6) for x in flash_out[-1][:4]],
},
"hbm_traffic_sweep": hbm_sweep,
}
def main() -> None:
print(
json.dumps(
{
"execution_mode": "CPU algorithm with controlled inputs; any latency/capacity coefficients are assumptions",
**build_report(),
},
indent=2,
)
)
if __name__ == "__main__":
main()The output comparison checks the attention calculation. The logical access report is a separate analysis of this loop order; it does not measure GPU memory traffic. The earlier linear-traffic and fixed-speedup claims have been removed.
| Quantity | Recorded value |
|---|---|
| Maximum numerical difference | 1.388e-16 |
| Reference retained score/probability elements | 4160 |
| Tiled retained full-matrix elements | 0 |
| Tiles visited | 10 |
Inspect the complete local result
{
"execution_mode": "CPU algorithm with controlled inputs; any latency/capacity coefficients are assumptions",
"equivalence_check": {
"seq_len": 64,
"head_dim": 16,
"block_q": 16,
"block_kv": 16,
"max_abs_diff": 1.388e-16,
"standard_materialized_elements": 4160,
"flash_materialized_hbm_elements": 0,
"flash_sram_tile_elements": 256,
"flash_tiles_visited": 10,
"sample_row0_first4": [
-0.267257,
0.731232,
-0.775466,
0.699456
],
"sample_last_row_first4": [
0.017991,
-0.007947,
0.059253,
0.070964
]
},
"hbm_traffic_sweep": [
{
"seq_len": 1024,
"num_heads": 32,
"head_dim": 128,
"standard_qkv_o_mb": 33.55,
"standard_nxn_matrix_hbm_mb": 268.44,
"standard_total_hbm_mb": 301.99,
"tiled_logical_access_mb_no_cache_reuse": 92.27,
"qkvo_compulsory_bytes_mb": 33.55,
"num_q_blocks": 8,
"accounting": "logical accesses, not measured HBM traffic; dense baseline and causal tiled loop differ in work"
},
{
"seq_len": 4096,
"num_heads": 32,
"head_dim": 128,
"standard_qkv_o_mb": 134.22,
"standard_nxn_matrix_hbm_mb": 4294.97,
"standard_total_hbm_mb": 4429.19,
"tiled_logical_access_mb_no_cache_reuse": 1174.41,
"qkvo_compulsory_bytes_mb": 134.22,
"num_q_blocks": 32,
"accounting": "logical accesses, not measured HBM traffic; dense baseline and causal tiled loop differ in work"
},
{
"seq_len": 16384,
"num_heads": 32,
"head_dim": 128,
"standard_qkv_o_mb": 536.87,
"standard_nxn_matrix_hbm_mb": 68719.48,
"standard_total_hbm_mb": 69256.35,
"tiled_logical_access_mb_no_cache_reuse": 17582.52,
"qkvo_compulsory_bytes_mb": 536.87,
"num_q_blocks": 128,
"accounting": "logical accesses, not measured HBM traffic; dense baseline and causal tiled loop differ in work"
},
{
"seq_len": 32768,
"num_heads": 32,
"head_dim": 128,
"standard_qkv_o_mb": 1073.74,
"standard_nxn_matrix_hbm_mb": 274877.91,
"standard_total_hbm_mb": 275951.65,
"tiled_logical_access_mb_no_cache_reuse": 69524.78,
"qkvo_compulsory_bytes_mb": 1073.74,
"num_q_blocks": 256,
"accounting": "logical accesses, not measured HBM traffic; dense baseline and causal tiled loop differ in work"
}
]
}
This image is an application-generated report of the saved program result, not a new inference run. Open the full-resolution report or read the raw result.
Separate storage savings from runtime
Tiling reduces intermediate storage and can improve data movement on a supported GPU kernel, but tile size, register pressure, cache reuse, precision, and backend support matter. A compulsory-byte lower bound is not actual HBM traffic. Measure the selected runtime backend before claiming a FlashAttention speedup.
Exponential overflow. Large scores can overflow a naive exponential. Subtracting the row maximum and maintaining the online scale keeps the reference calculation finite.
Accumulator rescaling. When the running maximum changes, leaving the old output accumulator unscaled changes the attention output; the comparison catches that error.
Read the executable checks
"""Unit and failure-boundary tests for standard vs tiled FlashAttention."""
from __future__ import annotations
import unittest
from lab import (
estimate_hbm_bytes,
make_qkv,
max_abs_diff,
standard_attention,
tiled_flash_attention,
)
class FlashAttentionTest(unittest.TestCase):
def test_tiled_flash_attention_matches_standard_attention_across_tile_sizes(
self,
) -> None:
q, k, v = make_qkv(seq_len=64, head_dim=16, seed=7)
std_out, _ = standard_attention(q, k, v, causal=True)
for bq, bkv in ((8, 8), (16, 16), (13, 19), (32, 8)):
flash_out, meta = tiled_flash_attention(
q, k, v, block_q=bq, block_kv=bkv, causal=True
)
self.assertLess(max_abs_diff(std_out, flash_out), 1e-12)
self.assertEqual(meta["materialized_nxn_elements"], 0)
def test_access_count_includes_repeated_kv_reads(self) -> None:
small = estimate_hbm_bytes(seq_len=4096)
large = estimate_hbm_bytes(seq_len=32768)
self.assertGreater(
small["tiled_logical_access_mb_no_cache_reuse"],
small["qkvo_compulsory_bytes_mb"],
)
self.assertGreater(
large["tiled_logical_access_mb_no_cache_reuse"]
/ small["tiled_logical_access_mb_no_cache_reuse"],
8,
)
self.assertNotIn("hbm_traffic_reduction_x", large)
def test_naive_softmax_overflows_without_max_shift(self) -> None:
"""Large attention logits overflow exp() when subtract_max=False, whereas online max shift stays finite."""
q = [[500.0, 500.0], [500.0, 500.0]]
k = [[1.0, 1.0], [2.0, 2.0]]
v = [[1.0, 0.0], [0.0, 1.0]]
with self.assertRaises(OverflowError):
standard_attention(q, k, v, causal=False, subtract_max=False)
safe_std, _ = standard_attention(q, k, v, causal=False, subtract_max=True)
safe_flash, _ = tiled_flash_attention(
q, k, v, block_q=1, block_kv=1, causal=False
)
self.assertLess(max_abs_diff(safe_std, safe_flash), 1e-12)
def test_forgetting_output_rescaling_across_tiles_corrupts_attention(self) -> None:
"""Omitting the exp(prev_max - new_max) rescaling on output_acc produces wrong attention outputs."""
q, k, v = make_qkv(seq_len=32, head_dim=8, seed=99)
# Emphasize a late key so running_max increases on a subsequent KV tile
k[-1] = [4.0] * 8
std_out, _ = standard_attention(q, k, v, causal=True)
correct_flash, _ = tiled_flash_attention(
q, k, v, block_q=8, block_kv=8, causal=True, rescale_previous_acc=True
)
buggy_flash, _ = tiled_flash_attention(
q, k, v, block_q=8, block_kv=8, causal=True, rescale_previous_acc=False
)
self.assertLess(max_abs_diff(std_out, correct_flash), 1e-12)
self.assertGreater(max_abs_diff(std_out, buggy_flash), 0.1)
if __name__ == "__main__":
unittest.main()Try a change and explain the result
Compare several tile sizes, including sizes that do not divide the sequence. Explain why the numerical result should agree even when the visitation order changes. Use the rescaling test to locate the first wrong assumption in an incorrect accumulator.
Read the explanation
The running state must use the same exponential scale as the new tile.
Select a supported attention backend
Prefer a supported attention backend in the serving runtime. Confirm mask semantics and numerical tolerance for the exact dtype. Do not identify a generic scaled-dot-product-attention call as FlashAttention unless the selected backend is known.
References
Sources checked October 24, 2026. Runtime-specific commands and results apply to the recorded versions; linked rolling documentation can change. The application scenarios are teaching examples, not claims about an author’s production incident.