Principles

Tokenization Basics BPE Algorithm GPT Tokenizers BPE Training Engineering

Model Architecture

From token ids to logits Embedding and LM Head

Attention Mechanisms

From Self-Attention to GQA Attention Sink

Position Encoding

Position Encoding Basics RoPE Math Derivation RoPE Implementation Length Extrapolation

GPU Programming Basics

GPU Architecture Basics Tensor Layout Triton Basics: Vector Add

Flash Attention Principles From Naive Implementation to Auto-Tuning Block Pointers and Multi-Dim Support Causal Masking Optimization Grouped Query Attention Backward Pass Implementation

Distributed Training

Data Parallelism ZeRO Optimizer Fully Sharded Data Parallel Tensor Parallelism Pipeline Parallelism Multi-Dimensional Hybrid Parallelism

Hands-on Training

Pretraining Data Tokenizer Training Model Architecture Data Pipeline Training Loop Monitoring and Validation

SystemsFlashAttention

From Naive Implementation to Auto-Tuning

Premium

Write your first Flash Attention kernel and use Auto-Tune for performance optimization.

Get code access

In the previous chapter, we derived the math behind Flash Attention (Tiling + Online Softmax). Now it is time to turn the math into code.

Core Loop Structure

Three questions sit between the formula and a working kernel, and it pays to settle them first.

Log in to continue reading

This is premium content. Please log in to access the full article.

Which loop is parallel, and which is serial. The derivation was written as "Outer Q, Inner K": the outer loop walks blocks of $Q$ , the inner loop walks blocks of $K, V$ . In Triton you will never see that outer loop, because the SPMD model parallelizes it implicitly: each Program claims one $Q$ block, and tl.program_id(0) is its index. The only loop written as a for is the inner one. Miss this and you will spend the whole kernel hunting for an outer loop that does not exist.

Which data stays put across iterations. A $Q$ block is loaded once and reused for the entire inner loop; $K$ and $V$ are re-loaded every round. That distinction decides whether a tl.load belongs inside or outside the loop, and it is precisely why Flash Attention drops HBM traffic from $O(N^2)$ to $O(N^2 d / M)$ .

How much state the accumulator carries. Naive softmax can compute a full row and normalize at the end. Tiled softmax cannot: each new $K$ block may raise the running maximum, which means everything accumulated so far has to be rescaled. So alongside the output accumulator acc, we carry a running max and a running sum — three pieces of state that survive across iterations.

With those three settled, the kernel below is just a line-by-line translation into Triton:

Flash Attention Principles

Through interactive visualizations, gain a deep understanding of Flash Attention's core techniques: the memory bottleneck, Online Softmax, and tiled matrix multiplication.

Block Pointers and Multi-Dim Support

Scale from single sequence to Batch/Head parallelism and simplify pointer math with block pointers.

Table of Contents

Core Loop Structure

Understanding Why tl.constexpr Is Necessary

Understanding Pointer Arithmetic

Interaction Guide

Verifying Numerical Correctness

Using Auto-Tuning to Find the Best Configuration

Introducing @triton.autotune

Pipeline Parallelism

Key Parameter Analysis

(

N2

d

/

M

)