Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Quadratic Attention Bottleneck › Worked Example: Tracing One Layer

Michael BrenndoerferUpdated June 21, 20253 min read

Language AI HandbookAll chapters

Part 17 of 67Efficient Attention

Chapter 117Quadratic Attention Bottleneck › Worked Example: Tracing One Layer

Reading section 8 of 13 · Worked Example: Tracing One Layer · 3 minSection 8 of 13 · 3 min

Section 8 of 13 · 3 min read

Quadratic Attention Bottleneck › Worked Example: Tracing One LayerOne section · Sections in this chapter8 / 13 ↓

Sections in this chapter13

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

To make all of these costs concrete, let's trace a single attention layer through every operation and compute the exact memory and FLOP cost at a specific sequence length. This exercise ties together everything discussed in this chapter.

Consider a single attention layer with the following configuration, representative of one layer in a GPT-2 Medium-style model:

Sequence length: $n = 2048$ tokens
Model dimension: $d = 1024$
Number of heads: $h = 16$
Head dimension: $d_k = d / h = 64$
Precision: fp16 (2 bytes per element)

In[17]:

Code

# Configuration
n = 2048
d = 1024
h = 16
d_k = d // h  # 64
dtype_bytes = 2  # fp16

# Step 1: Input tensor
input_tokens = n * d * dtype_bytes  # shape (n, d)

# Step 2: Q, K, V projections (three separate (d x d) weight matrices)
qkv_projection_flops = 3 * 2 * n * d * d  # 3 matrix mults of (n,d) x (d,d)
qkv_memory = 3 * n * d * dtype_bytes  # Q, K, V each of shape (n, d)

# Step 3: QK^T for each head
score_flops = h * 2 * n * n * d_k  # h heads, each (n,d_k) x (d_k,n)
score_memory = h * n * n * dtype_bytes  # h attention matrices of shape (n, n)

# Step 4: Softmax (in-place, no extra memory)
softmax_flops = h * 5 * n * n

# Step 5: Attention @ V
output_flops = h * 2 * n * n * d_k  # h heads, each (n,n) x (n,d_k)
output_memory = n * d * dtype_bytes  # shape (n, d)

# Step 6: Output projection (d x d weight matrix)
out_proj_flops = 2 * n * d * d
out_proj_memory = n * d * dtype_bytes

# Total
total_flops = (
    qkv_projection_flops
    + score_flops
    + softmax_flops
    + output_flops
    + out_proj_flops
)
total_memory_mb = (
    qkv_memory + score_memory + output_memory + out_proj_memory
) / (1024**2)
attention_memory_mb = score_memory / (1024**2)

Out[18]:

Console

Single Attention Layer Trace (n=2048, d=1024, h=16, fp16)

Step                                     FLOPs       Memory
------------------------------------------------------------
QKV Projections                 12,884,901,888  (weights excluded)
Score Computation (QK^T)         8,589,934,592       128.0 MB
Softmax Normalization              335,544,320    (in-place)
Output Computation (A*V)         8,589,934,592         4.0 MB
Output Projection                4,294,967,296  (weights excluded)
------------------------------------------------------------
TOTAL                           34,695,282,688       148.0 MB

Attention matrix memory alone: 128.0 MB
Attention matrix as % of attention memory: 86.5%

The breakdown reveals that even for just 2,048 tokens, the attention score matrices already consume a significant portion of working memory. Now consider that a 24-layer model runs this computation 24 times simultaneously during training, and all attention matrices must be retained for the backward pass. This is where the per-layer cost scales to multi-gigabyte territory for longer sequences.

Notice also that the QKV projections, while not $O(n^2)$ , are not trivial. At this sequence length, they consume more FLOPs than the attention core itself. This is the regime below the crossover point, where linear components still dominate. At 8,192 tokens (4x longer), the attention FLOPs would grow by 16x while the projection FLOPs grow by only 4x, pushing attention into dominance.