The attention scratch is still O(seq^2): 627 MB at 2048, 2.4 GB at 4096, against a model advertising 131072

[xoloki]

Split out of #187, which shared one attention scratch across the layers and closed there. This is the half it left.

The ceiling

block::reserve sizes scores and probs by the sequence, and they are the only things that grow as seq², where everything else grows as seq:

TinyLlama, fp16, one shared scratch (since #188)

  seq     scores+probs      total scratch
  2048       536.9 MB          627 MB
  4096         2.15 GB         2.4 GB
  8192         8.59 GB         9.2 GB

Before #188 this was 22x worse — one copy per layer — and a 2048-token prefill died with kIOGPUCommandBufferCallbackErrorOutOfMemory. Sharing bought a factor of 22 and did not change the exponent.

Llama 3.2 advertises 131072 tokens. Qwen 2.5 advertises 32768. Nothing near either is reachable, and the wall is quadratic so it arrives fast.

It is also the rest of the prefill slope

Prefill per token, after the tiling in #189 took the weights out of it:

  128   1.77 ms
  1024  2.47 ms
  2048  3.40 ms

Still rising, and the weights no longer are — so what is left growing is the attention, which is the same seq² the memory is. One change addresses both.

The shape of a fix

Process the queries in bands: take a slice of rows at a time, score them against the keys, softmax, weight the values, and reuse one scratch sized by the band rather than by the sequence. This is what flash-attention-style implementations do and for the same reason.

Bounds the memory by the band. Whether it is faster is a separate question — it trades a big allocation for more passes over the keys, and on a device where the attention kernels are already a small fraction of a decode step (attention_scores 6.2 us, attention_weighted 15.8 us) the win is at prefill and at length, not at decode.

Start with the mask, not with a long prompt

The arithmetic is where the bugs will be: a band has to know its own offset into the sequence so the causal mask stays right, and that is exactly the kind of off-by-one that produces fluent, plausible, wrong output rather than a crash.

causal_mask already takes a key offset — #187's fix relied on the distinction between "where the queries are" and "how many keys precede them", and got it wrong once before the uncached rope test caught it. A unit test on the mask at a band offset is cheap and is the thing most likely to fail.

A 4096-token end-to-end run is the confirmation, not the test, and it needs a model whose context allows it — TinyLlama's is 2048, so this wants Llama 3.2 or Qwen.

Current state, for whoever picks this up

Decode is 92.6 tok/s and 82% of the device (#190, closed). Prefill 512 is ~1.02 s. The scratch is shared, and model::scratch() reports its size in elements — added in #188 precisely so a caller sizing a context can ask.

[xoloki]

Measured on master at 2a0341d, for the record:

prefill    512   1.028s
          2048   6.836s     (3.34 ms/token, still rising)

decode  at ctx ~129   92.6 tok/s
        at ctx 1024   57.3 tok/s

Decode loses 38% between short context and 1024, which is the same seq² and is not something #190 addressed — that issue measured decode at short context throughout. So this is not only a prefill and memory problem; it is what a long conversation costs per token, which is the thing a user notices second after the first token appears.

Worth stating because it changes the value: chunking was filed as a memory ceiling with a prefill slope attached. It is also the largest remaining decode lever at any realistic context length.