Page 218 - 《软件学报》2026年第1期
P. 218
高彦杰 等: 大语言模型预训练系统关键技术综述 215
GPU 内的高带宽内存 (HBM) 以及主内存. 模型通常存储在 HBM 中, 内核负责将模型和数据加载到 SRAM, 再进
一步传输到寄存器文件和流式多处理器中进行运算. 其中, GPU 具有较强的浮点计算能力, 而访存带宽相对受限.
通过 Roofline [65] 模型对常见模型算子 (例如通用矩阵乘 GEMM) 的分析表明, 内存访问往往成为性能瓶颈. 基于这
一观察, 内核设计的性能优化目标转向尽可能减少对 HBM 的访问, 并最大化 SRAM 的缓存利用率.
GPU
SRAM: 19.0 TB/s (20 MB)
SRAM
GPU HBM HBM: 1.5 TB/s (40 GB)
CPU DRAM DRAM: 12.8 GB/s (>512 GB)
图 8 内存层级 [66]
由于自注意力的时间和内存复杂度与序列长度的平方成正比, Transformer 在处理长序列时运行缓慢, 且需要
大量内存. FlashAttention [66,67] 指出, 当前内核的一个缺陷在于注意力算法缺乏 I/O 感知能力, 也未充分考虑 GPU
各级内存之间的读写开销. 因此, 作者提出了一种 I/O 感知的注意力算法, 通过引入在线 Softmax 算法和内核融合
技术, 减少中间结果的 HBM 读写次数; 同时采用平铺策略, 并在反向传播阶段通过重计算进一步降低 I/O 开销. 此
外, 作者还将该算法扩展应用于稀疏注意力机制.
Q、K、V ∈ R N×d 被存储在 HBM 中, 内核读取 QKV 进行注意力计算, 注意力输出
注意力机制的输入为
O ∈ R N×d 被写入 HBM. FlashAttention 性能优化的设计目标是减少 HBM 访问的数量 (降低到低于 N 的二次方程
(
度), 其思路是通过分块 (tiling), 在线 Softmax 等技术让计算不需要物化留存空间复杂度为 O N 2 ) 的中间结果 QK
矩阵, 进而降低空间复杂度和访存. 作者采用了两种已建立的技术分块重计算来克服在 N 的二次方复杂度的
HBM 访问开销. 主要思想是将输入 Q、K、V 分割成块, 从慢速的 HBM 加载到快速的 SRAM 共享内存, 然后对
于这些块计算注意力输出, 通过在将每个块的输出按正确的归一化因子缩放后相加, 最终得到正确的结果. 为了达
到按块计算注意力, 需要对 Softmax 做在线化处理. Softmax 将 K 列耦合在一起, 因此作者用缩放的方法对大的
Q,K,V ∈ R N×d , N 是序列长度, d 是注意力头的维度, 通过
Softmax 进行分解. 对于标准的注意力机制计算, 给定输入
以下注意力计算产生输出 O ∈ R N×d [66] :
T
S = QK , P = Softmax(S ) ∈ R N×N , O = PV ∈ R N×d .
x ∈ R 分解后的在线 Softmax 计算如下 [66] :
B
对于向量
[ ] ∑ f (x)
m(x) B maxx i , f (x) B e x 1 −m(x) ...e x B −m(x) , l(x) B f(x) i , Softmax(x) B .
i l(x)
i
(1) (2) T
[
(1)
(2)
B
2B
对于向量 x , x ∈ R , 可以将连接后的向量 x = x x ] ∈ R 的 Softmax 分解为:
([ ]) ( ( ) ( )) [ ( (1) ( m( x (2) ( )]
)
)
(1) (2)
m(x) = m x x = max m x (1) ,m x (2) , f (x) = e m( x (1) −m(x) f x e )) −m(x) f x (2) ,
([ ]) ( ) ( ) f (x)
)
)
(1) (2)
l(x) = l x x = e m( x (1) −m(x) l x (1) +e m( x (2) −m(x) l x (2) , Softmax(x) = .
l(x)
2
标准的注意力实现将矩阵 S 和 P 实例化到高带宽内存 (HBM), 这需要 O(N ) 的内存. 通常情况下, N ≫ d (例
如, 对于 GPT-2, N = 1024 且 d = 64). 由于一些或大多数操作受制于内存 (例如 Softmax), 大量的内存访问导致较
(m(x),l(x)), 通过以上的分解后的 Softmax 计算方法, 可以一次
慢的墙钟时间. 因此, 如果追踪一些额外的统计信息
一个块地计算 Softmax. 因此, 其将输入 Q、K、V 分割成块, 计算 Softmax 值以及额外的统计信息, 然后合并结果.
这样通过分块, 内核融合等策略, 可以尽可能让每块在 GPU 共享内存完成全部的矩阵乘, Softmax, 遮罩, Dropout
和矩阵乘计算, 进而减少了重复的 HBM 数据读写, 同时为了减少 S 和 P 等中间结果常驻内存, 同时应用策略释放
(
中间结果并在反向传播时重算输入张量. 另一种惰性 Softmax 也是采用类似的思想, 通过在线化计算不缓存 O N 2 )
的中间结果.
