FlashAttention for Image Models

3.1 Tiling

Rather than computing the full \(S = QK^\top\) matrix at once, FlashAttention splits \(Q\), \(K\), and \(V\) into blocks small enough to fit into fast on-chip SRAM. For each pair of blocks (a block of \(Q\) and a block of \(K/V\)), it:

1. Loads the \(Q\)-block and \(K/V\)-block from HBM into SRAM once.
2. Computes the partial attention scores for just that block, entirely within SRAM.
3. Accumulates a running (partial) output and running softmax statistics.
4. Moves to the next \(K/V\) block, repeating, without ever writing the full \(N \times N\) score or probability matrix back to HBM.

This means the intermediate \(N \times N\) matrices \(S\) and \(P\) are never materialized in HBM at all — they exist only transiently, block by block, inside fast on-chip memory. This is the single biggest lever: FlashAttention reduces HBM accesses from

\[ \Theta(Nd + N^2) \]

in the standard algorithm to roughly

\[ O\!\left(\frac{N^2 d^2}{M}\right) \]

where \(M\) is the SRAM size — for typical head dimensions and SRAM sizes, this works out to roughly an order of magnitude (often cited around \(5\)–\(25\times\)) fewer memory accesses, which translates directly into wall-clock speedups even though the FLOP count is essentially unchanged.

3.2 Online (Streaming) Softmax

Tiling raises an obvious problem: softmax needs to normalize across the entire row of S (i.e., across all keys), but tiling only ever sees one block of keys at a time. The trick FlashAttention borrows and formalizes is the online softmax algorithm: instead of computing the softmax normalization constant in one pass over the whole row, it maintains a running maximum and running sum of exponentials as each new K/V block is processed, and rescales the accumulated output whenever the running maximum changes. This produces mathematically identical results to a full-row softmax while only ever needing to see one block of keys at a time, which is what makes single-pass, block-by-block computation possible without approximation.

3.3 Recomputation in the Backward Pass

Standard attention implementations save the \(N \times N\) attention matrix \(P\) from the forward pass so it can be reused during backpropagation to compute gradients. FlashAttention deliberately does not store \(P\). Instead, during the backward pass, it recomputes the needed blocks of \(S\) and \(P\) on the fly from \(Q\), \(K\), \(V\) and the small saved softmax statistics (the running max \(m\) and sum \(\ell\) per row), using the same fast SRAM-resident tiling strategy. This trades a modest amount of extra compute (recomputation) for a large reduction in memory reads/writes, and turns out to be a net win because HBM bandwidth, not FLOPs, is the bottleneck. This is the same “recomputation” idea used in gradient checkpointing, applied specifically to the attention matrix.

(See the equations in Section 4 for the exact form of \(S\), \(P\), \(m\), and \(\ell\) that get recomputed block-by-block.)

3.4 Net Effect

The combination of tiling and online softmax means FlashAttention computes mathematically exact attention (no approximation, unlike sparse or low-rank attention variants) while reducing memory usage from \(O(N^2)\) to \(O(N)\) and reducing wall-clock latency by multiples (commonly \(2\)–\(4\times\) in the original paper, more in later versions) due to far fewer HBM round trips. This is why FlashAttention was adopted almost universally rather than being treated as one option among many approximate-attention schemes: it changes how attention is computed, not what it computes.

5.1 FlashAttention (v1)

Introduced by Tri Dao et al. in “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness” (2022), this established the tiling + online-softmax + recomputation approach described above. It delivered exact attention with linear memory in sequence length and substantial wall-clock speedups (commonly cited as 2–4× over standard PyTorch attention, and larger gains at longer sequence lengths), immediately making longer-context training and higher-resolution vision workloads more tractable.

5.2 FlashAttention-2

FlashAttention-2 (Dao, 2023) kept the same IO-aware foundation but re-engineered the work partitioning and parallelism to better use GPU hardware:

• Reduced non-matmul FLOPs. The original algorithm spent relatively more time on softmax rescaling operations, which don’t map efficiently onto GPU tensor cores. FA2 restructured the inner loop to minimize this non-matmul overhead.
• Better parallelism across thread blocks. FA1 parallelized primarily over batch size and number of heads, which underutilizes the GPU when batch×heads is small relative to the number of streaming multiprocessors (a common situation for long-sequence workloads, exactly the regime high-resolution vision models push into). FA2 also parallelizes over the sequence dimension, keeping the GPU busier even at small batch sizes.
• Better work partitioning within a thread block, reducing communication between warps.

Net result: FlashAttention-2 is reported at roughly 1.7–3.0× faster than FlashAttention-1 (and 1.3–2.5× faster than a Triton implementation of FA1), reaching 50–73% of theoretical peak FLOPs/s on A100 GPUs, versus 3–10× faster than a standard (non-fused) attention implementation. This made FA2 the default choice integrated into PyTorch’s scaled_dot_product_attention, Hugging Face Transformers/Diffusers, and xFormers.

5.3 FlashAttention-3

FlashAttention-3 (Shah, Bikshandi, Dao et al., 2024) targets NVIDIA Hopper GPUs (H100/H800) specifically, exploiting hardware features not available on older architectures:

• Warp specialization and asynchrony. FA3 overlaps the (asynchronous) memory-copy and matmul instructions with the (synchronous) softmax computation, using Hopper’s asynchronous Tensor Memory Accelerator (TMA) and warp-group matrix multiply instructions (WGMMA) so that data movement and floating-point work happen concurrently rather than serially.
• TMA offload. The Tensor Memory Accelerator handles HBM↔︎SRAM data movement and index/bounds calculations in hardware, freeing up GPU registers that can then be used to increase tile size and further improve efficiency.
• Low-precision (FP8) support. FA3 adds block quantization and incoherent processing techniques so that attention can run in FP8 with much smaller accuracy loss than a naive FP8 cast would produce.

FlashAttention-3 requires an H100/H800-class GPU and CUDA ≥ 12.3 (12.8 recommended). Reported results include roughly 1.5–2.0× speedup over FlashAttention-2 on H100 in FP16, reaching up to ~740 TFLOPs/s (about 75% of the H100’s theoretical FP16 peak), and close to 1.2 PFLOPs/s when using FP8. FlashAttention-2 remains the appropriate choice on Ampere/Ada GPUs (A100, RTX 30/40-series), since FA3’s Hopper-specific hardware paths simply aren’t available there.

6.1 Vision Transformers (ViT)

A ViT tokenizes an image into a grid of patches (commonly \(14 \times 14\) or \(16 \times 16\) pixels), flattens each into an embedding, and runs full bidirectional self-attention across all patches at every layer. Unlike causal language models, ViT attention has no causal mask, so all \(N^2\) pairwise interactions are always computed — there’s no “skip the upper triangle” saving available. This makes ViTs particularly memory- and bandwidth-hungry at higher resolutions or with masked-autoencoder-style pretraining that uses large token counts. FlashAttention’s linear memory scaling directly enables training and fine-tuning ViTs at higher resolution, larger batch size, or with more patches (e.g., smaller patch size for finer-grained tokenization) than would otherwise fit in GPU memory.

6.2 Diffusion Models: U-Net Cross-Attention and Diffusion Transformers (DiT)

Text-to-image diffusion models (e.g., Stable Diffusion) use two kinds of attention repeatedly at every denoising step: self-attention over spatial locations within the image’s latent representation, and cross-attention between the image latents (queries) and text-embedding tokens from a language encoder (keys/values). Because denoising runs for many steps (often 20–50+), even small per-step attention costs multiply substantially over a full generation. FlashAttention (via the xFormers memory-efficient attention op, and later native PyTorch SDPA) was one of the first widely adopted speedups for Stable Diffusion inference, commonly cited as delivering up to ~100% (2×) throughput improvements and comparable memory reductions when swapped in for the default attention implementation.

Newer generation architectures increasingly replace the convolutional U-Net with a Diffusion Transformer (DiT) — a pure transformer operating directly on patchified latent tokens (used in Stable Diffusion 3, PixArt, Sora-style video models, and others). DiTs push token counts even higher than U-Nets (especially for high-resolution images, e.g., up to 4K in some PixArt variants, or multi-frame video), which makes the quadratic cost of self-attention the primary scaling bottleneck of the whole model. FlashAttention is effectively a prerequisite for DiT-based models to scale to high resolutions and long video sequences at all; without it, both memory and latency would grow unacceptably with resolution.

Cross-attention layers in diffusion models are slightly different from self-attention in that \(Q\) comes from one sequence (image latents) while \(K\) and \(V\) come from another (text tokens), so the two sequence lengths can differ. FlashAttention’s tiling scheme handles this transparently — the query and key/value blocks simply come from different tensors of possibly different lengths, with no algorithmic change required.

6.3 Multimodal / Vision-Language Models (CLIP and similar)

Contrastive and multimodal models like CLIP typically run a vision encoder (often a ViT) and a text encoder somewhat independently, so the direct win is the same ViT-side benefit described above. In more tightly fused vision-language architectures that interleave image and text tokens in a shared attention stack, FlashAttention’s ability to handle arbitrary (and even variable-length, padded, or block-diagonal) sequences efficiently is what makes joint image+text attention practical at scale, since the combined sequence length (image patches + text tokens) can be substantial.

6.4 Summary of the Vision-Specific Payoff

Across all these architectures, the pattern is the same: image models tend to produce long token sequences (patch grids, especially at high resolution), and attention layers in vision models are frequently fully bidirectional with no causal mask to exploit, so the full \(N^2\) cost is always paid. FlashAttention’s linear-memory, IO-aware computation is what makes it feasible to (a) increase resolution, (b) shrink patch size for finer detail, (c) increase batch size during training, and (d) run more diffusion steps or larger guidance/conditioning setups, all without hitting a memory wall or paying an outsized latency penalty.

7.1 PyTorch’s Built-in scaled_dot_product_attention (Easiest)

Since PyTorch 2.0, torch.nn.functional.scaled_dot_product_attention (SDPA) automatically dispatches to a fused, hardware-appropriate backend — including a FlashAttention kernel, a memory-efficient (xFormers-style) kernel, or a plain math fallback — depending on your GPU, dtype, and input shapes. This is the recommended default for most users because it requires no extra dependency and PyTorch handles backend selection for you.

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

class ViTSelfAttention(nn.Module):
    """Multi-head self-attention block for a Vision Transformer,
    using PyTorch's fused SDPA (dispatches to FlashAttention on
    supported hardware automatically)."""

    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = True):
        super().__init__()
        assert dim % num_heads == 0
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.proj = nn.Linear(dim, dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, C = x.shape  # batch, num_patches (+cls token), embed dim
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, B, heads, N, head_dim)
        q, k, v = qkv.unbind(0)

        # SDPA picks the best backend (FlashAttention / memory-efficient / math)
        # for the current device, dtype, and shapes. No causal mask for ViT.
        out = F.scaled_dot_product_attention(q, k, v, is_causal=False)

        out = out.transpose(1, 2).reshape(B, N, C)
        return self.proj(out)

To confirm which backend is actually being used (useful for debugging performance), you can scope the call with an explicit backend selector:

from torch.nn.attention import SDPBackend, sdpa_kernel

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    out = F.scaled_dot_product_attention(q, k, v)

7.2 The flash-attn Library (Explicit, Most Control)

For direct access to FlashAttention kernels — including features like variable-length (unpadded) batches, which are useful when batching images of different aspect ratios/token counts, or ALiBi/rotary position handling — install Tri Dao’s reference implementation:

pip install flash-attn --no-build-isolation

Then call it directly, using half precision (fp16/bf16), which the kernels require:

import torch
from flash_attn import flash_attn_func

# q, k, v: (batch, seqlen, num_heads, head_dim), dtype fp16 or bf16, on CUDA
B, N, H, D = 8, 4096, 12, 64  # e.g. a DiT-style block at higher resolution
q = torch.randn(B, N, H, D, device="cuda", dtype=torch.bfloat16)
k = torch.randn(B, N, H, D, device="cuda", dtype=torch.bfloat16)
v = torch.randn(B, N, H, D, device="cuda", dtype=torch.bfloat16)

out = flash_attn_func(q, k, v, dropout_p=0.0, causal=False)
# out: (B, N, H, D) -- identical (up to floating point tolerance) to
# standard softmax(QK^T / sqrt(D))V, computed with far less HBM traffic.

For diffusion cross-attention (differing Q and K/V sequence lengths), flash-attn also exposes flash_attn_func with independent query and key/value tensors of different lengths, or flash_attn_varlen_func for packed variable-length batches — handy for batching a mix of image resolutions or aspect ratios without padding waste.

7.3 xFormers memory_efficient_attention (Diffusion-Model Ecosystem)

Many diffusion pipelines (e.g., Hugging Face diffusers) historically defaulted to xFormers’ memory_efficient_attention, which wraps FlashAttention-style and other fused kernels behind a single call:

# pip install diffusers xformers
from diffusers import StableDiffusionPipeline
import torch

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
).to("cuda")

pipe.enable_xformers_memory_efficient_attention()
image = pipe("a watercolor painting of a lighthouse at dusk").images[0]

On recent diffusers/PyTorch versions, SDPA (section 7.1) is enabled automatically and typically matches xFormers’ speed and memory profile, making the explicit xFormers call less necessary than it once was — but it remains a good fallback for older stacks or GPUs where SDPA’s fused backends aren’t selected.

7.4 Minimal Benchmark Snippet

A quick way to sanity-check FlashAttention’s benefit on your own hardware — useful before committing a training run:

import torch, time
import torch.nn.functional as F

def bench(fn, q, k, v, iters=50):
    torch.cuda.synchronize()
    start = time.time()
    for _ in range(iters):
        fn(q, k, v)
    torch.cuda.synchronize()
    return (time.time() - start) / iters

B, H, N, D = 4, 12, 4096, 64  # N=4096 patches: e.g. a 1024x1024 image, 16px patches
q = torch.randn(B, H, N, D, device="cuda", dtype=torch.bfloat16)
k = torch.randn(B, H, N, D, device="cuda", dtype=torch.bfloat16)
v = torch.randn(B, H, N, D, device="cuda", dtype=torch.bfloat16)

def naive_attn(q, k, v):
    scale = q.shape[-1] ** -0.5
    attn = (q @ k.transpose(-2, -1) * scale).softmax(dim=-1)
    return attn @ v

def flash_attn(q, k, v):
    return F.scaled_dot_product_attention(q, k, v)

print("naive :", bench(naive_attn, q, k, v), "s/iter")
print("flash :", bench(flash_attn, q, k, v), "s/iter")

At \(N = 4096\) (roughly what a \(1024 \times 1024\) image at 16px patches produces), expect the naive path to be both markedly slower and to use substantially more peak memory (visible via torch.cuda.max_memory_allocated()), illustrating exactly the \(O(N^2)\) vs. \(O(N)\) memory gap discussed in Section 2.