CS336 LecTure 5 笔记补充

Course Goal

The main goal is : Make CUDA and GPUs less magic

In Deep Learning era，Compute leads to predictable performance。计算量能带来可预测的性能表现，通常情况下，计算量的增加会带来 LLM 性能的可预测提升。过去十年中，GPU 的 Parallel Scaling（并行扩展） 能力提升了超过 1000 倍，如果没有 GPU 的 Scaling，就不可能实现 LLM 的 Scaling。

From CPU to GPU then TPU

为什么单核性能停滞 (we don't use Dennard Scaling)

• Dennard Scaling（登纳德缩放定律）：早期晶体管缩小能降低功耗并提升频率。但在 20nm 以下物理极限下，量子隧穿效应导致漏电功耗激增（热墙问题）。单线程性能（CPU 频率）停滞。在九十年代的时候这种操作还是可行的，大约在二十一世纪初这种放缩遇到了瓶颈。
• Parallel Scaling（并行扩展）：为了突破瓶颈，计算能力转向超指数级增长，依靠增加核心数量（如 Tensor Cores）和低精度计算（FP8/FP16）来提升整体吞吐量。

GPU vs CPU

• CPU (中央处理器)：优化的是 Latency，降低延迟，即让每个线程尽快完成。拥有少量强大核心，擅长复杂逻辑控制。
• GPU (图形处理器)：优化的是 Throughput，Throughout 为吞吐量，即总体处理的数据量。拥有海量 Compute Units (ALUs，算术逻辑单元)，专为大规模并行运算而生。

可以理解为 CPU 是串联的，串行工作挨个执行，而 GPU 是并联的，并行工作同时执行。

GPU vs TPU

• 架构区别

在 GPU 中，GPU 的底基依然是 CUDA Cores（流处理器），它们负责执行基础的算术和逻辑运算。Tensor Core 是在 2017 年 Volta 架构才加入的专用加速单元。它与 CUDA Cores 是并行工作的，主要作为一个“协处理器”来加速密集线性代数计算如矩阵乘加，而数据准备、控制流等依然依赖 CUDA Cores。

在 TPU 中，TPU 从设计之初就是一个 ASIC 专用集成电路数学引擎。它的核心就是 MXU（Matrix Multiply Unit，矩阵乘法单元）。在 TPU 中，成千上万个计算单元被直接连接起来形成矩阵乘法阵列，除了最开始的加载数据和激活函数外，几乎不需要再访问存储单元。可以说，Tensor Core 就是 TPU 的心脏。

• 计算模式差异

在 GPU 中，Tensor Core 被定义为执行非常高效的矩阵乘法的微核心 Micro-core。它通过特定的指令如 HMMA/IMMA 来加速 4x4 或 16x16 这样的小块矩阵乘加运算。它依然需要遵循 GPU 传统的“加载-执行-存储”模型，频繁访问寄存器和共享内存。

TPU 核心采用的是 Systolic Array 脉动阵列，数据像血液在心脏泵抽取中一样在计算单元之间流动。矩阵数据被直接加载到脉动阵列中，阵列每秒可执行数万亿次乘加运算，数据在阵列内部流动完成计算，大幅降低了内存访问的频率和功耗。

• 设计差异

GPU 的 Tensor Core 是为了弥补通用 GPU 在处理 AI 矩阵运算时的效率不足而诞生的创新与突破。GPU 依然保留了大量用于图形渲染、纹理映射的硬件，Tensor Core 只是其庞大通用架构中的一个特化模块。

TPU 的 Tensor Core 体现了领域专用加速 Domain-Specific Acceleration。TPU 的设计理念类似于 RISC 精简指令集，通过限制功能的灵活性（比如它无法运行操作系统或处理通用程序），将芯片空间最大程度让渡给矩阵计算，从而获得极致的速度和能效比。

GPU 内存层次与执行模型

内存层次结构 Memory Hierarchy

Memory 距离 SM 越近，访问速度就越快。 必须遵循内存层级结构才能让程序运行更快。

内存层级	速度	容量 (以 A100 为例)	作用
Registers 寄存器	️ 最快	极小	线程私有，访问零成本
Shared Memory / L1 Cache	️ 极快	~192 KB / SM	同一个 Block 内共享，避免频繁访问全局内存，速度可提升约 8 倍
L2 Cache	快	~40 MB	芯片上 (On-die)，所有 SM 共享
Global Memory (HBM)	慢	~80 GB	芯片外部的显存，容量大但带宽是主要瓶颈

• 核心矛盾：Compute 能力（尤其是矩阵乘法）Scale 的速度超过了 Memory Bandwidth（内存带宽）的提升。
• 处理原则：尽量减少对 Global Memory 的访问次数。

执行模块

• Threads 线程：以并行方式工作。所有线程执行相同的指令，但输入不同。被称为 SIMT (Single Instruction, Multiple Threads 单指令多线程)。
• Blocks 线程块：由多个线程组成的 Group。每个 Block 在一个 SM 上运行，并拥有自己的 Shared Memory。
• Warp 线程束：线程总是以一个由连续编号的 32 个线程组成的 Warp 方式执行。Warp 是 GPU 调度和执行的最小基本单位。

内存通信规则

• 每个线程可以访问自己的 Register 和 Local Memory。
• Block 内的线程可以通过 Shared Memory 高速通信。
• 跨 Block 通信的数据必须通过 Global Memory 读写（速度较慢）。
• Constant Memory：适合存储不变常量（如超参数），所有线程共享且只读。

加速 GPU 的 Tirck

Control Divergence 控制发散 - 非内存瓶颈

• 原理：如果同一个 Warp 内的某些线程执行了 if 分支，而其他线程执行了 else 分支，GPU 就不能同时执行它们，必须顺序执行 (Sequentially) 各自的分支。
• 后果：未命中分支的线程处于空闲等待状态，产生 Control Divergence，严重降低 GPU 运行效率。
• 优化：在 Warp 中尽量避免产生条件分支 (Conditionals)。

Low Precision Computation 低精度计算

• 原理：通过 ReLU (max(0, x)) 举例：
  • Float 32：读取 x (4 bytes) + 写入结果 (4 bytes) = 8 bytes。Operations 为 1 次。Arithmetic Intensity (算术强度) = 8 bytes / FLOP。
  • Float 16：读取 (2 bytes) + 写入 (2 bytes) = 4 bytes。Operations 仍为 1 次。Arithmetic Intensity = 4 bytes / FLOP。
• 使用 FP16/BF16 替代 FP32，内存访问量减半，直接提升算术强度，同时 Tensor Core 计算吞吐翻倍。

Operator Fusion 算子融合

• 问题：如果有很多 Operations，拆分成多个 Kernel 会导致中间结果反复写入和读取 Global Memory。
• 解决方案：尽量将它们 Fusion (融合) 到一个 CUDA Kernel 中去运算，最后再传输到 Memory。将 Memory Bound 的碎片化操作转化为 Compute Bound。

Recomputation / Gradient Checkpointing 重计算

• 存在问题：在梯度回流 (Backpropagation) 时，存储 Activations（激活值）会导致大量的 Memory Read/Write 和显存爆炸。
• 解决方案：Throw away the activations, re-compute them（丢弃激活值重算）用约 30% 的额外计算代价，换取大量显存空间，从而允许增大 Batch Size。

Memory Coalescing 内存合并访问

• 原理：当一个 Warp 中的所有线程同时执行一个 Load Instruction，并且它们访问的内存地址落在同一个 Burst Section（突发段） 内时，这些内存访问可以被硬件合并为一次 DRAM 请求。
• 反例：如果 32 个线程随机访问显存，GPU 要发 32 次请求，带宽浪费 32 倍。

Tiling 分块

• 原理：Tiling 是对线程进行分组和排序，以尽量减少 Global Memory Access。
• 问题：在矩阵相乘中，如果 thread0,0 和 thread0,1 都要从 Global Memory 读取同一个数据 M0,0，显然会降低速度。
• 解决方案：将矩阵分割为很多的 Tiles，然后加载进 Shared Memory。
  • 优点 1：重复读取使用 Shared Memory 而非 Global Memory，速度快。
  • 优点 2：线程访问的内存地址更连续，符合 Memory Coalescing 优化原则。

注意：如果矩阵的 Size 不能被 Tile 的 Size 整除，就会导致稀疏矩阵 Sparse Matrix 的出现，造成计算效率严重下降（例如：能被 32 整除 vs 只能被 2 整除，性能完全是两个效果）。

解构 FlashAttention

Putting it together – unpacking FlashAttention

FlashAttention 是用来加速 Transformer 计算速度的一种手段，它集成了上述所有的 Trick

1. Tiling：将 Q, K, V 矩阵分块，加载到 SRAM（Shared Memory）中进行计算，避免将 \(O(N^2)\) 的 Attention 矩阵写入 HBM。
2. Online Softmax：允许在分块计算时，边遍历边动态更新全局最大值和归一化因子，使得分块计算在数学上完全等价，避免存储 \(N \times N\) 矩阵从而达到优化的目的。
3. Recomputation：反向传播时不读取巨大的 Attention 矩阵，而是通过重新执行前向的分块计算来获取梯度，解决了长上下文 Long Context 的显存瓶颈。