NVIDIA CUDA TileにおけるFlash Attentionのピークパフォーマンス調整

In this post, we dive into one of the most critical workloads in modern AI: Flash Attention, where you’ll learn:

How to implement Flash Attention using NVIDIA cuTile. Walk through the complete code for a production-ready implementation.

The “trap and rescue” optimization journey. This case study shows how naive optimizations (like just increasing tile size) can backfire, and how to fix them.

Advanced techniques like FMA patterns, fast math, loop splitting, and adaptive tiling for maximum performance.

Environment requirements:

CUDA 13.1 or higher

GPU architecture: Compute capability 8.X, 10.X, 11.X, 12.X (NVIDIA Ampere, NVIDIA Ada, NVIDIA Blackwell)

Python: 3.10 or higher

See the quickstart doc for more information on installing cuTile Python.

What is attention?

The attention mechanism is the computational heart of transformer models. Given a sequence of tokens, attention enables each token to “look at” every other token and decide how much to weigh their contributions. Mathematically, for input matrices Query (), Key (), and Value (), the output is:

Where:

The memory bandwidth problem

For a sequence length of (common in modern LLMs), the attention matrix contains million elements. In FP16, that’s 512 MB of intermediate storage per attention head, per batch item.

Standard attention implementations:

Compute the full attention matrix and write it to global memory (slow)

Apply softmax row-by-row

Read the matrix back and multiply by

This approach is memory-bound as the GPU spends most of its time waiting for data to move between HBM and compute units, rather than computing.

How Flash Attention solves the memory bandwidth problem

Flash Attention (introduced by Dao et al., 2022) is an IO-aware algorithm that never materializes the full matrix. Instead, it:

Tiles the computation: Processes in small blocks that fit in fast on-chip SMEM

Uses online softmax: Computes softmax incrementally without needing the full row

Fuses operations: Combines the matrix multiply and softmax into a single kernel pass

The result is a 2-4x speedup and significant memory savings, enabling longer context lengths.

Figure 1. Tiled Flash Attention computation

Understanding online softmax

The key algorithmic insight of Flash Attention is the online softmax trick. The numerically stable safe softmax requires knowing the maximum value across the entire row before computing:

But if we’re processing tiles, we don’t have access to the full row. Online softmax solves this by maintaining running statistics that can be updated incrementally.

The online softmax algorithm

We maintain two running values for each row:

: The maximum value seen so far (for numerical stability)

: The sum of exponentials seen so far (the softmax denominator)

When we process a new tile with values :

Update the maximum:

Compute correction factor: (rescales previous work)

Update the sum:

Update the accumulator:

is the matrix of the attention weights, and is the value matrix tile, corresponding to the Key tile of the current iteration. At the end, we normalize:

This enables us to compute an exact softmax without ever storing the full row.

Causal attention and grouped-query attention

Before diving into the implementation, let’s understand two important attention variants used in modern LLMs:

Causal attention

In autoregressive language models like GPT, LLaMA, and Claude, each token can only attend to previous tokens in the sequence, not future ones. This prevents “cheating” during training, where the model looks ahead to predict the next word.

Mathematically, we apply a triangular mask to the attention scores:

The masked attention becomes:

Adding to future positions ensures they become zero after softmax, effectively blocking information flow from future tokens.

Figure 2. Causal attention mask for four tokens

With causal masking, roughly half the attention matrix is masked (the upper triangle). We can skip computing these masked tiles entirely, providing a 2x algorithmic speedup. This is crucial for the K-loop splitting optimization.

Grouped-query attention

Standard multi-head attention has separate matrices for each attention head, leading to high memory usage:

Multi-head attention (MHA): 32 query heads → 32 K/V heads (1:1 ratio)

Grouped-query attention (GQA): 32 query heads → 4 K/V heads (8:1 ratio)

Multi-query attention (MQA): 32 query heads → 1 K/V head (32:1 ratio)

In GQA, multiple query heads share the same K/V heads. For example, with 32 query heads and 4 K/V heads:

Query heads 0-7 use K/V head 0

Query heads 8-15 use K/V head 1

Query heads 16-23 use K/V head 2

Query heads 24-31 use K/V head 3

This reduces K/V cache size by 8x during inference, critical for serving long-context models. Modern LLMs like LlamA 2, Llama 3, Mistral, and Qwen use GQA extensively.

When implementing in Flash Attention, each CUDA block computes attention for one query head, but loads the appropriate shared K/V head:

head_idx = bid_y % num_heads              kv_head_idx = head_idx // query_group_size

With a query group size of 8, query heads 0-7 all map to kv_head_idx = 0, sharing the same K/V tiles in memory.

Part 1: The flash attention kernel in CUDA Tile

Let’s implement Flash Attention step-by-step. Our baseline uses small 64×64 tiles and straightforward code—correct but not yet optimized.

1. Defining the kernel interface

In cuTile, the @ct.kernel decorator marks a Python function as a GPU kernel. We pass compile-time constants using ct.Constant[T] type annotations:

import mathimport cuda.tile as ctConstInt = ct.Constant[int]ConstBool = ct.Constant[bool]INV_LOG_2 = 1.0 / math.log(2)@ct.kernel()def fmha_kernel(    Q, K, V, Out,                  qk_scale: float,               input_pos: int,                TILE_D: ConstInt,              H: ConstInt,                   TILE_M: ConstInt,              TILE_N: ConstInt,              QUERY_GROUP_SIZE: ConstInt,    CAUSAL: ConstBool,             EVEN_K: ConstBool,         ):

1. Block ID mapping

Each CUDA block computes one tile of the output. Using ct.bid , we map the 2D grid to batch/head indices:

    bid_x = ct.bid(0)      bid_y = ct.bid(1)      batch_idx = bid_y // H    head_idx = bid_y % H    off_kv_h = head_idx // QUERY_GROUP_SIZE

1. Initializing accumulators

Before the main loop, we initialize the online softmax state and output accumulator:

    qk_scale = qk_scale * INV_LOG_2    offs_m = bid_x * TILE_M + ct.arange(TILE_M, dtype=ct.int32)    offs_m += input_pos    offs_m = offs_m[:, None]      offs_n_tile = ct.arange(TILE_N, dtype=ct.int32)    offs_n_tile = offs_n_tile[None, :]      m_i = ct.full((TILE_M, 1), -math.inf, dtype=ct.float32)      l_i = ct.full((TILE_M, 1), 0.0, dtype=ct.float32)            acc = ct.full((TILE_M, TILE_D), 0.0, dtype=ct.float32)

We use float32 for accumulators, even when inputs are float16 to maintain numerical precision during the iterative softmax computation.

1. Loading the query tile

The query tile is loaded once and reused across all K/V iterations:

q = ct.load(    Q,     index=(batch_idx, head_idx, bid_x, 0),     shape=(1, 1, TILE_M, TILE_D)).reshape((TILE_M, TILE_D))

The ct.load function handles boundary conditions automatically when the tile extends past the tensor edge.

1. The main loop over K/V tiles

This is the heart of Flash Attention. We iterate over K/V tiles:

 m_end = input_pos + (bid_x + 1) * TILE_M k_seqlen = K.shape[2] if CAUSAL:     Tc = ct.cdiv(min(m_end, k_seqlen), TILE_N) else:     Tc = ct.cdiv(k_seqlen, TILE_N) for j in range(0, Tc):     k = ct.load(         K,         index=(batch_idx, off_kv_h, 0, j),         shape=(1, 1, TILE_D, TILE_N),         order=(0, 1, 3, 2),           latency=2                 ).reshape((TILE_D, TILE_N))     qk = ct.full((TILE_M, TILE_N), 0.0, dtype=ct.float32)     qk = ct.mma(q, k, qk)

The order=(0,1,3,2) in the parameter tells cuTile load operation to use K transposed, and latency=2 hints that we can tolerate some latency (enabling better pipelining). Then we use the ct.mma=(q, k, k,qk) to perform the cuTile matrix multiply-accumulate.

1. Applying the causal mask

For autoregressive models (GPT, Llama, etc.), each token can only attend to previous tokens:

        if CAUSAL or not EVEN_K:            offs_n = j * TILE_N + offs_n_tile            mask = ct.full((TILE_M, TILE_N), True, dtype=ct.bool_)            if not EVEN_K:                mask = mask & (offs_n < k_seqlen)            if CAUSAL:                mask = mask & (offs_m >= offs_n)            mask = ct.where(mask, 0.0, -math.inf)            qk += mask

Adding -inf to masked positions ensures they become zero after softmax.

1. Online softmax update

Now we update our running softmax statistics:

     qk_max = ct.max(qk, axis=-1, keepdims=True)     qk_max_scaled = qk_max * qk_scale     m_ij = max(m_i, qk_max_scaled)     qk = qk * qk_scale     qk = qk - m_ij     p = ct.exp2(qk)     l_ij = ct.sum(p, axis=-1, keepdims=True)     alpha = ct.exp2(m_i - m_ij)       l_i = l_i * alpha     l_i = l_i + l_ij     acc = acc * alpha

1. Accumulating the output

Finally, we load the Value tile and accumulate:

        v = ct.load(            V,            index=(batch_idx, off_kv_h, j, 0),            shape=(1, 1, TILE_N, TILE_D),            latency=4        ).reshape((TILE_N, TILE_D))        p = p.astype(Q.dtype)        acc = ct.mma(p, v, acc)        m_i = m_ij

1. Final normalization and store

After processing all tiles, we normalize by the total sum and write the result:

 acc = ct.truediv(acc, l_i) acc = acc.reshape((1, 1, TILE_M, TILE_D)).astype(Out.dtype) ct.store(Out, index=(batch_idx, head_idx, bid_x, 0), tile=acc)

Launching the kernel: Host-side code

Now let’s look at the host-side code that launches the kernel:

import torchfrom math import ceildef tile_fmha(q, k, v, sm_scale=None, is_causal=True):    if sm_scale is None:        sm_scale = 1.0 / math.sqrt(q.size(-1))    batch_size, num_heads, seq_len, head_dim = q.shape    _, num_kv_heads, _, _ = k.shape    query_group_size = num_heads // num_kv_heads    q = q.contiguous()    k = k.contiguous()    v = v.contiguous()    o = torch.empty_like(q)    TILE_M, TILE_N = 64, 64    grid_x = ceil(seq_len / TILE_M)      grid_y = batch_size * num_heads      grid = (grid_x, grid_y, 1)    EVEN_K = (k.shape[2] % TILE_N) == 0    ct.launch(        torch.cuda.current_stream(),        grid,        fmha_kernel,        (q, k, v, o, sm_scale, 0, head_dim, num_heads,         TILE_M, TILE_N, query_group_size, is_causal, EVEN_K)    )    return o

This baseline with 64×64 tiles works correctly. But can we make it faster? Let’s find out.

Part 2: The “trap and rescue” optimization journey

We benchmark on the following configuration:

Hardware: NVIDIA B200

Batch: 4, Heads: 32, Head dimension: 128

Attention: Causal, Dtype: FP16

Sequence lengths: 1024, 2048, 4096, 8192, 16384

To interpret each step, we use Nsight Compute with a minimal section set:

LaunchStats

Occupancy

SpeedOfLight

ComputeWorkloadAnalysis

MemoryWorkloadAnalysis

Baseline performance

SeqLenThroughput (TFLOPS)1,0243302,0484414,0965118,19254616,384566Table 1. Baseline performance without any specific optimizations

This is our starting point with 64×64 tiles and no optimizations.

NCU insight (SeqLen=1024, B200):

Registers/thread: 128

Theoretical/achieved occupancy: 25% / 19.8%

Compute (SM) throughput: 37.8%

Memory throughput: 19.7%

Grid size: 2,048

1. The trap of larger tiles

A common intuition in GPU programming is “bigger tiles = better performance.” Larger tiles:

Amortize memory access overhead.

Improve L2 cache utilization.

Reduce kernel launch overhead per element.

So, let’s increase our tile size from 64×64 to 256×128:

TILE_M, TILE_N = 256, 128

The expected is better memory bandwidth utilization → faster performance. However, the result in TFLOPS are:

SeqLenBaseline (64×64)Larger tiles (256×128)Performance Degradation1,024330187-43%2,048441268-39%4,096511347-32%8,192546415-24%16,384566463-18%Table 2. Baseline performance compared to performance with larger tile sizes, showing degradation when using larger tile sizes

Performance degraded by 18-43% across all sequence lengths. This is the trap, where large tiles make performance worse.

Why does this happen?

Compute bottleneck: With more elements per tile, inefficient operations (separate mul/add, precise math) become the bottleneck.

Instruction overhead: More work per tile means more instructions before the next memory operation.

Lesson: Tile size and compute efficiency are interdependent. Large tiles only help if the computation is efficient enough to keep up.

NCU insight (SeqLen=1,024, NVIDIA B200):

Registers/thread jump to 168 (+31%), reducing theoretical occupancy to 18.75%

Achieved occupancy drops to 16.5%

Compute throughput collapses to 17.4% (the trap)

Memory throughput falls to 7.4%

Grid size shrinks to 512 (fewer blocks from larger tiles)

1. The rescue with fast math

One of the bottlenecks is special functions: exp2 (exponential) and truediv (division). By default, these are IEEE-754 precise—highly accurate, but slow.

For deep learning, we can trade a tiny bit of precision for massive speedups:

Before (precise operations):

p = ct.exp2(qk)alpha = ct.exp2(m_i - m_ij)acc = ct.truediv(acc, l_i)

After (fast math):

p = ct.exp2(qk, flush_to_zero=True)alpha = ct.exp2(m_i - m_ij, flush_to_zero=True)acc = ct.truediv(acc, l_i, flush_to_zero=True, rounding_mode=RMd.APPROX)

What these flags do:

flush_to_zero=True: Denormal numbers (extremely small values near zero) become exactly zero. This avoids slow microcode paths on the GPU.

rounding_mode=RMd.APPROX: Skips iterative refinement after initial hardware approximation.

With fast math, we’ve “rescued” the large tiles, and the results in TFLOPS are:

SeqLenLarger tiles (trap)Fast math (rescue)Improvement1,024187322+72%2,048268436+63%4,096347524+51%8,192415585+41%16,384463620+34%Table 3. Performance improvement when using two fast math optimizations

We now match or exceed the small-tile baseline, with 10-20% gains for longer sequences.

NCU insight (SeqLen=1,024, NVIDIA B200):

Registers/thread: 168 (unchanged)

Theoretical/achieved occupancy: 18.75% / 16.6% (unchanged)

Compute throughput rebounds to 24.0%

Memory throughput improves to 12.9%

1. K-loop split

For causal attention, we apply a triangular mask: each query can only attend to keys at earlier positions. In our baseline, we check if CAUSAL: mask… on every loop iteration.

But think about it: for a query tile at position 1000, most key tiles (0-900) need no masking at all. Only tiles near the diagonal need the mask. And tiles beyond the query position are completely masked (we can skip them entirely).

Figure 3. Tiled causal attention matrix (8 tiles per side)

The optimization splits the loop into phases:

mask_start = (input_pos + bid_x * TILE_M) // TILE_Nmask_start = min(mask_start, k_seqlen // TILE_N)if CAUSAL:    Tc = ct.cdiv(min(m_end, k_seqlen), TILE_N)else:    Tc = ct.cdiv(k_seqlen, TILE_N)for j in range(0, Tc):    if (CAUSAL or not EVEN_K) and j >= mask_start:        offs_n = j * TILE_N + offs_n_tile        mask = ct.full((TILE_M, TILE_N), True, dtype=ct.bool_)        if not EVEN_K:            mask = mask & (offs_n < k_seqlen)        if CAUSAL:            mask = mask & (offs_m >= offs_n)        mask = ct.where(mask, 0.0, -math.inf)        qk += mask

Why this matters: For a 16K sequence with 256-token tiles:

~50% of tiles are fully unmasked (no branch, no mask computation)

~1 tile per row is partially masked (full logic)

The rest are skipped entirely (early exit)

Result in TFLOPS:

SeqLenFast mathLoop splitImprovement1,024322373+16%2,048436552+27%4,096524684+31%8,192585770+32%16,384620813+31%Table 4. Performance improvement when using K-loop split optimization

This is the biggest single optimization—up to 32% speedup across all sequence lengths.

NCU insight (SeqLen=1,024, B200):

Registers/thread: 168 (unchanged)

Theoretical/achieved occupancy: 18.75% / 16.6% (unchanged)

Memory throughput improves to 14.5% (less wasted work)

Compute throughput remains 24.0% (work is more useful, not necessarily faster per cycle)

1. ProgramId remapping

One subtle optimization is reversing the block order for causal attention. When we process tiles in reverse (bottom-right to top-left), later-launched blocks have less work due to the causal mask. This improves load balancing and reduces tail effects.

Before (standard order):

bid_x = ct.bid(0)

After (reversed for causal):

if CAUSAL:    bid_x = NUM_M_BLOCKS - 1 - ct.bid(0)  else:    bid_x = ct.bid(0)

This small change improves wave scheduling, as blocks complete more uniformly across the GPU.

Result in TFLOPS:

SeqLenLoop splitRemappingImprovement1,024373377+1%2,048552560+1.5%4,096684696+1.8%8,192770781+1.5%16,384813835+2.6%Table 5. Performance improvement after remapping the block order of the tiles

A modest but consistent 1-3% gain, especially noticeable at longer sequences where tail effects matter most.

1. Autotuning

We’ve optimized large tiles, but there’s a catch: short sequences still prefer small tiles.

Why? With a 1,024-token sequence and 256-token tiles, we only have 4 tiles. That’s not enough to fully utilize all SMs on a B200. Smaller tiles (64×64) give us 16 tiles, better filling the GPU.

Rather than manually choosing a threshold, we can let cuTile’s autotuner benchmark multiple configurations and cache the best one for each input shape.

The autotuner approach:

def _fmha_autotune_configs():    gpu_capability = torch.cuda.get_device_capability()    if gpu_capability in [(12, 0), (12, 1)]:        yield SimpleNamespace(TILE_M=64, TILE_N=64, num_ctas=1, occupancy=2)    else:        yield SimpleNamespace(TILE_M=64, TILE_N=64, num_ctas=1, occupancy=2)        yield SimpleNamespace(TILE_M=128, TILE_N=128, num_ctas=1, occupancy=2)        yield SimpleNamespace(TILE_M=256, TILE_N=128, num_ctas=1, occupancy=1)

How to launch with autotuning:

Instead of calling ct.launch directly, use ct_experimental.autotune_launch:

import cuda.tile_experimental as ct_experimentaldef autotune_launch_fmha(    stream, q, k, v, o, sm_scale, input_pos,    hidden_size, num_heads, query_group_size, is_causal):    batch_size, _, q_len, _ = q.shape    def _grid_fn(cfg):        return (math.ceil(q_len / cfg.TILE_M), batch_size * num_heads, 1)    def _args_fn(cfg):        num_m_blocks = math.ceil(q_len / cfg.TILE_M)        even_k = (k.shape[2] % cfg.TILE_N) == 0        return (            q, k, v, o, sm_scale, input_pos,            hidden_size, num_heads, cfg.TILE_M, cfg.TILE_N,            query_group_size, is_causal, even_k, num_m_blocks,        )    ct_experimental.autotune_launch(        stream,        grid_fn=_grid_fn,        kernel=fmha_kernel,        args_fn=_args_fn,        hints_fn=lambda cfg: {"num_ctas": cfg.num_ctas, "occupancy": cfg.occupancy},        search_space=_fmha_autotune_configs,    )

Note: The autotuner API may be subject to change.

The autotuner works intelligently:

First call with seq_len=1024: Benchmarks all 3 configs, caches best one

First call with seq_len=2048: Benchmarks all 3 configs, caches best one

Subsequent calls: Uses cached config (zero overhead)

The cache key includes tensor shapes, so different sequence lengths automatically get different optimal configurations.

Result in TFLOPS:

SeqLenBaselineRemappingAutotuneSpeedup vs baseline1,0243303775481.66x2,0484415607081.61x4,0965116968171.60x8,1925467818871.62x16,3845668359181.62xTable 6. Original baseline compared to step 5 and to step 6 autotuned results

The autotuner discovers that 64×64 tiles are best for sequences ≤2,048, then transitions to larger tiles for longer sequences. This delivers 45% additional performance at short sequences compared to fixed large tiles, while maintaining peak performance at long sequences.

What the autotuner chose (on B200):

SeqLen 1,024: 64×64 tiles (high parallelism)

SeqLen 2,048: 64×64 or 128×128 tiles (balanced)

SeqLen 4,096+: 128×128 or 256×128 tiles (memory efficiency)

We now achieve optimal performance across all sequence lengths without manual tuning.

Summary: The optimization stack

OptimizationKey insightImpactBaseline (64×64)Correct but unoptimizedBaselineLarge tiles (256×128)TRAP: 18-43% slower!-18% to -43%+ Fast math (FTZ, APPROX)RESCUE: Large tiles now pay off+34% to +72% from trap+ K-loop splitBiggest single optimization+16% to +32%+ ProgramId remappingBetter load balancing+1% to +3%+ AutotuningOptimal tiles per sequence+10% to +45%Table 7. Step-by-step optimization results with performance impacts for each step

Final speedup: 1.60x-1.66x across all sequence lengths.

Getting started

Writing high-performance kernels is rarely about finding one “magic” setting. As we saw with the “trap and rescue”:

Optimizations are interdependent: Large tiles were slower until we fixed the math. You can’t evaluate tile size in isolation.

Math matters: Flags like flush_to_zero and APPROX are critical for unlocking Tensor Core throughput. Precise math is often overkill for deep learning.

Algorithmic wins compound: K-loop splitting gave us the biggest single improvement (up to 32%) by avoiding unnecessary work.

Autotuning beats manual heuristics: cuTile’s autotuner discovers optimal tile sizes per sequence length (64×64 for short sequences, 256×128 for long), delivering 10-45% gains over fixed configurations.

Cumulative effects are multiplicative: The full optimization stack delivers 1.60x-1.66x speedup across all sequence lengths—far more than any single optimization alone.

cuTile enables developers to express these optimizations—tiling, fast math controls, loop splitting, autotune—in clean, readable Python code while generating highly optimized PTX for NVIDIA GPUs.

You can find the completely optimized kernel in the TileGym repository. Happy hacking.