[CS336-05] GPU 基础

这篇是 Stanford CS336 Spring 2025 第 5 讲 Course Materials 的学习笔记。主题是 GPU：为什么它适合训练大模型，又为什么很多程序仍然跑不满 GPU。

1. GPU 和 CPU 的根本差异#

CPU 优化的是少量线程的低延迟执行。它有复杂控制、分支预测、大 cache、强单线程能力。

GPU 优化的是大量线程的吞吐。它有很多 SM，每个 SM 执行许多线程 block；线程以 warp 为单位运行，同一个 warp 内的线程执行同一条指令，只是数据不同。

因此 GPU 适合大规模规则计算，尤其是矩阵乘法；不适合频繁分支、随机访问和小而碎的任务。

2. Memory hierarchy 决定性能#

GPU 上不同位置的 memory 速度差异很大：register、shared memory、L1、L2、global memory。越靠近计算单元越快，但容量越小。

大模型训练的很多优化都可以归结为一句话：

尽量少从慢内存读写，尽量在快内存里复用数据。

这就是后面 kernel fusion、tiling、FlashAttention 的共同逻辑。

3. FLOPs 很快，memory 没那么快#

现代 GPU 的 tensor core 让低精度 matmul 极快，FLOPs 增长速度远高于显存带宽增长速度。因此很多操作不是 compute-bound，而是 memory-bound。

判断一个操作是否容易跑满 GPU，可以看 arithmetic intensity：

arithmetic intensity = FLOPs / bytes moved

矩阵乘法通常 intensity 高，容易利用 tensor core；elementwise 操作和 normalization 经常需要大量读写但计算很少，容易被内存带宽限制。

4. 让 GPU 跑快的几类技巧#

Lecture 5 总结了几类常见优化。

低精度计算减少每个数的 bytes，同时利用 tensor core。Operator fusion 把多个小 kernel 合并，避免中间结果反复写回 global memory。Recomputation 用额外计算换更少 activation 存储和读取。Coalesced memory access 让相邻线程访问连续内存。Tiling 把大矩阵拆成 block，放进 shared memory 反复复用。

这些技巧都服务于同一个目标：减少数据搬运，提高计算单元利用率。

5. FlashAttention 的系统直觉#

标准 attention 需要形成 N x N attention matrix。问题不只是 FLOPs，还有巨大中间矩阵的读写。

FlashAttention 的关键不是改变 attention 数学结果，而是改变执行方式：分块计算 attention，把 Q、K、V 的 tile 放进更快的 memory 中，并在线维护 softmax 所需的统计量，避免完整 materialize attention matrix。

这正好体现 GPU 优化原则：不让中间大矩阵频繁进出 global memory。

6. takeaway#

GPU 性能优化不是“把代码放到 CUDA 上就快”。它需要理解：

• GPU 是吞吐机器，不是低延迟机器；
• tensor core 很快，但 memory movement 往往是瓶颈；
• kernel fusion、tiling、recomputation 都是在减少慢内存访问；
• FlashAttention 这样的算法优化，本质上是系统感知的计算重排。

后续 Triton 和并行训练课程，都是在这个 GPU mental model 上继续展开。

参考#

• Stanford CS336 Spring 2025 Course Materials: lecture 5.pdf