CUDA 算子工程:手写 FlashAttention v2 之路

第 1 章 GPU 范式:从 CPU 思维到 SIMT 思维

作者 杨艺韬 · 5,105 字

第 1 章 GPU 范式:从 CPU 思维到 SIMT 思维

"Latency hiding is the central design principle of GPU architecture. Everything else follows from it." ——David Kirk & Wen-mei Hwu, Programming Massively Parallel Processors, 4th ed.

1.1 一个朴素的问题:GPU 比 CPU 快在哪

很多人第一次接触 GPU 编程时,脑子里都装着一个错误的类比:

"CPU 是 16 核的,GPU 是 10000 核的,所以 GPU 比 CPU 快 600 倍。"

这个类比听起来直观,但它几乎全错。

要看到错在哪,先把 2025 年的两款主流硬件摊开对比。一边是 Intel 第 5 代 Xeon Scalable(Emerald Rapids)8592+,一边是 NVIDIA H100 SXM5。

维度 Xeon 8592+ (CPU) H100 SXM5 (GPU)
核心数 64 132 SM × 64 FP32 cores = 8448
主频 3.9 GHz boost 1.83 GHz boost
FP32 算力 ~5 TFLOPs (AVX-512) ~67 TFLOPs (SIMT)
FP16 算力(Tensor 加速) ~10 TFLOPs (AMX) 989 TFLOPs (Tensor Core)
内存带宽 ~350 GB/s (DDR5-5600 × 8) 3350 GB/s (HBM3)
单线程指令延迟(典型) 1–5 cycles 4–24 cycles
缓存层级 L1 80KB / L2 4MB / L3 320MB L1+SMEM 256KB/SM / L2 50MB
TDP 350 W 700 W
价格(2025) ~$11,000 ~$30,000

CPU 数据来源:Intel® Xeon® Platinum 8592+ Product Specification (Q4 2023)。 GPU 数据来源:NVIDIA H100 Tensor Core GPU Architecture Whitepaper (2022)。

几个有趣的对比:

  1. GPU 的算力优势主要来自 Tensor Core:纯 FP32 SIMT 算力 H100 比 Xeon 高 13×,但 FP16 Tensor Core 对比 CPU 的 AMX 高接近 100×。如果你的工作负载用不到矩阵乘,GPU 的"算力优势"会被砍掉一个数量级。
  2. 内存带宽差距比算力差距更稳定:H100 是 Xeon 的 ~10×。这个比例从 Pascal 时代就大致维持。所以带宽优势是 GPU 真正的护城河——不是核心数。
  3. GPU 的单指令延迟反而更长:H100 的浮点加法、内存读取,单条指令的延迟(24+ cycles)比 CPU 的同类操作(4 cycles)要长。
  4. GPU 的 cache 容量小得离谱:每个 H100 SM 只有 256KB 的 L1+SMEM 容量,这是 Xeon 单核 L2 的 1/16。

第 3 条和第 4 条尤其反直觉——既然 GPU 单指令更慢、cache 更小,它凭什么"快"?

答案是:GPU 用大量并行的线程来掩盖单指令的慢。它不是让每个线程跑得快,而是让足够多的线程同时在跑,使得任何一个线程在等内存的时候,总有别的线程能算

这就是 GPU 设计哲学的第一个分水岭:延迟隐藏(Latency Hiding)

1.2 两条岔路:延迟优化 vs 吞吐优化

CPU 和 GPU 的差异不是"性能高低"的差异,而是"优化目标"的差异。

flowchart TB
  subgraph CPU [CPU · 延迟优化]
    C1[少量复杂核心]
    C2[深度乱序执行]
    C3[庞大分支预测]
    C4[多级大容量 cache]
    C5[强一致性内存模型]
    C1 --> C2 --> C3 --> C4 --> C5
  end
  subgraph GPU [GPU · 吞吐优化]
    G1[大量简单核心]
    G2[顺序执行]
    G3[无分支预测]
    G4[小容量但快速 SMEM]
    G5[弱一致性, 显式同步]
    G1 --> G2 --> G3 --> G4 --> G5
  end
  Goal1[目标: 单线程任务尽快完成] -.-> CPU
  Goal2[目标: 大量线程总吞吐最大化] -.-> GPU

把这两套设计哲学落到具体的硬件预算上,差异就更直观了。

1.2.1 CPU 是怎么花晶体管的

一颗现代 x86 CPU 核心的晶体管预算大致是这样分配的:

                    现代 CPU 单核晶体管分配(约 2-3 亿晶体管)
┌──────────────────────────────────────────────────┐
│ L1 / L2 cache:                          ~40%     │
│ 乱序执行 + 寄存器重命名:                ~20%     │
│ 分支预测器 + 分支目标缓冲:              ~10%     │
│ 解码器 (μop cache):                     ~10%     │
│ 实际算术单元 (ALU/FPU/AVX):             ~10%     │
│ 其他 (load/store unit, retire, etc.):   ~10%     │
└──────────────────────────────────────────────────┘

比例数据:综合 Intel/AMD 公开 die shot 与 Anand Lau Lai Shun's microarchitecture analysis,仅作量级参考。

注意最后一行:真正的算术单元只占了大约 10% 的晶体管。其余 90% 都在做"让单线程跑得更快"的辅助工作——预测下一条指令、并行执行多条无依赖指令、缓存可能要用的数据、记住每个寄存器最新的影子版本。

为什么花这么多晶体管在"辅助"上?因为 CPU 的设计目标是:给我一个串行任务,让它尽快跑完。当任务本身没有并行性(比如指针追踪、解析 JSON、跑一段单线程的业务逻辑),CPU 只能靠"在一条指令链上压榨更多并行"——也就是指令级并行(ILP, Instruction-Level Parallelism)。乱序执行、寄存器重命名、分支预测,全都是为 ILP 服务的。

1.2.2 GPU 是怎么花晶体管的

一个 Hopper SM 的晶体管预算分配截然不同:

                    Hopper SM 晶体管分配(约 30 亿晶体管 / SM × 132 SMs)
┌──────────────────────────────────────────────────┐
│ Tensor Core (4 个/SM):                   ~30%    │
│ FP32/INT32 SIMT cores (64 个/SM):        ~15%    │
│ 寄存器文件 (256KB/SM):                   ~25%    │
│ Shared Memory + L1 cache (256KB/SM):     ~15%    │
│ Warp scheduler (4 个/SM):                ~5%     │
│ Load/Store Unit, MMU, etc.:              ~10%    │
└──────────────────────────────────────────────────┘

比例数据:综合 NVIDIA Hopper Whitepaper 与公开 die shot 的量级估算。

一个鲜明对比:GPU 把约 45% 的晶体管花在了算术单元上(Tensor Core + SIMT cores),而 CPU 这个比例只有 ~10%。GPU 上几乎所有 CPU 引以为豪的"复杂逻辑"(乱序、分支预测、巨型 cache)全部砍掉,省下的晶体管全部投到算术单元和寄存器上。

但 GPU 砍掉这些"复杂逻辑",难道不会让性能掉得很惨吗?

会。GPU 的单指令性能确实比 CPU 差——一个 Hopper SIMT core 跑一条 FP32 加法需要 4 个周期,而 Xeon 跑一条 AVX-512 加法只要 1 个周期。GPU 单核根本打不过 CPU 单核。

但 GPU 的玩法从一开始就不是"单线程跑得快"。它的玩法是:同时跑成千上万个线程,让任何时刻都有线程可以执行。这就引出了 GPU 的核心范式——SIMT。

1.3 SIMT:GPU 的核心范式

SIMT 是 NVIDIA 在 2007 年 Tesla 架构的 G80 上正式引入的术语,全称是 Single Instruction, Multiple Threads——单指令多线程。

很多人把 SIMT 和 SIMD(Single Instruction, Multiple Data)混为一谈。它们确实相似,但差别足以把"会写 SIMD"的工程师在 GPU 上撞墙好几次。

flowchart LR
  subgraph SIMD [SIMD · 一个核控多条数据线]
    INS1[一条 SIMD 指令: VADD]
    INS1 --> LANE1[Lane 0]
    INS1 --> LANE2[Lane 1]
    INS1 --> LANE3[Lane 2]
    INS1 --> LANE4[Lane 3]
    LANE1 -.同步执行.-> RESULT1[结果寄存器]
    LANE2 -.同步执行.-> RESULT1
    LANE3 -.同步执行.-> RESULT1
    LANE4 -.同步执行.-> RESULT1
  end
  subgraph SIMT [SIMT · 一个 warp 调一组线程]
    WARP[一个 warp 中的 32 个线程]
    WARP --> T0[线程 0 有自己的寄存器]
    WARP --> T1[线程 1 有自己的寄存器]
    WARP --> T2[线程 2 有自己的寄存器]
    WARP --> TN[... 共 32 个线程]
    T0 -.独立的 PC, 独立的状态.-> SCHED[硬件调度器决定哪几个一起发指令]
    T1 -.独立的 PC, 独立的状态.-> SCHED
    T2 -.独立的 PC, 独立的状态.-> SCHED
    TN -.独立的 PC, 独立的状态.-> SCHED
  end

SIMD 和 SIMT 的关键差异有三点:

  1. 抽象层级不同:SIMD 是一条指令直接操作多个数据 lane,每个 lane 没有"自己的状态";SIMT 是一组(32 个)独立线程,每个线程有自己的程序计数器、自己的寄存器、自己的执行状态——只是它们恰好被硬件捆在一起、用同一条指令推进
  2. 编程模型不同:写 SIMD(比如 AVX-512、ARM SVE)需要显式地组织 lane、用 intrinsic 函数、对齐数据;写 SIMT 你就像在写普通的标量代码,由编译器和硬件去打包线程到 warp 里。
  3. divergence 处理不同:SIMD 遇到分支就需要程序员手写 mask;SIMT 硬件自动处理 divergence——当一个 warp 里 32 个线程走不同分支,硬件会自动让它们串行执行不同分支,但写代码的人不需要显式 mask。

第 3 点是 SIMT 模型最大的"魔法",也是它最大的"陷阱"。

1.3.1 Warp:SIMT 的最小调度单位

在 NVIDIA GPU 上,每 32 个线程组成一个 warp。warp 是硬件调度的最小单位——硬件每次给一个 warp 的 32 个线程发同一条指令。

为什么是 32?这个数字从 G80 一路沿用到 Hopper,没变过。原因是工程权衡:太小(比如 8)的 warp 会让指令带宽浪费在元数据上;太大(比如 128)的 warp 会让 divergence 代价太高。32 是 NVIDIA 工程师在 2006 年给出的折中。

每个 SM 上有 4 个 warp scheduler(Hopper 上是 4 个,Volta/Ampere 也是 4 个)。每个 warp scheduler 每周期可以发射一条指令到一个 warp。所以一个 SM 每周期最多发射 4 条指令到 4 个不同的 warp,对应 4 × 32 = 128 个线程的一条指令

flowchart TB
  SM[一个 Hopper SM]
  SM --> WS1[Warp Scheduler 0]
  SM --> WS2[Warp Scheduler 1]
  SM --> WS3[Warp Scheduler 2]
  SM --> WS4[Warp Scheduler 3]

  WS1 --> WP0[每周期发一条指令到 1 个 warp 的 32 个线程]
  WS2 --> WP1[每周期发一条指令到 1 个 warp 的 32 个线程]
  WS3 --> WP2[每周期发一条指令到 1 个 warp 的 32 个线程]
  WS4 --> WP3[每周期发一条指令到 1 个 warp 的 32 个线程]

但这只是 4 个 active warps。Hopper 一个 SM 的寄存器文件可以同时驻留 64 个 warps(2048 active threads/SM)。其余 60 个 warp 怎么办?它们处于"待命"状态——硬件维护着它们的 PC、寄存器、谓词,只在发现某个正在执行的 warp 卡住(比如等内存读完成)时,立刻切换到一个准备好的 warp 上跑

这个切换过程是零开销的——和 CPU 的线程切换不同(CPU 切换上下文要保存到内存,开销几百到几千个周期),GPU 的 warp 切换发生在硬件层面,所有 warp 的寄存器都常驻在 SM 的寄存器文件里。切换只是改一下 warp scheduler 的指针,0 周期。

这就是 GPU 延迟隐藏的本质:

  CPU 视角:                      GPU 视角:
  ┌──────────────────────┐      ┌──────────────────────┐
  │ Thread A 等内存(200c)│      │ Warp 0 等内存(800c)  │
  │ 此时核心闲置...      │      │   切到 Warp 1 算1c   │
  │ 内存到了, A 继续算   │      │   切到 Warp 2 算1c   │
  └──────────────────────┘      │   切到 Warp 3 算1c   │
                                │   ... 切 800 次 ...  │
                                │ Warp 0 内存到了      │
                                │ 此时核心一直没闲着!  │
                                └──────────────────────┘

CPU 看到一个线程等内存,会用复杂的乱序执行、prefetcher、缓存预测来填满这段空闲;GPU 看到一个 warp 等内存,直接换一个 warp 来跑。哪种思路更好?取决于你有多少独立的工作。

如果工作完全串行(一个长的依赖链),CPU 赢——它能把这条链跑得快。如果工作高度并行(几十万个独立的小任务),GPU 赢——它能让所有核心一刻不停。

LLM 推理与训练的核心算子(GEMM、Attention、LayerNorm)都属于后者。这是 GPU 在 LLM 时代成为绝对统治者的根本原因。

1.3.2 一个 warp 内的 divergence

SIMT 的"自动处理 divergence"听起来很美,但代价是真实的。看这段代码:

__global__ void diverge_kernel(int* arr) {
    int tid = threadIdx.x;
    if (tid % 2 == 0) {
        arr[tid] = work_a();   // 偶数线程做 work_a
    } else {
        arr[tid] = work_b();   // 奇数线程做 work_b
    }
}

一个 warp 里的 32 个线程,16 个走 if 分支,16 个走 else 分支。SIMT 硬件怎么办?

它会先让 16 个偶数线程执行 work_a(),此时 16 个奇数线程被 mask 掉(不执行,但占着 ALU 周期);然后让 16 个奇数线程执行 work_b(),此时偶数线程被 mask 掉

sequenceDiagram
    participant Hardware as Warp Scheduler
    participant Even as 16 偶数线程
    participant Odd as 16 奇数线程
    Note over Hardware: 进入 if 分支
    Hardware->>Even: 执行 work_a()
    Hardware->>Odd: mask 掉, 跟着走但不写结果
    Note over Hardware: 切到 else 分支
    Hardware->>Odd: 执行 work_b()
    Hardware->>Even: mask 掉, 跟着走但不写结果
    Note over Hardware: 两条路径都走完, warp 重新汇合

结果是:这个 kernel 的实际算力是非 divergence 情况的一半。因为同样的 32 个 ALU 周期,只有 16 个线程在干活。

这就是 SIMT 的代价。它不会编译报错,不会运行时崩溃,只会性能腰斩

在 Volta 之前(Maxwell/Pascal/Kepler),warp 内 divergence 后线程不会自动重新汇合到 reconvergence point——所有线程必须等 warp 退出函数才能重新一起跑。Volta 引入了 Independent Thread Scheduling,每个线程有了独立的 PC,理论上可以更细粒度地汇合。但这并不消除 divergence 的代价,只是让它更可控。

1.3.3 Memory Coalescing:访存的 SIMT 化

SIMT 不仅影响指令执行,还深刻影响内存访问。

GPU 的内存控制器以 transaction(事务)为单位读 HBM,每次最少读 32 字节。而一个 warp 里 32 个线程通常会同时发起内存读。如果这 32 个线程正好读连续的 32 个 4 字节元素(比如 float arr[i*32:i*32+32]),那么硬件可以把它们合并成 1 次 128 字节的 transaction,效率拉满。

但如果这 32 个线程读的地址东一榔头西一棒槌(比如 arr[idx[tid]]idx 是随机的),那就是 32 次独立的 transaction,带宽利用率掉到 1/32。

  Coalesced(合并访存):
  Thread:  0  1  2  3 ... 30 31
           ↓  ↓  ↓  ↓     ↓  ↓
  Address: 0  4  8  12... 120 124
           └────┬────────────┘
              1 次 128B 事务

  Uncoalesced(分散访存):
  Thread:  0      1     2    ...  31
           ↓      ↓     ↓         ↓
  Address: 0    1024  2048 ...  31744
           └─────┬─────┴─────────┘
              32 次独立事务(每次只用 4B/32B)

这个效应在 LLM 算子里非常常见。比如反向 attention 中 dQ 的累加,如果 thread layout 没设计好,每次 atomic add 都是 uncoalesced 的,性能可以差 10 倍。第 4 章会专门讲访存模式,第 16 章会讲怎么让 dQ 累加 coalesced。

1.4 案例:同样的 reduce,CPU 与 GPU 的两种思路

把上面所有的概念落到一个具体例子上:对一个长度为 N 的 float 数组求和

1.4.1 CPU 的思路

最直观的 CPU 写法:

float sum = 0.0f;
for (int i = 0; i < N; ++i) {
    sum += arr[i];
}

这段代码完全串行——sum += arr[i] 是一条依赖链,前一个加法结果是后一个的输入。但现代 CPU 不会真的让它串行:

充分优化后的 CPU 代码:

#include <omp.h>
#include <immintrin.h>

float sum_cpu(const float* arr, int N) {
    float total = 0.0f;
    #pragma omp parallel for reduction(+:total)
    for (int i = 0; i < N; i += 16) {
        __m512 v = _mm512_loadu_ps(&arr[i]);
        total += _mm512_reduce_add_ps(v);
    }
    return total;
}

这段代码已经把 OMP 多线程 + AVX-512 SIMD 都用上了。在 64 核 Xeon 上,对 1 亿个 float 求和大约要 30-50 ms。瓶颈通常是内存带宽(这是一个 ~0.25 FLOPs/byte 的访存密集型 kernel)。

1.4.2 GPU 的思路

GPU 上同样的任务,思路完全不同。一个朴素的 GPU 写法:

__global__ void reduce_naive(const float* arr, float* result, int N) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid < N) {
        atomicAdd(result, arr[tid]);
    }
}

这段代码有 N 个线程,每个线程读一个元素,然后做 atomicAdd 把它累加到全局结果上。能跑,但性能极差——所有线程在抢一个全局变量的 atomic 锁,串行化严重。

正确的 GPU reduce 思路是分层归约:

template <int BLOCK_SIZE>
__global__ void reduce_v2(const float* arr, float* block_sums, int N) {
    __shared__ float smem[BLOCK_SIZE];
    int tid = threadIdx.x;
    int gid = blockIdx.x * blockDim.x + tid;

    // 1. 每个线程读一个元素到 SMEM
    smem[tid] = (gid < N) ? arr[gid] : 0.0f;
    __syncthreads();

    // 2. block 内分层归约: 每一步把活跃线程数减半
    for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            smem[tid] += smem[tid + s];
        }
        __syncthreads();
    }

    // 3. 每个 block 输出一个部分和
    if (tid == 0) {
        block_sums[blockIdx.x] = smem[0];
    }
}

这段代码体现了 GPU 编程的几个核心思维:

  1. 大量线程协作:N 个线程同时读 N 个元素到 SMEM,用 __syncthreads() 同步。
  2. 分层归约:每一步活跃线程数减半(BLOCK_SIZE → BLOCK_SIZE/2 → ... → 1),充分利用并行性。
  3. Shared Memory 是关键:所有归约都在 SMEM 里完成,避免反复读写 HBM。
  4. 多级 launch:block 内归约出 block_sums,再 launch 一次 kernel 把 block_sums 归约到最终结果。

这个版本在 H100 上对 1 亿个 float 求和大约要 0.5-1 ms——比 CPU 快 50-100 倍。但这远不是最优。第 5 章我们会用 warp shuffle、cluster reduce、vectorized load 等技巧把它再提升 3-5 倍。

1.4.3 思维差异的本质

把两种实现摆在一起,能看出根本的思维差异:

维度 CPU 思路 GPU 思路
并行粒度 几十个线程 × SIMD lane 几十万个线程
同步模型 隐式(cache coherence) 显式(__syncthreads
内存层级 自动缓存(L1/L2/L3) 程序员管理(SMEM)
优化重点 ILP、SIMD、cache 命中 线程数、warp 协作、coalesced 访存
心智模型 "几个工人努力干活" "几万个工人协作分工"

CPU 思路更接近"我作为程序员控制一个高效的工人去完成任务";GPU 思路更接近"我设计一套规则,让几万个工人按规则自组织地完成任务"。这是从 CPU 到 GPU 最重要的思维迁移

1.5 心智模型:把 GPU 想成一个高度并行的厨房

为了让这一章的所有概念能在读者脑子里粘在一起,我建议建立一个比喻。

把 GPU 想成一个超大型厨房

flowchart TB
  subgraph Restaurant [GPU 餐厅]
    direction TB
    Kitchen1[厨房 1<br/>SM]
    Kitchen2[厨房 2<br/>SM]
    KitchenN[... 132 个厨房]

    subgraph K1 [厨房 1 内部]
      Chef[4 个领班<br/>4 个 Warp Scheduler]
      Team1[团队 1<br/>32 个厨师 = 1 warp]
      Team2[团队 2<br/>32 个厨师 = 1 warp]
      TeamN[...最多 64 团队]
      Chef -->|每周期分配一道菜| Team1
      Chef -->|每周期分配一道菜| Team2
      Chef -->|...| TeamN
      Pantry[食品柜<br/>SMEM 256KB]
      Fridge[冰箱<br/>L2 50MB]
      Warehouse[远处的仓库<br/>HBM 80GB]
    end
  end

这个比喻里:

用这个比喻能把很多 GPU 编程的"反直觉"现象解释通顺:

这个比喻不完美(任何比喻都不完美),但它能在你脑子里搭一个粗糙的骨架。后续章节遇到具体概念时,可以一直回到这个比喻上对照。

1.6 这一章的小结与下一章

这一章我们建立了 4 个核心认知:

  1. GPU 不是更多核的 CPU:它是为了高吞吐而完全不同设计的处理器。CPU 优化延迟,GPU 优化吞吐。
  2. SIMT 是 GPU 的核心范式:32 个线程为一个 warp,硬件给一个 warp 发同一条指令。warp 切换是零开销的,这是延迟隐藏的关键。
  3. Divergence 和 coalesced 访存是 SIMT 的两个常见陷阱:写代码时不会报错,但性能会无声地腰斩甚至砍到 1/32。
  4. GPU 编程的心智迁移核心是从"控制一个工人"到"设计协作规则":你不再为单个线程优化,而是为大群线程的协作设计模式。

这些认知是后续所有章节的地基。第 2 章我们会带读者深入 Hopper 的微架构,把 SM 内部的 Tensor Core、TMA、L1+SMEM、寄存器文件这些组件逐一打开,建立一份精确到具体硬件单元的地图。读完第 2 章,读者就能拿着 NVIDIA 官方 Whitepaper 的任何一页,对应到自己脑子里的 SM 模型上。

本章动手练习

  1. 在一台 H100 / A100 上写一个最朴素的 reduce kernel(用 atomicAdd),用 nvprof / Nsight Compute 测一下实际带宽利用率。预测一下结果,再看实际数字差距。
  2. 把 Section 1.4.2 的 reduce_v2 跑起来,对比朴素版本的性能差距。
  3. 阅读 NVIDIA H100 Whitepaper 第 2 章 "Hopper Streaming Multiprocessor",识别每一段对应到本章的哪个概念。