Attention Kernel Example

Implement Flash Attention v2 using CuTe DSL for efficient transformer inference

This guide demonstrates how to implement Flash Attention v2 using CUTLASS CuTe DSL. Flash Attention is a memory-efficient attention mechanism that reduces memory usage from O(N²) to O(N) through online softmax computation.

Overview

Flash Attention v2 computes scaled dot-product attention:

O = softmax(Q * K^T / sqrt(d)) * V

Where:

• Q: Query matrix (B×Sq×N×H)
• K: Key matrix (B×Sk×N×H)
• V: Value matrix (B×Sk×N×H)
• O: Output matrix (B×Sq×N×H)
• B: batch size, Sq/Sk: sequence length, N: number of heads, H: head dimension

Key Features

• Online softmax: Computes softmax incrementally without materializing full attention matrix
• Tiled computation: Processes attention in blocks to fit in shared memory
• Register pipeline: Overlaps shared memory loads with computation
• Causal masking: Optional support for autoregressive models

Architecture

┌─────────┐  ┌─────────┐
│    Q    │  │  K, V   │
└────┬────┘  └────┬────┘
     │            │
     ├─CpAsync────┤
     ↓            ↓
  ┌────────────────────┐
  │  Shared Memory     │
  │  sQ    sK    sV    │
  └──┬──────┬──────┬───┘
     │      │      │
     ↓      ↓      ↓
  ┌────────────────────┐
  │  Tensor Core MMA   │
  │    S = Q×K^T       │
  └─────────┬──────────┘
            ↓
  ┌──────────────────────┐
  │  Online Softmax      │
  │  - Update row_max    │
  │  - Update row_sum    │
  │  - Rescale prev O    │
  └─────────┬────────────┘
            ↓
  ┌──────────────────────┐
  │  Tensor Core MMA     │
  │    O = P×V           │
  └─────────┬────────────┘
            ↓
  ┌──────────────────────┐
  │  Normalize & Store   │
  └──────────────────────┘

Implementation

Running Examples

Standard Attention

python examples/python/CuTeDSL/ampere/flash_attention_v2.py \
  --dtype Float16 \
  --head_dim 128 \
  --m_block_size 128 \
  --n_block_size 128 \
  --num_threads 128 \
  --batch_size 1 \
  --seqlen_q 1280 \
  --seqlen_k 1536 \
  --num_head 16 \
  --softmax_scale 1.0

python examples/python/CuTeDSL/ampere/flash_attention_v2.py \
  --dtype Float16 \
  --head_dim 128 \
  --m_block_size 128 \
  --n_block_size 128 \
  --num_threads 128 \
  --batch_size 1 \
  --seqlen_q 1280 \
  --seqlen_k 1536 \
  --num_head 16 \
  --softmax_scale 1.0 \
  --is_causal

Performance Profiling

ncu python examples/python/CuTeDSL/ampere/flash_attention_v2.py \
  --dtype Float16 --head_dim 128 \
  --m_block_size 128 --n_block_size 128 \
  --num_threads 128 --batch_size 1 \
  --seqlen_q 1280 --seqlen_k 1536 \
  --num_head 16 --softmax_scale 1.0 \
  --is_causal --skip_ref_check

Key Concepts

Online Softmax Algorithm

The online softmax algorithm processes attention in blocks without materializing the full matrix:

Traditional approach (memory intensive):

S = Q @ K^T              # O(N²) memory
P = softmax(S)           # O(N²) memory
O = P @ V                # O(N²) memory

Flash Attention approach (memory efficient):

For each K/V block:
  1. Compute S_block = Q @ K_block^T
  2. Update row_max = max(row_max_prev, max(S_block))
  3. Rescale previous O by exp(row_max_prev - row_max)
  4. Update row_sum with new exp values
  5. Compute O += softmax(S_block) @ V_block

Finally: O = O / row_sum

This reduces memory from O(N²) to O(N).

Register Pipeline

The register pipeline overlaps shared memory loads with computation:

# Load first k-block
cute.copy(smem_tiled_copy_Q, tSsQ[0], tSrQ_view[0])
cute.copy(smem_tiled_copy_K, tSsK[0], tSrK_view[0])

for k in range(num_k_blocks):
    # Load next block while computing current
    k_next = (k + 1) % num_k_blocks
    cute.copy(smem_tiled_copy_Q, tSsQ[k_next], tSrQ_view[k_next])
    cute.copy(smem_tiled_copy_K, tSsK[k_next], tSrK_view[k_next])
    
    # Compute with current block
    cute.gemm(tiled_mma, acc_S, tSrQ[k], tSrK[k], acc_S)

Tile Size Selection

Choosing optimal tile sizes:

m_block_size (Query tile):

• Larger: Better arithmetic intensity, more registers
• Smaller: Lower SMEM usage, better for long sequences
• Typical: 64, 128

n_block_size (Key/Value tile):

• Affects online softmax accuracy
• Should match m_block_size for balanced computation
• Typical: 64, 128

head_dim:

• Usually fixed by model (64, 128, 256)
• Must be 16-byte aligned
• Pad if necessary

Constraint: (m_block_size * head_dim + 2 * n_block_size * head_dim) * sizeof(dtype) must fit in SMEM

Memory Analysis

Component	Standard Attention	Flash Attention v2
S matrix	O(B·N·Sq·Sk)	O(B·N·m_block·n_block)
P matrix	O(B·N·Sq·Sk)	O(B·N·m_block·n_block)
Total GMEM	O(N²)	O(N)
SMEM usage	Low	High (Q+K+V tiles)

Constraints

• Only fp16 and bf16 data types supported
• Head dimension must be 16-byte aligned (multiple of 8 elements)
• m_block_size * 2 must be divisible by num_threads
• Total SMEM: (m_block + 2*n_block) * head_dim * sizeof(dtype) < capacity
• Log-sum-exp for training backward pass not computed

Source Code

Flash Attention v2 (Ampere)

Complete source: examples/python/CuTeDSL/ampere/flash_attention_v2.py

Flash Attention (Hopper)

Hopper variant: examples/python/CuTeDSL/hopper/fmha.py