Scaling Book 入门第 3 章：内存层级与带宽 — 数据如何流动

本章目标：理解加速器内部的内存层级（HBM → Cache → 计算单元），以及数据搬运如何成为性能瓶颈。

对应原书：Chapter 2 (TPU internals) + Chapter 12 (GPU memory)
优先级：⭐⭐⭐ 高 | 建议时间：Day 3, 约 2 小时

3.1 数据搬运的层级

无论 GPU 还是 TPU，数据在使用前都必须经过多层”搬运”：

HBM（主内存，大但慢）
  ↓  HBM 带宽 ~1-9 TB/s
VMEM / L2 Cache（片上缓存，小但快）
  ↓  内部带宽 ~20-40 TB/s
MXU / Tensor Core（计算单元）

📋 背景知识：为什么需要内存层级？

计算机硬件有一个根本矛盾：容量大的内存必然慢，速度快的内存必然小。

层级	类比	容量	速度
寄存器	手边的便签纸	几 KB	最快
SMEM/VMEM	桌上的文件夹	几十-几百 KB	非常快
L2 Cache	抽屉里的文件	几十 MB	快
HBM	房间里的书架	几十-几百 GB	中等
CPU DRAM	隔壁仓库	几百 GB-TB	慢

矩阵乘法需要大量数据（权重 + 激活），但计算单元一次只能处理很小的 tile。内存层级的目标是让计算单元永远不等待数据——通过逐级 prefetch 和流水线。

TPU 的数据流

1. 权重和激活值存储在 HBM 中
2. 计算前，数据从 HBM 流式传入 VMEM（128 MB 片上缓存）
3. 从 VMEM 加载到 VREGs（向量寄存器），再送入 MXU
4. MXU 结果写回 VREGs → VMEM → HBM

关键：这个过程是流水线化（pipelined） 的 — 不需要等全部数据加载完才开始计算。

📋 背景知识：流水线（Pipelining）的概念

如果你需要处理一个很大的矩阵：

• 不用流水线：等全部数据从 HBM 搬到 VMEM，然后计算，再全部写回 → 时间 = T_load + T_compute + T_store
• 用流水线：边搬第 1 块边计算第 0 块边写回第 -1 块 → 时间 ≈ max(T_load, T_compute, T_store)

流水线是让计算和通信重叠的关键技术。如果 T_compute > T_load，MXU 不会空闲。

Double Buffering：为了让流水线不间断，VMEM 中通常分配两个 buffer——当 MXU 从 buffer A 读取数据计算时，DMA 引擎同时把下一个 chunk 从 HBM 加载到 buffer B。交替使用，永不停歇。

GPU 的数据流

GPU 的内存层级更深、更复杂：

HBM (80-192 GB, 3.35-9 TB/s)
  ↓
L2 Cache (50-126 MB, ~5.5 TB/s on H100, 全 SM 共享)
  ↓
SMEM / L1 Cache (256 kB/SM, 极高带宽, 程序员可控)
  ↓
Registers (256 kB/SM, 每个 thread 最多 256 个)
  ↓
Tensor Core / CUDA Core

GPU vs TPU 的关键差异：

TPU 的 VMEM 大且快（2× GPU L2 容量，10× GPU L2 带宽），这让 TPU 在某些推理场景中可以把整层权重放在 VMEM 里，极大减少 HBM 访问。

GPU 的 L2 Cache 是硬件自动管理的，程序员无法直接控制——一种 “spooky action at a distance”。你需要精心设计内存访问模式来利用 L2 cache，但 132 个 SM 同时竞争共享的 50 MB L2，很容易互相 evict。这也是为什么 GPU 编程中 coalesced memory access（合并内存访问）如此重要。

📋 背景知识：Coalesced Memory Access

GPU 中，同一个 warp（32 个线程）的内存访问会被合并。如果 32 个线程访问连续的 32 个 float，硬件将其合并为一次 128 字节的事务。但如果 32 个线程访问随机位置，则需要 32 次独立事务——带宽利用率降到 1/32！

这就是为什么 GPU 代码中数据布局（AoS vs SoA）如此重要。

3.2 TPU 内部架构

TensorCore 的三大组件

TPU 本质上是一个专门做矩阵乘法的机器。每个 TPU 芯片包含 1-2 个 TensorCore，每个 TensorCore 有三个核心组件：

┌─────────────────────────────────────────────┐
│                 TensorCore                    │
│                                              │
│  ┌─────────────┐  ┌────────┐  ┌──────────┐ │
│  │   MXU ×4    │  │  VPU   │  │   VMEM   │ │
│  │ (矩阵乘法)  │  │(向量op)│  │ (128 MB) │ │
│  │ 128×128     │  │8×128×4 │  │          │ │
│  │ systolic    │  │ ALUs   │  │   高BW   │ │
│  └─────────────┘  └────────┘  └──────────┘ │
│         ↑                ↑          ↑       │
│         └────────────────┴──────────┘       │
│                    VREGs (256 KB)             │
└─────────────────────────────────────────────┘
         ↕ HBM 带宽 (0.8-2.8 TB/s)
┌─────────────────────────────────────────────┐
│              HBM (16-96 GB)                  │
└─────────────────────────────────────────────┘

MXU（Matrix Multiply Unit）：核心矩阵乘法引擎。

• 128×128 的 systolic array（TPU v6e 为 256×256）
• 每 8 个时钟周期完成一次 bf16[8,128] × bf16[128,128] → f32[8,128] 乘法
• TPU v5e 单芯片 2 个 MXU → ~2×10¹⁴ bf16 FLOPs/s
• 这是 TPU 95%+ 的有效算力来源

VPU（Vector Processing Unit）：向量运算引擎。

• 处理 ReLU、LayerNorm、softmax 等逐元素操作
• 2D SIMD 架构：8（sublane）× 128（lane）× 4 ALUs
• FLOPs 能力约为 MXU 的 1/14（~1.4×10¹³ FLOPs/s）
• 包括 reduction（求和、求最大值）操作

VMEM（Vector Memory）：片上高速缓存。

• 128 MB（TPU v5e）— 比 GPU 的 SMEM 大 4×
• 带宽 ~22× HBM 带宽
• 程序员直接控制（不像 GPU L2 是自动缓存）
• 所有数据必须先到 VMEM 才能被 MXU/VPU 使用

Scalar Core（标量核心）

Scalar Core 是 TPU 的控制单元：

• 取指令、派发工作给 MXU 和 VPU
• 执行 HBM → VMEM 的 DMA 传输调度
• 单线程——每周期只能发起一个 DMA 请求

💡 TPU 的架构极端简单：一个 Scalar Core 控制 4 个 MXU + 1 个 VPU + 所有 DMA。这使得编译器的负担很重（必须精确安排所有流水线），但硬件效率极高。

3.3 Systolic Array 工作原理

MXU 的核心是一个 128×128 的 systolic array——16,384 个 ALU 排列成网格，每个 ALU 只做乘加（multiply-accumulate）。

工作流程

权重 W（RHS, 128×128）从上方逐对角线加载：
     ↓  ↓  ↓  ↓
   ┌──┬──┬──┬──┐
→  │  │  │  │  │  ← 输入 X（LHS, 8×128）从左方逐行送入
→  │  │  │  │  │
→  │  │  │  │  │
   └──┴──┴──┴──┘
     ↓  ↓  ↓  ↓
   结果从下方流出

每个 ALU 的操作：

1. 从左边接收一个激活值 $x$
2. 从上方接收部分和 $\text{partial_sum}$
3. 计算 $\text{partial_sum} + x \times w$（其中 $w$ 是预加载的权重）
4. 将新的部分和向下传递
5. 将激活值向右传递

流水线效率

时间步：  T0    T1    T2    T3    T4    T5    T6    ...
        ┌────┐
W 加载:  │load│ (对角线加载权重，需要 128 周期)
        └────┘
        ┌────┐┌────┐┌────┐┌────┐
X 流入:  │    ││ X0 ││ X1 ││ X2 │...  (每 8 周期一组新输入)
        └────┘└────┘└────┘└────┘
                     ┌────┐┌────┐
结果流出:             │ Y0 ││ Y1 │...
                     └────┘└────┘

• 初始气泡：加载权重 + 第一组输入 ≈ 128+8 = 136 周期
• 稳态：之后每 8 周期产出一组 [8,128] 的结果
• 矩阵必须 ≥ 128：如果矩阵某维度 < 128，systolic array 部分空闲 → 需要 padding

📋 背景知识：为什么 Systolic Array 如此高效？

矩阵乘法是极少数 $O(n^3)$ 计算 / $O(n^2)$ 数据 的算法之一。每个权重值被 8 个输入复用（沿行方向），每个输入值被 128 个权重复用（沿列方向）。这意味着：

• 一旦权重加载到 systolic array 中，就可以被多次使用而无需重新加载
• 数据在 ALU 之间直接传递（不经过内存），极大减少了内存访问
• 这也是为什么 matmul 天然适合做到 compute-bound——只要 batch 足够大

TPU v6e（Trillium）的变化

TPU v6e 将 systolic array 扩大到 256×256：

• FLOPs/周期 = 4×（面积翻倍 × 每个更大 tile 复用更多）
• 但也意味着张量的最小维度需要 ≥ 256 才能充分利用
• 适合更大的模型和 batch size

3.4 GPU SM 架构

SM（Streaming Multiprocessor）结构

GPU 的 SM 类似 TPU 的 TensorCore，但 数量多、单体小：

H100 GPU（132 个 SM）:
┌─────────── SM ×132 ─────────────┐
│  ┌──────────────────────────┐   │
│  │    4 × SM Subpartition    │   │
│  │  ┌─────┐ ┌──────┐ ┌───┐ │   │
│  │  │ TC  │ │32×FP32│ │Reg│ │   │
│  │  │     │ │CUDA   │ │16K│ │   │
│  │  │8×8×8│ │Cores  │ │   │ │   │
│  │  └─────┘ └──────┘ └───┘ │   │
│  └──────────────────────────┘   │
│  ┌──────────────────────────┐   │
│  │  SMEM / L1 Cache (256KB) │   │
│  └──────────────────────────┘   │
└─────────────────────────────────┘
          ↕ L2 Cache (50 MB, 共享)
          ↕ HBM (80 GB, 3.35 TB/s)

Tensor Core vs MXU

GPU 有更多更小的 matmul 单元 → 更灵活（可以同时做不同大小的 matmul），但也更难达到峰值利用率。

CUDA Core 与 Warp 调度

📋 背景知识：SIMT vs SIMD

• SIMD（TPU VPU）：所有 ALU 执行完全相同的操作，完全同步
• SIMT（GPU CUDA Core）：每个线程有自己的程序计数器（PC），但一个 warp（32 线程）同时执行相同指令

SIMT 的优势：可以优雅地处理条件分支（divergence）——不同线程走不同分支时，硬件自动 mask 不活跃的线程。代价：如果 32 个线程走了不同的 if/else 分支，两个分支都要执行（时间翻倍）。

这也是为什么 GPU 代码要避免 warp divergence——让同一个 warp 的线程尽量走相同路径。

每个 SM 可以有多达 64 个 resident warps（2048 线程），但每周期只执行其中几个。当一个 warp 等待内存时，调度器自动切换到另一个 warp → 用 warp 切换隐藏内存延迟（类似 CPU 的超线程但更极端）。

Occupancy（占用率）与性能

SM 的占用率定义为：

高占用率让调度器有更多 warp 可切换，更好地隐藏内存延迟。但占用率受限于三个因素：

🛠️ 实践：理解 GPU kernel 的 Occupancy

在 Megatron-LM 或 SGLang 中使用 Flash Attention kernel 时，可以用 NVIDIA 的 Occupancy Calculator 分析：

# 使用 ncu (Nsight Compute) 分析 kernel 的 occupancy
ncu --metrics sm__warps_active.avg.pct_of_peak_sustained_active \
    --metrics sm__maximum_warps_per_active_cycle_pct \
    python train.py

# 典型 Flash Attention kernel 的 occupancy:
# - 理论 occupancy: ~50% (受限于 SMEM 使用量)
# - 实际 achieved occupancy: ~45%
# - 但性能仍然很好，因为 compute-bound 操作不需要高 occupancy

关键洞察：Compute-bound 的操作（如 matmul）不需要高 occupancy——计算单元已经满负荷了，不需要靠 warp 切换来隐藏延迟。Memory-bound 操作（如 LayerNorm）才需要高 occupancy 来隐藏 HBM 延迟。

3.5 VMEM 预取与流水线重叠

权重预取策略

TPU 的 VMEM 预取是一个关键优化。在 Transformer 层的执行中：

时间 →
┌────────────────┬────────────────┬────────────────┐
│  Attention     │     FFN        │  Attention     │
│  (计算)        │    (计算)      │  (计算)        │
├────────────────┼────────────────┼────────────────┤
│ 预取 FFN 权重  │ 预取 Attn 权重 │ 预取 FFN 权重  │
│ HBM→VMEM      │  HBM→VMEM     │  HBM→VMEM     │
└────────────────┴────────────────┴────────────────┘

效果：当 Attention 计算完成时，FFN 权重已经在 VMEM 里了 → FFN 的 matmul 可以立即开始，不需要等 HBM 加载。

前提条件：

• 权重必须足够小（或经过分片）才能放进 VMEM（128 MB）
• Attention 的计算时间必须 ≥ FFN 权重的加载时间

VMEM 与算术强度

VMEM 带宽 ≈ 22× HBM 带宽。这意味着：

• 从 HBM 读取时：compute-bound 的临界算术强度 = FLOPs/s ÷ HBM BW ≈ 240
• 从 VMEM 读取时：compute-bound 的临界算术强度 = FLOPs/s ÷ VMEM BW ≈ 11

如果能把数据放在 VMEM 里，batch size 只需 11（而非 240）就能达到 compute-bound！ 这对推理特别有意义——小 batch 的推理通常是 memory-bound，但如果权重在 VMEM 中，就可以在极小 batch 下也达到 compute-bound。

GPU 的对应机制：L2 Cache Residency

GPU 没有 VMEM 预取的直接等价物，但 L2 cache 提供了类似的功能：

• L2 带宽 ≈ 1.6× HBM（5.5 TB/s vs 3.35 TB/s on H100）
• 但容量只有 50 MB → 只能缓存非常小的权重

NVIDIA 提供了 cudaAccessPolicyWindow API 让程序员”暗示”哪些数据应该留在 L2 中，但效果远不如 TPU 的 VMEM 直接控制。

🛠️ 实践：Megatron-LM 与内存层级

Megatron 中利用内存层级的关键配置：

--use-flash-attn              # Flash Attention: 利用 SMEM 做 attention tiling
--recompute-activations       # 重计算代替存储（trade compute for memory）
--distribute-saved-activations  # 将激活分布到多卡
--sequence-parallel            # 减少每卡激活量

Flash Attention 的本质就是利用内存层级差异：

• 标准 Attention：QK^T（N×N 矩阵）写入 HBM → softmax → 读回 → ×V → 写入 HBM
• Flash Attention：在 SMEM 中 tiling 完成 QK^T + softmax + ×V → 只写最终结果到 HBM
• 减少 HBM 访问量从 O(N²) 到 O(N)！

3.6 PCIe：CPU 与加速器之间的桥梁

加速器通过 PCIe 连接到 CPU 主机。这是整个带宽层级中最慢的环节之一。

带宽对比

PCIe 比 HBM 慢 ~67×！ 任何经过 PCIe 的操作都是潜在瓶颈。

PCIe 的用途

1. 数据加载：CPU DRAM → GPU HBM（训练数据）
2. Host Offload：把优化器状态/KV cache 卸载到 CPU 内存以节省 HBM
3. 跨节点通信路径：GPU → PCIe → NIC → InfiniBand → NIC → PCIe → GPU

PCIe 的算术强度分析

如果某个操作需要通过 PCIe 加载数据，它的有效”PCIe Roofline”是：

这意味着任何需要经过 PCIe 的操作，算术强度必须 > 19,800 才能不被 PCIe 瓶颈！标准 matmul 的算术强度只有 batch_size 级别——说明从 CPU 加载权重做 matmul 是完全不可接受的。

💡 这就是为什么 GPU offloading（如 DeepSpeed ZeRO-Infinity）只在特殊场景有价值——PCIe 太慢，只有极大的 matmul（算术强度 > 20000）才不被 PCIe 拖慢。实际中主要用于优化器状态 offload（不在关键路径上）。

🛠️ 实践：SGLang 中的 CPU Offloading

SGLang 的 RadixAttention 通过将冷 KV Cache 卸载到 CPU 来节省 GPU HBM：

# SGLang 的 KV Cache offloading 策略
# 当 GPU HBM 满时，将最久未使用的 KV Cache 搬到 CPU
# 关键约束：只有等下次该请求被调度时才需要搬回
# 搬回时间 = kv_size / PCIe_BW
#
# 例：1 个请求的 KV Cache (LLaMA 70B, seq=4096)
# = 2 × 80 layers × 8 heads × 128 dim × 4096 × 2 bytes
# = 1.28 GB
# 搬回时间 = 1.28 GB / 50 GB/s = 25.6 ms
#
# 这 25.6 ms 可以和其他请求的 decode 计算重叠！
# 所以 CPU offloading 在推理中比训练中更实用。

训练中 PCIe offloading 只适合优化器状态（非关键路径），但推理中通过 scheduling 可以隐藏 PCIe 延迟。

3.7 内存容量约束

除了带宽，内存容量也是关键约束。训练/推理一个模型需要在 HBM 中同时存储多种数据。

训练时的内存占用

激活值的占用取决于模型架构、batch size 和序列长度：

其中 $L$ = 层数，$s$ = 序列长度，$b$ = micro batch size，$h$ = hidden size，$a$ = attention heads。

LLaMA 70B 详细内存预算

单张 H100（80 GB）远远不够：

这就是为什么训练 70B 模型至少需要 14 张 H100——仅存放权重和优化器状态！加上激活值和通信缓冲区，实际需要更多。

推理时的内存占用

推理时不需要梯度和优化器，但引入了 KV Cache：

KV Cache 会随 batch size 和序列长度线性增长，很容易超过参数本身的内存占用！

📋 背景知识：KV Cache 为什么这么大？

在自回归生成中，每个 token 的 Attention 需要访问所有之前 token 的 Key 和 Value。为避免重复计算，这些中间值被缓存。对于 LLaMA 70B：

• 80 层，每层 8 个 KV heads（GQA），head_dim = 128
• 每 token 的 KV = $2 \times 80 \times 8 \times 128 \times 2$ bytes = 327,680 bytes ≈ 320 KB
• 4096 个 token 的完整 KV Cache = 4096 × 320 KB ≈ 1.28 GB（每请求）

当 batch=256 时，KV Cache 总量 = 256 × 1.28 GB ≈ 328 GB——比模型权重（140 GB）还大！

内存优化策略总览

🛠️ 实践：Megatron-LM 中的内存管理

Megatron-LM 提供了多种内存优化配置：

# ZeRO-1: 只分片优化器状态（最常用）
--use-distributed-optimizer

# Gradient Checkpointing: 只保留每层输入，反向时重计算
--recompute-activations
--recompute-granularity full    # 重计算整个 Transformer 层
--recompute-granularity selective  # 只重计算 attention（更常用）

# Sequence Parallelism: 沿序列维度分片激活值
--sequence-parallel

# 混合精度
--bf16                          # 使用 bf16 训练
--fp32-residual-connection      # 残差连接保持 fp32

实际内存占用估算（LLaMA 70B, 8×H100, TP=8）：

• 参数：140 GB / 8 = 17.5 GB/卡
• 优化器：840 GB / 8 = 105 GB（不分片时放不下！需要 ZeRO-1 + DP 分片）
• 激活值（selective recompute, mbs=1, seq=4096）：~8 GB/卡
• 总计：~30 GB/卡（含通信缓冲）→ 80 GB 足够

3.8 关键数字速记

你应该记住的几个量级（2024-2025 代硬件）：

带宽层级

容量层级

计算能力

记忆法：

• 带宽从上到下大约每级慢 10×
• 容量从上到下大约每级大 100-1000×
• 速度 × 容量 ≈ 常数（物理定律的体现）

💡 速算技巧

需要估算某操作是否 compute-bound？快速公式：

• 从 HBM 读取：临界 AI ≈ 240-570（取决于芯片）
• 从 VMEM 读取：临界 AI ≈ 11（TPU）
• 从网络读取：临界 AI ≈ FLOPs/s ÷ 网络 BW（通常 2000-20000）
• 从 PCIe 读取：临界 AI ≈ 20000（几乎不可能达到）

3.9 Worked Problems（习题与详解）

Problem 1：Flash Attention 的 HBM 访问量

题目：考虑标准 Attention 和 Flash Attention 对 HBM 的访问量。假设序列长度 $N = 4096$，head dimension $d = 128$，batch size $B = 1$，单个 attention head。所有数据为 bf16。

1. 标准 Attention 的 HBM 读写字节总量是多少？
2. Flash Attention（tile size $B_r = B_c = 256$）的 HBM 读写字节总量是多少？
3. 在 H100 上（HBM 3.35 TB/s），仅考虑内存访问时间，两者各需多长时间？

Problem 2：VMEM 预取的最大权重尺寸

题目：在 TPU v5e 上，VMEM 容量为 128 MB，HBM 带宽为 0.82 TB/s。

1. 如果一个 Transformer 层的 FFN 有两个权重矩阵 $W_1[D, 4D]$ 和 $W_2[4D, D]$（bf16），D 的最大值是多少才能让两个矩阵都放进 VMEM？
2. 如果 Attention 计算需要 $T_{attn}$ 时间，FFN 权重预取需要 $T_{prefetch}$ 时间，当 $D = 4096$ 时，batch size $B$ 至少为多少才能保证预取完全被计算隐藏？

Problem 3：Systolic Array 利用率

题目：TPU v5e 有 128×128 的 systolic array。考虑以下矩阵乘法：

1. bf16[64, 256] × bf16[256, 128]：MXU 利用率是多少？
2. bf16[8, 64] × bf16[64, 64]：MXU 利用率是多少？
3. 对于 bf16[8, D] × bf16[D, F]，D 和 F 至少为多少才能达到 100% 利用率？

Problem 4：端到端 TPU 数据搬运问题

题目（改编自原书 Question 6）：你有一个大矩阵 A: int8[128×1024, 128×1024]（约 16 GB），均匀分片在 TPU v5e 4×4 slice（16 个芯片，2 个 host）上，存储在 host DRAM 中。你想把整个矩阵收集到 TPU{0,0} 并与一个向量 bf16[8, 128×1024] 相乘。

数据：PCIe 带宽 15 GB/s/链路，ICI 带宽 45 GB/s/方向/链路（双向 90 GB/s），HBM 带宽 820 GB/s，bf16 FLOPs/s = 1.97×10¹⁴。4×4 slice 中每个 TPU 有 2 条 ICI 链路。

1. 最优的数据搬运方案是什么？
2. 每一步需要多长时间？
3. 总时间的下界和上界是多少？

关键要点

• 数据必须从 HBM → 片上缓存 → 计算单元，层层搬运；每层速度差 10×
• 流水线化让计算和搬运可以重叠，时间 ≈ max(计算, 搬运)；Double Buffering 实现无停顿流水线
• TPU：VMEM 128 MB、程序员直接控制、带宽 22× HBM；GPU：L2 50 MB、硬件自动管理
• MXU 是 128×128 systolic array，每 8 周期完成 bf16[8,128]×bf16[128,128]；张量维度需 ≥ 128 才能充分利用
• GPU SM 更多更小（132 个 SM × 4 Tensor Core），灵活但难达峰值；SIMT 支持分支但需避免 warp divergence
• VMEM 预取使临界算术强度从 240 降到 11，让小 batch 推理也能 compute-bound
• Flash Attention 将 HBM 访问从 $O(N^2)$ 降到 $O(N^2d^2/M)$，利用 SMEM 做 tiling
• PCIe 比 HBM 慢 ~67×（临界 AI ≈ 19,800），GPU offloading 仅适合非关键路径
• 训练 70B 模型需要每参数 ~16 bytes → 总计 1.1 TB → 至少 14 张 H100
• KV Cache 随 batch×seq 线性增长，可能超过模型参数本身的内存占用

进一步阅读

• 原书 Chapter 2: What Is a TPU → Memory & MXU 部分
• 原书 Chapter 12: What Is a GPU → Memory 部分
• 原书 Appendix B: How does a systolic array work
• Flash Attention 论文 (Dao et al., 2022) — HBM 访问量分析
• Flash Attention 2 (Dao, 2023) — 利用 warp 级并行进一步优化
• NVIDIA H100 Whitepaper — SM 架构详解
• Megatron-LM Activation Recomputation — Selective Recomputation 策略
• Systolic Array 动画 (fleetwood.dev) — 交互式理解脉动阵列
• How to Optimize a CUDA Matmul (Simon Boehm) — GPU 内存层级与 tiling 实战