Flash Attention in CUDA

Build a tiled, IO-aware Flash Attention kernel with online softmax and causal masking.

7 parts26 individual lessons2.2 estimated hours130 XP available

0%0 of 26 complete

No compressed chapters.

Every source step is its own lesson with intuition, concepts, correctly rendered MathJax mathematics, implementation, tests, mistakes, and a checkpoint.

PART 1

CUDA Primitives Warm-up

Implement basic elementwise and reduction kernels that form the building blocks for matrix and attention operations.

0/5

001vector add+5 XP 002scale array+5 XP 003elementwise exp+5 XP 004row max+5 XP 005row sum+5 XP

PART 2

Matrix Operations

Build the matmul, transpose, and dot-product utilities needed for attention scoring.

0/3

006dot product+5 XP 007matmul+5 XP 008transpose+5 XP

PART 3

Naive Attention Baseline

Compose QK^T scoring, row-wise softmax, and PV multiplication into a straightforward attention pipeline.

0/4

009qk scores+5 XP 010softmax rows+5 XP 011pv matmul+5 XP 012naive attention+5 XP

PART 4

Online Softmax Math

Implement the running-max and running-sum updates that allow softmax to be computed incrementally across tiles.

0/4

013online max+5 XP 014correction factor+5 XP 015update running sum+5 XP 016rescale output+5 XP

PART 5

Tiled Attention Building Blocks

Write the per-tile shared-memory routines for loading, scoring, reducing, exponentiating, and accumulating PV.

0/6

017load tile+5 XP 018tile scores+5 XP 019tile rowmax+5 XP 020tile exp+5 XP 021tile rowsum+5 XP 022accumulate pv+5 XP

PART 6

Fused Flash Attention Kernel

Assemble the tiled blocks into the full Flash Attention kernel and its host launcher.

0/2

023flash attention kernel+5 XP 024flash attention launcher+5 XP

PART 7

Causal Flash Attention

Extend the kernel with a causal mask for autoregressive attention.

0/2

025causal mask+5 XP 026flash attention causal kernel+5 XP