Lecture 6 · Kernel 与 Triton

CS336: Language Modeling from Scratch · Stanford · Spring 2025
📅 Apr 17 · 💻 可执行讲义 · 讲者 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 工程去逼近。

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

1.1 为什么需要 Benchmarking？

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

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

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

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 也常见——最小值最接近「无干扰时的真实成本」）：

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）：

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

这就需要 profiling。

1.3 Profiling：打开 PyTorch 黑箱

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

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 本身）：

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

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

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

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 代码里的魔法常数就是它。

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，中间值只存在于寄存器/共享内存：

# 未融合：每行都是一次完整的 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)

2.2 Kernel Launch Overhead

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

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: 异步返回

融合后：

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

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

3.1 CUDA 的执行模型

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

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["..."]

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

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 数，硬件自动排队调度。

3.2 手写 CUDA GeLU

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

__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 内归约与共享内存，复杂度陡增。

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，都交给编译器。

4.2 Triton GeLU

一个 Triton kernel 通常分两层：

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

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 内线程直接交换寄存器数据的指令，用于归约（比经共享内存中转更快）。

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

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

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 这种改变计算编排的算法，编译器发明不出来）。

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 写法：

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$ 矩阵。

6.2 Triton Softmax 的思路

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

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

@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）」。

各实现对比：

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

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

朴素三重循环：

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 节的算术强度分析在此兑现）。

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 变少，掩盖延迟的能力下降）。

8. 性能优化工作流

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

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 症状 → 诊断 → 处方）：

总结

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、张量并行、流水线并行与序列并行。

复习自测

start = time.time()
y = model(x)
print(time.time() - start)

参考资料

• 💻 lecture_06.py（官方可执行讲义）
• 📖 PyTorch Profiler Recipe
• 📖 PyTorch Benchmark Recipe
• 📖 Triton fused softmax tutorial
• 📖 Triton matmul tutorial — tiling 与分组遍历（L2 友好）的官方讲解
• 📄 Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations
• 📖 CUDA Parallel Thread Execution ISA
• 🌐 Horace He: Making Deep Learning Go Brrrr From First Principles
• 📖 NVIDIA Deep Learning Performance Guide
• 🧩 GPU Puzzles