Understanding FlashAttention Pt 1: Personal Notes

Sep 14, 2026 09:01 PM - 1 hour ago 1

0. Introduction

How IO-Aware Attention Makes Transformers Faster Without Approximating Attention

The mechanism, successful 3 words: Tiling + Online Softmax + Recomputation. Everything successful this handbook is elaboration connected that summary.

A method handbook connected nonstop tiled attention: GPU representation traffic, online softmax, guardant and backward passes, IO complexity, the improvement from FlashAttention-1 done FlashAttention-4, and existent model behavior.

0.1 How to Read This Handbook

This handbook was inspired by this tweet. Before the fix, present is what the modular attraction implementation looks like. Load $Q, K, V \in \mathbb{R}^{N \times d}$ successful HBM, then:

Read $Q, K$ from HBM, compute $S$, constitute $S$ to HBM.
Read $S$ from HBM, compute $P$, constitute $P$ to HBM.
Read $P, V$ by blocks from HBM, compute $O$, constitute $O$ to HBM.
Return $O$.

What stands retired to maine is the number of information trips to HBM. Every intermediate worth — $S$, $P$, $O$ — has to beryllium written retired and publication back. That is the problem FlashAttention is solving.

The handbook itself frames the taxable arsenic easiest to understand erstwhile 3 different questions are kept separate:

What mathematical usability is being computed? For dense attention, the target remains mean scaled dot-product attention.
How overmuch arithmetic does that usability require? Dense all-pairs query-key scoring remains quadratic successful series length.
How does the implementation move information done the GPU representation hierarchy? This is wherever FlashAttention changes the algorithmic execution dramatically.

The cardinal instruction I return from this framing is that wall-clock velocity is not wished by FLOP count alone. An algorithm tin execute fundamentally the aforesaid mathematical work, aliases moreover recompute intermediate values, and still tally faster because it moves acold little information to and from high-bandwidth memory.

Core distinction. Dense FlashAttention is an exact attraction algorithm: it does not switch softmax attraction pinch a low-rank, sparse, kernelized, aliases approximate formula. “Exact” refers to the mathematical attraction computation. Floating-point kernels tin still disagree by mini rounding effects because operations are reordered.

The connection exact is doing existent activity here. Exactness is simply a connection astir the mathematical function, not astir bitwise reproducibility. The kernel is free to reorder floating-point operations. It is not free to alteration the usability being computed.

0.2 Notation

For 1 attraction head, let

\[Q \in \mathbb{R}^{N_q \times d}, \quad K \in \mathbb{R}^{N_k \times d}, \quad V \in \mathbb{R}^{N_k \times d_v}.\]

For self-attention, typically $N_q = N_k = N$. The scaled people matrix is

\[S = \frac{QK^T}{\sqrt{d}} + B,\]

where $B$ represents an optional additive disguise aliases bias, and

\[P = \mathrm{softmax}\_{\mathrm{row}}(S), \quad O = PV.\]

Throughout, HBM refers to ample off-chip high-bandwidth GPU memory. On-chip representation is simply a wide school word for overmuch smaller, faster retention specified arsenic registers and shared memory/SRAM. Exact hardware specifications alteration by GPU generation.

The applicable quality I support coming backmost to:

HBM SRAM

slow	faster
large	smaller
off-chip	on-chip

System problem: wherever do each those intermediate values $(S, P, O)$ unrecorded while the GPU computes them?

That mobility — not the arithmetic — is what FlashAttention was built to answer.

Part 1: The Fundamental Problem

1.1 What FlashAttention really optimizes
1.2 The attraction equation is not the implementation
1.3 GPU representation level and why IO matters
1.4 Why materializing $S$ and $P$ is expensive
1.5 Dense arithmetic is still quadratic
1.6 Memory-efficient nonstop attraction predates FlashAttention

Part 2: The Mathematical Trick

2.1 Tiling queries, keys, and values
2.2 Softmax is the difficult portion of streaming
2.3 Online softmax from first principles
2.4 The blockwise merge recurrence
2.5 A complete numerical example

Part 3: Putting the Mathematics onto the GPU

3.1 FlashAttention guardant pass
3.2 Why dense FlashAttention is exact
3.3 IO complexity: what the theorem really says
3.4 Memory complexity: linear auxiliary state, not linear compute
3.5 Causal masking and tile skipping
3.6 Backward pass: recompute alternatively of save
3.7 Why much FLOPs tin still beryllium faster

Part 4: Architectural Compatibility

4.1 MHA, MQA, and GQA compatibility
4.2 Variable lengths, section attention, and dropout

Part 5: FlashAttention Evolution

5.1 FlashAttention-2: what changed
5.2 FA2 parallelism crossed series tiles
5.3 FA2 activity partitioning and non-matmul FLOPs
5.4 FlashAttention-3: the Hopper generation
5.5 FA3 asynchrony: overlap information movement, GEMM, and softmax
5.6 FA3 FP8: capacity without pretending precision is free
5.7 FlashAttention-4: the Blackwell generation
5.8 FA4 and asymmetric hardware scaling
5.9 FA4 implementation and existent status
5.10 FlashAttention-1 done -4 compared

Part 6: Using FlashAttention successful Frameworks

6.1 PyTorch scaled-dot-product attraction today
6.2 Exactness is not bitwise identity

Part 7: FlashAttention vs. Other Techniques

7.1 FlashAttention vs. PagedAttention
7.2 FlashAttention vs. sparse and linear attention

Part 8: Training vs. Inference

8.1 Training, prefill, and decode are different regimes

Part 9: Practical Engineering

9.1 Common implementation mistakes
9.2 Common misconceptions and the applicable intelligence model

Part 1: The Fundamental Problem

1.1 What FlashAttention Actually Optimizes

Start pinch the mean attraction function:

\[O = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}} + B\right) V.\]

FlashAttention is IO-aware. My moving definition:

Minimize information activity betwixt the different levels of GPU memory, alternatively than conscionable trying to trim the number of mathematical operations (FLOPs).

The velocity bottleneck successful modern AI hardware is often not really accelerated the GPU tin compute math, but really accelerated it tin read and constitute data. This is the memory-compute tradeoff.

A textbook implementation often makes this look for illustration 3 ample operations:

Form the people matrix $S$, use row-wise softmax to get $P$, past multiply by $V$. Mathematically that is fine. On a GPU, however, penning a immense intermediate matrix to HBM and reference it backmost tin beryllium acold much costly than the equation suggests.

FlashAttention’s cardinal publication is to make the algorithm IO-aware. It partitions the computation into tiles that fresh successful accelerated on-chip memory, streams blocks of $K$ and $V$, and maintains capable row-wise softmax authorities to nutrient the nonstop output without materializing the afloat $N \times N$ attraction matrix successful HBM.

What changes: the execution schedule, representation traffic, and stored intermediates.

What does not change: the dense scaled-dot-product attraction usability being evaluated.

This favoritism is why “FlashAttention is simply a faster benignant of attention” tin beryllium misleading. It is amended thought of arsenic an algorithm and kernel family for evaluating attraction efficiently connected accelerators. A exemplary tin usage causal masking, RoPE, MQA/GQA, aliases different attraction features and still usage a FlashAttention implementation underneath.

The original insubstantial contrasts this attack pinch approximate attraction methods that trim arithmetic by changing the mathematical problem. Dense FlashAttention does not make that trade. The aforesaid insubstantial besides introduced a block-sparse extension, but that sparse hold is simply a different lawsuit because omitting blocks changes which interactions are computed.

1.2 The Attention Equation Is Not the Implementation

The equation does not show you wherever tensors live.

A standard naive attraction implementation:

Calculate $S$, shop $S$.
Read $S$, cipher $P$, shop $P$.
Read $P$, cipher $O$.

The problem is the number of HBM information trips. A elemental materializing implementation does:

\[S \leftarrow QK^T / \sqrt{d}, \quad P \leftarrow \mathrm{softmax}(S), \quad O \leftarrow PV.\]

If $S$ is written to HBM aft the first matrix multiplication, publication for softmax, $P$ is written back, and past $P$ is publication again for the $PV$ multiplication, the GPU spends important clip moving an $N^2$ entity done memory.

FlashAttention avoids many times moving immense intermediates:

Calculate mini $S$ tile, softmax tile, usage tile pinch $V_j$, discard tile.
Calculate adjacent tile.

The cardinal questions I inquire erstwhile looking astatine immoderate equation:

How galore operations are required?
What information must move betwixt representation levels to execute those operations?

FlashAttention reuses the aforesaid dependencies. Blocks of $Q$, $K$, and $V$ are brought adjacent the compute units, people tiles are produced and consumed locally, and only compact row-wise statistic positive the output request to persist crossed tiles.

Algorithmic lesson. A computational chart is not a representation schedule. Writing $P = \mathrm{softmax}(QK^T)$ connected insubstantial does not require an implementation to shop each of $QK^T$ aliases $P$ successful off-chip representation astatine once.

This thought generalizes beyond attention. Fused kernels, tiling, recomputation, and usability scheduling often waste and acquisition a mini magnitude of other arithmetic for overmuch little activity of ample intermediates. On modern accelerators, that tin beryllium the correct waste and acquisition because matrix-multiply throughput has grown overmuch faster than galore different parts of the representation and execution hierarchy. FlashAttention-3 and -4 make that hardware dependence progressively explicit.

1.3 GPU Memory Hierarchy and Why IO Matters

GPUs expose a level alternatively than 1 azygous excavation of arsenic accelerated memory. The names and capacities alteration by architecture, but the intelligence exemplary is:

Registers / very section state
On-chip shared representation / SRAM
HBM / instrumentality memory

HBM is large, but information must recreation to the compute units. On-chip retention is overmuch smaller, but reuse location is overmuch cheaper. The original FlashAttention study models this asymmetry utilizing HBM and SRAM and explicitly optimizes the number of transfers betwixt them.

My breakdown of each level:

HBM:

Relatively large
Relatively slower to access
Physically farther from individual compute operations
Stores exemplary weights, $Q/K/V$, activations, ample tensors

SRAM / shared memory:

Much smaller
Faster
Cheaper reuse cost

Registers:

Even smaller
More local

Why tiling helps. Suppose $Q_i$ (a query tile) needs to interact pinch galore $K/V$ tiles. Instead of perpetually moving $Q_i$ backmost and forth, we tin support it adjacent to the compute units while processing:

\[K\_1, V\_1 \to K\_2, V\_2 \to K\_3, V\_3 \to \cdots\]

So 1 loaded $Q_i$ tin participate successful tons of computation. This is called reuse.

The intent of tiling is not simply “make tensors smaller.” It is to increase reuse while a tile is resident connected chip. A artifact of $Q$ tin interact pinch aggregate $K/V$ blocks earlier its partial softmax/output authorities is written back. Conversely, $K/V$ blocks tin beryllium streamed done query blocks according to the chosen schedule.

A kernel becomes IO-aware erstwhile the placement and activity of information are portion of the algorithm, alternatively than an afterthought near to a series of separately launched tensor operations.

Do not move this into a cosmopolitan motto that attraction is ever “memory-bound.” The bottleneck depends connected series length, caput dimension, dtype, disguise pattern, GPU generation, guardant vs. backward, and which kernel is running. FA3 and FA4 beryllium partially because, arsenic hardware changed, the ascendant bottlenecks changed too.

1.4 Why Materializing $S$ and $P$ Is Expensive

The quadratic intermediate becomes actual very quickly. Suppose a batch contains 1 sequence, pinch 32 attraction heads, series magnitude $N = 8192$, and a two-byte dtype specified arsenic FP16 aliases BF16. One dense tensor pinch shape

\[[1, 32, 8192, 8192]\]

contains $32 \times 8192^2$ elements. At 2 bytes per element, that is

\[32 \times 8192^2 \times 2 = 4{,}294{,}967{,}296 \text{ bytes} = 4 \text{ GiB}.\]

That is the size of 1 afloat score- aliases probability-like tensor for this example.

Concretely, my arithmetic:

Batch = 1
Heads = 32
seq_len = 8192
dtype = FP16/BF16
2 bytes/elem
tensor style = $[1, 32, 8192, 8192]$
$(32)(8192)(8192)$ elem $\times$ (2 bytes/elem) $= 4{,}294{,}967{,}296$ bytes $= 4$ GiB

Imagine reference and penning this overmuch data.

A naive decomposition whitethorn nutrient arrays of this standard astatine aggregate stages. This does not mean each modern model keeps some $S$ and $P$ live simultaneously, and compiler fusion tin already debar immoderate traffic. The illustration illustrates the basal problem: a dense $N^2$ intermediate is ample capable that many times penning and rereading it tin predominate representation usage and bandwidth.

Naive pipeline:

\[QK^T \to \text{huge } S \to \text{softmax} \to \text{huge } P \to PV\]

FlashAttention pipeline:

\[Q\_i K\_j^T \to \text{softmax tile} \to \text{tile} \times V\_j \to \text{discard}\]

FlashAttention avoids storing the afloat matrix successful HBM. It forms people tiles, applies the softmax update while the tile is connected chip, instantly uses those probabilities to accumulate the corresponding publication from $V$, past discards the tile.

Important wording: FlashAttention removes the request to materialize the afloat attraction matrix arsenic an off-chip intermediate. It does not region the logical pairwise interactions required by dense attention.

Important distinction. FlashAttention acknowledges that $N^2$ interactions beryllium and says: do not shop the full consequence of these interactions arsenic a elephantine intermediate if we tin devour each portion immediately.

This is besides why the representation use is particularly important during training, wherever naive autograd would different want ample intermediates for the backward pass.

1.5 Dense Arithmetic Is Still Quadratic

FlashAttention changes the representation schedule, not the mathematical function. The usability is still dense attention. That means:

For self-attention pinch $N$ tokens and caput magnitude $d$, forming each query-key scores requires activity proportional to $N^2 d$.

Multiplying the probabilities by values adds different dense pairwise matrix multiplication of the aforesaid wide order. FlashAttention reorganizes these operations, but it does not extremity evaluating the dense group of query-key interactions.

So 3 statements must not beryllium confused:

Arithmetic complexity: dense attraction remains $\mathcal{O}(N^2 d)$.
Large intermediate storage: FlashAttention avoids an $\mathcal{O}(N^2)$ materialized score/probability tensor successful HBM.
HBM traffic: the original insubstantial proves a little IO costs nether its two-level representation exemplary than modular materializing attention.

This is really a method tin make overmuch longer sequences applicable without making agelong discourse “free.” Doubling $N$ still astir quadruples the number of dense query-key pairs. FlashAttention chiefly attacks the data-movement and memory-footprint broadside of that computation.

If you want to trim the number of query-key pairs themselves, you request a different mathematical structure: for illustration sparsity, a section pattern, aliases a different attraction formulation. Those choices tin alteration exemplary behaviour and are conceptually abstracted from dense FlashAttention.

When reporting speedups, ever separate asymptotic arithmetic from measured runtime. A kernel tin go respective times faster astatine the aforesaid $\mathcal{O}(N^2 d)$ complexity because the changeless factors, occupancy, fusion, and representation postulation alteration dramatically.

My side-by-side summary:

Quality Dense FlashAttention

Arithmetic	$\mathcal{O}(N^2 d)$
Full attraction intermediates	Avoid $\mathcal{O}(N^2)$ storage
HBM traffic	Reduced substantially

1.6 Memory-Efficient Exact Attention Predates FlashAttention

It would beryllium historically inaccurate to opportunity FlashAttention first discovered that nonstop attraction tin debar quadratic memory.

Rabe and Staats showed earlier FlashAttention that attraction request not require $\mathcal{O}(N^2)$ representation pinch respect to series length. Their activity gave nonstop memory-efficient algorithms while retaining quadratic time, and a applicable accelerator implementation pinch subquadratic memory.

FlashAttention’s chopped publication was to move representation ratio into an definitive IO-aware GPU algorithm: tile the computation against the accelerator representation hierarchy, analyse HBM accesses, fuse the applicable operations, and show important wall-clock gains.

Online softmax besides has an earlier lineage. Milakov and Gimelshein described an online recurrence that computes the classical unchangeable softmax normalizer pinch less representation accesses. FlashAttention builds the aforesaid benignant of running-max/running-normalizer thought into tiled attention, while besides accumulating the value-weighted output.

A precise lineage is therefore:

stable / online softmax provides a streaming normalization tool,
earlier memory-efficient attraction shows nonstop attraction request not shop $N^2$ state,
FlashAttention co-designs tiling, softmax state, and GPU IO to make the attack accelerated successful practice.

The lineage, successful my shorthand:

\[\text{Stable softmax} \to \text{Online normalization} \to \text{Exact memory-efficient attention} \to \text{FlashAttention} \to \text{FlashAttention-2} \to \text{FlashAttention-3} \to \text{FlashAttention-4}\]

Rabe and Staats demonstrated nonstop memory-efficient attraction pinch quadratic computation but subquadratic representation (the earlier work). There is besides earlier activity connected online softmax by Milakov and Gimelshein.

FlashAttention’s cardinal publication was to bring together:

Tiling
Online softmax
Fused computation
GPU representation level awareness
IO complexity analysis

into a applicable high-performance algorithm.

This favoritism matters because “memory-efficient” does not automatically mean “IO-optimized for a peculiar hardware model.”

Part 2: The Mathematical Trick

2.1 Tiling Queries, Keys, and Values

Most important. This is the conception wherever the instrumentality lives.

Suppose $Q \in \mathbb{R}^{N_q \times d}$, $K \in \mathbb{R}^{N_k \times d}$, and $V \in \mathbb{R}^{N_k \times d_v}$. Instead of processing everything astatine once, disagreement them into blocks. For example:

\[Q: [Q\_1, Q\_2, Q\_3], \quad K, V: [K\_1, V\_1], [K\_2, V\_2], [K\_3, V\_3]\]

For 1 pair,

\[S\_{ij} \doteq \frac{Q\_i K\_j^T}{\sqrt{d}} + B\_{ij},\]

we process that people tile locally.

Then, for each query block:

Calculate scores
Softmax them
Multiply by $V_j$
Update the moving output
Discard the tile

The softmax utilized astatine each tile is

\[\mathrm{softmax}(x\_i) = \frac{e^{x\_i}}{\sum\_j e^{x\_j}}.\]

Matrix multiplication is easy to tile because it is fundamentally accumulation:

\[AB = \sum\_j A\_j B\_j,\]

so we tin cipher pieces and adhd them.

Softmax is harder because each constituent depends connected the full row. If we process the first block, we do not cognize the eventual denominator. Even worse, numerical stableness requires knowing the maximum.

Instead of forming each $N_q N_k$ scores astatine once, partition the matrices into blocks. For 1 query artifact $Q_i$ and 1 key/value artifact $(K_j, V_j)$, compute

\[S\_{ij} = \frac{Q\_i K\_j^T}{\sqrt{d}} + B\_{ij}.\]

The people tile $S_{ij}$ is mini capable to beryllium processed adjacent the compute units. Its publication is folded into moving statement statistic and an output accumulator, past the tile tin beryllium discarded.

Conceptually, for each query block:

load a tile of $Q_i$,
stream compatible $K/V$ tiles,
compute 1 people tile,
update row-wise softmax state,
accumulate the corresponding $V$ contribution,
move to the adjacent tile without penning a world $S$ aliases $P$ matrix.

Actual kernels take tile shapes and loop bid based connected on-chip capacity, caput dimension, GPU generation, causal structure, and activity partitioning. The original FlashAttention IO study uses tile sizes derived from SRAM capacity $M$.

Tiling unsocial is not enough. Matrix multiplication tiles constitute people because sums tin beryllium accumulated. Softmax couples each people successful a statement done a shared maximum and denominator, truthful we request a measurement to merge blocks without seeing the afloat statement astatine once.

That is the cardinal mathematical instrumentality successful the adjacent chapters.

2.2 Softmax Is the Difficult Part of Streaming

A numerically unchangeable softmax for 1 statement $x_1, \ldots, x_N$ uses

\[m = \max\_j x\_j, \qquad \ell = \sum\_j e^{x\_j - m}, \qquad p\_j = \frac{e^{x\_j - m}}{\ell}.\]

The subtraction by $m$ prevents overflow from ample affirmative logits. But it seems to create a streaming problem: really tin an early artifact beryllium normalized if a later artifact whitethorn incorporate a larger maximum?

The reply is to clasp capable statistic that tin beryllium rescaled erstwhile the maximum changes. Suppose the moving authorities aft earlier elements is $(m_{\text{old}}, \ell_{\text{old}})$ and a caller artifact has maximum $m_b$. Define

\[m\_{\text{new}} = \max(m\_{\text{old}}, m\_b).\]

Every publication accumulated nether the aged maximum tin beryllium converted to the caller reference by multiplying it by

\[\alpha \doteq e^{m\_{\text{old}} - m\_{\text{new}}}.\]

The caller artifact is evaluated comparative to the aforesaid $m_{\text{new}}$.

This is not an approximation. It is the identity

\[e^{x\_j - m\_{\text{old}}}\, e^{m\_{\text{old}} - m\_{\text{new}}} = e^{x\_j - m\_{\text{new}}}.\]

Mathematical insight. The moving maximum is simply a alteration of numerical reference point. When that reference changes, antecedently accumulated exponentials tin beryllium rescaled precisely successful existent arithmetic alternatively than recomputed from scratch.

The online-normalizer recurrence predates FlashAttention and provides the mathematical ground for streaming unchangeable softmax.

Worked example. Start pinch 2 utmost values to spot why the max subtraction matters:

\[x = [1000, 999].\]

Directly computing $e^{1000}$ overflows. Instead, group $m = 1000$, so

\[x - m = [1000 - 1000, \; 999 - 1000] = [0, -1],\]

and $e^{0} = 1$, $e^{-1} \approx 0.368$ are perfectly manageable.

Now watercourse 2 blocks. Let Block 1 $= [2, 1]$ and Block 2 $= [4, 3]$. After Block 1, $m_{\text{old}} = 2$. Block 1 is evaluated astatine $m_{\text{old}}$ arsenic $[e^{0}, e^{-1}]$. When Block 2 arrives pinch $m_b = 4$,

\[m\_{\text{new}} = \max(2, 4) = 4, \qquad \alpha = e^{2 - 4} = e^{-2}.\]

Rescaling Block 1 nether the caller reference:

\[[\alpha e^{0}, \; \alpha e^{-1}] = [e^{-2}, \; e^{-3}].\]

Block 2 astatine the caller reference:

\[[e^{4 - 4}, \; e^{3 - 4}] = [e^{0}, \; e^{-1}].\]

Putting it each together, pinch $x = [2, 1, 4, 3]$:

\[\mathrm{softmax}(x) = \frac{[e^{-2}, \; e^{-3}, \; e^{0}, \; e^{-1}]}{e^{-2} + e^{-3} + e^{0} + e^{-1}}.\]

The streaming computation for Block 1 is

\[\ell = \sum\_j e^{x\_j - m} = e^{0} + e^{-1} = 1 + e^{-1},\]

giving normalized outputs $[e^{0}, e^{-1}] / (1 + e^{-1})$. For Block 2, the moving normalizer is

\[\ell\_{\text{new}} = \alpha\, \ell\_{\text{old}} + \sum\_{j \in \text{Block 2}} e^{x\_j - m\_{\text{new}}} = e^{-2}(1 + e^{-1}) + e^{0} + e^{-1},\]

and the mixed numerator is $[\alpha e^{0}, \alpha e^{-1}, e^{0}, e^{-1}] = [e^{-2}, e^{-3}, e^{0}, e^{-1}]$, matching the afloat softmax computed successful 1 shot.

2.3 Online Softmax from First Principles

Process scalar logits $x_1, x_2, \ldots$ 1 astatine a time. Initialize

\[m\_0 = -\infty, \qquad \ell\_0 = 0.\]

After watching $x_j$, update

\[m\_j = \max(m\_{j-1}, x\_j), \qquad \ell\_j = \ell\_{j-1} e^{m\_{j-1} - m\_j} + e^{x\_j - m\_j}.\]

Milakov and Gimelshein show that this produces the aforesaid unchangeable softmax normalizer while requiring less passes complete the input than the accepted safe-softmax procedure.

For attention, we besides request the weighted worth sum. Introduce an unnormalized accumulator $a$:

\[a = \sum\_j e^{x\_j - m} v\_j.\]

When the maximum changes from $m$ to $m’$, rescale some $\ell$ and $a$ by $e^{m - m’}$. Then adhd the caller exponentials and worth contributions nether the caller reference. At the end,

\[o = \frac{a}{\ell}.\]

My moving type of the update:

\[a \doteq \sum\_j e^{x\_j - m} v\_j\]

When the maximum changes:

\[a\_{\text{old}} \to a\_{\text{old}}\, e^{m\_{\text{old}} - m\_{\text{new}}}\] \[\ell\_j \to \ell\_{j-1}\, e^{m\_{j-1} - m\_j} + e^{x\_j - m\_j}\]

Then adhd the caller worth contributions. At the end:

\[o = \frac{a}{\ell}.\]

For attention, $x_j$ is not a fixed input vector stored successful advance. Each artifact of logits is generated connected request from a matrix merchandise $Q K_j^T / \sqrt{d}$ positive mask/bias terms. The online recurrence lets the kernel devour that artifact immediately.

The aforesaid thought useful statement by statement and artifact by block, which is what makes a tiled nonstop softmax-attention guardant walk possible.

Mathematical penetration II. We do not request the full probability vector. We only request capable accusation to reconstruct its publication to the last output.

2.4 The Blockwise Merge Recurrence

For 1 query row, suppose the moving authorities aft immoderate cardinal blocks is

\[(m, \ell, a),\]

where $m$ is the maximum people seen truthful far, $\ell$ is the unchangeable softmax denominator nether that maximum, and $a \in \mathbb{R}^{d_v}$ is the unnormalized worth accumulator.

For a caller people artifact $s \in \mathbb{R}^b$ pinch matching values $V_b \in \mathbb{R}^{b \times d_v}$, let

\[m\_b = \max(s), \qquad m' = \max(m, m\_b),\] \[\alpha = e^{m - m'}, \qquad p = e^{s - m'}.\]

Then update

\[\ell' = \alpha \ell + \sum\_j p\_j,\] \[a' = \alpha a + p^T V\_b,\] \[m \leftarrow m', \qquad \ell \leftarrow \ell', \qquad a \leftarrow a'.\]

Finally,

\[o = a / \ell.\]

The mean of the moving state:

$m$ = max people seen truthful far
$\ell$ = unchangeable softmax denominator, exponential sum (normalizer)
$a$ = unnormalized worth accumulator

The blockwise recurrence successful my ain notation:

For 1 query row, support $(m, \ell, a)$.
Suppose the adjacent people artifact is $s = [s_1, s_2, \ldots, s_b]$ pinch corresponding values $V_b \in \mathbb{R}^{b \times d_v}$.
First cipher the artifact maximum: $m_b = \max(s)$, past update the world maximum: $m’ = \max(m, m_b)$.
Define: $\alpha = e^{m - m’}$, $p = e^{s - m’}$.
Then update: $\ell’ = \alpha \ell + \sum_j p_j$, $a’ = \alpha a + p^T V_b$.
Then group $m \leftarrow m’$, $\ell \leftarrow \ell’$, $a \leftarrow a’$.
Finally astatine the end: $o = a / \ell$.

For a artifact of query rows, $m$ and $\ell$ go row-wise vectors and $a$ becomes a matrix. Masks tin beryllium applied to the people tile earlier the exponentials, pinch masked positions contributing zero probability.

The recurrence is the algebraic logic tile boundaries do not alteration the dense softmax result. A different tiling changes the bid of floating-point operations, but not the intended real-arithmetic function.

FlashAttention’s published algorithms definitive balanced running-max / running-normalizer / output updates successful artifact form.

This recurrence is the bosom of tiled nonstop attention.

2.5 A Complete Numerical Example

Consider 1 already-scaled, unmasked attraction row

\[s = [2, 1, 4, 3]\]

with two-dimensional values

\[v\_1 = [1, 0], \quad v\_2 = [0, 1], \quad v\_3 = [2, 0], \quad v\_4 = [0, 2].\]

The world maximum is 4, truthful unchangeable unnormalized weights are

\[[e^{-2}, e^{-3}, e^{-1}, e^{-1}] \approx [0.135335, 0.049787, 1, 0.367879].\]

Their sum is

\[\ell \approx 1.55300179,\]

and afloat softmax gives

\[O \approx [1.37497284, 0.50582424].\]

Now process 2 blocks. For $[2, 1]$:

\[m\_1 = 2, \qquad \ell\_1 = 1 + e^{-1} = 1.36787944, \qquad a\_1 = [1, e^{-1}].\]

For the 2nd artifact $[4, 3]$, the caller maximum is 4, so

\[\alpha = e^{2 - 4} = e^{-2}.\]

Then

\[\ell\_2 = \alpha \ell\_1 + 1 + e^{-1} = 1.55300179,\] \[a\_2 = \alpha a\_1 + [2, 2 e^{-1}] \approx [2.13533528, 0.78554595].\]

Therefore

\[a\_2 / \ell\_2 \approx [1.37497284, 0.50582424],\]

matching the full-row computation.

The aged artifact was not revisited. Its publication was simply rescaled erstwhile a larger maximum appeared.

Tying everything together. Block 1 was ne'er recomputed erstwhile artifact 2 revealed a larger maximum; we rescaled the statistic from artifact 1. This is what makes streaming possible.

My step-by-step computation:

The people row: $s = [2, 1, 4, 3]$. Values: $v_1 = [1, 0]$, $v_2 = [0, 1]$, $v_3 = [2, 0]$, $v_4 = [0, 2]$.

Blocks: $[2, 1]$ and $[4, 3]$.

Block 1. Scores: $[2, 1]$. Maximum: $m_1 = 2$. Stable exponentials:

\[[e^{0}, e^{-1}] = [1, e^{-1}]\]

\[\ell\_1 = 1 + e^{-1} \approx 1.36788.\]

Value accumulator:

\[a\_1 = 1 \cdot [1, 0] + e^{-1} \cdot [0, 1] = [1, e^{-1}].\]

Block 2. Scores: $[4, 3]$. $m_b = 4$. $m’ = \max(2, 4) = 4$. $\alpha = e^{m - m’} = e^{2 - 4} = e^{-2}$.

\[\ell\_2 = \alpha \ell\_1 + e^{0} + e^{-1} = e^{-2}(1 + e^{-1}) + 1 + e^{-1} = e^{-2} + e^{-3} + 1 + e^{-1} \approx 1.553\] \[a\_2 = e^{-2} \cdot [1, e^{-1}] + [2 e^{0}, 2 e^{-1}] = [2 + e^{-2}, 2 e^{-1} + e^{-3}] \approx [2.135, 0.7855]\]

Output.

\[O\_2 = a\_2 / \ell\_2 \approx [1.375, 0.506]\]

Part 3: Putting the Mathematics onto the GPU

3.1 FlashAttention Forward Pass

A useful conceptual guardant walk is:

Partition $Q$ into query-row tiles and $K, V$ into key/value tiles.
For each query tile, initialize row-wise moving maxima, normalizers, and output accumulators.
Load a key/value tile and shape the section people tile $Q_i K_j^T / \sqrt{d}$.
Apply causal/local masks aliases additive biases that beryllium to this tile.
Compute the tile maximum, update the moving maximum, and rescale erstwhile state.
Exponentiate the existent tile comparative to the updated maximum.
Update the denominator and the value-weighted output accumulator.
Continue until each required cardinal tile has been contributed.
Normalize the accumulator row-wise and constitute the output.

The original FlashAttention algorithm chooses artifact sizes truthful the applicable tiles and authorities fresh successful on-chip SRAM, reducing trips to HBM.

An acquisition implementation tin reproduce the algebra successful a fewer lines of PyTorch, but specified codification is not a high-performance FlashAttention kernel. Production implementations dangle connected GPU-specific tiling, thread/warp scheduling, asynchronous copies, tensor-core instructions, and different low-level details.

The basal algorithmic thought is independent of 1 CUDA kernel: make a people tile, devour it instantly done online softmax and $V$ accumulation, and ne'er materialize the complete score/probability matrix successful HBM.

Minimal acquisition PyTorch.

def tiled_attention(q, k, v, block=128): scale = 1 / math.sqrt(q.shape[-1]) n_q, n_kv = q.shape[0], k.shape[0] O = torch.zeros_like(q) for i in range(0, n_q, block): q_blk = q[i:i+block] m = torch.full((q_blk.shape[0],), -float('inf')) l = torch.zeros((q_blk.shape[0],)) a = torch.zeros((q_blk.shape[0], v.shape[-1])) for j in range(0, n_kv, block): k_blk = k[j:j+block] v_blk = v[j:j+block] s = (q_blk @ k_blk.T) * scale m_new = torch.maximum(m, s.max(dim=-1).values) p = torch.exp(s - m_new[:, None]) alpha = torch.exp(m - m_new) l = alpha * l + p.sum(dim=-1) a = alpha[:, None] * a + p @ v_blk m = m_new O[i:i+block] = a / l[:, None] return O

Conceptually, for each KV block:

scores = $Q \times K_{block}^T$
update moving softmax, $\ell$
update output accumulator, $a$

This isn’t a high-performance FlashAttention kernel. Real implementations additionally exploit:

GPU-specific tiling
Registers
Shared memory
Warp scheduling
Tensor cores
Async copies
Specialized intrinsics
Occupancy optimization

This favoritism will go very important for FA-2/3/4.

Later FlashAttention generations support this semantic building while changing really activity is scheduled connected newer hardware.

My measurement database for the guardant pass:

Load $Q_i$.
Initialize $(m, \ell, a) = (-\infty, 0, 0)$ — $m = -\infty$, $\ell = 0$, $a = 0$.
Load a $K_j, V_j$ tile.
Compute the section people tile $S_{ij} = Q_i K_j^T / \sqrt{d}$.
Apply masks/biases $B_{ij} / M_{ij}$.
Find the tile maximum and update the moving maximum $\max(m, m_b)$.
Compute exponentials comparative to the caller maximum.
Update $\ell$ and $a$.
Move to the adjacent $K, V$ tile until each tile is completed/covered.
Normalize the accumulator row-wise $O = a / \ell$, past constitute $O$.

3.2 Why Dense FlashAttention Is Exact

Does FlashAttention approximate attention? For dense FlashAttention, no. The target remains

\[\mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}} + B\right) V.\]

My checklist for why this is exact:

No Q-K pairs are intentionally removed.
No low-rank approximation is introduced.
No replacement kernelized attraction usability replaces softmax.
This algorithm simply processes the aforesaid interactions successful blocks.

But exact has a qualification. Let maine illustrate.

Suppose 2 implementations compute $a + b + c$:

A: $(a + b) + c$
B: $a + (b + c)$

In nonstop mathematics,

\[(a + b) + c = a + (b + c).\]

But successful floating points, successful immoderate cases,

\[(a + b) + c \neq a + (b + c).\]

Tiling changes the simplification order. Fusion tin alteration the rounding behavior. Therefore:

\[\text{exact algorithm} \neq \text{bitwise identical output}.\]

This favoritism becomes important successful numerical testing.

Dense FlashAttention is called nonstop because the target usability remains

\[\mathrm{softmax}(QK^T / \sqrt{d} + B) V.\]

Tiling does not delete query-key pairs. Online softmax does not switch the exponential aliases normalization pinch different function. The artifact recurrence simply changes the bid successful which capable statistic are accumulated.

Three caveats support the connection exact precise:

Floating constituent is finite precision. Reordering additions and reductions tin nutrient mini numerical differences from different implementation. PyTorch explicitly warns that SDPA backends tin disagree because floating-point operations are fused and ordered differently.
Dropout is stochastic during training. Comparing 2 runs bit-for-bit requires matching random behaviour successful summation to mathematical attraction semantics.
Sparse variants are different. The original FlashAttention insubstantial besides presents block-sparse FlashAttention, which omits blocks and is truthful an approximate/sparse version comparative to afloat dense attention.

“Exact” does not mean “bitwise identical to each reference kernel.” It intends the algorithm is not intentionally changing dense softmax attraction to trim the mathematical work.

This favoritism matters erstwhile evaluating numerical tests. A sensible tolerance depends connected dtype, accumulation order, series length, and backend alternatively than requiring binary identity.

3.3 IO Complexity: What the Theorem Actually Says

Let maine get much theoretical.

The original FlashAttention insubstantial analyzes a two-level representation exemplary pinch HBM and on-chip SRAM of size $M$. Under the paper’s assumptions, including caput magnitude $d$ and

\[d \leq M \leq Nd,\]

standard materializing attraction requires approximately

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

HBM accesses, whereas FlashAttention requires

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

HBM accesses nether the specified regime.

The nonstop theorem has assumptions, truthful do not construe this arsenic “FlashAttention ever moves precisely this galore bytes.” It is an asymptotic consequence for a peculiar representation model.

Why does larger SRAM help? Suppose you person much on-chip memory. You tin fresh larger tiles. Larger tiles mean:

more information stays resident
more reuse
fewer HBM reloads

So expanding $M$ tin alteration the magnitude of HBM traffic. This is 1 logic GPU architecture matters truthful overmuch to FlashAttention performance.

The original FlashAttention insubstantial analyzes a two-level representation exemplary pinch HBM and on-chip SRAM of size $M$. Under the paper’s assumptions, including caput magnitude $d$ and $d \leq M \leq Nd$, modular materializing attraction requires $\Theta(Nd + N^2)$ HBM accesses, whereas FlashAttention requires $\Theta(N^2 d^2 / M)$ HBM accesses. The insubstantial besides proves an optimality consequence complete a scope of SRAM sizes successful this model.

Several specifications matter:

These are IO-complexity results successful a peculiar representation model, not a cosmopolitan byte count for each GPU.
The amount counts activity of scalar elements/words betwixt the modeled representation levels, up to asymptotic factors.
It is not the arithmetic complexity. Dense attraction still performs $O(N^2 d)$ work.
Increasing usable on-chip representation $M$ enables much reuse and reduces modeled HBM traffic.

A communal incorrect summary is “FlashAttention reduces attraction IO from $O(N^2)$ to $O(N)$.” The linear amount is the ample auxiliary representation footprint pinch respect to series length, not the wide HBM-access look above.

Intuitively, much usable on-chip representation lets larger moving tiles enactment resident and beryllium reused for much attraction activity earlier information must beryllium reloaded from HBM. In the theorem’s regime, that accrued reuse is why the HBM-access bound decreases arsenic $M$ grows.

3.4 Memory Complexity: Linear Auxiliary State, Not Linear Compute

Another important distinction: linear auxiliary authorities is not linear compute.

A naive attraction implementation whitethorn create $\mathcal{O}(N^2)$ attraction intermediates. FlashAttention does not. Instead it maintains:

$Q/K/V$ tiles
row-wise $m$
row-wise $\ell$
output accumulator $a$

The auxiliary attraction authorities scales astir linearly pinch series length, but linear representation is not linear computation. The computation is still

\[\mathcal{O}(N^2 d).\]

Example. For $N = 10{,}000$, location are approximately

\[10{,}000^2 = 100{,}000{,}000\]

Q-K interactions. FlashAttention does not make those disappear. However, it prevents you from needing a gigantic intermediate containing each of them.

Why do FlashAttention papers often opportunity representation becomes linear alternatively of quadratic?

The inputs and output already incorporate $O(Nd)$ elements. A materializing dense attraction implementation additionally creates $O(N^2)$ score/probability state. FlashAttention avoids storing those afloat matrices and retains only tiled moving authorities positive row-wise statistics. As a result, the other representation associated pinch the attraction cognition scales linearly pinch series magnitude alternatively than quadratically.

For training, the quality is particularly important because a straightforward backward walk mightiness different prevention a afloat probability matrix $P$. FlashAttention alternatively saves compact accusation specified arsenic output and row-wise normalization statistics, past recomputes score/probability tiles successful backward.

Linear representation does not connote linear runtime. The number of dense query-key interactions is still quadratic successful $N$.

Also debar claiming that full exemplary representation is $O(N)$. Other Transformer components devour activation memory, and autoregressive serving has KV-cache representation that grows pinch discourse length. FlashAttention is specifically changing really the attraction computation manages its intermediates.

One logic FlashAttention and activation checkpointing tin coexist: some waste and acquisition recomputation for reduced stored state, but astatine different scopes of the training graph.

3.5 Causal Masking and Tile Skipping

Consider autoregressive attention. Token $i$ cannot be to early token $j > i$.

The attraction matrix looks like

\[\begin{bmatrix} \checkmark & \times & \times \\ \checkmark & \checkmark & \times \\ \checkmark & \checkmark & \checkmark \end{bmatrix}\]

with mask

\[M = \begin{bmatrix} 0 & -\infty & -\infty \\ 0 & 0 & -\infty \\ 0 & 0 & 0 \end{bmatrix}.\]

Applying the disguise earlier softmax gives

\[\mathrm{softmax}(S') = \begin{bmatrix} \text{value} & 0 & 0 \\ \text{value} & \text{value} & 0 \\ \text{value} & \text{value} & \text{value} \end{bmatrix}.\]

With tiled attention, we tin skip full tiles:

Entire blocks supra the diagonal tin beryllium wholly skipped.
Blocks beneath the diagonal are afloat valid.
Diagonal blocks require element-level masking.

This is important because we are nary longer doing useless activity for evidently invalid early positions.

In causal self-attention, query position $i$ whitethorn not be to early cardinal positions $j > i$. A materialized disguise would beryllium different $N \times N$ object, but an businesslike tiled kernel tin logic astir the geometry of each tile.

For quadrate self-attention:

blocks strictly supra the causal bound are afloat masked and request not contribute,
blocks strictly beneath the bound are afloat valid,
only blocks intersecting the diagonal require element-level causal masking.

This avoids computing galore invalid tiles successful a causal kernel. FlashAttention-2 explicitly exploits causal structure, while existent implementations besides specify precise alignment rules for unequal query/key lengths.

Mask semantics are an API item that must not beryllium guessed. For example, existent PyTorch SDPA treats a Boolean attn_mask worth of True arsenic a position that participates successful attention, while different PyTorch disguise APIs usage different conventions.

Connection to section attention. Suppose each token tin only be to 128 adjacent tokens. Then galore tiles tin besides beryllium skipped. But now we person changed the mathematical attraction pattern: that is nary longer undistributed dense attention. FlashAttention tin beryllium the kernel executing the section pattern, but the exemplary itself is now doing sparse/local attention.

Local/sliding-window attraction tin likewise skip tiles extracurricular the permitted window. But erstwhile the exemplary intentionally restricts which pairs are attended, the model’s attraction pattern is sparse/local. FlashAttention tin beryllium the kernel utilized to execute that pattern, but it is nary longer the aforesaid mathematical problem arsenic unrestricted dense attention.

3.6 Backward Pass: Recompute Instead of Save

This is captious for training.

During the guardant pass, we do not prevention the full $P$ matrix ($N \times N$). We recompute it.

Forward:

Store compact info specified as:
- output
- row-wise normalization statistics
Do not shop $P \in \mathbb{R}^{N \times N}$.

Backward:

For each tile:
- a) Recompute $QK^T$
- b) Reconstruct the section probability values
- c) Calculate gradients
- d) Discard the tile
This trades much computation for little representation traffic/storage.

Conceptually: forward: don’t shop $P$. Backward: recompute $P$ tile-by-tile.

The gradient rehashing for $P$:

\[dV \mathrel{+}= P^T dO\] \[dP = dO \cdot V^T\] \[dS = P \odot (dP - D\_i[:, \text{None}])\]

and then:

\[dQ \mathrel{+}= dS \cdot K / \sqrt{d}\] \[dK \mathrel{+}= dS^T \cdot Q / \sqrt{d}\]

with due masking. Where

\[D\_i[:, \text{None}] = \sum\_j P\_{ij}\, dP\_{ij}\]

is the column/vector broadcast crossed each row.

A naive training implementation tin prevention $P = \mathrm{softmax}(S)$ for backward. FlashAttention avoids keeping that $N^2$ tensor successful HBM. Instead, it stores compact row-wise normalization accusation and recomputes score/probability tiles erstwhile gradients are needed.

For

\[S = QK^T / \sqrt{d}, \quad P = \mathrm{softmax}(S), \quad O = PV,\]

a useful row-wise personality is

\[D\_i = \sum\_r (dO\_i)\, O\_{ir} = \sum\_j P\_{ij}\, dP\_{ij}.\]

After recomputing a tile of $P$, the section derivatives tin beryllium expressed as

\[dV \mathrel{+}= P^T dO, \quad dP = dO\, V^T, \quad dS = P \odot (dP - D\_i[\text{:}, \text{None}]),\]

then

\[dQ \mathrel{+}= dS\, K / \sqrt{d}, \quad dK \mathrel{+}= dS^T\, Q / \sqrt{d}.\]

Masks connote zero probability/gradient publication for masked entries. FA2 stores a row-wise log-sum-exp amount that allows the probability tile to beryllium reconstructed stably from recomputed scores.

Backward recomputation is intentional. It spends other matrix-multiply activity to debar reference and penning a elephantine probability tensor, which tin beryllium a favorable waste and acquisition connected GPUs.

This is the memory-compute tradeoff applied to the backward pass.

3.7 Why More FLOPs Can Still Be Faster

This is 1 of the biggest lessons from FlashAttention.

Normally we think: less FLOPs $\Rightarrow$ faster. But connected GPUs, that’s incomplete.

Different operations person radically different throughput. Tensor cores are exceptionally bully astatine matrix multiplication. Other operations person different capacity profiles. These operations include:

exponentials
reductions
synchronization
shared-memory operations
memory transfers

Therefore, sometimes doing extra arithmetic is worthwhile if it eliminates costly representation traffic.

For example, see 2 options:

A: Compute $\to$ constitute immense $P$ to HBM $\to$ publication $P$ $\to$ compute.
B: Compute $\to$ discard $\to$ recompute later.

Option B performs much arithmetic. But it mightiness beryllium faster because it avoids moving a elephantine tensor done HBM.

This is simply a basal ML systems principle:

The costs of a FLOP depends connected what hardware executes it and what information activity surrounds it.

It is tempting to presume that less arithmetic operations ever connote little latency. Accelerator capacity breaks that intuition regularly.

Matrix multiplication maps exceptionally good to tensor cores. HBM traffic, synchronization, shared-memory traffic, exponentials, reductions, and kernel motorboat boundaries tin beryllium comparatively expensive.

FA2 makes this opposition explicit: 1 of its goals is to trim non-matmul FLOPs, because those operations do not bask the aforesaid throughput arsenic tensor-core GEMMs. The insubstantial besides improves activity partitioning truthful much of the GPU is occupied.

FA3 goes further connected Hopper by overlapping matrix multiplication, softmax, and information activity utilizing asynchronous hardware features. FA4 responds to Blackwell, wherever tensor-core throughput accrued faster than immoderate different resources, making exponentials and shared-memory postulation comparatively much important.

A amended capacity mobility is not simply “How galore FLOPs?” but “Which operations, connected which units, pinch what information movement, reuse, parallelism, and synchronization?”

This is the systems instruction that makes FlashAttention important beyond attraction itself: hardware ratio often comes from co-designing mathematical scheduling pinch the memory/execution hierarchy.

Part 4: Architectural Compatibility

FlashAttention is not tied to modular MHA.

Let

\[Q \in \mathbb{R}^{B \times N\_q \times H\_q \times d}, \quad K, V \in \mathbb{R}^{B \times N\_k \times H\_{kv} \times d}.\]

Three head-sharing regimes:

MHA: $H_q = H_{kv}$ $\Rightarrow$ each query caput has its ain K/V head.
MQA: $H_{kv} = 1$ $\Rightarrow$ each query heads stock 1 K/V head.
GQA: $1 < H_{kv} < H_q$ $\Rightarrow$ respective query heads stock K/V heads.

Vital distinction:

MHA/MQA/GQA defines the architecture and caput sharing.
FlashAttention defines businesslike execution of the attraction computation.

Thus MHA/MQA/GQA and FlashAttention tin coexist. Current implementations enforce style constraints: the number of query heads must beryllium divisible by the number of KV heads:

\[H\_q \bmod H\_{kv} = 0.\]

FlashAttention does not require each architecture to person the aforesaid number of query and key/value heads.

For mean MHA, $H_q = H_{kv}$. In MQA, $H_{kv} = 1$. In GQA, $1 < H_{kv} < H_q$. Current Dao-AILab kernels support MQA/GQA by passing less KV heads than query heads, pinch the request that the number of query heads beryllium divisible by the number of KV heads.

The kernel still evaluates attraction betwixt each query caput and its assigned KV head. Head sharing changes the architecture and KV-memory footprint. FlashAttention changes really the resulting attraction cognition is executed.

Orthogonal concepts:

MQA/GQA: really heads stock K/V projections,
RoPE: really position transforms Q/K,
FlashAttention: really attraction is scheduled and computed efficiently.

These tin beryllium utilized together.

Current PyTorch SDPA besides exposes enable_gqa. Its archiving labels GQA support experimental and imposes backend- and tensor-shape constraints, truthful accumulation codification should travel the nonstop version’s archiving alternatively than assuming cosmopolitan fused-kernel support.

4.2 Variable Lengths, Local Attention, and Dropout

Real systems aren’t always: aforesaid series magnitude + dense attraction + nary dropout.

Production implementations whitethorn support:

variable-length sequences
causal attention
sliding-window attraction (SWA)
dropout
MQA/GQA
ALiBi-style bias
KV-cache decoding (including optional RoPE handling)

But an important distinction: these are implementation capabilities, not basal properties of the FlashAttention mathematical idea.

Feature support depends on:

GPU
CUDA/ROCm backend
dtype
head dimension
mask
library version
kernel generation

PyTorch issue: scaled_dot_product_attention applies dropout according to the supplied dropout_p, truthful eval codification should explicitly usage 0.0 erstwhile dropout should beryllium disabled.

Production attraction seldom consists only of equal-length dense sequences pinch nary dropout. Current FlashAttention implementations support a broader characteristic set, but these are implementation capabilities, not properties of the mathematical thought itself.

The existent Dao-AILab repository documents kernels/interfaces for features including:

variable-length sequences,
causal attention,
local/sliding-window attention,
dropout successful training-oriented interfaces,
MQA/GQA,
ALiBi-style people bias successful applicable interfaces,
specialized incremental-decoding paths pinch KV cache, including optional RoPE handling.

Feature support differs by CUDA/ROCm backend and evolves complete time. For example, the repository documents abstracted NVIDIA and AMD backends pinch different implementation specifications and support matrices.

Dropout deserves a applicable warning. Current PyTorch scaled_dot_product_attention ever applies dropout according to its dropout_p argument, truthful callers must walk 0.0 during information erstwhile dropout should beryllium disabled.

Do not infer characteristic support from the sanction “FlashAttention.” Check the nonstop library, kernel generation, device, dtype, caput dimension, mask/bias, and training/inference way that will really run.