斯坦福 CS336 从零构建大模型 (2025 春) - 第五讲：GPU 与系统优化（Flash Attention）

斯坦福 CS336 第五讲的核心目标是打破 GPU 的“黑盒”魔法，深入理解底层硬件的运行机制以及如何编写极致高效的算法（如 Flash Attention）。课程详细解析了 GPU 的架构、内存墙瓶颈，并总结了提升大模型训练与推理效率的六大核心系统优化技巧。

以下是第五讲所有核心知识点的全景梳理：

一、 CPU 与 GPU 的本质区别

随着摩尔定律和 Dennard 缩放定律的失效，单线程 CPU 的性能增长已经停滞，深度学习被迫转向大规模并行计算。

• CPU（优化延迟 Latency）： 拥有庞大的控制单元（Control logic）和分支预测机制，核心目标是以极快的速度单行执行单个复杂任务。
• GPU（优化吞吐量 Throughput）： 控制逻辑极少，绝大部分芯片面积被让给了海量的计算单元（ALUs / 算术逻辑单元）。单个任务的延迟可能变高，但能在极短时间内并发完成海量任务。

二、 GPU 的物理架构与执行模型

为了写出高效的代码，必须理解 GPU 是如何组织计算的：

• 硬件层级： GPU 内部包含多个 SM（流式多处理器，Streaming Multiprocessor）。在 SM 内部，又包含了多个 SP（流处理器）和专门用于加速矩阵乘法（MatMul）的 Tensor Cores（张量核心）。(注：讲师也顺带提到了 Google 的 TPU，其架构理念极其相似，也包含专门做矩阵乘法的 MXU 单元。)
• 逻辑层级（SIMT 执行模型）： 任务被划分为不同的层级。Block（线程块） 被分配给各个独立的 SM 执行；在 Block 内部，线程被打包成 Warp（线程束，通常含 32 个线程）。同一个 Warp 中的所有 32 个线程必须在不同的数据上执行完全相同的指令（SIMT 模型）。

三、 终极瓶颈：算力富裕与“内存墙”

课程强调，在过去的演进中，GPU 的算力（FLOPs）暴增了约 100,000 倍，而内存带宽（数据搬运的速度）仅仅增长了约 100 倍。这意味着在绝大多数时候，系统的瓶颈不是算力不足，而是内存太慢（Memory bound）。

根据 GPU 的内存层级（Memory Hierarchy），数据的物理距离决定了访问速度：

• 极速内存（SM内部）： 寄存器（Registers）、L1 缓存、共享内存（Shared Memory）。访问极快，仅需约 20 个时钟周期。
• 龟速内存（SM外部）： 全局显存（Global Memory / HBM / DRAM）。访问需要漫长的 200 到 300 个时钟周期。

四、 六大 GPU 性能优化杀手锏

为了克服“内存墙”，压榨出最高效的矩阵乘法性能，工程师们总结出了以下技巧：

1. 消除条件分支 (Avoid Conditionals)

因为 GPU 采用 SIMT 模型，同一个 Warp 里的线程必须步调一致。如果代码中出现了 if/else 分支，GPU 无法同时执行两种逻辑。它只能先**强制挂起（休眠）**进入 else 的线程，执行完 if 后，再唤醒另一半执行 else。这会直接导致并行计算效率折半。

2. 降低数值精度 (Lower Precision)

使用 FP16、BF16 甚至是 FP8/INT8 格式。这不仅仅是为了减少显存占用，更关键的是将需要搬运的字节数砍半，相当于免费将内存带宽翻倍。在混合精度矩阵乘法中，输入可以是 16 位，但在内部进行累加时必须使用 32 位以防止精度丢失。

3. 算子融合 (Operator Fusion)

如果在 PyTorch 中直接写 sin2(x)+cos2(x)，系统会启动多次读写：计算一次 sin 并写回慢速全局显存，再读出来计算平方，往复多次。算子融合将这些连续的操作打包成一个“大操作”（Kernel），在极速计算单元内一次性全算完，最后只写回一次结果。目前像 torch.compile 这样的编译器会自动为你完成这种优化。

4. 激活值重计算 (Recomputation / Gradient Checkpointing)

这是一种**“用廉价的算力去换取昂贵的内存带宽”的策略。在反向传播时，我们本需要从慢速全局显存中读取前向传播保留下来的“激活值”。重计算策略选择直接不保存它们，而是在反向传播时利用局部极速内存实时重新计算一遍**这些激活值。虽然浪费了算力，但省去了海量的慢速内存读写，整体速度反而更快。

5. 内存合并访问 (Memory Coalescing & Burst Mode)

GPU 从全局显存读取数据时，不是逐个字节读取的，而是按块（Burst section）一锅端。如果一个 Warp 中的 32 个线程读取的是一段连续的内存地址，硬件会自动“合并”这些请求，一次拉取所有数据，吞吐量可提升 4 倍。如果乱序读取或按列遍历（导致访问非连续的 Burst section），访存速度会遭遇毁灭性打击。

6. 分块矩阵乘法 (Tiling)

对于巨大的矩阵乘法，如果直接算，大量元素会被反复从慢速显存中读取。Tiling 的思想是：将大矩阵切分成一个个小块（Tile），将一个小块整体加载到极速的“共享内存（Shared Memory）”中。在这个小块完成所有力所能及的乘加计算后，再换下一个小块。这能把对全局显存的访问量骤降 T（Tile尺寸）倍。

五、 解密波浪形的性能暴跌之谜 (Wave Quantization & Alignment)

讲师解释了为什么在测试矩阵乘法时，尺寸稍微改变一点，GPU 的利用率和速度就会出现断崖式下跌（呈现诡异的波浪图）：

• 分块对齐（Alignment）： 如果矩阵大小不是 Tile 大小或 Burst 块大小的整数倍，GPU 就会出现“读半块数据丢半块”的现象，导致访存翻倍。这也是为什么将词表大小强行填充（Pad）到 64 的整数倍能让模型提速 25% 的底层原因。
• 波次量化（Wave Quantization）： 假设你把任务切成了 120 个 Tile，而你的 H100 GPU 只有 108 个 SM。第一波 108 个 SM 跑得满满当当，但等它们跑完后，剩下的 12 个 Tile 只能占用 12 个 SM，导致剩余的 96 个 SM 处于完全闲置状态等待计算结束。这会引发利用率的断崖式暴跌。

六、 案例应用：解构 Flash Attention

Flash Attention 的成功完美综合了上述的底层优化原理：

• 痛点： 标准 Attention 需要计算 Q×K 得到一个庞大的 N×N 矩阵，然后对每一行进行 Softmax 归一化。这个过程必须把庞大的 N2 中间结果写入慢速 HBM，导致二次方的内存时间开销。
• Tiling（分块）： Flash Attention 将 K、Q、V 矩阵切成 Tile，复制到极速 SRAM 中局部计算，避开了全局读写。
• 在线 Softmax (Online Softmax)： 因为分块后无法直接获得整行的全局最大值，它采用了一种巧妙的数学技巧（追踪在线常数），使得可以在仅扫描当前块的情况下增量式地计算局部 Softmax，最后再统一调整归一化项。
• 重计算 (Recomputation)： 在反向传播时，它不存储庞大的 N2 注意力矩阵，而是按块在极速内存中实时重新计算，完美规避了内存墙限制。

---

七、第五讲复习题 (Lecture 5: GPUs)

一、 GPU 硬件与执行逻辑

1. CPU 与 GPU 的设计哲学： CPU 和 GPU 面对计算任务时，它们各自追求优化的核心目标（设计指标）有什么本质区别？
2. SIMT 模型与分支陷阱： 为什么在编写 GPU 高性能代码时，必须极力避免同一个线程束（Warp）内出现 if/else 等条件分支语句？

二、 内存墙危机 (The Memory Wall)

3. 算力与带宽的鸿沟： 过去几代硬件的演进中，出现了一个极其严重的不平衡，导致现代大语言模型几乎全都被卡在了“内存（Memory-bound）”上。请问这个不平衡是什么？
4. 内存层级的延迟惩罚： 在 GPU 的内存层级中，从极速内存（如寄存器、共享内存/L1）读取数据，与从外部龟速的全局显存（Global Memory / DRAM）读取数据，两者耗费的时钟周期大约相差多少倍？

三、 六大底层性能优化绝招

5. 数值降级的双倍红利 (Lower Precision)： 为什么将混合精度矩阵乘法的输入从 32 位（FP32）降为 16 位（FP16/BF16），能够相当于“免费将内存带宽翻了一倍”？
6. 算子融合 (Operator Fusion)： 如果我们在 PyTorch 中直接连续写 x = sin(x) 然后 y = x**2，底层运行效率会很低。算子融合机制是如何解决这个问题的？
7. 反直觉的重计算 (Recomputation)： 在反向传播时，工程师为什么宁愿浪费算力把前向传播产生过的激活值（Activations）重新再算一遍，也不愿意直接把它们保存在显存里供反向传播读取？
8. 内存合并访问 (Memory Coalescing)： 为什么在 GPU 编程中，强制要求连续的线程访问连续的内存地址（例如遍历矩阵时最好顺着行的方向而不是列的方向）？
9. 分块矩阵乘法 (Tiling)： 面对极大的矩阵乘法，Tiling 技术的核心思路是什么？它是如何利用 GPU 的物理内存层级的？

四、 案例应用：Flash Attention

10. 在线 Softmax (Online Softmax) 的数学魔法： Flash Attention 为了切块（Tiling）计算，必须在局部极速内存中完成 Softmax 操作。但传统 Softmax 必须拿到一整行的所有元素才能算出全局最大值和分母。它是用什么数学方法克服这个悖论的？

---

八、参考答案与知识点解析

1. CPU 与 GPU 的设计哲学？

答案： CPU 旨在优化延迟（Latency），它拥有庞大的控制单元和分支预测，致力于尽可能快地跑完单个任务序列。而 GPU 旨在优化吞吐量（Throughput），它牺牲了单个任务的延迟，将绝大部分芯片面积让给海量的计算单元（ALUs），以在整体上最快地并发完成成千上万个任务。

2. SIMT 模型与分支陷阱？

答案： GPU 采用 SIMT（单指令多线程）模型，同一个 Warp（一组 32 个线程）中的所有线程必须在同一时刻执行完全相同的指令。如果遇到 if/else 分支，GPU 无法同时执行两种逻辑，只能强制让进入 else 的线程休眠，先执行完 if，再唤醒另一半执行 else。这会直接导致该波次的并行计算效率折半。

3. 算力与带宽的鸿沟？

答案： 硬件演进历史中，GPU 的浮点计算能力（FLOPs）实现了惊人的超指数级爆发（提升了约 100,000 倍），而内存带宽（数据搬运的速度）的增长却非常缓慢（仅提升了约 100 倍）。这种算力远超带宽的鸿沟，导致昂贵的计算单元绝大多数时候都在无奈地等待缓慢的全局显存传送数据。

4. 内存层级的延迟惩罚？

答案： 紧贴计算核心的共享内存（Shared Memory / L1）极快，读取通常只需要约 20 个时钟周期。而穿过物理芯片去读取外部的全局显存（DRAM）则非常慢，需要漫长的 200 到 300 个时钟周期，两者相差高达 10 倍以上。

5. 数值降级的双倍红利？

答案： 在受限于内存带宽（Memory-bound）的系统中，真正的瓶颈是“往计算单元搬运了多少字节的数据”。将精度从 32 位降至 16 位，系统每次需要从慢速全局显存中搬运的字节数直接减少了一半，在相同的算力消耗下，这相当于免费把系统的数据吞吐带宽翻了一倍。

6. 算子融合的作用？

答案： 独立的连续操作会导致 GPU 每次算完一步就把中间结果写回慢速的全局显存，下一步再读出来，产生巨大的内存读写开销。算子融合（如 torch.compile 自动执行的优化）将多个操作合并成一个大的 CUDA Kernel，在极速计算单元（SM）的局部内存中一气呵成处理完所有步骤，最后只往全局显存写回一次最终结果。

7. 反直觉的重计算？

答案： 这个策略的核心逻辑是**“用极其廉价富裕的算力，去换取极其昂贵稀缺的内存带宽”**。保存并读取前向传播的庞大激活值需要耗费大量的慢速全局显存读写操作；相反，在反向传播时利用局部极速内存将这些激活值动态“实时重算”一遍，虽然耗费了计算，但省去了漫长的数据搬运，整体速度反而更快。

8. 内存合并访问的原理？

答案： 硬件设计上，全局显存（DRAM）并不是按单个字节吐出数据的，而是按“突发模式（Burst mode）”一次性返回一整个连续的内存块。如果相邻线程访问的是连续内存，硬件会自动把请求合并（Coalescing），一次拉取全部数据，吞吐量可飙升 4 倍。如果是乱序或跨列访问，就会读来大量无用数据，严重拖慢带宽。

9. Tiling 的核心思路？

答案： 如果直接算大矩阵乘法，大量元素会被从慢速显存中反复读取，造成灾难性的延迟。Tiling 的思路是将大矩阵切分为适合装入极速**共享内存（Shared Memory）**的小块（Tile），将小块整体加载后，在局部缓存中完成所有可能的乘加操作，然后再处理下一个小块，从而把全局显存的访问次数降低了 Tile 尺寸的倍数。

10. Flash Attention 与在线 Softmax？

答案： Flash Attention 为了不把庞大的二次方注意力矩阵写回慢速显存，利用了 在线 Softmax（Online Softmax） 算法。该数学技巧通过维护一系列增量更新的变量（追踪当前见过的最大值和归一化分母补偿项），使得网络可以在仅扫描当前块（Tile by Tile）的情况下，流式地、局部地算出精确的全局 Softmax 结果