Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

FlashAttention Algorithm › Tiling for SRAM

Michael BrenndoerferUpdated June 29, 20256 min read

Language AI HandbookAll chapters

Part 17 of 67Efficient Attention

Chapter 124FlashAttention Algorithm › Tiling for SRAM

Reading section 3 of 12 · Tiling for SRAM · 6 minSection 3 of 12 · 6 min

Section 3 of 12 · 6 min read

FlashAttention Algorithm › Tiling for SRAMOne section · Sections in this chapter3 / 12 ↓

Sections in this chapter12

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

The core insight of FlashAttention is to restructure the attention computation so that all intermediate values fit in SRAM rather than spilling to HBM. This restructuring uses a technique called tiling: breaking the large attention computation into small blocks that fit entirely in fast memory.

The Tiling StrategyLink Copied

Consider computing attention for a sequence of $n$ tokens. Standard attention materializes the entire $n \times n$ attention score matrix at once, which requires $O(n^2)$ memory. FlashAttention processes the computation in blocks of size $B_r \times B_c$ , where $B_r$ and $B_c$ are chosen so that the required data fits in SRAM.

Block Sizes

FlashAttention processes queries in blocks of $B_r$ rows and keys/values in blocks of $B_c$ columns. These block sizes are chosen based on SRAM capacity: typically $B_r = B_c = 64$ to $256$ depending on the GPU architecture and model dimension.

The algorithm proceeds as follows:

Load a block of $B_r$ queries into SRAM
For each block of $B_c$ keys and values, load them into SRAM, compute attention scores, update running statistics, and accumulate the output
Write the final output for these $B_r$ queries to HBM
Repeat for the next block of queries

The key point is that only blocks reside in SRAM at any time, never the full $n \times n$ matrix. The memory required is:

M_{\text{SRAM}} = O(B_r B_c + B_r d + B_c d)

where:

$B_r$ : the number of query rows in each block
$B_c$ : the number of key/value columns in each block
$d$ : the head dimension (size of each QKV vector)
$B_r B_c$ : memory for the block of attention scores
$B_r d$ : memory for query vectors and output accumulator
$B_c d$ : memory for key and value vectors

Since $B_r$ , $B_c$ , and $d$ are all constants chosen at compile time, this memory requirement is independent of sequence length $n$ .

In[7]:

Code

def compute_sram_requirements(block_r, block_c, d_model, dtype_bytes=2):
    """
    Compute SRAM memory requirements for FlashAttention tiling.

    Args:
        block_r: Query block size
        block_c: Key/value block size
        d_model: Model dimension
        dtype_bytes: Bytes per element

    Returns:
        Dictionary with memory requirements in KB
    """
    # Block of queries: B_r x d
    q_block = block_r * d_model * dtype_bytes

    # Block of keys: B_c x d
    k_block = block_c * d_model * dtype_bytes

    # Block of values: B_c x d
    v_block = block_c * d_model * dtype_bytes

    # Block of scores: B_r x B_c
    score_block = block_r * block_c * dtype_bytes

    # Running statistics: B_r (max and sum for softmax)
    stats = block_r * 2 * dtype_bytes

    # Output accumulator: B_r x d
    output_block = block_r * d_model * dtype_bytes

    total_bytes = (
        q_block + k_block + v_block + score_block + stats + output_block
    )

    return {
        "q_kb": q_block / 1024,
        "k_kb": k_block / 1024,
        "v_kb": v_block / 1024,
        "scores_kb": score_block / 1024,
        "stats_kb": stats / 1024,
        "output_kb": output_block / 1024,
        "total_kb": total_bytes / 1024,
    }


# Typical configurations
configs = [
    {"block_r": 64, "block_c": 64, "d_model": 64},  # Small head
    {"block_r": 128, "block_c": 128, "d_model": 64},  # Larger blocks
    {"block_r": 64, "block_c": 64, "d_model": 128},  # Larger head dimension
]

Out[8]:

Console

SRAM Requirements for FlashAttention Blocks:

    Br     Bc      d        Q        K        V    Score      Out      Total
---------------------------------------------------------------------------
    64     64     64     8.0K     8.0K     8.0K     8.0K     8.0K      40.2K
   128    128     64    16.0K    16.0K    16.0K    32.0K    16.0K      96.5K
    64     64    128    16.0K    16.0K    16.0K     8.0K    16.0K      72.2K

Typical GPU SRAM per streaming multiprocessor is around 192 KB. The table shows that even with generous block sizes of $128 \times 128$ and head dimension 64, the total SRAM usage stays under 50 KB, well within that capacity. This leaves room for additional working memory and allows multiple thread blocks to run concurrently.

Visualizing the Tiling PatternLink Copied

The tiling pattern creates a block structure in the attention computation. Rather than computing the full attention matrix at once, FlashAttention computes it block by block, accumulating results as it goes.

Out[9]:

Visualization

Square grid showing full attention matrix with all cells computed at once. — Standard attention materializes the entire n x n attention matrix in HBM. For long sequences, this matrix dominates memory usage.

Square grid divided into blocks with highlighting showing one block being processed at a time. — FlashAttention computes attention in blocks. Only one block resides in SRAM at a time, and results accumulate incrementally.

The visualization shows the fundamental difference. Standard attention fills the entire matrix (left), storing all $n^2$ values in HBM. FlashAttention processes one block at a time (right, with the current block highlighted), keeping only that block in fast SRAM. Completed blocks are not stored; instead, their contribution to the output is accumulated incrementally.