Runguo Li

Runguo Li

M.S. student in Information Science @ UIUC, advised by Prof. Minjia Zhang. Working on ML systems for large-scale models: MoE inference and expert streaming, distributed training and rollout, GPU compiler correctness.

第 5 讲：GPU

目录 · ← l4 · l6 →

第 5 讲：GPU

日期：4 月 13 日（周一，Spring 2026） | 讲师：Tatsu Hashimoto | 材料：lecture_05.pdf

概览

本讲为 GPU “去神秘化”：它和 CPU 有何不同、内部构造（SM、warp、内存层次），以及最关键的——为什么 GPU 会变慢 与 如何写出快算法：低精度计算、算子融合、重计算、显存合并访问（coalescing）、分块（tiling）。最后以 FlashAttention 作为综合案例拆解：KQV 矩阵乘的分块 + 在线（telescoping）softmax 技巧。

核心概念与定义

  • GPU 与 CPU 的区别：CPU 为”少数快线程”优化（延迟优先）；GPU 为”海量线程”优化（吞吐优先）。GPU 有大量小型 ALU、对分支支持弱、内存层次深。
    • 类比：CPU 是几位手艺精湛的工匠，每人做得极快；GPU 是上千名简单工人的流水线，总产量惊人——但每一步所有人都必须执行同一条指令（SIMT）。
  • 硬件结构：SM（流式多处理器）内部含大量 SP（流式处理器）来执行线程；一个线程块（thread block） 跑在一个 SM 上并拥有自己的共享内存；线程以warp（32 个连续编号的线程）为单位锁步执行（SIMT——单指令多线程）。
  • 内存层次：寄存器（最快）→ 共享内存/L1（在 SM 内部，比 DRAM 快约 8 倍但每字节成本高约 100 倍）→ L2（片上）→ HBM/global memory（GPU 旁的 DRAM 芯片）。算力（FLOPs）增长快于显存带宽——”内存墙（memory wall）”——因此”如何持续喂饱计算单元”是核心问题。
  • Tensor Core：专用矩阵乘电路（Volta 起），使矩阵乘比其他浮点运算快 10 倍以上。TPU 思路类似：轻量控制 + 大而快的矩阵乘单元 + 快内存，但核心数更少更大、没有 warp 概念（只有 block 模型）。
  • Roofline 模型：横轴为算术强度（FLOPs/byte），纵轴为达到的性能（FLOP/s）。拐点 = 加速器强度 = 峰值 FLOP/s ÷ 带宽。拐点左侧是 memory-bound（性能随强度线性上升）；右侧是 compute-bound（贴在峰值上）。
  • 让 GPU 变快的六种手段（摘自幻灯片）：
    1. 低精度计算：比特更少 = 搬运字节更少，提升算术强度（fp32 的 ReLU：8 字节/FLOP → fp16：4 字节/FLOP）。Tensor core 加速低/混合精度。前沿方向：FP8（E4M3/E5M2）、MXFP8（分块缩放、E8M0 缩放因子）、NVFP4。
    2. 算子融合（operator fusion）：把多个逐元素算子合成一个 kernel，避免数据在 HBM 之间来回搬运（例如 sin²x + cos²x 从 5 个 kernel 变成 1 个）。类比：由仓库传送带供料的工厂——不要把半成品每步都退回仓库；一次走完全部工序。
    3. 重计算（activation checkpointing）：不保存全部激活，反向时重算。常常是最优选择：3 个叠加 sigmoid 用重计算后显存访问降到 5/8。
    4. 显存合并（memory coalescing）：DRAM 以突发（burst，128 字节事务）读取；当一个 warp 的 32 个线程落在同一个突发内时访问被合并。对行主序矩阵而言，线程沿”行”方向移动是不合并的——经典性能陷阱。
    5. 分块（tiling）（最重要）：把输出矩阵切成 tile，把 A/B 的 tile 一次载入共享内存，被多个输出元素复用，并使访问合并。未分块的矩阵乘会从 global memory 读每个输入 N 次；分块后从 global memory 只读 N/T 次、从共享内存读 T 次——HBM 流量降低 T 倍。
    6. 避免控制发散（control divergence）（不是访存问题）：warp 内线程执行同一指令，条件分支会串行化（先走 A 路径再走 B 路径）——数据相关分支的隐形成本。
  • Wave quantization（波量化）：若线程块数量不能整除 SM 数量，最后一个波次会部分闲置（A100 有 108 个 SM；120 个 tile → 108 + 12）。这解释了”神秘”的周期性性能凹陷（1792→1793 的矩阵乘之谜）。
  • FlashAttention 拆解：
    • 第一步：KQV 矩阵乘的分块（blocked GEMM）——把 A/B tile 通过共享内存搬运。
    • 第二步：增量（在线）softmax——为了逐 tile 归一化，需维护滑动最大值并用 telescoping 修正，从而在不物化完整 S = QKᵀ 矩阵的前提下得到精确的 softmax 分母。
    • 反向传播：逐 tile 重算（不保存 attention 矩阵）。

代码示例：在线 softmax（FlashAttention 的核心）

代码（Python）：

import torch, math

def online_softmax_attention(Q, K, V, block_size=2):
    # Q,K,V: [n, d]（单头）；按 tile 处理打分矩阵的行
    n = Q.shape[0]
    acc = torch.zeros(n, V.shape[1])          # 加权和累加器
    m = torch.full((n,), -float("inf"))       # 滑动行最大值
    l = torch.zeros(n)                        # 滑动 exp 之和
    for j in range(0, n, block_size):
        S = Q @ K[j:j+block_size].T           # 打分 tile: [n, block]
        m_new = torch.maximum(m, S.max(dim=1).values)
        alpha = torch.exp(m - m_new)          # 重新缩放旧累加器
        P = torch.exp(S - m_new[:, None])     # 未归一化的 tile 概率
        acc = acc * alpha[:, None] + P @ V[j:j+block_size]
        l = l * alpha + P.sum(dim=1)
        m = m_new
    return acc / l[:, None]                   # 最终归一化

# 与朴素实现对照：
def naive(Q, K, V):
    S = Q @ K.T
    P = torch.softmax(S, dim=-1)
    return P @ V

代码做了什么： 在不物化完整 [n, n] 打分矩阵的前提下计算 softmax attention。它按 K/V 的列 tile 迭代，为每一行维护三个滑动量：最大值 m、指数和 l、加权累加器 acc。当新 tile 带来更大的最大值时，旧累加器用 exp(m_old − m_new) 重新缩放（telescoping 修正），因此最终的 acc / l 是精确的。

实现深挖：

  • 为什么要维护滑动最大值并重缩放：标准 softmax exp(S − max(S)) 需要先拿到整行；在线技巧让你可以流式处理 tile——用 exp(m_old − m_new) 缩放累加器，保证每个贡献都相对当前最大值被正确加权。这正是 FlashAttention 的前向，并与 KQV 矩阵乘融合在一起。
  • 为什么它带来 O(1)-block 显存：只有累加器（[n, d]）和 P 的 tile（[n, block]）活在寄存器/共享内存中；完整的 S 和 P 从不写入 HBM。
  • 为什么反向要重算：在反向重算 S 和 P 的 tile 就不必保存它们，代价是多一次 QKᵀ 类计算——即幻灯片里”用算力换显存”的取舍。

与作业的联系：作业 2 的核心任务是用 Triton 实现 FlashAttention-2（前向 和 反向），其中就包含这里的在线 softmax + 分块逻辑，再加上 mask 与 bias 处理。听课时的”分块 + 在线 softmax”是概念蓝图；作业里的 Triton 技巧（tl.dot、block pointer、反向重算）是机械实现。

关键要点

  1. GPU 是高度并行的 SIMT 机器：32 线程的 warp 锁步执行、线程块跑在带共享内存的 SM 上，且内存层次中真正的稀缺资源是带宽（而不只是 FLOPs）。
  2. 算力增长快于显存 → 必须最小化数据搬运：融合算子、合并访问、通过共享内存分块，并（有时）用重算代替保存。
  3. 低精度（fp16/bf16/fp8）提升算术强度并解锁 tensor core；roofline 模型告诉你处于 memory-bound 还是 compute-bound。
  4. 分块 + 在线 softmax = FlashAttention：把”看起来必然平方复杂度”的算子变成访存高效、全融合 kernel 的经典范例。
  5. 性能充满量化效应（波量化、对齐、bank conflict）——benchmark 与 profiling 必不可少，而像 1792→1793 这样的小改动可能带来非直观的大幅波动。

常见陷阱

  • 共享内存 bank conflict：32 个 bank、每周期每 bank 一次访问；跨步访问模式（例如读矩阵的列）会造成 32 路串行化。用 padding/swizzling 缓解。
  • 未合并的 HBM 访问：线程索引必须映射到连续地址（128 字节事务）；行主序下沿列方向遍历是经典杀手。
  • 波量化：网格尺寸尽量整除 SM 数量，避免最后一个波次闲置。
  • warp 发散：数据相关分支（例如 ReLU kernel 里的 if x < 0）会让两条路径都串行执行——理论无害，实践昂贵。
  • 寄存器膨胀导致低占用率：单线程使用超过约 160 个寄存器会减少 SM 能调度的 warp 数；thread coarsening（一个线程处理多个元素）有时是解药、有时是病因。
  • 以为 FLOPs 等于运行时间：同一个操作因融合与数据搬运差异，墙钟时间可能天差地别（第 3 讲的 RMSNorm 例子）。

复习题

  1. 问： 在同一块硬件上，为什么矩阵-向量乘是 memory-bound，而矩阵-矩阵乘是 compute-bound？
    • 答： 两者都要读 O(n²) 字节的矩阵，但矩阵-向量乘只做 O(n²) FLOPs（强度约 1），矩阵乘做 O(n³) FLOPs（强度约 n/3）。n=1024 时矩阵乘强度约 341 ≫ H100 的约 295（compute-bound），而矩阵-向量乘强度约 1 ≪ 295（memory-bound）。
  2. 问： 在线 softmax 如何在流式处理 tile 的同时保持结果精确？
    • 答： 它维护滑动最大值 m，并在最大值增大时把已累加的项乘以 exp(m_old − m_new)。这在代数上等价于”事后一次性减去最终最大值”——一种 telescoping 修正——因此最终累加器等于真实的 softmax 加权和。
  3. 问： 为什么分块能把矩阵乘的 HBM 流量降低 T 倍（T 为 tile 大小）？
    • 答： 每个输入元素只需在每个它参与的 tile 中载入共享内存一次（N/T 次，而不是 N 次），而在 tile 内部从高速共享内存读 T 次。全局显存读次数从每元素 O(N) 降到 O(N/T)。