Skand Tandon
Deep Implementation

Profiling attention memory growth

The KV cache grows with every token, and the attention scores matrix grows with the square of sequence length. Profiling both reveals why long-context inference is a memory problem before it is a compute one.

2025-11-04 deep implementation 8 min read

In the previous post I implemented scaled dot-product attention from scratch. The code works, but I left a claim unverified: that memory grows quadratically with sequence length for naive attention and linearly for the KV cache. Today I'm instrumenting both with PyTorch's memory profiling tools to see exactly where the bytes go.

The profiling setup

PyTorch provides torch.cuda.memory_allocated() and torch.cuda.max_memory_allocated() to track GPU memory. For finer-grained analysis, the memory snapshot API captures every allocation and free event. I'll use the simpler approach first.

import torch
import torch.nn.functional as F

def measure_attention_memory(seq_len, num_heads=32, head_dim=128, dtype=torch.float16):
    """Measure peak memory of naive attention at a given sequence length."""
    torch.cuda.reset_peak_memory_stats()
    torch.cuda.empty_cache()

    device = "cuda"
    Q = torch.randn(1, num_heads, seq_len, head_dim, device=device, dtype=dtype)
    K = torch.randn(1, num_heads, seq_len, head_dim, device=device, dtype=dtype)
    V = torch.randn(1, num_heads, seq_len, head_dim, device=device, dtype=dtype)

    mem_before = torch.cuda.memory_allocated()

    # Naive attention: materializes the full scores matrix
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (head_dim ** 0.5)
    mask = torch.triu(torch.ones(seq_len, seq_len, device=device, dtype=torch.bool), diagonal=1)
    scores = scores.masked_fill(mask, float("-inf"))
    weights = F.softmax(scores, dim=-1)
    output = torch.matmul(weights, V)

    mem_after = torch.cuda.max_memory_allocated()
    peak_mb = (mem_after - mem_before) / (1024 ** 2)

    # Clean up
    del Q, K, V, scores, mask, weights, output
    torch.cuda.empty_cache()

    return peak_mb

Quadratic scaling in action

Running this across a range of sequence lengths reveals the quadratic growth clearly:

# Sweep sequence lengths and record peak memory
results = []
for seq_len in [512, 1024, 2048, 4096, 8192]:
    mem_mb = measure_attention_memory(seq_len)
    results.append((seq_len, mem_mb))
    print(f"seq_len={seq_len:>5}  peak_mem={mem_mb:>8.1f} MB")

The dominant term is the scores matrix: shape (1, num_heads, seq_len, seq_len). In FP16, that's num_heads * seq_len^2 * 2 bytes. For 32 heads at sequence length 4096, that's 32 * 4096^2 * 2 = 1,073,741,824 bytes, or 1 GB. At 8192, it's 4 GB. At 16384, it's 16 GB, which doesn't fit on many GPUs with the model weights already loaded.

The softmax output (weights) has the same shape, so the peak memory is roughly 2 * num_heads * seq_len^2 * bytes_per_element when both tensors are alive simultaneously. This is the memory wall that FlashAttention breaks.

KV cache: linear but still large

The KV cache grows linearly with sequence length, not quadratically. But "linear" doesn't mean "small." Let me compute it for a realistic model.

def kv_cache_size_mb(seq_len, num_layers=80, num_kv_heads=8,
                     head_dim=128, dtype_bytes=2):
    """KV cache size for one request in a Llama-3.1-70B-like model."""
    # Per layer: K and V, each (num_kv_heads, seq_len, head_dim)
    per_layer = 2 * num_kv_heads * seq_len * head_dim * dtype_bytes
    total = num_layers * per_layer
    return total / (1024 ** 2)

for seq_len in [512, 1024, 2048, 4096, 8192, 16384, 32768]:
    mb = kv_cache_size_mb(seq_len)
    print(f"seq_len={seq_len:>5}  kv_cache={mb:>8.1f} MB")

For a 70B model with 80 layers, 8 KV heads (GQA), and head_dim 128 in FP16:

A single 128K-context request consumes 40 GB of KV cache in FP16. An H100 with 80 GB of HBM has the model weights (140 GB in FP16 across 8 GPUs, so 17.5 GB per GPU for TP=8) plus 40 GB of KV cache. That leaves very little room for batching. This is why KV cache compression (INT8, INT4, or even lower precision for cached values) and PagedAttention are so important.

Profiling with torch.profiler

For a more detailed view, torch.profiler with profile_memory=True shows every allocation and its call stack:

from torch.profiler import profile, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CUDA],
    profile_memory=True,
    record_shapes=True,
) as prof:
    Q = torch.randn(1, 32, 4096, 128, device="cuda", dtype=torch.float16)
    K = torch.randn(1, 32, 4096, 128, device="cuda", dtype=torch.float16)
    V = torch.randn(1, 32, 4096, 128, device="cuda", dtype=torch.float16)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / 11.31
    weights = F.softmax(scores, dim=-1)
    output = torch.matmul(weights, V)

print(prof.key_averages().table(sort_by="self_cuda_memory_usage", row_limit=10))

The output shows that the aten::mm call producing the scores matrix allocates the largest single tensor, followed by the softmax output. The final matmul (weights @ V) is much smaller because the output shape is (1, 32, 4096, 128) instead of (1, 32, 4096, 4096).

FlashAttention: where the memory goes

When using F.scaled_dot_product_attention with the FlashAttention backend, peak memory drops dramatically because the scores matrix is never fully materialized. Instead of O(N^2) memory, FlashAttention uses O(N) memory (plus some small buffers for the tiling).

def measure_flash_attention_memory(seq_len, num_heads=32, head_dim=128):
    torch.cuda.reset_peak_memory_stats()
    torch.cuda.empty_cache()

    Q = torch.randn(1, num_heads, seq_len, head_dim, device="cuda", dtype=torch.float16)
    K = torch.randn(1, num_heads, seq_len, head_dim, device="cuda", dtype=torch.float16)
    V = torch.randn(1, num_heads, seq_len, head_dim, device="cuda", dtype=torch.float16)

    mem_before = torch.cuda.memory_allocated()

    # PyTorch's SDPA dispatches to FlashAttention when available
    output = F.scaled_dot_product_attention(Q, K, V, is_causal=True)

    mem_after = torch.cuda.max_memory_allocated()
    return (mem_after - mem_before) / (1024 ** 2)

At sequence length 8192, naive attention peaks at roughly 4 GB for the scores matrix. FlashAttention peaks at a few hundred MB (the output tensor plus small working buffers). The compute is the same; only the memory pattern changes.

The practical rule

For sequence lengths above 2048, always use FlashAttention (via F.scaled_dot_product_attention or the flash-attn package). Below 2048, the naive implementation is fine and sometimes faster because the overhead of tiling outweighs the memory savings. Most serving engines make this choice automatically.

Memory budgeting for production

Understanding where memory goes lets you budget it properly for a serving deployment. For a 70B model on 8xH100 (80 GB each, TP=8):

This budget directly determines your maximum batch size and maximum sequence length. It's why INT8 quantization of both weights and KV cache is so appealing: halving the bytes doubles the capacity.

Next: the INT8 quantization pipeline, where I'll implement weight quantization from scratch and measure its impact on both memory and accuracy.