AI 大模型之 FlashAttention\_01

原创

ITPAPA
ITPAPA

ITPAPA

2026年10月13日 12:41
江苏

在小说阅读器读本章

去阅读

在公众号小说中沉浸阅读

FlashAttention

IO感知精确注意力     技术深化学习

范围：基本概念、技术演化、数学及计算机制、模块架构、主要配置参数、训练与推理流程、工程实现、性能评测、典型应用、风险与2026年最新公开进展。

1  技术定位与解决的问题

FlashAttention不是一种修改模型输出语义的近似注意力，而是IO-aware精确计算方法；用分块、在线softmax与重计算减少HBM与片上SRAM之间的访存。

问题背景：大规模模型的结构、计算、显存、训练稳定性与生成能力之间存在不同的取舍。本技术对应的改进点，应通过与恰当的基线比较才能判断。它属于“推理/运行时机制”，不应把单篇论文的实验数字误解为通用硬件指标。

2  结构组成与数据流

图1  FlashAttention 核心计算/数据流示意

•  步骤1：分块读取Q、K、V。

•  步骤2：SRAM中计算局部QKᵀ。

•  步骤3：在线维护softmax最大值与归一化和。

•  步骤4：分块累加加权V。

•  步骤5：输出并在反向传播时重计算。

3  底层原理与关键技术

3.1  机制设计与公式

普通实现显式物化N×N注意力矩阵，显存随序列平方增长；FlashAttention以分块状态规避这一物化。

3.2  复杂度与状态开销

核心在线Softmax状态包含running max、running denominator与归一化输出累计。

3.3  硬件执行与精度

算术工作量仍主要是QK与PV乘法，并不意味着注意力理论FLOPs由二次变线性。

架构识别原则：需要区分“数学形式”“训练目标”“GPU计算内核”“模型实例配置”“在线服务系统”。同名概念在这些层面的实现方式并不自动相同；评测必须说明层次、模型版本、数据精度、硬件与任务。

4  核心参数及工程含义

参数取值原则：除明确标注“论文示例”的值之外，本表列的是配置项与设计量，而非默认推荐值。实际有效范围受模型结构、训练方式、推理框架、显卡架构与可用显存约束。

4.1  定量公式与参数估算

在线Softmax状态更新：m\_new=max(m\_old, max(scores)); l\_new=exp(m\_old−m\_new)l\_old+sum(exp(scores−m\_new))。输出累计量以同一新max重标定。

4.2  可复核的示例计算

示例：序列4096时，普通显式单头Attention分数矩阵有约1,678万元素；FP16约32 MiB/头，仅为QK分数矩阵，不含梯度和其他中间值。FlashAttention避免完整物化。

4.3  核心伪代码与执行序列

for Q\_tile: initialize m,l,o; for KV\_tile: matmul→update(m,l,o); write normalized o。

4.4  易出错的工程边界

检查在线softmax与不稳定exp溢出的区别；kernel计时要单独区分prefill和decode。

5  实施流程与关键技术细节

5.1  训练准备

先用框架原生scaled\_dot\_product\_attention建立数值基线。

5.2  实现验证

启用flash backend，保证支持的head\_dim、mask和数据类型。

5.3  测试优化

训练同时评估forward/backward精度，检查dropout与mask路径。

5.4  交付评估

用nsys或PyTorch profiler比较HBM带宽、核启动数量、端到端吞吐。

6  技术验证与量化评估设计

图2  机制—参数—实现—应用的验证路径

实验控制：固定模型权重/训练Token预算、Tokenization、输入与输出长度分布、温度与采样配置、数据类型、GPU与软件版本；每组记录至少一次预热、多轮重复以及波动区间。结果按实际运行配置报告，避免跨研究条件直接比倍率。

7  典型应用场景

•  Transformer预训练及微调：先建立真实业务负载与基线，再决定是否引入该结构。

•  长上下文与高分辨率视觉模型：先建立真实业务负载与基线，再决定是否引入该结构。

•  带宽瓶颈明显的注意力计算：先建立真实业务负载与基线，再决定是否引入该结构。

8  局限 代价与容易误解的地方

•  不同GPU、张量形状下收益可能差别很大。

•  高阶内核对dtype、mask、head\_dim与显卡架构有限制。

•  数值舍入顺序不同，不能要求逐bit一致。

9  技术进展 截至2026年10月

FlashAttention-4于2026年发表，针对Blackwell GPU进一步协同优化kernel与硬件；属于FlashAttention系列后继，不应把FA4性能数据回填到原始FA1。

10  推荐实验记录模板

11  原始论文与权威资料

•  [1] https://arxiv.org/abs/2205.14135

•  [2] https://arxiv.org/abs/2603.05451

说明：所有实验数值应以所引文献的相应版本、配置与对照组为准。本文的示意图为重新绘制的知识结构图，不能作为作者官方架构图。

技术深化 在线Softmax推导与反向重计算

本章在原有概览基础上展开推导与实现。文中的尺寸和计算例为教学配置，实验代码验证明确标注的简化核心；论文方法、工程建议和生产实测分别说明。

瓦片状态如何保持精确语义

对一个查询行，设已处理键块的最大分数为m，指数和为l，未归一化的值累计为u。新块分数为s，值为V；更新最大值后，旧累计必须乘同一重标定系数，否则不同块指数处于不同尺度。最终输出为u/l。[1]

初始化m为负无穷、l与u为0。遮罩应在求最大值前施加；完全被遮蔽的一行会出现负无穷减负无穷，应由实现明确输出规则并避免NaN。FlashAttention的精确性指注意力数学运算保持一致，浮点累加次序改变后不要求逐位相等。

IO复杂度和算术复杂度要分开

普通分步实现可能把S×S分数与概率写入HBM再读取。分块实现把瓦片与行统计留在片上，将完整中间矩阵的读写替换为小状态读写。固定头维下，密集注意力前向仍有约4S²d的矩阵乘FLOPs；因果三角遮罩实际有效配对约减半。前向持久化额外统计可按每查询行一个logsumexp估计，不代表整个训练显存只为线性注意力状态。

Q瓦片扩大可复用更多K/V，但增加寄存器与共享内存占用；K瓦片扩大减少循环次数，也可能使占用率下降。必须把计算形状和设备资源一起考虑，而不是无限增大tile。

反向传播为什么可以重算

前向保存O与行logsumexp，反向按块重新计算概率P。对上游梯度dO，先算dV与dP，再用softmax雅可比的结构得到dS。它免去保存完整P，代价为额外计算。[1]

随后dQ由dS乘K得到，dK由dS转置乘Q得到，并带上缩放因子。dropout路径需要一致重现随机掩码；只保留种子但改变线程到元素的随机数映射，也可能让反向使用不同的掩码。

可复核实验

03\_online\_softmax.py按键序列分块实现单查询在线累计，与一次性FP64 softmax比较。实验支持非整除块长，故可检查最后一个瓦片。

m, l, u = -np.inf, 0.0, np.zeros(V.shape[1])

for start in range(0, len(K), tile):

s = K[start:start+tile] @ q / np.sqrt(len(q))

mn = max(m, s.max())

p = np.exp(s - mn)

a = np.exp(m - mn)

u = a \* u + p @ V[start:start+tile]

l = a \* l + p.sum()

m = mn

result = u / l

把分数放大至可能使直接exp溢出的范围，稳定参考实现和在线实现应保持有限值。再对多查询、因果遮罩与梯度做独立对照。这里的NumPy循环用于理解公式，不能用其CPU时间评价GPU FlashAttention。

性能验证方法

先分别固定Sq=Sk进行预填充、Sq=1进行解码。报告数据类型、头维、遮罩、dropout、前向或前向反向。GPU计时必须预热并同步，用CUDA事件或profiler，检查是否实际调用目标内核。使用框架自动选择后端时，还应保存后端分派日志。

常见故障的定位顺序

输出整体比例异常通常来自忘记最后除l或重复缩放；块边界误差来自遗漏旧u重标定；极值NaN来自掩码和指数处理；训练损失异常而推理正常，应重点检查dropout随机数和反向概率重建。算子变快而训练吞吐不变时，测量其他层与通信占比，再用Amdahl定律估算可获得的端到端收益。

图3 高亮示意两个先后处理的键块 完整分数矩阵只作解释不物化

本册配套实验验证记录

配套文件为 实验代码/03\_online\_softmax.py。本次使用NumPy CPU固定随机输入运行，执行通过。结果验证本章所标注的公式、索引、容量或事务边界，不代表完整模型训练效果或GPU性能实测。

tile 1 max\_abs\_error 0.0

tile 8 max\_abs\_error 0.0

tile 16 max\_abs\_error 0.0

tile 37 max\_abs\_error 0.0

工程实操 极大分数下的在线统计与无效行处理

固定输入与检查目标

单查询的三个分数依次为1000、1001、999，对应标量值1、3、7。每次只处理一个键，相当于键瓦片长度1。直接对这些分数求指数会超出常用浮点表示范围。

逐步计算与预期结果

1  首块处理1000后，统计为m=1000、l=1、u=1；分数的共同尺度由m记录。

2  第二块将m提高到1001，旧统计乘exp(-1)。于是l约1.367879、u约3.367879。

3  第三块不改变m，新权重为exp(-2)。最终l约1.503215、u约4.315226，输出u/l约2.870665。

4  若一行全部不可见，必须先检测无效行，再按接口约定返回规定值或报错。本教学检查选择拒绝无效行，不执行负无穷减负无穷。

接口与实现约束

前向统计的shape应只覆盖有效查询行，并保留能支持反向重建的信息。多查询实现不能在不同行之间共用m和l；ragged布局还要区分逻辑行号与packed行号。

梯度验证先固定无dropout参考，再对Q、K、V分别做数值或自动微分对照。若开启dropout，应把随机数映射视为前后向契约的一部分。错误处理与正式内核保持一致后，才比较性能。

故障核查表

扩展验证组合

瓦片长度取1、2、3、8，包含长度不足一个块和最后块不满；再测试全mask行、单可见键及重复最大值。

配套代码与结果

进阶实验/boundary\_cases.py 中的 c03 函数对应本节。可在进阶实验目录运行 python boundary\_cases.py，或执行 python -c \"from boundary\_cases import c03; print(c03())\" 单独检查。本节固定样例执行通过，结果保存在边界实验结果.json。

验证范围为正文明确标注的简化机制与错误反例。工程接口、扩展测试和故障注入属于本资料提出的实施建议；完整模型训练、GPU数值与吞吐、跨进程恢复仍须在目标环境验证。

训练与部署 注意力反向的完整链式关系

目标与验证范围

设O=PV，P为带mask分数的softmax，J为0.5倍O平方和。本例查询2行、键3行、头维2，第一行禁止读第三个键，第二行可读全部键。

解析关系与梯度路径

上游梯度U=O。先求dV=P转置U与dP=UV转置，再通过softmax结构求dS，最后dQ=dSK/sqrt(d)、dK=dS转置Q/sqrt(d)。被mask的概率为0，对应分数梯度也应为0。

随包对Q、K、V逐元素中心差分。若只验证dQ，错误的dK或dV归约仍可能进入训练。分数梯度逐行和应接近0，这来自softmax对共同分数平移的不变性。

训练配置与误用边界

FlashAttention重算概率时需要相同mask、行统计、缩放与随机dropout映射。前向与反向若使用不同有效长度，前向看似正确但参数更新会偏离参考。

混合精度实验先检查有限性，再设绝对与相对容差。梯度接近0时以绝对误差为主；梯度较大时同时报告相对误差与梯度范数。

部署验收与恢复

推理不需要反向保存的数据，训练则要计算激活、行统计和重算开销。部署显存和训练峰值不能直接套同一个缓存预算。

先通过小形状梯度测试，再扩展训练前向反向总时间。kernel前向变快，不自动代表优化器步吞吐提升。

配套检查与本次数值结果

训练部署实验/training\_checks.py 中的 t03 函数已执行通过。本例解析梯度与中心差分的最大绝对误差为 4.074e-11。完整输入、解析实现和实际输出见代码及训练与状态验证结果.json。

可在训练部署实验目录执行 python training\_checks.py，或执行 python -c \"from training\_checks import t03; print(t03())\"。所示结果来自NumPy CPU教学验证；目标模型的自动微分、低精度和分布式执行需另做对应验收。