How can attention store less?

Idea 1 of 3

Follow one weighted result

Start with query zero. A high score gives a key more weight, but the output still combines its value with the others. The table shows how the score row becomes weights. Each output component adds those weights multiplied by the corresponding component of every value row; the four-component result is shown beside it.

Run, then step from scores to weights to the result. Increasing the query multiplier multiplies all scores in this example. It makes the largest-score keys more influential; tied largest scores can share the weight.

import torch

n, d = 8, 4
q = torch.tensor([[(i + j) % 5 - 2 for j in range(d)]
                  for i in range(n)], dtype=torch.float32, device="cuda") * 1
k = torch.tensor([[((2 * i + j) % 7 - 3) / 2 for j in range(d)]
                  for i in range(n)], dtype=torch.float32, device="cuda")
v = torch.tensor([[(i + j) % 9 - 4 for j in range(d)]
                  for i in range(n)], dtype=torch.float32, device="cuda")
scores = q @ k.T / d**0.5
weights = torch.softmax(scores, dim=-1)
out = weights @ v
print(out[0])

Softmax is evaluated stably by subtracting the largest score before exponentiating. This preserves the normalised weights and avoids exponentiating a large positive score. Larger scores do not make the weighted result grow by the same multiplier.

Idea 2 of 3

Count the matrices between input and output

The straightforward expression stores an N × N score matrix and an N × N weight matrix. For each matrix this accounting counts one write and one later read. It assumes four-byte FP32 storage, even though browser arithmetic uses JavaScript doubles.

The two matrices therefore account for 4N² value transfers × 4 bytes = 16N² bytes. These are selected intermediate reads and writes, excluding Q, K, V, output, caching and all other traffic. They are not measured DRAM transfers.

import torch

n, d = 8, 4
q = torch.tensor([[(i + j) % 5 - 2 for j in range(d)]
                  for i in range(n)], dtype=torch.float32, device="cuda") * 1
k = torch.tensor([[((2 * i + j) % 7 - 3) / 2 for j in range(d)]
                  for i in range(n)], dtype=torch.float32, device="cuda")
v = torch.tensor([[(i + j) % 9 - 4 for j in range(d)]
                  for i in range(n)], dtype=torch.float32, device="cuda")
scores = q @ k.T / d**0.5
weights = torch.softmax(scores, dim=-1)
out = weights @ v
print(out[0])

Idea 3 of 3

Combine one key tile at a time

We can visit a few keys at a time, combine their contribution, then move to the next group. Keep a running weighted total and a normalisation total instead of writing full score and weight matrices. When a larger score arrives, earlier weights need a new common scale. Multiply both previous totals by the same factor before adding the new tile’s contributions. Their ratio is preserved. The result below shows these additions for output component zero.

The current output estimate uses only the keys visited so far. It need not equal the final answer until every key has been included. The final result below compares all queries against the full-matrix reference.

This experiment fixes N = 8 and width = 4. Its single-query scratch budget is tile width + 4 + 2 floats: one tile of scores, four accumulated output components and two normalisation scalars. It excludes inputs, output, tile rereads and the browser's own arrays and trace. It is not a production kernel's allocation or a total memory-traffic estimate.

No full score/weight matrices are written in the online algorithm's accounting. That removes the selected intermediate transfers from the previous experiment; Q, K, V and output still exist. The browser intentionally retains arrays for inspection. This is a small algorithm model, not GPU execution or a reproduction of a production FlashAttention kernel.