Flashattention

本文从数学原理出发，深入分析FlashAttention的核心思想、算法设计和各版本演进，通过详实的数学推导、直观的流程图表和具体的数值示例，帮助读者真正掌握这一革命性的Attention优化技术。 1. 问题的本质：传统Attention的根本瓶颈 1.1 传统Attention机制的计算模式 传统的Self-Attention机制遵循如下计算流程： $$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V $$ 让我们用具体数值来理解这个过程的复杂性： 示例场景：考虑一个典型的语言模型场景 序列长度：$n = 2048$（如GPT-2的上下文长度） 特征维度：$d_k = 64$（每个attention head的维度） 输入张量形状：$Q, K, V \in \mathbb{R}^{2048 \times 64}$ 第一步：计算注意力得分矩阵 $$S = \frac{QK^T}{\sqrt{d_k}} \in \mathbb{R}^{2048 \times 2048}$$ 这一步产生了一个 $2048 \times 2048 = 4,194,304$ 个元素的矩阵，以FP16精度存储需要约8MB内存。 第二步：Softmax归一化 $$P = \text{softmax}(S) \in \mathbb{R}^{2048 \times 2048}$$ Softmax计算需要： 计算每行的最大值：$m_i = \max_j S_{i,j}$ 计算指数和：$l_i = \sum_j e^{S_{i,j} - m_i}$ 归一化：$P_{i,j} = \frac{e^{S_{i,j} - m_i}}{l_i}$ 这又需要存储另一个 $2048 \times 2048$ 的矩阵。 ...