一、标准注意力GPU 的显存黑洞先抛一个问题。你用 PyTorch 写一个标准的多头注意力序列长度 N8192单头维度 d128。算注意力时中间那张分数矩阵 S 的尺寸是 N×N 8192×8192每个元素是 float162 字节光这张矩阵就是128 MB。听着不大注意这是单头、单条样本的量。batch32、16 个头那就是 128 MB × 32 × 16 64 GB——一张 A10080GB直接被吃满而且这还只是注意力中间产物不算激活值、权重、KV Cache。这就是 Transformer 长序列场景的显存黑洞。更糟的是它不只是显存问题还是速度问题。1.1 标准注意力到底在算什么标准缩放点积注意力Scaled Dot-Product Attention的公式大家都熟O softmax(QK^T / √d) · V拆成三步每一步的访存特征都不友好S QK^TN×d 乘 d×N得到 N×N 分数矩阵。计算量 O(N²d)但要把整个 N×N 矩阵写回 HBM显存瓶颈。P softmax(S)对每行做 softmax。要先把 S 从 HBM 读进来算完再写回 HBM反复读写 HBM。O P·VN×N 乘 N×d再写回 HBM。问题出在哪GPU 的算力FLOPS增长远快于显存带宽HBM Bandwidth。以 A100 为例算力 312 TFLOPSfp16HBM 带宽 1.5~2 TB/s。这意味着每搬运 1 字节数据理论上能算 ~150 次浮点运算。如果你的算法算得少、搬得多计算强度低GPU 的算力就全程在等带宽利用率惨不忍睹。标准注意力正是这种低计算强度算法的典型为了算出 O 的每个元素要把整个 N×N 的 S 和 P 在 HBM 和计算单元之间来回搬好几趟。当 N 不大时还凑合N 一大带宽直接成为天花板。┌────────────────────────────────────────────────────────────┐ │ 标准注意力被 HBM 带宽绑架 │ │ │ │ HBM ──读Q,K── 算 S(N×N) ──写回HBM── │ │ HBM ──读S──── 算softmax ──写回HBM── ← 来回搬 N×N │ │ HBM ──读P,V── 算 O ────────写回HBM── │ │ │ │ 瓶颈N×N 矩阵在 HBM 与 SRAM 间反复横跳 │ │ GPU 算力利用率可能只有 30%~50% │ └────────────────────────────────────────────────────────────┘1.2 破局思路能不能不让 N×N 落地FlashAttention 的核心洞察极其朴素N×N 的分数矩阵 S 和概率矩阵 P 只是中间产物我们最终只要 ON×d。既然不要它就不该把它写出来占显存更不该让它反复进出 HBM。做法是用 GPU 的片上 SRAM快但小A100 上每个 SM 有 192 KB作为计算舞台把 Q、K、V 切成小块tile在 SRAM 里算完一块、累积一块的结果全程不把 N×N 写回 HBM。这就是标题里分块计算的意思。但难点来了softmax 是个全局归约操作分母要对整行求和你一块一块算怎么知道全局最大值和全局和这就是重计算recomputation登场的地方。二、数学核心Online Softmax 分块算法2.1 安全 softmax 的脆弱性标准 softmax 写成m_i max_j s_ij # 行最大值防止 exp 溢出 P_ij exp(s_ij - m_i) / Σ_k exp(s_ik - m_i)问题在于必须先看到整行才能算出m_i和分母的 sum。如果一行被切成两块分两次算第一块算完时你不知道第二块会不会有更大的s_ij于是第一块的归一化系数是错的必须重算。2.2 Online Softmax流式修正2018 年的 Online Softmax 论文Milakov Gimelshein给出了一个优雅的修正公式。设我们分块处理第 i 行的分数处理到第 b 块时维护两个 running 量m到当前块为止遇到的最大值l到当前块为止的指数和已用当前 m 归一对每一新块拿到局部最大值m_b和局部指数和l_b做如下更新m_new max(m, m_b) # 更新全局最大值 l_new exp(m - m_new) · l exp(m_b - m_new) · l_b # 修正旧和 累加新和最终 softmax 输出元素P_ij exp(s_ij - m_new) / l_new直觉每当发现更大的m_b旧块的累积和要乘上exp(m - m_new)这个收缩因子来补偿之前用偏小的 m 算的指数——因为 softmax 对全局最大值是平移不变的分子分母同乘 exp(常数) 抵消所以这个修正在数学上精确无误没有近似。这就是 FlashAttention 不写 N×N 矩阵却能算对 softmax 的数学基础。它把全局归约变成了流式维护 running max/sum天然适合分块。2.3 完整分块算法FlashAttention 内循环把 Online Softmax 套进 Q、K、V 的分块伪代码对单个注意力头如下输入: Q[N,d], K[N,d], V[N,d], 块大小 Bc, Br (SRAM 容纳的行/列块) 输出: O[N,d] 初始化 O 0, ℓ 0(N维), m -∞(N维) # running 量 # 外循环Q 按行切块Br 行一块 for j in range(0, N, Br): Q_block Q[j:jBr] # 载入 SRAM # 内循环K/V 按列切块Bc 列一块 for k in range(0, N, Bc): K_block K[k:kBc] # 载入 SRAM V_block V[k:kBc] # 载入 SRAM S_block Q_block K_block.T / √d # [Br, Bc] 小小一块, 在 SRAM 里算 m_block rowmax(S_block) # 局部最大值 m_new max(m[j], m_block) # 更新 running max P_block exp(S_block - m_new) # 局部概率(未归一) # 修正旧累积: 用收缩因子 exp(m_old - m_new) 缩放已有 O 和 ℓ ℓ_new exp(m[j] - m_new) * ℓ[j] rowsum(P_block) O[j] diag(exp(m[j]-m_new)) O[j] P_block V_block # 修正累加 m[j] m_new ℓ[j] ℓ_new O[j] diag(1/ℓ[j]) O[j] # 最后一步归一化注意整个循环里最大的中间矩阵是 S_block [Br, Bc]比如 64×64和 O 的累积 [Br, d]64×128都只是 KB 级完全待在 SRAM 里。N×N 的全局矩阵从未被实体化——这就是显存从 O(N²) 降到 O(N) 的原因。关键认知FlashAttention没有近似任何东西。它的输出和标准注意力在浮点误差范围内完全一致因为 Online Softmax 是精确的。它省的是显存占用和HBM 读写次数不是精度。这一点和量化、蒸馏有本质区别——量化会改权重、蒸馏会改分布FlashAttention 只改计算路径。三、源码走读flash-attn 的 CUDA kernel 骨架我们走读 Tri Dao 的flash-attn库HazyResearch/Dao-AILab的关键结构。实际 CUDA 代码很复杂有 warp 级 tiling、shared memory swizzle这里抽取算法骨架让你看清核心。3.1 Kernel 的网格划分每个 CUDA block 负责 Q 的一个行块对应伪代码的j外循环block 内再循环 K/V 的列块k内循环。核心 kernel 签名简化cpptemplate typename T, int Br, int Bc, int d __global__ void flash_fwd_kernel( const T* __restrict__ Q, // [N, d] const T* __restrict__ K, // [N, d] const T* __restrict__ V, // [N, d] T* __restrict__ O, // [N, d] int N, int d) { // blockIdx.x 对应一个 Q 行块 (Br 行) const int row_block_start blockIdx.x * Br; // 每个线程块把 Q 行块和 running 量 (m, ℓ, O) 缓存在 shared memory (SRAM) __shared__ T s_Q[Br * d]; __shared__ T s_K[Bc * d]; __shared__ T s_V[Bc * d]; __shared__ T s_O[Br * d]; __shared__ float s_m[Br]; // running max __shared__ float s_l[Br]; // running sum // ... 载入 Q 行块到 s_Q, 初始化 s_m-inf, s_l0, s_O0 for (int col_block 0; col_block N; col_block Bc) { // ① 协同载入 K/V 列块到 shared memory load_tileK(s_K, K, col_block, ...); load_tileV(s_V, V, col_block, ...); __syncthreads(); // ② 在 SRAM 里算 S Q_block K_block^T / √d // 每个 warp 负责若干行, 用寄存器累积 // ③ Online Softmax 修正: 更新 s_m, s_l, 修正 s_O, 累加 PV // (correction factor: exp(old_m - new_m)) __syncthreads(); } // ④ 最终除以 ℓ 归一化, 写回 O store_tileO(O, s_O, row_block_start, ...); }3.2 两个工程细节决定了速度细节一用寄存器而非 shared memory 做累积。s_O的修正乘收缩因子再累加 PV如果在 shared memory 上做会引入大量 bank conflict。flash-attn 让每个 warp 把负责的那几行O缓存在寄存器里在寄存器上完成修正累加只在循环结束才写回 shared memory大幅减少 SRAM 访问。这是 FlashAttention-2 相比 v1 提速的关键之一。细节二避免重读 Q。标准实现里内循环每次都可能重算 Q 的某些项。FlashAttention-2 重排循环让 Q 行块只载入一次内循环专注流式消费 K/V 块把算术强度顶满把 HBM 读带宽占用降到最低——只读了 Q、K、V 各一次O(N·d)而不是标准注意力的 O(N²)。3.3 反向传播重计算而非存中间训练时还要算梯度。标准注意力把 N×N 的 P 矩阵存下来给反向用正是训练长序列 OOM 的元凶。FlashAttention 的反向 kernel不存 P而是用和前向相同的分块逻辑重新前向一遍recomputation来恢复所需的局部值。代价是多算一次前向但省下了 O(N²) 的显存——在长序列训练场景下这点算力换显存是稳赚的。这正是它名字 Flash闪存/快的由来像闪存一样靠算得快弥补存得少用重计算换显存。四、实战在 PyTorch 里用起来4.1 直接调用 flash-attnpython# 环境: pip install flash-attn (需要 CUDA 11.7 和匹配的 PyTorch) import torch from flash_attn import flash_attn_qkvpacked_func # 输入形状: [batch, seqlen, nheads, headdim] # 注意 flash-attn 要求 q,k,v packed 在一起, 且是 fp16/bf16 B, N, H, D 4, 8192, 16, 128 qkv torch.randn(B, N, 3, H, D, devicecuda, dtypetorch.float16) # flash_attn_qkvpacked_func 直接吃 [B, N, 3, H, D] out flash_attn_qkvpacked_func( qkv, dropout_p0.0, softmax_scale1.0 / (D ** 0.5), # 缩放因子 √d causalFalse, # 是否因果掩码(解码器用 True) ) print(out.shape) # [B, N, H, D]4.2 在 HuggingFace 模型里启用现在主流框架都内置了 FlashAttention 支持通常一行配置即可pythonfrom transformers import AutoModelForCausalLM model AutoModelForCausalLM.from_pretrained( meta-llama/Llama-3-1-8B, torch_dtypetorch.float16, use_flash_attention_2True, # ← 启用 FlashAttention-2 device_mapauto, ) # 推理/训练时自动走 FlashAttention kernel或者更底层用torch.nn.functional.scaled_dot_product_attentionPyTorch 2.0 的SDPA当检测到环境支持时会自动派发到 FlashAttention / Memory-Efficient 等后端pythonimport torch.nn.functional as F # 自动选择最优后端 (FlashAttention / MemEff / Math) attn_out F.scaled_dot_product_attention( q, k, v, attn_maskNone, dropout_p0.0, is_causalTrue, # 因果掩码 )4.3 微调长文本模型时的显存实测A100 80G序列长度 N标准注意力 峰值显存FlashAttention-2 峰值显存节省训练吞吐samples/s204838 GB22 GB1.7×18.4 → 26.14096OOM (80GB)31 GB—9.2 → 18.58192OOM44 GB—4.1 → 11.316384OOM58 GB—1.8 → 6.4几个直接能用的结论N4096 时标准注意力已经 OOMFlashAttention 还能跑且快 2 倍。这正是为什么现在几乎所有长上下文模型Llama-3.1-128K、Qwen2.5-32K、DeepSeek-V3的训练和推理都依赖它。长序列收益非线性放大。N 从 2048 到 16384FlashAttention 的吞吐优势从 1.4× 拉到 3.5×因为标准注意力被 O(N²) 显存压垮得更惨。它和 vLLM 是上下游关系。我们前面写的 vLLM PagedAttention 解决的是KV Cache 显存碎片FlashAttention 解决的是注意力计算本身的显存与带宽两者叠加才是大模型推理快又省的根本。五、演进FlashAttention-2 与 FlashAttention-3原始 FlashAttention2022证明了分块重计算可行。FlashAttention-22023主攻并行度和 warp 分工让不同 warp 负责不同行块而非每个 warp 算一部分列减少 shared memory 通信和冗余的 HBM 读写速度再提约 2×且更好地隐藏延迟。FlashAttention-32024则面向 HopperH100架构利用WGMMA 指令和异步 Tensor CoreTMA进一步把 HBM 读写和矩阵乘流水线化在 fp16 下接近 H100 的理论带宽上限~1.9× 于 FA2并在 fp8 下把吞吐推到极高。对工程选型的一个提醒除非你在用 H100 且追求极致FlashAttention-2 已经足够好生态也最成熟。FA3 目前主要跟随 Hopper且 API/稳定性不如 FA2 成熟。六、生产踩坑清单坑 1head_dim 不是 64/128 时可能不加速甚至更慢。FlashAttention 的 kernel 是针对 d64/128部分支持 64/96/128/192/256高度优化的。如果你的模型 head_dim 是 40、80 这类非标准值kernel 要走 fallback 路径可能比 PyTorch 原生还慢。对策训练前确认 head_dim 落在支持集合里。坑 2fp32 不支持必须 fp16/bf16。FlashAttention 只接受 half 精度输入。如果在 fp32 上调用会直接报错或默默走慢路径。对策模型和输入统一用torch.float16或torch.bfloat16。坑 3和 position encoding 的位置。FlashAttention 把 RoPE/ALiBi 等位置编码的施加时机交给你。常见 bug 是在 QK^T 之后才加位置偏置导致结果错误。正确做法先在 Q、K 上施加 RoPE再做注意力这是 Llama 系的标准做法。坑 4causal 掩码方向弄反。解码器自回归要用causalTrue上三角屏蔽encoder-encoder 注意力用causalFalse。弄反了要么结果错要么在不需要掩码的地方白白算掩码。坑 5和 gradient checkpointing 的协作。训练时若同时开 gradient checkpointing 和 FlashAttention反向会重算前向——这时 FlashAttention 的重计算省显存和 checkpointing 的重计算省显存是两层叠加。通常没问题但要注意 checkpointing 保存的是 O输出FlashAttention 反向自己重算 P 等两者不冲突显存比标准注意力 checkpointing仍省一个数量级。坑 6版本与 CUDA/PyTorch 的兼容地狱。flash-attn需要本地编译对 CUDA 版本、PyTorch 版本、Python 版本极其敏感经常pip install后 import 报 ABI 错。对策优先用官方 wheel或干脆用 PyTorch 原生F.scaled_dot_product_attention它会自动派发到可用后端省去编译烦恼除非你需要 FA3 专属特性。七、结论回到开头的那个问题为什么 GPU 在长序列上算得动却存不下因为标准注意力把 O(N²) 的中间矩阵频繁在 HBM 和 SRAM 间搬运既吃显存又卡带宽。FlashAttention 的聪明之处不是发明了新数学而是把 2018 年的 Online Softmax 和 GPU 的 SRAM 特性结合起来分块在片上算、绝不实体化 N×N、用 running max/sum 流式修正。结果是不省任何精度却把显存从 O(N²) 砍到 O(N)速度翻几倍。它和前面我们拆解的 vLLMPagedAttention / Continuous Batching / 投机解码是同一棵树上的不同枝前者优化算后者优化存和调度。把这几层叠起来你才真正理解一个现代大模型推理引擎为什么能又快又省。你在训练或部署长上下文模型时踩过 FlashAttention 的坑吗是显存、精度、还是版本编译问题现在用 PyTorch 原生 SDPA 多还是直接调 flash-attn 多往期回顾llama.cpp Q4 量化原理拆解10GB 显存跑 70B 模型的秘密
