AI 大模型之 FlashAttention\_02

原创

ITPAPA
ITPAPA

ITPAPA

2026年10月13日 12:43
江苏

在小说阅读器读本章

去阅读

在公众号小说中沉浸阅读

FlashAttention-2

更高并行度的注意力内核     技术深化学习

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

1  技术定位与解决的问题

FlashAttention-2在保留精确注意力与低HBM访问的基础上，改进线程块及warp间的工作划分，减少非矩阵乘运算，并让单个注意力头能够跨线程块并行。

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

2  结构组成与数据流

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

•  步骤1：Tile化Q/K/V数据。

•  步骤2：重新划分thread block并行度。

•  步骤3：降低非matmul运算开销。

•  步骤4：warp内重排局部计算。

•  步骤5：在线Softmax并输出。

3  底层原理与关键技术

3.1  机制设计与公式

与FA1同属精确注意力计算，但通过划分计算改善并行度与SM利用率。

3.2  复杂度与状态开销

降低softmax重标定等非matmul FLOPs，减轻计算路径瓶颈。

3.3  硬件执行与精度

不同shape的最佳block配置不同，kernel autotuning需注意可复现。

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

4  核心参数及工程含义

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

4.1  定量公式与参数估算

FA2仍遵从精确Softmax注意力 O(S²d) 算术复杂度；主要改变GPU工作分配、非GEMM指令与HBM读写。

4.2  可复核的示例计算

论文在A100上报告较FA1近2倍核心算子加速，并达到约50–73%理论峰值算力；这是指定任务/内核的结果。

4.3  核心伪代码与执行序列

增加沿序列方向的CTA并行→降低跨warp共享通信→重排行列分块→在线Softmax。

4.4  易出错的工程边界

检查thread block并行导致的寄存器压力、尾部瓦片浪费和不同head\_dim的回退实现。

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

5.1  训练准备

以FA1/torch SDPA为性能与误差对照。

5.2  实现验证

测试单头、小batch、长序列等容易低占用的形状。

5.3  测试优化

验证反向传播、中间梯度及因果mask。

5.4  交付评估

对GPU利用率、kernel占用率、整体训练tokens/s分别统计。

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

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

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

7  典型应用场景

•  大模型训练加速：先建立真实业务负载与基线，再决定是否引入该结构。

•  长序列注意力计算：先建立真实业务负载与基线，再决定是否引入该结构。

•  通用GPU推理框架的基础kernel：先建立真实业务负载与基线，再决定是否引入该结构。

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

•  论文中的A100理论峰值百分比不可跨卡直接照搬。

•  调用框架可能自动切换其他attention backend。

•  性能仍受模型其他层、内存和通信瓶颈制约。

9  技术进展 截至2026年10月

FlashAttention-3（2024）和FlashAttention-4（2026）分别着重Hopper异步流水化与Blackwell硬件特征；FA2仍是理解系列kernel设计的重要分界。

10  推荐实验记录模板

11  原始论文与权威资料

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

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

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

技术深化 线程块工作划分与GPU性能分析

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

从头级并行转向序列级并行

FA2仍计算同一密集注意力，重点在于让同一头的多个查询瓦片由不同CTA处理。若批量B=1、头数H=8、查询长度8192、查询瓦片128，CTA数量可达512；若只按批量和头分配，仅有8个CTA，许多SM会空闲。该数字是教学调度估算，实际内核还受方向、tile和调度策略影响。[1]

每个CTA独立完成自己查询行对应的完整键遍历，不需要多个CTA对同一输出行做一般性归约。这种沿查询方向的切分不同于split-K；后者会产生同一行的部分softmax状态，需要额外合并。

warp之间怎样减少通信

FA2前向采用sliced-Q思路，让不同warp处理不同查询行，共享K/V瓦片。相比多个warp沿K维产生同一输出的部分和，独立查询行更容易避免共享内存归约与同步。训练反向的dK/dV又需要跨查询累加，不能直接沿用前向无归约的结论。[1]

在线softmax中保留未归一化输出，把除法推迟到最终写回，可降低非矩阵乘指令。Tensor Core矩阵乘吞吐与标量exp、归约、除法吞吐不同，FLOPs相同并不表示硬件耗时相同。

资源预算模型

共享内存主要用于K/V瓦片和流水线缓冲，量级随BN、d和缓冲阶段数增加。寄存器承载查询、行统计和输出累计，量级随BM与d增加。tile变大增加数据复用，但每个SM能驻留的CTA可能减少，寄存器溢出还会产生local memory读写。

若注意力原先占训练时间40%，内核加速2倍，端到端上限约1.25倍，且假设其余部分不变。真实重叠通信和流水线可能改变f，应由profiler测量而不是猜测。

实验代码与计时纪律

04\_schedule.py只打印理论CTA数，用于解释占用率变化，不执行FA2。GPU实验需调用固定版本官方实现，保存依赖版本与真实内核名称。

for batch, heads, length in [(1, 8, 8192), (8, 32, 2048)]:

for tile\_q in (64, 128, 256):

ctas = batch \ heads \ ((length + tile\_q - 1) // tile\_q)

print(batch, heads, length, tile\_q, ctas)

性能网格包含头维64、128、256，序列512、2048、8192，批量1、8。先比较前向，再比较训练前向反向总时间；非整除长度用于检查尾瓦片。每组完成预热，记录多次GPU事件时间的中位数与尾部分布，并在测量窗口外检查数值误差。

如何读profiler结果

SM利用率低且CTA少，优先检查查询瓦片并行；CTA多而occupancy低，检查寄存器与共享内存驻留限制；大量barrier或共享内存流量，检查warp间归约；Tensor Core活跃但端到端仍慢，检查数据准备、MLP和通信。不能仅凭GPU总利用率判断FA2是否工作，高利用率也可能来自其他算子。

正确性与部署边界

causal的方形与非方形对齐必须按实现文档验证。底层kernel支持的mask与头维不能由SDPA接口名称推断。训练时还需比较dQ、dK、dV，极小分母与极大分数单独测试。若设备、编译器或框架升级，重新确认后端选择与性能，保留原结果用于回归，而不是把论文A100倍率作为服务验收值。

本册配套实验验证记录

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

batch\_heads\_seq\_tile\_ctas 1 8 8192 64 1024

batch\_heads\_seq\_tile\_ctas 1 8 8192 128 512

batch\_heads\_seq\_tile\_ctas 1 8 8192 256 256

batch\_heads\_seq\_tile\_ctas 8 32 2048 64 8192

工程实操 沿键切分的归约与沿查询切分的区别

固定输入与检查目标

分数为[1,2,4,0]，值为[2,5,1,9]。把四个键切成两块，分别保存局部m、l、u。该例解释split-K需要的归约，并不把它当成FA2常规前向的工作划分。

逐步计算与预期结果

1  第一块m=2，l约1.367879，u约5.735759；第二块m=4，l约1.018316，u约1.164841。

2  合并时取全局m=4，第一块的l和u都乘exp(2-4)，第二块保持原尺度，再求两块累计之和。

3  归一化后的正确输出约1.612955。如果先各块归一化，再直接平均两个输出，会得到约2.668533。

4  沿查询切分时各工作组处理不同输出行，通常没有本例这种同一行的部分状态归约；因此不能混用两种成本估算。

接口与实现约束

分割方案应写明切分轴、每组输出责任、局部统计布局和合并算子。改变split数量时必须保持输出语义一致，尾部键数不足时仍使用真实有效长度。

性能剖析把局部计算和归约分别标记。若split-K提高并行度，却增加状态写回与二次启动，短序列可能没有收益；执行策略应依据实际形状选择，而不是只依据设备SM数量。

故障核查表

扩展验证组合

同一输入分别切成1、2、3、4组，包含空尾块处理规则；保持dtype与遮罩相同，检查输出后再分析归约开销。

配套代码与结果

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

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

训练与部署 切分与归约的反向一致性

目标与验证范围

沿键切分的局部softmax状态合并，应同时保持前向函数和梯度。使用分数[1,2,4,0]、值[2,5,1,9]，目标J为归一化后的标量输出。

解析关系与梯度路径

合并后的softmax仍覆盖全部键。分数导数依赖全局输出y；只在各局部块内使用自己的输出，会遗漏块间归一化关系。

中心差分得到分数梯度约[0.016012,0.380898,-0.509337,0.112427]，总和接近0。改变切分数量应保持同一函数及梯度，不能仅对最后输出做宽松对照。

训练配置与误用边界

训练反向可能需要跨查询聚合dK/dV，其归约方向与前向split不同。调度文档应分别写明前向分片、反向责任区和累加精度。

原子累加会改变浮点顺序；可用确定性参考验证数学语义，再对生产归约设合理容差。容差不能用来容忍遗漏完整分支。

部署验收与恢复

性能记录分成局部计算、状态合并、反向归约和工作区，才能解释split带来的并行收益与额外成本。

本例分数不同且远离最大值切换点，适合中心差分；极值相等和全mask行另做实现边界测试。

配套检查与本次数值结果

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

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