✨ Awesome-Anything

Awesome-AI-Researchs

Awesome-GitHub-Repos

Awesome-AI-Examples

🖼️ Digital Image Processing

DIP I Image Demosaicing and Histogram Manipulation

DIP II Image Denoising

DIP III Edge Detection

DIP IV Digital Half-toning

DIP V Geometric Image Modification

🍃 LLM Components

1.1 Tokenization

1.2 Embedding

1.3 Attention

1.4 FFN, Residual Addition, and LN

1.5 Positional Encoding

1.6 Architecture & Decoding Policy

🌱 LLM Pre-training

2.1 Pre-training Definition

2.2 Pre-training Data

2.3 Pre-training Workflow

2.4 Continued Pre-training (CPT)

2.5 Pre-training Evaluation

☘️ LLM Post-Training

3.1 Supervised Fine-Tuning (SFT)

3.2 Reinforcement Learning

3.3 RLHF: Reinforcement Learning from Human Feedback

3.4 Optimizing PPO-Based Algorithms

3.5 Direct Preference Optimization (DPO)

🍀 LLM Popular Models

🪴 LLM Applications

5.1 Prompt Engineering

5.2 LLM-based Agents

5.3 Retrieval-Augmented Generation (RAG)

5.4 Enterprise-level RAG Practice (I)

5.5 Enterprise-level RAG Practice (II)

🌿 LLM Optimization

6.1 GPU Usage Analysis

6.2 Latency Analysis for Inference

6.3 FlashAttention

6.4 PagedAttention

6.5 Packing for Efficient Training

6.6 PEFT

🌾 LLM Compression

7.1 Quantization

7.2 Pruning

7.3 Knowledge Distillation

🌵 LLM Hands-on Practice

8.1 nanoGPT: Training a GPT from Scratch

8.2 Instruction Fine-Tuning Qwen3-0.6B: A Minimal Working Example

8.3 OpenRLHF: SFT Qwen3-0.6B and Code Analysis

8.4 OpenRLHF: Reward Model Training and Code Analysis

8.5 OpenRLHF: PPO Tutorial and Code Analysis

🌴 LLM Must-read Papers

9.1 Technical Reports Explained

9.2 MedLLM

9.3 o1 Techinical Paper

9.4 Prompting

9.5 LLM-as-Judge

🌳 LLM Q&A

10.1 Why GRPO training loss starts at 0

10.2 VLM-R1 Reproduction Results

10.3 KL(P||Q) or KL(Q||P)?

10.4 Left Padding or Right Padding?

10.5 What is QK-Norm?

10.6 Why LLM training rarely adopts Dropout?

10.7 Dropout and DropPath

🐝 VLM Image Encoders

11.0 Introduction

11.1 Vision Transformer

📝 MISC.

🤔 Large Language Models Pass the Turing Test.

OpenAI ChatGPT Progression

⏳ Claude Sonnet 4 Test with Sandtris

🔍 What DeepSeek-R1's Surge Teaches Us About SEO

💰 CivitAI’s Payment Issue

🎓 arXivSub, an Academia Assistant

6.3 FlashAttention

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

🎈

FlashAttention proposes an acceleration technique that significantly speeds up computation, saves memory, and achieves IO-aware precise attention. It can effectively alleviate the computational and storage complexity that grows quadratically with the sequence length N in transformers, improving resource usage and efficiency.

FlashAttention can be applied during both training and inference, with key features including:

Faster Computation: FlashAttention doesn’t reduce FLOPs but reduces the number of HBM accesses by being IO-aware, thus speeding up computation. This is achieved by using tiling and operator fusion.

Memory-efficient: FlashAttention utilizes extra statistics to change the computation order in the attention mechanism, avoiding the need to materialize the full attention matrix and reducing the space complexity from O(N²) to O(N).

Exact Attention: Unlike approximate methods like sparse attentions, FlashAttention performs exact computation via tiling. The results generated by FlashAttention are mathematically equivalent to standard attention.

6.3.1 Computation and Memory Limitations

💡

The computational complexity and memory complexity of the self-attention blocks both scale quadratically with the sequence length N.

Many approximate methods have attempted to reduce the computation and memory demands of attention. For example, sparsification or low-rank methods try to lower the complexity from quadratic to linear or sub-quadratic.

However, these methods have not seen widespread adoption, mainly because they focus on reducing FLOPs but overlook the memory access bottleneck.

On modern GPUs, computation speed has far outpaced memory access speed, and in transformers, the major bottleneck for computational operations is memory access.

📎

Some Concepts

FLOPS: Floating point operations per second. It measures how many floating-point operations a processor can perform per second. It’s a standard metric for hardware performance.

FLOPs: Floating point operations. It refers to the count of floating-point operations needed by an algorithm or model, representing the algorithm’s complexity.

Math bandwidth: The number of math operations a processor can execute per second, commonly measured in OPS (operations/second). If we use floating-point operations, the unit becomes FLOPS.

Note that math bandwidth measures the rate at which math operations are performed, while compute bandwidth refers to the overall throughput of the processor, including both math and other operations.

Memory bandwidth: The amount of data a processor can read from or write to memory per second, usually measured in bytes/second.

Math-bound: When performance is limited by compute capability. Typical cases include large matrix multiplications and deep-channel convolutional operations.

Memory-bound: When performance is limited by memory bandwidth. Element-wise operations like activation functions, dropout, and mask are often memory-bound. Reduction operations are also often memory-bound, for example: softmax, batch normalization, and layer normalization.

The softmax and normalization themself do not reduce vectors to scalers, but they involve reduction operations like: summation for softmax, mean & variance for normalization

For the self-attention block, aside from the matrix multiplication being math-bound, most other operations like softmax, dropout, and mask are memory-bound.

GPU Memory Hierarchy

GPU memory consists of various types and sizes of memory blocks, each with different read/write speeds. The closer to the core, the faster the speed.

On-Chip Memory: Mainly refers to caches and some specialized storage (e.g., texture memory). It features "small capacity but high bandwidth." In this context, SRAM belongs to on-chip memory. Its capacity is only about 20MB, but the bandwidth can reach up to 19TB/s.
Off-Chip Memory: Mainly refers to global memory, which is essentially the GPU's display memory, a.k.a. Video RAM (VRAM). It is characterized by "large capacity but low bandwidth." Here, HBM belongs to off-chip memory. For example, an A100 GPU has 40GB (or 80GB) of HBM, but the bandwidth compared to SRAM is much lower, at about 1.5TB/s (or 1.9TB/s). Therefore, minimizing HBM read/write operations and making full use of the faster SRAM for computation is crucial for performance.

GPU Execution

GPUs run operations through a large number of threads, each executing a kernel. The typical three-step execution process is:

Each kernel loads input data from the slower HBM into the faster SRAM.
Perform computation inside SRAM.
Write the computation results from SRAM back to HBM.

Kernel Fusion

For operations that are heavily memory-bound, a common optimization method is kernel fusion. It speeds up execution by avoiding the need to load data from HBM repeatedly. The idea is to merge multiple operations into one, reducing the number of HBM reads/writes.

For example, if I have two calculations, A and B. The traditional way would compute A first, write the result to memory, then load it back and compute B.

However, if A’s result is used immediately by B, and both can fit into SRAM, then fusing A and B into one kernel will save a lot of time on memory reads/writes.

<ins/>

6.3.2 Standard Attention and Safe Softmax

Standard Attention

First, define q, k, v: batch_size is set to 1, seq_len is set to N, and emb_size is set to d.

We only focus on the attention part here, ignoring dropout, mask, etc. Below is the standard attention computation process: ,

Safe Softmax

For float32 and bfloat16, when , becomes inf, leading to overflow issues. To ensure numerical stability, we usually subtract the maximum value during the computation. This is mathematically equivalent to the standard softmax, called safe softmax:

6.3.3 FlashAttention Forward Process

Tiling

🧩

Tiling

The read/write bandwidth of SRAM is an order of magnitude higher than that of HBM, but its capacity is much smaller. By fusing multiple operations into a single kernel operation, we can reduce the number of times we access HBM, thus improving the efficiency of memory-bound operations.

However, since SRAM has limited capacity, it may not be possible to complete the entire computation in one go. We must perform tiling to ensure that the working set for each tile fits within SRAM.

Difficulty of Tiling

Notice that the original steps of the attention mechanism are matrix multiplication --> scale --> mask --> softmax --> dropout --> matrix multiplication. While matrix multiplication and element-wise operations are relatively easy to tile, tiling becomes very difficult for softmax because it requires the full input data for each row.

How FlashAttention Works

It introduces extra statistics , to achieve decoupling.

Safe Softmax

Decoupled Softmax

where the , which is why we multiply to it to get the correct first-half values for ; same reason holds for . Similarly, you know why there is also an exponential term in .

By maintaining two extra statistics, we can achieve a tiled computation of softmax. At the same time, note that multiple softmax can be computed in parallel on the GPU, improving computational efficiency.

Kernel Fusion

Fusing mask and dropout into the forward pass (below is for one tile):

Tiling allows all attention operations to be executed with a single CUDA kernel

Data is loaded from HBM into SRAM
All computation operations (matrix multiplication, masking, softmax, dropout, matrix multiplication) are performed in SRAM
Then the results are written back to HBM

By fusing multiple operations into one kernel, it avoids repeatedly reading and writing data from HBM. Above is a very simplified version; next, we will dive deeper into the details of how the final output is aggregated.

🌻

A Simple Example

For the vector [1, 2, 3, 4], we compute the softmax by splitting it into two blocks: [1, 2] and [3, 4].

Compute block1:

Compute block2:

Aggregrate into final softmax output:

Above is a simple example showing the computation of decoupled softmax. However, in real-world case, we cannot really first computing all and before we calculating the softmax because this way doesn’t fully decouple the tiles, requiring more storage for intermediate results and more IO operations. Next, we’ll combine the full calculation of an attention tile into a single kernel operation.

<ins/>

Full Forward Process

Original Attention (w/o mask & dropout)

Matrices Q, K, V are stored in HBM with dimensions .

Load Q and K from HBM into SRAM.

Compute .

Write S back to HBM.

Load S from HBM into SRAM.

Compute .

Write P back to HBM.

Load P and V from HBM into SRAM.

Compute .

Write O back to HBM.

Return O.

Our Goal

For each tile of Q, K, V, we denote , , , which are stored with , in HBM

Load all to SRAM.

Compute

Compute

Compute , ,

Write , , back to HBM.

In the above, by simplifying the detailed computation, we can see the idea behind the tiling that all computations for attention in a tile/block are done before writing back to HBM.

Let’s see the full steps by following the official algorithms

We already have Q, K, V in HBM, and the SRAM only has size of M, then we initialize the O, l, and m, in HBM too.

Set the tile/block size:

for K and V
for Q and O

Why taking minimum with ? Because we will generate intermediate results such as in shape of . We don’t want this take space larger than . Of course you can set it as other values as long as all parameters properly fit in the SRAM.

First, split the Q matrix into blocks, where each block has a shape of . It’s easy to understand that stores the query information of tokens.

Similarly, we do the splitting to matrix into blocks, where each block has a shape of . Also, stores the key information of tokens.
Split the V matrix into blocks, where each block has a shape of . And stores the value information of tokens.

The O, l and m are also split accordingly into blocks, each of which has shapes of , , and , respectively.

Below shows a single tile/block computation for index and , but the actual updates involving the softmax are a bit more complicated than it looks.

For Loop: outer loop for and inner loop for

We traverse through the on and in the outer loop; and traverse through the on , , , and in the inner loop. You may notice that, to calculate a single , we actually need to compute a weighted sum on all , so in each tile, the computed is not the actual value of it, and we need to iteratively update it until the end of the traversal of . This is the same case for and .

For each and , we first calculate the , which is an intermediate result has a shape of

Next, we do the safe softmax (for )

As you remember:
Here we are doing:

When we are calculating , we don’t have access to , but we can still take the max with the previous ones; so coming to the end (i.e. until all is traversed), we can get the final

Here we cache the and wait until the end of computation for this tile/block to update the , because the from the last round will be used for calculating and updating later

Now, let's do the safe softmax (for )

Remember that the is the summation of all elements after the safe exponential operation

Here we first do the exponential to get using the calculated

Then calculate

Clearly, this is based on , which is not the actually , but just like in updating , here we can also update the interatively

Previous is based on previous round’s , so we multiply it by to replace it with current ; in this around, the is calculated based on , so we multiply it with to replace it with

Still, we wait until the end of this tile’s computation to update the , because the calculation for needs last round’s and

Now, let’s update the

Originally, to get , we need the and the for all , but the is intermediate results and all we have are and . Therefore, we have to use last round’s and and current and to update , iteratively.

The final formula is

This formula is hard to understand because it’s the outcome of a long derivation. To understand it, we have to derive it from the very beginning.

📌

Derivation for

where the represents the slice in . Note we start from 1 instead of 0.

Just a heads up: is the safe exponential (one element in the numerator from the saft softmax, actually before adjusted by ) and the is the safe softmax

Congrats! You’ve understood the most complex part. Now let’s update the , and , so we can use it in the next round of

<ins/>

A simple implementation w/ C++

You can easily run this code on Colab, and observe results like

Note that in official implementation, if you are using it for training models, FlashAttention does not save the intermediate results needed for backpropagation. Instead, they are recomputed during backward.

<ins/>

Computation and GPU Memory

🛠

Computational Complexity

In line 9 of the algorithm, we have , where , . The computational complex of calculating is .

In line 12 of the algorithm, we have , where , . Similarly, the computational cost here is .

Next, the total computation cost is simply how many times steps (1) and (2) are executed. Combining the above, the forward computational cost of flash attention is:

Similarly, the backward computational cost is also summarized in the paper as .

Overall, the computational complexity of FlashAttention is similar to the standard Attention, with increased FLOPs due to recomputation during backpropagation.

Memory Complexity

Compared to standard attention, if we don't consider the output , FlashAttention only needs to store , with a memory requirement of .

Standard attention, however, needs to store , with a memory requirement of .

Clearly, compared to standard attention, Flash Attention significantly reduces memory usage.

IO Complexity

Let's first look at line 6 of the code. In each outer loop, we load blocks of K and V. After all outer loops are finished, it’s like we load the complete K and V where , the IO complexity here is:

Now let's look at line 8 of the code. In each inner loop, we also load parts of blocks. Since are relatively small (IO complexity is ), we can ignore them and just consider . In each fixed outer loop, all inner loop iterations together will fully traverse where . At the same time, we will loop through times. Therefore, the IO complexity here is:

Loading back to HBM will have an overall IO complexity:

Therefore, the total IO complexity of flash attention is:

The paper mentions that usually falls between 64 and 128, and the is about 100KB, so:

Thus, we can see that the IO complexity of flash attention is significantly better than of standard attention.

6.2 Latency Analysis for Inference

6.4 PagedAttention

Catalog

Last update: 2025-04-28

🧬 AI Learning Blog 🧬

Anticipated updates:

Must-read papers, and Hands-on practice in general
Digital Image Processing/Mathematics
Vision Language Model