第 5 章 · 约 24 分钟
算子与运行时
如何组织执行,少搬、少提交、少等待?
阅读中·实验未完成·考核未交
执行
提交、搬移、同步是三种不同的税
同一条数学公式,可以有完全不同的数据访问。主机把算子提交给加速器,数据在全局内存、片上存储和寄存器之间搬移,网格结束时往往还要同步。把 softmax 写成三个独立算子,就会把整块注意力分数写回 HBM 再读出来——这是在交临时数据的税。
分块决定一份权重能在片上复用多少次。块太大,片上放不下;块太小,启动开销和重复读取上升。FlashAttention 的核心不是“更快的注意力公式”,而是把在线归约、分块和融合绑在一起,让中间结果不必掉出片上。
未融合
QKᵀ
写回 HBM
softmax
写回 HBM
AV
融合
分块 · 在线归约 · 输出写一次
融合改变的是 R,不一定改变 F
相邻算子融合后,数学上的 FLOPs 几乎不变,但 HBM 流量下降,算术强度上升,屋顶线上的点向右移。解码器里每层的 Norm、RoPE、残差都是融合的候选:它们自己的运算量不大,单独跑时却要各付一次启动和一次同步。
编译器负责选择分块、布局和融合;运行时负责提交、重放和动态形状。Prefill 的序列长度每次不同,decode 的 batch 也在变。一张静态图可以消掉启动开销,却要为动态形状准备多个版本,或在运行时走 eager。选择本身也是成本。
实验
算子融合
额外 HBM 流量 0 B。融合不改变渐近 FLOPs,只改变 R。
0.02 ms
带宽墙
读取 R/β100%
计算 F/Π0%
下限取两者较大者。较短的一根无论再缩短,只要没超过另一根,总时间就不动。
考核
第 5 章考核
3 题
1.FlashAttention 主要减少的是?
2.把 RMSNorm 和线性层拆成两个核,最常见的代价是?
3.动态形状让运行时为难,是因为?
全部作答后交卷。