From FlashAttention to FlashAttention 2

Introduction

This is a code-level walkthrough of how FlashAttention 2 is implemented, and why it’s faster than FlashAttention 1. We’ll trace the entire path from the Python API down to the CUDA kernels, referencing the Dao-AILab/flash-attention repository.

Background: CUDA Programming & CUTLASS

Before diving into FlashAttention, let’s cover the GPU programming concepts that the implementation relies on.

GPU Architecture: SMs, Warps, and Memory Hierarchy

A GPU is made up of many SMs (Streaming Multiprocessors). An A100 has 108 SMs, an H100 has 132. Each SM contains:

One SM (Streaming Multiprocessor)
┌──────────────────────────────────────┐
│  CUDA Cores (FP32/INT32)             │  ← scalar math
│  Tensor Cores                        │  ← matrix multiply (mma.sync)
│  Load/Store Units                    │  ← memory access
│  Warp Schedulers                     │  ← pick which warp runs next
│                                      │
│  Shared Memory / L1 Cache (164KB)    │  ← on-chip, fast
│  Register File (256KB)               │  ← per-thread, fastest
└──────────────────────────────────────┘

A warp is a group of 32 threads that execute in lockstep. A CTA (Cooperative Thread Array) — also called a thread block — is a group of warps (e.g., 4 warps = 128 threads) that share the same shared memory and run on the same SM.

Warp size (32 threads) is fixed by hardware — not controllable. The number of warps per CTA (kNWarps) is the tunable knob, which determines the thread block size: kNThreads = kNWarps * 32. For example, kNWarps=4 means 128 threads, kNWarps=8 means 256 threads. Larger head dimensions need more warps because the tiles are bigger — more data to load in parallel, more compute to keep the Tensor Cores fed.

Multiple CTAs can run simultaneously on one SM, but each CTA gets its own private chunk of shared memory — they cannot see each other’s. The SM’s shared memory pool is partitioned between them:

SM shared memory pool: 164KB (A100)
┌─────────────────────────────────┐
│  CTA 0: 96KB smem               │
├─────────────────────────────────┤
│  CTA 1: 48KB smem               │  ← only if total ≤ 164KB
├─────────────────────────────────┤
│  unused: 20KB                    │
└─────────────────────────────────┘

This is why tile size and occupancy (CTAs per SM) trade off: smaller tiles use less shared memory, allowing more CTAs per SM, which means better latency hiding when one CTA stalls on memory.

Memory Hierarchy: HBM vs Shared Memory vs Registers

                Speed        Size         Scope
Registers       fastest      256KB/SM     per-thread
Shared Memory   ~19 TB/s     164KB/SM     per-CTA (private)
HBM             ~2 TB/s      40-80GB      global (all CTAs)

The whole motivation for FlashAttention is that standard attention writes the $N \times N$ matrix to HBM, which is ~10x slower than shared memory. By keeping intermediate results in shared memory and registers, we avoid this bottleneck.

Tensor Cores

Each SM contains Tensor Cores — specialized hardware units for matrix multiply. On A100 (3rd gen), each SM has 4 Tensor Cores. When a warp issues an mma.sync instruction, a Tensor Core executes a small matrix multiply (e.g., 16x8x16) in one shot — much faster than doing it with scalar CUDA cores.

Tiles: Chunking Matrices to Fit On-Chip

A tile is a block-sized chunk of a matrix that fits in shared memory or registers — a logical concept, not a hardware construct. In code, it’s just a pointer + shape + stride into the full matrix.

When matrices are too large for on-chip memory, we process them in tiles:

Full Q: (N, d)              Full K: (N, d)
┌──────────────────┐        ┌──────────────────┐
│                  │        │                  │
│                  │        │  K tile          │
│  Q tile          │        │  (kBlockN × d)   │
│  (kBlockM × d)   │        │                  │
│                  │        ├──────────────────┤
│                  │        │                  │
├──────────────────┤        │                  │
│                  │        │                  │
│                  │        └──────────────────┘
└──────────────────┘

One thread block loads one Q tile into shared memory, then iterates over K/V tiles:

For each K tile:
    Load K tile into smem              ← kBlockN × d
    Compute S_tile = Q_tile @ K_tile^T ← kBlockM × kBlockN (in registers)
    Load V tile into smem              ← kBlockN × d
    Accumulate O_tile += softmax(S_tile) @ V_tile

With kBlockM=128, kBlockN=64, kHeadDim=128:

• Q tile: 128 x 128 elements in smem
• K tile: 64 x 128 elements in smem
• V tile: 64 x 128 elements in smem
• S tile: 128 x 64 — computed in registers, never written to HBM

The tile sizes (kBlockM, kBlockN) are chosen by benchmarking the tradeoff between competing constraints:

• Shared memory capacity: The tile must fit. smem ≈ (kBlockM * kHeadDim + 2 * kBlockN * kHeadDim) * 2 bytes.
• Occupancy: Smaller tiles → less smem per CTA → more CTAs per SM → better latency hiding.
• Compute efficiency: Larger tiles amortize softmax/masking overhead over more matmul FLOPs.
• Causal masking waste: Roughly half the S tile is masked out. Square tiles (kBlockM ≈ kBlockN) waste less.
• Register pressure: The output accumulator acc_o lives in registers — too large and it spills to slow local memory.

Row-Major vs Column-Major

For a 2D matrix, there are two ways to lay it out in linear memory:

Row-major — elements in the same row are contiguous:

Matrix:     Memory:
[a b c]     [a b c d e f]
[d e f]

addr(i, j) = i * num_cols + j
stride = (num_cols, 1)

Column-major — elements in the same column are contiguous:

Matrix:     Memory:
[a b c]     [a d b e c f]
[d e f]

addr(i, j) = j * num_rows + i
stride = (1, num_rows)

The key difference is which stride is 1. In CuTe notation: row-major is Stride<num_cols, _1>, column-major is Stride<_1, num_rows>. This matters because GPU memory accesses are coalesced when adjacent threads read adjacent (stride-1) addresses.

Background: FlashAttention 1

Standard attention computes $O = \text{softmax}(QK^T / \sqrt{d})V$ by materializing the full $N \times N$ attention matrix $S = QK^T$ in GPU HBM. This is both memory-hungry ($O(N^2)$) and IO-bound — the GPU spends most of its time shuttling data between HBM and the compute units.

FlashAttention 1 (Dao et al., 2022) solves this with two key ideas:

Tiling

Instead of computing the full $S$ matrix, Q is split into blocks of $B_r$ rows, and K/V are split into blocks of $B_c$ rows. Each thread block loads one Q block into shared memory, then iterates over K/V blocks, computing a $B_r \times B_c$ tile of $S$ at a time. The full $N \times N$ matrix is never materialized in HBM. (See Tiles above for details on how tiling works.)

Online Softmax

The challenge with tiling is that softmax requires the full row of $S$ to normalize. FlashAttention uses the “online softmax” trick: as we process each K block, we maintain running statistics — the row-wise max $m_i$ and the sum of exponentials $\ell_i$. When a new K block produces a new max, the running output accumulator $O$ and sum $\ell$ are rescaled:

\[m_i^{\text{new}} = \max(m_i^{\text{old}}, \tilde{m}_i)\] \[\ell_i^{\text{new}} = e^{m_i^{\text{old}} - m_i^{\text{new}}} \ell_i^{\text{old}} + e^{\tilde{m}_i - m_i^{\text{new}}} \tilde{\ell}_i\] \[O_i^{\text{new}} = \frac{\ell_i^{\text{old}} e^{m_i^{\text{old}} - m_i^{\text{new}}}}{\ell_i^{\text{new}}} O_i^{\text{old}} + \frac{e^{\tilde{m}_i - m_i^{\text{new}}}}{\ell_i^{\text{new}}} \tilde{P}_i V_j\]

This produces numerically exact attention with $O(N)$ HBM accesses per row instead of $O(N^2)$.

Backward Pass: Recomputation

Instead of storing the full $S$ matrix for the backward pass (which would cost $O(N^2)$ memory), FlashAttention stores only the output $O$ and the logsumexp $L = m + \log(\ell)$ per row. During the backward pass, it recomputes $S = QK^T$ on-the-fly. This trades extra FLOPs for a massive reduction in memory — and since attention is IO-bound anyway, the extra compute is nearly free.

FlashAttention 2: What Changed

FlashAttention 2 (Dao, 2023) achieves ~2x speedup over FA1 through three algorithmic improvements:

1. Reducing Non-Matmul FLOPs

FA1 performs the online softmax rescaling (exponentiation, max, sum) on every element of $S$, which adds up to significant non-matmul FLOPs. FA2 restructures the algorithm to delay normalization: instead of normalizing $O$ after each K block, it accumulates the un-normalized result and normalizes only once at the very end. This reduces the rescaling work from $O(N^2 d)$ to $O(N^2 + Nd)$.

2. Swapping the Loop Order — Parallelism over Sequence Length

FA1’s outer loop iterates over K/V blocks and the inner loop over Q blocks. This means dQ accumulation requires atomic adds. FA2 swaps the loops: the outer loop is over Q blocks (each one assigned to a thread block), and the inner loop iterates over K/V blocks. Each thread block now owns its Q block exclusively and accumulates $O$ in registers without atomics.

3. Better Work Partitioning Between Warps

FA1 splits work across warps along both Q and K dimensions, requiring shared memory communication. FA2 splits warps only along the Q rows dimension (not K columns). Each warp computes the full $S = QK^T$ for its Q rows against all K columns, then independently multiplies with $V$. This eliminates the cross-warp synchronization through shared memory that FA1 needed for the softmax reduction.

Code-Level Deep Dive

Now let’s trace through the actual implementation. All links reference the Dao-AILab/flash-attention OSS repo.

Layer 1: Python API

The user-facing entry point is flash_attn_func in flash_attn/flash_attn_interface.py. It accepts Q, K, V tensors in shape (batch, seqlen, nheads, headdim) along with options like causal, dropout_p, window_size, softcap, and alibi_slopes.

def flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False,
                    window_size=(-1, -1), softcap=0.0, alibi_slopes=None,
                    deterministic=False, return_attn_probs=False):

This wraps FlashAttnFunc, a torch.autograd.Function that:

1. Forward: Pads headdim to a multiple of 8, calls the C++ forward kernel via flash_attn_2_cuda.fwd(...), and saves Q, K, V, O, softmax_lse, and rng_state for backward.
2. Backward: Allocates dQ, dK, dV, pads dout similarly, and calls flash_attn_2_cuda.bwd(...).

There’s also a variable-length variant flash_attn_varlen_func that takes cu_seqlens_q / cu_seqlens_k for packed batching without padding.

Layer 2: C++ Entry Point

csrc/flash_attn/flash_api.cpp is the C++/CUDA bridge. The key function is run_mha_fwd, which dispatches based on dtype (fp16/bf16) and head dimension:

void run_mha_fwd(Flash_fwd_params &params, cudaStream_t stream, bool force_split_kernel=false) {
    FP16_SWITCH(!params.is_bf16, [&] {
        HEADDIM_SWITCH(params.d, [&] {
            BOOL_SWITCH(params.is_causal, Is_causal, [&] {
                if (params.num_splits <= 1 && !force_split_kernel) {
                    run_mha_fwd_<elem_type, kHeadDim, Is_causal>(params, stream);
                } else {
                    run_mha_fwd_splitkv_dispatch<elem_type, kHeadDim, Is_causal>(params, stream);
                }
            });
        });
    });
}

The HEADDIM_SWITCH macro dispatches to pre-compiled kernels for each supported head dimension (32, 64, 96, 128, 192, 256). Each instantiation lives in its own .cu file (e.g., flash_fwd_hdim128_bf16_causal_sm80.cu) to parallelize compilation.

Layer 3: Kernel Launch Configuration

csrc/flash_attn/src/flash_fwd_launch_template.h defines the grid/block configuration:

const int num_m_block = (params.seqlen_q + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM;
dim3 grid(num_m_block, params.b, params.h);

The grid has three dimensions: Q sequence blocks x batch x heads. Each thread block owns one Q tile of kBlockM rows — this is the FA2 “parallelize over Q” strategy. The typical block sizes for common head dimensions are:

Head dim	kBlockM	kBlockN	Warps	Notes
64	128	128	4	Full utilization
128	128	64	4	A100 default
128	64	64	4	sm86/89 causal
192	128	64	8
256	128	64	8	A100 (128KB smem)
256	64	64	4	H100 (2 CTAs/SM)

These values (including kNWarps — see GPU Architecture) are set at the launch site via Flash_fwd_kernel_traits<kHeadDim, kBlockM, kBlockN, kNWarps>, and different GPU architectures get different configurations based on profiling.

Layer 4: Kernel Traits — The MMA and Memory Configuration

csrc/flash_attn/src/kernel_traits.h defines Flash_fwd_kernel_traits — the template parameters that control everything about how the kernel maps to hardware:

using TiledMma = TiledMMA<
    MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,    // Tensor Core instruction
    Layout<Shape<Int<kNWarps>, _1, _1>>,         // All warps stacked in M dimension
    Tile<Int<16 * kNWarps>, _16, _16>>;

This is a CuTe construct that describes how to map a large matrix multiply onto Tensor Cores. Let’s break down the three parts:

Part 1: MMA_Atom — The Hardware Instruction

MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>

This is one mma.sync PTX instruction that computes a 16 x 8 x 16 matrix multiply: 16 rows of A (Q), 8 columns of B (K), inner dimension 16. Inputs are fp16/bf16, output accumulates in fp32. One warp (32 threads) executes it collaboratively, with each thread holding a few elements of the result in registers.

The _TN suffix indicates the memory layout of the operands. For a 2D matrix, there are two ways to lay it out in linear memory:

• Row-major (stride = (num_cols, 1)): elements in the same row are contiguous. In CuTe: Stride<num_cols, _1>.
• Column-major (stride = (1, num_rows)): elements in the same column are contiguous. In CuTe: Stride<_1, num_rows>.

So _TN means: A is Transposed (column-major input read as row-major), B is Normal (column-major). This matters for GPU performance because memory accesses are coalesced when adjacent threads read adjacent (stride-1) addresses.

Part 2: Warp Layout — The FA2 Design Choice

Layout<Shape<Int<kNWarps>, _1, _1>>   // e.g., 4 x 1 x 1

This controls how warps are assigned to the M, N, K dimensions. With kNWarps=4:

• 4 warps along M (Q rows) — each warp handles a different set of Q rows
• 1 warp along N (K columns) — each warp covers all K columns
• 1 warp along K (inner dimension)

This is the defining FA2 design choice. All warps are stacked vertically on different Q rows. Each warp independently computes the full $S = QK^T$ for its Q rows against all K columns, then independently multiplies with V. No cross-warp communication needed for softmax.

Compare with what FA1 would use: Layout<Shape<_2, _2, _1>> — 2 warps on M, 2 on N. The warps splitting N would need shared memory to communicate partial softmax results.

FA2: kNWarps=4, Layout<4,1,1>          FA1 style: Layout<2,2,1>

S tile (128 x 64)                      S tile (128 x 64)
┌──────────────────────┐               ┌───────────┬───────────┐
│ Warp 0: rows 0-31    │ ← full row   │ Warp 0    │ Warp 1    │
│                      │               │           │           │
├──────────────────────┤               ├───────────┼───────────┤
│ Warp 1: rows 32-63   │               │ Warp 2    │ Warp 3    │
├──────────────────────┤               │           │           │
│ Warp 2: rows 64-95   │               └───────────┴───────────┘
├──────────────────────┤               No warp has a full row!
│ Warp 3: rows 96-127  │               Must communicate via smem
└──────────────────────┘               for softmax reduction.
Each warp has full row →
softmax independently!

Part 3: Tile Size — Output Per MMA Step

Tile<Int<16 * kNWarps>, _16, _16>    // e.g., 64 x 16 x 16

With kNWarps=4, one MMA step across all warps produces a 64 x 16 output:

• 4 warps x 16 rows each = 64 rows of Q
• 16 columns of K (each warp issues two 16x8 mma.sync instructions side by side)
• Inner dimension 16 (one chunk of head_dim consumed per step)

But the full S tile is kBlockM x kBlockN = 128 x 64. So the kernel iterates:

Full S tile: 128 × 64

Step layout (each box = one TiledMMA invocation = 64×16):

         K cols: 0-15    16-31    32-47    48-63
        ┌────────┬────────┬────────┬────────┐
M rows  │ step   │ step   │ step   │ step   │
 0-63   │ (0,0)  │ (0,1)  │ (0,2)  │ (0,3)  │
        ├────────┼────────┼────────┼────────┤
M rows  │ step   │ step   │ step   │ step   │
 64-127 │ (1,0)  │ (1,1)  │ (1,2)  │ (1,3)  │
        └────────┴────────┴────────┴────────┘

128/64 = 2 steps along M
 64/16 = 4 steps along N
Total: 8 positional steps

Each positional step also loops over the K (head_dim) dimension: kHeadDim/16 = 128/16 = 8 inner steps. These are all #pragma unroll‘d at compile time (the sizes are constexpr), so the compiler emits a flat sequence of mma.sync instructions with no loop overhead.

Shared Memory Layout: Swizzling to Avoid Bank Conflicts

using SmemLayoutAtomQ = decltype(
    composition(Swizzle<kSwizzle, 3, 3>{},
                Layout<Shape<_8, Int<kBlockKSmem>>,
                       Stride<Int<kBlockKSmem>, _1>>{}));

The base layout is Layout<Shape<_8, 64>, Stride<64, _1>> (when kBlockKSmem=64): 8 rows, 64 columns, row-major. Row i, col j maps to address i*64 + j.

The problem without swizzling: GPU shared memory has 32 banks (4 bytes each). When a warp reads a column (same j, different i), threads access 0*64+j, 1*64+j, 2*64+j, .... Since 64 % 32 == 0, every row hits the same bank — a 32-way bank conflict that serializes the access.

Swizzle<3, 3, 3> fixes this by XOR-ing bits of the row index into the column address:

addr(row, col) = row * 64 + (col XOR (row << 3))

Now each row’s XOR offset sends it to a different bank, eliminating the conflict entirely.

The full Q shared memory layout tiles this 8x64 atom to cover the full kBlockM x kHeadDim tile:

using SmemLayoutQ = decltype(tile_to_shape(
    SmemLayoutAtomQ{},
    Shape<Int<kBlockM>, Int<kHeadDim>>{}));  // e.g., 128 x 128

Total shared memory holds one Q tile + two KV tiles:

static constexpr int kSmemSize = Share_Q_K_smem
    ? std::max(kSmemQSize, kSmemKVSize)
    : kSmemQSize + kSmemKVSize;

For hdim=128, kBlockM=128, kBlockN=64, this is about 96KB.

Global memory loads use SM80_CP_ASYNC_CACHEGLOBAL — asynchronous copies that bypass L1 cache, overlapping data movement with computation.

Layer 5: The Forward Kernel — The Inner Loop

csrc/flash_attn/src/flash_fwd_kernel.h contains the main forward function compute_attn_1rowblock. Here’s the logical structure:

Setup

// Each thread block handles one Q block
const int m_block = blockIdx.x;  // which Q tile
const int bidb = blockIdx.y;     // batch index
const int bidh = blockIdx.z;     // head index

The function loads one Q block into shared memory and initializes the output accumulator acc_o (in fp32 registers) to zero:

Tensor acc_o = partition_fragment_C(tiled_mma, Shape<Int<kBlockM>, Int<kHeadDim>>{});
clear(acc_o);

The Main Loop (Iterating over K/V blocks)

The kernel loops over K/V blocks in reverse order (from the last block down to block 0). Reverse iteration is an optimization — the last blocks need masking (for causal or uneven sequence lengths), and by handling them first, the remaining iterations skip masking entirely.

The loop has two phases:

Phase 1: Masked iterations — handles the last ceil(kBlockM/kBlockN) K blocks that may need causal masking:

for (int masking_step = 0; masking_step < n_masking_steps; ++masking_step, --n_block) {
    // 1. Compute S = Q @ K^T (into acc_s in registers)
    flash::gemm(acc_s, tSrQ, tSrK, ...);

    // 2. Apply causal mask to acc_s
    mask.apply_mask(acc_s, ...);

    // 3. Online softmax: rescale acc_o by new max, then exp2 on acc_s
    softmax.softmax_rescale_o(acc_s, acc_o, scale_softmax_log2);

    // 4. Convert acc_s from fp32 to fp16/bf16 -> rP
    Tensor rP = convert_type<Element>(acc_s);

    // 5. Compute O += P @ V (accumulate into acc_o)
    flash::gemm_rs(acc_o, tOrP, tOrVt, ...);
}

Phase 2: Non-masked iterations — the fast path, no masking needed:

for (; n_block >= n_block_min; --n_block) {
    // Same as above, but without mask.apply_mask()
    flash::gemm(acc_s, ...);
    softmax.softmax_rescale_o(acc_s, acc_o, ...);
    Tensor rP = convert_type<Element>(acc_s);
    flash::gemm_rs(acc_o, tOrP, tOrVt, ...);
}

Each iteration overlaps data loads with compute using the cp_async pipeline: while computing the current KV block, the next KV block is being asynchronously loaded from HBM into shared memory.

Epilogue: Normalize and Write O

After the loop, the accumulated output is normalized by dividing by the softmax sum:

Tensor lse = softmax.normalize_softmax_lse(acc_o, params.scale_softmax);

Then acc_o is converted from fp32 to fp16/bf16 and written to HBM through shared memory:

Tensor rO = convert_type<Element>(acc_o);
// Write to shared memory, then to global memory
cute::copy(smem_tiled_copy_O, taccOrO, taccOsO);
cute::copy(gmem_tiled_copy_O, tOsO, tOrO);  // smem -> registers -> gmem

The logsumexp lse = max + log(sum) is also stored — it’s needed for the backward pass.

Layer 6: The Online Softmax Implementation

csrc/flash_attn/src/softmax.h implements the Softmax struct with the critical softmax_rescale_o method.

First iteration (Is_first=true):

reduce_max<true>(scores, row_max);                    // row_max = max(S_row)
scale_apply_exp2(scores, row_max, softmax_scale_log2); // S = exp2(S * scale - max * scale)
reduce_sum<true>(scores, row_sum);                     // row_sum = sum(exp(S))

Subsequent iterations (Is_first=false):

// Save old max, compute new max
cute::copy(row_max, scores_max_prev);
reduce_max<false>(scores, row_max);  // row_max = max(old_max, new_max)

// Rescale old accumulator and sum by exp(old_max - new_max)
for (int mi = 0; mi < size(row_max); ++mi) {
    float scores_scale = exp2f((scores_max_prev(mi) - row_max(mi)) * softmax_scale_log2);
    row_sum(mi) *= scores_scale;
    for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) {
        acc_o_rowcol(mi, ni) *= scores_scale;  // Rescale old output
    }
}

// Exponentiate new scores and accumulate sum
scale_apply_exp2(scores, row_max, softmax_scale_log2);
reduce_sum<false>(scores, row_sum);

A subtle optimization: the code uses exp2f (base-2 exponential) instead of expf (base-e). This maps directly to a single PTX instruction and is faster. The scaling factor softmax_scale_log2 = softmax_scale * log2(e) converts accordingly.

The reduce_max function uses Allreduce<4> — a warp-level reduction across 4 threads (a “quad”), matching the Tensor Core MMA layout where 4 threads share the same row. This avoids needing shared memory for the reduction.

Layer 7: The Backward Kernel

csrc/flash_attn/src/flash_bwd_kernel.h implements compute_dq_dk_dv_1colblock. The backward pass is more complex because it computes three gradients.

Key difference from forward: The backward kernel’s outer loop is over K/V blocks (each thread block owns one K/V tile), and the inner loop is over Q blocks. This is because dK and dV can be accumulated in registers, while dQ must be accumulated across multiple thread blocks (requiring either atomics or a split-accumulate strategy).

// Thread block owns one K/V column block
const int n_block = blockIdx.x;

// Accumulators for dK, dV in registers
Tensor acc_dk = partition_fragment_C(tiled_mma_dkv, Shape<Int<kBlockN>, Int<kHeadDim>>{});
Tensor acc_dv = partition_fragment_C(tiled_mma_dkv, Shape<Int<kBlockN>, Int<kHeadDim>>{});
clear(acc_dv);
clear(acc_dk);

The inner loop iterates over Q blocks (in reverse):

for (; m_block >= m_block_min; --m_block) {
    // 1. Recompute S = Q @ K^T (this is the recomputation from forward)
    flash::gemm(acc_s, tSrQ, tSrK, ...);

    // 2. Recompute P = softmax(S) using saved LSE
    scale_apply_exp2<false>(scores, lse, params.scale_softmax_log2);

    // 3. Compute dP = dO @ V^T
    flash::gemm(acc_dp, tdPrdO, tdPrV, ...);

    // 4. Compute dS = P * (dP - dP_sum) where dP_sum = rowsum(dO * O)
    for (int mi ...) {
        for (int ni ...) {
            dS(mi, ni) = scores(mi, ni) * (dS(mi, ni) - dP_sum(mi));
        }
    }

    // 5. Accumulate dV += P^T @ dO
    flash::gemm(acc_dv, tdVrPt, tdVrdO, ...);

    // 6. Accumulate dK += dS^T @ Q
    flash::gemm(acc_dk, tdKrdSt, tdKrQt, ...);

    // 7. Compute dQ = dS @ K (accumulated to dQ_accum, possibly with atomics)
    flash::gemm(acc_dq, tdQrdS, tdQrKt, ...);
}

The backward uses three different TiledMMA configurations:

• TiledMmaSdP: For computing S and dP (warps split AtomLayoutMSdP x (kNWarps/AtomLayoutMSdP))
• TiledMmadKV: For computing dK and dV (warps split AtomLayoutNdKV x (kNWarps/AtomLayoutNdKV))
• TiledMmadQ: For computing dQ (warps split AtomLayoutMdQ x (kNWarps/AtomLayoutMdQ))

For dQ accumulation, when running with multiple thread blocks per K/V column (sequence parallel mode), the code uses atomicAdd:

if (!Seq_parallel) {
    cute::copy(gmem_tiled_copy_dQaccum, acc_dq_reshaped, tdQgdQaccum);
} else {
    for (int i = 0; i < size(acc_dq); ++i) { atomicAdd(&tdQgdQaccum(i), acc_dq(i)); }
}

There’s also a deterministic mode where each thread block writes to a separate dQ accumulator buffer (indexed by blockIdx.x * dq_accum_split_stride), and a final reduction combines them. This avoids atomics and guarantees bitwise reproducibility.

Layer 8: Multi-Query / Grouped-Query Attention (MQA/GQA)

GQA support is surprisingly simple. The h_h_k_ratio parameter (h / h_k) maps multiple Q heads to the same K/V head:

Tensor gK = local_tile(mK(_, bidh / params.h_h_k_ratio, _), ...);
Tensor gV = local_tile(mV(_, bidh / params.h_h_k_ratio, _), ...);

The grid still launches one thread block per (Q head, batch), but multiple Q heads read from the same K/V head.

Summary

Aspect	FlashAttention 1	FlashAttention 2
Outer loop	K/V blocks	Q blocks (one per thread block)
Warp partitioning	Split across Q and K dims	Split only across Q rows
Cross-warp communication	Shared memory reduction	None (each warp has full K column)
Output normalization	After each K block	Once at the end
Non-matmul FLOPs	$O(N^2 d)$ rescaling	$O(N^2 + Nd)$ rescaling
Backward dQ accumulation	Atomic adds	Atomic or deterministic split-accumulate
Typical speedup	-	~2x over FA1

The code’s extensive use of NVIDIA’s CuTe library (from CUTLASS) for tensor layout algebra and Tensor Core MMA instructions keeps the kernel portable across head dimensions and hardware generations, while the swizzled shared memory layouts eliminate bank conflicts that would otherwise degrade throughput.