CS336 Lecture 5-6: GPUs & Kernels, Triton

为什么这两讲至关重要

在已经学过的课程中，L1-L4 讲的是"模型长什么样"，L7-L8 讲的是"多卡怎么协作"。L5-L6 卡在两者之间——单卡内部到底在发生什么。

如果你不理解 GPU 的硬件结构，你就无法理解：

• 为什么 Flash Attention 是过去五年 LLM 领域最重要的系统优化
• 为什么大矩阵乘法是"compute-bound"而 Attention 是"memory-bound"
• 为什么你写的 PyTorch 代码有时 GPU 利用率只有 30%
• ZeRO-3 的通信开销到底卡在哪里（看懂 L7-8 那块没懂的）

Percy 在 L6 的原话很直白：“You cannot optimize what you don’t measure, and you cannot measure what you don’t understand.”

Lecture 5：GPU 架构解剖（Tatsu）

5.1 从一张 H100 说起

H100 的基本参数：

• HBM（显存）：80GB，带宽 ~3.35 TB/s
• SRAM（片上共享内存）：每 SM 约 256KB，总带宽 ~几十 TB/s（比 HBM 快一个数量级）
• 计算能力：BF16 下 ~990 TFLOPS

关键不在于这些数字有多大，而在于数字之间的比例。

5.2 内存层次：为什么"内存墙"是核心问题

GPU 的内存是分层的，越靠近计算单元越快、越小：

寄存器 (Registers)   ← 最快，每个线程私有，容量极小
    ↓
SRAM / Shared Memory ← 片上，同一 block 内线程共享，~256KB/SM
    ↓
L2 Cache             ← 片上，所有 SM 共享，~50MB
    ↓
HBM (显存/Global Mem) ← 片外，容量大但慢，80GB，带宽 ~3.35TB/s

核心矛盾：计算单元的速度远超数据搬运的速度。

• H100 计算能力：~990 TFLOPS（BF16）
• H100 内存带宽：~3.35 TB/s

H100 的算力/带宽比约为 300 次运算/Byte。这意味着，每读取 1 Byte 数据必须做 300 次运算才能跑满算力。反观一个极端情况：如果每个 FP16 数（2 Byte）只做 1 次运算，那每 Byte 仅 0.5 次，远低于 300 的门槛，GPU 利用率将极其低下。

实际中大部分操作远达不到这个比值，所以 GPU 大多数时间不是在计算，而是在等数据。这就是所谓的内存墙（Memory Wall）。

🔗 跟 L2 的知识串起来：L2 你学了 FLOPs 估算和混合精度训练。FP16 省一半内存，不仅省显存容量，更关键的是——省了一半的显存带宽。带宽才是真正的稀缺资源。

5.3 GPU 执行模型详解：Thread → Warp → Block → Grid

想象 GPU 是一家超级工厂，要处理一亿个零件（数据）。工厂的组织层级：

Grid（整个工厂的一次订单）
  └── Block 0（车间 0）── 固定在 SM 0 上
  │     ├── Warp 0（流水线 0）── 32 个工人同时拧螺丝
  │     ├── Warp 1（流水线 1）── 32 个工人同时拧螺丝
  │     └── ...（最多几十个 Warp）
  └── Block 1（车间 1）── 固定在 SM 1 上
  └── Block 2（车间 2）── 固定在 SM 2 上
  └── ...

核心原则：工人（Thread）很多，但工厂靠流水线（Warp） 来组织，靠车间（Block） 来分配资源。

第一层：Thread（线程）

什么是 Thread？ Thread 是 CUDA 编程模型里最小的逻辑执行单元。你写 CUDA/Triton 代码时，心里想的通常是"每个线程处理一个数据元素"。

// CUDA 伪代码
__global__ void add(float* a, float* b, float* c) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;  // 我是第几个线程？
    c[i] = a[i] + b[i];  // 我负责算第 i 个加法
}

但硬件上，Thread 不是"物理核心"。GPU 有成千上万个 CUDA Core，但线程数量远大于 Core 数量（比如 H100 有 14592 个 FP32 Core，但可以同时驻留数十万个线程）。

原因：GPU 是吞吐量处理器，不是延迟处理器。它不靠单个线程跑得快，而是靠同时塞入海量线程，让 Core 永远有活干。

第二层：Warp（线程束）—— 最重要的概念

32 个线程 = 1 个 Warp。Warp 是 GPU 硬件上真正的执行单位。32 个线程被"捆绑"在一起，必须同时执行同一条指令。

这就是 SIMT（Single Instruction, Multiple Threads）：

• CPU 的 SIMD：一条指令处理多个数据（向量寄存器）。
• GPU 的 SIMT：一条指令广播给 32 个线程，每个线程用自己的数据执行。

类比：32 人划龙舟

• 32 个桨手（线程）坐在一条船上（Warp）。
• 鼓手（Warp Scheduler）喊一声"划"，32 人同时划桨。
• 如果有人想"不划"（if 分支），他不能真的停下——他必须假装划（mask 掉），但船还在走。

⚠️ Volta 之后的 Independent Thread Scheduling：Pre-Volta 架构中 Warp 是调度的原子单位（整个 Warp 共用一个程序计数器）。Volta（2017）引入 Independent Thread Scheduling 后，Warp 内每个线程有独立的程序计数器和状态，调度粒度更细。但对初学者而言，Warp 仍是执行的基本分组——“基本调度单位"即可，“原子"一词过于绝对。

Warp Divergence（分支发散）：性能杀手

__global__ void bad_kernel(float* x) {
    int i = threadIdx.x;
    if (x[i] > 0) {
        x[i] = sqrt(x[i]);   // 路径 A
    } else {
        x[i] = x[i] * x[i];  // 路径 B
    }
}

硬件实际执行：

1. Warp Scheduler 先广播"执行路径 A"的指令，32 个线程中 x[i] > 0 的线程真正干活，其余的空闲等待（mask 掉）。
2. 然后广播"执行路径 B"的指令，刚才空闲的线程干活，刚才干活的线程空闲等待。
3. 两条路径串行执行完，才算完。

代价：如果 Warp 内 16 个走 A、16 个走 B，执行时间 ≈ 2 倍。

⚠️ 注意：Warp 内的分支发散只发生在同一个 Warp 内部。如果 Block 0 的 Warp 0 全走 A，Block 0 的 Warp 1 全走 B，这两个 Warp 是并行执行的，互不干扰。

第三层：Block（线程块）

Block 是资源分配的单位。一个 Block 内的所有 Warp 必须跑在**同一个 SM（Streaming Multiprocessor）**上。Block 内的线程可以：

• 共享 SRAM / Shared Memory（通常 ~256KB/SM，Block 内按需分配）
• 同步：__syncthreads()（Block 内所有线程到达此点才能继续）

一个 SM 能塞多少个 Block？不是 1 个，而是多个。 取决于三个资源限制：

实际例子：如果 Block 大小 = 256 线程（8 个 Warp），每个 Block 用 32KB Shared Memory：

• Warp 限制：64 / 8 = 8 个 Block
• Shared Memory 限制：256KB / 32KB = 8 个 Block
• 结果：8 个 Block/SM，共 2048 线程同时驻留

第四层：Grid（网格）

Grid 是一次 Kernel 调用的所有 Block 的集合。Block 之间默认不共享内存、不直接同步（除非用原子操作或 Cooperative Groups）。

GPU 调度器会把 Grid 里的 Block 分发到空闲的 SM上。如果你的 GPU 有 132 个 SM（H100），至少可以有 132 个 Block 同时执行（实际更多，因为 1 SM 可驻留多 Block）。

关键机制：Latency Hiding（延迟隐藏）

这是 GPU 高性能的核心秘密。

问题：从 HBM 读数据需要 几百个时钟周期。如果 Warp 0 的 32 个线程都在等数据，那这 32 个 CUDA Core 就空闲了？

不。Warp Scheduler 会立刻切换到 Warp 1，让 Warp 1 的线程用 Core 计算。Warp 1 也等数据了？切换到 Warp 2……直到所有 Warp 都在等数据，SM 才真正空闲。

类比：医院挂号窗口

• 窗口（CUDA Core）只有 1 个。
• 病人（Warp）有 32 个。
• 病人 A 去拍 X 光（访存），需要等 30 分钟。
• 窗口不叫 A 傻等，而是立刻叫 B 来看病。
• 等 32 个病人全去拍片了，窗口才空闲。

结论：SM 上驻留的 Warp 越多，越能隐藏内存延迟。这也是为什么 Occupancy（占用率）是重要的性能指标。

代码到硬件的完整映射

# PyTorch 代码
y = torch.relu(x)  # 假设 x 有 1,048,576 个元素

底层发生了什么：

1. PyTorch 调用 CUDA Kernel relu。
2. Grid 被划分：比如每个 Block 用 256 线程，需要 4096 个 Block。
3. 每个 Block 被分配到某个 SM。
4. 每个 Block 的 256 线程被分成 8 个 Warp（256/32=8）。
5. Warp Scheduler 在 SM 上调度这 8 个 Warp，可能和其他 Block 的 Warp 一起竞争 Core。
6. 每个 Warp 的 32 线程同时执行 max(0, x[i])。
7. 如果某个 Warp 在等 x[i] 从 HBM 加载，Warp Scheduler 切到另一个 Warp。

常见误区澄清

Thread 是你编程时的视角，Warp 是硬件执行的视角，Block 是资源分配的视角，Grid 是任务调度的视角。 GPU 的高性能来自于：用海量 Warp 填满 SM，让 Warp Scheduler 在访存延迟期间不断切换，保证计算单元永远有活干。

5.4 Tensor Core：Transformer 的真正引擎

先问一个问题：深度学习到底在算什么？

Transformer、CNN、RNN，底层最核心的操作都是矩阵乘法：

• Attention 里的 $QK^T$、$PV$
• FFN 里的 $XW_1$、$XW_2$
• 卷积展开后也是矩阵乘

一个 $4096 \times 4096$ 的矩阵乘法，包含几百亿次乘加运算。如果用最原始的方式做，非常慢。

CUDA Core vs Tensor Core：手算 vs 专用计算器

CUDA Core（标量运算）是 GPU 里通用的浮点运算单元，一次只做一个标量的运算：

# CUDA Core 的视角
c = a * b + d   # 1 次乘加（FMA），处理 1 个数

就像用手算做矩阵乘法：算 $C_{ij} = \sum_k A_{ik} \times B_{kj}$，一次算一个元素。

Tensor Core（矩阵运算）是 NVIDIA 从 Volta 架构（V100）开始引入的专用电路，一次指令完成一个小矩阵的乘加：

这不是手算，而是一个专用计算器——按一下按钮，直接吐出 $4\times4$ 的结果。

具体算一下：一个 $4\times4 \times 4\times4$ 矩阵乘，结果有 16 个元素，每个元素需要 4 次乘加（对 $k$ 求和），总共 $16 \times 4 = 64$ 次乘加。

Tensor Core 一个指令周期 = 64 次乘加。CUDA Core 一个指令周期 = 1 次乘加。

这还没完——现代 Tensor Core（A100/H100）实际处理的矩阵块比 $4\times4$ 更大（如 $8\times8\times4$ 或 $16\times16\times16$），一个指令周期等于几千次乘加。

吞吐量对比

作为对比，用 CUDA Core 做 FP32：

• A100 CUDA Core FP32 峰值约 19.5 TFLOPS
• Tensor Core FP16 是 CUDA Core FP32 的 16 倍

这意味着什么？ 如果你用 PyTorch 默认的 float32 跑大矩阵乘法，PyTorch 调用 CUDA Core，只能吃到 ~20 TFLOPS。用 FP16/BF16，PyTorch 自动调用 cuBLAS 走 Tensor Core，吃到 312 TFLOPS（A100）。差了一个数量级。

为什么 FP8 比 FP16 更快？不只是省内存

很多人以为 FP8 训练快是因为"显存占用减半，带宽省了”。这只是一半原因。

更关键的原因是 Tensor Core 本身：

• FP8 数据位宽是 FP16 的一半
• 同样的晶体管面积、同样的能耗，Tensor Core 在 FP8 模式下可以塞入两倍数量的运算单元，或每个周期处理两倍的数据通路
• 所以 H100 FP8 的 Tensor Core 吞吐量（~1980 TFLOPS）是 FP16（~989 TFLOPS）的 2 倍

这是硬件层面的翻倍，不是软件层面的优化。

混合精度训练到底快在哪？

结合 5.2 节的内存墙和这里的 Tensor Core：

混合精度训练快的两层原因：

1. 内存侧：显存占用减半，HBM 读写量减半，缓解内存墙。
2. 计算侧：Tensor Core 的 FP16 吞吐量是 CUDA Core FP32 的 16 倍——这是更主要的加速来源。

FP8 训练（H100）：

1. 内存侧再减半。
2. 计算侧再翻倍（相对 FP16）。
3. 所以 FP8 不只是"省显存”，而是让 Tensor Core 本身跑得更快。

5.5 GPU 性能的"波浪形"波动从哪来

你用 torch.matmul 做不同尺寸的矩阵乘法时会发现：并不是尺寸越大效率越高，而是在某些"甜点尺寸"上效率突然飙升，在其他尺寸上掉下来。

原因三个层面：

① Tiling 与 SM 占用：边缘浪费

背景：GPU 不会一口气算完整个大矩阵。它把矩阵切成小块（Tile），每个 SM 领一块去算。H100 上 cuBLAS 的 tile size 通常是 128×128 或 256×256。

问题：除不尽怎么办？ 假设你要算 $C = A \times B$，矩阵尺寸是 300×300：

• Tile size = 128
• 300 ÷ 128 = 2.34… 除不尽
• GPU 只能分配 3×3 = 9 个 tile 来覆盖（2 个 tile 只够 256，不够 300）

实际需要：300×300
Tile 覆盖：384×384（3 × 128 = 384）

┌────────┬────────┬────────┐
│  128   │  128   │  128   │  ← 最右边一列（44列）是浪费的
├────────┼────────┼────────┤
│  128   │  128   │  128   │  ← 最下面一行（44行）是浪费的
├────────┼────────┼────────┤
│  128   │  128   │  128   │  ← 右下角（44×44）完全浪费
└────────┴────────┴────────┘

实际有效计算占比：

近 40% 的算力在算空气。

对比"甜点尺寸"：如果矩阵是 256×256（刚好 2×128），没有边缘浪费，利用率接近 100%，性能突然飙升。

这就是"波浪"的第一个来源：尺寸刚好对齐 tile size 时效率最高，差一点就要多开一整圈 tile，大量算力浪费在边缘。

② Bank Conflict：共享内存冲突

SRAM 的物理结构：SM 上的 Shared Memory 被分成 32 个 bank（可以想象成 32 个抽屉）。规则：每个时钟周期，每个 bank 只能服务一个地址。

什么时候会冲突？ 一个 Warp 有 32 个线程。理想情况是 32 个线程各访问不同的 bank，32 个抽屉同时打开，一次搞定。

噩梦情况：32 个线程访问了同一个 bank 的不同地址。

线程 0:  读 bank 0, 地址 0
线程 1:  读 bank 0, 地址 1   ← 冲突！bank 0 一次只能服务一个
线程 2:  读 bank 0, 地址 2   ← 排队
...
线程 31: 读 bank 0, 地址 31  ← 排到第 32 个周期

结果：本来 1 个周期能完成的读取，变成 32 个周期串行执行。性能直接打骨折。

为什么跟矩阵尺寸有关？ 不同的矩阵尺寸导致不同的数据排布（stride）。某些尺寸下线程访问的地址恰好均匀分布在 32 个 bank 上（无冲突），某些尺寸下恰好撞在同一个 bank 上（严重冲突）。

注：这是 Triton 编译器帮你自动优化的部分，但底层硬件仍然受这个规律支配。

③ Occupancy（占用率）：SRAM 与并行度的零和博弈

回顾 5.3：一个 SM 能同时驻留的 Warp 数有上限。Warp 越多，越能隐藏内存延迟（Latency Hiding）。

SRAM 的零和博弈：每个 Block 可以申请一定量的 Shared Memory。如果一个 Block 申请太多：

• SM 的 Shared Memory 总量固定（如 256KB）
• 1 个 Block 占了 200KB → 这个 SM 只能同时放 1 个 Block
• 1 个 Block 可能只有 4 个 Warp → SM 上只有 4 个 Warp 在跑
• Warp 太少 → 内存延迟藏不住 → 计算单元空转

为什么跟矩阵尺寸有关？ 大矩阵乘法通常需要更大的 tile，意味着每个 Block 需要更多 SRAM 来缓存 A/B 的子块：

在"SRAM 使用量"和"并行 Block 数"之间有个最优 trade-off。

Lecture 6：Kernels & Triton（Percy）

根据课程源代码 lecture_06.py，L6 的实际内容按顺序分为六大块：

1. Benchmarking & Profiling
2. Kernel Fusion（仓库/工厂类比）
3. CUDA 手写 Kernel（以 GeLU 为例，看 gelu.cu）
4. Triton 编程（GeLU → Softmax → Matmul 难度递增）
5. PyTorch Compilation（torch.compile 作为"不自己写 kernel"的选项）
6. PTX Assembly（看 Triton 编译出了什么）

⚠️ 注意：Flash Attention 不是 L6 的主体内容，它是 Assignment 2 的作业。L6 只教到 Triton 编程基础（elementwise → row-wise reduction → tiled matmul），FA 需要你自己在作业里组合这些技能。

6.1 Benchmarking & Profiling：不要猜，要测

Percy 在 L6 的第一个大 section 就是 profiling，不是 GPU 架构。核心信息：

“Spec sheets are marketing material. There is no substitute for benchmarking your own workload on your own hardware.”

L6 展示了完整的 profiling 工作流：

1. Benchmarking（测总时间）：对 MLP 分别改变 batch_size、num_layers、dim，看性能如何 scaling——结果往往不符合你的直觉，因为 CUDA kernel 行为是非均匀的
2. Profiling（看哪里花了时间）：用 torch.profiler 看具体调用了哪些 CUDA kernel

with torch.profiler.profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]
) as prof:
    run()
table = prof.key_averages().table(sort_by="cuda_time_total", row_limit=20)

学会读 kernel 名字：比如 cutlass_80_simt_sgemm_256x128_8x4_nn_align1

• cutlass = NVIDIA 的线性代数 CUDA 库
• 256x128 = tile size
• sgemm = single precision GEMM
• nn = no transpose
• 名字本身就告诉你了底层怎么 tiling 的

6.2 Kernel Fusion：为什么"合并操作"能省时间

这是 CS336 在 L6 引用的核心直觉，来自 Horace He 的博客：

DRAM 是仓库，SRAM 是工厂。

仓库（DRAM/HBM）：什么都存，但运到工厂很慢
工厂的桌面（SRAM）：空间很小，但拿取极快

假设你要做 y = tanh(0.5 * x * (1 + 0.5 * x))——三个操作：

Naive（不融合）：

读 x 从 DRAM → 算 0.5*x → 写回 DRAM
读 0.5*x 从 DRAM → 算 1+0.5*x → 写回 DRAM
读两次从 DRAM → 算 tanh → 写回 DRAM

6 次 DRAM 访问

Fused（融合）：

读 x 从 DRAM → 算全部 → 写回 DRAM

2 次 DRAM 访问

核心原则：尽可能让数据留在 SRAM 上，把所有能一起算的都算完，再写回去。

L6 用 GeLU 函数做了实际 benchmark：

结论：fused kernel 的核心价值就是减少 DRAM/HBM 访问次数。这跟之前说的"Attention 是 memory-bound"是同一套逻辑。

6.3 什么是 Kernel

Kernel 就是在 GPU 上执行的函数。你在 CPU 上调用它，它在 GPU 上跑。

# PyTorch 的一个操作背后就是一个（或多个）kernel
c = torch.matmul(a, b)  # 调用 cuBLAS 的矩阵乘法 kernel
c = F.softmax(x, dim=-1)  # 调用 softmax kernel

CUDA 的编程模型是 thread-level 的：你写代码时想的是"每个线程做什么"。这对简单操作还行，对复杂操作（如 Attention）就非常痛苦——你需要手动管理共享内存分配、bank conflict、coalesced memory access、寄存器压力……

Triton 的编程模型是 block-level 的：你写代码时想的是"每个 block 做什么"，Triton 编译器帮你处理 block 内部的线程调度和内存优化。

CUDA:  你告诉每个线程做什么 → 手动管理共享内存、同步、bank conflict
Triton: 你告诉每个 block 做什么 → 编译器自动优化线程级细节

6.4 CUDA 手写 Kernel：以 GeLU 为例

L6 用 load_inline 在 Python 中嵌入了一段 CUDA 代码（来自课程 gelu.cu）：

#include <math.h>
#include <torch/extension.h>

__global__ void gelu_kernel(float* in, float* out, int num_elements) {
    // 每个线程负责一个元素
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < num_elements) {
        out[i] = 0.5 * in[i] * (1.0 + tanh(
            0.79788456 * (in[i] + 0.044715 * in[i] * in[i] * in[i])));
    }
}

torch::Tensor gelu(torch::Tensor x) {
    torch::Tensor y = torch::empty_like(x);
    int num_elements = x.numel();
    int block_size = 1024;
    int num_blocks = (num_elements + block_size - 1) / block_size;
    gelu_kernel<<<num_blocks, block_size>>>(x.data_ptr<float>(), y.data_ptr<float>(), num_elements);
    return y;
}

这里演示了 CUDA 编程的标准模式：

1. Grid → Block → Thread：<<<num_blocks, block_size>>>
2. 每个线程通过 blockIdx.x * blockDim.x + threadIdx.x 知道自己是谁
3. 边界检查 if (i < num_elements)

所以 CUDA 编程的本质就是：你把 N 个元素的问题并行化到 N 个线程，每个线程处理 i，然后 GPU 硬件把这些线程分组到 SM 上自动调度。 对于 element-wise 操作这很简单，但对于需要线程间通信的操作（softmax、attention），你就得亲自动手管理 shared memory 了。

6.5 Triton 编程：从简单到复杂的三个层级

层级一：Elementwise（GeLU）— 最简单

L6 的实际 Triton GeLU 实现 —— 每个 block 处理连续的一批元素：

@triton.jit
def triton_gelu_kernel(x_ptr, y_ptr, num_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < num_elements                      # 边界处理
    
    x = tl.load(x_ptr + offsets, mask=mask)            # 从 HBM 加载
    
    # GeLU 计算
    a = 0.79788456 * (x + 0.044715 * x * x * x)
    exp = tl.exp(2 * a)
    tanh = (exp - 1) / (exp + 1)
    y = 0.5 * x * (1 + tanh)
    
    tl.store(y_ptr + offsets, y, mask=mask)            # 写回 HBM

对比 CUDA 版本：

• CUDA：blockIdx.x * blockDim.x + threadIdx.x 手动算索引
• Triton：tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)，语义相同但更清晰
• Triton 自动处理了 memory coalescing，CUDA 版本也要手动（但这里因为连续访问所以自动就 coalesced）

层级二：Row-wise Reduction（Softmax）

L6 的 Triton softmax —— 每个 block 负责一行，行内做 max/sum reduction：

@triton.jit
def triton_softmax_kernel(x_ptr, y_ptr, x_row_stride, y_row_stride,
                          num_cols, BLOCK_SIZE: tl.constexpr):
    row_idx = tl.program_id(0)                         # 每个 block 处理一行
    col_offsets = tl.arange(0, BLOCK_SIZE)
    
    # 读这一行
    x_start_ptr = x_ptr + row_idx * x_row_stride
    x_row = tl.load(x_start_ptr + col_offsets,
                    mask=col_offsets < num_cols, other=float("-inf"))
    
    # softmax = exp(x - max) / sum(exp(x - max))
    x_row = x_row - tl.max(x_row, axis=0)              # 行内 max，一行干完
    numerator = tl.exp(x_row)
    denominator = tl.sum(numerator, axis=0)             # 行内 sum
    y_row = numerator / denominator
    
    # 写回
    y_start_ptr = y_ptr + row_idx * y_row_stride
    tl.store(y_start_ptr + col_offsets, y_row, mask=col_offsets < num_cols)

关键差异：Softmax 比 GeLU 复杂，因为每个输出元素依赖同一行的所有输入元素。这里 Triton 的 tl.max(..., axis=0) 和 tl.sum(..., axis=0) 自动处理了行内规约。

层级三：Tiled Matmul — 最难

矩阵乘法的 Triton 实现是课程 matmul tutorial 的内容，核心在于：

1. 把 C = A @ B 沿 K 维度分块
2. 每次加载 A 的一块和 B 的一块到 SRAM
3. 在 SRAM 里做 mini-matrix multiplication 累加
4. 用 tl.dot(a, b) 调用 Tensor Core

那完整的 Triton matmul 代码我上一版写的差不多了，这里不重复。

6.6 三种 Triton 用例的对比

理解这三种模式很重要，因为 Flash Attention 就是 “Tiled Matmul” + “Online Softmax” 的组合。

6.7 torch.compile：第四种选择

L6 还介绍了 torch.compile，它本质上自动做了你在 6.2 学的事：把多个 PyTorch 操作融合成一个 kernel。

compiled_gelu = torch.compile(manual_gelu)   # 把你手写的 Python 函数自动编译

它的性能通常介于手写 CUDA/Triton 和纯 PyTorch 之间，好处是完全不需要写底层代码。

6.8 PTX Assembly：看你的代码真正变成了什么

PTX（Parallel Thread Execution）是 CUDA 的"汇编语言"。Triton 编译后会生成 PTX。

L6 演示了打印 PTX 来分析 Trito n 实际做了什么（这是 debug 性能问题的终极手段）：

list(kernel.cache[0].values())[0].asm["ptx"]

关键观察：

• ld.global.* / st.global.* = 读写 global memory（HBM）
• %ctaid.x = block index, %tid.x = thread index
• %f* = floating point registers, %r* = integer registers
• Triton 一个线程处理 8 个元素（thread coarsening）——这是编译器自动做的

6.9 Bonus：Flash Attention 在 CS336 里怎么学

根据课程实际结构：

• L6 不直接教 Flash Attention。L6 教你写 Triton（elementwise → softmax → matmul），为做 A2 做准备
• Assignment 2 才是你真正实现 Flash Attention 的地方，用 Triton 写一个简化版（或部分实现）的 FA kernel
• A2 还包含 Profiling（分析 MLP 各层瓶颈）和 DDP（多卡分布式训练）

理解 kernel fusion 的 warehouse/factory 类比，理解 Triton 的三种通信模式，你就可以去读 Flash Attention 论文了。本质上 FA 就是把 Attention 的 QK^T → softmax → PV 融合进一个 tiled kernel，用 online softmax 避免存储 O(N²) 中间矩阵——正是你在 L6 学到的"tiled matmul + row-wise softmax"的组合技巧。

Roofline 模型：统一理解 GPU 性能

Roofline 模型用一张图回答"我的操作受什么限制"：

纵轴：可达到的 FLOPs/s（对数尺度）
横轴：算术强度 = FLOPs / Bytes（对数尺度）

         ┌──────────────────────────────┐
         │   Memory-Bound     │Compute-Bound│
         │      (斜坡)        │  (平台)    │
 峰值 ── ┼─────────────────────┤────────── │  ← GPU 理论算力上限
         │     ╱              │             │
         │    ╱               │             │
         │   ╱                │             │
         │  ╱                 │             │
         │ ╱                  │             │
         │╱                   │             │
         └────────────────────┴──────────────┘
                   算术强度 →

• Memory-bound 区域（左侧斜坡）：你的操作在等待数据，加带宽就能加速
• Compute-bound 区域（右侧平台）：你的操作在等待计算，加带宽无用

具体到 Transformer 的各个操作：

这解释了 Flash Attention 为什么有效：Attention 是 memory-bound 的，瓶颈是 HBM 带宽而不是计算。Flash Attention 消除了 O(N²) 的 HBM 读写，相当于拿掉了最大的带宽消耗源。

这也解释了为什么 FFN 越来越大：FFN 的矩阵乘是 compute-bound，扩大 FFN 能更充分利用 GPU 的算力。Attention 是 memory-bound，扩大序列长度会让瓶颈更严重。

CS336 Assignment 2 要做什么

根据课程公开信息，A2 包含三部分：

1. Profiling：用 PyTorch Profiler 分析一个简单 Transformer 的 GPU 利用率，找出哪些操作在等数据、哪些在等计算
2. Triton 实现 Flash Attention：手写一个简化但正确的 Flash Attention kernel，验证对 online softmax 的理解
3. DDP 多卡训练：在 2-4 张卡上跑通分布式数据并行，观察通信开销

本节小结

学完这两讲再回头看 L7-L8 的 ZeRO-3：ZeRO-3 之所以"通信等待变久了"，本质是因为参数分片后每次前向都要 All-Gather 参数，All-Gather 是 memory-bound 的通信操作——你从"算力瓶颈"变成了"带宽瓶颈"。理解了 L5-L6 的内存/带宽分析框架，ZeRO 各阶段的 trade-off 就一目了然了。

参考资料

• CS336 GitHub
• Triton 官方文档 — Matrix Multiplication Tutorial
• Flash Attention 论文 (Dao et al., 2022)
• Flash Attention 2 论文 (Dao, 2023)
• Roofline Model 原论文 (Williams et al., 2009)
• CUDA C++ Programming Guide — Memory Hierarchy