Skip to content

Completing the Output Tile by Tile

Return to the article

This follows the FlashAttention-2 forward pass. The steps are separated for explanation and do not prescribe the kernel’s instruction order.

Two Q rows on the left, four K vectors above, four V vectors below, and output accumulator vectors on the right. K and V are transposed to arrange tokens horizontally. The center shows only the current 2×2 tile of scores or unnormalized exponentials. Q and output positions stay fixed while traversing two K/V tiles. The full grid maps processing positions and does not store the full score matrix. This educational example omits scaling and masking. Current step: Input and output positions.

Current step SVG Current step PNG · 2400 px Representative scene PNG

Click the figure to open it at full size. Focus the step selector to navigate with arrow keys.