Inference Engineering · Post 03

FlashAttention, one query block at a time

1 / 1

In HBM — the full tensors

4 tokens, head dimension 3, blocks of 2. The 1/√d scale is left off so the scores stay whole numbers.

The score matrix, which is never stored

Shown only so you can see where the live tile sits. Solid cells are on chip right now; everything faint has either been discarded or not computed yet.

On chip — SRAM

Accumulators for the current query block

How O was built this step

Written back to HBM

resident on chip this step's tile changed this step ← → arrow keys to step