# Lecture 6 · Kernel 与 Triton

> **CS336: Language Modeling from Scratch** · Stanford · Spring 2025
> 📅 Apr 17 · 💻 [可执行讲义](https://github.com/stanford-cs336/spring2025-lectures/blob/main/lecture_06.py) · 讲者 Tatsunori Hashimoto

---

## 承上启下

**上一讲（[[Lecture 5 · GPU]]）**：
- **硬件解剖**：SM（流多处理器）、SIMT 执行模型与 warp 锁步、内存层级（HBM → L2 → SRAM/共享内存 → 寄存器）
- **性能瓶颈**：算力 vs 带宽，roofline 模型 $P = \min(\pi, \beta \cdot \text{AI})$，算术强度（FLOPs/byte）决定算子是 compute-bound 还是 memory-bound
- **优化策略**：低精度、tiling（利用共享内存）、内存合并、算子融合、重计算
- **FlashAttention**：online softmax + tiling，将注意力的额外内存从 $O(N^2)$ 降到 $O(N)$

**本讲（Kernel 与 Triton）**：
上一讲从硬件视角回答「**应该**怎么优化」，本讲从软件视角回答「怎么**真正动手**写出快的代码」。关键挑战在于：PyTorch 的高层抽象（一行 `x + y`）与硬件之间隔着巨大的性能黑箱——同样的数学计算，融合与否可能相差 **5 倍以上**。本讲的方法论（benchmark → profile → 写 kernel）也是 [[Lecture 2 · PyTorch 与资源核算]] 里「纸面核算 FLOPs」的实证闭环：算出来的上限，要靠测量和 kernel 工程去逼近。

> [!note] 一句话总览
> **Benchmarking 找瓶颈 → profiling 看内部 → 写 kernel 融合算子**。同一个函数有五条实现路径：手写 PyTorch 公式（慢，拆成多个 kernel）→ `torch.compile`（自动融合）→ CUDA C++（快但难）→ **Triton**（Python 般简单、接近 CUDA 般快）→ PyTorch 内置算子（工程优化到极致）。会选择抽象层级，比会写某一种 kernel 更重要。

> [!note] 本讲实验环境
> 讲义在 **A100** 上运行：108 个 SM、80 GB HBM、40 MB L2 缓存、每 SM 192 KB L1/共享内存（上一讲的 H100 则是 132 个 SM、50 MB L2、每 SM 228 KB）。具体数字不同，但内存金字塔的结构和一切优化原理完全一致。

---

## 1. Benchmarking 与 Profiling：性能优化的起点

### 1.1 为什么需要 Benchmarking？

**Benchmarking（基准测试）**：测量某段代码的**端到端墙上时钟时间（wall-clock time）**。

它回答的是：
- 这个实现到底快不快？
- 换一个 batch size / hidden dim / 层数，速度怎么变？
- 优化前后是否真的变快？

> [!warning] 易错点：不要靠感觉、也不要靠 spec sheet 优化性能
> GPU 性能高度依赖硬件、驱动、CUDA 版本、PyTorch 版本、输入形状和 dtype。论文、规格表、博客给的数字都只是参考，**最终必须在自己的环境上跑 benchmark**。这也是讲义把 benchmarking 放在所有 kernel 内容之前的原因：没有测量，一切优化都是盲人摸象。

**讲义中的示例对象**：一个简单 MLP（也是作业里资源核算的老朋友）。

```python
class MLP(nn.Module):
    def __init__(self, dim: int, num_layers: int):
        super().__init__()
        self.layers = nn.ModuleList([
            nn.Linear(dim, dim) for _ in range(num_layers)
        ])

    def forward(self, x: torch.Tensor):
        for layer in self.layers:
            x = torch.nn.functional.gelu(layer(x))
        return x
```

Benchmark 的基本流程（讲义版取多次的平均值；取 `min` 也常见——最小值最接近「无干扰时的真实成本」）：

```python
def benchmark(description: str, run: Callable, num_warmups: int = 1, num_trials: int = 3):
    # 1. warmup：让 CUDA kernel 编译、缓存、allocator 都进入稳定状态
    for _ in range(num_warmups):
        run()
    torch.cuda.synchronize()

    # 2. 正式计时：CUDA 默认异步，所以每次 run 之后都要 synchronize
    times = []
    for _ in range(num_trials):
        start = time.time()
        run()
        torch.cuda.synchronize()
        times.append((time.time() - start) * 1000)   # 毫秒

    print(description, f"{sum(times) / len(times):.2f} ms")
```

**为什么要 warmup？** 第一次调用会触发很多一次性开销，混进计时里会严重高估：
- kernel 首次加载、JIT 编译（`torch.compile`/Triton 的编译开销都发生在第一次调用）
- PyTorch CUDA caching allocator 首次向驱动申请大块显存（之后从池子里拿就快了）
- GPU 有动态频率调节，冷启动时可能尚未升频

**为什么要 `torch.cuda.synchronize()`？**
- CUDA kernel launch 是**异步**的：CPU 把任务提交进队列后立即返回，GPU 还在后台执行。
- 不同步的话，计到的可能只是「提交任务」的微秒级时间，而不是 GPU 真正的执行时间——这是 GPU 计时最经典的错误。

### 1.2 Benchmark 看什么？

讲义依次缩放了几个维度（矩阵乘法 dim 从 1024 扫到 16384；MLP 用 dim=256/2048、layers=4~64、batch=128~1024、steps=2~10）：

| 维度 | 观察点 | 含义 |
| :--- | :--- | :--- |
| `num_steps` | 时间随步数线性增长 | 基本验证计时代码没错（sanity check）|
| `num_layers` | 层数越多越慢 | 每层 linear + GeLU 都增加 kernel 数量 |
| `batch_size` | batch 变大不一定成比例变慢 | 大 batch 提高矩阵乘法的算术强度与 SM 占用 |
| `dim` | hidden dim 变化时耗时会跳变 | 不同形状会派发到不同的底层 CUDA kernel |

> [!tip] 直觉：性能不是形状的连续函数
> `dim=127`、`dim=128`、`dim=129` 的耗时可能完全不同，因为底层 kernel 的 tile 大小、内存对齐、能否走 Tensor Core 路径都随形状离散地跳变。这就是为什么工程实践中模型维度总取 64/128 的倍数——不是玄学，是为了踩中高速 kernel 的甜点位（参见 [[Lecture 3 · 架构与超参数]] 里的维度选择惯例）。

**Benchmark 的局限**：
- 它只告诉你「慢」，不告诉你「为什么慢」；
- 它看不到一行 PyTorch 背后调用了哪些 kernel；
- 它无法区分瓶颈是 CPU launch overhead、GPU 计算、HBM 读写还是同步点。

这就需要 profiling。

### 1.3 Profiling：打开 PyTorch 黑箱

**Profiling（性能剖析）**：把一次运行拆开，记录每个操作、每个 CUDA kernel 的 CPU/GPU 时间、调用次数和调用栈。

```python
from torch.profiler import profile, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    with_stack=True,
) as prof:
    output = model(input)
    loss = output.sum()
    loss.backward()

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
```

Profiler 能回答：
- 一个 PyTorch op 对应几个 CUDA kernel？
- 时间主要花在 `aten::mm`、`aten::add`、`aten::gelu` 还是 `cudaLaunchKernel`？
- 小矩阵时是不是 launch overhead 比计算本身还大？
- MLP 训练里前向、反向、优化器分别占多少？

**讲义逐个 profile 的操作及观察**（先用 50 ms 的 `sleep` 校准 profiler 本身）：

| 操作 | Profiler 暴露出的事实 | 启发 |
| :--- | :--- | :--- |
| `x + y` | 一个 elementwise kernel | 简单但 memory-bound |
| `x @ w`（dim=2048） | 调用 `cutlass_80_simt_sgemm_256x128_8x4_nn_align1` 之类的 GEMM kernel | kernel 名里就写着 tile 形状（256×128）；GEMM 已被高度优化 |
| `x @ w`（dim=128） | 派发到**另一个** GEMM kernel | 同一个 op、不同形状 → 不同 kernel，印证「性能不连续」 |
| `torch.cdist` | 分解成 `matmul`、`pow`、`sum`、`cat` 等一串 kernel | 高层 API 可能隐藏一整串 kernel |
| 手写 `gelu` 公式 | `mul`/`add`/`tanh`/`pow` 拆成多个 kernel | 每个中间结果都要往返 HBM |
| PyTorch `gelu` | 单个 fused kernel | 融合能显著省内存带宽 |

> [!tip] 直觉
> Benchmark 是体温计（「发烧了」），profiler 是 X 光（「炎症在哪」）。性能优化不是玄学，是**证据链**：先确认慢、再定位慢在哪、然后才谈优化手段——顺序反了就会把时间花在不痛不痒的地方。

---

## 2. Kernel Fusion：把工厂搬到仓库旁边

### 2.1 为什么多个小算子会慢？

回忆 [[Lecture 5 · GPU]] 的内存层级：

```mermaid
graph LR
    HBM["HBM / DRAM<br/>大但慢"] --> SRAM["SRAM / Shared Memory<br/>小但快"]
    SRAM --> REG["Register<br/>最快"]
```

逐元素操作本身 FLOPs 很少，真正贵的是数据搬运：
- 每个 kernel 从 HBM 读入完整输入；
- 把中间结果写回 HBM；
- 下一个 kernel 再把同一批数据从 HBM 读回来；
- 每次 kernel launch 还有固定的 CPU 端开销。

以手写 GeLU 为例。GeLU 的精确定义是 $\operatorname{GELU}(x) = x \cdot \Phi(x)$（$\Phi$ 为标准正态的累积分布函数），常用 tanh 近似：

$$\operatorname{GELU}(x) \approx \frac{x}{2}\left(1 + \tanh\!\left(\sqrt{\tfrac{2}{\pi}}\left(x + 0.044715\,x^{3}\right)\right)\right)$$

其中 $\sqrt{2/\pi} \approx 0.79788456$——后面 CUDA/Triton 代码里的魔法常数就是它。

```python
def manual_gelu(x):
    return 0.5 * x * (1.0 + torch.tanh(
        math.sqrt(2.0 / math.pi) * (x + 0.044715 * x ** 3)
    ))
```

这看起来是一行 Python，但 eager 模式的 PyTorch 会把它拆成 `pow`、`mul`、`add`、`tanh`、再 `mul` 等**七八个 kernel**，每一步都可能读写整个张量——内存流量放大近一个数量级，而这是个纯 memory-bound 的算子。

**Fusion（算子融合）**：把这串操作合成一个 kernel，中间值只存在于寄存器/共享内存：

```python
# 未融合：每行都是一次完整的 HBM 读写
tmp1 = x * x
tmp2 = tmp1 * x
tmp3 = x + 0.044715 * tmp2
tmp4 = tanh(...)
y = 0.5 * x * (1 + tmp4)

# 融合后：一个 kernel 内完成，读一次 x、写一次 y
y = fused_gelu(x)
```

| 实现 | Kernel 数 | HBM 往返 | 讲义观察 |
| :--- | :---: | :---: | :--- |
| 手写公式 | 多个 | 多次 | 最慢 |
| `torch.nn.functional.gelu` | 1 个 | 1 读 + 1 写 | 很快 |
| `torch.compile(manual_gelu)` | 1 个 Triton fused kernel | 1 读 + 1 写 | 接近 PyTorch 内置 |
| 手写 CUDA | 1 个 | 1 读 + 1 写 | 快，但开发成本高 |
| 手写 Triton | 1 个 | 1 读 + 1 写 | Python 写法，接近 CUDA |

> [!tip] 直觉
> Fusion 的本质**不是少做计算，而是少搬数据**——把「工厂」（计算）搬到「仓库」（数据所在的寄存器）旁边，而不是每道工序都把半成品运回总仓再取出来。对 memory-bound 算子，砍掉一次 HBM 往返比省几个 FLOPs 值钱得多。上一讲的 roofline 语言：融合让算子的算术强度乘上融合的算子个数。

### 2.2 Kernel Launch Overhead

每次 kernel launch 都有固定的 CPU 端开销（约 5-10 μs 量级）。对一个跑几毫秒的大 GEMM 这可以忽略；但对几微秒就算完的小 elementwise op，**launch 本身可能比计算还贵**——GPU 大部分时间在等 CPU 派活。

```mermaid
sequenceDiagram
    participant CPU
    participant GPU
    CPU->>GPU: launch kernel 1
    GPU-->>CPU: 异步返回
    CPU->>GPU: launch kernel 2
    GPU-->>CPU: 异步返回
    CPU->>GPU: launch kernel 3
    GPU-->>CPU: 异步返回
```

融合后：

```mermaid
sequenceDiagram
    participant CPU
    participant GPU
    CPU->>GPU: launch fused kernel
    GPU-->>CPU: 异步返回
```

> [!warning] 易错点
> 小模型、小 batch、短序列时 launch overhead 占比最高（profiler 里表现为 CPU 时间远大于 CUDA 时间）；大模型训练时 HBM 带宽和跨卡通信通常更关键。除了融合，**CUDA Graphs**（把一整串 launch 录制后一次性重放，`torch.compile` 的 reduce-overhead 模式会用）也是消除 launch 开销的利器。

---

## 3. CUDA Kernel：最底层的手工控制

### 3.1 CUDA 的执行模型

CUDA 的核心思想：**你写「单个线程要做什么」，GPU 帮你并行启动海量线程**。为了把海量线程组织起来，CUDA 提供了两级层次：

```mermaid
graph TB
    GRID["Grid<br/>一组 thread blocks"] --> B0["Block 0"]
    GRID --> B1["Block 1"]
    GRID --> B2["Block 2"]
    B0 --> T00["thread 0"]
    B0 --> T01["thread 1"]
    B0 --> T02["..."]
```

| 概念 | 含义 | 直觉 |
| :--- | :--- | :--- |
| **Thread** | 最小执行单元 | 处理一个或几个元素 |
| **Thread block** | 一组 thread，整体调度到同一个 SM | 块内可共享 shared memory、可用 `__syncthreads()` 同步 |
| **Grid** | 一次 kernel launch 的所有 block | 总工作量 |
| `blockIdx` | 当前 block 编号 | 定位粗粒度任务 |
| `threadIdx` | 当前 block 内线程编号 | 定位细粒度元素 |
| `blockDim` | 每个 block 的线程数 | 计算全局 index 的乘数 |

计算一维张量全局索引的经典写法：

```cpp
int i = blockIdx.x * blockDim.x + threadIdx.x;
```

这个两级结构直接映射硬件：**block ↔ SM**（一个 block 整体驻留在一个 SM 上，块内线程才能共享该 SM 的 shared memory）、**thread ↔ CUDA core**（以 32 线程的 warp 为单位锁步执行）。跨 block 之间**不能**同步、不共享 shared memory——这个限制换来了任意规模的可扩展性：block 数可以远超 SM 数，硬件自动排队调度。

> [!note] Wave quantization：block 要开够
> 讲义给的经验法则：**block 总数至少是 SM 数的 4 倍**（A100 即 ≥432 个）。如果 block 数只比 SM 数多一点点（如 110 个 block、108 个 SM），第一「波」占满全卡、第二波只有 2 个 block 在跑——尾巴上的利用率极低。多开 block 让调度器能把波与波之间的缝隙填满。

### 3.2 手写 CUDA GeLU

讲义中的 CUDA GeLU 形状如下（源码写在 `gelu.cu`，用 PyTorch 的 `load_inline` 即时编译并绑定成 Python 函数）：

```cpp
__global__ void gelu_kernel(float* in, float* out, int num_elements) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < num_elements) {
        float x = in[i];
        out[i] = 0.5 * x * (
            1.0 + tanh(0.79788456 * (x + 0.044715 * x * x * x))
        );
    }
}

inline unsigned int cdiv(unsigned int a, unsigned int b) {
    return (a + b - 1) / b;   // 向上取整除法：保证覆盖所有元素
}

torch::Tensor gelu(torch::Tensor x) {
    TORCH_CHECK(x.device().is_cuda());
    TORCH_CHECK(x.is_contiguous());

    torch::Tensor y = torch::empty_like(x);
    int num_elements = x.numel();
    int block_size = 1024;
    int num_blocks = cdiv(num_elements, block_size);

    gelu_kernel<<<num_blocks, block_size>>>(
        x.data_ptr<float>(), y.data_ptr<float>(), num_elements
    );
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return y;
}
```

逐条拆解关键点：
- `i < num_elements`：最后一个 block 通常不满（元素总数不是 1024 的倍数），越界线程直接跳过——这是 CUDA 里最常见的边界处理模式。
- `block_size=1024` 是每 block 线程数（硬件上限也是 1024）；`num_blocks = ⌈num_elements / block_size⌉` 决定 grid 大小，两者用 `<<<grid, block>>>` 语法传入。
- `TORCH_CHECK(x.is_contiguous())`：这个 kernel 假设元素在内存中连续排布；非连续张量（转置视图等）需要 stride-aware 的索引计算，否则结果错误——上一讲「转置不改内存布局」陷阱的又一现场。
- `C10_CUDA_KERNEL_LAUNCH_CHECK()`：CUDA 错误是异步暴露的，不主动检查会在很远的地方才炸，极难定位。

**CUDA 的代价**：
- 写 C++/CUDA，编译慢、调试痛苦（错误常常延迟异步出现）。
- Thread 索引、memory coalescing、shared memory、同步、occupancy 全部手动负责。
- Elementwise 算子还算简单；softmax、matmul、RMSNorm 这类**跨元素聚合**的算子，要手工编排 warp 内归约与共享内存，复杂度陡增。

> [!tip] 直觉
> CUDA 给你的是性能控制权的上限，同时也是心智负担的上限。它的定位类似汇编之于 C：**该会读、该能写小段，但日常工作应该发生在更高的抽象层**——这正是 Triton 存在的理由。

---

## 4. Triton：Python 里的 GPU Kernel DSL

### 4.1 Triton 想解决什么？

Triton 由 OpenAI 开源（Tillet et al., 2019），目标是让 GPU kernel 编程接近 Python 的体验，同时保留接近 CUDA 的性能。它的关键设计决策是**改变编程单位**：CUDA 让你写「一个 thread 做什么」，Triton 让你写「一个 block（tile）做什么」——块内怎么分线程、怎么合并访存、怎么用 shared memory，都交给编译器。

| 维度 | CUDA | Triton |
| :--- | :--- | :--- |
| 编程单位 | 单个 thread | 一个 block / tile |
| 语言 | C++/CUDA | Python DSL |
| Memory coalescing | 手动设计 | 编译器负责 |
| Shared memory 管理 | 手动 | 编译器负责 |
| SM 内调度 | 手动/隐式 | 编译器负责 |
| 跨 block 调度 | 人工设 grid | 仍需人工设 grid |

> [!tip] 直觉
> Triton 的抽象层级刚好卡在甜点位：比 PyTorch 低——你能控制 tile 大小、数据什么时候进 SRAM；比 CUDA 高——不必给 32 个线程逐一安排工作。经验规律：**thread 级的琐碎决策（coalescing、bank conflict）编译器已经做得和专家一样好，而 tile 级的算法决策（怎么分块、什么顺序遍历）人类仍然更强**——Triton 恰好把前者自动化、把后者留给你。

### 4.2 Triton GeLU

一个 Triton kernel 通常分两层：

1. **Python wrapper**：分配输出、设置 grid、调用 kernel；
2. **`@triton.jit` kernel**：写 tile 级计算。

```python
def triton_gelu(x: torch.Tensor):
    y = torch.empty_like(x)
    num_elements = x.numel()
    block_size = 1024
    grid = (triton.cdiv(num_elements, block_size),)
    triton_gelu_kernel[grid](x, y, num_elements, BLOCK_SIZE=block_size)
    return y

@triton.jit
def triton_gelu_kernel(x_ptr, y_ptr, num_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < num_elements

    x = tl.load(x_ptr + offsets, mask=mask)
    y = 0.5 * x * (1.0 + tl.tanh(0.79788456 * (x + 0.044715 * x * x * x)))
    tl.store(y_ptr + offsets, y, mask=mask)
```

对照 CUDA 版，能看清 Triton 的习惯用法：
- `tl.program_id(0)` 相当于 `blockIdx.x`——但注意**没有 `threadIdx`**，因为你操作的是整个 tile；
- `tl.arange(0, BLOCK_SIZE)` 一次性生成整个 tile 的偏移向量，对它的运算是「张量式」的（类似 NumPy 向量化）；
- `mask` 承担了 CUDA 里 `if (i < num_elements)` 的边界职责；
- `tl.load` / `tl.store` 是**显式**的内存访问——HBM 和片上的边界在代码里看得见，这正是性能可控的来源；
- `BLOCK_SIZE: tl.constexpr` 标记编译期常量，编译器据此展开循环、静态分配寄存器。

讲义的实测观察：
- Triton GeLU 与 PyTorch 内置 fused GeLU 速度接近；
- 比手写 Python 公式（未融合）快很多；
- 对纯 elementwise 算子，Triton 不一定碾压手写 CUDA——**它的真正优势在于让「写复杂 tile 逻辑」（归约、融合链、attention）的成本降低一个数量级**。

### 4.3 PTX：Triton 最终也要落到硬件

Triton 编译产物是 **PTX（Parallel Thread Execution）**——NVIDIA 的虚拟指令集，类似 GPU 的汇编中间表示（再由驱动编译成具体架构的机器码 SASS）。讲义打印了生成的 PTX，能读到的信息：

- `ld.global.*` / `st.global.*`：对 global memory（HBM）的读写指令——数一数就知道这个 kernel 的内存流量；
- `%ctaid.x` / `%tid.x`：block index / thread index——Triton 替你生成了 CUDA 里手写的那套索引；
- `%f*` / `%r*`：浮点/整数寄存器；
- **Thread coarsening（线程粗化）**：讲义的 PTX 显示编译器让**每个线程处理 8 个元素**而不是 1 个——摊薄索引计算的开销、增加指令级并行；
- softmax 的 PTX 里会出现 `shfl.sync.*`：**warp 内线程直接交换寄存器数据**的指令，用于归约（比经共享内存中转更快）。

> [!note] 抽象层泄漏（leaky abstraction）
> 你写的是 Python/Triton，但性能最终取决于编译器把它降到 PTX、再落到 SM/warp/寄存器/HBM 上的方式。高手和普通用户的差别，往往就是**遇到反常性能时愿不愿意往下看一层**——看 kernel 名、看 PTX、看 Nsight 的计数器。

---

## 5. `torch.compile`：让编译器帮你融合

PyTorch 2.0 之后，`torch.compile` 可以把 Python/PyTorch 代码捕获成计算图，交给编译器后端优化：

```python
def manual_gelu(x):
    return 0.5 * x * (1.0 + torch.tanh(
        math.sqrt(2.0 / math.pi) * (x + 0.044715 * x ** 3)
    ))

compiled_gelu = torch.compile(manual_gelu)
```

Profiler 里可以看到类似：

```
Torch-Compiled Region
triton_poi_fused_add_mul_tanh_0
```

这两行揭示了完整流水线：
- **TorchDynamo** 在字节码层面捕获 Python 代码、构建计算图；
- **Inductor** 后端把图编译成 **fused Triton kernel**——kernel 名字里的 `fused_add_mul_tanh` 直白地写着它融合了哪些算子；
- 也就是说，`torch.compile` 生成的正是你在第 4 节手写的那类 Triton 代码。

**优点**：
- 不用手写 Triton/CUDA，一行装饰器；
- 对 elementwise 链、reduction、常见 pattern 的融合非常有效；
- 对已有 PyTorch 模型侵入性小。

**限制**：
- **Graph break**：动态 shape、依赖张量取值的 Python 控制流（`if x.sum() > 0`）、不可捕获的 op 会把图切碎，切碎处退回 eager、融合收益消失；
- 编译本身有开销（首次调用秒级到分钟级），短脚本/频繁变 shape 的场景不一定划算；
- 生成的 kernel 融合策略偏保守，不一定比领域专家手写的强（FlashAttention 这种改变计算编排的算法，编译器发明不出来）。

> [!tip] 直觉：优化的性价比阶梯
> 先写**清晰的 PyTorch**并 benchmark；如果 profiler 显示大量小 kernel，先试 `torch.compile`（近乎免费）；只有编译器搞不定的热点，才值得手写 Triton；确认 Triton 也到不了目标，才考虑 CUDA。每往下一层，开发成本涨一个量级——**把人力花在 profiler 证明值得的地方**。

---

## 6. Fused Softmax：从逐元素到行归约

### 6.1 Softmax 的内存问题

Softmax 对矩阵 $x \in \mathbb{R}^{M \times N}$ 逐行归一化（减最大值是防 $e^{x}$ 溢出的数值稳定技巧）：

$$\operatorname{softmax}(x)_{ij} = \frac{e^{x_{ij} - m_i}}{\sum_{k=1}^{N} e^{x_{ik} - m_i}}, \qquad m_i = \max_k x_{ik}$$

朴素 PyTorch 写法：

```python
def manual_softmax(x):
    x = x - x.max(dim=-1, keepdim=True).values
    numerator = torch.exp(x)
    denominator = numerator.sum(dim=-1, keepdim=True)
    return numerator / denominator
```

这会拆成 5 个 kernel：`max`（归约）→ `sub` → `exp` → `sum`（归约）→ `div`，每一步都完整读写 $M \times N$ 矩阵。

> [!example] 算一算：手写 softmax 浪费了多少带宽？
> 对 $M \times N$ 矩阵逐 kernel 计数（单位：元素个数）：
> - `max`：读 $MN$，写 $M$；`sub`：读 $MN + M$，写 $MN$；`exp`：读 $MN$，写 $MN$；`sum`：读 $MN$，写 $M$；`div`：读 $MN + M$，写 $MN$
> - 合计：读约 $5MN$、写约 $3MN$，总流量 $\approx 8MN$
> - 而理想的融合实现只需**读 $MN$ + 写 $MN$ = $2MN$**
>
> 对这个彻底 memory-bound 的算子，内存流量差 4 倍 ≈ 速度差 4 倍——这就是讲义说 fused softmax「理论加速 4×」的来源。

### 6.2 Triton Softmax 的思路

如果一行放得进片上内存，就可以让**一个 Triton program 负责一整行**：

1. 把整行加载到 SRAM/寄存器（一次 HBM 读）；
2. 在片上做 `max` 归约、减、`exp`、`sum` 归约；
3. 归一化后一次写回（一次 HBM 写）。

```python
@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)              # 一个 program 处理一行
    col_offsets = tl.arange(0, BLOCK_SIZE)

    x_offsets = row_idx * x_row_stride + col_offsets
    mask = col_offsets < num_cols

    row = tl.load(x_ptr + x_offsets, mask=mask, other=-float("inf"))
    row = row - tl.max(row, axis=0)         # 片上归约
    numerator = tl.exp(row)
    denominator = tl.sum(numerator, axis=0) # 片上归约
    y = numerator / denominator

    y_offsets = row_idx * y_row_stride + col_offsets
    tl.store(y_ptr + y_offsets, y, mask=mask)
```

实现细节：
- `BLOCK_SIZE` 取 `triton.next_power_of_2(num_cols)`——一个 tile 覆盖整行，多出的位置用 `mask` + `other=-inf` 填充（$e^{-\infty} = 0$，不影响 max 与 sum）；
- 传 `x_row_stride` 而不是假设连续，是对非连续布局的基本防御；
- grid 大小就是行数 $M$——行之间天然独立，正是「尴尬并行（embarrassingly parallel）」。

**各实现对比**：

| 实现 | Kernel 数 | 中间结果是否落 HBM | 讲义结果 |
| :--- | :---: | :---: | :--- |
| 手写 PyTorch | 5 个 | 是 | 慢（~4× 内存流量）|
| `torch.compile` | 1 个 fused Triton kernel | 否 | 快 |
| PyTorch `softmax` | 1 个优化 kernel | 否 | 快 |
| 手写 Triton | 1 个 kernel | 否 | 快，且可继续定制 |

> [!tip] 直觉
> Softmax 比 GeLU 更有代表性：它不是纯 elementwise，而是**行内归约（row-wise reduction）**——CUDA 里要手工编排 warp shuffle 和共享内存，Triton 里一句 `tl.max(row, axis=0)` 完事。这正是 Triton「tile 级编程」抽象的价值所在。局限也随之而来：**一行必须放得进片上内存**。当行长到放不下（比如 attention 里 $N$ 很大的分数矩阵行），就得分块处理、用上一讲的 **online softmax** 增量维护 $(m, \ell)$——沿这条路走到底，就是 FlashAttention（[[Lecture 5 · GPU]] 第 6 节）。

---

## 7. Matmul：为什么最难也最重要

讲义最后回到 matmul，因为：
- Transformer 绝大部分 FLOPs 在矩阵乘法；
- GEMM 通常 compute-bound，是 Tensor Core 的主战场；
- 高性能 matmul 需要 tiling、shared memory、register blocking、warp 级 MMA、数据布局等一整套工程，是检验 kernel 功力的终极关卡。

朴素三重循环：

```python
for i in range(M):
    for j in range(N):
        acc = 0
        for k in range(K):
            acc += A[i, k] * B[k, j]
        C[i, j] = acc
```

**读写计数**：朴素实现从 HBM 读约 $MKN$ 个元素（每个输出都独立读一遍 $A$ 的行和 $B$ 的列）、写 $MN$ 个。**Tiling** 让一个 block 算一个 $T \times T$ 输出块：沿 $K$ 维每步协作加载两个 tile 进共享内存、块内复用，总读取量降为约 $2MNK/T$——**缩小 $T$ 倍**（上一讲 4.2 节的算术强度分析在此兑现）。

```mermaid
graph LR
    A["A tile<br/>M_tile × K_tile"] --> SRAM["Shared Memory"]
    B["B tile<br/>K_tile × N_tile"] --> SRAM
    SRAM --> TC["Tensor Core / MMA"]
    TC --> C["C tile<br/>M_tile × N_tile"]
```

Triton 官方 matmul 教程还有一层妙招——**分组遍历（grouped ordering）**：输出块不按行优先顺序算，而是按小组（比如 9 个一组）沿 L2 友好的顺序遍历。行优先顺序算 9 个输出块要从 HBM 载入 90 个输入 tile，分组顺序只需 54 个——**同样的计算，L2 命中率不同，带宽消耗差近一倍**。这说明即使「compute-bound」的 GEMM，访存编排仍然决定能否真正吃满 Tensor Core。

为什么 matmul 难写：
- 要让 A/B tile 在 shared memory 中充分复用，还要用寄存器再分一层块（register blocking）；
- 要避免 shared memory bank conflict（多线程撞同一个 bank 就串行化）；
- 要让 global memory 访问 coalesced；
- 要把计算映射到 Tensor Core 支持的 tile 形状（如 16×16×16 的 MMA 单元）；
- 要平衡寄存器用量与 occupancy（寄存器用太多，SM 上驻留的 warp 变少，掩盖延迟的能力下降）。

> [!warning] 易错点：不要重写 matmul
> 讲义在 dim=16384 上对比了 Triton matmul 与 PyTorch（cuBLAS）——工业级库经过多年调优，普通人手写几乎不可能超越。正确的姿势：**通用 GEMM 用 cuBLAS/CUTLASS/PyTorch 内置；只有当你需要「结构特殊的融合」时才自己写**——典型如 fused `gelu(A @ B)`（GEMM 的结果还在寄存器里时顺手做完激活再写回，省一轮 HBM 往返）、fused dequantize + matmul（推理量化场景，见 [[Lecture 10 · 推理]]）。此时 Triton 的可编程性才是不可替代的。

---

## 8. 性能优化工作流

本讲真正想训练的是一个**工作流**，而不是记住某个 kernel 的写法：

```mermaid
flowchart LR
    A["写清晰 PyTorch"] --> B["Benchmark<br/>确认慢不慢"]
    B --> C["Profiler<br/>定位慢在哪里"]
    C --> D{"瓶颈类型"}
    D -->|"多个小 kernel"| E["Fusion<br/>torch.compile / Triton"]
    D -->|"memory-bound"| F["减少 HBM 往返<br/>tiling / fusion"]
    D -->|"compute-bound"| G["用更强 kernel<br/>Tensor Core / cuBLAS"]
    D -->|"通信/同步"| H["减少同步点<br/>异步/重叠"]
    E --> I["重新 benchmark"]
    F --> I
    G --> I
    H --> I
```

**常见判断表**（profiler 症状 → 诊断 → 处方）：

| 现象 | 可能原因 | 首选动作 |
| :--- | :--- | :--- |
| profiler 里很多 `aten::add/mul/tanh` | 小算子未融合 | `torch.compile` 或 Triton |
| CUDA 时间很小但 CPU 时间大 | kernel launch overhead | 融合、增大 batch、CUDA graphs |
| `aten::mm` 占大头 | GEMM 主导（正常）| 检查 dtype、shape 对齐、Tensor Core 路径 |
| `cudaDeviceSynchronize` 很多 | CPU/GPU 同步过多 | 少用 `.item()` / `print(tensor)` |
| softmax / norm 慢 | memory-bound reduction | fused kernel |
| 显存读写过多 | 中间结果落 HBM | fusion / recomputation |

> [!tip] 直觉
> 这张流程图和判断表合起来就是一句话：**测量决定方向，抽象层级决定手段**。单卡 kernel 优化到位之后，新的瓶颈会出现在卡与卡之间——那是下一讲的战场。

---

## 总结

```mermaid
mindmap
  root((Kernel 与 Triton))
    Benchmark
      端到端时间
      warmup
      cuda synchronize
      缩放 batch dim layers
    Profiling
      torch.profiler
      CPU / CUDA time
      kernel 名称
      调用栈
    Fusion
      减少 HBM 往返
      GeLU
      Softmax
      launch overhead
    CUDA
      thread
      block
      grid
      手动 shared memory
      控制强但难
    Triton
      block 级编程
      Python DSL
      tl.load / tl.store
      编译成 PTX
      thread coarsening
    torch.compile
      Dynamo 捕获图
      Inductor 生成 Triton
      自动融合
      graph break 限制
    Matmul
      tiling 分块复用
      分组遍历 L2 友好
      fused epilogue
      别自己重写 GEMM
```

**关键要点**：

1. **性能必须实测**
   - Spec sheet 只是上限；真实速度取决于 shape、dtype、kernel 派发、版本和硬件。
   - Benchmark 要 warmup（排除编译与初始化的一次性开销），并用 `torch.cuda.synchronize()` 处理 CUDA 异步——否则测到的只是「提交任务」的时间。

2. **Profiler 是打开 PyTorch 黑箱的工具**
   - 一行高层 API 可能对应一串 kernel（`cdist` → matmul + pow + sum + cat）。
   - Kernel 名称本身就是信息：`cutlass_80_simt_sgemm_256x128_...` 写着 tile 形状，`triton_poi_fused_add_mul_tanh` 写着融合了什么。

3. **Fusion 的核心收益是少搬数据**
   - 对 GeLU、softmax、LayerNorm 这类 memory-bound 算子，少读写 HBM 比少做 FLOPs 更关键；手写 softmax 与 fused 版的内存流量差约 4 倍。
   - 手写 Python 公式会被拆成多个 kernel；PyTorch 内置、`torch.compile`、Triton 都能融合。

4. **CUDA 控制力强，Triton 性价比高**
   - CUDA 以 thread 为单位编程，适合极限优化，但 coalescing/共享内存/同步全部手动。
   - Triton 以 block/tile 为单位编程，把 thread 级琐事交给编译器，人类专注 tile 级算法——写 fused elementwise、行归约、定制 attention 的最佳层级。

5. **`torch.compile` 是第一道自动优化**
   - Dynamo 捕获图、Inductor 生成 fused Triton kernel；先试它，再考虑手写。
   - Graph break、动态 shape、编译开销是三大限制；改变计算编排的算法级优化（FlashAttention）编译器给不了。

6. **抽象层会泄漏**
   - PyTorch → Triton → PTX → SM/warp/register/HBM，每一层都可能左右性能（PTX 里能看到 thread coarsening、`shfl.sync` 归约）。
   - 真正的系统优化能力是：知道什么时候留在高层、什么时候下探一层，以及**用 profiler 证据决定**这件事。

**下一讲预告（[[Lecture 7 · 并行 I：基础]]）**：

单 GPU 的 kernel 优化只是第一步。训练更大的模型会撞上**单卡内存和算力天花板**（上一讲算过：BF16 + AdamW 全参数训练，单卡 80 GB 连 7B 都放不下），下一讲进入分布式训练：从 collective communication 原语（all-reduce、all-gather）到数据并行、ZeRO/FSDP、张量并行、流水线并行与序列并行。

---

## 复习自测

> [!question]- Q1：下面这段计时代码测出的时间为什么严重偏小？该怎么改？
> ```python
> start = time.time()
> y = model(x)
> print(time.time() - start)
> ```
> CUDA kernel launch 是异步的：CPU 提交任务后立即返回，GPU 还没算完，`time.time()` 只测到了「提交任务」的开销。修正：在计时结束前调用 `torch.cuda.synchronize()` 等 GPU 排空队列；另外首次调用包含编译/分配等一次性开销，正式计时前应先 warmup 几轮。

> [!question]- Q2：算一算——对 $M \times N$ 矩阵，手写 5 步 softmax 与 fused softmax 的 HBM 流量各是多少？速度大约差几倍？
> 手写版 5 个 kernel（max/sub/exp/sum/div）合计读约 $5MN$、写约 $3MN$，总流量 $\approx 8MN$ 个元素；fused 版读 $MN$ + 写 $MN$ = $2MN$。softmax 是 memory-bound（每元素只有一次 exp 和除法），耗时近似正比于内存流量，所以理论差距约 **4 倍**——与讲义实测量级一致。

> [!question]- Q3：概念辨析——CUDA 和 Triton 的编程单位分别是什么？这个差别为什么恰好是 Triton 的甜点位？
> CUDA 写「单个 thread 做什么」（`threadIdx` 定位到元素），Triton 写「一个 block/tile 做什么」（`tl.program_id` + `tl.arange` 操作整块）。thread 级的机械决策——memory coalescing、shared memory 编排、bank conflict——编译器已能做到专家水平，自动化掉纯赚；而 tile 级的算法决策（分块大小、遍历顺序、融合什么）人类更强，保留给程序员。于是 Triton 用一句 `tl.max(row, axis=0)` 替代了 CUDA 里手写的 warp shuffle 归约。

> [!question]- Q4：`torch.compile` 在什么情况下拿不到融合收益？为什么它永远发明不了 FlashAttention？
> Graph break 时失效：依赖张量取值的 Python 控制流、动态 shape、不可捕获的 op 会把图切碎，断点处退回 eager 执行。更根本地，编译器做的是**保语义的局部重写**（把相邻的 elementwise 链合成一个 kernel），而 FlashAttention 是**算法级重构**——用 online softmax 的数学恒等变换改变整个计算的编排顺序，这种「先证明一个恒等式再重排计算」的创造超出了模式匹配式编译优化的能力范围。

> [!question]- Q5：既然 cuBLAS 的 matmul 几乎不可战胜，什么时候仍然值得用 Triton 自己写一个含 matmul 的 kernel？
> 当你需要**结构特殊的融合**而库没有提供时：比如 fused `gelu(A @ B)`——GEMM 结果还在寄存器/SRAM 里时顺手完成激活再写回，省掉一轮 $MN$ 的 HBM 往返；或推理时的 fused dequantize + matmul。此时收益不来自「matmul 本身更快」（通常还略慢于 cuBLAS），而来自**消灭了 GEMM 前后额外算子的内存往返**。通用、无融合需求的 GEMM 永远直接调库。

---

## 参考资料

- 💻 [lecture_06.py（官方可执行讲义）](https://github.com/stanford-cs336/spring2025-lectures/blob/main/lecture_06.py)
- 📖 [PyTorch Profiler Recipe](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html)
- 📖 [PyTorch Benchmark Recipe](https://pytorch.org/tutorials/recipes/recipes/benchmark.html)
- 📖 [Triton fused softmax tutorial](https://triton-lang.org/main/getting-started/tutorials/02-fused-softmax.html)
- 📖 [Triton matmul tutorial](https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html) — tiling 与分组遍历（L2 友好）的官方讲解
- 📄 [Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations](https://www.eecs.harvard.edu/~htk/publication/2019-mapl-tillet-kung-cox.pdf)
- 📖 [CUDA Parallel Thread Execution ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html)
- 🌐 [Horace He: Making Deep Learning Go Brrrr From First Principles](https://horace.io/brrr_intro.html)
- 📖 [NVIDIA Deep Learning Performance Guide](https://docs.nvidia.com/deeplearning/performance/dl-performance-gpu-background/index.html)
- 🧩 [GPU Puzzles](https://github.com/srush/gpu-puzzles)

> [!tip] 🎓 Lecture 6 核心收获
> 高性能不是「会写 CUDA」这么粗暴，而是**先测量、再剖析、再选择正确的抽象层**：PyTorch 写得清楚，`torch.compile` 先试；热点还不够快，再用 Triton；真到极限场景，才下到 CUDA/PTX。