第 6 讲 · 高级 CUDA：Warp 特化 · Advanced CUDA Programming: Warp Specialization

对应 Maps to：CMU 15-779（Fall 2025，Zhihao Jia），Lecture 6 · Wed 9/17 · Advanced CUDA Programming: Warp Specialization。 阅读材料：原课这一讲没有指定必读论文。本讲的”论文精读”选了与主题最相关的两篇：J. Shah et al., FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision（NeurIPS 2024）；M. Bauer, H. Cook, B. Khailany, CudaDMA: Optimizing GPU Memory Bandwidth via Warp Specialization（SC 2011）。推荐：NVIDIA Hopper Architecture In-Depth（白皮书 / 技术博客）；CUTLASS 文档中关于 Hopper warp-specialized GEMM 的部分；M. Bauer, S. Treichler, A. Aiken, Singe: Leveraging Warp Specialization for High Performance on GPUs（PPoPP 2014）。 依据说明：本讲以原课 slides 06-warp-specialization.pdf 的主线组织：SM 架构与 warp 调度（活动 / 停顿 / 可发射）；数据从主机到设备再到 SM；软件流水线的两种实现；cp.async 与 Ampere 上的多级流水矩阵乘法；多级流水的问题；warp 特化（空间划分）与 CudaDMA；warp 特化 + 双缓冲；Hopper 的 TMA 与 setmaxnreg；Hopper 上的 warp 特化矩阵乘法与乒乓调度；FlashAttention-3 的 warpgroup 内重叠。论文数字引自论文原文（只做转述）。离散事件模拟器 smsim.py、其中的周期参数（部分是假设值）和全部实验都是本讲义自己设计的，不代表课堂上讲了什么。 配套代码 Code：code/lecture06/（4 个脚本 + 共享模块 smsim.py；只用 Python 标准库；每个脚本 < 1 秒） 预计阅读时间：3–4 小时（含运行代码）

0. 本讲导航 · How to read this lecture

0.1 学习目标

学完这一讲，你应该能够：

1. 说清楚 warp 的三种状态（活动、停顿、可发射），以及为什么”靠很多 warp 隐藏延迟”在 Tensor Core 内核里不够用。
2. 用 Little 定律估算”要有多少数据在路上”才能让访存管线不停；用周期账判断一个分块矩阵乘法是计算受限还是访存受限。
3. 比较软件流水线的两种实现（寄存器中转、cp.async 多级流水），说出多级流水的问题：级数随延迟增长、发射开销在关键路径上、寄存器压力、代码复杂。
4. 解释 warp 特化：生产者 warp 和消费者 warpgroup 怎样通过环形缓冲区和屏障（full / empty）交接数据，为什么它把发射开销移出关键路径。
5. 说出 Hopper 为 warp 特化提供的硬件支持：TMA、带事务计数的 mbarrier、异步的 wgmma、setmaxnreg 寄存器重分配，并会做寄存器预算。
6. 解释乒乓调度在 persistent GEMM（藏收尾）和 FlashAttention-3（藏 softmax）里各自消除了哪段空白，以及为什么 exp 的吞吐会成为下一代 GPU 上注意力的瓶颈。

0.2 先修与本讲范围

本讲不写真正的 CUDA：没有 GPU 无法验证。我们用伪代码说明结构，用 smsim.py 在时间轴上模拟”谁在等谁”。

0.3 本讲地图

§1  问题：Tensor Core 为什么在等
§2  SM 与 warp 调度复习：为什么"多开 warp"藏不住 Tensor Core 内核的延迟
§3  数据从哪里来：主机 → 显存 → SM；Little 定律                     ← 代码 01（演示 1）
§4  软件流水线：寄存器中转、cp.async 多级流水，以及它们的问题           ← 代码 01（演示 2–3）
§5  warp 特化：生产者—消费者、环形缓冲区与屏障                         ← 代码 02（演示 1–2、4）
§6  Hopper 的硬件支持：TMA、mbarrier、wgmma、setmaxnreg               ← 代码 02（演示 3）
§7  Hopper GEMM：persistent kernel 与乒乓调度                         ← 代码 03
§8  FlashAttention-3：把 softmax 藏到矩阵乘法后面                      ← 代码 04
§9  论文精读：FlashAttention-3；CudaDMA
§10 小结  §11 常见误区  §12 练习  §13 延伸阅读  §14 下一讲预告

0.4 关于模拟器的数字

smsim.py 把一个 SM 抽象成四种共享资源——TC（Tensor Core）、MEM（访存管线）、SFU（exp 等特殊函数）、FMA（CUDA core）——和若干执行者（warp / warpgroup）。周期参数的来源：

具体数字取决于假设，但”同样硬件下哪种安排更好”的结论在代码 02 演示 4 换参数后依然成立。这是模型，不是测量，代码 04 会把它和 FlashAttention-3 的实测并排放在一起。

1. 问题：Tensor Core 为什么在等 · The problem

第 5 讲的结尾留下一个问题：按 roofline 估算，FlashAttention 在 H100 上应该是计算受限的；但 FlashAttention-3 论文报告，FlashAttention-2 在 H100 上只用到峰值的约 35%。缺的 65% 去了哪里？

把一个 SM 想象成一个厨房。Tensor Core 是一台极快的切菜机，每秒能切成千上万份；食材从仓库（显存）运过来，路上要走一段时间（延迟）。如果厨师（warp）的做法是”去仓库拿一筐菜 → 等它运到 → 切 → 再去拿下一筐”，切菜机大部分时间都停着。

GPU 每一代都在给”让切菜机一直有活干”提供新工具：

这条时间线的主题只有一个词：异步（asynchrony）。计算越快，越需要让”搬数据”和”算”同时进行；而要同时进行，就要回答”谁来搬、谁来算、怎样交接”。warp 特化是 Hopper 时代的标准答案。

2. SM 与 warp 调度复习 · SM architecture and warp scheduling

2.1 warp 的三种状态

第 3 讲讲过：一个 SM 分成 4 个子分区（sub-partition），每个子分区有一个 warp 调度器，每个周期最多从它管辖的 warp 里挑一个发射一条指令。驻留在 SM 上的 warp 处于三种状态之一：

• 活动（active）：已经分配到这个 SM 上（有寄存器、有共享内存），还没执行完；
• 停顿（stalled）：活动但这一周期不能发射——在等内存、等前一条指令的结果、等屏障；
• 可发射（eligible）：活动且下一条指令的操作数都准备好了。

调度器每个周期从可发射的 warp 里挑一个。只要每个周期都有至少一个可发射的 warp，硬件就不会空闲——这就是用线程级并行（TLP）隐藏延迟：一个 warp 在等内存时，切换到另一个 warp，切换本身不花时间（第 3 讲 04 号脚本的调度动画）。

2.2 为什么这招在 Tensor Core 内核里不够用

靠很多 warp 隐藏延迟有一个前提：SM 上要驻留足够多的 warp。但高性能矩阵乘法要用大块（第 4 讲：块越大，复用越多，访存越少）：

• 一个 128×256 的输出块，FP32 累加器有 32,768 个数，分给 256 个线程每人 128 个寄存器；
• 多级流水的共享内存：每级 48 KB，4 级 192 KB——H100 一个 SM 最多 228 KB。

结果是一个 SM 上只能放一个 CTA，只有 8–12 个 warp。第 3 讲的占用率计算告诉我们，这远远不够用 TLP 隐藏几百个周期的延迟。

所以 Tensor Core 内核走的是另一条路：在一个 warp 内部制造并行（指令级并行 ILP / 异步）——发出加载后不等它完成，接着去做别的事；需要数据时再等。这就是软件流水线和 warp 特化要做的事。

3. 数据从哪里来 · Where the data comes from

3.1 主机 → 显存 → SM

数据到达 Tensor Core 要经过好几段：

主机内存 ──PCIe / NVLink-C2C──▶ HBM 显存 ──▶ L2 缓存 ──▶ SM：L1 / 共享内存 ──▶ 寄存器 ──▶ Tensor Core
          (cudaMemcpy，第 3 讲)     3.35 TB/s      50 MB        每 SM 228 KB          每 SM 256 KB

第 3 讲讲过第一段（cudaMemcpy）。本讲关心的是后几段：从 HBM / L2 到共享内存，再到 Tensor Core。在 Volta 之前，数据要先进寄存器再由线程写进共享内存；Ampere 的 cp.async 让数据直接从全局内存进共享内存；Hopper 的 TMA 让一个线程就能发起整块的拷贝；Hopper 的 wgmma 还能直接从共享内存读操作数，不必先搬进寄存器。

3.2 一个分块的周期账（代码 01 演示 1）

以 H100 上一个 CTA 计算 128×256 的输出块为例（BF16 输入、FP32 累加，沿 K 方向每步 64）：

$K = 4096$ 时共 64 步。理想情况（Tensor Core 从不等待）是 $64 \times 1024 = 65{,}536$ 周期，约 35.8 µs。

3.3 Little 定律：要有多少数据”在路上”

排队论里的 Little 定律：系统里平均的”在途量” = 吞吐 × 延迟。用在访存上：

也就是说，要让访存管线一直满负荷，任何时刻都要有约 44 KB（约 0.9 个块）的请求已经发出、还没到达。如果每次都等上一块到了才发下一块，管线大部分时间是空的。隐藏延迟 = 提前发出足够多的请求。 无论是靠多个 warp（TLP）还是靠一个 warp 提前发多个请求（流水线），本质都是增加在途量。

4. 软件流水线 · Software pipelining

4.1 起点：同步加载

第 4 讲的分块矩阵乘法：

for (k = 0; k < K; k += BK) {
    load_tile(A_smem, A + ..., k);  load_tile(B_smem, B + ..., k);   // 每个线程读几个数、写进共享内存
    __syncthreads();
    mma(acc, A_smem, B_smem);
    __syncthreads();                                                // 下一轮会覆盖共享内存
}

每一步都要等”发起 100 + 搬运 768 + 延迟 700”个周期，Tensor Core 才能算 1024 个周期。代码 01 的结果：利用率 39.2%。时间线上 TC 行一段忙、一大段空：

TC    |·········MMMMMMM·········MMMMMM··········MMMMMM·········MMMMMM··········|
MEM   |·LLLL···········LLLLL···········LLLL···········LLLLL···········LLLL·····|

4.2 实现一：寄存器中转的双缓冲

Ampere 之前的标准做法：计算第 $k$ 步之前，先发出第 $k+1$ 步的全局加载，数据进寄存器；第 $k$ 步算完后，把寄存器里的数据写进共享内存（另一份缓冲），再进入下一步。

regs = load_global(k = 0);  store_smem(buf[0], regs);  __syncthreads();
for (k = 0; k < nk; k++) {
    if (k + 1 < nk) regs = load_global(k + 1);      // 发出就走，不等
    mma(acc, buf[k % 2]);                           // 用当前缓冲区计算
    if (k + 1 < nk) { store_smem(buf[(k+1) % 2], regs); __syncthreads(); }   // 等数据到、写进另一份
}

加载的延迟被 MMA 盖住了，利用率升到 57.0%。但有两个代价：

• “寄存器 → 共享内存”这一段还在关键路径上；
• 暂存一整块要占大量寄存器：48 KB ÷ 128 个线程 = 每线程 96 个寄存器，和累加器抢地方。

4.3 实现二：cp.async 多级流水（Ampere）

Ampere 引入 cp.async：数据直接从全局内存拷到共享内存，不经过寄存器，发出后线程可以继续执行。配合”提交组 / 等待组”，可以让多步的加载同时在路上：

// S 级流水：共享内存里有 S 份缓冲区
for (s = 0; s < S - 1; s++) { cp_async_tile(buf[s], k = s); cp_async_commit_group(); }   // 序幕 prologue
for (k = 0; k < nk; k++) {
    if (k + S - 1 < nk) cp_async_tile(buf[(k + S - 1) % S], k + S - 1);               // 提前 S-1 步发出
    cp_async_commit_group();
    cp_async_wait_group(S - 1);       // 等到"第 k 步那一组"已经到达（最多还有 S-1 组在路上）
    __syncthreads();
    mma(acc, buf[k % S]);
    __syncthreads();
}

代码 01 演示 2 的结果（K = 4096，64 步）：

3 级以后就不涨了。原因在稳态：每一步，计算 warp 要先花 100 个周期发起下一批 cp.async，再做 1024 个周期的 MMA——上限是 $1024 / 1124 = 91.1\%$。发起加载的时间是计算 warp 自己付的。

4.4 多级流水的问题

1. 级数随延迟增长。代码 01 演示 3：要让利用率离最好结果不到 1 个百分点，延迟 300 / 700 / 1500 / 3000 周期时分别需要 2 / 3 / 4 / 5 级。5 级要 240 KB，超过了 H100 每个 SM 的 228 KB。经验公式：$S \ge 1 + \lceil (\text{延迟} + \text{搬运} + \text{发射}) / (\text{MMA} + \text{发射}) \rceil$（练习 2）。
2. 发射开销在关键路径上。同一个 warp 既要发起加载、又要发起 MMA，两件事抢发射槽；地址计算还要占寄存器。
3. 代码复杂。序幕、稳态、尾声（drain）三段，wait_group 的参数稍有差错就是数据竞争。
4. 一刀切。所有 warp 做同样的事，所以每个 warp 的寄存器数也一样——但”搬数据”和”做 MMA”需要的资源天差地别（§6.4）。

5. warp 特化 · Warp specialization

5.1 想法：按职责划分 warp

多级流水是时间上的划分：同一个 warp 轮流做”搬”和”算”。warp 特化是空间上的划分：一部分 warp 只搬，另一部分只算。

生产者 warp（producer）           环形缓冲区（S 份）            消费者 warpgroup（consumer）
────────────────────            ┌───┬───┬───┬───┐            ──────────────────────────
等 empty[slot]                   │ 0 │ 1 │ 2 │ 3 │            等 full[slot]
发起加载 → 到达时 full[slot] ++   └───┴───┴───┴───┘            MMA
slot = (slot + 1) % S                                         empty[slot] ++（缓冲区还给生产者）

两组 warp 通过每个缓冲区的两个屏障交接：

• full[slot]：数据到了，消费者可以用；
• empty[slot]：消费者用完了，生产者可以覆盖。

这个想法并不新。2011 年的 CudaDMA（§9.2）就提出让一部分 warp 专门做”DMA”（数据搬运），另一部分专门计算；2014 年的 Singe 编译器把 warp 特化用在燃烧化学的计算内核上。它在 Hopper 上成为主流，是因为 Hopper 的硬件让生产者变得极其”轻”（§6）。

5.2 伪代码

if (warpgroup_id == 0) {                       // 生产者
    for (i = 0; i < n_k; i++) {
        slot = i % S;  phase = (i / S) & 1;
        wait(empty[slot], phase);             // 等这个缓冲区被消费者释放
        tma_load(buf[slot], A, B, k = i, full[slot]);   // 一条 TMA 指令；数据到达时硬件自动在 full[slot] 上到达
    }
} else {                                       // 消费者
    for (i = 0; i < n_k; i++) {
        slot = i % S;  phase = (i / S) & 1;
        wait(full[slot], phase);              // 等数据到
        wgmma(acc, buf[slot]);  wgmma_wait();
        arrive(empty[slot]);                  // 用完了，还给生产者
    }
    epilogue(acc);
}

注意 phase：同一个屏障被循环使用，每用一轮翻转一次相位，等待者要等”正确的那一轮”。smsim.py 里用”到达次数 ≥ 第几次使用”来表示同样的意思。

5.3 效果（代码 02）

同样 4 级缓冲、K = 4096：

时间线对比（每格约 167 周期）：

warp 特化   TC |·········MMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMMM|
多级流水    TC |·········MMMMMMMMMMMMM·MMMMMM·MMMMMMMMMMMMMMMMMMM·MMMMMM·MMMMMMMMMMMMMMM|

多级流水每隔一两步就有一小段空白，那是计算 warp 在发 cp.async。warp 特化的 TC 行几乎连续。剩下的 4% 是开头第一次加载和结尾的收尾。

发射开销的影响（代码 02 演示 2）：

两点值得注意：

• 发射开销为 0 时，两种方案一样好。warp 特化的价值不是魔法，而是把”计算 warp 必须亲自做的杂事”挪走了。
• 即使生产者也用 cp.async、每步发射 800 个周期，warp 特化仍有 94.9%：只要生产者每步的工作不超过消费者每步的 MMA，它就能几乎完全藏在后面。

它不能创造带宽（代码 02 演示 4）：带宽减半后每步搬运 1536 周期 > MMA 1024 周期，两种方案都是 64.0%，访存管线利用率 96%。warp 特化把已有的带宽和计算重叠起来；如果访存本身就是瓶颈，要靠更大的块（更多复用）或更好的 L2 命中。

6. Hopper 的硬件支持 · Hopper features

6.1 TMA：一个线程搬一整块

TMA（Tensor Memory Accelerator） 是 Hopper 新增的硬件单元。程序先在主机上创建一个”张量描述符”（全局内存里张量的形状、步长、元素类型、要搬的块的形状、共享内存里的排布方式），之后在内核里一个线程发一条指令，就能把一个多维的块从全局内存搬到共享内存（或反过来）：

• 地址计算由硬件完成，不占线程的指令和寄存器——这就是模型里 issue_tma = 10 远小于 issue_cp = 100 的原因；
• 自动处理越界（越界的部分填 0），不需要每个线程判断边界；
• 可以按指定的”swizzle”模式写进共享内存，让后续读取没有 bank 冲突（第 4 讲 §6 用 padding 做的事，现在由硬件做）；
• 在线程块集群（cluster）里还能多播：一次读取，同时写进集群里几个 SM 的共享内存，减少重复读取 L2。

6.2 事务型 mbarrier

生产者发出 TMA 后，怎样让消费者知道”数据到齐了”？Hopper 的共享内存屏障（mbarrier）可以记录预期的字节数：生产者先声明”这个屏障要等 48 KB”，TMA 每写进一部分数据就自动扣减，扣到 0 时屏障翻转相位，等待它的消费者被唤醒。整个过程不需要任何线程去”检查”数据是否到了。

6.3 异步的 wgmma

Hopper 的 Tensor Core 指令 wgmma 以 warpgroup（4 个连续的 warp，128 个线程）为单位发出，操作数 A、B 可以直接从共享内存读，而且是异步的：发出后线程可以继续执行别的指令（例如 softmax），之后再用 wgmma.wait_group 等待结果。这让”在一个 warpgroup 内部重叠 GEMM 和别的计算”成为可能（§8.3）。

6.4 setmaxnreg：按角色重新分配寄存器

一个 SM 有 65,536 个 32 位寄存器。CUDA 在编译时给内核里的每个线程分配同样多的寄存器。但 warp 特化里两种角色的需求天差地别：

• 生产者只发 TMA：几个地址、几个循环变量，几十个寄存器足够；
• 消费者要存累加器：一个 128×256 的输出块由两个消费者 warpgroup（256 个线程）分担，每个线程光累加器就要 128 个寄存器。

Hopper 的 PTX 指令 setmaxnreg 让一个 warpgroup 在运行时主动减少（生产者）或增加（消费者）自己的寄存器上限。代码 02 演示 3（1 个生产者 + 2 个消费者，共 384 个线程）：

平均分配时消费者只剩 40 个寄存器放 A/B 片段、地址、softmax 的中间量（FlashAttention 里），很容易溢出到本地内存（实际是显存，非常慢）。把生产者的寄存器”让”给消费者后，消费者多了 64–72 个。这是按职责划分 warp 带来的另一个好处：每种 warp 的资源可以按需配置。

7. Hopper GEMM：persistent kernel 与乒乓调度 · Pingpong

7.1 persistent kernel 与收尾

前面只算了一个输出块，收尾（epilogue：FP32 累加器转成 BF16、写回显存）只发生一次。真实的 Hopper GEMM 常写成 persistent kernel：只启动和 SM 数量一样多的 CTA，每个 CTA 常驻一个 SM，依次计算很多个输出块（由一个”块调度器”分配）。好处是省掉反复启动 CTA 的开销、可以跨块重叠（第 7 讲会把这个想法推到极致）。

问题是：每个块都有一次收尾，收尾期间 Tensor Core 没活干。 只有 1 个消费者时，利用率大约是 主循环 / (主循环 + 收尾)。

7.2 两种安排：合作式与乒乓

CUTLASS 的 Hopper GEMM 有两种双消费者安排：

• 合作式（cooperative）：两个消费者 warpgroup 合算同一个输出块，各负责一半的行。每个线程的累加器少一半；但两者同时进入收尾，收尾期间 Tensor Core 空闲。
• 乒乓（pingpong）：两个消费者各自负责不同的输出块，轮流使用 Tensor Core。一个在收尾时，另一个正好在主循环里做 MMA。CUTLASS 用一个”有序屏障”保证两者的主循环轮流进行。

7.3 效果（代码 03）

一个 CTA 依次计算 8 个 128×256 的输出块，收尾约 1,280 周期：

K = 512 时的时间线：

1 个消费者  TC |··MMMMMMMM·MMMMMMMMM·MMMMMMMMM·MMMMMMMM··MMMMMMMM··MMMMMMMM·MMMMMMMMM·MMMMMMMMM·|
乒乓        TC |··MMMMMMMMMmmmmmmmmmmMMMMMMMMMMmmmmmmmmmMMMMMMMMMMmmmmmmmmmMMMMMMMMMMmmmmmmmmmm·|

（M、m 分别是两个消费者的 MMA。）

收尾越重（在收尾里融合激活函数、量化、加残差……），乒乓越有价值：K = 1024、收尾 ×8 时，1 个消费者掉到 61.1%，乒乓仍有 91.8%。乒乓剩下的损失主要是最后一个块的收尾——它后面没有主循环可以重叠（只扣这一段的上限是 92.8%）。

模拟器的一个简化（代码 README 第 4 条）：为了对比方便，乒乓的每个消费者也用 128×256 的块。这需要每线程 256 个累加器寄存器，现实中放不下。真实的乒乓内核里每个 warpgroup 的块更小，每步 MMA 相应变短，但”收尾被另一个的主循环盖住”的结论不变。

什么时候用哪种？ K 大时两者差别很小，合作式的累加器更省寄存器、可以用更大的块；K 小或收尾重时，乒乓更好。这类选择通常交给自动调优（第 9 讲）。

8. FlashAttention-3：把 softmax 藏到矩阵乘法后面 · FlashAttention-3

8.1 注意力里的”收尾”在每一步都有

GEMM 的收尾每个输出块只有一次；注意力的内层循环里，每个 K/V 块都夹着一次 softmax：

for 每个 K/V 块 j：
    S = Q K_jᵀ          ← GEMM0（Tensor Core）
    P = exp(S - m) ...  ← softmax（SFU 做 exp，CUDA core 做最大值、求和、缩放）
    O += P V_j          ← GEMM1（Tensor Core）

代码 04 演示 1 的周期账（一个 warpgroup 负责 64 行 query，每块 128 个 key，head dim 128，H100）：

exp 要花两个 GEMM 时间的一半。 用论文的数字核对：Tensor Core 989 TFLOP/s，exp 只有 3.9 TFLOP/s，吞吐差约 254 倍；每个分数要做 $4d = 512$ 次矩阵乘法运算、1 次 exp，时间比 $512 / 254 \approx$ 0.50。

串行执行时，每 1,024 个周期的 GEMM 要等 768 个周期的 softmax：Tensor Core 利用率 57.1%。

8.2 warpgroup 之间的乒乓

两个消费者 warpgroup 各负责一半的 query 行，各自串行地跑”GEMM0 → softmax → GEMM1”。一个在做 softmax 时，另一个的两个 GEMM 正好可以用 Tensor Core。代码 04：98.5%。

TC    |AAAAAABBBBBB···AAAAAAAAAAAAABBBBBBBBBBBBAAAAAAAAAAAABBBBBBBBBBBBBAAAAAAA|
SFU   |······xxxxxxxxxxxx··········xxxxxx······xxxxxx······xxxxxx·······xxxxxx·|

FlashAttention-3 论文用命名屏障强制两个 warpgroup 的 GEMM 轮流执行，并报告：在 head dim 128、序列长度 8192 的 FP16 前向上，乒乓把吞吐从约 570 提到 620–640 TFLOP/s。

8.3 warpgroup 内的两级流水

只有一个 warpgroup 时也能重叠一部分：wgmma 是异步的，所以可以先发出下一块的 GEMM0，再做这一块的 softmax：

S_0 = Q K_0ᵀ（异步发出）
for j：
    等 S_j
    异步发出 S_{j+1} = Q K_{j+1}ᵀ     ← 和下面的 softmax 同时进行
    P_j = softmax(S_j)
    O += P_j V_j                       ← 要用 P_j，只能等 softmax 做完

GEMM0 和 softmax 重叠了，GEMM1 仍要等。代码 04：79.0%。粗略估算：每块的时间 ≈ max(GEMM0, softmax) + GEMM1 = max(512, 768) + 512 = 1,280 周期，利用率 1,024 / 1,280 = 80%（练习 5）。两种都用：99.2%。

8.4 模型和实测的差距

论文的消融设置：非因果、FP16、批 4、序列 8448、16 头、head dim 128。注意两边的”只有 X”并不完全对应（论文里”没有 warp 特化”意味着加载也由计算 warp 发起）。

模型明显乐观：它假设 K/V 的加载完全被生产者藏住、屏障和同步不花时间、寄存器足够、两个 warpgroup 的 GEMM 能无缝衔接。真实内核里每一项都有代价，所以每种重叠只贡献十几个百分点，合起来也只有 67%。模型能告诉你为什么要重叠、哪种重叠消除哪段空白；它不能预测真实的 TFLOP/s。 这是做系统研究时对待简单模型的正确态度。

8.5 换一代硬件：exp 越来越贵

代码 04 演示 3 改变”每 SM 每周期的 Tensor Core FLOP”，保持 exp 吞吐不变：

A100 上 exp 不是大问题，这也是 FlashAttention-2 在 A100 上不需要这些技巧就能到 50–73% 的原因之一。B200 的假设下，exp 和 GEMM 一样贵：两种重叠都用时 Tensor Core 还能到 96%，但 SFU 也满载了——两种资源同时成为瓶颈，没有任何余量。FlashAttention-4 的对策是在 CUDA core 上用多项式近似计算一部分 exp，给 SFU 减负（CMU11-868 第 21 讲有展开）。

这个趋势是第 1 讲”算力增长不均衡”的又一个例子：Tensor Core 的吞吐每一代翻倍，别的单元没有跟上，瓶颈就转移到它们身上。

9. 论文精读 · Paper reading

9.1 FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision（Shah et al., NeurIPS 2024）

1. 这篇论文要解决什么问题？ FlashAttention-2 在 H100 上只用到峰值的约 35%。它是为 Ampere 设计的：没有利用 Hopper 的异步特性（TMA、异步 wgmma），也没有利用 FP8 低精度。

2. 核心想法是什么？ 三个技术：(1) warp 特化：生产者 warp 用 TMA 加载 Q、K、V，消费者 warpgroup 计算，通过环形缓冲区和屏障交接，并用 setmaxnreg 把寄存器让给消费者；加上 warpgroup 之间的乒乓调度。(2) GEMM 与 softmax 交错：在一个 warpgroup 内部，下一块的 GEMM0 与这一块的 softmax 重叠（两级流水）。(3) FP8 的块量化与不相干处理：按块量化，并用随机正交变换（Hadamard 变换加随机符号）把离群值”摊开”，降低量化误差。

3. 怎样验证？ H100 SXM5 上，比 FlashAttention-2 快 1.5–2.0 倍；FP16 最高约 740 TFLOP/s（约 75% 的峰值）；FP8 接近 1.2 PFLOP/s；FP8 的数值误差比一个基线 FP8 注意力低 2.6 倍。消融实验（本讲 §8.4）说明两种重叠缺一不可：缺少任何一种都会掉到 570–582 TFLOP/s。

4. 从系统的角度看，这篇论文的贡献是什么？ 它是”把硬件的新特性系统地用到一个关键内核上”的范例。论文本身的算法思想（分块、在线 softmax）和 FA2 一样，新的是执行方式：谁搬数据、谁算、什么时候等。它也清楚地展示了一个趋势：随着 Tensor Core 越来越快，内核设计的难点从”减少访存”转向”重叠异构的计算单元”（Tensor Core、SFU、CUDA core、TMA）。

5. 局限与后续？ 实现与 Hopper 强绑定，代码复杂（依赖 CUTLASS 的 CuTe 抽象）；实验以前向为主，低精度注意力对大规模训练的影响仍是开放问题。手写这种内核的成本，正是第 8 讲（Triton 等 DSL）和第 11 讲（Mirage 等超优化器）想降低的。Blackwell 上 exp 与 GEMM 的比例进一步恶化，引出了 FlashAttention-4 的新技巧。

9.2 CudaDMA: Optimizing GPU Memory Bandwidth via Warp Specialization（Bauer, Cook, Khailany, SC 2011）

1. 问题：在 Fermi 时代的 GPU 上，把数据从全局内存搬进共享内存的代码和计算代码混在同一批 warp 里，访存模式受计算的线程布局约束，难以合并，也难以和计算重叠。

2. 想法：把一个线程块里的 warp 分成两类——DMA warp 专门搬数据，计算 warp 专门计算；用一组库接口（类似 DMA 引擎的”开始传输 / 等待完成”）和命名屏障协调两者。这样搬运的线程布局可以独立地为合并访存优化，计算 warp 也不用分心。

3. 意义：这是 warp 特化思想的早期系统化表述。它当时需要用软件（命名屏障）模拟 Hopper 后来用硬件实现的东西（TMA、mbarrier）。读这篇论文可以看到：好的软件抽象往往先于硬件出现，硬件随后把它固化下来。 后续的 Singe 编译器（PPoPP 2014）把 warp 特化推广到更一般的计算，让编译器自动给不同的 warp 分配不同的任务。

9.3 两篇论文之间的联系

CudaDMA 提出”按职责划分 warp”时，硬件并不直接支持：DMA warp 自己也要一条条发出访存指令，还要占用寄存器。十几年后，Hopper 的 TMA 让生产者”轻”到只需要一个线程，setmaxnreg 让它把寄存器让出来，异步 wgmma 让消费者也能边等边算——FlashAttention-3 正是在这些硬件上把同一个思想用到了极致。本讲代码 02 演示 2 的结论在两篇论文之间搭了一座桥：发射开销为 0 时，warp 特化和多级流水一样好；warp 特化的价值随着”搬数据的杂事”的多少而变化，而硬件的演进正在不断改变这个量。

10. 小结 · Summary

1. Tensor Core 太快了，数据跟不上；高性能内核的块大、寄存器和共享内存多，一个 SM 上只能放很少的 warp，靠多开 warp 隐藏延迟（TLP）不够用，必须在 warp 内部异步执行。
2. Little 定律：要让访存管线满负荷，在途数据 = 带宽 × 延迟（模型里约 44 KB / SM）。
3. 软件流水线：寄存器中转（57.0%，每线程多 96 个寄存器）→ cp.async 多级流水（3 级 88.0%）。多级流水的问题：级数随延迟增长、发射开销在关键路径上（上限 91.1%）、代码复杂、资源一刀切。
4. warp 特化把”搬”和”算”交给不同的 warp，通过环形缓冲区 + full/empty 屏障交接（96.0%）；它把发射开销移出关键路径，但不能创造带宽。
5. Hopper 的支持：TMA（一个线程搬一整块，硬件算地址、处理越界和 swizzle）、事务型 mbarrier、异步 wgmma、setmaxnreg（生产者 24–40、消费者 232–240 个寄存器）。
6. 乒乓调度让两个消费者轮流使用 Tensor Core：在 persistent GEMM 里藏收尾（K = 256 时 73.7% → 92.2%），在 FlashAttention-3 里藏 softmax（模型 57.1% → 98.5%；论文实测各贡献十几个百分点，合起来 661 TFLOP/s）。
7. 每一代 Tensor Core 都在变快，别的单元没跟上：exp 在 H100 上是 GEMM 的一半，在 B200 的假设下和 GEMM 一样贵。

11. 常见误区 · Common misconceptions

• “占用率越高越好。” 对访存密集的简单内核基本成立；对 Tensor Core 内核，大块带来的复用比占用率更重要。高性能 GEMM / FlashAttention 内核的占用率往往很低（一个 SM 一个 CTA），延迟靠异步和流水线隐藏。
• “多级流水的级数越多越好。” 级数够盖住延迟就行（代码 01：3 级以后不再提升），多了只会占共享内存、限制块大小。
• “warp 特化让内核更快，所以总该用。” 发射开销为 0 时它和多级流水一样（代码 02 演示 2）；访存受限时它也救不了（演示 4）。它的价值取决于”搬数据的杂事”有多重、寄存器需求是否悬殊。
• “warp 特化就是多开几个 warp。” 不是。多开 warp（TLP）是让同样的工作有更多份可以切换；warp 特化是让不同的warp 做不同的工作，互相配合。
• “TMA 让访存更快。” TMA 不增加带宽；它减少的是发起访存的指令和寄存器开销，并把地址计算、越界处理、swizzle 交给硬件。
• “乒乓调度提高了 Tensor Core 的峰值。” 峰值不变；乒乓减少的是 Tensor Core 的空闲时间（收尾、softmax 期间）。
• “模拟器说 99%，所以真实内核能到 99%。” 代码 04 的教训：简单模型适合比较方案、定位空白，不能预测绝对性能。

12. 练习 · Exercises（含参考答案）

练习 1　用 Little 定律估算：H100 全卡带宽 3.35 TB/s，访存延迟按 700 个周期（1.83 GHz）算，全卡要有多少字节在途？平均每个 SM 多少？

答案：延迟 700 / 1.83 GHz ≈ 383 ns。在途字节 = 3.35 × 10¹² × 383 × 10⁻⁹ ≈ 1.28 MB，每个 SM 约 9.7 KB。（模型里每个 SM 用的是 64 B/周期的 L2 带宽，所以要 44 KB；访问 HBM 时带宽更低，在途量也更少，但延迟通常更长。）

练习 2　推导多级流水需要的级数：设每步发射 $I$、MMA $C$、搬运 $X$、延迟 $L$ 个周期。稳态下每步用时 $C + I$。第 $k$ 步的加载在第 $k - S + 1$ 步开始时发出，要在第 $k$ 步开始前到达。证明需要 $S \ge 1 + \lceil (I + X + L) / (C + I) \rceil$，并用代码 01 的参数（$I = 100$、$C = 1024$、$X = 768$）算延迟 1500 周期时的 $S$。

答案：第 $k$ 步的加载发出后要经过 $I + X + L$ 才能用；从发出到第 $k$ 步开始，中间隔了 $S - 1$ 步，每步 $C + I$。要求 $(S - 1)(C + I) \ge I + X + L$。代入：$(100 + 768 + 1500) / 1124 = 2.1$，向上取整 3，$S \ge 4$。与演示 3 的结果一致（4 级）。这个公式忽略了访存管线排队，是一个保守估计。

练习 3　一个 warp 特化内核用 1 个生产者 warpgroup 和 3 个消费者 warpgroup。生产者用 setmaxnreg 降到 24 个寄存器，消费者最多能升到多少（必须是 8 的倍数）？

答案：$128 \times 24 + 384 \times r \le 65{,}536$ → $r \le 162.67$，取 8 的倍数 160。

练习 4　代码 03 里 1 个消费者、K = 256 时利用率 73.7%，估算上限 76.2%。这两个数为什么不一样？乒乓为什么到不了 100%？

答案：上限只考虑了每块的收尾；实际还有开头第一次加载的等待（约 延迟 + 搬运 ≈ 1,500 周期）。乒乓到不了 100% 是因为开头的加载和最后一个块的收尾无法和别的主循环重叠。

练习 5　按 §8.3 的粗略估算方法，算 FlashAttention-3 组内两级流水在 B200 假设（GEMM0 = GEMM1 = 256，exp = 512，其余 256）下的利用率，和代码 04 的 49.6% 比较。

答案：每块 ≈ max(GEMM0, exp + 其余) + GEMM1 = max(256, 768) + 256 = 1,024，Tensor Core 忙 512，利用率 50%，与模拟的 49.6% 一致。

练习 6　为什么说”高性能 GEMM 内核的占用率往往很低”不是缺点？用一个 SM 上的寄存器和共享内存解释。

答案：128×256 的块要 192 KB 共享内存（4 级）和每线程 200 多个寄存器，一个 SM 只放得下一个 CTA（约 12 个 warp）。但大块意味着每个字节被复用更多次，访存需求更低；延迟由异步流水线隐藏，不需要很多 warp。把块缩小来提高占用率，会增加访存量，往往更慢。

练习 7　warp 特化里，如果生产者每步的工作（发射 + 等待 empty）比消费者每步的 MMA 还慢，会发生什么？用代码 02 演示 2 的思路设计一个实验来验证。

答案：消费者会在 full 屏障上等，Tensor Core 利用率下降到 MMA / 生产者每步时间。实验：把 issue_cp 设成大于 mma（例如 1500），比较 run_warp_specialized(..., use_tma=False) 的利用率是否接近 1024/1500 ≈ 68%。

练习 8　FlashAttention-3 论文说它的 FP8 版本”接近 1.2 PFLOP/s”。H100 的 FP8 稠密峰值是 1,979 TFLOP/s。这个利用率比 FP16 的 75% 低还是高？用 §8.5 的思路解释为什么。

答案：约 1.2 / 1.979 ≈ 61%，比 FP16 的 75% 低。FP8 让 Tensor Core 快了一倍，但 exp 的吞吐不变，exp / GEMM 的比例从 0.5 变成约 1.0——相当于在 H100 上提前遇到了 §8.5 里 B200 的问题。

13. 延伸阅读 · Further reading

• NVIDIA, NVIDIA H100 Tensor Core GPU Architecture（白皮书）——TMA、线程块集群、分布式共享内存、异步事务屏障的官方描述。
• CUTLASS 文档：Efficient GEMM in CUDA、Hopper 的 warp-specialized / pingpong / cooperative 内核调度说明。
• M. Bauer, S. Treichler, A. Aiken, Singe: Leveraging Warp Specialization for High Performance on GPUs（PPoPP 2014）。
• J. Shah et al., FlashAttention-3（2024）——§9.1。
• PTX ISA 文档中关于 cp.async、cp.async.bulk.tensor、mbarrier、wgmma、setmaxnreg 的章节。
• CMU11-868 第 21 讲：FlashAttention-3 的乒乓、FP8 不相干处理、FlashAttention-4 的多项式 exp。

14. 下一讲预告 · Next lecture

本讲在一个 SM、一个内核的范围里消除空白：让 Tensor Core 不等数据、不等收尾、不等 softmax。第 7 讲《高级 CUDA：Mega-Kernel》把视野扩大到整个模型：小批量解码时，一个 Llama 的前向要启动上百个内核，每个内核之间都有启动开销和”全局屏障”，下一层的权重要等上一层完全结束才能开始加载。如果把整个模型编译成一个常驻的内核，在 GPU 内部运行一个任务调度器，用细粒度的依赖代替内核边界，会发生什么？我们会换一个粒度——从”一个 SM 上的周期”到”132 个 SM 上的任务”——用一个新的执行模型 megasim.py 比较”逐算子内核”和”mega-kernel”两种执行方式的时延。