课程简介
CS336 Lecture 5 · GPUs, TPUs — 中文 Lecture Note
本笔记基于 Stanford CS336 (2026) Lecture 5 编写，主题为 GPU 与 TPU 硬件模型、GPU 性能优化技巧及 FlashAttention 原理。目标读者为有编程基础、希望深入理解语言模型训练/推理底层硬件的学习者。
学习目标
完成本讲后，你应该能够：
- 描述 GPU 的核心硬件结构（SM、内存层次、warp），并将其与 CPU 的设计哲学进行对比。
- 解释 GPU 的 SIMT 执行模型及其对控制流（如 if 语句）的语义影响。
- 列出六种提升 GPU 利用率的核心技巧，并说明各自解决的问题。
- 理解 roofline 模型，并能解释"内存受限"与"计算受限"的划分。
- 说明 FlashAttention 如何通过对注意力计算进行分块（tiling）、重计算（recomputation）与算子融合，实现高效的内存利用。
- 理解低精度计算（如 FP8、MXFP8）在现代 GPU 中的实现方式与权衡。
课程主线
本讲分为三个部分：
- 第一部分：GPU 硬件与编程模型 — 建立 GPU 的思维模型（CPU 对比、SM、内存层次、SIMT、TPU 映射）。
- 第二部分：六个 GPU 加速技巧 — 从控制流发散到低精度、算子融合、重计算、合并内存访问和分块。
- 第三部分：FlashAttention 复盘 — 将前述技巧综合应用于注意力机制，实现内存高效的注意力计算。
最终目标：理解课程开篇展示的矩阵乘法吞吐量曲线图——为何在特定矩阵维度下性能出现周期性暴跌，以及如何通过调整参数（如 padding）规避这些问题。
核心概念
1. GPU 硬件模型
- CPU vs GPU：CPU 面向低延迟、复杂控制流，拥有少数 ALU 与大型控制单元；GPU 面向高吞吐量，拥有大量轻量级计算核心（SM），单个任务延迟高，但总吞吐量极大。
- SM（Streaming Multiprocessor）：GPU 的基本计算单元，类似于"核心"。A100 有 128 个 SM。每个 SM 包含流处理器（SP），可并行执行线程。
- 内存层次：
- 寄存器（最快、最本地化）
- 共享内存（SM 内，可编程，近 L1 缓存）
- L1/L2 缓存（硬件管理）
- 全局内存 / HBM（DRAM，容量大，速度慢）
- 常量内存、主机内存（CPU 侧，可卸载）
关键点：全局内存访问比 L1/共享内存慢约 10 倍（A100 实测）。因此，所有优化策略的核心都是：尽可能把数据留在靠近计算单元的内存层级。
2. GPU 编程模型
- 线程（Thread）：轻量级执行单元，遵循 SIMT（单指令多线程） 模型——同一 warp 中所有线程执行相同指令，但输入数据不同。
- Block：一组线程，保证被调度到同一个 SM 上，可访问共享内存。
- Warp：GPU 的调度单元，32 个连续编号线程构成一个 warp。调度器以 warp 为单位分发指令。
控制发散（Divergence）：在 SIMT 下，if/else 的两个分支都会被所有线程执行，不满足条件的线程被屏蔽（空转）。因此，GPU 代码中应尽量避免分支，改用掩码/乘法实现条件逻辑（如 ReLU）。
注：一个 block 中的 warp 数量因硬件而异，没有固定值。
3. TPU 与 GPU 的趋同演化
- TPU 在高层结构上与 GPU 高度相似：均有矩阵乘法单元（GPU：Tensor Core；TPU：MXU）、并行向量单元和内存层次（HBM + 共享内存 SMEM）。
- 主要差异在规模与灵活性：
- GPU：大量（约 528）小矩阵乘法单元，灵活性高
- TPU：少量（约 2~8）大矩阵乘法单元，锁死在大矩阵乘法上
- 命名陷阱：TPU 的"张量核心"指处理器单元；GPU 的"Tensor Core"指矩阵乘法电路。
4. Roofline 模型与操作强度
- Roofline 模型：横轴为操作强度（每字节内存移动对应的 FLOPs），纵轴为可达吞吐量。
- 曲线分为上升段（内存受限）和平坦段（计算受限）。训练/推理时应处于平坦段——即计算单元已饱和。
- 操作强度（Arithmetic Intensity） = FLOPs / Bytes moved。提高操作强度是优化的核心方向（例：ReLU 在 FP32 下为 1 FLOP / 8 bytes，操作强度极低）。
六个 GPU 加速技巧
技巧 1：控制流最小化（避免发散）
- 原因：if/else 分支在 SIMT 下两个分支都执行，空转等待。
- 实践：用掩码乘法替代分支（如 ReLU 使用
x * (x > 0)）。
技巧 2：低精度计算（Low Precision）
- 动机：降低每个操作的数据字节数，缓解内存带宽压力；同时张量核心在低位宽下吞吐量大幅提升。
- 原理：矩阵乘法中，输入以低精度（如 BF16）送入，乘积以全精度（FP32）累加，输出转回 FP32。
- 挑战：不同算子对精度敏感度不同（如 softmax、exp 常需 FP32）；低位宽易上溢/下溢，需要缩放因子。
- MXFP8：分块量化，每 32 个元素共享一个 E8M0 缩放因子。代价是转置/量化操作复杂——实践中需为每个量化矩阵保留原始与转置两份拷贝。
- FP4（MXFP4）：仅 16 个值（-6 到 6），每 16 个元素一个 E4M3 缩放因子。已有训练尝试，但尚未成为主流。
- 结论：低精度可节省 20–30% 训练时间，但非免费——量化/反量化开销需权衡。第一层和最后一层通常最难量化（易导致训练不稳定）。
技巧 3：算子融合（Operator Fusion）
- 思想：将多个算子合并为一个 kernel，减少中间结果在全局内存中的读写次数。
- 例子：
sin²x + cos²x在 PyTorch 中产生 5 个中间操作。融合后只读一次全局内存、写一次结果。 - 工具：Torch Compiler / JAX Compiler 可自动完成简单融合；复杂融合（如 FlashAttention）需手动设计。
技巧 4：重计算（Recomputation）
- 思路：丢弃反向传播所需的中间激活值，在反向时重新计算。
- 收益：内存访问次数减少（8 次 → 5 次），计算量不变（或略增）。
- 适用场景：计算充裕但内存带宽紧张时（例如注意力得分矩阵的中间值）。
技巧 5：合并内存访问（Coalesced Access）
- 原理：DRAM 以突发（burst）模式工作——访问一个地址会一次性返回一个连续段的数据（约 128 字节），只要访问落在同一段内就是"免费"的。
- 要求：warp 中所有线程的读地址应落在同一个（或相邻）突发段内。
- 反例：行主序矩阵按列访问时，每个线程落在不同突发窗口，导致严重低效。
- 实践：尽量按行（连续地址）访问矩阵；无法满足时对矩阵进行 padding，使行对齐到突发边界。
- 案例：nanoGPT 将词表 50,257 → 50,304（对齐到 64 的倍数），获得约 25% 加速。
技巧 6：分块（Tiling）
- 思想：将数据切分为块（tile），加载到共享内存后反复使用，大幅减少全局内存访问次数。
- 效果：矩阵乘法中，无分块时每个输入元素从全局内存读取 N 次；分块大小 T 时减少到 N/T 次。
- 约束：分块大小须与矩阵维度、突发段边界、共享内存容量对齐；不整除时会出现"瘦长块"或波次量化问题。
- 自动化：PyTorch 的
max-autotune编译器会尝试不同分块大小以找到最优。
波次量化（Wave Quantization）
- 现象：举例 A100（108 SMs）上使用 256×128 分块，矩阵从 1792 增加到 1793 时，块数从 98 跳到 120，超过 SM 数，需两轮调度。第二轮只有 12 个块，大部分 SM 空闲，性能骤降。
- 启示：矩阵维度应尽量为 SM 数量的倍数（或其约数）。
FlashAttention 解析：所有技巧的综合
注意力计算的核心操作：
S = Q @ K^TP = softmax(S)O = P @ V
朴素实现中，S（N×N）和 P（N×N）都要写入全局内存，内存复杂度为 O(N²)。
FlashAttention 的两大关键手段：
① 在线 Softmax（Online Softmax）
- 经典 softmax 需要先找最大值（求指数），再计算归一化，需要两遍扫描。
- 在线 softmax：边扫描边追踪 running max 和 running sum。
- 遇到新最大值 m'，将已累加的 sum 按
sum *= exp(m_old - m_new)修正，再累加exp(x - m_new)。
- 遇到新最大值 m'，将已累加的 sum 按
- 效果：每个块（tile）可独立计算 softmax，无需等待全局最大值，从而实现分块式注意力计算。
② 分块 + 重计算
- 将 Q、K、V 切成块（例如 Q 块、K 块、V 块），逐块加载到共享内存。
- 计算每个 Q 块对所有 K 块的部分注意力，在线更新 running max/sum，并将结果以加权方式累加到输出块。
- 反向传播时不保存注意力得分矩阵（N×N），而是逐块重新计算——这就是重计算技巧的直接应用。
- 最终，全局内存访问量从 O(N²) 降到 O(N² / T)（T 为分块大小），实现内存高效。
③ 结果
- FlashAttention 实现了一个融合的、内存高效的注意力 kernel，显著提升训练/推理吞吐量，是大模型推理（长上下文）的关键优化之一。
关键公式
- 操作强度（Arithmetic Intensity）：
- 无分块矩阵乘法全局内存读取次数（每个元素）：
次 - 分块后（分块大小为 T）：
次全局读取 + T 次共享内存读取 - 在线 softmax 修正：
（其中 为归一化累加和， 为 running max）
课程作业关联
- 作业 1/2 中会涉及与 CUDA Mode / GPU 书中练习相似的矩阵乘法分块优化。
- FlashAttention 的分块与在线 softmax 部分将是作业中的核心编程任务。
- 建议先独立完成一个朴素矩阵乘法 kernel，再逐步加入分块、合并访问与 padding，观察吞吐量变化。
常见误区
- 认为 GPU 只是"更多核心的 CPU" — GPU 的调度单位是 warp，分支会导致隐式空转。
- 忽略内存层次 — 全局内存访问的代价远高于 L1/共享内存，性能瓶颈常在内存而非计算。
- 以为低精度是"免费的" — 量化/反量化有额外开销，且第一/最后一层难以量化。
- 分块大小越大越好 — 需要对齐突发段与 SM 数量；过大可能导致波次量化。
- FlashAttention 只是矩阵乘法的优化 — 其核心难点在于 softmax 的在线计算，这是注意力分块的前提。
- 混淆 TPU 与 GPU 的"Tensor Core" — 前者指处理器，后者指矩阵乘法单元。
复习清单
- [ ] 能画出 GPU 内存层次结构并列出各级延迟相对大小
- [ ] 解释 SIMT 模型：warp 内所有线程执行相同指令
- [ ] 描述共享内存（可编程）与 L1 缓存（自动）的区别
- [ ] 计算 ReLU 在 FP32 下的操作强度（1 FLOP / 8 bytes）
- [ ] 解释 MXFP8 的缩放因子机制及转置问题
- [ ] 给出一个算子融合的实例（如 sin²x + cos²x 的融合 kernel）
- [ ] 描述重计算如何减少内存访问次数（8 → 5 例子）
- [ ] 定义合并内存访问（Coalesced Access），并说明 DRAM 突发（burst）原理
- [ ] 解释分块的收益公式（N/T 次全局读取）及波次量化现象
- [ ] 推导在线 softmax 的更新公式
- [ ] 总结 FlashAttention 的完整流程：分块 → 在线 softmax → 重计算
总结
本讲的核心信息：系统性能的关键不在计算而在数据移动。现代 GPU 算力增长远超内存带宽增长，因此所有优化策略都围绕"减少全局内存访问、提高数据复用"展开。六个技巧（控制流最小化、低精度、算子融合、重计算、合并访问、分块）是这一原则的具体体现。最后，FlashAttention 正是这一原则的集大成者——通过分块、在线 softmax 与重计算，将注意力从 O(N²) 内存访问降低到接近 O(N)，成为大模型训练与推理不可或缺的基石。
GPU 与 TPU：从硬件心智模型到 FlashAttention 的系统之旅
CS336 Language Modeling from Scratch 第五讲深入探讨了 GPU 和 TPU 的硬件架构、性能优化技巧，以及如何将这些知识整合起来理解 FlashAttention。本文基于课程逐字稿整理，带你建立从硬件到算法的完整认知链条。
引言：当魔法遇见工程
深度学习模型常常被形容为“魔法”，但 GPU 同样如此——尤其是当你第一次接触它的时候。一张展示矩阵乘法吞吐量随矩阵维度变化的图表可能看起来令人困惑：为什么在某些特定规模下，性能会突然暴跌？为什么增加一个维度会导致吞吐量显著下降？这些“魔法”背后其实都有清晰的工程原理。
在 CS336 的学习路径中，前两讲专注于神经网络本身的理论基础。而本讲标志着课程进入系统部分——这是整个技术栈中最容易理解的部分之一，因为你可以推理所有的组成部分，并有合乎逻辑的步骤引导你得到结果。但正如课程讲师所说：“当你第一次接触 GPU 时，感觉可能会有点像在跟一个既神奇又非常奇怪的设备打交道。”
理解硬件不是系统方向研究者的专利。即使你只关注架构设计或一般性地推理这些事，也需要理解模型将如何在硬件上执行，才能设计出有效的模型。正如 Percy 在第一讲中提到的，扩展的一个关键部分就是确保你有效地利用资源——不理解你的系统，你就永远无法高效地使用资源。
从 CPU 到 GPU：设计哲学的转变
为什么频率提升不再有效
回顾计算历史，大约在 2000 年代之前，让计算机变快的主要方式是提高时钟速度——晶体管变小，时钟速度变快，这种被称为“丹纳德缩放”的规律在一段时间内表现良好。但随后这种扩展方式达到了极限：晶体管变小的速度仍在持续，但让晶体管变小并不一定意味着运行得更快，其中有根本的物理原因。
那么，在这个无法单纯通过频率提升来加速的世界里，我们如何让大语言模型继续扩展呢？答案是并行性。这不仅是 GPU 的故事，也是并行计算范式的核心：不再通过让时钟越来越快来串行地加速事情，而是横向扩展，让更多的东西并行地执行指令。
GPU 的硬件心智模型
CPU 是为快速、串行执行而设计的，适用于有复杂分支逻辑和控制流的场景。它拥有非常大的控制单元，但只有少数几个 ALU（算术逻辑单元）。它的设计目标是低延迟——从收到指令到完成指令之间的时间应该非常短。
GPU 则完全不同。它更看重吞吐量——也许单个任务完成得并不快，但所有任务的总吞吐量非常大。实现这一点的关键在于拥有海量、轻量级的核心。GPU 拥有成百上千个计算单元，集成在一个芯片上，可以并行执行。
在 GPU 中，基本单元是 SM（Streaming Multiprocessor，流式多处理器） 。每个 SM 都不是一个单一的计算单元——它内部包含流处理器（Stream Processor），这些流处理器可以并行地执行不同的线程。一个 A100 有 128 个 SM，每个 SM 都可以独立编程并执行不同的任务。
但 GPU 的故事中非常重要的部分不仅仅是计算部分，还有内存部分。现代硬件和 LLM 优化实际上是由内存来定义的。
内存层次结构：速度与容量的权衡
GPU 中有不同级别的缓存和内存：
- 寄存器（Registers） ：容量最小但速度最快
- L1 缓存和共享内存（Shared Memory） ：非常快，大约 20-30 个周期即可读取
- L2 缓存：速度慢得多
- 全局内存（Global Memory）或高带宽内存（HBM） ：比 L1 缓存慢约 10 倍
当 H200 宣称“我有 144 GB 内存”时，它指的是全局内存。物理上，全局内存是位于计算单元之外的内存芯片，而 L2 和 L1 缓存位于芯片本身内部——这就是为什么它们的延迟要好得多。
你可能会问：既然共享内存这么好，为什么不把整个芯片都做成共享内存呢？答案是成本：贵几百倍，而且更耗电。你需要这种内存层次结构，并有效地利用它。
需要注意的是，共享内存和 L1 缓存虽然物理上接近，但工作方式不同。L1 缓存是自动运行的，用于存储最近访问的数据元素；而共享内存是可编程的——你实际上会与它交互，可以往里面放东西，也可以从里面取东西。
GPU 的编程模型：线程、Block 和 Warp
GPU 的执行模型中有三个重要的角色：
线程（Thread） 遵循 SIMT（Single Instruction, Multiple Threads，单指令多线程） 模型——所有线程执行完全相同的指令，但可以有不同的输入。这对代码效率有重要影响：如果线程可以做各种各样不同的事情，编程就会变得非常困难。GPU 实际上是在可编程性和效率之间的一种权衡。
Block（线程块） 是一组线程，保证在单个 SM 上运行。Block 可以访问一个共享的内存池，这在我们讨论 tiling（分块）时变得非常重要。
Warp（线程束） 是 GPU 内部运行时的调度单元，由 32 个连续编号的线程组成。这减少了调度器决定运行哪些线程的开销。
控制发散：if 语句的代价
由于 SIMT 模型的限制，在 GPU 代码中写 if 语句会有意想不到的代价。每个线程必须执行相同的指令，所以实际上每个线程都会执行 if 和 else 两个分支——但如果你不属于某个分支，你只会屏蔽掉那个计算，处于闲置状态。这被称为控制发散（Control Divergence） 。
这就是为什么 GPU 代码中实现 ReLU 时，你通常会看到在乘以零或者乘以掩码，而不是写一个只更新一半的 if 语句——因为乘法是同时全部进行的，而 if 语句可能需要等待两个时钟周期来执行两个分支。
TPU：趋同进化的另一种路径
TPU 在高层次上与 GPU 非常相似——这有点像是趋同进化：如果你想构建一个能效高的机器学习加速器，你最终会到达同样的设计方向。
TPU 更加针对机器学习工作负载进行了优化，有更轻量级的控制单元和更大的矩阵乘法单元。它们的内存结构与 GPU 非常相似：包含慢速的高带宽内存（HBM） 和一个更快的本地内存称为 SMEM（共享内存） 。
真正有趣的是，GPU 中的每个概念在 TPU 中都有一个对应的概念。GPU 中的张量核心（矩阵乘法单元）几乎与 TPU 中的 MXU（Matrix Multiply Unit） 完全对应——它们实际上使用相同的底层电路，都是所谓的脉动阵列（Systolic Array） ，通过流式输入输出来执行矩阵乘法。
但有一个令人困惑的命名问题：TPU 称它们的流式多处理器为“张量核心”，而 GPU 称它们的矩阵乘法单元为“张量核心”。如果你听到“TPU 的张量核心”，那指的是处理器；如果听到“GPU 的张量核心”，指的是矩阵乘法单元。
设计上的主要区别在于规模与数量的权衡：一个 A100 可能有约 128 个 SM 和 528 个矩阵乘法单元，而一个 TPU 只有 2 个处理器和 8 个矩阵乘法单元，但每个单元规模更大。这给了 GPU 更多灵活性——可以编程让它们做不同的事情——而 TPU 则被锁定在做大型矩阵乘法上。课程中提到了一个实际例子：在做批次大小扫描时，批次大小降到 64 就停止了，因为 TPU 的张量核心拒绝接受任何小于 64 维输入的矩阵乘法。
为什么矩阵乘法是特权操作
从 V100 开始，NVIDIA 决定在硬件中直接加入矩阵乘法单元（张量核心）。一旦那个硬件存在了，矩阵乘法就成为机器学习中一个享有特权的操作：在可并行化但非矩阵乘法的操作与矩阵乘法能完成的操作量之间，存在巨大的吞吐量差距（超过 10 倍）。这就是为什么任何能随计算规模扩展的机器学习架构都必然包含矩阵乘法——这是你能够真正有效获得大量计算吞吐量的唯一方式。
计算的戏剧性变化：不同组件的非同步扩展
GPU 的不同组件随着时间以不同的速度扩展。总体的计算吞吐量（灰色线）增长最快；内存带宽（绿色线）是数据传输速度，增长相对缓慢；并行性（蓝色线）是连接不同 GPU 设备的速度，也增长得非常慢。
这意味着什么？在 GPU 编程的早期阶段，你可能不会太多考虑内存问题，因为计算速度不比内存传输快多少。但如今，计算和内存之间的差距越来越大——这就是为什么本讲要讨论的很多优化都是内存优化。
六个关键技巧：让 GPU 跑得更快
技巧一：低精度计算
算力扩展中相当重要的一部分来自数字表示的精简。从 FP32 到 BF16 再到 INT8，你一直在降低数字精度。少了一半的比特，移动的比特也就少了一半，内存移动量就减少了。
一个简单的例子：在大小为 n 的向量上做 ReLU，如果用 FP32，读取需要 4 字节，写入需要 4 字节，总共 8 字节内存访问，却只有 1 个 FLOP 的工作量——8 字节每 FLOP。降低精度直接减半了这个比例。
但低精度在实际中如何工作？当你做矩阵乘法时，你会在取乘积之前进行降精度转换，而在累加部分和时以全精度进行累加。真正让低精度训练成熟的关键在于弄清楚哪些操作需要使用哪种数字格式：矩阵乘法中的权重和激活可以低精度，但 softmax 可能需要 FP32。
FP8 是低精度领域的下一个前沿。一旦到了 FP8，不再有一个“标准”格式——有 E4M3（4 位指数、3 位尾数）和 E5M2（更多指数位、更少尾数位）等不同选择。
更先进的 MXFP8 格式使用多个缩放因子：矩阵中的每 32 个元素对应一个缩放因子。这带来了一个有趣的问题：矩阵转置变得非常复杂，因为需要重新量化。实践中，训练时会为每个量化矩阵创建两份拷贝——一份给原始矩阵，一份给转置矩阵。这就是绕过问题的方式。
技巧二：算子融合（Operator Fusion）
想象 GPU 是一个工厂：内存仓库、小型计算工厂、传送带在来回运送东西。如果有许多不同的操作需要顺序执行，就要反复运送中间数据，产生大量内存瓶颈。如果只有一个巨大的工厂，能接收所有原材料然后输出成品，就只需要为内存带宽付两次费用。
这就是算子融合。一个简单的例子：计算 sin² + cos²。朴素实现中每个操作（sin、cos、平方、加法）都相当于一个工厂，每次调用都从全局内存读写数据。融合后的单个 GPU 内核只从全局内存读取一次，在 SM 内部完成所有操作后写回一次。
编译器（如 Torch Compiler 或 JAX Compiler）可以自动完成这类基础操作的融合。高级版本的融合有时需要手动干预。
技巧三：重计算（Recomputation）
通常在反向传播中，我们向前传播时保存激活值，然后在反向传播时使用这些激活值。但如果我们扔掉所有激活值，在反向传播时重新计算它们呢？
想象一个由三个 sigmoid 组成的简单计算图。朴素实现需要读取 x 一次，写入 S2、S1、out 三次，然后反向传播再读取三次、写入一次——总共八次内存访问。如果丢弃所有中间激活值，只存储最终输出，反向传播时临时重新正向传播并即时计算激活值，那么只需要两次内存读取和一次写入——总共五次内存访问。
用额外的计算换取更少的内存访问。在计算资源充裕但内存昂贵的情况下，这是很有价值的技巧。
技巧四：合并内存访问（Coalesced Memory Access）
全局内存是 DRAM，由于结构原因以突发（burst）模式工作。访问某个位置时，同一突发段（burst section）内的其他位置也可以免费读取。这意味着，如果读取是连续排列的、彼此紧挨着，速度会快得多。
对于行主序（row major）的矩阵，如果线程试图在每个时钟周期读取每一列，每个元素都会落在不同的突发窗口内——非常低效。但如果翻转过来读取行，每个线程都可以在一次内存读取中获得需要的数据。
这就解释了为什么矩阵的大小和排列方式如此重要。理解这一点，你就能明白为什么 Karpathy 的 nanoGPT 在把词汇表大小从 50,257 增加到 50,304 时会获得 25% 的加速——因为填充（padding）可以让突发传输段的边界与分块大小对齐。
技巧五：分块（Tiling）
这是最重要的一个技巧。核心思想是：利用内存层次结构，尽可能多地完成工作。把你需要的东西切分成块（tile），每个块加载到共享内存中，一旦加载完成就反复从该块读取，尽可能多地做计算。
在朴素的 N×N 矩阵乘法中，每个输入元素会被从全局内存读取 N 次。使用大小为 T 的分块后，每个输入元素只需从全局内存读取 N/T 次，然后在共享内存中被访问 T 次。全局内存访问量减少了 T 倍。
但分块也有微妙的问题。想象一个 128×128 的分块，如果矩阵是 256×256，一切完美——四个分块完美切割。但如果矩阵增加 1 个维度，就会产生两个非常细长的、几乎什么都没有的分块。分块大小的选择非常复杂，需要优化。
PyTorch 的 max-autotune 编译器选项就是干这个的：尝试所有不同的分块大小应用到你的矩阵上，找出哪个实际上更快。
技巧六：处理波次量化（Wave Quantization）
回到那个神秘的性能下降图表。为什么从 1792 到 1793，矩阵乘法吞吐量会暴跌？
使用 256×128 的块大小，1792 时有 98 个块。但 1793 时，两个维度上的块数都增加了一个，变成 120 个块。A100 有 108 个 SM——98 个块可以放进去，但 120 个块放不进去。处理完 108 个 SM 的活儿后，还剩 12 个块要处理，而此时大部分 GPU 都空闲着。仅仅因为你把块的大小增加了一，就需要再等一整轮循环来处理新一批块——这就是波次量化。
整合：FlashAttention
现在我们有了理解 FlashAttention 的所有元素。FlashAttention 是注意力机制方面的一大改进，而且完全是系统层面的改进——它把 PyTorch 简单的注意力实现变成了一个单一的、非常巧妙融合的内核。
注意力计算包含几个部分：Q、K、V 三个矩阵乘法和中间的一个 softmax。矩阵乘法我们知道怎么做——分块矩阵乘法。但 FlashAttention 的棘手之处在于 softmax 是全局操作，把所有元素都联系在一起。
关键技巧在于在线计算 softmax（Online Softmax） ：你一边计算一边计算归一化的 softmax。每次遇到一个比以前见过的更大的数字，就交换出最大值并修正之前的结果。除此之外，就是对之前所有内容进行在线的、累积的指数求和。因为 softmax 可以在线地、逐块地计算，我们就可以一个块一个块地处理注意力计算，而无需看到其余的分块。
FlashAttention 还包含重新计算的部分：不保存 N×N 大小的注意力激活值，而是在反向传播时逐个分块重新计算。
正如论文所述，他们只做了两件事：应用了分块和重计算，以在内存访问次数上达到次二次级别。
结语：为什么硬件理解至关重要
理解硬件，也就是 GPU，是支撑规模和算力的基础。不透彻理解到最底层的细节，就无法完全掌握为什么我们会做出某些决策。课程讲师强调：“我不希望你们成为那种机械照搬、盲目给矩阵配 32 个乘法器的人。我希望你们真正理解我们为什么这样做。”
当前算力增长的方式意味着我们需要思考矩阵乘法——这些是算术密度非常高的核心操作。由于计算与内存之间的差距，我们需要非常仔细地考虑数据移动。像 FlashAttention 这样的东西，以及一些非常硬件感知的架构工作，在做架构设计时，仔细考虑硬件对未来系统获得良好性能至关重要。
本讲在课程中的位置很特殊：它从神经网络的理论基础转向了系统实现，为后续的并行化、扩展和推理优化打下基础。Percy 的下一讲将涵盖并行化的更高级主题，但 GPU 的底层理解是所有上层优化的根基。理解这张图是什么、为什么看起来是那样的——每一处奇怪的凹陷和隆起背后，都是内存层次结构、分块策略和硬件特性的具体体现。
课程资源：
CS336 Lecture 5：完整中文逐字稿
- 视频：官方视频
- 来源：英文人工字幕提取后翻译为中文；保留段落起始时间戳，只修正明显的断句、重复和语序。
- 说明：这是逐字稿，不是讲义摘要；相邻且语义连续的片段已合并为阅读段落；术语保留必要英文。
正文
[00:05.120] 好的，我们开始吧。现在我们要进入这门课的系统部分。过去两节课我们一直在讲神经网络的内容，现在我们要开始深入探讨GPU、并行化和推理。我认为系统部分，在某种程度上，会是整个技术栈中最容易理解的部分之一，因为你可以推理所有的组成部分，而且会有合乎逻辑的步骤引导你得到结果。所以这非常不同，也很有成就感。但当你第一次接触GPU时，感觉可能会有点像在跟一个既神奇又非常奇怪的设备打交道。我不太确定。做个课堂小调查吧，你们中有多少人上过任何类型的GPU编程课？好的，这相当不错。大概四分之一，可能略少一点。但我认为当你第一次接触这些东西时，它确实有点像魔法。所以理解它很有好处。我认为深度神经网络感觉像魔法，但GPU同样也可以。
[01:12.110] 有时人们会给你看这样的图。我想向你们承诺的是，在本次课程结束时，你们会理解这张图是什么，以及为什么这张图看起来是那样的。这张图展示的是你在矩阵乘法中获得的吞吐量，其中x轴是你的矩阵维度。你可能会想，嗯，也许更大的矩阵乘法总是能获得更高的吞吐量，因为你有更多的工作要做。但实际上，你会看到所有这些疯狂的模式，在特定规模下，你的矩阵乘法会比其他的慢得多。为什么？你需要理解这一点。你还需要理解一些关于如何从概念上思考GPU的基础知识，以便构建快速算法。即使你不是系统方向的人——我猜你们中可能不是所有人，甚至不是大多数人是系统方向的人——但如果你在考虑架构，或者只是在一般性地推理这些事情，你需要理解你的模型将如何在硬件模型上执行，才能设计出有效的模型。
[02:07.790] 回顾Percy的第一节课，扩展的一个关键部分就是确保你有效地利用资源。不理解你的系统，你就永远无法高效地使用资源。而我自己并不是系统方向的人。没想到吧。所以在我开始之前，我想确保向你们推荐一些我认为非常好、非常有用，并且我在准备这节课时使用过的资源。Horace He有非常出色的关于GPU工作原理的博客和解释。还有一个爱好者团体叫CUDA Mode，他们热衷于讨论GPU和CUDA内核等等。另外还有一本TPU书，以及——什么？现在叫GPU Mode了。他们改了名字。好吧，GPU Mode。还有那本TPU书，但现在它也是Google那些人写的一本GPU书了，实际上，这真是一个令人难以置信的资源。所以如果你们对我今天讲的任何内容感到非常兴奋或感兴趣，我强烈建议你们去看看那本书。
[03:06.520] 而且他们还有练习可以做。你们的作业里会有一些和那些练习很相似的内容。所以你们可以稍微预览一下我们之后会给你们布置什么。好的。那么这节课将分为三个部分。第一部分，我认为，只是关于GPU的广泛入门材料。这主要是硬件模型，以及GPU的编程方式是如何运作的，因为GPU是一个比较陌生的设备。我希望你们从原理上理解它们与CPU的不同。一旦你们理解了这些，我们将讨论，我认为是，六个不同的技巧，重要的技巧，来让GPU跑得更快。这些只是你如何开始调整代码以便在GPU上高效运行的一些构建模块。然后最后，我们将所有这些整合在一起的部分是，我们会逐步讲解FlashAttention。
[03:59.000] 这将是我们的胜利时刻。我们会说，实际上，到此为止，我们已经理解了重新发明我们自己的FlashAttention并理解所有不同组件所需的一切。所以这就是我们在课程结束时要努力达成的目标。我们还没有讲到扩展，但从Percy的入门课或者一般地跟随语言模型的发展，我们知道如果我们手头有更多算力，我们就能得到更好的语言模型。这是我们运作的硬通货。所以如果我们有更快的硬件，或者我们有更好的利用率，或者我们有更多的芯片，那么我们就能改进并行化。所有这些都能推动进步。所以从某种意义上说，过去几年的游戏规则就是，我们如何能调动更多的算力来训练模型，我想在过去一年左右，则是如何调动更多的算力来推理这些模型。
[04:52.770] 所以算力极其重要。而系统是其中的一部分。如果我们在90年代谈论这个，我想我们会关注的是时钟速度。我们会说，嗯，让计算机变快的方式是，我们有一个串行执行指令的CPU，我们让这些指令变得越来越快。我的意思是，这种晶体管变小、时钟速度变快的扩展类型，被称为丹纳德缩放，在一段时间内表现良好。实际上，大约在2000年代，这种扩展方式就达到了极限。还有一个有趣的题外话。晶体管的数量，当然，一直在持续增加。但让晶体管变小并不一定意味着你的晶体管运行得更快。它们并不会以更高的时钟速度运行。在某个节点之后，你实际上无法继续从越来越小的晶体管中获益，这其中有根本的物理原因。
[05:44.290] 所以这种让你的时钟更快、让指令执行更快的老方法行不通了。那么在这个无法单纯通过频率提升来加速的世界里，我们如何让LLM继续扩展呢？这就是GPU和一般意义上的并行扩展登场的地方。Percy下节课会讲更高级的主题。但之后，就是并行性。并行性和GPU都属于同一个范式的一部分。不是通过让你的时钟越来越快、越来越快来串行地加速事情，而是你要横向扩展，让更多的东西并行地执行你的指令。所以这就是GPU的故事。这一张幻灯片，我认为，一张图，就总结了并行扩展的状态，那就是在K20和M40的早期，你拥有按今天标准来看可观但有点可怜的FLOPs。
[06:39.670] 然后突然间，在P100和V100时代，这条曲线开始起飞。你每年能执行的浮点运算量得到了超指数的增长。这张图在我们能够调动的算力量方面，是一个相当了不起的图。所以GPU的扩展确实推动了我们所见到的几乎所有的算力扩展。我们将讨论这里的各个组成部分。我在这里先强调一下，以便让你们有个印象。在2017年的V100时代，我们开始有了张量核心。这是重大的硬件创新之一。然后在接下来的几年里，我们看到了FPGA中的结构化稀疏性，以及所有这些更低位宽的数值格式也在推动我们的FLOPs向前发展。我们稍后会详细看看这到底意味着什么。那么，GPU 与 CPU 有何不同？那些已经接触过 GPU 的同学应该已经知道了。我们一会儿就会讲到更高级的内容。
[07:38.620] 但对于其余同学来说，我认为理解它们的区别所体现出的模型，或者说某种意义上的设计哲学，是非常重要的。CPU 是为快速、串行执行而设计的，适用于有复杂分支逻辑、复杂条件判断、复杂控制流的场景。正因为如此，CPU 通常拥有非常大的控制单元，并且只有少数几个 ALU（算术逻辑单元）来执行数学或其他运算。但它的设计目标就是非常快速地运行。所以你想要尽可能快，并且延迟要低。从你收到一条指令到完成这条指令之间的时间应该非常非常短。这就是 CPU 的一般设计哲学。相比之下，GPU 更看重的是吞吐量。我认为右边这张图非常形象地说明了，当你改变这种设计哲学时，最终会得到什么样的结果。
[08:26.142] 在 GPU 中，也许你分派一个任务，它需要很长时间才能完成。可能是你处理一个任务，然后停下来，再去处理另一个任务，然后再回来处理之前的任务。但总的来说，尽管你的单个任务完成得并不快，但你所有任务的总吞吐量是非常大的。而实现这一点的关键在于拥有海量、轻量级的核心。GPU 拥有成百上千个计算单元，它们被集成在一个芯片上。所有这些计算单元都可以并行执行。所以，与 CPU 相比，这是一种非常不同的编程和思维方式。我想揭开 GPU 的神秘面纱。我不希望你们把它想象成 NVIDIA-SMI 输出里的一堆文本。我希望你们能够在脑海中想象出 GPU 的硬件结构。我认为这张图通常被用来形象地说明 GPU 的内部结构。
[09:21.428] 我认为这是描述 GPU 内部运作方式的绝佳逻辑图。GPU 的基本单元是 SM（流式多处理器）。流式多处理器有点像一种核心。它是一个独立的计算单元，拥有可以加速其计算的子组件。并且它可以访问某些特定类型的内存。所以，基于这种并行性的理念，每个 SM 并不是一个单一的计算单元。它内部包含流处理器，这些流处理器可以并行地执行不同的线程。但 SM 是离散的计算单元。我们稍后讲到编程模型时会明白为什么。你几乎可以把每个 SM 看作一个核心。而 GPU 拥有大量的 SM。一个 A100 有 128 个 SM。每个 SM 都可以独立编程，并且可以独立执行不同的任务。
[10:15.560] 它们还可以访问外部的全局内存。还有另一件我们将会看到的事情——那就是 GPU 的计算部分。但我们在整个讲座中会看到，这个故事中非常非常重要的一部分不仅仅是计算部分，还有内存部分。现代硬件和 LLM 优化实际上是由内存定义的。因此，我们需要了解这些内存是什么，以及它们位于何处。当我们谈论内存时，有不同种类的内存。与 GPU 连接的有不同级别的缓存。还有寄存器，它们容量更小但速度更快。然后你还有全局内存或高带宽内存。从这个表格可以看出——这是一些人做的基准测试，我记得这是在 A100 上的延迟数据——L1 缓存和共享内存非常快。大约在 20 到 30 个周期内，你就可以从中读取数据。
[11:14.220] 如果访问 L2 缓存，速度就慢得多；而如果你访问全局内存，那比 L1 缓存要慢 10 倍。这对你的速度来说是一个巨大的差异。这种全局内存通常就是高层程序员所认为的 GPU 内存。当 H200 告诉你我有 144 GB 内存时，它指的是全局内存。而不是共享内存或 L1 缓存。那么，物理上这些内存位于何处呢？通常，全局内存是那些位于计算单元之外的内存芯片，而 L2 和 L1 缓存则物理上存在于芯片本身内部。这是一张 GPU 芯片的图片。你可以看到所有这些 SM。这些绿色的块就是 SM。你可以在这里看到 L2 缓存。而 L1 缓存则位于每个 SM 内部。所以它们在物理上更靠近计算单元。
[12:07.070] 这就是为什么它们的延迟要好得多。你可能会说，好吧，既然共享内存这么好，为什么不把整个芯片都做成共享内存呢？那要贵得多，贵几百倍。而且根据这些内存的制造方式，它也更耗电。所以这就是为什么如果你要构建一个实用的加速器，通常你不会只想拥有巨量共享内存。你需要这种内存层次结构，并有效地利用它。我认为最近被 NVIDIA 收购的 Grok，其设计就是只使用这些巨量的 SRAM。对于推理等某些工作负载来说，这非常有用。但尽管如此，对于大多数加速器来说，你必须尊重这种内存层次结构，才能让任何东西跑得快。是的。共享内存在 L1 缓存里吗？是的。所以共享内存和 L1 缓存，我想，它们的工作方式是不同的。缓存实际上只是一个缓存东西，用于存储最近访问的数据元素。
[13:10.255] 共享内存则是一个可编程的元素，你实际上会与它交互，你可以往里面放东西，也可以从里面取东西。好的。明白了。所以这就是内存的存储部分。请记住这一点，因为内存层次结构会不断被提及。这是一个非常重要的对象。现在，我们已经建立了硬件的思维模型。接下来我们将开始讨论 GPU 的软件和编程模型。GPU 的执行模型中有三个重要的角色。一个是线程。如果你熟悉任何形式的多并行编程，你就会知道线程是什么。线程是一个轻量级的东西，它并行运行并执行一些工作。在 GPU 中，线程的一个重要特点是它们遵循 SIMT（单指令多线程）模型，这意味着你拥有的所有线程都执行完全相同的指令。它们会有不同的输入，这意味着线程实际上可以执行不同的有用工作，但它们执行的指令必须是相同的。
[14:11.960] 这对你的代码效率以及编程模型都有重要影响。如果线程可以做各种各样不同的事情，编程就会变得非常困难。GPU 是在可编程性和效率之间的一种权衡。接下来，block（线程块）是一组线程。block 之所以是一个相关的对象，是因为一个 block 被保证在单个 SM 上运行。记住，SM 有共享内存。所以你可以把线程看作是并行执行工作的。而 block 可以访问一个共享的内存池。当我们稍后讨论一种叫做 tiling（分块）的概念时，这一点将变得非常重要，因为这些 block 将重用相同的内存片段。最后，warp（线程束）本质上是 GPU 内部运行时的调度单元。线程将以 32 个连续编号的线程为一组来执行。这就是一个 warp。
[15:04.430] 这减少了调度器决定运行哪些线程的开销。但本质上，这三个术语是GPU最重要的概念。我可能会说，哦，我们将有不同的 warp 在执行。这就是我在这里的意思。这是一个由32个线程组成的基本单元，它们一起执行。对。你提到所有线程执行相同的指令。是 block 中的所有线程还是 warp 中的所有线程？一个 warp 中的所有线程都会被执行。很好。是的。所以调度器决定下一条指令将发送给哪个 warp。没错。或者抱歉。调度器决定哪个 warp 接下来执行，以及执行什么任务等等。好的。那么在 GPU 的内存模型中，设备代码可以从寄存器中读取。寄存器是一些最快、最本地化的项目。这可能存储类似数组起始和结束地址这样的内存地址。
[16:04.530] 这是你可能在寄存器中存储内容的地方。然后你有本地内存。这也是本地的，就在你的单个 SM 内。你还有共享内存，可以在你的 block 内访问。所以如果你有需要在线程之间传递的东西，或者你想要跨线程复用的信息，这些都将保存在共享内存中。当然，你还可以从全局内存读取，也就是远端的 DRAM，只要你愿意承担等待数据从全局内存传输过来的延迟代价。还有常量内存，我觉得在各种编程工作中用得不太多，你也可以使用。你甚至可以在主机内存之间来回传输。当我说主机内存时，指的是位于机器上的 CPU 内存。如果你需要卸载超出 GPU 容量的数据，你可以这样做。这就是你可以在内存方面做的一系列操作。
[16:58.800] 而这个内存模型的关键点，也是最重要的部分，就是一旦你超出共享内存的范围，一切都会变慢。因此，分组 block 以减少全局内存读取次数将成为贯穿本次讲座的核心策略。现在，我将用三张幻灯片、两张幻灯片来讲解 TPU。本次讲座的大部分内容是关于 GPU 的，主要是因为我觉得如果你们出去编写代码、构建模型，大多数人可能会在 GPU 上进行。但我认为讨论 TPU 非常重要，一方面是因为我觉得 TPU 很酷而且越来越流行，另一方面也因为 TPU 是 GPU 的另一种演进路径。你可以看到什么是相同的，什么不是。了解 TPU 的工作方式也非常有趣。事实上，非常有趣的是，TPU 在高层次上与 GPU 非常相似。
[17:58.400] 这有点像是趋同进化。如果你想构建一个能效高的机器学习加速器，你最终会到达同样的设计方向。所以你在头脑中应该记住的 TPU 和 GPU 之间核心思维差异是，TPU 在某种意义上更简单。它们更加针对机器学习工作负载进行了优化。它们有轻量级的控制单元，并且有更大的矩阵乘法单元。所以它们有和 GPU 相同的基本结构，我接下来会讲到，它们会有一个专门的电路来做矩阵乘法。它们会有能够执行并行向量运算的组件，还会有一个控制系统。它们有非常相似的内存结构，其中包含慢速内存，称为高带宽内存，还有一个更快的本地内存，称为 SMEM 或共享内存。
[18:50.830] 所以这看起来和 GPU 的架构非常非常相似。真正的区别在于这些组件的大小以及这些操作的灵活性。还有，哦，我今天不深入讨论的主要区别是这些加速器的网络连接。所以如果你思考 CPU 和 GPU 之间的最大区别是什么，实际上最大的区别在于网络连接，而不是单个芯片本身，因为芯片实际上就是为了矩阵乘法而生的。这来自 JAX 书籍，这真是一本很棒的参考资料。但他们有这两个表格。他们有一个关于 GPU 的部分，描述了如何基本上将 GPU 的概念映射到 TPU 的概念。今天我讲的左边这些内容，你都可以直接映射到 TPU 上。但在本次讲座的这一点上，最重要的事情是注意，GPU 中的每个概念在 TPU 中都有一个对应的概念。
[19:46.620] 这种映射实际上非常精确。GPU 中的张量核心，也就是矩阵乘法单元，几乎与 MXU 完全对应，MXU 是张量核心中的矩阵乘法单元。如果我们深入探讨非常细节的问题，它们实际上使用相同的底层电路。它们都是所谓的脉动阵列，通过流式输入输出来执行矩阵乘法。所以架构非常相似。那么它们是如何——哦，抱歉。是的。你知道 GPU 和 TPU 上典型的 [听不清] 尺寸有多大吗？哦，我不知道，我一时想不起来。抱歉。我需要去查一下。是的。所以这两者之间的主要区别在于我们所拥有的不同单元的数量。哦，实际上，在我继续之前，我要提一件非常非常令人困惑的事情。
[20:33.980] 抱歉。TPU 称它们的流式多处理器为张量核心，GPU 称它们的矩阵乘法单元为张量核心。它们的命名完全一样。所以你必须根据上下文来区分它们。如果有人说 TPU 的张量核心，那指的是处理器。如果有人说 GPU 的张量核心，他们指的是矩阵乘法。所以如果你看一个 GPU，像 A100 这样的典型 GPU 可能有大约 132 个流式多处理器。这是很多处理器。一个 TPU 只有两个单元。只有两个独立的单元在并行运行。那么 GPU，如果我们只看矩阵乘法单元，会有 528 个矩阵乘法单元，而 TPU 只有 8 个。
[21:19.840] 这些单元的规模大致相当，看起来是可比较的对象。为什么会这样呢？嗯，TPU 依赖的是更大但数量更少的矩阵乘法单元，而 GPU 依赖的是更小但数量多得多的矩阵乘法单元。你可以想象，如果你有很多较小的矩阵乘法单元，这会给你更多的灵活性。你可以编程让它们做这个做那个。而对于 TPU，你被锁定在做大型矩阵乘法上。所以在我最近写的一篇论文中有一个有趣的情况，我们做了批次大小的扫描，结果一路降到 64。它就在 64 停止了。为什么？因为张量核心拒绝接受任何小于 64 维输入的矩阵乘法。它必须是大的矩阵乘法。所以这是主要区别。但正如你所看到的，GPU 和 TPU 有很多相似之处。它们都依赖于快速内存和慢速内存之间的内存层次结构。
[22:12.370] 它们都有矩阵乘法单元。正因为如此，我今天教你们的核心理念将适用于 GPU 和 TPU。基本思想总是相同的，就是出于这个原因。这确实是趋同进化。能够经济有效地分配内存的方式是有限的，快速执行矩阵乘法的方式也是有限的。有些设计决策是不同的。但除此之外，你最终会到达同样的地方。好的，酷。所以这就是TPU。对，很好。而GPU，我认为，已经取得了巨大的成功，因为它使得像矩阵乘法这类事情可以轻松扩展。如果你想要更高的吞吐量，你只需要不断增加SM。只要你的内存带宽能够处理这些SM，你就能获得越来越高的吞吐量。编程是我们不太考虑的事情，但它其实出奇地简单。
[23:02.150] 如果你试图编写更底层的、看起来像机器码的东西，比如PTX，这实际上是可行的，因为这一切都是SIMT。这并不是说你要对每一个线程单独编程。你只是说，好的，这是我的指令，这是我的输入，现在在所有输入上运行它们。这可能是你们已经熟悉的事情，比如使用或应用函数式编程。最后，我还没有怎么谈论线程。但无论何时你在进行GPU编程，这些线程总是非常轻量级的。它们可以随时被停止、启动。调度器可以决定，好的，我有另一个更适合运行的线程束。所以我要停下你，然后启动这个。交换这些任务的速度非常快。如果某些任务停滞或具有某些降低吞吐量的属性，这就能实现更高的吞吐量。所以这就是GPU作为硬件模型的优势。它已经取得了巨大的成功。
[23:57.580] 我认为，即使在张量核心实现之前的早期，人们就已经意识到，商品化的图形硬件，这种大规模的并行性，对于科学计算会非常有效。所以这是一篇关于使用图形硬件进行快速矩阵乘法的早期论文之一。这真的是非常非常酷的黑客式工作。人们想明白了，嗯，我们有这些已实现的着色器，但我们可以编写特定的着色器，实际上能给我们带来矩阵乘法。而且不同的渲染设置实际上能给我们带来更快的矩阵乘法。这是你能做的非常酷的计算机科学的事情。但如今，你不需要这样做了。你不需要手动编程你的GPU来执行矩阵乘法，因为从V100开始，NVIDIA决定：这里有块硬件，它将为你执行矩阵乘法。
[24:48.330] 一旦那个硬件存在了，矩阵乘法就成为机器学习中一个享有特权的操作。所以在这一点上，在可并行化但非矩阵乘法的操作与你在矩阵乘法中能完成的操作量之间，存在一个巨大的吞吐量差距。所以这就是为什么我们看到的任何近期能随计算规模扩展的机器学习架构，都会包含矩阵乘法的一个原因，因为这是你能够真正有效获得大量计算吞吐量的唯一方式，比你在GPU上能执行的任何其他浮点运算快10倍以上。然后我要讲的最后一个硬件相关内容——我认为这是这一部分的最后一张幻灯片——是GPU的不同组件随着时间以不同的速度扩展。我们知道GPU正变得越来越快。就像CPU一样，总体的计算吞吐量正变得越来越快。那就是灰色的线。
[25:54.010] 那相当快。我们在总浮点运算方面获得越来越快的改进。绿色的是我们的内存带宽。这是数据传输的速度。它的增长相对缓慢。蓝色的是并行性。这是连接不同GPU设备的速度。这也增长得非常缓慢。这意味着什么呢？这意味着如果你处在GPU编程的早期阶段，你可能不会太多考虑内存问题，因为你的计算速度并不比内存传输快多少。但现在，随着我们在这个图的右边走得越来越远，我们将看到计算和内存之间的差距越来越大。这就是为什么我今天要讲的很多优化都是内存优化。随着我们进一步向右走，我们会越来越多地看到内存和通信瓶颈。因此，要充分利用我们现有的硬件，将需要越来越多的努力。所以这就是我们对GPU非常基础的概述。我不可能用20分钟的讲座就让你完全理解GPU硬件是什么。但希望你们有了一个初步的认识。
[26:56.920] 心智模型应该是：GPU是极其极其并行的。它们同时执行指令。计算规模非常快，真的非常非常快，甚至比内存还要快，这意味着内存非常重要。因为内存很重要，我们所做的一切都必须尊重内存层次结构，我们做尽可能多的事情让它进入共享内存而不是全局内存。所以这希望是你们对GPU的理解。如果你们有这个理解，至少你们将掌握必要的基础知识，以便内化这次讲座的下一部分。好的。在进入下一部分之前，我可以在这里暂停一下，以防大家有硬件问题或其他任何问题。好的。所以在V100之前，没有矩阵乘法单元？是的。没有专门用于矩阵乘法的张量核心。对的。张量核心是V100的一个补充。在那之前，你需要手动编程。
[27:50.860] 当然，你有那么多SM和那么多ALU，你仍然可以获得相当好的吞吐量。但是一个专门的矩阵乘法电路确实改变了这个游戏规则。哦，是的。后面的同学。共享内存和L1缓存之间有什么区别？缓存是自动运行的。你无法控制L1缓存。（听不清）说L1只是因为物理距离近。还是说它还有其他不同的地方？是的，我的意思是，物理距离是其中一个很大的因素。我的意思是，如果我没记错的话，它们都是SRAM。这也是全局内存慢得多的原因之一。但物理距离和所使用的互连是导致速度较慢的原因。是的。顺便说一句，这并不是必然的。另外一点是，TPU中相当于L2缓存的东西要快得多，因为他们在芯片上做了不同的设计权衡来实现这一点。
[28:42.830] 那是他们的卖点之一。是的。如果计算和内存之间的差距持续增长，你认为这会影响未来芯片的布局方式吗？是的。或者可能是未来芯片的预算分配。有一件我真的很不想谈但现在你们逼我谈的事情是，推理领域的硬件实际上更加疯狂。所以人们开始做的一件事，我想是去年，叫做prefill和decode分离架构，其中prefill部分，这是非常重的矩阵乘法，在一个芯片上，而推理的decode部分，更受内存带宽限制，在另一个芯片上。而且有些模型，我想是Step 1-3，这是一个中国的开源模型，甚至有一个尝试分离不同层的方法。所以注意力机制会去往一个为某件事优化的芯片或加速器，然后MLP会去往另一个不同的加速器。
[29:37.810] 所以你开始看到，如果这个差距继续增长，任何利用你非常宝贵的快速内存的聪明技巧都会变得非常有价值。你开始在推理中非常明显地看到这一点，推理比训练更加受内存限制。是的，很好。哦，对。一个 block 里有多少个 warp？我知道吗？我可不认为我能脱口而出，我觉得这也是跟硬件相关的。好的，很好。那么这节课的第二部分就是，现在你们已经对 GPU 是什么以及它如何工作有了一个大概的了解。现在我们要讨论所有那些让工作负载在 GPU 上跑得飞快的基本技巧。这些是一些基本组件，你可以把它们组合起来让事情变得更快。为了保持我们的动力，我本节结束时的目标之一就是向你们解释这张图。
[30:32.700] 所以这张图，正如我之前说的，看起来非常简单。横轴是我正在相乘的方阵的大小。我在相乘两个方阵。纵轴是我每秒处理多少条目。所以这是吞吐量。随着我向右移动，总的来说，吞吐量在上升。因为我有更大的矩阵，我有更多的工作要做。所以如果我有更多的工作，我就能让我的硬件得到更充分的利用。但这里到处都是各种奇怪的图案。我们想要理解为什么这里会出现这些奇怪的图案，以及我们如何才能在这条曲线上处在最顶端而不是底端。这就是我将在第二部分要谈论的内容。Percy 已经提到了这一点。这就是 roofline 模型。但我之前展示的那张图看起来有点像这个经典的 roofline 模型。
[31:19.370] 这看起来像是有一个对角线部分，然后是一个平坦的部分。这正是我们所期望的。Roofline 模型表明，在某个点之前，我们会受限于内存。我们是内存受限的。所以无论你给我多少计算能力，我的吞吐量都不会上升。在某个点，你有足够多的工作量对应于每次内存移动，以至于我可以让我的计算单元完全饱和。超过那个点之后，拥有越来越大的工作单元就不会再帮助我了。无论你给我多少额外的工作，那都没用，因为我已经完全饱和了我的矩阵乘法单元。所以这就是我们从经典 roofline 模型中所期望的，当我们有更快的内存时，我们可以把这个往上推。但先是内存受限，然后是计算受限。这大致是我们预期的。
[32:12.170] 现在，如果你想确保你的代码在 GPU 上高效运行，实际上你需要做的是，你需要处在这个平坦的部分。一旦我们到了这个平坦的部分，就没有别的可以做了。我们达到了最大吞吐量。所以我们希望避免自己处在这个对角线部分，这意味着我们想要增加操作强度。我们每次内存读取所做的计算量需要足够高。所以这就是我们需要做的。我们将讨论六个不同的技巧。实际上，其中一些并不是关于内存的技巧，只有第一个是。这六个不同的技巧可能是理解如何让 GPU 跑得快所必需的。第一个技巧的风格有点不同，但其余的都会带有这种尽可能减少内存移动的风格。第一个是 GPU 特有的东西。
[33:00.942] 它与我之前提到的其中一件事有关，即 SIMT 模型。所以在 GPU 中，每个线程都在执行相同的事情。那么，如果你在一个 GPU 上执行的代码里写一个 If 语句，会发生什么？在 CPU 上，会发生什么很清楚。在 CPU 上，你要么选择其中一个分支，要么选择另一个分支，你将执行那个分支。然后你就继续干你的事。在 GPU 中，你会做一些非常不同且稍微有点反直觉的事情。每个线程都必须执行相同的指令。所以实际上，每个线程都会执行 If 和 Else 两个分支。但如果你不属于 If 分支，你只会屏蔽掉那个计算。你将只是闲置。所以每当你有一个条件语句时，你对执行的思维模型看起来有点像这样。你的一部分线程将根据这个 If 语句分叉。
[33:49.062] 其中一些将进入下面的分支，一些将进入上面的分支。下面分支的先执行，上面分支的等待。然后上面分支的执行，下面分支的等待。所以为了执行一个 If 语句，你实际上是把这两个分支都走了一遍。然后一些线程只是在等待，如果它们处在分支的错误一侧，就什么都不做。这被称为控制发散。这在某些方面很好，因为它简化了编程。不好的方面在于，你现在有了这些大的计算间隙，你的 GPU 在空转。这就是为什么你真的不想写 If 语句的原因之一。如果你看 GPU 代码，如果你有一个 ReLU，你通常会看到你在乘以零或者乘以掩码，而不是写一个只更新一半的 If 语句，因为乘法是同时全部进行的。
[34:38.460] 而一个 If 语句可能实际上不得不等待两个时钟周期或者类似这样，两个操作。所以这就是第一件事。但是，这和所有那些内存的事情是两回事，那些内存问题是这堂课的主要叙事或者说主要故事。第一个技巧，你其实不需要做太多事情，但我认为它是最重要的一个，而且在硬件方面，我认为 NVIDIA 投入了最多精力，那就是低精度计算。如果你看这张图——我把 Bill Dally 又请回了幻灯片上——这个超指数增长，这到底是从哪里来的？实际上，其中相当重要的一部分来自数字表示。你从 FP32 开始，然后到 BF16，然后到 Int8。你一直在降低你的数字精度。当然，如果你少了一半的比特，你移动的比特也就少了一半。
[35:30.390] 所以你的内存移动量就减少了。这样你至少在某种程度上解决了内存瓶颈。这是一个很傻的例子。但我就是想把这个概念牢牢钉进去。假设，我在一个大小为 n 的向量上做一个 ReLU。如果我使用的是 float 32，我需要做一次读取来读入我的数字。在我计算出 ReLU 之后，我需要做一次写入。那将是每个操作 4 字节，我的内存访问总共有 8 字节。我只有一个 FLOP 的工作量。所以那是 8 字节每 FLOP。如果我降低精度，当然，我把每个操作的字节数减半了。我把每个 FLOP 需要移动的字节数减半了。抱歉，这是算术强度的倒数。算术强度应该是这个的倒数。那么，这是一个非常简单的例子。但实际上，低精度在现代 GPU 中是如何工作的，实际上非常复杂，而且根据你觉得事物有多酷，这要么非常酷，要么非常疯狂。
[36:34.110] 所以当你做矩阵乘法时，比如说，你有一个以较低精度运行的张量核心，这在实践中实际上是如何工作的？嗯，实际上你并不会把所有东西都保持在低精度。你会在取乘积之前进行降精度转换。而当你在矩阵乘法中累加部分和时，你会以全精度进行累加。然后你可能以 FP32 格式输出你的输出或结果。我认为，低精度实际上有点像是门黑艺术，部分原因不是因为你可以随便降低精度。任何人都能做到这一点。你必须弄清楚你的哪些操作需要使用哪种数字格式。所以对于矩阵乘法，也许你的权重和激活都可以是低精度的。如果我正在做 softmax，也许我需要在 FP32 下进行。或者如果我正在做指数运算，也许那也需要 FP32。
[37:22.160] 或者也许我可以直接用BF16就够了。所以我认为，真正让低精度训练达到今天这样的成熟状态，确实花了好几年的时间。而且其中大量的工作是缓慢的、循序渐进的、经验性的，就是去摸索哪些运算可以降精度，怎么降，以及如何在尽可能让训练和推理的大部分环节都保持低精度的同时，维持训练的稳定性。人们一直在把这件事往前推进。我想一旦你考虑到高速内存的成本有多高，你就会意识到，要是能把精度再砍掉一半就好了。大家真的、真的很想做到这一点。所以FP8可以说是低精度领域的下一个前沿。一旦到了FP8，就不再有一个大家都用的、标准的、规范化的格式了。
[38:17.030] 如果你用16位，那你很可能大部分场景都在用BF16。但对于FP8，会有E4M3，也就是四位指数位和三位尾数位。或者也可能有E5M2，指数位更多，尾数位更少。因为这些东西用在不同的地方，一旦位宽低到这个程度，就没有一个放之四海而皆准的方案了。但我觉得真正有意思的是，随着新硬件一起出现的那些更先进的低精度数值格式。所以我简单讲讲MXFP8，因为它既很酷，而且它给训练带来的影响，至少可以说，相当奇特。在FP8里，传统的训练方式是你所有的，比如说，激活值都是8位的。所以就在上面这里。然后你会有一个FP32的缩放因子，用来把这个值上下调整。为什么需要缩放因子呢？因为比如说你只有四位指数位。
[39:17.870] 指数位这么少，很快就会上溢或者下溢。所以你需要缩放因子来让它保持在合适的范围内，这样才能保留一些信息。然后人们意识到，这并不是最优的，因为一个矩阵内部可能有量级差异非常大的部分。我序列的一部分可能比另一部分有大得多的激活值。所以与其用一个单一的缩放因子，为什么不用很多个缩放因子呢？所以MXFP8，以及这些更先进的、带多缩放因子的格式，现在都有很多个缩放因子。每一个缩放因子对应着相应的、用颜色标记的子矩阵。所以这个矩阵中的前四个元素对应第一个缩放因子，以此类推。现在我有更多缩放因子了，而这些缩放因子本身会是更低精度的。它们也将会是8位的。它们会是E8M0格式。所以他们只有指数位。
[40:15.420] 所以它们都是2的幂次方，用8位表示。所以你必须在这些方面做权衡。那么，我刚才想说什么来着？好的。所以MXFP8在元素本身会使用更多的尾数位。缩放因子它们自己是FPA。而且你的矩阵里每32个元素就会有一个缩放因子。好的，我在这里暂停一下。先别读第三个要点。现在，想一想。这个设计有什么问题？如果你做了这么精巧的设计，最大的问题是什么？这里最大的问题是——比如说，我想转置这个矩阵。转置后的矩阵，它的缩放模式，可以说，跟原矩阵的缩放模式是不一样的。以前矩阵转置很简单。我直接转置这个矩阵就行了。不需要重新量化。什么都不用做。但转置这个矩阵现在变得非常复杂，因为我需要可能重新量化整个矩阵，才能符合我那种每32个元素一个缩放因子的模式。
[41:16.320] 所以当你真正使用MXFP8的时候，内部发生的事情——我觉得这既疯狂又有点酷——就是当你用这个东西训练模型时，它实际上会为每个量化矩阵创建两份拷贝，一份给原始矩阵，一份给转置矩阵。所以如果你任何时候想转置，转置版本就已经准备好了。这就是绕过这个问题的方法，我觉得这是个又疯狂又酷的事情。这里的重点也许不是强调这些GPU的冷知识。重点是，一旦你开始接触这些FPA、FP4格式，事情就变得非常奇特了，因为你丢失的指数位和尾数位数量是相当极端的。所以你必须非常巧妙地处理这件事。而且就像我说的，如果你真的要去训练一个模型，FPA，我认为很多前沿实验室为了算力收益正在这么做，你必须得非常小心。
[42:08.670] 这张图确实展示了我刚才说的。每次你量化，你会用一种方式量化原始矩阵，然后你会转置一次再量化。所以你会同时拥有两份拷贝。然后另一部分，就是我几页幻灯片之前告诉你的，你只量化那些你认为安全的特定层。这是另一个你必须通过试错来摸索的事情。但一旦你把所有这些组合起来，你就能得到训练FPA的巨大好处，那就是当你在做那些矩阵乘法的时候，你很可能能节省20%到30%，甚至更多，取决于你矩阵的大小。不幸的是，它不会带来2倍的加速，因为你必须做所有这些量化操作。做这个有很大的开销。所以它不是免费的。是的，我记得读到过第一层和最后一层非常难量化。
[42:53.760] 你对这个有什么直觉上的理解吗？没错，说得对。我对第一层没有太多直觉。至于最后一层。我觉得，难量化是因为它是一个一阶的因素。它驱动了很大一部分的损失。所以我觉得如果你试图量化最后一层，你会同时遇到不稳定性和损失大幅增加的问题。我之所以把这个放在上面，几乎是因为我觉得这太离谱了。下一个前沿，这个已经被实现为MXFP4了。MXFP8的全部内容可以在这张幻灯片上展示给你。就是负6到6。你拥有所有可能的值。但现在你理解了MXFP8，你也能理解这个东西是怎么运作的了。你会有一个结构化的块，每个块是这些四位中的一个。每16个这样的数字，你会有1个缩放因子。那个缩放因子本身会是FP8。
[43:47.080] 这些缩放因子是E4M3格式的。已经有一篇论文用FP4训练过了。我还没听说有人真正在现实中用FP4成功训练过大型、严肃、真正的模型。但我认为这正在到来。下一代模型很可能会用FP4。所以这就是低精度的情况。我觉得你确实能把极限推得这么远，最后只剩下正6到负6，却还能用这个精度下的矩阵做有趣的事情，这真的很酷。哦，是的，有个问题。所以这只适用于矩阵运算吗，还是说你可以只应用于单个标量步骤？是的，你可以量化任何东西。比如，你可以量化ReLU之后的激活值。只是没有那么大必要而已。比如像矩阵乘法，如果你量化它，你可以获得非常大的吞吐量提升。这就是为什么你看那个MXFP8训练的例子，全都是矩阵乘法被量化到了FP8。
[44:52.570] 对于大多数[听不清]来说，这个开销可能不值得。这跟硬件有关系。不是说矩阵乘法硬件部分只做这个。嗯，它不是只做这个的。是的，你总是可以用更低精度的方式来做，这会减少内存移动。这确实没错。但对其他方面来说，这就不太说得通了。这有什么硬件限制吗？这不是限制。但硬件确实为所有这些提供了支持，所以这是内置在硬件里的。其中很多东西都内置于硬件中，旨在加速这些操作。是的，你把这一点定位为对内存带宽的改进，或者说减少内存带宽需求。但就这一点而言，有什么好处吗？是的，当然有，确实如此。对于计算而言，仅仅对量化后的数字做乘法，就能获得近乎线性的提升。但因为你必须进行量化和反量化，所以这些好处会被稀释掉。
[45:43.110] 是的。你需要训练缩放因子吗？缩放因子是一个取决于具体实现的选择。根据不同的库，你可能会得到不同的量化和反量化方法。“训练”这个词用在这里有点重。但你可能会有量化因子，它会扫描数据并选取最大值、最小值。或者你可能会有缩放因子，它查看历史的运行统计数据，然后使用这些统计信息。我不会真正把这些称为训练。在推理时，你可能会更精确地拟合它们。但那里没有梯度。在推理时，你可能会。在训练时，你不会。是的。在结构化稀疏方面有没有做过什么工作？是的。因为我猜英伟达有过。是的。在结构化稀疏方面已经做了很多工作。我的意思是，MoE，我认为，是人们在使用的结构化稀疏操作的一个成功例子。
[46:27.110] [听不清] 是的，没错。所以我的意思是，已经做了很多工作。Chris Ré 在多种不同类型的结构化矩阵方面做了很多工作。我认为在计算收益和表示损失之间的权衡方面，这些工作的效果已经趋于平淡了。我不认为这是个坏主意。但从经验上看，我认为它未必取得了预期的成功。哦，是的。我们一直在讨论压缩模型，讨论量化模型已经几十年了。目前的主流方法是什么？是在全精度或半精度下训练，然后蒸馏成一个量化模型，这样你就能在苹果电脑之类的东西上运行它了吗？现在的情况是怎样的？[听不清]？是的。我的意思是，这结合了几种方法——你很可能想要训练一个更大的模型，然后进行量化压缩。有各种关于量化、扩展定律的论文，人们已经写了很多，精确地描述了这里面的权衡取舍。
[47:22.880] 我认为目前最先进的方法涉及到在某种程度上结合量化感知训练，以及模型扩展与各种类型的训练后量化。是的，我觉得这里面有很多东西还没有被完全理解透彻。即使是我交谈过的业界团队，他们也在谈论仍在进行量化方面的科学研究。是的，好的。那么下一个技巧——这个更像是概念性的编程层面的东西——是算子融合。这是个相当简单的想法。但你会惊讶于这件事在多少情况下还没有被实现。把 GPU 想象成一个工厂。你有你的内存仓库。你还有你的小型计算工厂。而且你有一条传送带在来回运送东西。如果你扩大你的计算规模，你就会被你的小传送带卡住瓶颈。如果你有很多操作，这种情况会更糟。
[48:10.780] 想象一下，我有很多不同的操作需要执行。我要把我的原材料运到我的工厂，然后来回、来回、再来回地运送。这会带来很多内存瓶颈，你需要为此付费，并且需要耗费大量双向的内存带宽。那么，如果我只有一个巨大的工厂，它能接收所有原材料，然后把我的成品运回来，这样不是更好吗？这样，我只需要为内存带宽付两次费用，这是个非常简单的想法。这就是算子融合。这种情况非常常见。如果你编写一个非常简单的计算，比如说，计算正弦平方加余弦平方，那么 PyTorch 的计算图会是什么样子？它看起来像这样。你有你的输入 x，我计算它的正弦，我计算它的余弦。我把它两边都平方，然后我把它们加起来，就得到了我的输出。天真地看，如果你不做任何智能的优化，那么这里的每一个操作都相当于一个工厂。
[49:03.230] 每调用其中一个操作，你都会从全局内存中读取和写入数据。而且每次这样做你都会产生相当多的内存开销。所以这相当痛苦。你可以想象，与其这样做，为什么不直接编写一段 GPU 代码，让它只从全局内存读取一次，然后在我之前谈到的 SM 内部完成所有操作。然后在 SM 内部完成后，我再将它写回全局内存，这一个内核就完成了所有这些操作。这被称为“融合”。任何像这样简单的东西——这是一个非常简单的融合。这就像一个计算图，我可以直接把它压缩成一个单一单元——编译器可以自动为你完成这件事。所以 Torch 编译器（它是 PyTorch 的编译器），或者 JAX 编译器（它是 JAX 的编译器，并且与 JAX 集成得更深入），这两者都会将这类基础操作融合成一次单一的 CUDA 调用。
[50:01.580] 我稍后会更详细地谈谈融合。高级版本有时需要手动干预。但对于简单版本，你只需认为编译会帮你解决这些问题。第三个想法是重计算。我们知道在反向传播中，我们会怎么做呢，嗯，我们向前传播，在树状结构上向上计算。这是来自 Percy 的 221 课程的内容。你沿着树向上进行。同时你保存那些激活值。你存储这些小的激活数值。然后当我反向传播时，我会使用我存储的激活值来计算沿着计算图向下传递的反向信号。我们都对反向传播有一定了解。所以我们都明白这是什么。但同时，我们现在看待这个问题不是从数学的角度，即如何计算梯度，而是从系统的角度来看待它。
[50:48.297] 我们如何最小化内存使用？好吧，如果我们想想会发生什么，假设我有一个简单的计算，x -> sigmoid -> sigmoid -> sigmoid -> out，这里的每一个 sigmoid 操作都会在计算图中产生一些需要追踪的残差或激活值。所以在我第一个操作之后，我有 S2，第二个操作之后，有 S1。然后当我要进行反向传播时，它会消耗掉 S2、S1，以及我的输出，或者抱歉，是输出处的梯度。我会将这个梯度反向传播，一直推到我输入 x 的导数。所以如果我数一下内存读写的次数，那就是读取 x 一次，写入 S2、S1、out 三次。
[51:26.400] 然后在反向传播时，情况就反过来了。所以这总共是八次内存读写，算术强度非常低，因为我只是堆叠了一堆 sigmoid 函数。那么，我还能做些什么呢？我能做的是，我可以直接扔掉所有的激活值。通常你不会这么做。我为什么要重新计算某些东西呢？但想象一下，你处于一个计算成本超级、超级便宜的世界。你只有非常充裕的计算资源，你的计算单元没有全部用上。但你的内存非常昂贵，任何一次读写都有很大的代价。那么你可能会更喜欢这个计算图：我读取一个 X，然后我执行 sigmoid、sigmoid、sigmoid，得到输出。
[52:04.098] 我会存储这个输出。我必须写入这个，因为这是我计算图的输出。但现在在反向传播时，我的反向计算图会得到 D_out。并且我会临时地重新正向传播 X，然后在我需要的地方即时计算出激活值。所以现在对于 D_out 和 Dx，这将是两次内存读取，以及一次针对 Dx 的内存写入。那么我就从总共八次内存访问降到了五次。现在我的总内存访问量是原来的 5/8，而计算量保持不变。但我额外花了一点计算资源，节省了一点内存带宽。这是另一个非常重要的技巧。我们通常不会想到这一点，因为一般来说，我们只关心如何计算梯度。但一旦你有了系统思维，你就会开始考虑重计算（recomputation）。好的，第四个技巧。这又是一个关于内存的知识点，但令人惊讶的是它变得很重要，同时也说明了为什么理解你所使用的硬件模型有时很重要。
[53:05.330] 全局内存（global memory），也就是位于芯片外部的大容量内存，是 DRAM。而 DRAM，由于其结构方式，实际上是以突发（burst）模式工作的。这意味着，当我访问，比如说，第一个位置时，我实际上得到的不仅仅是第一个数值。我可以免费读取这前四个值。这是因为内存单元的排列方式，存在这些内存单元块。然后当我激活一个放大器时，我就可以几乎免费地读出所有这些值。它的并行结构方式允许我读取内存块。作为一个硬件模型，这是合理的。这被称为突发段（burst section）。所以每当一个段被访问时，所有其他位置也可以免费被传送出来。
[53:51.680] 可以说，你不必为它们付费。所以你可能会有比如 128 字节的突发大小。这意味着一次读取可能会返回 128 字节的数据，只要它们位于同一个连续块内。这意味着，如果你的读取是连续排列的，彼此紧挨在一起，那么速度会快得多。用一个术语来说，这种内存访问被称为合并的（coalesced）。这个词的意思是，如果一个线程束（warp）中的所有线程都落在同一个突发段内。那这意味着，假设我有三个线程。我所有的三个线程，抱歉，四个线程，四个线程，而我的四个线程位于一个突发段内，或者说它们正在从这个突发段读取数据。那么这就是合并访问，因为它们用满了整个这个突发段。并且我充分利用了 DRAM 内存这个非常好的特性。如果我有多个线程束，那么我会有多个不同的突发段。
[54:53.170] 但是你可以通过将内存读取分组在一起来高效地读取内存。是的。为什么内存访问会出现这种模式，像这个 T 形状的模式？哪一个，像这样的吗？这个。是的。据我所知，内存单元排列在一个网格中。而决定读取哪些单元的放大器是按列操作的。为此激活电压通常是一个需要花费大量时间的过程。所以一旦你选定了列，读取该列内的多个元素就会相对便宜得多。但这是为了选择、读取上面那个红色的单元格。是的，没错。所以如果我想读这个，比如说，我选中了这一列。或者至少在这个图里，列是便宜的。但如果我试图读取这一行，我可能会免费得到这一整行，因为我寻址了整个这一行。
[55:50.965] 你不会免费得到一个 T 形的。是的。好的。很好。你可能会说，Tatsu，你为什么要跟我讲 DRAM 的冷知识？我告诉你这些 DRAM 知识的原因是——想想矩阵。对于矩阵，你通常想要读取大块的数据，因为你要做矩阵乘法（matmul）。而矩阵乘法中某些类型的内存访问是具有优势的。它们比其他访问方式结构要好得多。有时你可以利用这一点来实现非常非常快的内存读取。所以我记得的口诀或者说要点是，如果你有线程，并且你有一个行主序（row major）的矩阵，即按行排列，索引递增方向是沿着行的，那么沿着主轴线移动的线程不是合并访问。
[56:39.200] 为了让它更直观——实际上，我喜欢右边这张图。想象我有一个 4 乘 4 的矩阵。这是行主序的。所以它的顺序是黄色块在最前面，红色块其次，以此类推。这是内存地址中的线性顺序。这是我要存储的矩阵。现在，想象我有一些线程，试图在每个时钟周期读取每一列。所以我读一列，向右移动，再读一列，再向右移动。如果在第一次迭代时我这样做，我会访问到这个元素、这个元素、这个元素、这个元素。注意，这些元素中的每一个都落在不同的突发窗口内。而且我已经读取了整个向量，可以说，或者整个矩阵，抱歉，可以说，尽管我想要的只是一列。
[57:26.620] 现在，想象一下如果我反过来，如果我翻转这个矩阵，我可以读取整个块。每个线程都可以在一次内存读取中获得它需要的数据。所以合并访问可以显著提升从大型连续结构（如矩阵）读取数据的性能。这就是为什么你可能需要考虑行、列以及其他排序方式的一个原因。我想谈的最后一个——这是个大问题，也许是（目前提到的这些想法里）最重要的一个，我把它留到了最后——就是分块（tiling）的思想。我认为这个可能对性能的影响最大。我认为，分块比任何其他想法都更能体现“利用内存层级结构尽可能多地完成工作”这一理念。我们想要做的是，尽可能地分组我们的内存访问。共享内存（shared memory）的访问应该一起进行，并且应该在共享内存上完成。之所以叫分块，是因为我要把我需要的东西切分。
[58:28.150] 我要对我的空间进行分块。每个块（tile）将被加载到我的共享内存中。一旦这个块被加载，我将反复地从该块读取，并尽可能多地做计算。让我们考虑矩阵乘法的情况。我想将这个矩阵和这个矩阵相乘。右下方的这个矩阵是我的输出。这是我要写出结果矩阵的地方。这是比较标准的表示方法。现在，如果我想做这个，我在遍历线程时会多次读取——抱歉。同一个条目会被多次读取。准确地说，在朴素的 N 乘 N 矩阵乘法中，每个条目会被读取 N 次。所以你可能会想，如果我多次读取这个条目，比如说 M₀，为什么我不能把所有这些内存访问分组到一起，而不是每次都访问全局内存呢？矩阵的一个好处是，存在一种非常自然的分组方式。
[59:22.810] 这基本上就是把矩阵切成子矩阵，然后加载它们，再对它们进行操作。这里是一个很简单的算法。我要把矩阵切成 2 乘 2 的子矩阵。然后我要加载不同的块。所以我首先加载 M₀,₀ 块，也就是左上角这个。我要加载 N₀,₀ 块，也就是上面那个矩阵的左上角。我要把它们加载到共享内存中。所以我要为从全局内存缓慢传输这些数据支付一次代价。但是一旦它们在共享内存中，我就可以非常快地读写这些块。这个输出块目前也会在共享内存中。我将完全在这些子矩阵上进行运算。所以我会将这两个子矩阵相乘。并且我会把部分和放入我的输出子矩阵中。这样我就能够在快速的共享内存中完成大量的读写操作。
[01:00:14.300] 然后等计算完成后，我可以把结果写回去。这里面的数学是什么样的呢？好吧，在没有分块（tiling）的情况下，我已经告诉过你们了，当我在做大尺寸 N 乘 N 的矩阵乘法时，输入中的每个元素都会被从全局内存中读取 N 次。如果我使用大小为 T 的分块，那么每个输入元素需要从全局内存中读取 N/T 次，然后在每个块内部再被访问 T 次。这很合理，考虑一下极端情况。如果 T 等于 N，那么我从全局内存读取它一次，然后在共享内存中访问它 N 次。这很棒，因为全局内存访问量减少了 T 倍。如果你有非常大的共享内存分块，这意味着在像矩阵乘法这样的操作上，你可以极大地减少全局内存访问。
[01:00:58.160] 好的。行。是的。你这么做是不是因为没法把整个块都放进去？是的，是的。没错。但如果你在构建那些非常花哨的替代加速器，比如像什么 RPU 之类的，全都是 SRAM，那么是的，你直接把整个矩阵放进去就行了。那样整个计算过程会非常快，但是非常昂贵，非常非常昂贵。是的。而且这种情况在 GPU 层面也会发生，因为有接触式充电（contact charging），我甚至无法把一个 [听不清] 装进一个里面。然后你就不断叠加那个单元，一次又一次。是这样运作的吗？你的意思是，你是在问这种拆分是否也发生在加速器层面？[听不清] 在 GPU 之间进行分块。
[01:01:41.390] 这个 [听不清] 将会是一个分块。没错。是的。所以数据并行（data parallel），你可以这样理解。如果我的数据在这个轴向上，权重在这个轴向上，如果我沿着其中一个轴进行切割，那就有点像是数据并行；张量并行（tensor parallel）则是另一种并行方式。所有这些并行方式在某些方面看起来都有点像分块。是的。好的。哦，对了。那么对于之前那个问题，我能否假设像 PyTorch 这样的库会自动完成这些操作？是的，是的。就像你调用的任何矩阵乘法内核（matmul kernel）都会在底层进行分块。是的，没错。如果你在做一些非常规的操作，你可能需要考虑分块之类的、正在发生的事情。
[01:02:18.490] 你会希望能够局部地重用你的读取。但我得说，即使是 PyTorch 或类似的工具也不是完美的。所以我们稍后会讨论一些会出现的不足之处。所以，分块作为一个想法非常非常强大，但它也很复杂。实际上，你可能会遇到一些由分块引起的非常微妙的问题。想象一下，我有一个大小为 128, 128 的分块。如果我的矩阵是 256 乘 256，那一切都很好。我有四个分块，完美切割。现在，想象一下我把矩阵的维度增加 1。现在，我生成了两个非常细长的分块，里面基本上什么都没有。所以现在你意识到，实际上，这个分块大小的选择其实非常复杂。
[01:03:06.020] 根据矩阵的不同，我可能需要调整分块大小，以便在内存带宽和这些巨大的、什么都不做的空白区域之间进行权衡。分块大小实际上是一个需要你去优化的东西。所以，合并内存访问（coalesced memory accesses）、你加速器中共享内存的大小，以及你矩阵的尺寸，所有这些都会影响哪个分块是最优的。如果你们了解 PyTorch 中的不同选项，那里有一个叫做 max-autotune 的编译器选项。我不知道你们有多少人用过这个或听说过这个。你打开那个标志。在接下来的 15 分钟里，PyTorch 会运行一系列没完没了的基准测试。然后你会意识到，实际上，它正在做的是，尝试所有不同大小的分块应用到你的矩阵上，并试图找出哪一个实际上更快。
[01:03:53.140] 你会发现，通过真正优化这些东西，你能获得非常可观的性能提升。你会发现另一个与合并（coalescing）相关的事情是，如果你的矩阵尺寸不合适，你尝试读取一个分块时可能会触发大量的内存读取。所以想象一下，我有一个理想的世界，我有，额，让我们说，这个蓝色的是行。我想读取这个漂亮的分块。而我的突发传输段（burst sections）的宽度就是这个分块的宽度。在这种情况下，如果我做四次读取，就搞定了。因为突发传输段的原因，我读取了整个子矩阵，所以通过合并（coalescing）实现了良好的内存读取。现在，想象一下我把矩阵大小移动 1。现在我把这第一行给挤过去了一点。
[01:04:37.400] 所以这个蓝色框现在移到了这里。这样一来，我无法通过合并（coalesce）我的读取来一次性读取一个分块，因为现在所有这些数据都是参差不齐且不对齐的。所以我需要至少读取两个分块大小的数据才能读取出这一个分块。所以，一旦你开始考虑合并（coalescing）和所有其他这些事情，优化分块读取实际上会涉及到很多复杂性。因此，在这种情况下，你需要做的就是对矩阵进行填充（padding），基本上是为了让突发传输段的边界与你所拥有的分块大小对齐。这就是为什么你最终会看到这样的东西，看起来非常神秘。我认为这是我最喜欢的奇怪例子之一。Andrej Karpathy，在他做 nanoGPT 速度挑战的时候，他说，当我把词汇表大小从 50,257 增加到 50,304 时，我获得了 25% 的加速。
[01:05:31.250] 为了获得加速而不得不进行填充，这有点奇怪。但一旦你意识到像这样的事情，或许就说得通了。填充可以开始移动行和列的位置，这样你实际上可以获得非常高效的读取，而在其他时候，你则做不到。好的，请讲。有没有一个通常应该考虑的通用分块大小？你不需要特别去考虑分块大小本身，因为分块，通常你应该把它看作是一个偏底层系统层面的东西，除非你自己在编写内核（kernels）。但你确实需要考虑你矩阵的大小。你希望你的矩阵大小能够被分块大小整除。对吧。所以你希望它是——我稍后会详细说这个，但对，一般来说——是 2 的幂，最好也能被 32 整除。
[01:06:15.690] 好的。那么现在我们可以来谈谈这里的这个矩阵谜题了。我们基本上已经拥有了理解这个问题的所有必要信息。所以我们将逐一分析。我们已经讨论过计算强度（compute intensity）和屋顶线图（roofline plot）。你需要为每次内存读取提供足够的计算量，以便能够使你的计算单元达到饱和（saturate）。这就是推动这个对角线上升趋势的原因。所以这一条已经解决了。现在，让我们谈谈其他的线条。第一个是分块（tiling）。如果我们看这些线条中的每一个，并根据矩阵的可整除性来给它们上色，你会发现蓝色的那些基本上只能被 1 整除。
[01:06:55.570] 它们甚至不是偶数。橙色的可以被 2 整除。绿色的是被 8 整除，红色是 16，然后是紫色是 32。你可以看到，那些只能被 1 或只是 2 整除的那些矩阵，实际上吞吐量非常差。但一旦你开始达到更高的倍数，你基本上就不会再有收益递减——抱歉，是不会有性能损失了。16 或 32 表现的一样好。这并不仅仅是因为 2 的幂有什么魔力或类似的原因。这是因为 16 和 32 已经足够大了，以至于它具有正确的突发窗口（burst window）属性。它能恰好容纳整个突发窗口。因此，你可以为你的整个分块实现合并（coalesced）读取。
[01:07:38.620] 而且它对齐方式有重大影响。现在，如果我们看看剩下的部分，如果我们看这条橙色的线，我们有时会看到这种奇怪的周期性行为，性能会突然下降。这非常奇怪，因为我只是在做矩阵乘法。为什么仅仅增加一个额外的维度就会突然让吞吐量暴跌？我认为，利用我们学到的所有知识来理解这个现象，其实是挺有意思的。Percy 把这称为 GPU 冷知识。我想他不太喜欢这个，但我很喜欢。那么，为什么从这里到这里，吞吐量会下降这么多呢？嗯，它恰好发生在 1792 到 1793 之间。这是为什么呢？嗯，我使用的是 256 x 128 的块大小。这意味着——有多少个块呢？在 1792 时是 98 个块。但一旦我增加到 1793，实际上，我将两个维度上的块数都增加了一个。
[01:08:35.130] 所以我现在有 120 个块。之前有 98 个块，现在有 120 个块。这确实不好，但也不至于这么糟糕。这个下降是巨大的。事实证明，之所以如此戏剧化，是因为 A100 有 108 个 SM。98 个块可以放进 108 个 SM 里。但 120 个块放不进 108 个 SM。所以当你处理完 108 个 SM 的活儿后，你发现，嗯，我还有 12 个块要处理。我想，我得处理它们。而此时你的大部分 GPU 都空闲着。所以你会再次遇到这种非常奇怪的现象，仅仅因为你把块的大小增加了一，突然间，你就需要再等一整轮循环来处理一批新的块。
[01:09:17.250] 这被称为波次量化。以上就是第二部分的全部内容。我们如何在 GPU 上让机器学习负载跑得更快？嗯，很大程度上，这个故事的核心就是关于内存，内存，内存。那么，我们如何减少必须进行的内存访问次数呢？好吧，让我们将内存读取合并起来。让我们将操作融合在一起，这样就不需要读那么多次。对于那些无论如何都必须读取的数据，我们要重用内存，所以把数据移到共享内存中并尽可能重用，这就是分块。最后，也许你可以用其他资源来换取内存。那就是量化和重计算。所以所有这些想法都非常紧密地围绕着这个核心观念：内存是我们的瓶颈，我们该如何应对。
[01:10:04.080] 好的。最后这部分有什么问题吗？是的。我觉得我们可以直接增大 SRAM 来摆脱 DRAM。那太贵了。我是说，有两个原因。一是制造 SRAM 本身就很贵，而且也很难，因为它必须在物理上（与计算单元）靠近。信号传播很困难。我认为第三个原因是，SRAM 和 DRAM 只是不同的存储机制。SRAM 需要一直供电才能保持数据，非常耗电。所以如果你考虑总拥有成本和能耗，你可能不想要一个巨大的 SRAM 芯片。你真的应该只在需要的地方使用它。是的。这些东西能扩展到超过 32 吗？或者说，以 64 为块大小操作会更好吗？
[01:10:51.960] 你的意思是，比如如果我用更大的块会有帮助吗？是的。不。因为我觉得只要它能被你所谓的突发大小整除，那就没问题。这是一个整除性的属性。在其他现象中，你可能会从更大的整除因子中获得好处，但这里不是这种情况。是的。为什么是 108？那只是 SM 的数量。你得去问黄仁勋了。我无法为你解答。是的。让我们回到硬件方面，抱歉，不是这部分。我刚刚脑子里突然想到一个问题。你们研究过 GPU 和 Cerebras 的晶圆级引擎相比吗？是的。我觉得那边那位（指某人）对晶圆级（技术）比我了解得多。我们在上次讲座后讨论过，那个编译过程有多困难，因为如果你有晶圆级的东西，那么就会有波干扰效应和各种各样可怕的、乱七八糟的事情发生。
[01:11:46.170] 那是我完全不熟悉的世界。但我认为它比我们这里讨论的要复杂得多。好吧，酷。好的。那么现在我认为我们可以谈谈某种意义上的压轴戏了，那就是，我们已经学了很多关于 GPU 的知识。如果我们能把所有这些知识联系起来，最终能够重现或理解我认为非常重要且先进的东西，那就太好了，那就是 Flash Attention。Flash Attention 是注意力机制能力方面，我认为，一段时间以来的一大改进。而且它完全是系统层面的改进。它从 PyTorch 简单的注意力实现变成了一个单一的、非常巧妙融合的核函数。这来自原始论文。你会看到，在服务这些模型时，延迟有了显著的改善。而且正如我在两节课前向你们展示的，Flash Attention 是一种很好的、内存高效的计算方法。
[01:12:33.950] 所以你可以做越来越大规模的注意力计算。而且你可以看到，很多收益来自于从 HBM（即全局内存）传输的内存量的减少。所以他们到底做了什么？如果你读这篇论文，他们会说，好吧，我们只做了两件事。我们应用了分块和重计算。然后我们以一种能够在内存访问次数上达到次二次级别的方式来实现它。我想我们现在已经理解了所有这些组成部分。让我们看看实际发生了什么。记住，注意力计算包含几个部分。它有 K， Q， V。我们先用 K 和 Q 做矩阵乘法。然后最后用 V 做矩阵乘法。中间还有一个 softmax。所以这些是我们需要处理的操作——三个矩阵乘法和一个 softmax。我们有矩阵乘法。我们现在知道怎么做了。所以我们要做分块矩阵乘法。论文中的图 1 实际上就是分块矩阵乘法。我们以前见过这个。
[01:13:34.700] 我们把输入，也就是 K 和 Q，切分开来。然后我们在一个块上逐块地做乘法。计算完成后，我们把结果写出去。这里没有任何花哨或特别的地方。Flash Attention 比较棘手的地方，你们会在作业里看到，是对于整个结果有一个大的全局 softmax。所以我们不能就这么简单地做，因为 softmax 是一个全局操作，把所有元素都联系在一起。所以它把不同的块连接在一起。如何对注意力进行分块可能看起来不那么明显。关键的技巧是——一旦你理解了这一点，这确实是核心的想法或观察——一旦你看到了这个，剩下的就非常简单了。
[01:14:15.020] 通常，为了计算 softmax，你会按照图里展示的这种方式来做。也就是，你会先对所有的元素求指数。为了数值稳定性，你会减去最大值。然后你就得到了标准的 softmax 归一化。现在，有一种不同的方法可以在线计算这个。我认为一旦你意识到这一点，它其实很简单，那就是，你打算一边计算，一边计算归一化的 softmax。每次你遇到一个比你以前见过的更大的数字，你就要交换出那个最大值。因此，你有一种在线计算 softmax 的方式。当然，唯一棘手的地方在于最大值，而我们处理最大值的方法是，我们始终跟踪最大值。
[01:14:54.217] 然后每次遇到新的最大值时，我们就交换它并修正之前的结果。除此之外，这实际上就是对之前看到的所有内容进行在线的、累积的指数求和。当然，在最后，你只需要把你累积的所有结果除以你在线计算的那个累加值。关于 softmax 最重要的一点是，因为它是可以在线的、逐块或逐段地进行计算，所以我们就可以一个块一个块地（在分块矩阵乘法中）计算它。我不需要看到其余的分块就能计算出这个。我可以为我的分块计算 softmax。然后我可以保存部分结果，也许将它们写回全局内存，然后继续处理后续工作。这样我们就能把所有不同的交互部分拆解成更简单的组件。我认为这来自 FlashAttention 2 或者 3，应该是 2。我喜欢这个可视化，它展示了哪些部分正在被分块处理。
[01:15:46.090] 虚线框表示在 SRAM 中完成的分块操作。蓝色框表示位于高带宽内存（即全局内存）中的部分。你会看到——比如说，我们计算 K 和 Q 的内积。这些是分块矩阵乘法。它们存储在 SRAM 中，我们将其保留在 SRAM 中。然后我们计算这个指数运算。类似地，在这个分块中，我们将计算这个运行中的 softmax。接下来我会移动到下一个分块，做同样的操作。但现在我要随着推进保持 softmax 的部分运行和。我可以将它们保存在共享内存、寄存器或其他地方。需要跟踪的变量非常少。最后，在我计算完所有这些指数化数值和归一化因子之后，只需要简单地以分块形式与 V 相乘，最后再做除法。
[01:16:33.710] 这让我能够将注意力操作切分成这些小块，然后依次执行。最后一步是融合（fusion）。你要尽可能将两者同时进行，避免来回搬运数据。这个我不会详细讲解。但你可以做的是，在反向传播时重新计算你需要的所有内容。Flash Attention 的最后一部分是，如果你保存了激活值，你仍然需要保存这个 N 平方大小的激活值用于注意力计算。与其保存这些，不如全部丢弃，然后在反向传播时逐个分块重新计算。所以它也包含了重新计算的部分。这样就构成了 flash attention。现在我们已经看到了所有不同的技巧是如何运作的。
[01:17:19.780] 基本上就是这些。你看到的很多关于高效实现和如何用好 GPU 的技巧，归根结底都是这类想法。所以总结一下，理解硬件非常重要。硬件，在我们的场景中也就是 GPU，是支撑规模和算力的基础。我们需要真正理解到最底层的细节，才能完全掌握为什么我们会做出某些决策。我不希望你们成为那种机械照搬、盲目给矩阵配 32 个乘法器的人。我希望你们真正理解我们为什么这样做。在我看来，这需要一直追溯到硬件层面。另外，当前算力增长的方式意味着我们确实需要思考矩阵乘法。这些是算术密度非常高的核心操作。由于计算与内存之间的差距，我们需要非常仔细地考虑数据移动。
[01:18:13.360] 最后，像 flash attention 这样的东西，我觉得非常酷。还有一些非常硬件感知的架构工作。在做架构设计这类事情时，仔细考虑硬件对未来系统获得良好性能至关重要。谢谢。[掌声]
CS336 2026 · Lecture 5 · GPUs, TPUs
- 视频：官方视频
- 字幕：en-US · manual captions
- 处理：连续字幕合并为自然段；每段保留首条字幕的起始时间戳
- 状态：原语言字幕稿，待校对
正文
[00:05.120] OK. We will get started. Now we're going to start the systems portion of the class. So we've been doing a bunch of neural network stuff for the last two lectures. And now we're going to start getting into GPUs, parallelization, and inference. And systems, I think, is going to be, I think, part of the stack that is most reasonable in a way, in the sense that you can reason through all the pieces and there will be kind of logical steps that you can go through to get the result. So it's very different and satisfying. But when you first start interacting with GPUs, it will feel a little bit like both magic and a very strange device that you're interacting with. I don't really know. As a kind of class survey, how many of you have taken a GPU programming class of any kind? OK. So this is pretty good. That's a fourth, maybe a little bit less. But I think when you first interact with these things, it's actually a little bit of magic. So it's good to understand. And I think deep neural networks feel like magic, but also, GPUs can too.
[01:12.110] Sometimes people will show you plots like this. I think my promise to you is that by the end of this lecture, you will understand what this plot is or why this plot looks the way it does. And this is how much throughput you're getting on a matrix multiply where the x-axis is your dimension. And you might think, well, maybe bigger matrix multiplies are always more throughput because you have more work to do. Well, you've got all these kind of crazy patterns happening where, at certain sizes, your matmuls are much, much slower than others. Why? You'll need to understand that. You also need to understand some basics of how to mentally think about a GPU in order to build fast algorithms. And even if you're not a systems person-- I think probably not all of you, or even a majority of you, are systems people. But if you're thinking about architectures or just generally reasoning about these things, you need to understand the hardware model of how your model is going to execute in order to make an effective model.
[02:07.790] Thinking back to Percy's initial lecture, a big part of scaling is making sure you use your resources effectively. And without understanding your systems, you will never be efficient at using your resources. And I am not a systems person. Surprise. And so before I start, I want to make sure to point you towards resources that I think are very good and very useful and that I used in making this lecture. Horace He has really incredible blogs and explainers of how GPUs work. There's an enthusiast group, CUDA Mode, that loves to talk about GPUs and CUDA kernels and so on. And there's a TPU and-- what? GPU Mode now. They changed their name. OK, fine, GPU Mode. And the TPU book, but now it's also a GPU book by the Google folks, actually, this is a really incredible resource. So if you're really excited or interested by any of the stuff I talk about today, I would strongly encourage you to look at that book.
[03:06.520] And they also have exercises that you can work through. And your assignment will have things that look a little bit like those exercises. So you can get a little bit of a preview of what we will send you to later. Great OK. So this lecture is going to be in three parts. The first part is really, I would say, just broad introductory material on GPUs. This is going to be just the hardware model, just how the programming of the GPU kind of works because the GPU is kind of a foreign device. And I want you to understand philosophically how they're different from CPUs. And once you do that, we're going to talk about, I think, six different tricks, important tricks, to get a GPU to go fast. These are just kind of building blocks of how you can start to manipulate your code in order to make it run efficiently on a GPU. And then finally, the last part where we put all of this together is going to be, we're going to walk through FlashAttention.
[03:59.000] And this will be our victory lap. We will say, actually, at this point, we understand everything we need in order to reinvent our own FlashAttention and to understand all the different components. So that's going to be the goal that we're going to be building towards by the end of lecture. So we haven't gotten to scaling yet, but from Percy's introductory lecture or just generally following language modeling, we know that if we have more compute on hand, we can get better language models. This is the currency by which we operate. So if we have faster hardware or if we have better utilization or we have more chips, and then we have improved parallelization. All of these can drive progress. And so in some sense, the game for the last few years has been, how can we bring more compute to bear in training models, and I guess in the last year or so, how can we bring more compute to bear in the inference of these models.
[04:52.770] So compute is just extremely important. And systems is part of that. So if we were talking about this in the '90s, I think the things that we would be focusing on is clock speed. We would be saying things like, well, how we make computers go fast is, we have a CPU that executes instructions serially, and we make those instructions go faster and faster. I mean, this type of scaling where transistors got smaller and clock speeds got faster, called Dennard scaling, held quite well for a while. Really, the 2000s was about when this kind of scaling tapped out. And one interesting side note. The number of transistors, of course, has continued to go up. But making transistors smaller does not necessarily mean that your transistors go faster. They don't move at a faster clock speed. There's fundamental physical reasons why you don't actually continue to get benefits of smaller and smaller transistors at a certain point.
[05:44.290] So this old approach of making your clocks faster, making instructions execute faster, is not going to work. So how do we get LLMs to continue to scale in this world where we can't just make things faster in this frequency sense? And that's where the GPU shows up and parallel scaling, in general. Percy's going to talk about more advanced topics next lecture. But after that, it's going to be parallelism. Parallelism and GPUs are both kind part of the same paradigm. Instead of making things go faster serially with your clock going faster and faster and faster, instead you're going to scale horizontally, have more things executing your instructions in parallel. And so this is the story of GPUs. This one slide, really, I think, or one picture, summarizes the state of parallel scaling, which is, in the early days of K20s and M40s, you had respectable but somewhat pitiful by today's standards, flops.
[06:39.670] And all of a sudden, around P100s and V100s, this curve takes off. And you get super exponential scaling in the amount of floating-point operations that you can do year by year. This is kind of a remarkable plot in the amount of compute that we can bring to bear. And so GPU scaling has really driven almost all of the compute scaling that we've seen. And we'll talk about all of the different components that we have here. I'll highlight them here just to put them into your head. At the V100 era in 2017, we start getting tensor cores. That's one of the big hardware innovations. And then in the next few years, we're seeing structured sparsity in FPGA and all these lower number formats also driving our flops forward. And we're going to see what that means in detail in a little bit. So how is the GPU different from a CPU? Those that have interacted with the GPU already know this. We'll move on to more advanced things in a minute.
[07:38.620] But for the rest of you, I think it's very important to understand the model, the philosophy in some sense, of how they're different. So a CPU is designed for quick, serial execution where you have complex branching logic, complex conditionals, complex control flows. Because of this, a CPU, generally speaking, has really big control units. And it has a few ALUs that can do mathematical or other operations. But this is meant to run very quickly. So you want to go as fast as possible and have low latency. The time between when you get an instruction and when you complete it should be very, very short. That's the general CPU design philosophy. In contrast, the GPU is really about throughput. And the plot on the right, I think, is a great pictorial illustration of, when you change these philosophies, what you end up getting.
[08:26.142] So in a GPU, maybe you dispatch a task, it takes a long time for it to finish. It might be that you process a task, and then you stop, and then you go work on another task, and then you go back and work on another task. But in general, even though your tasks don't finish quickly, you're getting large aggregate throughput across all of your tasks. And what enables this is having tons and tons of lightweight cores. So a GPU has hundreds and hundreds of compute units that are packed into a chip. And all of these can execute in parallel. So this is a very different mode of programming and thinking compared to a CPU. And I want to demystify what a GPU is. I don't want you to think of it as a set of pieces of text on NVIDIA-SMI. I want you to be able to visualize what the hardware is of a GPU. I think this picture is, I think, usually used as a pictorial illustration of what a GPU is.
[09:21.428] I think this is a great logical picture of what's happening inside of a GPU. So a GPU, the basic unit of a GPU is SM, a streaming multiprocessor. And a streaming multiprocessor is kind of a core. It's its own independent compute unit that has subcomponents that it can use to accelerate its computation. And it has access to certain kinds of memories that it can connect to. So each SM, building on this idea of parallelism, it's not a single compute unit. It has streaming processors inside of it that can then execute different threads in parallel. But an SM is the discrete compute unit. We'll see why later when I get to the programming model. Each SM, you can think of as a core, almost. And GPUs have tons of SMS. A A100 will have 128 SMs. Each of these can be independently programmed, and they can independently execute different jobs.
[10:15.560] And they have access to global memory outside of it. And the other thing that we will see-- that's the compute story of a GPU. But we will see throughout this lecture that a really, really important part of this story is not just the compute part but the memory part. Modern hardware and LLM optimization is really defined by the memory. And so we need to understand what the memories are and where they live. So when we talk about memory, there's different kinds. There's different levels of caches that are connected to a GPU. There are registers, which are even more limited and even faster. And then you have global or high-bandwidth memory. And as you can see from this table-- this is a benchmarking that some folks did of, I think this was at A100 latencies-- the L1 and shared memory, these are very fast. In something like 20 to 30 cycles you, can read and get data from them.
[11:14.220] If you're hitting the L2 cache, much slower, and if you're hitting global memory, that's 10 times slower than the L1 cache. This is a big difference in your speed. And this global memory is kind of what you normally think of high-level programmer as a GPU memory. When H200 tells you I've got 144 gigs of memory, you're talking about global memory. You're not talking about shared or L1 cache. Now, physically, where do these things live? Well, global memory usually are these kind of memory chips that live outside of your compute unit, whereas L2 and L1 cache, they kind of live physically in the chip itself. So this is the picture of a GPU chip. And you see all of these SMs. These green blocks are SMs. And you see L2 right here. And the L1 is inside of each of SMs. So they're physically much closer to the compute units.
[12:07.070] That's why their latency is so much better. And you might say, OK, then if shared memory is so nice, why don't you build a whole chip out of shared memory? It's much more expensive, hundreds of times. It's also much more power hungry from the way that these memories are built. And so this is the reason why if you're building a practical accelerator, often you don't want to just have a giant amount of shared memory. You want to have this hierarchy of memory and use it effectively. I think Grok, which was recently acquired by NVIDIA, was a design where you just have these giant amounts of SRAM. And for certain workloads like inference, this is very helpful. But still, for most accelerators, you're going to have this hierarchy that you're going to really have to respect in order to make anything run fast. Yes. [INAUDIBLE] be shared memory in the L1 cache. Yeah. So the shared memory in the L1 cache, I guess, operate differently. The cache is really just a cache thing that is storing recently accessed data elements.
[13:10.255] Shared memory is a programmable element that you'll actually interact with, that you can put things in and out of. OK. Cool. So that's the memory store. And keep that in mind because it's going to keep coming back, the memory hierarchy. That's a really important object. Now, we have the hardware mental model in place. And now we're going to start talking about the software and programming model of a GPU. So there's three different important players that live in the execution model of a GPU. So there's a thread. And if you're familiar with multiparallel programming of any kind, you know what a thread is. A thread is a lightweight thing that's running in parallel and doing some work. In a GPU, an important thing about the threads is they follow the SIMT model, which means that all of the threads that you have are executing the same exact instructions. They're going to have different inputs, which means that the threads can actually do separate useful work, but the instructions they execute have to all be the same.
[14:11.960] And this has important implications for both the efficiency of your code as well as the programming model. If the threads can do all sorts of different stuff, it becomes very hard to program. GPUs are this tradeoff between programmability and efficiency. Now, a block is a group of threads. And the reason why a block is a relevant object is that a block is guaranteed to run on a single SM. And remember, an SM has a shared memory. And so you can think of a thread as doing work in parallel. Blocks have access to a shared pool of memory. That will become very important later when we talk about this idea called tiling, where these blocks are going to be reusing the same pieces of memory. Finally, a warp is essentially the scheduling unit of what is happening inside the GPU. Threads are going to execute in groups of 32 consecutively numbered threads. That is a warp.
[15:04.430] This decreases the overhead of a scheduler deciding which threads to run. But essentially, these are the important three terminology items for GPU. I might say something like, oh, we're going to have different warps executing. That's what I mean here. That's a basic unit of 32 threads that execute together. Yes. You mentioned the all threads execute same instructions. Is it all threads in a block or all threads in a warp? All threads in a warp are going to get executed. Cool. Yeah. So then the scheduler decides which warp the next instruction goes to. That's right. Or sorry. The scheduler decides which warp is going to be executing next and on what assignment and so on. Cool. And so in the memory model of a GPU, device code can read from registers. Registers are some of the fastest and most local items. This might be storing things like memory addresses of where you might need your arrays start and end.
[16:04.530] That's a place where you might store things in registers. Then you've got local memory. This is local, too, within your single SM. You've got shared memory that can be accessed within your block. So if you have something you want to transfer across threads or you want to have information that is going to be reused across threads, those are going to live in shared memory. You can, of course, then read out, or sorry, read from global memory, the DRAM that's far away, as long as you're willing to pay the latency hit of waiting for data to come to you from global memory. And there's also constant memory, which I don't think I've seen very much used in various programming things, that you can use as well. Now, you can even transfer to and back from host memory. When I say host memory, that means CPU memory that lives on the machine. You can do this if you need to offload beyond the size of your GPU. This is the set of things that you can do in your memory.
[16:58.800] And really, the key thing, the most important part of this memory model, is that as soon as you go outside of your shared memory, things are going to be slow. So grouping the blocks to decrease the amount of global memory reads is going to be the name of the game throughout this lecture. Now, I'm going to talk for three slides, two slides, about TPUs. The majority of this lecture is about GPUs, mainly because I think if you all go out and you write your code and you build your models, it's probably going to be on GPUs for most of you. But I think it is very important that we talk about TPUs, both because I think TPUs are cool and becoming more popular, but also because TPUs are the alternative evolution of GPUs. And you can see what's the same, what's not. It's very interesting to see how TPUs work as well. So actually, what's very interesting is that TPUs are very similar to GPUs at a high level.
[17:58.400] This is kind of convergent evolution. If you want to build a energy-efficient accelerator that does machine learning, you end up at the same place. So the core mental difference between a TPU and a GPU that you should keep in mind is that TPUs are in some sense simpler. They're more optimized for the machine learning workload. They have lightweight control units. And they have much bigger matrix multiply units. And so they have the same basic structure that I'll talk about for GPUs, where they will have a specialized circuit to multiply matrices. They will have components that can do parallel vector operations. And they will have some control system. They have a very similar-looking memory structure where they have slow memory, which is called high bandwidth memory, and then they have an even faster local memory called SMEM or shared memory.
[18:50.830] So this looks very, very similar to the architecture of a GPU. Really, the difference is in how these things are sized and the flexibility of these operations. And oh, and the main difference that I won't get into today is the networking for these accelerators. So if you're thinking about what is the big difference between a CPU and a GPU, actually, the biggest differences come in networking, not in the individual chips because, really, the chips, they live to just multiply matrices. And this is from the JAX book, which is a really lovely resource. But they have these two tables. They have a section on GPUs where they describe how to essentially map GPU ideas to TPU ideas. And whatever I say today about what's on the left here, you can map directly to the TPU. But the important thing at this point in the lecture to notice is that for every concept in the GPU, there's a corresponding concept in the TPU.
[19:46.620] And this mapping is actually pretty precise. A tensor core in a GPU, which is the matrix multiply unit, corresponds almost exactly to an MXU, which is a matrix multiply unit in the tensor core. And if we're getting into really nitty-gritty details, they actually use the same underlying circuitry. They're both what's called a systolic array, which kind of streams data in and out to do a matrix multiply. So the architecture is very similar. Now, how are they-- oh, sorry. Yes. Do you know what a typical [INAUDIBLE] size is on the GPU versus TPU? Oh, I do not know that off the top of my head. I'm sorry. I'll have to look that up. Yes. So the main difference between these two is the number of different units that we have. Oh, actually, I'll mention one very, very confusing thing before I move on.
[20:33.980] Sorry. TPUs call their streaming multiprocessors tensor cores GPUs call their matrix multiply units tensor cores. They are named exactly the same thing. So you will have to disambiguate them based on context. If someone says tensor core for TPU, that means a processor. If someone says tensor core for GPU, they mean a matrix multiply. So if you look at a GPU, a typical GPU like A100 might have something 132 streaming multiprocessors. That's a lot of processors. A TPU will only have two units. There's only two separate units running in parallel. Now, the GPU, if we just look at the matrix multiply units, will have 528 matrix multiply units, whereas a TPU will have just 8.
[21:19.840] And these are roughly scaled to be comparable-looking objects. Why is it that it's this way? Well, the TPU is relying on much bigger, smaller numbers of matrix multiply units, whereas the GPU is relying on smaller but much more numerous matrix multiply units. And you can imagine that if you have a small number of many matrix multiply units, it's going to give you more flexibility. You can program them to do this and that. Whereas a TPU, you're locked in to big, big matrix multiplies. And so there was a funny situation in a recent paper that I wrote where we did a batch size sweep, and it goes all the way down to 64. And it stops at 64. Why is that? Because the tensor core refuses to accept anything smaller than a 64-dimensional input. It has to be big matrix multiplies. So this is the main difference. But as you can see, GPUs and TPUs share a lot of similarities. They rely on a memory hierarchy between fast memory and slow memory.
[22:12.370] And they have a matrix multiply unit. And because of that, the core concepts that I'm teaching you today is going to transfer over between GPU and TPU. The basic ideas are always going to be the same for that reason. And this is really convergent evolution. There are only so many ways to cost-effectively allocate memory. And there's only so many ways to multiply a matrix very fast. There's some design decisions that are different. But otherwise, you end up at the same place. OK. Cool. So that's TPUs. And right, good. And GPUs, I think, have been enormously successful because it allows for easy scalability for things like matrix multiplies. If you want more throughput, you just keep adding SMs. As long as your memory bandwidth can handle those SMs, you get more and more throughput. The programming is something we don't think about, but it's deceptively easy.
[23:02.150] If you were to try to write lower level, machine code-looking stuff like TTX, it's actually doable because it's all SIMT. It's not like you're programming every single thread. You're saying, OK, here are my instructions, here are my inputs, now run them on all of my inputs. It's something that maybe you all are already familiar with things like applying or using functional programming. Finally, I haven't talked much about threads. But whenever you're doing GPU programming, these threads are always very lightweight. They can be stopped, started at any time. The scheduler can decide that, OK, I have another warp that's better to run. So I'm going to stop you and start this one. And it's very quick to swap in these jobs. That allows for much higher throughput if certain jobs are stalled or have some property that's lowering throughput. So that's the benefit of GPUs as a hardware model. It's been enormously successful.
[23:57.580] And I think even in the early days before tensor cores were implemented, people realized that commodity graphics hardware, this massive parallelism would be very effective for scientific computing. So this is one of the original papers on using graphics hardware for fast matrix multiplies. And this is a really, really cool hacker work. People figure it out, well, we have these shaders that are implemented, but we can program specific shaders that will actually give us matrix multiplies. And different rendering settings actually give us faster matrix multiplies. And that's just really cool computer science things you can do. But nowadays, you do not have to do this. You do not have to manually program your GPUs to do matrix multiplies because as of the V100, NVIDIA decided here is a piece of hardware that will do your matrix multiplies for you.
[24:48.330] And once that hardware existed, matmuls became the one privileged operation in machine learning. So at this point, there's a gigantic gap in the throughput that you're going to get between parallelizable but nonmatrix multiply operations versus amount of operations you can do in a matrix multiply. So this is one of the reasons why any near future machine learning architecture that we see that scales with compute is going to have a matrix multiply in it because this is the one way that you can really effectively get a lot of compute throughput, more than 10 times faster than any other floating-point operation that you can do in a GPU. And then the last hardware thing that I'll end with-- I think this is the last slide of this section-- is that different components of the GPU are scaling at different rates over time. And we know that the GPUs are getting faster and faster. This is just like a thing, as well as CPUs, overall compute throughput is getting faster and faster. So that's the gray line.
[25:54.010] That's pretty fast. We're getting faster and faster improvements in the total floating-point operation. Green is our memory bandwidth. This is how fast we can transfer data. That's growing comparatively slowly. Blue is parallelism. This is how fast we can connect different GPU devices. That's also going very slowly. And what does this mean? This means that if you're out here in the early days of GPU programming, probably you weren't thinking much about memory because your compute wasn't that much faster than your memory transfer. But now as we go more and more to the right of this plot, we're going to see a bigger and bigger gap between compute and memory. And that's why a lot of the optimization I'm talking about today are memory optimizations. As we go further to the right, we're going to see memory and communication bottlenecks more and more. And so it's going to take more and more work to fully utilize the hardware that we have. So that was our very basic overview of GPUs. I can't possibly give you a 20-minute lecture that fully helps you understand what the GPU hardware is. But hopefully, you have a flavor.
[26:56.920] The mental model should be GPUs are massively, massively parallel. They execute instructions all at the same time. Compute scale is really fast, really, really fast, and even faster than memory, which means memory is really important. And because memory is important, everything we do has to respect the memory hierarchy, as much stuff that we do goes into shared memory and not global memory. So this is hopefully your understanding of GPUs. And if you have this understanding, at least you'll understand all the basics necessary to internalize the next part of this lecture. OK. I can pause for a moment here in case people have hardware questions or anything else before I move on to the next part. Yes. So before v100s, there was no matrix multiply unit? Yeah. There was no tensor core that specialized to a matmul. Yeah. The tensor core is a vt100 edition. Before that, you would be programming manually.
[27:50.860] Of course, you have so many SMs and so many ALUs that you can still do quite good throughput. But a specialized circuit for matrix multiply really changes the game for that. Oh, yeah. Back there. What's the difference between the shared memory and the L1 cache? The cache just operates on its own. You don't get to control the L1 cache. [INAUDIBLE] saw that L1 just because of the physical distance. Or does it also differ some other way? Yeah, I mean, physical distance is a big part of it. I mean, if I believe, if I remember right, they're both SRAM. That's of the reasons why global memory is much slower. But the physical distance and the interconnects that are used are the reasons why it's slower. Yeah. This is not inherent, by the way. One additional side thing is that the equivalent to L2 cache in TPUs is much faster because they make different design tradeoffs on silicon to enable that.
[28:42.830] That's one of their selling points. Yes. If the compute memory gap keeps growing, do you think that affects the way future chips will be laid out? Yeah. Or maybe the future chip budget splits. One of the things that I really didn't want to talk about but now you forced me to, is the inference landscape for hardware is actually even more crazy. So one of the things that people started doing, I think last year, was something called prefill decode disaggregation, where the prefill part, which is very heavy matrix multiplies, is one chip, and then the decode part for your inference, which is more memory bandwidth-limited, is another chip. And some of the models, I think, this was step 1-3, which is a Chinese open-source model, even has a thing where they try to disaggregate different layers. So an attention will go to one chip optimized for one thing or one accelerator, and then MLPs will go to a different accelerator.
[29:37.810] And so you're starting to see, if this gap grows, any clever tricks to utilize your very precious fast memory becomes very valuable. And you start to see this very heavily in inference, which is even more memory-bound than training. Yeah. . Cool. Oh, yeah. How many warps are in a block? Do I know that? I don't think that on top of my head, I think it's also hardware-dependent. OK. Good. So the second part of this lecture is now you all have a general sense of what a GPU is and how it works. Now we are going to talk about all of the basic tricks to make workloads fast on a GPU. These are kind of basic components that you can put together to make things fast. And to keep us motivated, one of my goals by the end of this section is to explain this plot to you.
[30:32.700] So this plot, as I said before, has a very simple picture. It's the size of the matrix, a square matrix that I'm multiplying. I'm multiplying two square matrices. And this y-axis is how many items I'm processing per second. So this is the throughput. And as I'm going to the right, in general, throughput is going up. Because I have bigger matrices, I have more work to do. And so if I have a lot more work, I can keep my hardware more utilized. But there's all sorts of strange patterns here. And we want to understand why are there strange patterns here and how can we be at the very top of that curve rather than the bottom. So that's what I'm going to be talking about in this second part. And Percy already mentioned this. This is the roofline model. But the plot I showed before looks a little bit like this classic roofline model.
[31:19.370] This looks like there's a diagonal part and then a flat part. And this is what we expect. The roofline model says, up until a certain point, we will be memory-limited. We're memory-limited. So no matter how much compute you give me, my throughput is not going to go up. At some point, you have enough work per unit of memory movement that I can keep my compute units completely saturated. And past that point, having bigger and bigger units of work does not help me anymore. No matter how much extra work you give me, that doesn't help because I've fully saturated my matrix multiply units. So this is what we expect out of the classic roofline model, where as we have faster memory, we can push this upwards. But there's a memory-bound and then a compute-bound This is kind of what we expect.
[32:12.170] Now, if you want to make sure that your code runs efficiently on a GPU, really, what you need to do is, you need to be on this flat part. As soon as we're at this flat part, there's nothing else to do. We're at maximum throughput. So we want to prevent ourselves from being in this diagonal part, which means we want to increase the operational intensity. The amount of computation we do per memory read needs to be high enough. So that's what we need to do. And there's six different tricks that we will talk about. Actually, some of these are not memory tricks, the first one only. There are six different tricks that are maybe necessary to understand to make GPUs go fast. The first one is a little bit different in flavor, but then the rest of them will have this flavor of trying to minimize memory movement as much as possible. The first one is a GPU-specific thing.
[33:00.942] And it relates to one of the things I said before, the SIMT model. So in a GPU, every thread is executing the same thing. Now, what happens if you write an If statement on a piece of code that executes on a GPU? On a CPU, it's quite clear what happens. On a CPU, you will either pick one branch or the other, you will execute that branch. And you will move on with your life. In a GPU, you will do something actually quite different and slightly counterintuitive. Every thread has to execute the same instruction. So actually, every thread will execute both If and the Else. But if you're not part of If branch, you will just mask out that computation. You'll just sit idle. And so whenever you have a conditional, your mental model for the execution kind of looks like this. Part of your threads will diverge based on this If statement.
[33:49.062] Some of them will go to the bottom branch, some of them will go to the top branch. And the bottom branch ones will execute first, top branch ones will wait. And then the top branch ones will execute, and the bottom branch ones will wait. So in order to execute an If statement, you're actually walking through both of these. And then some of the threads are just waiting, doing nothing if they're on the wrong side of the branch. And this is called control divergence. And this is nice in some ways because it simplifies programming. It's not nice in that you now have these big compute gaps where your GPU is doing nothing. This is one of the reasons why you really do not want to write If statements. If you look at GPU code, if you have a ReLU, you'll often find you're multiplying with zeros or multiplying with masks instead of doing an If statement where you're updating half because a multiplication is going to all go at once.
[34:38.460] And an If statement might actually have to wait two clock cycles or something like this, two operations. So that's the first thing. But, that's a separate story from all this memory stuff, which is the main narrative or the main story of this. The first one, which you don't really have to do too much for, but I think is one of the most important ones, and hardware-wise, I think where NVIDIA is investing the most effort, is low-precision computation. And if you look at this plot-- I've brought Bill Dally back to the slide-- this super exponential growth, where does this really come from? Actually, a pretty nontrivial part of this is number representations. You start at FP32, you go to BF16, and then you go to Int8. You just keep decreasing the precision of your numbers. And of course, if you have half the bits, you have half the bits to move.
[35:30.390] And so you have your memory movement. And so you've solved the memory bottleneck for at least a little bit. This is a dumb example. But I'm just going to really drill it in there. Let's say, I have a ReLU on a vector of size n. If I'm on float 32, I need to do one read to read my number. I need to do one write after I've computed my ReLU. That's going to be 4 bytes per operation, total of 8 bytes for my memory access. I only have one FLOP worth of operation. So that's 8 bytes for FLOP. If I have the precision, of course, I've halved the number of bytes per op. I've halved the number of bytes that I need to move per FLOP. Sorry, this is the inverse of arithmetic intensity. Arithmetic intensity should be reverse of this. Now, that's a very simple example. But actually, how low precision works in a modern GPU is actually very complicated and actually either very cool or very crazy, depending on how you think about how cool things are.
[36:34.110] So when you do a matrix multiply, let's say, you have a tensor core that's operating in lower precision, how does that actually work in practice? Well, you're actually not going to keep everything in low precision. You're going to downcast before you take your product. And when you sum up your partial summands in your matrix multiply, you're going to do this in full precision. And then you might emit your outputs or your results in FP32. And part of, I think, the reason why low precision is actually kind of a black art is not because you can just decrease precision. Anyone can do that. You have to figure out which of your operations has to be in which number format. So matrix multiplies, maybe both your weights and activations can be low precision. If I'm doing a softmax, maybe I need to do that in FP32. Or if I'm doing an exponential, maybe that also needs to be FP32.
[37:22.160] Or maybe I can get away with BF16. So I think it took several years, really, to get low-precision training to be the state that it is today. And a lot of that work was slow, incremental, empirical work figuring out which of these operations we can downcast, how we can downcast them, and how we can preserve stability of training while making as much of our training and inference low precision as possible. People have been pushing this further and further. I think once you think about how hard it is to pay for fast memory, you realize, if only I can cut my precision by another factor of 2. People really, really want to do that. And so FP8 was the next frontier, I would say, of low precision. Once you get to FP8, there's no longer a standard, canonical format that everyone uses.
[38:17.030] If you're using 16 bits, you're probably using BF16 for a lot of what you're doing. But with FP8, there's going to be E4M3, which stands for four bits of exponent and three of mantissa. Or there's going to be E5M2, which is more exponent, fewer mantissa. Because these are used in different places, there's no one-size-fits-all once you start to get to so few bits. But I think the thing that's really interesting is the more advanced, low-precision number formats that are coming out with the new hardware. So I'll talk briefly about MXFP8 because it's both cool and what it does to training is strange, to say the least. So in FP8, the traditional version of training, what you have is that all of, let's say, activations are in 8 bits, So this is up here. And then you're going to have a scaling factor in FP32 that's going to allow you to scale this up and down. Why do you have a scaling factor? The scaling factor is necessary because you only have four bits of exponent, let's say.
[39:17.870] That's so few bits of exponent you're going to quickly over or underflow. So you need the scaling factor to keep you in the right range so you can keep some information. Well, people realized this is not optimal because a single matrix might have very different magnitudes. One part of my sequence might have much bigger activations than another. So instead of having a single scaling factor, why don't we have many scaling factors? So MXFP8 and these multiscaling factor formats that are much more advanced now have many scaling factors. Each of these scaling factors corresponds to the relevant colored submatrix. So the first four elements in this matrix correspond to this first scaling factor and so on and so forth. Now I have more scaling factors, these scaling factors are actually going to be in lower precision. These are going to be 8 bits as well. They're going to be E8M0. So they only have exponent bits.
[40:15.420] So they're all powers of 2 in 8 bits. So you have to make these tradeoffs. Now, what I was going to say? OK. So MXFP8 is going to use more mantissa in the elements. Scale factors are themselves FPA. And you're going to have one scale factor every 32 elements in your matrix. OK. I'm going to pause here. Don't read the third thing. Now, think for a moment. What is the problem with this design? What is the big issue if you do something as clever as this? So the big issue here is-- let's say, I want to transpose this matrix. The transpose of this matrix does not have the same pattern, so to speak, scaling pattern as this matrix. And so now it used to be that a transpose of a matrix was simple. I just transpose this matrix. I don't have to requantize. I don't have to do anything. Transposing this matrix is very complicated because I now have to potentially requantize this whole matrix to fit this 1 out of 32 pattern that I have.
[41:16.320] And so what happens internally when you actually use MXFP8-- and I found this to be crazy but also kind of cool-- is that when you're training your model with this thing, it actually creates two copies of every quantized matrix, one for your original matrix, one for the transpose. So if you ever want to transpose, you have the transposed version waiting for you. And that's how you get around this, which I think is a crazy and cool thing. The point here is maybe not to emphasize the GPU trivia. But the point here is that once you start getting to these FPA, FP4 formats, things start getting pretty exotic because the amount of exponent and mantissa loss that you have are pretty extreme. So you have to be very clever about how you do this. And as I said, if you were to go and really train a model, FPA, which I think a lot of frontier labs are doing for the compute benefits, you have to be very careful.
[42:08.670] This plot does show what I was saying. So every time you quantize, you're going to quantize one way, and you're going to transpose and quantize. So you're going to have both copies. And then the other part, which I told you about a few slides ago, is you're only going to quantize certain layers that you think are safe to do so. And this is another trial and error thing that you have to figure out. But once you have all this together, you have the big benefits of training an FPA, which is that when you're doing those matrix multiplies, you're probably going to get 20% to 30% savings, maybe more, depending on the size of your matrix. Unfortunately, it's not going to be a 2x speedup because you have to do all these quantization operations. There's a big overhead doing this. So it's not free. Yes, I remember reading that the first layer and the last layer are very hard to quantize.
[42:53.760] Do you have any intuition about that? That's right, that's right. I do not have intuitions for the first layer. The last layer. I think, being hard to quantize is because it is a first-order factor. It drives a ton of the loss. And so I think you end up getting both instability and big loss increases if you attempt to quantize the very last layer. I put this up there almost because I think it's so ridiculous. The next frontier, this is already implemented as MXFP4. The entirety of MXFP8 can be shown to you in this slide. It's negative 6 to 6. You have all the possible values. But now that you understand MXFP8, you can also understand how this thing works. You're going to have a structured block, where each block is going to be one of these four bits. Every 16 of these numbers, you're going to have one scaling factor. And that scaling factor will itself be FP8.
[43:47.080] This is a E4M3 scaling for those scaling factors. There has been a paper that has trained with FP4. I don't think I've heard of anyone successfully training in the wild, real, big, serious models with FP4. But I think this is coming. The next generation of models will probably be FP4. So that was low precision. I think it's really cool that you can really push the limits so far that you're ending up with just plus 6 to negative 6, and yet you can do interesting things with matrices in that precision. Oh, yeah, there's a question. So is this only for matrices that you can apply this, not just single scalar step? Yeah. So you can quantize anything. You can quantize the activations after a ReLU, for example. But there just isn't as much of a need to. So like a matmul, if you quantize that,you can get really big throughput improvements. And that's why if you look at the MXFP8 training one, these are all matmuls that are getting FP8 quantized.
[44:52.570] The overhead is probably not worth it for most of [INAUDIBLE]. This has something to do with the hardware. It's not like the matmul hardware part only does this. Well, it's not exclusive to that. Yeah. You can always do lower precision stuff. And then it will decrease your memory movement. This is true. But for the others, it doesn't quite make sense. This isn't any hardware limitations? It's not a limitation. But there is hardware support for all of this. So it is baked in. A bunch of this is baked into the hardware to accelerate it. Yeah. You positioned this as an improvement for memory bandwidth or loss, requiring less memory bandwidth. But is there a benefit in this one? Yes, yes, for sure. For compute, you get basically linear improvements for just multiplying the quantized numbers. But the fact that you have to quantize and dequantize means that the benefits are more diluted.
[45:43.110] Yeah. Do you have to train the scale factor? The scaling factor is an implementation-dependent choice. Depending on the library, for example, you might get different ways of quantizing and dequantizing. Training is a strong word. But you might have quantization factors that scan through and pick the max, min. Or you might have scaling factors that look at historical running statistics and then use those. I wouldn't really call them training per se. At inference, you might fit them more precisely. So there's no gradients. At inference time, you might. At training time, you would not. Yeah. Has anything been done with structured sparsity? Yeah. Because I guess NVIDIA had it. Yeah. There's a lot that's been done with structured sparsity. And I mean, MOV's, I think, one successful example of a structured sparse operation that people use.
[46:27.110] [INAUDIBLE] Yeah, right. So I mean a lot has been done. Chris Re here has done a lot with lots of different kinds of structured matrices. I think the benefits of that in terms of the compute benefits versus the representation losses, I think has washed out. I don't think it's a bad idea. But empirically, it hasn't, I think, necessarily panned out. Oh, yeah. A lot of squeeze that we've been talking about quantizing models for decades. What is current approach? Is it train in full precision or half precision and then distill into a quantized model that you can run on your Apple computer and things like that? What it's like [INAUDIBLE]? Yeah. I mean, there's a combination of-- you probably want to train a bigger model and then quantize down. There's various kinds of quantization, scaling laws, and papers that people have written, precisely characterizing the tradeoffs here.
[47:22.880] I think the state-of-the-art involves combining both quantization-aware training, to some extent, and model scaling with post-training quantization of various kinds. Yeah. There's a lot that I think is not fully well understood. Even industry teams that I've talked to are talking about still doing science of quantization. Yeah. OK. So the other trick-- this one is much more of a conceptual programming thing-- is operator fusion. It's fairly simple idea. But you'll be surprised at how much this doesn't already happen. So think of a GPU as a factory. You've got your memory warehouse. And you've got your little compute factory. And you've got a little belt conveyor going back and forth. If you scale up your compute, you're going to be bottlenecked by your little belt conveyor. And this will be worse if you have many operations.
[48:10.780] Imagine I've got lots of different operations I need to do. And I'm going to ship my raw materials to my factory and back and forth and back and forth. This is a lot of memory bottleneck that you're paying for and lots of duplex bidirectional memory bandwidth. Now, wouldn't it be better if I just had a one giant factory that would take all the raw materials and ship back my finished product? This way, I only have to pay for memory bandwidth twice, very simple idea. This is operator fusion. And this appears very often. If you write a very simple, let's say, computation of sine squared plus cosine squared, what does the PyTorch computation graph look like? It looks like this. You got your x, I do my sine, I do my cosine. I square it on both sides, I add them. And I get my output. Naively, if you don't do anything smart, what this will do is each of these will be a factory, so to speak.
[49:03.230] You will read and write from global memory each time one of these is called. And you will incur quite a bit of memory cost every time you do so. So this is pretty painful. And you can imagine instead of doing that, why don't I just write GPU code that reads once from global memory, does the whole thing inside the SM, as I talked about before. And then after it's done inside the SM, I'm going to write it back to global memory, a single for the kernel that does all of these operations. This is called fusion. And anything that's easy like this-- this is a very easy fusion. It's like a graph that I can just squish down into a single unit-- a compiler can automatically do this for you. So Torch compiler, which is a compiler for PyTorch, or JAX compiler, which is JAX's compiler, which is more deeply integrated into JAX, both of these will fuse these kinds of basic things into a single CUDA call.
[50:01.580] I'll talk a little bit more about fusion later. The advanced version sometimes require manual intervention. But the simple versions, you can just think of as compilation will solve your problems. The third idea is recomputation. So we know in back propagation, what do we do, well, we go forwards, upwards up on the tree. This is from Percy's 221 class. You go upwards on the tree. And you keep track of activations. You store these little activation numbers. And then when I go backwards, I'm going to use the activations that I store to compute the backward signal down my tree. We all know backprop to some extent. So we all know what this is. But also, now we're looking at this not from the perspective of mathematics, how do we compute a gradient. We're thinking about this from the perspective of systems.
[50:48.297] How do we minimize our memory? Well, if we think about what's happening, let's say, I have something simple I have x sigmoid, sigmoid, sigmoid out, each of these sigmoids is going to generate in the computation graph some of residual or some of activation that I need to track. So after my first operation, I have S2, second operation, S1. And then when I need to do my backwards pass it will consume S2, S1, and my output, or sorry, the gradient at my output. And I'm going to push this backwards towards the derivative of my input. So if I count the number of memory reads and writes, it's going to be one memory read for x, three memory writes for S2, S1 out.
[51:26.400] And then on the backwards pass, this is flipped. So this is a total of eight memory reads and writes, very low arithmetic intensity because I'm just stacking a bunch of sigmoids. Now, what can I do instead? Well, what I can do is, I can just throw away all of the activations. Normally, you wouldn't do this. Why would I want to recompute something? But imagine you're in a world where computation is super, super cheap. You just have abundant computation. You're not using all of your units. But your memory is very expensive. Any read and write has a lot of penalty. Then you might prefer this computation graph where I read an X, and then I'm going to do sigmoid, sigmoid, sigmoid, output.
[52:04.098] I'm going to store this output. I have to write this because that's the output of my graph. But now in the backwards pass, my backwards graph is going to get in D out. And on the fly, I'm going to reforward the X and then compute the activations on the fly right where I need them. So now this is going to be two memory reads for the D out and Dx and one memory right for the Dx. So I've gone from eight total memory accesses to just five. So now I have 5/8 the total memory accesses, same computation. But I've spent a little bit of extra compute and saved a little bit of extra memory. This is another very important trick. We don't normally think about this because, usually, we're concerned with just computing the gradient. But once you have a systems mindset, you start to think about recomputation. Cool. Trick number four. This is another memory trivia thing but surprisingly becomes important and also why sometimes it's important to understand the hardware model of what you're interacting with.
[53:05.330] So global memory, which is the big memory that lives outside the chip, is DRAM. And DRAM, because of the way that it is structured, actually comes in bursts. What that means is, when I access, let's say, the very first position, I don't actually get just the first value. I can read off all four of these first values for free. And that's because the way the memory cells are lined up, there's these blocks of memory units. And then when I activate an amplifier, I can just read out all of these almost for free. The way that it's structured in parallel allows me to read blocks of memory. That makes sense as a hardware model. And so this is called a burst section. And so whenever a section is accessed, all the other locations can also get delivered for free.
[53:51.680] And you won't pay for them, so to speak. And so you might have something like 128 bytes worth of a burst. That means a single read will return, potentially, 128 bytes worth of data, as long as it lives in the same contiguous block. So this means that if you have your reads arranged contiguously where they're all together, it's going to go much faster. And to use a term, you say that a memory access is coalesced. So that word means if all of the threads in a warp fall in the same burst. So what that means is let's, say, I have three threads. And all three of my threads, sorry, four threads, four threads, and my four threads live inside of a single burst section, or they're reading from this burst section. Then this is coalesced because they're using up all of this entire burst section. And I'm fully utilizing this very nice property of DRAM memory. And so if I have multiple warps, I'm going to have different multiple burst sections.
[54:53.170] But you can efficiently read out of memory by grouping together your memory reads. Yeah. Why does the memory access have this pattern, like this T shape pattern? Which one, like this? This one. Yeah. So as far as I understand, the memory cells are arranged in a grid. And your amplifiers are deciding which cells to read, operate at a column-wise level. And activating the voltage for this is often a thing that takes a significant amount of time. So once you've selected the column, readouts of multiple elements within that column are much comparatively cheaper. But this is to select, to read that upper red cell. Yeah, right. So if I want to read this, for example, I'm selecting this. Or maybe it's columns that are cheap in this diagram at least. But if I'm trying to read this, I might get this whole row for free because I've addressed this whole row.
[55:50.965] You don't get a T for free. Yeah. OK. Good. So you're saying like, Tatsu, why are you telling me DRAM trivia? The reason why I'm telling you DRAM trivia is because-- think about matrices. Matrices, you often want to read in big blocks because you're going to do matmuls. And certain kinds of memory accesses on matmuls are privileged. They're much more nicely structured than others. And sometimes you can take advantage of this to get very, very fast memory reads. So the mnemonic or the thing that I remember is, if you have threads and you have a row major matrix, so it's sorted where each row is how your indexes are increasing, threads that move along the major axis are not coalesced.
[56:39.200] And to really make this visual-- actually, I like this diagram on the right. Imagine I have this 4 by 4 matrix. This is row major. So it's sorted in the order of the yellow block comes first, the red block comes next, and so on and so forth. This is the linear ordering in memory address. This is the matrix that I'm trying to store. Now, imagine that I have threads that are trying to read each column, each clock cycle. So I read a column, I move to the right, I read a column, I move to the right. If I'm doing that at the very first iteration, I hit this element, this element, this element, this element. Notice that each one of these limit a different burst window. And I've read the entire vector, so to speak, or the entire matrix, sorry, so to speak, even though all I wanted was a single column.
[57:26.620] Now, imagine if I was reversed, if I was flipping this matrix, I could read the entire block. Every thread would get its read in a single read of the memory. So coalescing can give you significant improvements in reads from large, contiguous structures like matrices. So this is one reason why you might want to think a little bit about row versus column versus other kinds of ordering. The last one that I want to talk about-- and this is the big one and maybe the most important idea, I've saved it for last-- is the idea of tiling. And I think this one has probably the big impact on performance. Tiling, I think, more than any other idea, is this idea of trying to do as much as you can with the memory hierarchy. So what we want to do is, we want to group our memory accesses as much as possible. And shared memory accesses should be done together, and they should be done on shared memory. So the reason why it's called tiling is I'm going to cut up the things that I need.
[58:28.150] I'm going to tile my space. And each tile will be loaded into my shared memory. And once that tile is loaded, I'm going to repeatedly read from that tile and do as much as I can. So let's think about the matrix multiply case. I want to multiply this matrix and this matrix. And this matrix on the right bottom, this is my output. This is where I'm writing my matrix out. This is a pretty standard notation for this. Now, if I want to do this, I'm going to be reading as I'm reading through the threads multiple-- sorry. The same entry will be read multiple times. And to be precise, each entry will be read n times in an N by N matrix multiply, naively. So you might think, if I'm reading this entry, let's say, M0, multiple times, instead of hitting global memory each time, why can't I group all of those memory accesses together? And the nice thing about matrices is that there's a very natural grouping.
[59:22.810] And this is to essentially cut the matrix up into submatrices and then load them and then operate upon them. So here is a very simple algorithm. I'm going to cut out my matrices into 2 by 2 submatrices. And then I'm going to load different tiles. So I'm going to first load the M 0,0 tile, the top left here. I'm going to load the N 0,0 tile, the top left of my top matrix. And I'm going to load those into shared memory. So I'm going to pay the cost of transferring those slowly from global memory once. But once they're in shared memory, I can read and write very fast with those tiles. This output tile will also be, for now, in shared memory. And I will operate on these submatrices entirely. So I'll multiply these two submatrices. And I'll put the partial sum into my output submatrix. So that allows me to do a bunch of reads and writes all in the shared fast memory.
[01:00:14.300] And then I can write it back out when I'm done. What does the math look like? Well, in the nontiled case, I already told you, each element of my input is going to have to be read n times when I do this big N by N matrix multiply. If I tile them with a tile size of T, then each input is going to be read N over T times from global memory and then T times within each tile. And this makes sense, the extreme. If T is equal to N, I read it once from global memory and access it N times in shared memory. And this is great because the T times reduction in global memory access. If you have big, big shared memory tiles, it means that you can dramatically reduce your global memory access on something like a matmul.
[01:00:58.160] Cool. OK. Yes. Are you doing this because you can't put that entire block? Yes, yes. That's right. But if you were building one of these very fancy alternative accelerators like rock or something, which is all SRAM, then yeah, you just put the whole matrix in there. And you're very fast for the whole thing but very expensive, very, very expensive. Yeah. And then this also happens at the GPU level too because there's contact charging where I can't even fit a [INAUDIBLE] into one. And then you just keep adding that cell again and then again. Is that how it works? You mean, like, are you asking whether this kind of splitting happens at the accelerator level? [INAUDIBLE] tile across the GPUs.
[01:01:41.390] This [INAUDIBLE] will be one tile. That's right. Yeah. So data parallel, you can think of as this. If I have data on this axis and weights here, if I cut it on one of the axes that's kind of data parallel, tensor parallel could be another kind of parallelism. All these parallelism look a little bit like tiling in some ways. Yeah. Cool. OK. Oh, yeah. So for the previous one, can I assume that a library like PyTorch will just do that automatically? Yes, yes. So like any matmul kernel that you call will do tiling underneath. Yeah, that's right. If you're doing some of exotic operation, you might want to think about tiling like, things happening.
[01:02:18.490] You want to be able to locally reuse your reads. But I will say even PyTorch or something is not perfect. And so we will talk about some of the imperfections that will happen. So tiling is very, very powerful as an idea, but it's also complex. You can actually end up with some really subtle things happening with tiles. So imagine I have a tile size of 128, 128. If my matrices are 256 by 256, life is good. I've got four tiles, perfectly cut. Now, imagine I increase my coordinate size by one. Now, I've generated two very skinny tiles with basically nothing inside of them. So now you realize, actually, this tile sizing thing is actually pretty complex.
[01:03:06.020] Depending on my matrices, I might need to adjust my tile sizes to somehow trade off my memory bandwidth versus this gigantic empty space where I'm doing nothing. The tile sizes is actually a thing that you optimize. So coalesced memory accesses, the size of your shared memory in your accelerator, and the sizes of your matrices all affect which tiles are optimal. And if you know about different options in PyTorch, there's a compiler option called max-autotune. I don't how many of you have used this or about this. You turn on that flag. For the next 15 minutes, PyTorch is going to run an endless series of benchmarking tests. And you'll realize, actually, what it's doing is, it's trying all sorts of different tile sizes on your matrices and trying to figure out which one is actually faster.
[01:03:53.140] And you'll find that you'll get nontrivial benefits from actually optimizing these things. The other thing that you'll find related to coalescing is, or if your matrices are the wrong size, what you can do is that your attempt to read a tile might actually trigger a lot of memory reads. So imagine I have this nice world where I have,N let's say, this blue one is the row. And I want to read this nice tile. And my burst sections are the width of the tile. In this case, if I do four reads, I'm good to go. Because of the burst sections, I read this entire submatrix, so good memory reads from coalescing. Now, imagine I shift the matrix size by one. Now I've bumped this first row over.
[01:04:37.400] So this blue box now lives over here. And there's no way for me to coalesce my reads to read a single tile in one go because now all these things are ragged and misaligned. So I need do at least two tiles worth of reads to read out this one tile. So there's actually a lot of complexities involved in optimizing tile reads once you start thinking about coalescing and all these other things. And so in this case, what you have to do is padding in order to line up, essentially, the boundaries of the burst sections with the tile sizes that you have. And this is why you end up getting stuff like this, which seems very mysterious. I think this is one of my favorite weird ones. Andrej Karpathy, when he was doing nanoGPT speedruns, said well, when I increase my vocab size to 50,257 to 50,304, I got a 25% speedup.
[01:05:31.250] And it's kind of weird to have to pad in order to get a speedup. But maybe it makes sense once you realize things like this. Padding can start to shift the rows and columns so that you actually get very efficient reads, whereas at other times, you do not. Yes, go ahead. Is there a usual tile size that you should be taking into account? You don't want to think about tile sizes per se, because tiles, you should normally think of as a lower level systems thing, unless you're writing your own kernels. But you do want to think about the size of your matrices. But you want your size of your matrices to be divisible by the tile size. Right. So you want it to be-- I'll talk about that in a moment, but right-- generally, powers of 2, ideally divisible also by 32.
[01:06:15.690] OK. So now we can talk about this matrix mystery over here. We have essentially all of the pieces that we need to understand this. So we'll walk through this one by one. We've already talked about compute intensity, the roofline plot. You need to have enough work for each memory read in order for you to be able to saturate your compute unit. That's what's driving this diagonal upward trend. So that one is done. Now, let's talk about some of the other ones. First one's tiling. If we look at each of these lines and we color them based on, essentially, the divisibility of the matrices, you will find that the blue ones are essentially only divisible by 1.
[01:06:55.570] They're not even even. Orange ones are divisible by 2. Green is divisible by 8, red is 16, and then purple is 32. And you see that the ones that are not divisible by 1 or just 1 or just 2 actually have very bad throughput. But once you start to get to higher multiples, you start to essentially get no diminishing-- sorry, no penalties. 16 or 32 perform equally well. And it's not because of somehow powers of 2 are magic or something like that. It's because 16 and 32 are sufficiently large that it has the right burst window property. It fits the whole burst window in. And so you can get coalesced, reads into your entire tile.
[01:07:38.620] And it has a major impact on alignment. Now, if we look at what's remaining, if we look at this orange line, we'll sometimes see this weird periodic behavior where, suddenly, performance will drop. And this is very weird because I'm just multiplying matrices. Why does adding a single extra dimension suddenly kill my throughput? And this is, I think, using all of the things that we learned is kind of fun to get to this end. Percy calls this GPU trivia. So I think he hates this one, but I love this one. So why is it that going from here to here, I get such a big drop in throughput? Well, it happens between 1792 to 1793. And why does this happen? Well, I'm using a tile size of 256 by 128. So that means-- how many tiles are there? Well, there's 98 at 1792. But once I bumped myself up to 1793, well, actually I've increased both tile sizes by one.
[01:08:35.130] So I now have 120 tiles. So I have 98 tiles before, and I have 120 tiles now. That's bad, but it's not this bad. This drop is gigantic. Well, it turns out the reason why it's so dramatic is because the A100 has 108 SMs. And 98 tiles fits in 108 SMs. But 120 tiles does not fit in 108 SMs. And so you finish doing 108 SMs and you're, well, I've got 12 more of these guys. I got to, I guess, process them. And most of your GPUs sitting idle. So you can end up with these very weird effects once again, where by bumping up your tile size by one, suddenly, you need to wait a whole other cycle for you to process a bunch of tiles.
[01:09:17.250] And that's called wave quantization. So that was all for part two. How do we make machine learning workloads go fast on a GPU? Well, a lot of the story is really just about memory, memory, memory. So how do we reduce the amount of memory access we have to do? Well, let's coalesce our memory reads together. Let's fuse our operations together so we don't have to read as much. And for the things that we have to read no matter what, well, let's reuse the memory, so move things to shared memory and reuse it as much as possible, which is tiling. And then finally, maybe you can trade other resources that you have for memory. And that's quantization and recomputation. So all these ideas very much relate to this idea of memory is our bottleneck, how do we deal with that.
[01:10:04.080] Cool. Any questions for the last thing? Yeah. I think we just bump up SRAM so to get rid of DRAM. Very expensive. I mean, there's two reasons. One is it's just expensive to build SRAM, also very difficult because it has to be physically close. Signal propagation is hard. And I think a third reason is that SRAM and DRAM are just different storage mechanisms. SRAM needs to be powered the whole time to hold value, very energy hungry. And so if you're thinking about total cost of ownership, energy consumption, maybe you don't want a giant SRAM chip. You really want to be using it where you need it. Yes. Those things scale up past 32, or if you like, operate in chunks of 64.
[01:10:51.960] You mean like, does it help if I go higher? Yeah. No. Because I think it's kind of as long as it divides by your burst size, then you're good to go. It's a divisor property. There are other phenomena where you might get benefits from larger divisors but not here. Yes. Why 108? That's just the number of SMs. You will have to ask Jensen for that. I do not have an answer for you. Yes. So going back to the hardware side, sorry, not this part. It just occurred in my mind. Have you studied GPU versus Cerebras' wafer-scale engine? Yeah. I mean, I think sitting over there knows much more about wafer scale stuff than I do. We were talking about, after last lecture, how difficult the compilation was because if you have something wafer scale, then you have wave interference effects and all sorts of horrible things that happen.
[01:11:46.170] That's a world that I'm not familiar with. But I think is much more complicated than what we have here. OK. Cool. All right. So now I think we can talk about, in some sense, the finale, which is, we've learned a bunch of things about GPUs. It would be nice if we can tie that all together to essentially be able to reproduce or understand something that I think is really important and advanced, which is flash attention. So flash attention is one of the big, I think, improvements in attention capabilities in a while. And it's all systems. So it's going from just PyTorch's naive implementation of attention to a single, very cleverly fused kernel. This is from the original paper. But you've got this dramatic improvement in the latency of serving these things. And as I showed you in, I think, two lectures ago, flash attention is a nice, memory-efficient way of computing things.
[01:12:33.950] So you can do bigger and bigger attention. And you can see, a lot of the gains are going to come from improvements in the amount of memory that is transferred from HBM, which is the global memory. So what do they actually do? If you read the paper, they'll say, OK, we just do two things. We apply tiling and recomputation. And then we do it in a way that allows us to do things subquadratically in the number of memory accesses. I think we understand all of those components now. So let's see what's actually happening. So remember, attention computation has a couple of things. It has a K, Q, V. And we do a matrix multiply with the K's and the Q's. And then we have a matrix multiply with a V at the end. And I have a softmax in the middle. So those are the operations that we need to deal with-- three matrix multiplies and one softmax. So we have matrix multiplies. We how to do them now. So we're going to do a tiled matrix multiply. Figure 1 from the paper is literally just tiled matrix multiplies. We've seen this before.
[01:13:34.700] We're taking our inputs, our K's and Q's. We're cutting them up. And we're just multiplying them chunk by chunk on a tile. And then after we do the computation, we write them out. Nothing fancy or special here. The thing that is tricky about flash attention, and you'll see this in the assignment, is there's a big global softmax on the whole thing. So we can't just do this naively because the softmax is a global operation that ties everyone together. So it connects all of the different tiles together. It might not seem obvious how to tile attention. The key trick, and once you see this, this really the core, I guess, idea or observation, once you see this, the rest is very simple.
[01:14:15.020] Normally, to compute a softmax, you would do it this way, as shown in this plot here, which is, you would exponentiate all of your elements. You would subtract the maximum for numerical stability reasons. And then you would normalize that standard softmax. Now, there's a different way of computing this online. And I think once you realize it, it's kind of simple, which is, you're going to, as you go, compute the normalized softmax. And every time you encounter a bigger number than what you've seen before, you're going to swap out the maximum. And so you have an online way of computing the softmax as you go. The only tricky thing, of course, being the max, and the way we deal with the max is, we keep track of the max.
[01:14:54.217] And then we swap it out and correct for it every time we see a new maximum. Otherwise, this is really just an online running sum of an exponential from everything that we've seen before. And of course, at the end, you'll just divide everything by the accumulator that you've computed online. Now, the important thing about the softmax is because it's online and it goes kind of block by block or piece by piece, we can compute this tile by tile. I don't need to see the rest of the tiles in order to compute this. I can compute the softmax for my tile. And then I can save the partial results, maybe write them back to global memory, and then just continue on with my day. So this allows us to now break up all the different interacting pieces into something simpler. I think this was from FlashAttention 2 or maybe 3. I think it's 2. I like this visualization, which shows you which parts are being tiled.
[01:15:46.090] So the dashed blocks are being done, tiled in SRAM. The blue boxes are the ones that are in high-bandwidth memory, the global memory. And what you see is-- let's say, we do the K, Q inner product. These are tiled matrix multiplies. And they're stored in SRAM. And we keep them in the SRAM. And we compute this exponential. And then similarly, in this tile, we're going to compute this running softmax. And I'm going to move to my next tile. And I'm going to do the same thing. But now I'm going to keep the partial running sums of my softmax as I go. I can keep them in shared memory or in registers or elsewhere. This is a very small number of things to keep track of. At the end, after I've computed all these exponentiated quantities and my normalizers, it's very simple to just multiply it with a V, once again at the tile form, and then divide at the very end.
[01:16:33.710] This allows me to essentially take the attention operation, cut it up into these small tiled chunks, and then go. Finally, you've got fusion. You're going to do both of these together as much as you can without going back and forth. And then I'm not going to cover this or talk about it. But what you can do is, you can, in the backward pass, recompute everything that you need. So the last piece of flash attention is, if you saved your activations, you still have this N squared-sized activation that you're going to need to save for your attention. Instead of saving that, throw it all away and recompute everything tile by tile on the backward pass. So it's got that recomputation piece as well. That gives you flash attention. Now we've seen all the different tricks in action.
[01:17:19.780] And that's essentially all there is. A lot of the tricks that you see for efficient implementation and how to use GPUs well boil down to those kinds of ideas. So to put it all back together, important to understand the hardware. The hardware, in our case, the GPU, is powering scale and compute. And we need to understand, really, all the way to the low-level details to fully grasp why we make certain decisions. I don't want you guys to be the kind of people that cargo cult, 32 multipliers for your matrices. I want you to really understand why we do those things. And that goes all the way down, in my opinion, to the hardware level. The other thing is that the current way that compute is growing means that we really want to think about matmuls. Those are the core operations that are very arithmetically dense. And because of the gap between compute and memory, we need to think very carefully about data movement.
[01:18:13.360] And then finally, things like flash attention, I think, are really cool. There's also architecture work that's very hardware-aware. Thinking carefully about the hardware when we do things like architecture design is critical to good performance for future systems. Thanks. [APPLAUSE]