{
  "schemaVersion": 1,
  "slug": "flash-attention",
  "title": "Compute attention one tile at a time",
  "className": "AttentionSchedules",
  "question": "Can exact attention avoid storing the entire score matrix?",
  "description": "Compare dense attention with an online-softmax reference, then replay the tiles and running statistics that produce the same result.",
  "category": "Attention systems",
  "tags": [
    "FlashAttention",
    "Tiling",
    "Online softmax",
    "Memory"
  ],
  "curator": "TensorViz",
  "attribution": "Original TensorViz teaching example. PyTorch provides the underlying operators.",
  "license": "No separate redistribution license has been declared for these project examples.",
  "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.",
  "openingStep": "tiles",
  "provenance": {
    "capturedAt": "2026-09-17T06:12:56.641732+00:00",
    "productRevision": "cf0ab1c9c58c1dc640235b742fa032cade00b43c",
    "sourceSha256": "d084ad3ac3b5fb59a7807c26bade601b86151badc89c965d3889d1217a73ef10",
    "python": "3.13.13",
    "torch": "2.7.1",
    "execution": {
      "inputs": {
        "args": [
          {
            "shape": [
              1,
              6,
              4
            ],
            "dtype": "float32",
            "fill": "randn"
          },
          {
            "shape": [
              1,
              6,
              4
            ],
            "dtype": "float32",
            "fill": "randn"
          },
          {
            "shape": [
              1,
              6,
              4
            ],
            "dtype": "float32",
            "fill": "randn"
          }
        ]
      },
      "device": "cpu",
      "mode": "eval",
      "seed": 0,
      "deterministic": true
    },
    "outcome": "passed",
    "outputShape": [
      2,
      1,
      6,
      4
    ],
    "method": "TensorViz multifile projection and isolated execution worker",
    "scope": "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.",
    "numericalChecks": [
      "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"
    ]
  },
  "layers": [
    {
      "id": "input",
      "label": "input",
      "op": "Input",
      "kind": "input",
      "parent": null,
      "params": {
        "shape": "1,6,4"
      }
    },
    {
      "id": "dense",
      "label": "dense",
      "op": "Module",
      "kind": "module",
      "parent": null,
      "params": {},
      "source": {
        "line": 54,
        "endLine": 54
      }
    },
    {
      "id": "tiled",
      "label": "tiled",
      "op": "Module",
      "kind": "module",
      "parent": null,
      "params": {},
      "source": {
        "line": 55,
        "endLine": 55
      }
    },
    {
      "id": "stack@AttentionSchedules#0",
      "label": "stack@AttentionSchedules#0",
      "op": "torch.stack",
      "kind": "layer",
      "parent": null,
      "params": {},
      "source": {
        "line": 58,
        "endLine": 58
      }
    },
    {
      "id": "output",
      "label": "output",
      "op": "Output",
      "kind": "output",
      "parent": null,
      "params": {}
    },
    {
      "id": "dense.softmax",
      "label": "dense.softmax",
      "op": "Softmax",
      "kind": "layer",
      "parent": "dense",
      "params": {},
      "source": {
        "line": 15,
        "endLine": 15
      }
    },
    {
      "id": "python@tiled#2",
      "label": "python@tiled#2",
      "op": "Python",
      "kind": "layer",
      "parent": "tiled",
      "params": {},
      "source": {
        "line": 34,
        "endLine": 34
      }
    },
    {
      "id": "python@tiled#0",
      "label": "python@tiled#0",
      "op": "Python",
      "kind": "layer",
      "parent": "tiled",
      "params": {},
      "source": {
        "line": 32,
        "endLine": 32
      }
    },
    {
      "id": "python@tiled#3",
      "label": "python@tiled#3",
      "op": "Python",
      "kind": "layer",
      "parent": "tiled",
      "params": {},
      "source": {
        "line": 39,
        "endLine": 39
      }
    },
    {
      "id": "sub@tiled#1",
      "label": "sub@tiled#1",
      "op": "torch.sub",
      "kind": "layer",
      "parent": "tiled",
      "params": {},
      "source": {
        "line": 41,
        "endLine": 41
      }
    },
    {
      "id": "python@tiled#5",
      "label": "python@tiled#5",
      "op": "Python",
      "kind": "layer",
      "parent": "tiled",
      "params": {},
      "source": {
        "line": 41,
        "endLine": 41
      }
    },
    {
      "id": "sub@tiled#0",
      "label": "sub@tiled#0",
      "op": "torch.sub",
      "kind": "layer",
      "parent": "tiled",
      "params": {},
      "source": {
        "line": 40,
        "endLine": 40
      }
    },
    {
      "id": "python@tiled#4",
      "label": "python@tiled#4",
      "op": "Python",
      "kind": "layer",
      "parent": "tiled",
      "params": {},
      "source": {
        "line": 40,
        "endLine": 40
      }
    },
    {
      "id": "mul@tiled#1",
      "label": "mul@tiled#1",
      "op": "torch.mul",
      "kind": "layer",
      "parent": "tiled",
      "params": {},
      "source": {
        "line": 45,
        "endLine": 45
      }
    },
    {
      "id": "add@tiled#1",
      "label": "add@tiled#1",
      "op": "torch.add",
      "kind": "layer",
      "parent": "tiled",
      "params": {},
      "source": {
        "line": 45,
        "endLine": 45
      }
    },
    {
      "id": "python@tiled#1",
      "label": "python@tiled#1",
      "op": "Python",
      "kind": "layer",
      "parent": "tiled",
      "params": {},
      "source": {
        "line": 33,
        "endLine": 33
      }
    },
    {
      "id": "mul@tiled#0",
      "label": "mul@tiled#0",
      "op": "torch.mul",
      "kind": "layer",
      "parent": "tiled",
      "params": {},
      "source": {
        "line": 42,
        "endLine": 44
      }
    },
    {
      "id": "add@tiled#0",
      "label": "add@tiled#0",
      "op": "torch.add",
      "kind": "layer",
      "parent": "tiled",
      "params": {},
      "source": {
        "line": 42,
        "endLine": 44
      }
    },
    {
      "id": "div@tiled#0",
      "label": "div@tiled#0",
      "op": "torch.div",
      "kind": "layer",
      "parent": "tiled",
      "params": {},
      "source": {
        "line": 47,
        "endLine": 47
      }
    },
    {
      "id": "python@tiled#6",
      "label": "python@tiled#6",
      "op": "Python",
      "kind": "layer",
      "parent": "tiled",
      "params": {},
      "source": {
        "line": 47,
        "endLine": 47
      }
    }
  ],
  "connections": [
    {
      "id": "edge-1",
      "source": "python@tiled#0",
      "target": "python@tiled#1",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-2",
      "source": "python@tiled#0",
      "target": "python@tiled#3",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-3",
      "source": "python@tiled#0",
      "target": "sub@tiled#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-4",
      "source": "python@tiled#3",
      "target": "sub@tiled#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-5",
      "source": "sub@tiled#0",
      "target": "python@tiled#4",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-6",
      "source": "python@tiled#3",
      "target": "sub@tiled#1",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-7",
      "source": "sub@tiled#1",
      "target": "python@tiled#5",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-8",
      "source": "python@tiled#1",
      "target": "mul@tiled#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-9",
      "source": "python@tiled#4",
      "target": "mul@tiled#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-10",
      "source": "mul@tiled#0",
      "target": "add@tiled#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-11",
      "source": "python@tiled#5",
      "target": "add@tiled#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-12",
      "source": "python@tiled#2",
      "target": "mul@tiled#1",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-13",
      "source": "python@tiled#4",
      "target": "mul@tiled#1",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-14",
      "source": "mul@tiled#1",
      "target": "add@tiled#1",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-15",
      "source": "python@tiled#5",
      "target": "add@tiled#1",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-16",
      "source": "add@tiled#1",
      "target": "div@tiled#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-17",
      "source": "add@tiled#0",
      "target": "div@tiled#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-18",
      "source": "div@tiled#0",
      "target": "python@tiled#6",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-19",
      "source": "input",
      "target": "dense.softmax",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-20",
      "source": "input",
      "target": "dense",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-21",
      "source": "input",
      "target": "python@tiled#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-22",
      "source": "input",
      "target": "python@tiled#2",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-23",
      "source": "input",
      "target": "python@tiled#3",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-24",
      "source": "input",
      "target": "sub@tiled#1",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-25",
      "source": "input",
      "target": "add@tiled#1",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-26",
      "source": "input",
      "target": "tiled",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-27",
      "source": "dense.softmax",
      "target": "stack@AttentionSchedules#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-28",
      "source": "dense",
      "target": "stack@AttentionSchedules#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-29",
      "source": "python@tiled#6",
      "target": "stack@AttentionSchedules#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-30",
      "source": "tiled",
      "target": "stack@AttentionSchedules#0",
      "sourcePort": null,
      "targetPort": null
    },
    {
      "id": "edge-31",
      "source": "stack@AttentionSchedules#0",
      "target": "output",
      "sourcePort": null,
      "targetPort": null
    }
  ],
  "steps": [
    {
      "id": "dense",
      "target": "dense",
      "title": "Start with ordinary exact attention",
      "note": "Q, K and V each have six rows and four features. Dense attention materializes all 36 scores and normalizes across all six keys per query."
    },
    {
      "id": "coupled",
      "target": "dense.softmax",
      "title": "Notice what couples the tiles",
      "note": "A row's softmax denominator needs contributions from every key. Normalizing each tile independently and summing would give the wrong answer. The tiled route must retain running normalization statistics."
    },
    {
      "id": "tiles",
      "target": "tiled",
      "title": "Update the running maximum and denominator",
      "note": "The nested loop computes a small score tile and rescales earlier contributions whenever the running maximum increases. It retains a denominator and weighted numerator for each current query. Open the source to follow the exact update."
    },
    {
      "id": "equivalent",
      "target": "output",
      "title": "Compare the final outputs",
      "note": "The output stacks dense and tiled results, each [1, 6, 4]. Numerical checks compare multiple tile sizes, uneven final tiles, large logits and gradients. The recorded schedule below shows the partial answer becoming complete."
    }
  ],
  "milestone": {
    "name": "FlashAttention",
    "year": 2022,
    "kind": "Inference method",
    "fidelity": "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.",
    "papers": [
      {
        "title": "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness",
        "url": "https://arxiv.org/abs/2205.14135"
      }
    ]
  },
  "experiment": {
    "kind": "schedule",
    "title": "Finish attention without retaining every score.",
    "description": "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.",
    "controlLabel": "Tile update",
    "cases": [
      {
        "id": "tile-1",
        "label": "Query rows 0–1 · key columns 0–1",
        "target": "tiled",
        "note": "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.",
        "schedule": {
          "rows": 6,
          "columns": 6,
          "rowStart": 0,
          "rowCount": 2,
          "columnStart": 0,
          "columnCount": 2,
          "completedColumns": 0
        },
        "vectors": [
          {
            "label": "Current score tile · row-major",
            "values": [
              0.5993776922264088,
              -1.1169048401112192,
              -1.9624242954894378,
              1.2289214186836044
            ]
          },
          {
            "label": "Running row maximum",
            "values": [
              0.5993776922264088,
              1.2289214186836044
            ]
          },
          {
            "label": "Running denominator",
            "values": [
              1.1797330593686322,
              1.0411165026321278
            ]
          },
          {
            "label": "Partial output · query 0",
            "values": [
              0.653822080394852,
              1.9052249031129327,
              0.9410291873451526,
              -0.03712522410369098
            ]
          },
          {
            "label": "Final dense output · query 0",
            "values": [
              -0.0674336192777835,
              0.3616644543013671,
              -0.032788449072674745,
              0.12249373372748368
            ]
          }
        ],
        "metrics": [
          {
            "label": "Score elements in current tile",
            "value": 4
          },
          {
            "label": "Elements in a full score matrix",
            "value": 36
          },
          {
            "label": "Max difference from final output · rows 0–1",
            "value": 1.5435604488115655
          }
        ]
      },
      {
        "id": "tile-2",
        "label": "Query rows 0–1 · key columns 2–3",
        "target": "tiled",
        "note": "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.",
        "schedule": {
          "rows": 6,
          "columns": 6,
          "rowStart": 0,
          "rowCount": 2,
          "columnStart": 2,
          "columnCount": 2,
          "completedColumns": 2
        },
        "vectors": [
          {
            "label": "Current score tile · row-major",
            "values": [
              0.29103417274308757,
              0.19164018146104933,
              -0.3675986147612462,
              0.14482992940218867
            ]
          },
          {
            "label": "Running row maximum",
            "values": [
              0.5993776922264088,
              1.2289214186836044
            ]
          },
          {
            "label": "Running denominator",
            "values": [
              2.579549414534565,
              1.5819257504585613
            ]
          },
          {
            "label": "Partial output · query 0",
            "values": [
              0.9386033502738842,
              1.2957160835682766,
              0.4384212344431347,
              -0.13421498495469766
            ]
          },
          {
            "label": "Final dense output · query 0",
            "values": [
              -0.0674336192777835,
              0.3616644543013671,
              -0.032788449072674745,
              0.12249373372748368
            ]
          }
        ],
        "metrics": [
          {
            "label": "Score elements in current tile",
            "value": 4
          },
          {
            "label": "Elements in a full score matrix",
            "value": 36
          },
          {
            "label": "Max difference from final output · rows 0–1",
            "value": 1.0060369695516678
          }
        ]
      },
      {
        "id": "tile-3",
        "label": "Query rows 0–1 · key columns 4–5",
        "target": "tiled",
        "note": "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.",
        "schedule": {
          "rows": 6,
          "columns": 6,
          "rowStart": 0,
          "rowCount": 2,
          "columnStart": 4,
          "columnCount": 2,
          "completedColumns": 4
        },
        "vectors": [
          {
            "label": "Current score tile · row-major",
            "values": [
              0.5578407733358283,
              1.0498077316430776,
              -2.1762740656526693,
              -0.16504958786188773
            ]
          },
          {
            "label": "Running row maximum",
            "values": [
              1.0498077316430776,
              1.2289214186836044
            ]
          },
          {
            "label": "Running denominator",
            "values": [
              3.255508719247094,
              1.8632142659525093
            ]
          },
          {
            "label": "Partial output · query 0",
            "values": [
              -0.06743361927778342,
              0.3616644543013672,
              -0.03278844907267477,
              0.12249373372748366
            ]
          },
          {
            "label": "Final dense output · query 0",
            "values": [
              -0.0674336192777835,
              0.3616644543013671,
              -0.032788449072674745,
              0.12249373372748368
            ]
          }
        ],
        "metrics": [
          {
            "label": "Score elements in current tile",
            "value": 4
          },
          {
            "label": "Elements in a full score matrix",
            "value": 36
          },
          {
            "label": "Max difference from final output · rows 0–1",
            "value": 2.220446049250313e-16
          }
        ]
      }
    ]
  },
  "revision": "f7f9d18206e7873a",
  "files": {
    "graph.tensorviz.json": {
      "url": "/models/flash-attention/f7f9d18206e7873a/graph.tensorviz.json",
      "sha256": "7d2100a83756b37273d26d23b1b4b029ee8f73396aa9ba3bf4921271a2316ed5",
      "bytes": 56985
    },
    "model.py": {
      "url": "/models/flash-attention/f7f9d18206e7873a/model.py",
      "sha256": "d084ad3ac3b5fb59a7807c26bade601b86151badc89c965d3889d1217a73ef10",
      "bytes": 6991
    },
    "run.py": {
      "url": "/models/flash-attention/f7f9d18206e7873a/run.py",
      "sha256": "f45be815c649175c5886ae9d50b57db7f22b0c50ad1ef271d07779731ff00541",
      "bytes": 812
    },
    "requirements.txt": {
      "url": "/models/flash-attention/f7f9d18206e7873a/requirements.txt",
      "sha256": "edb86016b42cb3cc6cdc0c490aab929d3bb098ced4515394f0b55766033e0daf",
      "bytes": 91
    },
    "preview.svg": {
      "url": "/models/flash-attention/f7f9d18206e7873a/preview.svg",
      "sha256": "c0b64a2d91987925ce6c9dd996f1d2791dd37b5c8c05106219aa16b42e359044",
      "bytes": 1456
    }
  }
}