{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Online Softmax Without the Attention Matrix\n",
        "\n",
        "Running maxima and rescaled sums recover exact attention arithmetic without retaining every score. A small CPU example makes the invariant explicit.\n",
        "\n",
        "Original CPU online-softmax illustration for one query row and scalar values; not a FlashAttention GPU kernel.\n",
        "\n",
        "The code uses Python 3 and the standard library. Saved output is from the checked CPU run.\n"
      ],
      "id": "flashattention-online-softmax--0"
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "source": [
        "from math import exp, log\n",
        "\n",
        "scores = [1000 + log(i) for i in range(1, 5)]\n",
        "values = [10.0, 20.0, 30.0, 40.0]\n",
        "maximum = max(scores)\n",
        "weights = [exp(score - maximum) for score in scores]\n",
        "dense = sum(w * v for w, v in zip(weights, values)) / sum(weights)\n",
        "\n",
        "def streaming(tile_size):\n",
        "    running_max = float(\"-inf\")\n",
        "    denominator = 0.0\n",
        "    numerator = 0.0\n",
        "    for start in range(0, len(scores), tile_size):\n",
        "        tile_scores = scores[start:start + tile_size]\n",
        "        tile_values = values[start:start + tile_size]\n",
        "        new_max = max(running_max, max(tile_scores))\n",
        "        correction = exp(running_max - new_max)\n",
        "        local_weights = [exp(s - new_max) for s in tile_scores]\n",
        "        denominator = denominator * correction + sum(local_weights)\n",
        "        numerator = numerator * correction + sum(\n",
        "            w * v for w, v in zip(local_weights, tile_values))\n",
        "        running_max = new_max\n",
        "    return numerator / denominator\n",
        "\n",
        "for tile_size in (1, 2, 4):\n",
        "    result = streaming(tile_size)\n",
        "    assert abs(result - dense) < 1e-10\n",
        "    print(f\"tile={tile_size}: output={result:.6f}\")\n",
        "assert abs(dense - 30.0) < 1e-10\n",
        "length, block = 4096, 64\n",
        "print(f\"score_entries: full_matrix={length * length}, one_row_tile={block}\")\n"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tile=1: output=30.000000\n",
            "tile=2: output=30.000000\n",
            "tile=4: output=30.000000\n",
            "score_entries: full_matrix=16777216, one_row_tile=64\n"
          ]
        }
      ],
      "id": "flashattention-online-softmax--1"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Sources\n",
        "\n",
        "- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness: https://arxiv.org/html/2205.14135v2\n"
      ],
      "id": "flashattention-online-softmax--2"
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.12.14"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}