FlashAttention from Scratch: From Softmax to Blackwell

FlashAttention computes attention while avoiding the large intermediate matrices that a straightforward implementation writes to GPU memory. Its central mathematical idea fits in a few equations. Understanding its modern implementations requires following those equations through a GPU’s memory hierarchy, parallel execution, and numerical formats.

We will build that understanding one step at a time. First compute an attention row by hand. Then turn it into a streaming calculation, map that calculation onto GPU tiles, derive backward, and investigate why each hardware generation needs a different execution schedule.

Research snapshot: September 27, 2026. The path reaches FlashAttention-4, programmable attention, and the September 2026 MXFP8 and FP4 developments. “State of the art” here describes techniques and evidence available on that date; the fastest implementation depends on workload, hardware, precision, and software version.

1. How to use this tutorial

You need basic Python, matrix multiplication, exponentials, and sums to begin. GPU programming becomes necessary in Section 6; derivatives become necessary in Section 8. You can run the mathematical examples on a CPU with PyTorch. Keep the hardware experiments for a compatible GPU.

Each stage pairs explanations and worked examples with a checkpoint or experiment. Use the progression below to decide what to study and implement next.

At each checkpoint, explain the result before optimizing it. If you cannot explain a five-token example, a profiler trace from a 32K-token workload will not resolve the confusion.

A short historical map

Dates identify papers or announcements, not the first appearance of every implementation. FA4 code preceded its March 2026 paper, as the author post explains. The history also has parallel branches: PagedAttention and FlashInfer solve serving problems and are not numbered successors to FA2.

2. Begin with one attention row

For one batch item and one head, define:

Learned projections usually produce $Q$, $K$, and $V$ from hidden states. We start after those projections. Attention is

The mask $M_{ij}$ is zero for an allowed pair and $-\infty$ otherwise. Each row has its own softmax. A key receives a large weight when its score is large relative to other keys for that query. This is the scaled dot-product operation introduced in the Transformer paper.

Why divide by $\sqrt d$? Under the illustrative assumption that query and key components are independent, zero-mean, and unit-variance, their dot product has variance $d$. Scaling by $1/\sqrt d$ keeps its variance near one. Actual learned tensors need not satisfy these assumptions exactly; the calculation explains the choice of scale.

Work through a complete example

Let $d=2$ and use one query:

After scaling, the scores are $s=[1,2,3]$. Their softmax is approximately

The output is a weighted average:

Values need not be scalars. If each value has 64 components, the same three weights combine all 64 components independently.

Stabilize softmax before doing anything clever

Computing $e^{s_j}$ directly can overflow. Subtracting the same constant $c$ from every score leaves the ratio unchanged:

Choose $c=m=\max_j s_j$. Every exponential argument is then nonpositive. For the example, compute $[e^{-2},e^{-1},1]$ and divide by their sum. For scores $[1001,1002,1003]$, compute exactly the same shifted quantities.

Checkpoint: Add 1000 to every score. Explain why the output is unchanged. Then mask the third key: the remaining weights must be renormalized; simply setting the old third weight to zero gives the wrong answer.

3. Why straightforward attention is expensive

A direct implementation writes scores, reads them for softmax, writes probabilities, and reads those probabilities for the value product:

flowchart TD
    QK["Q and K"] --> MM1["Matrix multiply"]
    MM1 --> S["Write full score matrix"]
    S --> SM["Read scores; softmax"]
    SM --> P["Write full probability matrix"]
    P --> MM2["Read probabilities and V; matmul"]
    MM2 --> O["Output"]

For $B$ batch items and $H$ heads, one dense score tensor has $BHN_qN_k$ elements. With $B=1$, $H=32$, $N_q=N_k=8192$, and two bytes per element:

That is one intermediate. If scores and probabilities occupy distinct buffers, their combined storage is 8 GiB. FP32 storage doubles those numbers. By comparison, when $d=d_v=128$, each of $Q$, $K$, $V$, and $O$ occupies only 64 MiB at two bytes per element.

An SM, or streaming multiprocessor, executes thread blocks. A CTA is another name for a thread block. NVIDIA threads execute in groups of 32 called warps. A tensor core accelerates matrix multiply-accumulate operations, abbreviated MMA.

FlashAttention combines GPU-aware tiling, online normalization, and backward recomputation to avoid writing full $S$ and $P$ matrices. The operation still evaluates dense query–key interactions; eliminating quadratic intermediate storage does not eliminate quadratic arithmetic. FlashAttention paper

With equal value and key dimensions, the two matrix products require approximately

floating-point operations when multiply and add count separately. Causal attention has fewer valid pairs, but still grows quadratically when both sequence dimensions grow together.

Checkpoint: If sequence length doubles, which tensors grow fourfold and which grow twofold? Why does this say nothing yet about how many times a kernel reads them?

4. Derive streaming attention

The obstacle to tiling is softmax: each probability depends on a denominator covering the whole row. We need a summary of earlier keys that can be updated when later keys arrive.

4.1 Maintain three quantities

For one query and the set $A$ of keys already processed, maintain

$m_A$ and $\ell_A$ are scalars; $u_A$ has $d_v$ components. The output over these keys is $u_A/\ell_A$.

The online-normalizer paper develops streaming softmax statistics. Adding the weighted sum $u$ makes the state useful for attention: we need the final weighted average, so earlier individual probabilities can disappear.

Suppose a new tile $T$ arrives. Set

Every old exponential can be expressed on the new scale:

Therefore

This is the entire recurrence. Both old sums receive the same correction, so they remain expressed in the same units. After the final tile, return $o=u/\ell$.

4.2 Stream the numerical example

Process scores $[1,2]$ first, with values $[10,20]$:

Now score $3$ and value $30$ arrive. The maximum rises to 3, so $\alpha=e^{-1}$:

Dividing gives $25.7521$, matching the dense calculation. We discarded the first two scores without losing the information needed for the final weighted average.

4.3 Merge independently computed partitions

Given states $(m_A,\ell_A,u_A)$ and $(m_B,\ell_B,u_B)$ for disjoint key subsets, define

In exact arithmetic, different merge trees give the same result. Floating-point order can change rounding. This composition underlies split-KV decoding; the FlashInfer attention-state tutorial explains the recursive formulation.

If partitions return normalized outputs and natural-log normalizers $L_A=\log\sum_{j\in A}e^{s_j}$, merge with

Partial outputs cannot generally be averaged equally. Each is weighted by its partition’s share of total softmax mass. If both partitions are empty for a row, handle that row separately.

4.4 Handle the empty state deliberately

Initialize $m=-\infty$, $\ell=0$, and $u=0$. An entirely masked tile has maximum $-\infty$. Subtracting $-\infty-(-\infty)$ produces NaN, so the implementation must guard this case.

Our examples define a fully masked row to return zero output and log-sum-exp $-\infty$. This is an explicit API convention; ordinary softmax over such a row has a zero denominator. We assume all unmasked input scores are finite.

Checkpoint: Merge three partitions in two orders. Explain why they require an identical query and scoring rule. Shared cached keys do not make query-dependent attention outputs interchangeable between requests.

5. Run the math before writing a GPU kernel

This complete, CPU-compatible forward example supports one head, different query/key lengths, different value dimensions, and an optional causal mask. Copy the Python block into a file and run it with PyTorch installed. It computes in FP64 so rounding does not obscure the recurrence.

q_start and k_start are absolute token positions. Allowed causal pairs satisfy key_position <= query_position. This makes slicing a KV partition unambiguous.

import math
import torch


def prepare(q, k, v):
    assert q.ndim == k.ndim == v.ndim == 2
    assert q.shape[1] == k.shape[1] and k.shape[0] == v.shape[0]
    assert min(*q.shape, *k.shape, *v.shape) > 0
    assert q.device == k.device == v.device
    return tuple(x.to(torch.float64) for x in (q, k, v))


def score_tile(q, k, causal, q_start, k_start):
    s = (q @ k.T) / math.sqrt(q.shape[1])
    if causal:
        qi = torch.arange(q.shape[0], device=q.device) + q_start
        kj = torch.arange(k.shape[0], device=k.device) + k_start
        s = s.masked_fill(kj[None, :] > qi[:, None], -torch.inf)
    return s


def finish(m, ell, u):
    # Our contract: zero output and -inf LSE for a fully masked row.
    denom = torch.where(ell > 0, ell, torch.ones_like(ell))
    out = u / denom[:, None]
    lse = torch.where(ell > 0, m + denom.log(), -torch.inf)
    return out, lse


def dense_attention(q, k, v, causal=False, q_start=0, k_start=0):
    q, k, v = prepare(q, k, v)
    s = score_tile(q, k, causal, q_start, k_start)
    m = s.amax(dim=-1)
    # Zero is only a computational shift for entirely masked rows.
    shift = torch.where(torch.isfinite(m), m, torch.zeros_like(m))
    p = torch.exp(s - shift[:, None])
    return finish(m, p.sum(dim=-1), p @ v)


def streaming_attention(q, k, v, block_q=32, block_k=64,
                        causal=False, q_start=0, k_start=0):
    q, k, v = prepare(q, k, v)
    assert block_q > 0 and block_k > 0
    outputs, normalizers = [], []
    for i in range(0, q.shape[0], block_q):
        qt = q[i:i + block_q]
        m = torch.full((qt.shape[0],), -torch.inf,
                       dtype=q.dtype, device=q.device)
        ell = torch.zeros_like(m)
        u = torch.zeros((qt.shape[0], v.shape[1]),
                        dtype=q.dtype, device=q.device)
        for j in range(0, k.shape[0], block_k):
            kt, vt = k[j:j + block_k], v[j:j + block_k]
            s = score_tile(qt, kt, causal, q_start + i, k_start + j)
            new_m = torch.maximum(m, s.amax(dim=-1))
            shift = torch.where(torch.isfinite(new_m), new_m,
                                torch.zeros_like(new_m))
            alpha = torch.exp(m - shift)
            p = torch.exp(s - shift[:, None])
            ell = alpha * ell + p.sum(dim=-1)
            u = alpha[:, None] * u + p @ vt
            m = new_m
        out, lse = finish(m, ell, u)
        outputs.append(out)
        normalizers.append(lse)
    return torch.cat(outputs), torch.cat(normalizers)


torch.manual_seed(7)
q = torch.randn(5, 16)
k = torch.randn(13, 16)
v = torch.randn(13, 8)
# Queries are the final five positions of a thirteen-token cache.
ref, ref_lse = dense_attention(q, k, v, causal=True, q_start=8)
out, lse = streaming_attention(q, k, v, block_q=3, block_k=4,
                                causal=True, q_start=8)
torch.testing.assert_close(out, ref, atol=1e-10, rtol=1e-10)
torch.testing.assert_close(lse, ref_lse, atol=1e-10, rtol=1e-10)
print("Streaming output and LSE match the dense reference.")

p inside the loop is an unnormalized exponential tile. Dividing at the end avoids repeated normalization. The largest score tile is block_q × block_k; the output accumulator is block_q × d_v.

Python loops validate the equations but do not provide GPU fusion: on CUDA, the individual PyTorch operations would launch separate kernels and store intermediates in device memory. This forward-only memory description also assumes no autograd recording; ordinary autograd through the loop may save intermediates. Section 8 explains how a designed backward avoids that.

Rectangular causal masking is a correctness issue

For two queries that are the final positions of a five-token cache, the allowed-pair mask is

             key position
             0  1  2  3  4
query 3      1  1  1  1  0
query 4      1  1  1  1  1

Using query positions 0 and 1 produces a different operation. PyTorch SDPA documents an upper-left rectangular is_causal=True mask; the FlashAttention package uses bottom-right alignment from version 2.1. Construct the desired mask explicitly when comparing them. Set SDPA’s dropout_p=0.0 for these no-dropout examples. SDPA contract, FlashAttention API changelog

Experiment: Change tile sizes, use odd sequence lengths, and multiply queries and keys by 20. Try q_start=-2 so some rows are fully masked. Those outputs should be zero, their LSE should be $-\infty$, and valid rows should still match.

6. FlashAttention-1: make the algorithm IO-aware

The 2022 step was to exploit GPU memory: load tiles into on-chip storage, immediately consume score tiles in the value product, and retain only the statistics and partial output needed later. Read the original paper and the tiling/recomputation portion of Tri Dao’s Stanford talk, around 22–30 minutes.

6.1 Understand tile sizes and live storage

Let $B_r$ be the query rows and $B_c$ the key rows in a tile. One iteration computes

then updates $u$ using the exponential tile and $V_t$. For $B_r=B_c=64$ and $d=d_v=128$, a simplified inventory is:

This is a liveness exercise, not a launch-ready shared-memory allocation. Registers and shared memory have different budgets. A kernel may overwrite scores with probabilities, retain more fragments, or double-buffer loads. Instruction layouts and compiler allocation determine actual placement.

A larger query tile reuses each KV tile across more queries but enlarges the accumulator. A larger key tile reduces loop iterations but enlarges the score tile. These are competing costs.

6.2 Separate the idea from the historical loop order

The original FA1 paper’s forward schedule puts the KV loop outside the query loop:

for each KV tile:
    load K_tile and V_tile
    for each Q tile:
        load Q_tile and its running output/statistics
        compute scores; update online attention
        store updated output/statistics

It avoids full score/probability storage but can move running output state repeatedly. Our reference instead keeps one query tile through the KV loop, a useful stepping stone to FA2. It is not a literal reproduction of FA1’s schedule. FA1, Algorithm 1

For the query-owned traversal, the state flow looks like this:

flowchart TD
    Q["Keep one Q tile"] --> S["Compute a score tile"]
    KV["Load next K and V tile"] --> S
    S --> P["Compute shifted exponentials"]
    P --> U["Update m, ell, and weighted sum u"]
    U --> More{"More KV tiles?"}
    More -->|Yes| KV
    More -->|No| O["Normalize u / ell and write output"]

The loop carries the small row state. Scores and exponentials are temporary tiles whose storage can be reused after consumption.

6.3 Build a small fused forward

Start with contiguous FP16/BF16 inputs, FP32 accumulation, one head dimension, noncausal attention, and no dropout. Assign a thread block to one query tile. Keep its accumulator and statistics on-chip through the KV loop.

parallel task (batch, head, query_tile):
    load Q_tile
    initialize m, ell, u
    for each KV tile:
        load K_tile and V_tile
        S = matmul(Q_tile, transpose(K_tile)) * scale
        apply boundary and attention masks
        update m, ell, u using the online recurrence
    store u / ell
    optionally store m + log(ell)

This is algorithmic pseudocode. Map its lines to loads, reductions, MMA fragments, and stores using the executable Triton fused-attention tutorial, which implements an FA2-style algorithm. Isolate its ordinary forward path before the later hardware-specific branches.

For CuTe DSL, first establish which elements belong to each thread and which layouts the two MMAs expect. Converting the first MMA accumulator into the second MMA operand is part of the work. Matrix shapes alone do not specify that layout.

If using base-2 exponentials, convert consistently:

An LSE stored as $L_2=\log_2\sum_j2^{s_{2,j}}$ converts to natural-log LSE via $L=L_2\ln2$. Mixing conventions breaks backward and partition merging.

Checkpoint: Identify every global-memory write. There should be no full $N_q\times N_k$ score or probability output. Inputs need not be read only once: different query tiles may reload the same KV data.

7. FlashAttention-2: expose independent work

FA2, published in 2023, targets non-matmul arithmetic, insufficient parallelism across query tiles, and communication between warps. FA2 paper, Section 3

7.1 Delay normalization

If we store a normalized partial output $o$, each update is

Storing $u=\ell o$ removes repeated denominator multiplications and divisions. Our reference already does this. Reductions and exponentials still exist; eliminating divisions does not eliminate all non-matmul work.

7.2 Parallelize output rows

The approximate number of query-tile tasks is

One batch, eight query heads, $N_q=8192$, and $B_r=128$ provide 512 tasks. Parallelizing only batch and heads would provide eight. Query tiles need no output reduction between them because they own different rows.

More tasks can improve utilization, but register demand, shared memory, and work per task still constrain execution.

7.3 Divide query rows between warps

FA2 assigns distinct query rows to warps while reusing KV data. Each warp accumulates its own output rows. An earlier split-K arrangement has warps compute partial contributions to the same output, requiring communication and a reduction. FA2 avoids that within-block reduction. FA2 work partitioning

Within-block KV split:            Within-block query-row split:
warp 0: part of every output      warp 0: complete rows  0–15
warp 1: part of every output      warp 1: complete rows 16–31
warp 2: part of every output      warp 2: complete rows 32–47
warp 3: part of every output      warp 3: complete rows 48–63
         reduction needed                 disjoint output ownership

The row ranges illustrate ownership, not a prescribed tensor-core layout. Distinguish this from splitting KV across independent blocks during decode: that later split accepts a reduction to expose more work.

For square causal attention, classify tiles as entirely allowed, entirely masked, or crossing the diagonal. Skip masked tiles and apply elementwise masking only at the boundary. Earlier queries now have less work, creating load imbalance.

Experiment: Sweep tile sizes and warp count. Record latency, registers, shared memory, and spills. Explain one case where a larger tile loses despite better reuse: which resource cost increased?

8. Backward: spend arithmetic to save memory

Let $G=\partial\mathcal L/\partial O$ be the incoming gradient. We derive the no-dropout case with a fixed mask and scale $c=1/\sqrt d$.

From $O=PV$:

The softmax Jacobian gives

Define $\delta_i=\sum_tP_{it}dP_{it}$. Substituting $dP_{it}=G_i\cdot V_t$ gives

We can compute $\delta$ from the saved output and incoming gradient. Finally,

For an allowed pair, reconstruct probabilities from a saved row normalizer:

Save $O$ and row statistics, retain or recover inputs, and reconstruct one probability tile at a time. The FA1 paper introduces the recomputation strategy; FA2 develops the LSE and gradient formulation further.

A backward implementation you can check

Append this block after the forward example. It uses sequential accumulation to make gradient ownership easy to understand.

def backward_reference(q, k, v, out, lse, grad_out,
                       block_q=32, block_k=64,
                       causal=False, q_start=0, k_start=0):
    q, k, v = prepare(q, k, v)
    out, lse, grad_out = (x.to(torch.float64)
                          for x in (out, lse, grad_out))
    assert block_q > 0 and block_k > 0
    dq, dk, dv = torch.zeros_like(q), torch.zeros_like(k), torch.zeros_like(v)
    delta = (out * grad_out).sum(dim=-1)
    scale = 1.0 / math.sqrt(q.shape[1])
    for i in range(0, q.shape[0], block_q):
        rows = slice(i, i + block_q)
        qt, gt = q[rows], grad_out[rows]
        li = lse[rows]
        # For a fully masked row, exp(-inf - 0) is safely zero.
        shift = torch.where(torch.isfinite(li), li, torch.zeros_like(li))
        for j in range(0, k.shape[0], block_k):
            cols = slice(j, j + block_k)
            kt, vt = k[cols], v[cols]
            s = score_tile(qt, kt, causal, q_start + i, k_start + j)
            p = torch.exp(s - shift[:, None])
            dp = gt @ vt.T
            ds = p * (dp - delta[rows, None])
            dq[rows] += scale * (ds @ kt)
            dk[cols] += scale * (ds.T @ qt)
            dv[cols] += p.T @ gt
    return dq, dk, dv


# Compare with autograd through an ordinary dense softmax.
torch.manual_seed(8)
qg = torch.randn(5, 7, dtype=torch.float64, requires_grad=True)
kg = torch.randn(9, 7, dtype=torch.float64, requires_grad=True)
vg = torch.randn(9, 3, dtype=torch.float64, requires_grad=True)
s = score_tile(qg, kg, causal=True, q_start=4, k_start=0)
og = torch.softmax(s, dim=-1) @ vg
g = torch.randn_like(og)
expected = torch.autograd.grad(og, (qg, kg, vg), g)
with torch.no_grad():
    o, l = streaming_attention(qg, kg, vg, block_q=3, block_k=4,
                               causal=True, q_start=4)
    actual = backward_reference(qg, kg, vg, o, l, g,
                                block_q=3, block_k=4,
                                causal=True, q_start=4)
for result, reference in zip(actual, expected):
    torch.testing.assert_close(result, reference, atol=1e-10, rtol=1e-10)
print("Recomputed backward matches dense autograd.")

A GPU backward must handle competing ownership requirements. A query-owned block can finish $dQ$ locally, but its $dK$ and $dV$ contributions overlap with other blocks. A KV-owned block can finish $dK$ and $dV$, but $dQ$ needs a reduction across KV blocks. Choose ownership, workspace, and synchronization deliberately.

Recomputing backward has five matrix products: reconstruct scores, compute $dP$, compute $dV$, compute $dQ$, and compute $dK$. More arithmetic can still be faster when it avoids enough memory traffic. This is an attention-specific instance of activation recomputation.

Dropout requires backward to reproduce the same random decisions and scaling. GQA requires accumulating contributions from several query heads into shared KV heads. Neither feature is implemented here.

Checkpoint: Derive $\delta_i=G_i\cdot O_i$ without constructing $dP$. Draw the reduction axes for $dQ$, $dK$, and $dV$ before studying FA4’s backward pipeline.

9. The inference branch: decode, GQA, and paged KV

Prefill processes many prompt queries, often with $N_q\approx N_k$. Ordinary single-token decode has $N_q=1$ and a growing $N_k$. Query-tile parallelism largely disappears for each head during decode.

9.1 Split the long dimension that remains

Flash-Decoding splits the KV sequence into independent partitions. Each task computes a partial output and normalization statistics; a reduction combines them using Section 4’s merge. The extra tasks can occupy otherwise idle SMs. They also introduce partial-result storage and reduction overhead, so the split count must depend on workload. Flash-Decoding

Append this runnable example after the forward block. The two partitions have different absolute key offsets, which matters for causal masking.

def merge_outputs(oa, la, ob, lb):
    l = torch.logaddexp(la, lb)
    shift = torch.where(torch.isfinite(l), l, torch.zeros_like(l))
    wa, wb = torch.exp(la - shift), torch.exp(lb - shift)
    return wa[:, None] * oa + wb[:, None] * ob, l


torch.manual_seed(9)
q = torch.randn(2, 8)
k = torch.randn(11, 8)
v = torch.randn(11, 5)
oa, la = streaming_attention(q, k[:4], v[:4], causal=True, q_start=9)
ob, lb = streaming_attention(q, k[4:], v[4:], causal=True,
                             q_start=9, k_start=4)
merged, merged_lse = merge_outputs(oa, la, ob, lb)
ref, ref_lse = dense_attention(q, k, v, causal=True, q_start=9)
torch.testing.assert_close(merged, ref, atol=1e-10, rtol=1e-10)
torch.testing.assert_close(merged_lse, ref_lse, atol=1e-10, rtol=1e-10)
print("Split-KV merge matches full attention.")

This demonstrates the reduction’s mathematics. Its Python calls are sequential; a GPU implementation supplies the parallel execution.

9.2 GQA changes the reuse opportunity

Multi-head attention normally has one KV head per query head. Grouped-query attention, or GQA, shares a KV head across $r=H_q/H_{kv}$ query heads. Multi-query attention is the case $H_{kv}=1$.

For equal-sized groups, query head $h$ maps to KV head $\lfloor h/r\rfloor$. A reference can repeat KV tensors to check the result, but a performance implementation should reuse the original storage.

For one-token decode with equal key/value dimensions and $s$ bytes per KV element, an idealized traffic model gives

At two bytes per element and $r=1$, that is only one FLOP per KV byte. GQA increases the opportunity for reuse, but the kernel must realize it. Packing query heads that share KV into one work tile can help fill a matrix instruction’s otherwise mostly empty query dimension. This is illustrated in the FA3/CUTLASS lecture, around 40–44 minutes.

The estimate excludes repeated loads, metadata, scale tensors, outputs, and cache effects. It explains why a compute-efficient prefill kernel can still be a poor decode choice.

9.3 Paging changes where keys live

A serving system need not allocate one contiguous maximum-length cache per request. With page size $P$, logical token $t$ maps to

PagedAttention concerns allocation and access to these KV blocks, including memory sharing and fragmentation. A FlashAttention-style computation can consume a paged cache. Paging and streaming softmax address different problems. PagedAttention paper

FlashInfer provides inference kernels and planning interfaces for workload-dependent layouts and execution. A backend family name does not uniquely identify a single kernel. Measure the selected implementation and separate planning from execution. FlashInfer paper, KV layout tutorial

Experiment: First permute physical pages while preserving the logical page table; outputs must remain unchanged. Then benchmark different split counts at several batch sizes and context lengths. Find a case where split-KV overhead exceeds its parallelism benefit.

10. FlashAttention-3: overlap work on Hopper

Before this stage, learn double buffering with a small matrix multiplication: while consuming one input buffer, fill another. The mathematical dependency stays unchanged; independent work moves earlier in time.

FA3, introduced in 2024, develops this idea for Hopper. Its main ingredients are asynchronous data movement, asynchronous tensor-core execution, and careful overlap with softmax. Study the FA3 paper, author explanation, and CUTLASS lecture.

10.1 Learn the hardware vocabulary

These are Hopper mechanisms; consult the Hopper tuning guide. An asynchronous instruction can return control before its data is ready. Issuing an MMA and consuming its accumulator are separate events.

10.2 Draw dependencies before drawing overlap

For one KV tile, the dependency is unavoidable:

K ready → QK matmul complete → softmax complete → PV matmul complete
                                                   ↑
                                                V ready

For different tiles or query work, other operations can proceed concurrently. A conceptual schedule might look like this:

This is a dependency sketch, not a cycle-accurate FA3 schedule. The running output must be rescaled at the correct time before accumulating contributions on a new scale. A shared-memory slot cannot be reused while an asynchronous operation still reads it.

FA3 uses producer/consumer specialization and overlap both across and within consumer warpgroups. Additional in-flight score state increases register demand, so deeper pipelines can lose their advantage through spills or reduced residency. FA3 pipeline design

10.3 Treat FP8 as a second lesson

FA3 also studies FP8, including block quantization and incoherent processing to reduce outlier-related error. The latter can transform $Q$ and $K$ with the same orthogonal matrix $R$:

Before rounding, scores are unchanged. After quantization, spreading a large component across coordinates can change the error distribution. Layout restrictions and converting intermediate probabilities into valid MMA operands add implementation work. FP8 therefore changes both arithmetic and data preparation. FA3 numerical techniques

Checkpoint: Annotate an actual Hopper mainloop with “issued,” “ready,” “consumed,” and “safe to reuse.” Then disable one supported overlap mechanism and measure its effect. A timeline explanation should identify the dependency that became exposed.

11. FlashAttention-4: optimize the resources left behind

The FA4 paper appeared in March 2026 and focuses on datacenter Blackwell. Its premise is asymmetric scaling: matrix throughput grows faster than exponential throughput and shared-memory bandwidth. The algorithm and pipeline must account for the resources that have become relatively more expensive. FA4 paper

11.1 Build a resource model before changing the kernel

Suppose a tile has $M$ queries, $N$ keys, and equal key/value dimension $d$. Count roughly $4MNd$ matrix FLOPs and $MN$ exponentials. If matrix throughput is $R_{\mathrm{mma}}$ and exponential throughput is $R_{\exp}$, estimate

Their ratio is

This gives a useful prediction: faster tensor cores or a smaller head dimension make exponentials relatively more significant. Increasing both sequence tile dimensions does not, by itself, change this simplified ratio.

Using the FA4 paper’s illustrative B200 rates of 8192 BF16 matrix FLOPs/cycle/SM and 16 exponentials/cycle/SM, a $128^3$ tile needs about 1024 cycles for its two matrix products and another 1024 cycles of exponential throughput. Its simplified MMA operand reads require about 768 shared-memory cycles. These are resource lower bounds, not a measured kernel duration. FA4 resource analysis, Section 3.1.1

Overlap can move execution toward the largest resource demand instead of their sum, but dependencies, reductions, loads, and barriers limit that ideal.

11.2 TMEM changes intermediate placement

Datacenter Blackwell introduces tensor memory, TMEM, used by its tcgen05 tensor-core operations. Accumulators can live there rather than occupying the issuing threads’ registers. TMEM is a distinct on-chip storage space with instruction-specific access and synchronization rules. It is neither global memory nor interchangeable with shared memory. Blackwell tuning guide

FA4 overlaps two query tiles, assigns specialized work to load, MMA, softmax, and output-correction roles, and reuses TMEM regions as intermediates retire. Some exponentials run through polynomial evaluation on ordinary arithmetic units alongside hardware exponential instructions. Colfax FA4 explanation

The general exponential identity is

The fractional argument lies in $[0,1)$, where a polynomial can approximate the second factor. This spends additional arithmetic and registers to reduce pressure on the exponential unit. It is useful only if those other resources have room and the resulting numerical error is acceptable.

11.3 Derive conditional rescaling correctly

FA4 also avoids many output rescaling operations. To understand why this can work, return to softmax’s freedom to choose an offset. The maximum is a safe choice, but any common finite reference $r$ gives the same exact-arithmetic ratio:

Here is a teaching formulation using natural exponentials. For a new tile with maximum $t$, retain $r$ while $t-r\leq\tau$; otherwise move the reference up. Both accumulated quantities always use the same reference:

r_new = r                           if t - r <= tau
        max(r, t)                   otherwise
a = exp(r - r_new)
p = exp(scores - r_new)
ell_new = a * ell + sum(p)
u_new   = a * u   + p @ V_tile

Initialize the reference from the first valid tile; handle empty tiles separately. If $r$ does not change, $a=1$, so the correction multiply can be skipped. For example, changing scores from maximum 2 to maximum 3 does not require a new reference: exponentials of up to $e^1$ are perfectly finite. The final ratio is still correct if numerator and denominator both retain reference 2.

The threshold limits how far exponentials can grow above one. If scores use base 2, a threshold of 8 means a factor of $2^8=256$; the same factor corresponds to $8\ln2$ with natural exponentials. Thresholds and warp-level decisions depend on formats and implementation. See the rescaling discussion in Ted Zadouri’s FA4 talk, around 14–17 minutes.

Do not update the denominator to a new scale while leaving the numerator on the old one. Final division cannot repair that inconsistency. This derivation explains the invariant; it does not reproduce every detail of the production correction pipeline.

11.4 Backward makes shared-memory traffic visible

Backward needs five matrix products and several differently oriented intermediates. FA4 arranges transposed probability and score-gradient tiles so some MMA operands can remain in TMEM. Its paired-CTA backward design reduces shared-memory operand traffic, while exchanging data between CTAs to satisfy the $dQ$ reduction. It also reduces global atomic accumulation work. Colfax FA4 backward explanation

Use Section 8’s equations to make an original storage-lifetime table:

Two objects can share physical storage only when their lifetimes do not overlap. An asynchronous instruction can extend a lifetime past the line that issued it. When an optimized kernel stalls, investigate whether it is waiting for data, a free buffer, or completion of an operation that still owns storage.

11.5 Scheduling and implementation language matter

Causal query tiles have unequal work. Launching long tiles earlier can reduce the tail where only a few SMs remain busy. Variable-length batches add another source of imbalance. Persistent kernels let a physical block process multiple logical work tiles, making scheduling a separate design choice.

FA4 is implemented in CuTe DSL. Python is its authoring language; compiled GPU instructions execute the kernel. Inspect the upstream dispatch and architecture-specific implementations instead of assuming a package name specifies one pipeline. The author post also notes that newer cuDNN versions incorporated related techniques, so benchmark current eligible cuDNN paths. FA4 author post

Experiment: Pick one tile configuration and make a table of matrix work, exponentials, shared-memory traffic, and live storage. Predict the effect of one change before measuring it. An unsuccessful optimization is useful if its regression follows your resource model.

12. Programmable attention: FlexAttention meets FA4

Once the base kernel works, models ask for variations: a positional bias, a sliding window, or isolation between packed documents. Two changes cover many cases:

1. Modify a score before softmax: $S'_{ij}=f(S_{ij},i,j,\ldots)$.
2. Specify which query–key pairs are allowed.

For example, a causal window of width $w$ permits