From Softmax Attention to FlashAttention-2

Introduction

This post derives scaled dot-product attention, turns stable softmax into an online recurrence, and implements the same one-pass forward algorithm in Triton and CuTe DSL without materializing the full attention matrix.

The complete executable lesson is available in Day 12 of the CuTe DSL curriculum. The standalone implementation is in solution.py.

The goal is not to reproduce every production FlashAttention-2 optimization. Instead, the focus is the algorithmic core:

Q-row ownership
    -> stream K/V tiles
    -> update online (m, l, U) state
    -> normalize once
    -> never write the N x N score matrix to HBM

1. What softmax attention computes

For one batch-head, let:

• Q have shape (N_q, D);
• K and V have shape (N_k, D);
• D be the head dimension.

Scaled dot-product attention is:

\[S = \frac{QK^T}{\sqrt{D}}, \qquad P_{ij} = \frac{\exp(S_{ij})} {\sum_{r=0}^{N_k-1}\exp(S_{ir})}, \qquad O = PV.\]

Each query row compares with every key row. Softmax converts the scores into positive weights that sum to one, and the output is a weighted combination of the value rows.

Stable softmax

The numerically stable form subtracts each score row’s maximum:

\[m_i = \max_j S_{ij}, \qquad P_{ij} = \frac{\exp(S_{ij}-m_i)} {\sum_r \exp(S_{ir}-m_i)}.\]

Subtracting the maximum does not change the result because the same factor $e^{-m_i}$ appears in every numerator and in the denominator. It prevents large positive scores from overflowing exp.

2. Why ordinary attention is memory-hungry

A direct implementation commonly creates:

scores        = Q @ K.T       # N_q x N_k
probabilities = softmax(...)  # N_q x N_k
output        = probabilities @ V

When N_q = N_k = N, the score and probability matrices contain N^2 elements even though the final output contains only N * D. Writing and rereading those quadratic intermediates from HBM is the main problem FlashAttention removes.

FlashAttention keeps only a small score tile on chip:

one Q row/block
    for each K/V tile:
        score tile = Q tile @ K tile.T
        update online softmax state
        update output accumulator with V tile
    normalize once and store O

The complete N x N matrix is never stored in HBM.

3. The one-pass softmax attention equation

For one query row, maintain three values after processing some K/V tiles:

\[m = \max_{\text{processed }j} s_j,\] \[\ell = \sum_{\text{processed }j} e^{s_j-m},\] \[U = \sum_{\text{processed }j} e^{s_j-m}v_j.\]

U is deliberately unnormalized. The current attention output would be U / l, but FlashAttention-2 delays that division until all tiles have been processed.

For a new score tile $S_t$ and matching value tile $V_t$, compute:

\[\widetilde{m}_t = \operatorname{rowmax}(S_t), \qquad m' = \max(m, \widetilde{m}_t),\] \[\alpha = e^{m-m'}, \qquad P_t = e^{S_t-m'},\] \[\ell' = \alpha\ell + \operatorname{rowsum}(P_t),\] \[U' = \alpha U + P_tV_t.\]

After the final tile:

\[O = \frac{U}{\ell}.\]

These updates preserve the invariant:

\[\ell' = \sum_{\text{all processed }j}e^{s_j-m'}, \qquad U' = \sum_{\text{all processed }j}e^{s_j-m'}v_j.\]

The factor alpha rescales every old term from the old maximum to the new maximum. This is why the algorithm remains stable even when a later tile contains a much larger score.

Verify the recurrence before writing a kernel

import math
import torch


def online_attention_reference(q, k, v, key_tile=16):
    scale = 1.0 / math.sqrt(q.shape[-1])
    running_max = torch.tensor(
        -torch.inf, device=q.device, dtype=torch.float32
    )
    running_sum = torch.tensor(
        0.0, device=q.device, dtype=torch.float32
    )
    output_accumulator = torch.zeros_like(q)

    for start in range(0, k.shape[0], key_tile):
        k_tile = k[start : start + key_tile]
        v_tile = v[start : start + key_tile]
        scores = (k_tile @ q) * scale

        new_max = torch.maximum(running_max, scores.max())
        previous_scale = torch.exp(running_max - new_max)
        probabilities = torch.exp(scores - new_max)

        running_sum = (
            running_sum * previous_scale + probabilities.sum()
        )
        output_accumulator = (
            output_accumulator * previous_scale
            + probabilities @ v_tile
        )
        running_max = new_max

    return output_accumulator / running_sum

For any key-tile size, this produces the same result as:

direct = torch.softmax(
    k @ q / math.sqrt(q.shape[-1]), dim=0
) @ v

4. What FlashAttention-2 changes

FlashAttention-1 already introduced tiling and online softmax. FlashAttention-2 improves the work assignment and removes unnecessary non-matmul operations.

4.1 Delay output normalization

The normalized FA1 recurrence can be written as:

\[O' = \frac{\alpha\ell}{\ell'}O + \frac{P_tV_t}{\ell'}.\]

FA2 instead accumulates the unnormalized U:

\[U' = \alpha U + P_tV_t,\]

and performs O = U / l once at the end. This avoids repeatedly dividing and rescaling every output element.

4.2 Assign each Q block to one CTA

The outer parallel dimension is the Q sequence:

CTA 0 owns Q rows   0..Br-1
CTA 1 owns Q rows  Br..2Br-1
...

Each CTA loops over all K/V tiles and exclusively owns its output tile. The forward pass therefore needs no atomic accumulation.

4.3 Partition warps only along Q rows

FA2 stacks warps along Q rows rather than splitting a softmax row across the K-column dimension:

CTA query block
├── warp 0 -> complete softmax rows
├── warp 1 -> complete softmax rows
├── warp 2 -> complete softmax rows
└── warp 3 -> complete softmax rows

Each warp sees the complete K-column extent for its Q rows, avoiding cross-warp communication for the softmax reduction.

Production kernels additionally use tensor-core MMA, swizzled shared memory, asynchronous K/V pipelines, FP16/BF16 inputs, causal masks, dropout, and a specialized backward pass. The implementations below isolate the forward algorithm first.

5. Correctness and benchmark contract

Both kernels use contiguous tensors with shape:

[batch_heads, sequence_length, head_dim]

The Triton path uses FP16 Q/K/V/output with FP32 online-softmax and output accumulation. The CuTe path keeps Q/K/V/output in FP32 memory, but converts matrix-operand fragments to TF32 and executes both QK^T and PV with tensor-core TiledMma.

The direct reference materializes the score matrix only for testing. It computes in FP32 and then converts to the input dtype:

def attention_reference(q, k, v):
    scale = 1.0 / math.sqrt(q.shape[-1])
    scores = (
        torch.matmul(q.float(), k.float().transpose(-1, -2))
        * scale
    )
    output = torch.matmul(
        torch.softmax(scores, dim=-1), v.float()
    )
    return output.to(q.dtype)

Correctness tests use:

TEST_SHAPES = (
    (2, 17, 32),
    (3, 33, 64),
    (4, 64, 64),
)

Lengths 17 and 33 deliberately exercise partial key tiles and partially populated Q blocks.

Each benchmark:

• checks correctness before timing;
• reuses compiled callables;
• uses triton.testing.do_bench;
• reports the two attention matmuls:

\[\text{FLOPs} \approx 4 B_h N^2 D.\]

The count omits exponentials, reductions, and rescaling, so TFLOP/s is an attention convention rather than a complete instruction count.

6. Implementations

Both implementations apply the same (m, l, U) invariant to matrix tiles. They differ in storage dtype, tile shape, and how tensor-core ownership is expressed.

6.1 Triton implementation

This implementation follows the forward path in Triton Tutorial 06: Fused Attention, with a deliberately narrow feature set:

• FlashAttention-2 forward only;
• non-causal attention;
• FP16 Q/K/V and FP16 output;
• FP32 online-softmax and output accumulation;
• four warps, with no warp specialization;
• no FP8, tensor descriptors, autotuning, or backward kernel.

One program owns a block of 64 Q rows. It keeps that Q tile resident and processes K/V blocks of 32 or 64 rows:

scores = Q_tile @ K_tile.T / sqrt(D)
m_new = max(m, rowmax(scores))
alpha = exp2(m - m_new)
p = exp2(scores - m_new[:, None])
l = alpha * l + rowsum(p)
U = alpha[:, None] * U + p.to(fp16) @ V_tile
m = m_new

tl.dot lowers both matrix products to tensor-core operations. U is a BLOCK_M x D FP32 accumulator. The BLOCK_M x BLOCK_N score tile stays on chip, so the kernel never creates an N x N output.

The kernel uses base-2 exponentiation like the official tutorial. Multiplying the softmax scale by log2(e) preserves the same mathematical softmax while mapping to exp2.

import triton
import triton.language as tl


@triton.jit
def triton_flash_attention2_kernel(
    q_ptr,
    k_ptr,
    v_ptr,
    o_ptr,
    softmax_scale,
    N_CTX: tl.constexpr,
    HEAD_DIM: tl.constexpr,
    BLOCK_M: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    query_block = tl.program_id(0)
    batch_head = tl.program_id(1)
    query_indices = (
        query_block * BLOCK_M + tl.arange(0, BLOCK_M)
    )
    key_offsets = tl.arange(0, BLOCK_N)
    dimensions = tl.arange(0, HEAD_DIM)
    query_mask = query_indices < N_CTX
    tensor_base = batch_head * N_CTX * HEAD_DIM
    q_offsets = (
        tensor_base
        + query_indices[:, None] * HEAD_DIM
        + dimensions[None, :]
    )
    q = tl.load(
        q_ptr + q_offsets,
        mask=query_mask[:, None],
        other=0.0,
    )

    running_max = (
        tl.zeros((BLOCK_M,), dtype=tl.float32)
        - float("inf")
    )
    running_sum = (
        tl.zeros((BLOCK_M,), dtype=tl.float32) + 1.0
    )
    output_accumulator = tl.zeros(
        (BLOCK_M, HEAD_DIM), dtype=tl.float32
    )
    qk_scale = softmax_scale * 1.44269504

    for key_start in tl.range(0, N_CTX, BLOCK_N):
        key_indices = key_start + key_offsets
        key_mask = key_indices < N_CTX
        kv_offsets = (
            tensor_base
            + key_indices[:, None] * HEAD_DIM
            + dimensions[None, :]
        )

        keys = tl.load(
            k_ptr + kv_offsets,
            mask=key_mask[:, None],
            other=0.0,
        )
        scores = tl.dot(q, tl.trans(keys))
        scores = tl.where(
            key_mask[None, :],
            scores * qk_scale,
            -1.0e6,
        )

        tile_max = tl.max(scores, axis=1)
        new_max = tl.maximum(running_max, tile_max)
        probabilities = tl.math.exp2(
            scores - new_max[:, None]
        )
        previous_scale = tl.math.exp2(
            running_max - new_max
        )
        tile_sum = tl.sum(probabilities, axis=1)

        values = tl.load(
            v_ptr + kv_offsets,
            mask=key_mask[:, None],
            other=0.0,
        )
        output_accumulator *= previous_scale[:, None]
        output_accumulator = tl.dot(
            probabilities.to(tl.float16),
            values,
            output_accumulator,
        )
        running_sum = (
            running_sum * previous_scale + tile_sum
        )
        running_max = new_max

    tl.store(
        o_ptr + q_offsets,
        (output_accumulator / running_sum[:, None]).to(
            tl.float16
        ),
        mask=query_mask[:, None],
    )

The launcher specializes sequence length, head dimension, and tile sizes:

def launch_triton_attention(q, k, v, output):
    batch_heads, sequence_length, head_dim = q.shape
    block_m = 64
    block_n = min(64, head_dim)
    grid = (
        triton.cdiv(sequence_length, block_m),
        batch_heads,
    )

    triton_flash_attention2_kernel[grid](
        q,
        k,
        v,
        output,
        1.0 / math.sqrt(head_dim),
        N_CTX=sequence_length,
        HEAD_DIM=head_dim,
        BLOCK_M=block_m,
        BLOCK_N=block_n,
        num_warps=4,
        num_stages=3,
    )

Triton correctness

def test_triton_attention():
    for index, shape in enumerate(TEST_SHAPES):
        torch.manual_seed(20261015 + index)
        q = torch.randn(
            shape, device="cuda", dtype=torch.float16
        )
        k = torch.randn_like(q)
        v = torch.randn_like(q)

        output = torch.empty_like(q)
        launch_triton_attention(q, k, v, output)
        torch.cuda.synchronize()

        torch.testing.assert_close(
            output,
            attention_reference(q, k, v),
            atol=2e-2,
            rtol=2e-2,
        )

The Triton kernel passes all three test shapes.

Triton benchmark

Measured on an NVIDIA B300:

(batch_heads, N, D)	Triton ms	Triton TFLOP/s
(16, 256, 64)	0.0102	26.23
(16, 512, 64)	0.0123	87.33

This is still narrower than the production tutorial—there is no causal path, autotuning, descriptors, FP8, backward, or warp specialization—but it now has the defining blockwise tensor-core FA2 forward dataflow.

6.2 CuTe DSL implementation

The optimized CuTe lesson keeps FP32 tensors at the API boundary but replaces the scalar dot products with TF32 tensor-core TiledMma. One warp owns a 16-row Q tile and streams 16-key K/V tiles:

one-warp CTA
├── Q tile:     16 x D, loaded once
├── K/V tile:  16 x D, replaced every iteration
├── scores:    16 x 16 = TF32 TiledMMA(Q, K^T)
├── softmax:   two 16-lane row reductions at a time
└── tile U:    16 x D  = TF32 TiledMMA(P, V)

The MMA atom is the Ampere-and-later warp instruction:

mma_op = cute.nvgpu.warp.MmaTF32Op((16, 8, 8))
tiled_mma = cute.make_tiled_mma(mma_op)

It wraps:

mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32

Q/K/V remain FP32 in global and shared memory. make_fragment_A/B creates TF32 operand fragments for this MMA atom; copying the FP32 partitions into those fragments performs the TF32 conversion. Accumulators remain FP32.

The core dataflow is:

WARP_SIZE = 32
QUERY_TILE = 16
KEY_TILE = 16


def half_warp_reduce_sum(value):
    for offset in (8, 4, 2, 1):
        value += cute.arch.shuffle_sync_bfly(
            value,
            offset=offset,
            mask=-1,
            mask_and_clamp=31,
        )
    return value


def half_warp_reduce_max(value):
    for offset in (8, 4, 2, 1):
        value = cute.arch.fmax(
            value,
            cute.arch.shuffle_sync_bfly(
                value,
                offset=offset,
                mask=-1,
                mask_and_clamp=31,
            ),
        )
    return value


@cute.jit
def __call__(self, q, k, v, output):
    mma_op = cute.nvgpu.warp.MmaTF32Op((16, 8, 8))
    tiled_mma = cute.make_tiled_mma(mma_op)
    self.kernel(q, k, v, output, tiled_mma).launch(
        grid=(
            cute.ceil_div(self.sequence_length, QUERY_TILE),
            self.batch_heads,
            1,
        ),
        block=(WARP_SIZE, 1, 1),
    )


@cute.kernel
def kernel(self, q, k, v, output, tiled_mma):
    lane, _, _ = cute.arch.thread_idx()
    query_block, batch_head, _ = cute.arch.block_idx()

    # Shared tiles use FP32 storage.
    s_q = allocate_fp32((QUERY_TILE, self.head_dim))
    s_k = allocate_fp32((KEY_TILE, self.head_dim))
    s_scores = allocate_fp32((QUERY_TILE, KEY_TILE))
    s_probabilities = allocate_fp32((QUERY_TILE, KEY_TILE))
    s_v = allocate_fp32((self.head_dim, KEY_TILE))
    s_tile_output = allocate_fp32((QUERY_TILE, self.head_dim))
    s_output = allocate_fp32((QUERY_TILE, self.head_dim))

    thr_mma = tiled_mma.get_slice(lane)

    # QK^T partitions and TF32 operand fragments.
    t_q = thr_mma.partition_A(s_q)
    t_k = thr_mma.partition_B(s_k)
    t_scores = thr_mma.partition_C(s_scores)
    r_q = tiled_mma.make_fragment_A(t_q)
    r_k = tiled_mma.make_fragment_B(t_k)
    r_scores = tiled_mma.make_fragment_C(t_scores)

    # PV partitions and TF32 operand fragments.
    t_p = thr_mma.partition_A(s_probabilities)
    t_v = thr_mma.partition_B(s_v)
    t_tile_output = thr_mma.partition_C(s_tile_output)
    r_p = tiled_mma.make_fragment_A(t_p)
    r_v = tiled_mma.make_fragment_B(t_v)
    r_tile_output = tiled_mma.make_fragment_C(t_tile_output)

    load_q_once(s_q, q, query_block, batch_head)

    for key_start in range(0, self.sequence_length, KEY_TILE):
        load_masked_kv_tile(s_k, s_v, k, v, key_start)

        r_scores.fill(0.0)
        for k_block in range(score_k_blocks):
            cute.autovec_copy(t_q[..., k_block], r_q[..., k_block])
            cute.autovec_copy(t_k[..., k_block], r_k[..., k_block])
            cute.gemm(
                tiled_mma,
                r_scores,
                r_q[..., k_block],
                r_k[..., k_block],
                r_scores,
            )
        store_fragment(r_scores, s_scores)

        # Two 16-lane groups reduce two score rows simultaneously.
        tile_max = half_warp_reduce_max(score)
        new_max = cute.arch.fmax(running_max, tile_max)
        alpha = cute.math.exp(running_max - new_max)
        probability = cute.math.exp(score - new_max)
        tile_sum = half_warp_reduce_sum(probability)
        running_sum = alpha * running_sum + tile_sum

        r_tile_output.fill(0.0)
        for k_block in range(output_k_blocks):
            cute.autovec_copy(t_p[..., k_block], r_p[..., k_block])
            cute.autovec_copy(t_v[..., k_block], r_v[..., k_block])
            cute.gemm(
                tiled_mma,
                r_tile_output,
                r_p[..., k_block],
                r_v[..., k_block],
                r_tile_output,
            )
        store_fragment(r_tile_output, s_tile_output)
        s_output = alpha * s_output + s_tile_output

    output = s_output / running_sum

The executable notebook spells out allocate_fp32, masked staging, fragment stores, and per-row state indexing. The important precision contract is:

global/shared Q, K, V, P, O: FP32
MMA A/B fragments:           TF32
MMA and softmax accumulators: FP32

After QK^T, each half warp reduces one score row across the 16-key tile:

\[m' = \max(m, \operatorname{rowmax}(S_t)),\] \[\ell' = e^{m-m'}\ell + \operatorname{rowsum}(e^{S_t-m'}),\] \[U' = e^{m-m'}U + \operatorname{TF32MMA}(e^{S_t-m'}, V_t).\]

Partial Q and K/V tiles are masked, and the output is normalized once after the final tile.

Previous scalar FP32 baseline

The first version assigned one Q row to each warp, processed one key per recurrence update, and used lane FMAs plus a warp shuffle: ```python import cutlass import cutlass.cute as cute from cutlass.cute.runtime import from_dlpack WARP_SIZE = 32 WARPS_PER_CTA = 4 def warp_reduce_sum(value: cutlass.Float32): for offset in (16, 8, 4, 2, 1): value += cute.arch.shuffle_sync_bfly( value, offset=offset, mask=-1, mask_and_clamp=31, ) return value class SimplifiedFlashAttention2: def __init__( self, batch_heads, sequence_length, head_dim, ): self.batch_heads = batch_heads self.sequence_length = sequence_length self.head_dim = head_dim self.values_per_lane = ( head_dim + WARP_SIZE - 1 ) // WARP_SIZE self.softmax_scale = 1.0 / math.sqrt(head_dim) @cute.jit def __call__( self, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, output: cute.Tensor, ): self.kernel(q, k, v, output).launch( grid=( cute.ceil_div( self.sequence_length, WARPS_PER_CTA, ), self.batch_heads, 1, ), block=( WARP_SIZE * WARPS_PER_CTA, 1, 1, ), ) @cute.kernel def kernel( self, q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, output: cute.Tensor, ): thread_idx, _, _ = cute.arch.thread_idx() query_block, batch_head, _ = ( cute.arch.block_idx() ) warp_idx = thread_idx // WARP_SIZE lane = thread_idx % WARP_SIZE query_idx = ( query_block * WARPS_PER_CTA + warp_idx ) q_values = cute.make_rmem_tensor( (self.values_per_lane,), cutlass.Float32, ) output_accumulator = cute.make_rmem_tensor( (self.values_per_lane,), cutlass.Float32, ) q_values.fill(0.0) output_accumulator.fill(0.0) if query_idx < self.sequence_length: for item in cutlass.range_constexpr( self.values_per_lane ): dimension = lane + item * WARP_SIZE if dimension < self.head_dim: q_values[item] = q[ batch_head, query_idx, dimension, ] running_max = -cutlass.Float32.inf running_sum = cutlass.Float32.zero for key_idx in cutlass.range( self.sequence_length ): partial_score = cutlass.Float32.zero for item in cutlass.range_constexpr( self.values_per_lane ): dimension = lane + item * WARP_SIZE if dimension < self.head_dim: partial_score += ( q_values[item] * k[ batch_head, key_idx, dimension, ] ) score = ( warp_reduce_sum(partial_score) * self.softmax_scale ) new_max = cute.arch.fmax( running_max, score ) previous_scale = cute.math.exp( running_max - new_max ) probability = cute.math.exp( score - new_max ) running_sum = ( running_sum * previous_scale + probability ) for item in cutlass.range_constexpr( self.values_per_lane ): dimension = lane + item * WARP_SIZE if dimension < self.head_dim: output_accumulator[item] = ( output_accumulator[item] * previous_scale + probability * v[ batch_head, key_idx, dimension, ] ) running_max = new_max for item in cutlass.range_constexpr( self.values_per_lane ): dimension = lane + item * WARP_SIZE if dimension < self.head_dim: output[ batch_head, query_idx, dimension, ] = ( output_accumulator[item] / running_sum ) ```

The current TF32 TiledMMA host wrapper still accepts ordinary FP32 PyTorch tensors through DLPack and compiles a kernel specialized for the input shape: specialized for the input shape:

def prepare_cute_attention(q, k, v):
    output = torch.empty_like(q)
    q_cute = from_dlpack(q, assumed_align=16)
    k_cute = from_dlpack(k, assumed_align=16)
    v_cute = from_dlpack(v, assumed_align=16)
    output_cute = from_dlpack(
        output, assumed_align=16
    )

    compiled = cute.compile(
        SimplifiedFlashAttention2(*q.shape),
        q_cute,
        k_cute,
        v_cute,
        output_cute,
    )
    return output, compiled, (
        q_cute,
        k_cute,
        v_cute,
        output_cute,
    )

CuTe correctness

def test_cute_attention():
    for index, shape in enumerate(TEST_SHAPES):
        torch.manual_seed(20261015 + index)
        q = torch.randn(shape, device="cuda")
        k = torch.randn_like(q)
        v = torch.randn_like(q)

        output, compiled, arguments = (
            prepare_cute_attention(q, k, v)
        )
        compiled(*arguments)
        torch.cuda.synchronize()

        torch.testing.assert_close(
            output,
            attention_reference(q, k, v),
            atol=3e-2,
            rtol=3e-2,
        )

The CuTe kernel passes the same sequence/head shapes against the FP32 reference. The tolerance accounts for TF32 operand rounding in both matrix products.

CuTe benchmark

Measured on the same NVIDIA B300:

(batch_heads, N, D)	CuTe ms	CuTe TFLOP/s
(16, 256, 64)	0.0716	3.75
(16, 512, 64)	0.1370	7.84

The previous scalar FP32 CuTe path reached about 2.86 and 4.29 TFLOP/s. TF32 TiledMMA improves these shapes by approximately 1.31x and 1.83x while preserving FP32 storage.

7. Comparing the implementations

Property	Triton	CuTe DSL
Q ownership	One program per 64-row Q block	One warp per 16-row Q tile
K/V update size	32 or 64 keys	16 keys
Dot product	FP16 tensor-core tl.dot	TF32 tensor-core TiledMma
Inputs/output	FP16 / FP16	FP32 / FP32
Output state	BLOCK_M x D FP32 registers	16 x D FP32 shared tile
Output normalization	Once after K/V loop	Once after K/V loop
Score matrix in HBM	Never	Never
Tensor cores	Yes	Yes
Shared K/V staging	No	Yes

Both implementations remove the defining quadratic intermediate and implement the same stable online recurrence, but they now target different lesson stages.

The main difference is how ownership is expressed:

• Triton uses matrix tiles and tensor-core lowering for a focused FP16 FA2 forward implementation.
• CuTe explicitly maps 16-row/16-key tiles to TF32 MMA fragments, then exposes the shared-memory score epilogue and half-warp softmax reductions.

8. Next optimization steps

The next version should preserve the exact (m, l, U) equations while replacing one mechanism at a time:

1. Keep the persistent CuTe output tile in register fragments instead of shared memory.
2. Assign multiple warps to a larger Q tile while reusing each staged K/V tile.
3. Double-buffer K/V and overlap global-to-shared copies with current-tile compute.
4. Fuse score-fragment epilogue work more tightly with row reductions.
5. Add causal masks and broader head-dimension support.
6. Store logsumexp and implement backward recomputation.
7. Replace warp MMA with TMA plus tcgen05.mma/TMEM on B300.

Production FlashAttention-2 combines all of these with architecture-specific tile sizes, register-pressure control, and carefully overlapped memory movement. The one-pass recurrence stays the same.