MODEL EXPLORER / ATTENTION SYSTEMS
Compute attention one tile at a time
Can exact attention avoid storing the entire score matrix?
RECORDED NUMERICAL EXAMPLE
Finish attention without retaining every score.
Replay one query block through three key/value tiles. A fused GPU implementation keeps tiles in SRAM and avoids full score/probability writes to HBM. This unfused CPU reference verifies the online-softmax arithmetic; it makes no GPU memory or speed measurement.
The current score tile is 2 × 2. Earlier tiles are discarded after contributing to the running maximum, denominator and weighted numerator. The partial output uses only the keys seen so far; it becomes the full attention output after the third tile. The diagram is a schedule, not a stored six-by-six score tensor.
Current query rows 0–1, key columns 0–1. 0 earlier key columns have already contributed to this query block.
Current tile Already accumulated Other tiles
Earlier score tiles are discarded. Only their running statistics and weighted output contribution remain.
Bars share one scale within this example. Values are rounded to four significant digits.
- Score elements in current tile
- 4
- Elements in a full score matrix
- 36
- Max difference from final output · rows 0–1
- 1.544
Verified by the local recipe
- Tiled attention matches dense attention across five tile sizes and uneven final tiles
- Running-max correction stays finite and correct for large logits
- CPU autograd gradients agree with dense attention; this is not the paper's memory-efficient backward
Download the source and run.py to reproduce these checks. Choosing an example here replays recorded values.
THE SAME MODEL, ON YOUR MACHINE
Pick up where
the graph leaves off.
Compare dense attention with an online-softmax reference, then replay the tiles and running statistics that produce the same result.
Use in VS Code ↗1. Save these files in one folder
2. Reproduce the example
Use Python 3.13 in a dedicated environment, matching the tested recipe. These commands are for macOS / Linux shells.
python3.13 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt
python run.pyExpected output: [2,1,6,4]
Input 1: [1,6,4] · float32 randn
Input 2: [1,6,4] · float32 randn
Input 3: [1,6,4] · float32 randn · CPU / eval · seed 0
3. Open the source in TensorViz
Install the preview, open this folder in local VS Code, and choose Open as Graph above AttentionSchedules in model.py. Select your Python environment. Save before choosing Validate.
Installation and supported setup ↗
The recipe was run with Python 3.13.13 and Torch 2.7.1. This is an untrained example; a successful shape check does not measure accuracy.
ABOUT THIS EXAMPLE
Compare dense attention with an online-softmax reference, then replay the tiles and running statistics that produce the same result.
Inference method · 2022
FlashAttention
Six-token, width-four online-softmax teaching reference. Query-major loop with two-by-two tiles; the original paper's Algorithm 1 uses a key-major loop and a fused GPU kernel.
Source, capture & limitations +
The CPU reference teaches the forward algorithm. Its unfused PyTorch operations and generic autograd do not reproduce FlashAttention's GPU memory behavior or efficient backward. Real FlashAttention also benefits training, not just inference. The numerical panel counts score elements, not total allocated memory.
Original TensorViz teaching example. PyTorch provides the underlying operators.
No separate redistribution license has been declared for these project examples.
- Content revision
f7f9d18206e7873a- Source SHA-256
d084ad3ac3b5fb59a7807c26bade601b86151badc89c965d3889d1217a73ef10- Captured
- 2026-09-17 · Python 3.13.13 / Torch 2.7.1
Unmasked CPU dense/tiled forwards and autograd arithmetic checks across tile sizes. No fused kernel, memory-efficient custom backward, GPU benchmark or measured memory traffic.
Read the model manifest ↗