Flash Attention

Flash Attention解决了什么问题？ 标准注意力计算需要将整个N×N的注意力矩阵存储在GPU高带宽内存（HBM）中。对于长序列，这个矩阵非常大——128K序列长度的注意力矩阵需要约64GB HBM。GPU核心（SM）与HBM之间的数据搬运成为瓶颈。 Flash Attention的核心创新：不在HBM中实例化完整注意力矩阵，而是在SRAM中分块计算。 GPU内存层次 理解Flash Attention需要先理解GPU的内存层次： SRAM (片上共享内存) ├── 延迟: ~20 cycles ├── 带宽: ~19 TB/s (A100) └── 容量: ~192KB per SM HBM (高带宽内存) ├── 延迟: ~200+ cycles ├── 带宽: ~2 TB/s (A100) └── 容量: 80GB (A100) 标准注意力的问题：在HBM中读写O(n²)大小的矩阵，受限于2TB/s的HBM带宽。 Flash Attention的解决思路：将计算分块，每块在SRAM中完成，减少HBM访问次数。 算法原理 标准注意力计算 S = Q @ K^T / √d # [N, N] 注意力分数 P = softmax(S) # [N, N] 归一化 O = P @ V # [N, d] 输出 问题：S和P是N×N矩阵，需要完整存储在HBM中。 ...

深度解析注意力机制的变体与演进 注意力机制是Transformer架构的灵魂。自2017年"Attention is All You Need"发表以来，围绕标准Self-Attention的改进工作从未停止。在2026年，这个领域已经发展出一棵枝繁叶茂的技术树。本文将系统梳理注意力机制的主要演进方向和关键变体。 标准Self-Attention回顾 标准注意力机制的核心公式 O = softmax(QK^T / √d_k) V，其计算复杂度为O(n²d)，其中n是序列长度，d是特征维度。这个n²复杂度是所有注意力优化的核心敌人——当序列长度从2K增长到128K，计算量增长4000倍。 标准注意力还有两个特性值得注意：一是它是一种"全局注意力"，每个位置都和所有位置交互；二是它是"内容相关"的，注意力权重完全由Q和K的内容决定。后续的许多变体正是从打破这两个特性入手的。 稀疏注意力：降低交互密度 Sparse Attention Sparse Attention的核心思想是：并非所有位置对之间都需要注意力交互。通过限制每个位置只关注局部窗口和少数"全局位置"，计算复杂度可以从O(n²)降到O(n√n)或O(n log n)。 Longformer和BigBird是这个方向的代表。Longformer使用滑动窗口加少量全局token的策略，适合长文档处理。BigBird在滑动窗口基础上增加了随机连接，从理论上证明了随机性对表达能力的贡献。 Block Sparse Attention 将注意力矩阵分块，只计算部分块的注意力。GPT-4系列模型使用的便是这种策略——将序列分成固定大小的块，块内全注意力，块间稀疏注意力。这种方法在硬件实现上效率很高，因为矩阵分块运算天然适合GPU的Tensor Core。 线性注意力：打破二次方魔咒 线性注意力是数学上最优雅的改进方向。其核心思想是将softmax(QK^T)V改写为φ(Q)(φ(K)^T V)的形式，利用矩阵乘法结合律，将复杂度从O(n²d)降到O(nd²)。当d远小于n时（这在大模型中通常成立），这是一个巨大的加速。 Performer Performer使用随机特征映射（Random Features）来近似softmax函数。其优势是理论保证好——近似误差可以控制。但实践中，近似精度和下游任务性能之间存在权衡。 Linear Transformer Linear Transformer使用ELU+1作为核函数替代softmax，更加简洁。虽然近似程度不如Performer，但计算效率更高，且在某些任务上表现意外地好。 线性注意力的共同问题是：信息瓶颈。当序列长度远大于特征维度时，线性注意力的表达能力受限，因为所有信息需要压缩到d×d的矩阵中。这导致线性注意力在需要精确位置信息的任务上表现较差。 Flash Attention：硬件感知的优化 Flash Attention不是一个数学上的改进，而是一个系统工程的突破。它通过利用GPU的SRAM层次结构，将注意力计算分块进行，避免在HBM（高带宽内存）和SRAM之间反复搬运数据。 虽然理论上仍然是O(n²)复杂度，但实际运行速度快2-4倍，内存使用大幅降低。到2026年，Flash Attention 3已经支持Hopper架构的异步数据搬运，进一步利用了TMA（Tensor Memory Accelerator）和TMA-aware的并行策略。 Flash Attention的成功给我们的启示是：在大模型时代，算法优化不能只看渐近复杂度，必须考虑硬件特性。一个"理论复杂度更高但硬件友好"的算法，在实践中可能优于"理论复杂度更低但硬件不友好"的算法。 长上下文注意力 RoPE与位置外推 旋转位置编码（RoPE）本身不是注意力变体，但它极大影响了长上下文注意力的工作方式。RoPE的优势是可以通过插值实现位置外推——用32K位置训练的模型可以外推到128K甚至更长。 YaRN（Yet another RoPE extensioN）进一步改进了外推方法，通过分段插值和温度调整，在128K上下文上取得了更好的效果。到2026年，大多数支持长上下文的主流模型都采用了基于RoPE的外推方案。 Ring Attention Ring Attention解决了超长上下文的分布式计算问题。它将序列分布在多个GPU上，每个GPU计算局部的注意力，然后通过环形通信传递中间结果。这使得在有限显存下处理1M+长度序列成为可能。 注意力机制的效率-表达力权衡 回顾所有这些变体，我们能看到一个共同的权衡：计算效率 vs 表达能力。 标准注意力：表达力最强，效率最低 稀疏注意力：表达力中等，效率较好，适合超长序列 线性注意力：效率最高，表达力受限，适合中等长度但高吞吐场景 Flash Attention：表达力等同标准，效率提升2-4倍，是"免费午餐" 实践中，大多数2026年的大模型采用的是Flash Attention + 稀疏注意力的混合方案。在8K-32K的标准上下文窗口内使用Flash Attention，在32K+的扩展窗口上切换到稀疏模式。 ...

注意力计算的性能瓶颈 标准注意力计算的核心瓶颈不在于计算量，而在于显存访问。注意力矩阵 QK^T 的形状为 [n, n]，对于 n=8192 的序列，仅这个矩阵就需要256MB显存（FP32）。GPU的计算速度远超显存带宽——A100的计算能力为312 TFLOPS，而HBM带宽仅2TB/s。 这就是经典的"内存墙"问题：注意力计算大部分时间不是在算，而是在等数据搬运。 Flash Attention V1：分块计算的突破 Tri Dao在2022年提出的Flash Attention V1核心思想是：通过分块计算避免实例化完整的注意力矩阵。 在线Softmax Flash Attention的关键创新是在线Softmax算法。标准Softmax需要先遍历一次数据求最大值，再遍历一次计算指数和归一化。在线Softmax通过维护一个运行最大值和运行和，在单次遍历中完成Softmax计算： def online_softmax(blocks): """在线Softmax：逐块更新""" m = float('-inf') # 运行最大值 l = 0.0 # 运行和 result = None for block in blocks: block_max = block.max() # 更新全局最大值 new_m = max(m, block_max) # 修正之前的累积值 alpha = math.exp(m - new_m) beta = math.exp(block_max - new_m) l = alpha * l + beta * block.sum() if result is None: result = beta * block else: result = alpha * result + beta * block m = new_m return result / l 分块策略 将Q、K、V分别切分为块，每次只加载一小块到SRAM（GPU片上高速缓存）中计算。外层循环遍历K/V块，内层循环遍历Q块： for each K_j, V_j block: load K_j, V_j to SRAM for each Q_i block: load Q_i to SRAM compute S = Q_i @ K_j^T update online softmax statistics compute partial O = P @ V_j write O_i back to HBM 这种策略使得HBM访问量从O(n²)降低到O(n²d/M)，其中M是SRAM大小。对于典型配置，这相当于减少了10-20倍的显存访问。 Flash Attention V2：效率的进一步提升 V2在V1基础上做了几个关键优化： 减少非矩阵乘法运算 GPU上矩阵乘法（GEMM）的效率远高于其他运算（如rescale、softmax）。V2重新组织了计算顺序，将rescale操作推迟到最后，减少了中间的rescale次数。 并行化改进 V1主要沿序列维度并行化，V2增加了批次和头维度的并行化，使得在长序列场景下能更好地利用GPU的并行能力。 前向和后向的分配优化 V2重新分配了前向和后向传播中的工作负载，减少了线程块之间的同步开销。 Flash Attention V3：FP8与异步流水线 2024年发布的Flash Attention V3针对H100 GPU的FP8张量核心和异步执行特性进行了深度优化。 ...

Ring Attention：突破 GPU 显存墙 当上下文窗口从 4K 扩展到 1M Token 时，单 GPU 的 KV Cache 存储和注意力计算已经无法满足需求。Ring Attention（环形注意力）通过将注意力计算分布到多 GPU，实现了超长上下文的训练和推理。 一、长上下文的挑战 1.1 注意力计算复杂度 标准注意力的计算和显存需求都是 $O(n^2)$： 对于 $n = 1,048,576$（1M Token）： 注意力矩阵：$10^{12}$ 元素 = 8 TB（FP16） 单 GPU A100 80GB 显存：只能放下 0.01 的注意力矩阵 即使使用 Flash Attention 减少 HBM 访问，单 GPU 也无法存储完整的注意力矩阵。 1.2 KV Cache 显存需求 每个 Token 需要存储 Key 和 Value 向量。以 DeepSeek V4（MLA 压缩后）为例： $$\text{KV Cache/Token} = 512 \times 2 \times 2 \text{ bytes} = 2 \text{ KB}$$ ...

Flash Attention 3：让 GPU 跑满的注意力计算 Flash Attention 系列是近年来大模型工程领域最重要的优化之一。从 Flash Attention 1 到 3，每代都在逼近 GPU 硬件的理论极限。2026 年，Flash Attention 3 已经成为所有主流大模型推理和训练的标配。本文将深入解析其原理。 一、问题：标准注意力的内存瓶颈 1.1 GPU 内存层次 现代 GPU（如 H100）有复杂的内存层次： ┌─────────────────────────────────────────────┐ │ GPU 内存层次 │ ├─────────────────────────────────────────────┤ │ │ │ ┌─────────┐ 延迟: ~20 cycles │ │ │Register │ 带宽: ~30 TB/s │ │ │(SRAM) │ 容量: 256 KB/SM │ │ └─────────┘ │ │ ↑ │ │ ┌─────────┐ 延迟: ~200 cycles │ │ │L2 Cache │ 带宽: ~12 TB/s │ │ │ │ 容量: 50 MB │ │ └─────────┘ │ │ ↑ │ │ ┌─────────┐ 延迟: ~400+ cycles │ │ │ HBM │ 带宽: ~3.35 TB/s (H100) │ │ │(显存) │ 容量: 80 GB │ │ └─────────┘ │ │ │ └─────────────────────────────────────────────┘ 关键洞察：HBM 带宽只有 SRAM 的 1/9，但标准注意力几乎完全在 HBM 上操作。 ...

Flash Attention 系列回顾 Flash Attention 1：IO-Aware 的革命 标准注意力计算的问题是显存访问过多。标准实现： # 标准注意力（显存不友好） def standard_attention(Q, K, V): # Q, K, V: (batch, heads, seq, dim) # 需要 O(N^2) 显存访问来生成 attention matrix S = Q @ K.T # (N, N) 矩阵，需要全部存储在 HBM P = softmax(S) # 需要读取整个 S 矩阵 O = P @ V # 需要读取 V 和 P # 总 HBM 访问：O(N^2 * d) 只读，O(N^2) 写 return O HBM（High Bandwidth Memory）的带宽约 3 TB/s，而 SRAM（L2/L1 缓存）的带宽约 19 TB/s。Flash Attention 的核心洞察是：减少 HBM 访问，尽可能在 SRAM 中完成计算。 Flash Attention 1 通过 tiling 策略，将注意力分解为若干块，使每个块能放进 SRAM： ...

为什么需要 Flash Attention？ 标准 Attention 的瓶颈不在计算，而在显存 IO。考虑 $n = 8192$, $d = 128$ 的注意力计算： 操作 HBM 读写量 SRAM 计算 $S = QK^T$ 读 Q(8MB) + 读 K(8MB) + 写 S(512MB) ~16 GFLOPs $P = \text{softmax}(S)$ 读 S(512MB) + 写 P(512MB) ~8 MFLOPs $O = PV$ 读 P(512MB) + 读 V(8MB) + 写 O(8MB) ~16 GFLOPs 总计 ~1.6 GB ~32 GFLOPs A100 的 HBM 带宽 2 TB/s，SRAM 带宽 19 TB/s。计算只需 ~0.02ms，但 IO 需要 ~0.8ms——IO 是瓶颈的 40 倍。 ...