Attention from scratch + roofline

Open in Colab ▶️ Run this notebook in Colab

Attention from scratch + roofline#

Track 01 - Inference · Notebook 02 · Runtime: ≈5 min on CPU or GPU

Prerequisites: 07_gpu/01 (GPU architecture tour), 01_inference/01 (KV cache).

Papers: Dao et al. 2022 (2205.14135); Williams et al. 2009 Roofline.

What#

Implement scaled dot-product attention from scratch, verify against F.scaled_dot_product_attention, derive the FLOPs/bytes/arithmetic- intensity formulas, and classify prefill vs decode on the roofline.

This is the analytical foundation for every attention optimisation that follows (PagedAttention, FlashAttention, GQA/MLA): once you know the AI of each stage you know where the ceiling is before you write a line of kernel code.

from llm_systems_cookbook.nb import bootstrap

import math

import torch
import torch.nn.functional as F

s = bootstrap("01_inference_02_attention_roofline")

Naive attention, matched against SDPA#

attn(Q, K, V) = softmax(Q K^T / sqrt(d)) V

def naive_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
                    causal: bool = True) -> torch.Tensor:
    B, H, T, D = q.shape
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(D)
    if causal:
        mask = torch.triu(torch.ones(T, T, device=q.device, dtype=torch.bool), diagonal=1)
        scores = scores.masked_fill(mask, float("-inf"))
    return torch.softmax(scores, dim=-1) @ v


torch.manual_seed(0)
B, H, T, D = 1, 4, 128, 32
q = torch.randn(B, H, T, D)
k = torch.randn(B, H, T, D)
v = torch.randn(B, H, T, D)

y_ref = F.scaled_dot_product_attention(q, k, v, is_causal=True)
y_ours = naive_attention(q, k, v, causal=True)
max_err = (y_ref - y_ours).abs().max().item()
print(f"naive vs SDPA max abs err = {max_err:.3e}")
s.check("naive_matches_sdpa_1e_minus_3", lambda: max_err < 1e-3, msg=f"err={max_err:.3e}")

FLOPs / bytes / arithmetic intensity#

Two matmuls (QK^T and attn_weights @ V) dominate:

flops  = 4 * B * H * T * T * D
bytes  = dtype_bytes * B * H * D * (T + 2T + T)   # Q + K + V + O

def attention_flops(batch: int, heads: int, seq: int, head_dim: int) -> int:
    return 4 * batch * heads * seq * seq * head_dim


def attention_bytes(batch: int, heads: int, seq: int, head_dim: int,
                     dtype_bytes: int = 2) -> int:
    return dtype_bytes * batch * heads * head_dim * (seq + 2 * seq + seq)


def arithmetic_intensity(batch: int, heads: int, seq: int, head_dim: int,
                          dtype_bytes: int = 2) -> float:
    return attention_flops(batch, heads, seq, head_dim) / \
           attention_bytes(batch, heads, seq, head_dim, dtype_bytes)


for seq in (64, 256, 1024, 4096):
    ai = arithmetic_intensity(1, 32, seq, 128)
    print(f"T={seq:>5}  AI = {ai:>7.1f} FLOPs/byte")

# Doubling T should roughly double AI (scores matrix grows quadratically,
# but we only count the matmul ops; bytes scale linearly).
ai_1k = arithmetic_intensity(1, 32, 1024, 128)
ai_2k = arithmetic_intensity(1, 32, 2048, 128)
print(f"AI(T=2k) / AI(T=1k) = {ai_2k / ai_1k:.2f}  (expect ~2)")

s.check(
    "ai_doubles_with_seq_len",
    lambda: 1.8 <= ai_2k / ai_1k <= 2.2,
    msg=f"ratio = {ai_2k / ai_1k:.3f}",
)

Classify prefill vs decode#

Prefill attention at N=1024 has enough intensity to be compute-bound. Decode (N=1 in the query dim; full seq_k) has AI ≈ 2 FLOPs/byte - memory-bound everywhere.

def decode_attention_intensity(batch: int, n_kv_heads: int, head_dim: int,
                                seq_k: int, dtype_bytes: int = 2) -> float:
    flops = 4 * batch * n_kv_heads * 1 * seq_k * head_dim
    bytes_ = dtype_bytes * batch * n_kv_heads * head_dim * (1 + 2 * seq_k + 1)
    return flops / bytes_


T4_RIDGE = 203  # FLOPs/byte (from the GPU arch tour)
A100_RIDGE = 153

prefill_ai = arithmetic_intensity(1, 32, 1024, 128)
decode_ai = decode_attention_intensity(1, 8, 128, 2048)
print(f"prefill@N=1024 AI = {prefill_ai:.1f}  (T4 ridge = {T4_RIDGE})")
print(f"decode@N=1 kv=2048 AI = {decode_ai:.1f}")

s.check(
    "prefill_compute_bound_on_T4",
    lambda: prefill_ai > T4_RIDGE,
    msg=f"prefill AI={prefill_ai:.1f}  ridge={T4_RIDGE}",
)
s.check(
    "decode_memory_bound",
    lambda: decode_ai < 20,
    msg=f"decode AI={decode_ai:.1f}",
)
s.check(
    "prefill_compute_bound_on_A100_too",
    lambda: prefill_ai > A100_RIDGE,
    msg=f"prefill AI={prefill_ai:.1f}  A100 ridge={A100_RIDGE}",
)

Empirical O(T²) memory growth#

Run attention at doubling sequence lengths and record torch.cuda.max_memory_allocated (or skip on CPU). Memory should grow quadratically in T - this is the thing FlashAttention fixes.

if torch.cuda.is_available():
    peaks: dict[int, float] = {}
    for T_ in (256, 512, 1024):
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()
        q = torch.randn(1, 8, T_, 64, device="cuda", dtype=torch.float16)
        k = torch.randn_like(q)
        v = torch.randn_like(q)
        with torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.MATH]):
            _ = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        torch.cuda.synchronize()
        peaks[T_] = torch.cuda.max_memory_allocated() / 1024**2

    for T_, p in peaks.items():
        print(f"T={T_:>4}  peak memory = {p:.1f} MiB")
    growth = peaks[1024] / peaks[256]
    print(f"ratio T=1024/T=256 = {growth:.2f}  (~16 expected for O(T^2))")
    s.check(
        "empirical_memory_quadratic_in_T",
        lambda: 8 <= growth <= 25,
        msg=f"growth ratio = {growth:.2f}",
    )
else:
    s.skip("empirical_memory_quadratic_in_T", "no CUDA (peak memory metric requires CUDA)")

Attention intensity on the roofline#

Prefill’s arithmetic intensity grows linearly with T; decode stays pinned near 2 FLOPs/byte regardless of KV length. Plot both curves against the T4 ridge - the crossover is what separates compute-bound prefill from memory-bound decode.

import matplotlib.pyplot as plt

seqs = [64, 128, 256, 512, 1024, 2048, 4096, 8192]
prefill_curve = [arithmetic_intensity(1, 32, T_, 128) for T_ in seqs]
decode_curve  = [decode_attention_intensity(1, 8, 128, T_) for T_ in seqs]

fig, ax = plt.subplots(figsize=(6.5, 3.6))
ax.loglog(seqs, prefill_curve, "o-", label="prefill (Q,K,V all length T)")
ax.loglog(seqs, decode_curve,  "o-", label="decode (Q length 1, KV length T)")
ax.axhline(T4_RIDGE,  color="tab:gray",  linestyle=":",  label=f"T4 ridge ≈ {T4_RIDGE}")
ax.axhline(A100_RIDGE, color="tab:olive", linestyle="--", label=f"A100 ridge ≈ {A100_RIDGE}")
ax.fill_between(seqs, 1, min(T4_RIDGE, A100_RIDGE), color="tab:blue",   alpha=0.08)
ax.fill_between(seqs, max(T4_RIDGE, A100_RIDGE), max(prefill_curve) * 2,
                color="tab:orange", alpha=0.08)
ax.set_xlabel("sequence length T")
ax.set_ylabel("arithmetic intensity (FLOPs / byte)")
ax.set_title("attention AI vs T - prefill crosses the ridge, decode never does")
ax.legend(fontsize=8, loc="center right")
ax.grid(True, which="both", alpha=0.3)
fig.tight_layout()
plt.show()

Exercises#

  1. Redo the AI calculation for grouped-query attention with num_kv_heads = n_head / 4. Bytes drop by 4× on the KV terms; intensity on decode rises proportionally.

  2. Plot a roofline heatmap: x-axis = context length (64..8192), y-axis = batch size (1..64), colour = achievable TFLOPs on T4 and H100. Shows the “sweet spot” batch at each context.

  3. Add causal masking to the FLOP count (halve the scores for the upper triangle) and compare.

References#

  • Dao et al. 2022 §3 (FlashAttention) for the memory formula we reproduced above.

  • Horace He, Making Deep Learning Go Brrrr, for the roofline vocabulary.

s.summary()
s.save()