CS336 · 第 5 课

GPU、TPU 架构与性能

同样数量的 FLOPs，为什么有的算子接近硬件峰值，有的却大部分时间都在等待数据？本讲从计算单元、memory hierarchy 和执行模型出发，建立判断 kernel 瓶颈的硬件视角，并比较 GPU 与 TPU 的共同目标和不同取舍。

这一讲到底解决什么问题？

Lecture 2 已经教会我们核算 FLOPs 和 bytes，但公式还没有解释：为什么相同 FLOPs 在不同算子上会得到完全不同的实际时间？答案取决于数据在硬件中的位置，以及程序能否让足够多的计算同时进行。

HBM → L2 → shared memory / L1 → registers → arithmetic units
大/慢                                                   小/快

kernel workload
├─ math FLOPs
├─ memory traffic bytes
├─ parallel work / synchronization
└─ resource use per block（registers, shared memory）
     ↓
occupancy 与 achieved throughput

先看 GPU 为什么适合这类 workload

CPU vs GPU

A100 SM (Streaming Multiprocessor) 结构:
├── 108 个 SM
├── 每个 SM 包含:
│   ├── 64 FP32 CUDA cores (INT32 也可)
│   ├── 4 个 Tensor Cores (第三代)
│   │   └── 支持 TF32, FP16, BF16, FP64 等矩阵指令模式
│   ├── 192 KB L1 cache / Shared Memory (可配置)
│   └── 寄存器文件 (256 KB)
├── 40 MB L2 cache (共享)
└── 80 GB HBM2e (高带宽内存)
    带宽: 2 TB/s

H100 SM (Streaming Multiprocessor) 结构:
├── H100 SXM 产品启用 132 个 SM
├── 每个 SM 包含:
│   ├── 128 FP32 CUDA cores
│   ├── 4 个 Tensor Cores (第四代)
│   │   └── 支持 FP8, FP16, BF16, TF32
│   ├── 256 KB unified L1 / Texture / Shared Memory
│   │   └── 其中 shared-memory carveout 最高 228 KB/SM
│   └── 寄存器文件 (256 KB)
├── 50 MB L2 cache
└── 80 GB HBM3
    带宽: 3.35 TB/s

GPU 内存层次

从慢到快，从大到小

内存带宽 vs 计算速度

H100 (SXM5) 关键数字:
  计算峰值 (bf16): 989 TFLOP/s
  内存带宽:       3.35 TB/s
  Ridge point:    989e12 / 3.35e12 ≈ 295 FLOP/byte

  这意味着: 在只考虑峰值算力与 HBM 带宽的理想 Roofline 中，
  算术强度约高于 295 FLOP/byte 才可能进入 compute-bound 区域。
  它描述的是吞吐平衡点，不是单次访存延迟。

GPU 执行模型

Thread, Block, Grid

# GPU 编程的基本抽象
Grid
 └── Thread Block (CTA)  ← 在单个 SM 上执行
      └── Warp (32 threads)  ← 以 SIMD 方式执行
           └── Thread        ← 基本执行单位

# 为什么是 Warp？
# Warp 调度器向 32 个线程发出共同指令，各线程保有独立状态
# NVIDIA 将这种抽象称为 SIMT，而不是传统向量机意义的 SIMD

SIMT 执行模型

SIMT = Single Instruction, Multiple Thread

# 坏: 分支发散
if threadIdx.x % 2 == 0:
    do_A()  # 16 个线程执行
else:
    do_B()  # 16 个线程执行
# 两条路径可能被分别执行，不能再让全部 lane 同时做有效工作

# 好: 无分支
# 让所有线程走相同路径

Occupancy

Occupancy = 活跃 warp 数 / 最大支持 warp 数

高 occupancy 意味着 SM 有更多的 warp 可以调度——当一个 warp 等待内存时，SM 可以切换到另一个 warp。

Roofline Model

概念

Roofline model 直观地展示一个操作的性能上限：

性能 (FLOP/s)
    ↑
    │        ┌──────── 计算上限 (compute ceiling)
    │       /
    │      /  ← 斜率 = 内存带宽
    │     /
    │    /
    │───/───────────────→ 算术强度 (FLOP/byte)
    │  ridge point

• 低算术强度：在斜线上 → memory-bound，性能 = 带宽 × 算术强度
• 高算术强度：在平顶下 → compute-bound，性能 = 计算峰值

常见操作的 Roofline 位置

优化策略

Tensor Cores

什么是 Tensor Core？

Tensor Core 是 GPU 上的专用硬件单元，专门做矩阵乘法：

使用 Tensor Core 的条件

• 数据类型与累加模式必须由该代 Tensor Core 支持；FP32 输入常通过 TF32 模式执行，而不是保留完整 23-bit fraction 参与乘法
• 现代 cuBLAS 即使维度未对齐也可能使用 Tensor Core，但按字节和 tile 对齐通常效率更高
• 最佳维度倍数随数据类型、GPU、库版本和所选 kernel 改变，不能简化成永远是 8 或 16

TF32：两全其美

内存访问模式

Coalesced Access（合并访问）

Global memory 请求会被合并为尽可能少的内存事务；现代 NVIDIA GPU 的缓存行由 32-byte sectors 组成。warp 访问连续且对齐的地址时，所需 sector 更少：

Shared Memory Bank Conflicts

Shared memory 被分为 32 个 bank，每个 bank 宽 4 字节：

GPU 与 TPU：相同约束，不同执行抽象

TPU：相似目标，不同执行模型

Spring 2026 官方 Lecture 5 明确把 TPU 放进 GPU 旁边比较。两者都把轻量控制、大型矩阵单元、片上快速存储与 HBM组合在一起，但 TPU 的核心抽象更偏向编译器调度的大块矩阵程序，而不是 GPU 的 warp/SIMT 线程层次。

CourseStack 扩展边界：systolic array、JAX/XLA、ICI/NVLink 细节是用于建立工程对照的补充，不应被误读为 Lecture 5 对具体 TPU generation 的完整规格说明。官方讲义的核心结论是 GPU 与 TPU 都受矩阵单元、HBM 带宽和互联约束，但编程抽象与网络组织不同。具体 sharding 选择在第 7–8 课继续展开。

仅用于展示 units：GPT-3 175B 报告的训练计算量约 3.14e23 FLOPs，
# 1024 张 GPU 全程保持 MFU=0.5，且忽略停机、评估和 I/O。

# A100 (312 TFLOP/s dense bf16, MFU=0.5):
#   单卡实际: 156 TFLOP/s
#   1024 卡: 1.6e17 FLOP/s
#   时间: ~22.8 天

# H100 (989 TFLOP/s dense bf16, MFU=0.5):
#   单卡实际: 495 TFLOP/s
#   1024 卡: 5.07e17 FLOP/s
#   时间: ~7.2 天

# 峰值 dense bf16 比值约 3.17×；真实端到端加速不会只由峰值算力决定。

多 GPU 互联

节点内互联

节点间互联

通信层次

// 最底层，最灵活，最复杂
__global__ void vector_add(float *a, float *b, float *c, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        c[i] = a[i] + b[i];
    }
}

// 启动 kernel
vector_add<<<num_blocks, block_size>>>(d_a, d_b, d_c, n);

# Python 风格的 GPU kernel DSL；仍需设计 program shape 与访存
@triton.jit
def vector_add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offsets < n
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    tl.store(out_ptr + offsets, x + y, mask=mask)

# 不手写 kernel，让编译器尝试图优化与融合
@torch.compile
def fused_operation(x, w):
    return gelu(x @ w)  # 是否做 epilogue/kernel fusion 取决于后端与形状

知识检查

如果只能记住 5 件事

1. GPU 通过大量并行执行单元追求 throughput；性能来自并行度，不是单个 thread 更快。
2. Registers、shared memory/L1、L2 与 HBM 构成容量和速度不同的数据层次，优化往往是在减少更慢层级的流量。
3. Roofline 用 arithmetic intensity 连接 FLOPs 与 bytes：瓶颈位置决定应该优化计算还是数据移动。
4. Tensor Core 峰值只有在 dtype、shape、layout 和软件实现匹配时才有意义；硬件规格不是应用性能保证。
5. GPU 与 TPU 都围绕矩阵单元、片上存储、HBM 和互联设计，但执行抽象与编译器分工不同。

延伸阅读