FlashAttention 3 - A Worklog[WIP]

Hello all. This blog post is about how to implement Flash Attention 3, specifically with CuteDSL, targeting H100 GPUs. The main goal of this work is to create a worklog on incrementally building and adding features to each kernel, to clearly and cleanly understand how we can get to Model FLOPs Utilization(MFU).

Much of this work is inspired by some of the best people whose content I really enjoyed in the GPU kernel optimization and ML performance engineering domain — specifically Simon, Aleksa Gordic, Kapil Sharma, and Hamza for their similar worklogs, which made me want to create one like this in the first place. Since this is going to be my take on such a worklog, I wanted to do something not available on the internet, or at least not in the form of a worklog.

• How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklog - Simon Boehm
• Inside NVIDIA GPUs: Anatomy of high performance matmul kernels - Aleksa Gordić
• Learn CUTLASS the hard way! - Kapil Sharma
• Worklog: Optimising GEMM on NVIDIA H100 for cuBLAS-like Performance (WIP) - Hamza

At present, most ML/AI architectures depend on two basic primitives: GEMM and Attention. I’m taking Attention(specifically Flash Attention) and exploring its kernel design and how to push it toward MFU. Below is the rough plan.

Flash Attention 3(FA3) Worklog:

• Attention math
• Online softmax
• GPU memory and compute hierarchy
• Tiling
• FA2
• CuTe DSL fundamentals and primitives
• K1 — naive attention
• K2 — tiled softmax
• K3 — using TMA
• K4 — using WGMMA
• K5 — warp specialization
• K6 — ping-pong scheduling
• K7 — FP8
• K8 — FP8 + incoherent processing
• Flash decode

The plan is to follow the above as one section each. Throughout this worklog, I’ll provide good intuition along with the details on how things work or what we need to make them work. Since each of these can be its own topic, I’ll also add separate notes for many of them, which you can refer to if you want to go in-depth.

This is a work in progress, marked with the “WIP” tag. Once it’s complete, I’ll remove that.

I sincerely thank all the many people who planted this idea in me. I hope this will be helpful for anyone who wants to understand FlashAttention 3 (FA3) from first principles. All the code will be made available here: FlashAttention3

I’ve made sure there aren’t any conceptual mistakes. If you find any, please let me know and I’m happy to correct them and make this useful for everyone. If you’re in this domain, feel free to ping me; I’m happy to collaborate on projects around GPU kernels and ML perf engineering.

I’ll add references along each section, instead of tagging them at the end. I’m assuming basic ML/AI and GPU knowledge to get the best out of this blog. All the analysis will be on H100 SXM GPUs, since FA3 was targeted at the Hopper architecture and aims to squeeze out its full utilization through Hopper-specific features like TMA, WGMMA, asynchronous operations, and FP8 support. One point worth mentioning here: there isn’t a huge algorithmic difference between FA2 and FA3, but a lot of engineering work went into FA3 to take advantage of Hopper. FA2 reached around 70% of theoretical max FLOPS on Ampere (A100) but didn’t exploit Hopper-specific features, while FA3 reaches around 75% on H100 in FP16, and with FP8 it gets close to 1.2 PFLOPS with 2.6x smaller numerical error than baseline FP8 attention. More on this in the upcoming sections.

Also, you’ll come across a lot of acronyms throughout this blog. I expand each term on its first use, so if you hit one you’re unsure about, a quick Ctrl+F will take you back to where it was introduced.

That’s it on the intro. Let’s get started.

Reference of FlashAttention-3 Paper - Link

Attention Math

Attention is the core of transformers, and transformers are the core of large language models (LLMs). Transformers consist of two primitives — attention, followed by feed-forward networks (FFNs) — surrounded by norms and residual connections. Below is an image that captures this.

Attention gives you the context of what is being referred to and meant in a sentence, and the FFN predicts the next token based on this context. Both are critical to each other and together form the backbone of present-day LLMs.

Now, if we look at attention in more detail, it’s composed of three operations:

Refer to Jay Alammar’s transformers blog if you’re not sure of the attention mechanism and what happens under the hood.

Here, (1) and (3) are matrix multiplications. In GPU terminology, we call them GEMM (General Matrix Multiply), and (2) is a softmax along the last dimension. This matters because softmax was a major bottleneck before, and FlashAttention made algorithmic changes to how it’s computed, while staying numerically exact.

Note: There is a scale factor of \(1/\sqrt{d}\) applied to \(S\), which we’re ignoring here for simplicity. The scale keeps the dot products from growing too large as \(d\) increases, which would otherwise push the softmax into saturated, low-gradient regions. In code, we’ll definitely include it. There is also a causal mask that needs to be applied, to make sure each query attends only to keys prior to it and not after. This is needed for next-token prediction.

The reason for optimizing attention is that materializing the intermediate \(S\) and \(P\) matrices is expensive. Why? Think in terms of dimensions and how much memory we’d need to get this computation done.

Let’s do a simple calculation. Let \(Q\), \(K\), and \(V\) be of shape \((b, s, d)\) → (batch, sequence length, model dim). Then \(S = QK^\top\) has shape \((b, s_1, s_2)\). Now apply softmax. What happens when the sequence length is 1 million and batch size is 1? Assuming FP32 (4 bytes per element):

And that’s just to materialize one \(S\) matrix. We still need to do the softmax and another matmul on top of that. Attention is also computed many times, since LLMs have hundreds of transformer blocks. Most GPUs don’t have 4 TB of memory, and even if they did, it still wouldn’t be feasible — we need space for intermediate results and bookkeeping too.

This is the major bottleneck FlashAttention addresses. At its core, it does two things: online softmax and tiling. Together, these ensure we never materialize the full \(s \times s\) matrix in GPU memory, while producing the exact same numerical result as if we had. This is what made it possible to train models with longer sequence lengths, which gave us more context and helped LLMs become much more useful.

There are a lot of details in and around this, but we’ll focus on FlashAttention 3 going forward, with all necessary information and intuition along the way for each optimization and topic. With a sense of what attention does and why we need to optimize it, we’ll next talk about online softmax.

A side note:

If you want to understand how LLMs work more clearly, try implementing one. Check out the Stanford CS336 course and follow their assignments — they’re self-contained and really push you to have a clear understanding of all the internals of LLMs. I did it as part of my coursework at NYU, and I’m glad to have taken that course. Highly recommend it.

Also, take a look at this: Transformer block accounting and the below image — think through it to understand the memory and FLOPs counts for a transformer block. This gives you a picture of what it takes to run one transformer block, so you can appreciate all the engineering that’s gone into making it fast and efficient.

Online Softmax

First, let’s revisit softmax and where it’s used.

Given a list of numbers \(\{x_1, x_2, \ldots, x_N\}\), for each element we compute \(e^{x_i} / \sum_j e^{x_j}\) and that produces a new list. Why do we do this? To make sure the new list sums to 1. Softmax is applied as the final step so that the probabilities over all choices add up to one. In the case of attention, we apply softmax to \(QK^\top\). These are the scores given to each key for a given query, and all the scores for a given query must add up to 1. We then use these scores to weight the value matrix through the \(PV\) matmul.

Now, one more important aspect: along which dimension of \(S = QK^\top\) do we apply softmax? The shape of \(S\) is \((b, s_1, s_2)\), where \(s_1\) is the query sequence length and \(s_2\) is the key sequence length. The last dimension is the key’s sequence length, and softmax is applied along this one, so that for each query we get the importance of every key and can extract the context.

With the softmax concept in place, we’ll proceed to online softmax. Here, instead of processing the full list at once, we process it in chunks and rescale the partial results as and when needed. At the end, we get the exact same result as if we had processed the whole list at once. To process in chunks and rescale, we need to keep a few temporary variables for bookkeeping. The reason we want this: our bottleneck is that we can’t hold the full list of numbers in memory, and chunk-by-chunk processing avoids that, at the cost of a little bookkeeping. That’s the gist of online softmax. Now let’s look at the actual math.

The math

Before getting to online softmax, we need to address an overflow issue with exponentials. The largest representable value in float16 (fp16) is 65504. For large \(x\), \(e^x\) overflows this range. To mitigate this, we use a trick known as safe softmax: subtract the max before exponentiating.

Subtracting \(m\) guarantees \(x_i - m \le 0\), so every exponent is non-positive, which is the safe region for the exponential.

Note: this can make some terms very small. If \(x_i\) is far below the max, \(e^{x_i - m}\) can underflow to 0 in fp16. Unlike overflow, this is harmless: a term that underflows was already negligibly small, so it contributes essentially nothing to the sum or to the final probability. Overflow produces inf and corrupts the whole result, whereas underflow just drops a value that rounds to zero anyway.

3-pass safe softmax

Notation: \(m_i = \max_{j=1}^{i} x_j\) with \(m_0 = -\infty\); \(d_i = \sum_{j=1}^{i} e^{x_j - m_N}\) with \(d_0 = 0\) (\(d_N\) is the safe-softmax denominator); \(a_i\) is the final softmax value.

This algorithm iterates over \([1, N]\) three times. In the context of self-attention, the \(x_i\) are the pre-softmax logits computed by \(QK^\top\). When \(N\) is large, all \(N\) logits don’t fit in memory at once, so we need to break the computation into chunks rather than processing the full row in one shot. That is exactly what online softmax does (the GPU memory hierarchy that makes this necessary is covered in the next section).

From safe softmax to online softmax

If we fuse equations (5), (6), and (7) into a single loop, we reduce the number of global memory passes from 3 to 1. Unfortunately, we cannot fuse (5) and (6) directly, because (6) depends on \(m_N\), which isn’t known until the first loop finishes.

The fix: define a surrogate sequence

as a stand-in for the original \(d_i := \sum_{j=1}^{i} e^{x_j - m_N}\). This removes the dependency on \(m_N\), and the \(N\)-th terms of both sequences are identical: \(d_N = d_N'\). So we can safely replace \(d_N\) in equation (7) with \(d_N'\).

We can also find a recurrence between \(d_i'\) and \(d_{i-1}'\):

This recurrence only depends on \(m_i\) and \(m_{i-1}\), so we can compute \(m_i\) and \(d_i'\) together in the same loop.

2-pass online softmax

It still takes two passes. And we can reduce it to a single pass to minimize global I/O. We’ll see that in the FA2 section, after we’ve covered the GPU hierarchy and tiling.

References

• From Online Softmax to FlashAttention - by Zihao Ye (notation used in this section is from here)

GPU Memory and Compute Hierarchy

We left online softmax at two passes. The piece still missing is where those chunks live and how we move them between memory levels. To understand that, and to understand why we process attention in chunks at all, we first need a mental model of how a GPU stores and computes data. That is what this section builds.

We’ll build the memory hierarchy first, then the compute side. Tiling, which combines the two, follows in the next section.

The memory hierarchy

Why a hierarchy? Memory on a GPU is layered based on functionality, bandwidth, size, and access pattern. It is very similar to a CPU’s hierarchy of HDD, RAM, cache, and registers. On a GPU, the levels are: global memory (GMEM, also called HBM, high-bandwidth memory, or off-chip memory), L2 cache, shared memory (SMEM, on-chip memory) and L1 cache, and registers. This ordering goes from largest to smallest in size, lowest to highest in bandwidth, and highest to lowest in latency.

Some rough numbers to keep in mind, for the H100 SXM specifically:

• GMEM / HBM3: ~80 GB capacity, ~3.35 TB/s bandwidth, hundreds of cycles of latency.
• L2 cache: ~50 MB, shared across the whole GPU.
• SMEM / L1: up to 228 KB per SM, configurable as a split between shared memory and L1 cache in code. Bandwidth is roughly an order of magnitude higher than HBM, with latency in the tens of cycles.
• Registers: 65536 32-bit registers per SM (often written as 64K), the fastest storage, allocated per thread.

The point of these numbers is the gap: each step down the hierarchy is dramatically faster but dramatically smaller. That tension is what every kernel optimization has to navigate.

There is a clear tradeoff in designing kernels around these hardware characteristics. This is one of the reasons I love this domain: you can clearly see what you gain and what you give up. One example is occupancy. We want to maximize the number of warps resident per SM, but to do so we need enough resources (registers, shared memory) available on that SM, and resource usage depends on how you framed the algorithm and how you use memory and registers. Balancing this is what it takes to achieve high utilization.

I highly recommend Aleksa Gordic’s writeup on NVIDIA GPU anatomy if you want more depth before proceeding.

The compute hierarchy

We’ve summarized the memory hierarchy. Let’s do a similar one for compute.

Kernels are launched as a grid, which contains multiple blocks, each of which is composed of multiple warps, and each warp is 32 threads. The number of threads per warp is fixed by hardware; warps are not directly exposed as a software unit you tune. Blocks and grids are the software-accessible, tunable units. Multiple blocks can be resident within an SM, depending on the block’s resource requirements and the SM’s available resources. All threads within a block can share that SM’s shared memory to communicate and exchange data.

That covers the structure of parallel compute. The compute units themselves, inside each SM, are: CUDA cores, Tensor cores, and Special Function Units (SFUs). Each is built for specific use cases: CUDA cores for general arithmetic, Tensor cores for matrix multiply and fused multiply-add (FMA), and SFUs for transcendental functions like exponentials.

Because these are separate hardware units, they can run in parallel. I mention this because it is one of the optimizations we’ll apply to make attention faster: use the Tensor cores for the matmuls while the CUDA cores or SFUs simultaneously handle the softmax. This is just to plant the idea; we’ll return to it in the warp specialization and ping-pong sections.

You can reference the NVIDIA H100 GPU Whitepaper for more details.

SIMT: the execution model

One more foundational idea, and arguably the foundation of GPU programming: SIMD/SIMT, Single Instruction Multiple Data / Single Instruction Multiple Threads. We issue a single instruction and apply it across many data elements at once. This is the core reason GPUs accelerate deep learning so well.

Matrix multiplication is a natural fit for this model. Each element of the output matrix is an independent dot product: element \((0,0)\) of the result is the dot product of row 0 of \(A\) with column 0 of \(B\), element \((0,1)\) is row 0 of \(A\) with column 1 of \(B\), and so on. In a basic model where one thread computes one output element, every thread runs the same instruction (a dot product) but operates on different data. That is SIMT.

This is also the bridge to the next section. The block/warp/thread structure above is how work is grouped; tiling is how we shape the data so that grouping maps cleanly onto the memory hierarchy.

Tiling

Before defining tiling, let’s see why we need it, and why it is the thing that lets us combine the memory and compute hierarchies. There are two reasons.

First, we can’t fit all the data in fast memory at once, so we break it into chunks and process them piece by piece. (This is exactly the constraint that forced online softmax into chunked processing.)

Second, we want to avoid repeated transfers from GMEM to SMEM. Moving data is where much of the time goes. HBM bandwidth, while large in absolute terms, is far lower than the rate at which the compute units can consume data, so without care the compute units sit idle waiting for data to arrive. These are memory stalls, and we want to minimize them. Since SMEM is small, we can only hold a small portion of the data at a time, which again pushes us toward breaking the data into small pieces.

Why “tiles”? Because we chunk in 2D, along both rows and columns, so each piece is a tile. The Ampere and Hopper Tensor cores consume matrix fragments directly, so tiling has become the norm, and modern libraries like Triton and CuTe expose tiles rather than scalar elements as their basic primitive.

How tiling works

Here is the standard framing, on a generic GEMM \(C = A \times B\) where \(A\) is \(M \times K\), \(B\) is \(K \times N\), and \(C\) is \(M \times N\).

Instead of one thread computing one full output element (one full dot product over the entire \(K\) dimension), we organize the work as:

• An outer loop over output tiles: the \(M \times N\) output is divided into tiles of size \(T_M \times T_N\). Each block of threads is responsible for one output tile.
• An inner loop over the \(K\) dimension: the shared dimension \(K\) is also split into chunks of size \(T_K\). For each step, the block loads one \(T_M \times T_K\) tile of \(A\) and one \(T_K \times T_N\) tile of \(B\) from GMEM into SMEM, multiplies them, and accumulates the partial result into registers.

After the inner loop finishes walking the full \(K\) dimension, the accumulator in registers holds the complete output tile, which is then written back to GMEM once.

In pseudocode:

# C = A @ B
# A: (M, K), B: (K, N), C: (M, N)
# tile sizes: T_M, T_N, T_K

for each output tile (i, j) in the (M/T_M) x (N/T_N) grid:      # outer loop, parallel across blocks
    acc = zeros(T_M, T_N)                                       # lives in registers

    for k_step in range(0, K, T_K):                             # inner loop over K
        A_tile = load A[i*T_M : (i+1)*T_M, k_step : k_step+T_K]  # GMEM -> SMEM
        B_tile = load B[k_step : k_step+T_K, j*T_N : (j+1)*T_N]  # GMEM -> SMEM
        acc += A_tile @ B_tile                                  # compute on SMEM data

    C[i*T_M : (i+1)*T_M, j*T_N : (j+1)*T_N] = acc               # SMEM/registers -> GMEM

The inner \(K\)-loop is pure accumulation: each \(K\)-step loads a fresh pair of tiles and adds its partial product into acc. The loop walks the shared dimension in \(K/T_K\) steps, and only after the last step does acc hold a complete output tile.

Compare this with the naive version, where each output element is computed independently and every access goes to GMEM:

# naive: one thread, one output element
for i in range(M):
    for j in range(N):
        acc = 0
        for k in range(K):
            acc += A[i, k] * B[k, j]    # every read hits GMEM
        C[i, j] = acc

To tie this back to the compute hierarchy from the previous section, here is how the tiling concepts map onto the GPU’s parallel execution units:

The key thing this mapping makes clear: the grid/block split is what gives us parallelism across output tiles, while the SMEM/register split is what gives us data reuse within a tile. Tiling works because it exploits both at once. (The warp-level detail, how a warp’s threads cooperate with the Tensor cores on a fragment, is deferred to the WGMMA section.)

We’ve described what the loops do and how the work maps to hardware. The next subsection makes the payoff concrete: it counts the memory traffic and shows exactly why this structure is faster.

Why tiling reduces memory traffic

The intuition that tiling helps becomes concrete once we count. We’ll count GMEM traffic in elements rather than bytes; the bytes-per-element factor is the same for both versions, so it cancels in the final ratio. For simplicity, assume square tiles of size \(T \times T\) (so \(T_M = T_N = T_K = T\)).

Naive version. Each output element of \(C\) is one dot product over the full \(K\) dimension. Computing it reads \(K\) elements of \(A\) and \(K\) elements of \(B\) from GMEM, so \(2K\) reads per output element. There are \(M \cdot N\) output elements, so total GMEM traffic is:

The useful compute is one multiply and one add per term, over \(K\) terms, for each of the \(M \cdot N\) outputs:

The arithmetic intensity is the ratio of compute to memory traffic:

One FLOP per element moved. The kernel is entirely memory-bound: the compute units do a single operation for every value dragged across the slow GMEM boundary.

Tiled version. Now the output is divided into \((M/T) \times (N/T)\) tiles. For each output tile, the inner loop walks the \(K\) dimension in \(K/T\) steps. At each step it loads one \(T \times T\) tile of \(A\) and one \(T \times T\) tile of \(B\), which is \(2T^2\) elements. So the GMEM traffic for a single output tile is:

There are \(\dfrac{M \cdot N}{T^2}\) output tiles in total, so:

The compute \(W\) is unchanged, \(2 \cdot M \cdot N \cdot K\) FLOPs, since tiling reorganizes the work but does not remove any of it. The arithmetic intensity becomes:

The ratio. Comparing the two:

Tiling improves arithmetic intensity by a factor of \(T\), the tile size. Equivalently, it cuts GMEM traffic by a factor of \(T\) for the same amount of compute.

Where the factor of \(T\) comes from. It is worth seeing why the ratio is exactly \(T\), because the reuse happens at a finer level than “load tile A, load tile B.” The line acc += A_tile @ B_tile looks like a single operation, but multiplying two \(T \times T\) tiles is itself a \(T \times T \times T\) triple loop: every row of A_tile must be dotted with every column of B_tile. So a single row of A_tile, once it sits in SMEM, is read \(T\) times, once for each output column, rather than being fetched from GMEM \(T\) separate times. Likewise each column of B_tile is reused across all \(T\) rows. In numbers: the tile-tile multiply does \(2T^3\) FLOPs of compute on only \(2T^2\) loaded elements, so each loaded element is touched \(T\) times. That is the factor of \(T\). Note that this reuse lives inside the tile-tile multiply, not across \(K\)-steps: each \(K\)-step still loads a fresh pair of tiles, but each loaded tile, once in SMEM, earns \(T\) uses before being discarded.

What the factor of \(T\) buys you. To see why this matters, compare arithmetic intensity against real hardware. For H100 SXM in BF16, the Tensor cores deliver roughly 990 TFLOP/s of compute, while HBM3 delivers roughly 3.35 TB/s of bandwidth. To stay consistent with our element-based traffic count, we convert that bandwidth into elements per second: at 2 bytes per BF16 element, 3.35 TB/s is about \(1.675 \times 10^{12}\) elements/s. Dividing compute by bandwidth gives the hardware’s break-even arithmetic intensity:

This is the roofline model: a kernel’s attainable performance is capped by whichever constraint binds first, memory bandwidth or peak compute. The crossover is the ridge point: any kernel below it is memory-bound, any kernel above it is compute-bound. A naive GEMM sits at \(I = 1\), nearly 600× below the threshold. The Tensor cores can compute almost 600 times faster than HBM can feed them, so they spend the overwhelming majority of their time idle, waiting on data. Tiling raises \(I\) from 1 toward \(T\), moving the kernel up the roofline, toward the regime where compute, not memory, is the limit.

This is the whole game. Tiling does not reduce compute; it reduces how many times each value crosses the slow GMEM boundary, moving the kernel from memory-bound (\(I = 1\)) toward compute-bound (\(I = T\)), which is the regime where the Tensor cores can actually be kept busy.

We’ll cover how this same tiling idea applies to attention, where the second matmul (\(P \times V\)) depends on the softmax output of the first, in the FA2 section. The attention case has an extra wrinkle that plain GEMM does not, and online softmax is what resolves it.

For a deeper, code-level walkthrough of GEMM tiling, see Simon’s GEMM worklog.

References

• How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklog - Simon Boehm
• Inside NVIDIA GPUs: Anatomy of high performance matmul kernels - Aleksa Gordić
• NVIDIA H100 GPU Whitepaper

FlashAttention-2 (FA2)

Now that we have all the basic building blocks we need to construct FlashAttention, let’s do it.

First and foremost: FA1 and FA2 are very similar. They differ only in loop ordering and in what is kept in shared memory. We’ll look at FA2 directly. The algorithm at its core is also the same for FA3, so getting it clear here lets us relate to and build on it for the rest of the blog.

Before we get into the GPU-specific details, let’s do the math-heavy part first. The UWashington paper by Zihao Ye does a very good job of this, so we’ll use it as our reference here.

A quick note on notation before the derivation begins. We adopt the paper’s bracket-index style for this section, because shapes will matter at every step and bracket notation makes them self-documenting:

• \(Q[k, :]\): the \(k\)-th row of the \(Q\) matrix, shape \((1, d)\).
• \(K^\top[:, i]\): the \(i\)-th column of \(K^\top\), shape \((d, 1)\).
• \(V[i, :]\): the \(i\)-th row of the \(V\) matrix, shape \((1, d)\).
• \(O[k, :]\): the \(k\)-th row of the output matrix \(O\), shape \((1, d)\).

The running scalars and row vectors carry over from the online-softmax section: \(m_i\) is the running max, \(d_i'\) is the running (surrogate) softmax denominator (both scalars), and we now add \(o_i'\), the running (surrogate) output accumulator (a row vector of shape \((1, d)\)). The computation of every output row is independent, so we explain it for a single row.

Also, as in the online-softmax section, we drop the \(1/\sqrt{d}\) softmax scale and the causal mask here for simplicity. Both are present in real code: the scale keeps the dot products from saturating the softmax as \(d\) grows, and the causal mask ensures each query attends only to keys at or before its position. We omit them only to keep the recurrence readable.

From multi-pass to single-pass attention

In the online-softmax section we asked whether softmax could be done in one pass and the answer was no. But in self-attention our final target is not the attention-score matrix \(A\). It is the output matrix \(O = A V\). So the better question is: can we find a one-pass recurrence for \(O\) instead?

With the notation in hand from the section intro, we add three more symbols specific to this derivation:

• \(x_i\): a scalar, the \(i\)-th attention score. Comes from a row times a column: \(x_i = Q[k, :]\, K^\top[:, i]\).
• \(a_i\): a scalar, the \(i\)-th softmax weight.
• \(o_i\): a row vector of shape \((1, d)\), the partial aggregation \(\sum_{j=1}^{i} a_j\, V[j, :]\). This is the running output state.

We can now write the per-row attention computation as a recurrence. We reuse the 2-pass online softmax from before, and add a second loop that accumulates the output row.

Multi-pass self-attention.

A few things to read off this:

• In equation (10), \(a_i\) is scalar: \(x_i, m_N, d_N'\) are all scalars.
• In equation (11), \(a_i\, V[i, :]\) is a scalar times a row vector, producing a row vector of shape \((1, d)\). Adding it to \(o_{i-1}\) (also \((1, d)\)) gives \(o_i\) of shape \((1, d)\). So \(o_i\) accumulates a sum of \(d\)-dimensional row vectors, one per loop iteration.
• After the loop, \(o_N = \sum_{j=1}^{N} a_j\, V[j, :]\), which is the \(k\)-th row of \(O\).

This is multi-pass: the second loop cannot start until the first finishes, because \(a_i\) in equation (10) depends on \(m_N\) and \(d_N'\), which are only known after the first loop completes.

Fusing into a single pass

To fuse the two loops we have to remove the dependency on \(m_N\) and \(d_N'\). Start by substituting the definition of \(a_i\) from equation (10) into equation (11):

Same dependency problem: \(m_N\) and \(d_N'\) aren’t available inside the loop. We play the same surrogate trick we used for \(d_i'\) in the online-softmax section: define a surrogate sequence \(o_i'\) that uses the running values \(m_i\) and \(d_i'\) instead of the final \(m_N\) and \(d_N'\):

The \(N\)-th terms of the two sequences agree. Setting \(i = N\) in the surrogate definition, the running \(m_i\) and \(d_i'\) become the final \(m_N\) and \(d_N'\):

The right-hand side is exactly the definition of \(o_N\) from equation (12). So computing the surrogate \(o_N'\) gives us the true output. The reason this works is the same as for \(d_i'\): only the final term of the sequence is needed, and at the final term the running statistics have caught up to the global ones. Revisit the online-softmax section for the fuller argument.

Next we find a recurrence between \(o_i'\) and \(o_{i-1}'\). The goal is to write \(o_i'\) in terms of \(o_{i-1}'\), a constant, and the data of step \(i\). We walk it step by step.

Start from the definition of \(o_i'\) and split off the last term (\(j = i\)):

Now we need to massage the first sum so it matches the definition of \(o_{i-1}'\). By definition, \(o_{i-1}' = \sum_{j=1}^{i-1} \frac{e^{x_j - m_{i-1}}}{d_{i-1}'} V[j, :]\). The first sum above has \(m_i\) and \(d_i'\) where we’d like \(m_{i-1}\) and \(d_{i-1}'\). We change them one at a time.

First, rewrite the exponent. Using \(e^{x_j - m_i} = e^{x_j - m_{i-1} + m_{i-1} - m_i} = e^{x_j - m_{i-1}}\, e^{m_{i-1} - m_i}\):

The factor \(e^{m_{i-1} - m_i}\) doesn’t depend on \(j\), so it can come out of the sum:

Now we still have \(d_i'\) where we’d like \(d_{i-1}'\). Multiply and divide by \(d_{i-1}'\):

The sum is now exactly \(o_{i-1}'\). Substituting:

Plug back into the split form of \(o_i'\):

The recurrence (13) depends only on \(d_i', d_{i-1}', m_i, m_{i-1}, x_i,\) and \(V[i, :]\). Every one of these is available within a single loop iteration, so we can fuse the entire attention computation into one pass:

FlashAttention (single-pass).

A note on the shapes inside the loop: every operation on the right side of the \(o_i'\) update is a scalar coefficient multiplying either \(o_{i-1}'\) (a \((1, d)\) row) or \(V[i, :]\) (a \((1, d)\) row). The result is a \((1, d)\) row, matching \(o_i'\)’s shape. The expensive part of one iteration is the row addition; the rest is scalar arithmetic.

The states \(x_i, m_i, d_i',\) and \(o_i'\) all have small footprints that fit comfortably in GPU shared memory: three scalars and one \((1, d)\) row per query. This is the property that makes FlashAttention work. We never materialize the full \(N \times N\) score matrix, only a handful of running states per row.

Tiling the algorithm

The single-pass algorithm processes the entire sequence as a loop over individual key/value indices. Tiling changes how many rows we process per loop step, in two independent ways. We do them one at a time. Step 1 follows the paper exactly: tile the key/value sequence dimension. Step 2 extends the same idea to query rows, which lifts us into the form the real kernel uses (Algorithm 1).

Step 1: tiling the key/value rows.

The paper does this step first because all operations in the single-pass algorithm are associative: nothing forces us to process one \(K^\top\) column and one \(V\) row at a time. We can process a block of them in one loop step, as long as the recurrence’s max-tracking and rescaling still work.

We do not touch the query row. \(Q[k, :]\) stays whole. What we tile is the key/value sequence dimension: instead of walking it one row at a time, we walk it \(b\) rows at a time, where \(b\) is the tile size.

New notations for this step (matching the paper):

• \(b\): the block size of the tile.
• \(\#\text{tiles} = N / b\): number of tiles in one row of the score matrix.
• \(x_i\): now a row vector of shape \((1, b)\), holding the dot products of \(Q[k, :]\) against the \(i\)-th tile of \(K^\top\). So \(x_i = Q[k, :]\, K^\top[:, (i-1)b : ib]\), covering columns \((i-1)b + 1\) through \(ib\).
• \(x_i[j]\): the \(j\)-th entry of \(x_i\) (scalar). It’s the score of \(Q[k, :]\) against the \(j\)-th key row inside tile \(i\).
• \(m_i^{(\text{local})}\): scalar, the local maximum within \(x_i\).
• \(m_i, d_i', o_i'\): same running states as before (two scalars and one row vector of shape \((1, d)\)). They now update per tile instead of per row.

With those, the tiled algorithm is:

Reading it: \(m_i^{(\text{local})}\) is the max within the current tile only, and the regular \(m_i\) folds that local max into the running global max. The \(d_i'\) update is the same recurrence as before, except the new contribution is a sum over the \(b\) entries of \(x_i\) instead of a single \(e^{x_i - m_i}\). The \(o_i'\) update similarly: the per-step contribution is now a sum of \(b\) scaled value rows. Nothing in the algorithm structure changed; we just process \(b\) entries per loop step instead of one.

This is the form the paper stops at. The states \(x_i, m_i, d_i', o_i'\) all still have small footprints that fit comfortably in GPU shared memory (the largest is \(o_i'\), one row vector of \(d\) elements). One loop step now consumes a whole \((b, d)\) tile of \(K^\top\) and \(V\), which is what lets each tile be loaded into SMEM once and fully consumed there.

Step 2: tiling the query rows.

The paper’s algorithm stops at K/V tiling, but the real kernel (Algorithm 1 below) also tiles query rows. So we take one more step.

So far we have processed one query row \(Q[k, :]\). But every output row is computed independently: row \(k\) of \(O\) never depends on row \(k'\). That independence means we don’t need a new recurrence to handle many query rows. We just run the Step 1 recurrence for a block of query rows at once, and every per-row quantity gains a row dimension. The paper’s \(b\) becomes \(B_c\) (the K/V tile size, “c” for the column dimension of the score matrix), and we introduce \(B_r\) (the Q tile size, “r” for the row dimension).

Concretely, lift each shape:

• The query block becomes \(Q_{[\,\cdot\,]}\) of shape \((B_r, d)\) instead of one row.
• The score \(x_i\) becomes a matrix of shape \((B_r, B_c)\), where row \(r\) contains the scores of query row \(r\) against all \(B_c\) key rows of tile \(i\).
• The running statistics \(m_i\) and \(d_i'\) become vectors of shape \((B_r,)\), one value per query row.
• The running output \(o_i'\) becomes a block of shape \((B_r, d)\).
• \(V_{[(i-1)B_c : iB_c]}\) is the \(i\)-th K/V tile of shape \((B_c, d)\), same as Step 1.

The scalar reductions inside the loop become per-row reductions: the max over a tile becomes a rowmax (reduce the \(B_c\) axis, one max per query row), and the sum becomes a rowsum. The recurrence is otherwise identical:

A note on the shapes inside the loop: \(m_i\) is a \((B_r,)\) vector, so in \(x_i - m_i\) it broadcasts across the \(B_c\) columns of the \((B_r, B_c)\) score matrix. In the \(o_i'\) update, the per-row scalars \(\frac{d_{i-1}'}{d_i'} e^{m_{i-1} - m_i}\) and \(\frac{1}{d_i'}\) multiply each query row of the output block. \(\text{diag}(\cdot)\) is the cleanest notation for this: \(\text{diag}(v)\) turns a length-\(B_r\) vector \(v\) into a \(B_r \times B_r\) diagonal matrix, so that \(\text{diag}(v)\, M\) scales row \(r\) of matrix \(M\) by entry \(v[r]\). There is no new algebra here, just the Step 1 recurrence with a row index attached.

With both tilings in place, one loop step now processes a full \((B_r, B_c)\) score tile: a block of query rows against a block of key/value rows, with all running state held as small \((B_r,)\) vectors and one \((B_r, d)\) output block. This is the form that maps directly onto the GPU, which is the next subsection.

Mapping to hardware

The general mechanics of tiling, the grid/block/warp split, and SMEM-versus-register reuse were covered in the tiling section. Here we cover what is specific to attention, and one important point up front: the hardware mapping does not line up one-to-one with the clean mental model from the tiled algorithm above. The algorithm is a sequential recurrence over rows and tiles; the hardware is a mesh of warps and Tensor cores that only goes fast when work is shaped to fit them. Bridging that gap is exactly what the kernel sections (K1–K8) are about. Here we set up the correct picture so the algorithm walkthrough that follows is straightforward.

The outer loop: blocks over query tiles.

This part does match the mental model. The work is parallelized over query blocks: each thread block takes one \(B_r \times d\) block of query rows and computes the corresponding \(B_r \times d\) block of output rows. Output rows are independent (the property we used to tile Q), so blocks are fully independent and the GPU schedules them across SMs with no inter-block communication. This is the outer loop, line 3 of Algorithm 1.

The inner loop: where the mental model and the hardware diverge.

In the tiled algorithm, the inner loop reads as a simple sequential sweep: for each key/value tile, compute the score tile, update \(m\), \(d'\), \(o'\), move on. It is tempting to map this as “one thread (or one warp) owns one query row and runs the recurrence for it.” That model is intuitive but wrong, and it is worth seeing why, because the real mapping is what every later optimization builds on.

The inner loop’s heavy work is two matmuls:

• \(S = Q_i K_j^\top\), producing a \(B_r \times B_c\) score tile (the first matmul),
• \(\tilde{P} V_j\), producing the \(B_r \times d\) output contribution (the second matmul).

A matmul tile on a modern GPU is not computed by independent per-row threads. It is computed cooperatively by warps through the Tensor cores. A Tensor core consumes a fragment: a small matrix shard whose elements are spread across all 32 lanes of a warp. A single thread does not hold a full row of the result, it holds a few scattered elements of the fragment. So the \(B_r\) query rows of the block are partitioned across the block’s warps as a byproduct of the matmul fragment layout, not by an explicit “row \(r\) → thread \(r\)” assignment.

This has a direct consequence for the softmax. The rowmax and rowsum in the recurrence are reductions along a row, over the \(B_c\) score values for a given query row. But after the first matmul, those \(B_c\) values for one row are spread across multiple lanes of a warp. Computing the max or sum for that row therefore requires warp-level coordination, lanes exchanging partial results through warp-shuffle instructions, not a single thread looping over the row. So the inner loop is genuinely cooperative: warps work together on the matmul tiles, and again on the cross-lane reductions for the softmax statistics, and again on the rescaling of the running output.

So the corrected picture of one inner-loop step is: the block’s warps cooperatively compute the \(S\) tile on the Tensor cores, cooperatively reduce it row-wise (warp shuffles) to get \(m^{(\text{local})}\) and the denominator update, cooperatively rescale the running output, and cooperatively compute the \(\tilde{P} V\) matmul on the Tensor cores. The sequential recurrence is still exactly what executes, it is just executed by a team of warps per step rather than by one thread per row.

The precise fragment layout, how many warps cooperate, and how the two matmuls are scheduled relative to the softmax, depend on the GPU generation. On Hopper this is the warpgroup-level WGMMA instruction, and getting it right (overlapping the matmuls with the softmax, keeping the Tensor cores fed) is the substance of the later kernel sections. We will cover WGMMA and that scheduling in the K4 section; here it is enough to hold the model that the inner loop is warp-cooperative, not per-thread.

The sweep direction itself is simple: a block walks the \(K^\top\) tiles left to right and the \(V\) tiles top to bottom, updating the running state \(m\), \(d'\), \(O\) after each tile.

Why the footprint is what makes this work.

The SRAM footprint of a block depends only on the block sizes (\(B_r\), \(B_c\)) and the head dimension \(d\), never on the sequence length \(L\). At any moment a block holds only its current Q tile, one K tile, one V tile, and the small \((B_r,)\) running-state vectors, all sized by the block parameters and \(d\). The \(N \times N\) score matrix is never materialized. Since shared memory is small (228 KB/SM on H100) and \(L\) can be very large (16k or more), this \(L\)-independence is exactly what lets FlashAttention scale to long context without running out of memory.

This connects back to the occupancy tradeoff from the memory and compute hierarchy section: because the footprint is fixed by the block parameters and \(d\), the block size becomes a tuning knob. Larger blocks mean more reuse per tile but more SRAM and registers per block, which lowers how many blocks can be resident per SM. That balance is exactly the kind of tradeoff we’ll revisit when we profile real kernels.

With this mapping in mind, blocks over query tiles, warp-cooperative inner loop over key/value tiles, two Tensor-core matmuls bridged by a warp-reduced softmax, the full algorithm reads directly off the hardware.

The full algorithm: FlashAttention-2 forward pass

We have now built the algorithm twice: as a per-row recurrence, and as the \(B_r \times B_c\) tiled form that maps onto the hardware. Algorithm 1 below is that same algorithm exactly as the FA2 paper states it. It is kept in the paper’s own notation so it can serve as a faithful reference; the mapping back to the blog’s notation follows.

Algorithm 1 — FlashAttention-2 forward pass

Require: Matrices \(Q, K, V \in \mathbb{R}^{N \times d}\) in HBM, block sizes \(B_c\), \(B_r\).

1. Divide \(Q\) into \(T_r = \lceil N / B_r \rceil\) blocks \(Q_1, \ldots, Q_{T_r}\) of size \(B_r \times d\) each, and divide \(K, V\) into \(T_c = \lceil N / B_c \rceil\) blocks \(K_1, \ldots, K_{T_c}\) and \(V_1, \ldots, V_{T_c}\), of size \(B_c \times d\) each.
2. Divide the output \(O \in \mathbb{R}^{N \times d}\) into \(T_r\) blocks \(O_1, \ldots, O_{T_r}\) of size \(B_r \times d\) each, and divide the logsumexp \(L\) into \(T_r\) blocks \(L_1, \ldots, L_{T_r}\) of size \(B_r\) each.
3. for \(1 \le i \le T_r\) do
4.     Load \(Q_i\) from HBM to on-chip SRAM.
5.     On chip, initialize \(O_i^{(0)} = (0)_{B_r \times d}\), \(\ell_i^{(0)} = (0)_{B_r}\), \(m_i^{(0)} = (-\infty)_{B_r}\).
6.     for \(1 \le j \le T_c\) do
7.         Load \(K_j, V_j\) from HBM to on-chip SRAM.
8.         On chip, compute \(S_i^{(j)} = Q_i K_j^\top \in \mathbb{R}^{B_r \times B_c}\).
9.         On chip, compute \(m_i^{(j)} = \max\!\left(m_i^{(j-1)},\ \text{rowmax}(S_i^{(j)})\right) \in \mathbb{R}^{B_r}\), \(\tilde{P}_i^{(j)} = \exp\!\left(S_i^{(j)} - m_i^{(j)}\right) \in \mathbb{R}^{B_r \times B_c}\) (pointwise), \(\ell_i^{(j)} = e^{\,m_i^{(j-1)} - m_i^{(j)}} \ell_i^{(j-1)} + \text{rowsum}(\tilde{P}_i^{(j)}) \in \mathbb{R}^{B_r}\).
10.        On chip, compute \(O_i^{(j)} = \text{diag}\!\left(e^{\,m_i^{(j-1)} - m_i^{(j)}}\right) O_i^{(j-1)} + \tilde{P}_i^{(j)} V_j\).
11.     end for
12.     On chip, compute \(O_i = \text{diag}\!\left(\ell_i^{(T_c)}\right)^{-1} O_i^{(T_c)}\).
13.     On chip, compute \(L_i = m_i^{(T_c)} + \log\!\left(\ell_i^{(T_c)}\right)\).
14.     Write \(O_i\) to HBM as the \(i\)-th block of \(O\).
15.     Write \(L_i\) to HBM as the \(i\)-th block of \(L\).
16. end for
17. Return the output \(O\) and the logsumexp \(L\).

Reading the algorithm with what we’ve built. Every line maps onto something we’ve already derived. The paper’s notation lines up with the blog’s as follows: \(S_i^{(j)}\) is our score tile \(x_i\), \(\ell\) is our denominator \(d_i'\), \(m_i^{(j)}\) is our running max \(m_i\), \(\tilde{P}\) is the unnormalized \(e^{x_i - m_i}\) term, and \(O_i^{(j)}\) is our running output \(o_i'\). The superscript \((j)\) is the inner-loop step.

• Lines 1–2 are the tiling setup: \(Q\) is cut into \(T_r\) row blocks of size \(B_r\), and \(K, V\) into \(T_c\) row blocks of size \(B_c\), the two tile sizes from the Q-tiling and K/V-tiling steps above.
• Line 3 iterates over query blocks; line 6 iterates over key/value tiles. The hardware actors behind these loops, one independent thread block per query block, a warp-cooperative sweep inside, were covered in the mapping subsection.
• Lines 4, 7 are the GMEM → SRAM loads: one Q block once per outer iteration, one K and V block per inner step.
• Line 8 is the first matmul, \(S = QK^\top\), producing a \(B_r \times B_c\) score tile.
• Line 9 is the online softmax, batched over \(B_r\) rows. The rowmax and rowsum here are exactly the warp-reduced steps from the hardware mapping above: a row’s \(B_c\) score values are spread across warp lanes, so these reductions run as warp shuffles. \(\ell\) is the running denominator \(d_i'\), and its update is the same surrogate recurrence we derived.
• Line 10 is the output recurrence. \(O_i^{(j)}\) is the running output \(o_i'\). The \(\text{diag}(e^{m^{(j-1)} - m^{(j)}})\) factor rescales the previous partial output when the max changes, and \(\tilde{P}_i^{(j)} V_j\) is the second matmul, \(P \times V\), adding the current tile’s contribution.

The one real difference: rescaling is deferred. Look closely at line 10 versus our earlier \(o_i'\) recurrence. Our recurrence divided by \(d_i'\) at every step. Algorithm 1 does not divide inside the inner loop. \(O_i^{(j)}\) is kept unnormalized throughout the sweep, rescaled only by the max-correction factor \(e^{m^{(j-1)} - m^{(j)}}\). The division by the denominator happens exactly once, on line 12, after the inner loop finishes: \(O_i = \text{diag}(\ell_i^{(T_c)})^{-1} O_i^{(T_c)}\).

This is the key FA1 → FA2 change. FA1 rescaled the output by the denominator on every inner step; FA2 moves that division out of the inner loop and applies it a single time at the end. The inner loop then contains only matmuls and the cheap max-correction, with the expensive per-step normalization removed. Fewer non-matmul operations on the hot path means the Tensor cores stay busier, and combined with the loop reordering (outer loop over query blocks, which parallelizes cleanly across thread blocks), this is most of where FA2’s speedup over FA1 comes from.

Line 13, the logsumexp. \(L_i = m_i^{(T_c)} + \log(\ell_i^{(T_c)})\) is only needed for the backward pass, where it lets the gradient recompute the softmax without storing the full probability matrix. Since this blog covers only the forward pass, we include \(L\) for completeness, to keep Algorithm 1 faithful to the paper, but we won’t use it further.

Note: in this section we only cover the forward pass, since inference is the main thing we want to cover. Training additionally needs the backward pass to propagate gradients and update parameters; for that, please refer to the paper.

References

• FlashAttention-2 paper, Tri Dao

CuTe DSL fundamentals and primitives

We’ve built the FlashAttention algorithm and tiled it for the GPU. From here on, every kernel we write (K1 through K8) will be written in CuTe Python DSL, the surface we use to express layouts, tensors, and Tensor-Core operations on Hopper. This section is the foundation: enough CuTe to read and understand the kernels that follow, built from first principles. Deep, kernel-specific mechanics (TMA hardware behavior, WGMMA fragment internals, swizzle and core-matrix details) are deferred to the kernel sections where they actually matter. Here we cover the abstractions.

We’ll go: what CuTe is and why it exists, then layouts, then tensors, then partitioning, then cute.gemm and cute.copy at a high level. Note: I initially added this part here, but it was very long, so I had to move it to a separate blog post to keep the content structured while still explaining most of it. So, pls refer to this CuTe DSL for the full in-depth ideas and understanding.

Experimental setup

Before getting into K1, two practical notes. All measurements in this worklog run on the same hardware and software stack (H100 SXM, fixed CuTe DSL and PyTorch versions, etc.), and all kernels are launched through a single shared harness with the same benchmark and correctness-check conventions. Rather than recapping that in the blog, the details live in the project README: hardware and software versions, the repo layout, the workload sweep (sequence lengths, head dims, causal flag), the tensor layout convention, and how to run any kernel from the command line.

The repo is at: GitHub.

We’ll introduce just enough of the harness inline when K1 lands (one example command, a one-line note on what CHECK=1 does, and what TFLOPS gets reported). From K2 onward, the same conventions apply unchanged. If anything in the setup is unclear while reading a kernel section, the README is the reference.

With that, on to K1.

The kernels are being worked on, and the content for K1-K8 will be available soon, one by one.