LogoCookLLM Docs
LogoCookLLM Docs
HomeCookLLM

Principles

Tokenization
Tokenization BasicsBPE AlgorithmGPT TokenizersBPE Training Engineering
Model Architecture
Transformer LM
From token ids to logitsEmbedding and LM Head
Attention Mechanisms
From Self-Attention to GQAAttention Sink
Position Encoding
Position Encoding BasicsRoPE Math DerivationRoPE ImplementationLength Extrapolation
GPU Programming Basics
GPU Architecture BasicsTensor LayoutTriton Basics: Vector Add
FlashAttention
Flash Attention PrinciplesFrom Naive Implementation to Auto-TuningBlock Pointers and Multi-Dim SupportCausal Masking OptimizationGrouped Query AttentionBackward Pass Implementation
Distributed Training
Data ParallelismZeRO OptimizerFully Sharded Data ParallelTensor ParallelismPipeline ParallelismMulti-Dimensional Hybrid Parallelism

Hands-on Training

Overview
Pretraining
Pretraining DataTokenizer TrainingModel ArchitectureData PipelineTraining LoopMonitoring and Validation
X (Twitter)
SystemsFlashAttention

Flash Attention Principles

Premium

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

Log in to continue reading

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

Flash Attention

Deeply understand Flash Attention principles and Triton implementation

From Naive Implementation to Auto-Tuning

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

Table of Contents

The Memory Bottleneck in Standard Attention
GPU Memory Hierarchy: SRAM and HBM
The Logical Trap in the Standard Implementation
The Bandwidth Gap Between SRAM and HBM
Comparison of Speed Differences
The Essence of the Bottleneck: IO-bound
The Capacity Limit of SRAM
Physical Constraints and Cost
The Capacity Limit
The Core Idea: Optimizing IO Complexity
Avoiding Spilling the Intermediate Matrix to Memory
The Principle of Online Softmax
The Limitations of Offline Algorithms
Online Algorithms and Dynamic Correction
Deriving the Correction Formula
Numerical Demonstration: Using the Sequence [3, 2, 5, 1] as an Example
A Summary of the Mathematical Principle Behind Flash Attention
Tiled Matrix Multiplication (Tiling)
Why Do We Need to "Tile"?
Visual Demonstration: The Tiled Computation Flow
Key Things to Observe
Combining Tiling with Attention
Comparing Loop Strategies: V1 vs V2
Interactive Guide to the Diagram
Softmax Correction in Tiled Attention
The Naive Implementation: The Limitations of Local Softmax
The Solution: Online Rescaling
Initialization
Inner Loop: Traversing the K-Blocks
Final Step: Normalization
Complete Algorithm Pseudocode
The Full Picture of the Algorithm
Summary