▶️ Run this notebook in Colab

FlashAttention-2 in an attention layer

Track 01 - Inference · Notebook 05 · Runtime: ≈20 min on Ampere+ GPU

Prerequisites: 07_gpu/04 (Triton FA2 kernel), 01_inference/01 (KV cache).

Hardware: Ampere or newer (compute capability ≥ 8.0). T4 (Turing, cc 7.5) is not supported because the Triton FA2 kernel relies on Ampere-era features: async global-to-shared copies (cp.async) and the larger per-SM shared-memory budget needed for FA2’s tile sizes. T4’s tensor cores can do FP16 matmul, but the kernel’s pipelining assumptions don’t hold there. Without a suitable GPU the kernel checks are skipped.

Papers: Dao et al. 2022 (2205.14135), Dao 2023 (2307.08691).

---

What

Drops the Triton FA2 forward kernel into an nn.Module attention layer and compares it against F.scaled_dot_product_attention at realistic LLM dimensions (B=1, H=32, T=2048, D=128). This is the minimal thing needed before you’d swap the layer into a real model.

Three things to check:

1. Numerical match to SDPA. Max abs error < 1e-3 in FP16.

2. Memory growth. FA2 is linear in sequence length; naive attention is quadratic. As T goes from 512 to 2048 (4×), FA2 peak memory grows ≤ 5× while naive grows ~16×.

3. Throughput. FA2 ≥ 3× naive (MATH backend) at T=2048.

from llm_systems_cookbook.nb import bootstrap

import math
import time

import torch
import torch.nn as nn
import torch.nn.functional as F

IS_CUDA = torch.cuda.is_available()
IS_AMPERE_PLUS = IS_CUDA and torch.cuda.get_device_capability(0) >= (8, 0)

s = bootstrap("01_inference_05_flashattention2_triton")

Shape algebra (runs on CPU)

For standard attention with B=1, H=32, T=2048, D=128 in FP16, the T × T score matrix alone is 2048 * 2048 * 2 bytes = 8 MiB per head, 256 MiB per batch. FA2 never materialises this matrix - it streams blocks and keeps a running (m, l, O) state of O(T * D) size.

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


def flash_running_state_bytes(batch: int, heads: int, seq: int, head_dim: int,
                               dtype_bytes: int = 2) -> int:
    # m, l are (B, H, T) in FP32; running O is (B, H, T, D) in FP32.
    return 4 * batch * heads * seq * (2 + head_dim)


naive = score_matrix_bytes(1, 32, 2048) / 1024**2
flash = flash_running_state_bytes(1, 32, 2048, 128) / 1024**2
print(f"naive T×T score matrix alone:    {naive:7.1f} MiB")
print(f"FA2 running (m, l, O) state:      {flash:7.1f} MiB")
print(f"ratio:                             {naive / flash:.1f}x")

s.check(
    "flash_state_smaller_than_naive_score_matrix",
    lambda: flash < naive / 2,
    msg=f"flash={flash:.1f} MiB  naive={naive:.1f} MiB",
)

The FA2 kernel (reproduced for self-contained execution)

Same kernel as 07_gpu/04 - reproduced here so this notebook runs standalone. See that notebook for the derivation.

triton = None
tl = None
flash_attn_triton = None

if IS_AMPERE_PLUS:
    try:
        import triton
        import triton.language as tl

        @triton.jit
        def _fa2_fwd(
            Q, K, V, O,
            sm_scale,
            stride_qb, stride_qh, stride_qm, stride_qk,
            stride_kb, stride_kh, stride_kn, stride_kk,
            stride_vb, stride_vh, stride_vn, stride_vk,
            stride_ob, stride_oh, stride_om, stride_ok,
            Z, H, N_CTX,
            BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, HEAD_DIM: tl.constexpr,
            CAUSAL: tl.constexpr,
        ):
            start_m = tl.program_id(0)
            off_hz = tl.program_id(1)
            off_z = off_hz // H
            off_h = off_hz % H

            q_offset = off_z * stride_qb + off_h * stride_qh
            k_offset = off_z * stride_kb + off_h * stride_kh
            v_offset = off_z * stride_vb + off_h * stride_vh
            o_offset = off_z * stride_ob + off_h * stride_oh

            offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
            offs_n = tl.arange(0, BLOCK_N)
            offs_k = tl.arange(0, HEAD_DIM)

            q_ptrs = Q + q_offset + offs_m[:, None] * stride_qm + offs_k[None, :] * stride_qk
            q = tl.load(q_ptrs, mask=offs_m[:, None] < N_CTX, other=0.0)

            m_i = tl.zeros((BLOCK_M,), dtype=tl.float32) - float("inf")
            l_i = tl.zeros((BLOCK_M,), dtype=tl.float32)
            acc = tl.zeros((BLOCK_M, HEAD_DIM), dtype=tl.float32)

            end_n = (start_m + 1) * BLOCK_M if CAUSAL else N_CTX
            for start_n in range(0, end_n, BLOCK_N):
                cur_n = start_n + offs_n
                k_ptrs = K + k_offset + cur_n[:, None] * stride_kn + offs_k[None, :] * stride_kk
                v_ptrs = V + v_offset + cur_n[:, None] * stride_vn + offs_k[None, :] * stride_vk
                k_blk = tl.load(k_ptrs, mask=cur_n[:, None] < N_CTX, other=0.0)
                v_blk = tl.load(v_ptrs, mask=cur_n[:, None] < N_CTX, other=0.0)
                qk = tl.dot(q, tl.trans(k_blk)) * sm_scale
                if CAUSAL:
                    qk = tl.where(offs_m[:, None] >= cur_n[None, :], qk, float("-inf"))
                qk = tl.where(cur_n[None, :] < N_CTX, qk, float("-inf"))
                m_new = tl.maximum(m_i, tl.max(qk, axis=1))
                alpha = tl.exp(m_i - m_new)
                p = tl.exp(qk - m_new[:, None])
                l_i = l_i * alpha + tl.sum(p, axis=1)
                acc = acc * alpha[:, None] + tl.dot(p.to(v_blk.dtype), v_blk)
                m_i = m_new

            acc = acc / l_i[:, None]
            o_ptrs = O + o_offset + offs_m[:, None] * stride_om + offs_k[None, :] * stride_ok
            tl.store(o_ptrs, acc.to(O.dtype.element_ty), mask=offs_m[:, None] < N_CTX)

        def flash_attn_triton(q, k, v, causal=True):
            B, H, T, D = q.shape
            o = torch.empty_like(q)
            BLOCK_M, BLOCK_N = 64, 64
            grid = (triton.cdiv(T, BLOCK_M), B * H)
            _fa2_fwd[grid](
                q, k, v, o, 1.0 / math.sqrt(D),
                q.stride(0), q.stride(1), q.stride(2), q.stride(3),
                k.stride(0), k.stride(1), k.stride(2), k.stride(3),
                v.stride(0), v.stride(1), v.stride(2), v.stride(3),
                o.stride(0), o.stride(1), o.stride(2), o.stride(3),
                B, H, T,
                BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, HEAD_DIM=D,
                CAUSAL=causal,
            )
            return o

        print("FA2 kernel ready")
    except Exception as e:  # noqa: BLE001
        print(f"Triton unavailable: {type(e).__name__}: {e}")
        flash_attn_triton = None

Attention layer using the kernel

FA2Attention is a standard causal self-attention nn.Module with the Triton kernel replacing the inner QK^T / softmax / AV path. Everything else (q/k/v projections, output projection) is plain PyTorch.

class FA2Attention(nn.Module):
    def __init__(self, d_model: int, num_heads: int) -> None:
        super().__init__()
        assert d_model % num_heads == 0
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.qkv = nn.Linear(d_model, d_model * 3, bias=False)
        self.o = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, _ = x.shape
        qkv = self.qkv(x).view(B, T, 3, self.num_heads, self.head_dim)
        q, k, v = qkv.unbind(dim=2)  # each (B, T, H, D)
        q = q.transpose(1, 2).contiguous()
        k = k.transpose(1, 2).contiguous()
        v = v.transpose(1, 2).contiguous()
        out = flash_attn_triton(q, k, v, causal=True)
        return self.o(out.transpose(1, 2).reshape(B, T, -1))


class SDPAAttention(nn.Module):
    '''Reference layer using F.scaled_dot_product_attention (MATH backend).'''

    def __init__(self, d_model: int, num_heads: int) -> None:
        super().__init__()
        assert d_model % num_heads == 0
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.qkv = nn.Linear(d_model, d_model * 3, bias=False)
        self.o = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, _ = x.shape
        qkv = self.qkv(x).view(B, T, 3, self.num_heads, self.head_dim)
        q, k, v = qkv.unbind(dim=2)
        q = q.transpose(1, 2).contiguous()
        k = k.transpose(1, 2).contiguous()
        v = v.transpose(1, 2).contiguous()
        with torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.MATH]):
            out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.o(out.transpose(1, 2).reshape(B, T, -1))

Numerical equivalence

Seed-match the two layers, run the same input through both, check max abs error. Both layers share the same Q, K, V, O projection weights, so any difference is kernel-driven.

if IS_AMPERE_PLUS and flash_attn_triton is not None:
    D_MODEL, NUM_HEADS, T_SMALL = 1024, 16, 256
    torch.manual_seed(0)
    fa2 = FA2Attention(D_MODEL, NUM_HEADS).cuda().half()
    torch.manual_seed(0)
    sdpa = SDPAAttention(D_MODEL, NUM_HEADS).cuda().half()

    x = torch.randn(1, T_SMALL, D_MODEL, device="cuda", dtype=torch.float16)
    with torch.inference_mode():
        y_fa = fa2(x)
        y_sdpa = sdpa(x)
    max_err = (y_fa - y_sdpa).abs().max().item()
    print(f"max abs error, FA2 vs SDPA: {max_err:.3e}")

    s.check(
        "fa2_matches_sdpa_within_1e_minus_3",
        lambda: max_err < 1e-3,
        msg=f"max abs err = {max_err:.3e}",
    )
else:
    s.skip("fa2_matches_sdpa_within_1e_minus_3", "no Ampere+ / no Triton")

Memory growth

Peak max_memory_allocated at T=512 and T=2048 for both backends. The ratio tells you which is quadratic in T and which is linear.

def peak_mb_for(T: int, layer: nn.Module) -> float:
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    x = torch.randn(1, T, layer.head_dim * layer.num_heads, device="cuda", dtype=torch.float16)
    with torch.inference_mode():
        layer(x)
    return torch.cuda.max_memory_allocated() / 1024**2


if IS_AMPERE_PLUS and flash_attn_triton is not None:
    D_MODEL, NUM_HEADS = 1024, 16
    torch.manual_seed(0)
    fa2 = FA2Attention(D_MODEL, NUM_HEADS).cuda().half()
    torch.manual_seed(0)
    sdpa = SDPAAttention(D_MODEL, NUM_HEADS).cuda().half()

    mem: dict[str, dict[int, float]] = {"fa2": {}, "sdpa": {}}
    for T in (512, 2048):
        mem["fa2"][T]  = peak_mb_for(T, fa2)
        mem["sdpa"][T] = peak_mb_for(T, sdpa)
    for T in (512, 2048):
        print(f"T={T:>4}   FA2={mem['fa2'][T]:7.1f} MiB   SDPA={mem['sdpa'][T]:7.1f} MiB")

    fa2_growth  = mem["fa2"][2048]  / max(mem["fa2"][512], 1e-6)
    sdpa_growth = mem["sdpa"][2048] / max(mem["sdpa"][512], 1e-6)
    print(f"FA2 growth {fa2_growth:.1f}x   SDPA growth {sdpa_growth:.1f}x   (T grew 4x)")

    s.check(
        "fa2_memory_linear_in_T",
        lambda: fa2_growth <= 6.0,
        msg=f"FA2 grew {fa2_growth:.1f}x when T grew 4x",
    )
    s.check(
        "fa2_beats_sdpa_in_peak_memory",
        lambda: mem["fa2"][2048] < mem["sdpa"][2048],
        msg=f"FA2 {mem['fa2'][2048]:.1f} MiB vs SDPA {mem['sdpa'][2048]:.1f} MiB at T=2048",
    )
else:
    s.skip("fa2_memory_linear_in_T",           "no Ampere+ / no Triton")
    s.skip("fa2_beats_sdpa_in_peak_memory",    "no Ampere+ / no Triton")

Throughput

Same input, same layers, just different kernels.

def bench(layer: nn.Module, x: torch.Tensor, n_iter: int = 20) -> float:
    with torch.inference_mode():
        for _ in range(3):
            layer(x)
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(n_iter):
            layer(x)
        torch.cuda.synchronize()
    return n_iter / (time.perf_counter() - t0)


if IS_AMPERE_PLUS and flash_attn_triton is not None:
    D_MODEL, NUM_HEADS, T_LARGE = 1024, 16, 2048
    torch.manual_seed(0)
    fa2 = FA2Attention(D_MODEL, NUM_HEADS).cuda().half()
    torch.manual_seed(0)
    sdpa = SDPAAttention(D_MODEL, NUM_HEADS).cuda().half()
    x = torch.randn(1, T_LARGE, D_MODEL, device="cuda", dtype=torch.float16)
    its_fa2 = bench(fa2, x)
    its_sdpa = bench(sdpa, x)
    print(f"FA2  = {its_fa2:6.1f} it/s")
    print(f"SDPA = {its_sdpa:6.1f} it/s")

    s.benchmark(
        "fa2_at_least_3x_naive_at_T2048",
        measured=its_fa2,
        baseline=its_sdpa,
        must_beat=3.0,
        unit="it/s",
        higher_is_better=True,
    )
else:
    s.skip("fa2_at_least_3x_naive_at_T2048", "no Ampere+ / no Triton")

Memory curve + speedup bar

Left: peak allocation vs sequence length - SDPA’s MATH backend materialises the T × T score matrix, so it climbs as T² on the log axis; FA2’s streaming state rises linearly. Right: FA2 vs SDPA throughput at T=2048, the point where the memory gap is already an order of magnitude.

import matplotlib.pyplot as plt

if IS_AMPERE_PLUS and flash_attn_triton is not None:
    sweep_Ts = [256, 512, 1024, 2048]
    fa_peaks, sdpa_peaks = [], []
    for T_ in sweep_Ts:
        fa_peaks.append(peak_mb_for(T_, fa2))
        sdpa_peaks.append(peak_mb_for(T_, sdpa))

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.4))
    ax1.plot(sweep_Ts, sdpa_peaks, "o-", label="SDPA (MATH)")
    ax1.plot(sweep_Ts, fa_peaks,   "o-", label="FA2 (triton)")
    ax1.set_xscale("log")
    ax1.set_yscale("log")
    ax1.set_xlabel("sequence length T")
    ax1.set_ylabel("peak memory (MiB)")
    ax1.set_title("peak memory vs T")
    ax1.legend()
    ax1.grid(True, which="both", alpha=0.3)

    ratio = its_fa2 / its_sdpa
    ax2.bar(["SDPA (MATH)", "FA2 (triton)"], [its_sdpa, its_fa2],
            color=["tab:gray", "tab:blue"])
    ax2.set_ylabel("iterations / s")
    ax2.set_title(f"latency @ T=2048 - FA2 = {ratio:.2f}x SDPA")
    ax2.grid(True, axis="y", alpha=0.3)
    fig.tight_layout()
    plt.show()
else:
    print("skipped - no Ampere+ / no Triton.")

Exercises

1. Swap FA2Attention into a minimal GPT-2 block (LN → attn → LN → MLP) and compare end-to-end forward latency vs an all-SDPA version at T=2048.

2. Autotune BLOCK_M, BLOCK_N, num_warps via @triton.autotune for your specific GPU.

3. Drop in flash-attn (the official package) and verify our Triton kernel is within 2× of its throughput. If not, the usual gap is either autotuning or missing pipelining; see Dao 2023 §3.3.

References

• flash-attn official: Dao-AILab/flash-attention

• vLLM’s call site: vllm/attention/backends/flash_attn.py.

s.summary()
s.save()