CS336 2026 · Lecture 6已发布

Kernels, Triton

日期: Wed April 15
讲师: Percy
课程资料: 9 项
内容产物: 4 项

课程简介

本讲把上一讲的硬件直觉落到算子实现：学习如何建立正确的 benchmark、使用 profiler 定位瓶颈，并从访存、并行粒度和融合机会分析 Kernel。随后用 Triton 编写向量加法、Softmax 和矩阵乘等算子，理解 program、block、mask、tiling 与自动调优的配合方式。

课程资料lecture_06.py打开链接↗论文Triton - An Intermediate Language and Compiler for Tiled Neural Network Computations核心论文↗文档Triton Fused Softmax Tutorial动手实践↗文档PyTorch Profiler Recipe性能分析↗文档Triton Matrix Multiplication Tutorial动手实践↗文档PyTorch Benchmark Recipe基准测试↗文档CUDA Parallel Thread Execution ISA执行模型↗技术文章Introducing Triton - Open-Source GPU Programming for Neural Networks项目介绍↗视频课程视频课程视频↗

Assignment 1: Basics截止

Assignment 2: Systems已发布

CS336 Lecture 6：Kernels、Triton 与 GPU 性能优化

学习目标

完成本讲后，你应该能够：

理解 GPU 的硬件层次结构（寄存器、共享内存、L2 缓存、HBM）及其对性能的影响
解释线程、线程块（thread block）、网格（grid）和 warp 之间的关系
识别影响 GPU 性能的关键因素：控制发散、占用率、存储体冲突、内存合并
进行正确的基准测试（benchmarking）和性能分析（profiling）
使用 Triton 编写三种类型的内核：逐元素操作（GeLU）、按行归约（softmax、row sum）、矩阵乘法（tiling）

课程主线

本讲有两个核心部分：

回顾 GPU 硬件与编程模型：建立对 GPU 性能瓶颈的直觉
动手编写 Triton 内核：从最简单的逐元素操作出发，逐步过渡到需要线程间通信的归约和矩阵乘法

贯穿始终的主线是：编程模型负责正确性，硬件理解负责性能。Triton 的价值在于它让你以"线程块"为思考单位，自动处理大部分底层细节，但你仍需理解硬件才能真正优化性能。

一、GPU 硬件回顾

1.1 内存层次结构

GPU 的存储层次从快到慢、从小到大依次为：

层次	本讲建立的性能直觉
寄存器（每 SM）	线程私有、最快；讲义以 B200 约 65,000 个寄存器（约 256 KB）为例。
L1 缓存与共享内存（每 SM）	两者共享同一物理内存；共享内存可由程序员显式使用。
L2 缓存	芯片级共享，比每个 SM 内的存储层次更大。
HBM	容量最大、跨代增长明显，但在这条层级中访问最慢。

课程强调的是容量与带宽大致负相关：越靠近计算单元的存储越快，容量通常越小。

注意：B200 还有张量内存（tensor memory, TMEM），位于寄存器和共享内存之间，供张量核心使用，但对程序员不可见。

1.2 编程模型

三个层次

线程（thread）：在数据的一小部分上执行代码
线程块（thread block）/ CTA：一组线程，共享同一块共享内存
网格（grid）：线程块的集合，一次内核启动对应一个网格

（H100/B200 还有线程块簇（thread block cluster），支持分布式共享内存，本讲不展开。）

为什么需要线程块？

逐元素操作（如 GeLU）用线程直接处理每个元素就够了，线程之间不需要通信。
但 softmax、矩阵乘法这类操作需要线程间通信。如果全部通过 HBM 完成，代价极高。
线程块让一组线程共享 SM 上的共享内存，实现块内通信。

核心思想：线程块被调度到某个 SM 上执行，从 HBM 读取数据到共享内存，处理完后写回 HBM。

1.3 编程模型与硬件的交互

编程模型本身很简单：知道线程、线程块、网格即可写出正确的代码。但性能高度依赖硬件细节。

Warp：锁步执行的线程组

每个 warp = 32 个线程。64 线程的线程块 = 2 个 warp。
warp 内所有线程必须以锁步（lockstep）方式执行同一条指令。
控制发散（control divergence）：如果 warp 内不同线程走不同分支（if A else B），必须串行执行，严重降低效率。
SM 同时运行多个 warp，一个 warp 等待 HBM 时，调度器零成本切换到另一个 warp——这是延迟隐藏的关键机制。

占用率（Occupancy）

每个线程最多使用 255 个寄存器；SM 寄存器总量固定（如 65536 个）。
每个线程用寄存器越多，能同时调度的线程越少 → 占用率降低。
占用率低不一定是坏事：如果每个线程做更多工作（线程粗化，thread coarsening），可能反而更好。

计算示例：128 线程的线程块，每线程 160 个寄存器：

每块需 128 × 160 = 20,480 个寄存器
SM 最多同时容纳 65,536 ÷ 20,480 ≈ 3 个块 = 384 线程 = 12 个 warp
若 SM 最多支持 64 个 warp，占用率 = 12/64 = 18%

存储体冲突（Bank Conflicts）

共享内存分为 32 个存储体（bank），每个 4 字节宽。
每个时钟周期每个 bank 只能被一个线程访问。
多个线程访问同一 bank → 串行化 → 存储体冲突。
最坏情况：32 个线程同时访问矩阵第一列 → 32 路冲突。
矩阵乘法中不可避免（访问 A 的行和 B 的列），解决方案是 swizzling（重新排列共享内存布局）。

内存合并（Memory Coalescing）

warp 的 32 个线程访问 HBM 时，访问被合并为 128 字节的缓存行事务。
理想情况：所有线程访问同一缓存行 → 一次事务完成。
访问列时，大部分加载的数据未被使用 → 浪费带宽。

区分的要点：存储体冲突是共享内存的问题；内存合并是 HBM 的问题。

块占用率（Block Occupancy）

线程块以"波次"（wave）调度到 SM。
若线程块数不是 SM 数的整数倍，最后一波只有少量块 → 部分 SM 空闲。
例：B200 有 148 个 SM，启动 160 个块 → 第二波只有 12 个块。
解决思路：让线程块数量整除 SM 数量。

二、基准测试与性能分析

2.1 方法论：成功配方

对代码做基准测试和性能分析
做出修改
再次基准测试和性能分析

关键原则：在编写任何内核之前，先测量现有代码，找出瓶颈在哪里。

2.2 基准测试（Benchmarking）：测量端到端时间

基准测试只回答"花了多久"，不回答"时间花在哪里"。

正确做法

python

# 预热（warmup）：排除惰性编译等一次性开销
for _ in range(num_warmups):
    run()
torch.cuda.synchronize()  # 必须等待异步 CUDA 操作完成

# 多次计时，取平均
for trial in range(num_trials):
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    start_event.record()
    run()
    end_event.record()
    torch.cuda.synchronize()
    times.append(start_event.elapsed_time(end_event))

注意事项：

必须预热：首次运行可能包含编译时间，而稳态性能才是你关心的。
使用 CUDA events 而非 CPU 计时：避免测量到 CPU 启动内核的开销。
多次计时取平均：消除方差。更严格的做法是看分布（如 P95）。

基准测试的用途：观察扩展性

矩阵乘法的时间随维度增加应呈三次方增长。但注意：维度较小时（如 < 2000），时间几乎不变——GPU 是为大矩阵乘法设计的，小矩阵会很低效。

2.3 性能分析（Profiling）：定位瓶颈

性能分析告诉你时间花在哪里，还能揭示"引擎盖下实际发生了什么"。

PyTorch 内置 profiler

python

with torch.profiler.profile(activities=[ProfilerActivity.CUDA]) as prof:
    run()
    torch.cuda.synchronize()
print(prof.key_averages().table(sort_by="cuda_time_total"))

观察结果

a + b：调用名为 kernel at CUDA functor add 的内核——一个张量加法内核。
a @ b：调用 Cutlass 内核，名称形如 cutlass3x_sm100_simt_sgemm_f32_f32_f32_f32_f32_64x64x16...
- cutlass：NVIDIA 的线性代数 CUDA 库
- sm100：Blackwell 架构（B200）
- f32：单精度浮点
- 64x64x16：分块（tile）形状

重要发现：改变张量维度会触发不同的内核（如从 64x64x16 变为 32x32x16）。PyTorch 在底层为你做了大量决策。

三、案例研究：GeLU 的三种实现

3.1 GeLU 函数

GeLU 激活函数（tanh 近似）：

gelu(x) = 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x³)))

3.2 三种实现

Naïve PyTorch 实现：直接写出公式，每一步是独立内核
PyTorch 内置实现：torch.nn.functional.gelu(x, approximate="tanh")
编译实现：torch.compile(naive_gelu)

基准测试中，Naïve 版本明显慢于内置实现与编译实现；重点不在某个单点数值，而在于三者的数学结果相同、底层 kernel 数量和内存流量却不同。

3.3 为什么性能差异这么大？

用 profiler 查看：

Naïve 实现：调用多个内核（加法、tanh、乘法等）。每次内核调用都从 HBM 读取数据、写回结果，然后下一个内核再读取……没有内核融合。
内置实现：只有一个 GeLU CUDA 内核。
编译实现：torch.compile 生成单个 Triton 内核，将所有操作融合。

核心概念——内核融合（kernel fusion）：将多个逐元素操作合并为一个内核，每个元素只从 HBM 读一次、写一次，而非多次往返。

四、Triton 入门

4.1 CUDA vs Triton

CUDA（NVIDIA 开发）：

思维模型：每个线程做什么
优点：细粒度控制，接近底层
缺点：线程间通信需要手动管理同步（synchronize）、共享内存分配

Triton（OpenAI 开发）：

思维模型：每个线程块做什么
优点：不需要显式处理线程同步、共享内存分配；编译器自动生成 PTX
适合从零开始的课程和大多数训练/推理场景

4.2 Triton 的概念框架

思考"块做什么"：一个块将数据加载到共享内存 → 操作 → 写回全局内存
是"巨型矩阵操作"（PyTorch 风格）和"单线程操作"（CUDA 风格）之间的中间点
实际上你写的是 Python 代码，使用 @triton.jit 装饰器

五、Triton 内核示例

5.1 GeLU：逐元素操作

python

@triton.jit
def triton_gelu_kernel(x_ptr, y_ptr, num_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)  # 标识当前块
    start = pid * BLOCK_SIZE
    offsets = start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < num_elements  # 处理边界情况

    x = tl.load(x_ptr + offsets, mask=mask)
    # 计算 GeLU...
    tl.store(y_ptr + offsets, y, mask=mask)

启动方式：

python

num_blocks = triton.cdiv(num_elements, BLOCK_SIZE)
kernel[(num_blocks,)](x, y, num_elements, BLOCK_SIZE=BLOCK_SIZE)

关键点：

内核函数没有返回值，需要提前分配输出张量 y
方括号中的 (num_blocks,) 定义网格形状
tl.program_id 告诉块"我是谁"
tl.arange(0, BLOCK_SIZE) 生成 0 到 BLOCK_SIZE-1 的索引
mask 处理张量大小不能整除块大小的情况
指针就是整数，用偏移量访问
计算部分看起来几乎像普通 PyTorch——这是 Triton 的设计目标

5.2 底层发生了什么：PTX

Triton 编译器将内核编译为 PTX（Parallel Thread Execution）——GPU 的中间汇编语言。

观察 PTX 代码可以学到：

ld.global.* / st.global.*：从 HBM 读写
%ctaid.x：块索引；%tid.x：线程索引
%f*：浮点寄存器；%r*：整数寄存器
每个线程实际处理 8 个元素——编译器自动做了线程粗化（thread coarsening）

注意：PTX 代码只编译一次，所有线程执行同一份代码，通过线程 ID 区分自己。

5.3 Softmax：按行归约

朴素 PyTorch 版本需要约 5 次 MN 读取、3 次 MN 写入（max、subtract、exp、sum、divide），而理论只需要一次读、一次写。

Triton 版本：让每个块负责一行。

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)
    col_offsets = tl.arange(0, BLOCK_SIZE)

    x_ptrs = x_ptr + row_idx * x_row_stride + col_offsets
    x_row = tl.load(x_ptrs, mask=col_offsets < num_cols, other=float("-inf"))

    x_row = x_row - tl.max(x_row, axis=0)  # 数值稳定性
    numerator = tl.exp(x_row)
    denominator = tl.sum(numerator, axis=0)
    y_row = numerator / denominator

    tl.store(y_ptr + row_idx * y_row_stride + col_offsets, y_row,
             mask=col_offsets < num_cols)

关键点：

假设每行能放进一个块（BLOCK_SIZE >= num_cols）
stride 告诉如何将二维索引映射到线性内存
tl.max(x_row, axis=0) 沿列做归约——Triton 处理了内部的通信和同步
块之间没有共享内存，天然独立

5.4 Row Sum：一行装不下的归约

当行长度超过块大小时（如 4096 列，块大小 1024），需要在单个线程块内部循环分片。

python

@triton.jit
def row_sum_kernel(x_ptr, out_ptr, N, BLOCK_SIZE: tl.constexpr):
    row = tl.program_id(0)
    acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32)  # 每个线程的累加器

    for start in range(0, N, BLOCK_SIZE):  # 遍历所有分片
        cols = start + tl.arange(0, BLOCK_SIZE)
        mask = cols < N
        x = tl.load(x_ptr + row * N + cols, mask=mask, other=0.0)
        acc += x

    result = tl.sum(acc, axis=0)  # 最终的线程间归约
    tl.store(out_ptr + row, result)

概念区别：这里循环遍历的是同一行的不同分片（tile），不是独立的块。块对应整行，每个块处理整行的所有分片。

5.5 矩阵乘法：Tiling 的艺术

朴素方法的问题

固定输出元素 C[m,n]，对每个 k 从 HBM 读取 A[m,k] 和 B[k,n]：

读取次数 ~ M×K×N，算术强度 O(1)
大量冗余读取：C4 和 C5 都用到 A4、A5、A6

理想化方法

把整个 A 和 B 加载到共享内存：

读取次数降至 M×K + K×N，算术强度 O(N)（理想情况）
问题：A 和 B 通常太大，放不进共享内存

Tiling 方法（分块）

核心思想：在全局层面是朴素的，在局部层面是理想化的。

将输出矩阵 C 划分为输出分块（tile），每个分块对应一个线程块
对每个输出分块，扫描 A 的行分块和 B 的列分块：
- 将对应的 A 分块、B 分块从 HBM 加载到共享内存
- 做矩阵乘法，累加到部分和
完成后将输出分块写回 HBM

算术强度提升到 O(tile_size)，通常不够理想 O(N)，但已足够好。

额外收益——内核融合：既然在写内核，可以顺便应用逐元素激活函数（如 ReLU），在写回 HBM 之前完成，无需额外往返。

python

@triton.jit
def matmul_relu_kernel(a_ptr, b_ptr, c_ptr, M, N, K,
                       stride_am, stride_ak, stride_bk, stride_bn,
                       stride_cm, stride_cn,
                       BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
                       BLOCK_K: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    indices_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    indices_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    indices_k = tl.arange(0, BLOCK_K)

    a_ptrs = a_ptr + indices_m[:, None] * stride_am + indices_k[None, :] * stride_ak
    b_ptrs = b_ptr + indices_k[:, None] * stride_bk + indices_n[None, :] * stride_bn

    acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)

    for k in range(0, K, BLOCK_K):
        a = tl.load(a_ptrs, mask=(indices_m[:, None] < M) & (indices_k[None, :] + k < K), other=0.0)
        b = tl.load(b_ptrs, mask=(indices_k[:, None] + k < K) & (indices_n[None, :] < N), other=0.0)
        acc += tl.dot(a, b)
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    acc = tl.maximum(acc, 0.0)  # ReLU，融合激活

    c_ptrs = c_ptr + indices_m[:, None] * stride_cm + indices_n[None, :] * stride_cn
    tl.store(c_ptrs, acc, mask=(indices_m[:, None] < M) & (indices_n[None, :] < N))

回顾 stride 的作用：张量在内存中是线性化的。stride 告诉你多维索引 (row, col) 如何映射到线性索引：index = row * stride_row + col * stride_col。

六、与课程作业的关联

本讲内容直接服务于作业中的内核实现任务。你需要：

用 profiler（Nsight）分析 PyTorch 代码，理解底层内核调度
为注意力机制编写 Triton 内核（Flash Attention）
利用 tl.load / tl.store / tl.dot / tl.reduce 等 Triton 原语

本讲覆盖的 GeLU、softmax、矩阵乘法是 Flash Attention 的三个构成要素：

逐元素操作（GeLU）→ 理解如何加载、计算、存储
按行归约（softmax）→ 理解块内归约
矩阵乘法 tiling → 理解如何将大操作分解到多个线程块

七、常见误区

把 HBM 读写当免费：每次内核启动都伴随 HBM 往返。减少内核数量（融合）是优化第一要务。
认为占用率越高越好：低占用率不一定是坏事。有时让每个线程做更多工作（线程粗化）反而更好。
混淆存储体冲突和内存合并：前者是共享内存问题，后者是 HBM 问题。解决方案完全不同。
忘记预热或不同步：基准测试必须预热、多次计时、显式 torch.cuda.synchronize()，否则数据不可信。
认为 Triton 自动处理所有事情：Triton 处理线程块内部的同步和共享内存，但你要思考如何划分块、如何管理掩码、如何避免边界错误。
块大小越大越好：块太大可能放不进共享内存，或导致存储体冲突；块太小无法有效利用硬件。需要实验。
忽略尾波效应：线程块数量应尽量整除 SM 数量，否则最后一波会有 SM 空闲。

八、复习清单

硬件与模型

[ ] 能否画出 GPU 内存层次图，并说出每层的容量、带宽量级？
[ ] 能否解释 thread、thread block、grid、warp 的关系？
[ ] 为什么需要线程块，而不是只用一个线程网格？

性能因素

[ ] 什么是控制发散？为什么应该尽量避免？
[ ] 什么是占用率？如何计算？为什么高占用率不等于高性能？
[ ] 什么是存储体冲突？最坏情况是什么？swizzling 是什么？
[ ] 什么是内存合并？什么访问模式是最佳的？
[ ] 什么是波次量化（wave quantization）问题？

基准测试与分析

[ ] 基准测试前为什么必须预热？
[ ] 为什么用 CUDA events 而不是 CPU 计时？
[ ] 为什么 torch.cuda.synchronize() 是必需的？
[ ] profiler 能告诉你什么？如何从内核名称推断实现细节？

Triton 编程

[ ] Triton 的思维模型与 CUDA 有何不同？
[ ] tl.program_id、tl.arange、tl.load、tl.store 各自做什么？
[ ] 为什么要用 mask？什么时候需要？
[ ] Triton 生成的 PTX 代码中，%ctaid.x 和 %tid.x 分别代表什么？

算法模式

[ ] 逐元素操作（GeLU）的 Triton 内核结构是怎样的？
[ ] 按行归约（softmax）为什么适合用"一行一块"策略？
[ ] 一行太大时怎么办？（分片循环 + 累加器）
[ ] 矩阵乘法的 naive、idealized、tiling 三种方法有何区别？
[ ] Tiling 如何影响算术强度？
[ ] 什么情况下可以在写回 HBM 前融合激活函数？

从 PyTorch 到 Triton：理解 GPU 内核编写的核心思维

导语

当我们在 PyTorch 中写下 x @ y 或 torch.softmax(x, dim=-1) 时，我们很少会思考底层到底发生了什么。但真正需要榨取性能时，你就必须直面 GPU 的硬件现实。本文是斯坦福 CS336《Language Modeling from Scratch》第六讲的完整梳理，将带你从硬件架构出发，逐步走到 Triton 内核的编写实践。

上一讲我们建立了 GPU 性能的宏观认知，这一讲则深入代码层面。我们会看到：为什么简单的 PyTorch 表达式可能性能糟糕？为什么需要理解 warp、共享内存和存储体冲突？以及如何用 Triton 编写 GeLU、Softmax 和矩阵乘法内核。最终，这些知识将直接服务于你实现 Flash Attention 和优化 Transformer 训练的目标。

一、GPU 硬件回顾：容量与带宽的权衡

让我们从一张典型的 NVIDIA GPU 简化示意图开始。课程将 A100、H100、B200 视为同一设计思路的不同代际：每张卡都有大约一两百个流式多处理器（Streaming Multiprocessors，SMs），代际变化更多体现在存储容量、带宽与具体硬件功能上。

每个 SM 内有寄存器，以及共享同一物理内存的 L1 缓存和共享内存。
L2 缓存由整个芯片共享；高带宽内存（HBM）容量最大，并且跨代增长明显。
课程用 B200 作为例子：每个 SM 约有 65,000 个寄存器，合计约 256 KB。

这里的关键规律是：容量与带宽大致负相关。寄存器最快，随后是 L1、L2，HBM 在这条层级中最慢；优化时首先要避免不必要地往返较慢的存储层次。

二、编程模型：线程、线程块与网格

GPU 的编程模型可以用三层结构概括：

线程（Thread）：在数据的一小部分上执行代码的最小单元。
线程块（Thread Block），也称为并发线程数组（Concurrent Thread Arrays，CTAs）：一组线程，它们可以访问同一块共享内存。
网格（Grid）：所有线程块的集合。

当你启动一个内核时，实际上是启动了一个由线程和线程块组成的网格，让它们全部并行执行。

为什么需要线程块？

如果你只想做逐元素操作（比如 GeLU），每个线程处理一个元素就够了。但对于需要线程间通信的操作（如 softmax 或矩阵乘法），这种视图远远不够。原因在于：如果每个计算单元都直接读写 HBM，你会被内存带宽限制死死卡住。

正确的做法是使用共享内存（Shared Memory），它位于 SM 本地，存取速度远快于 HBM。一个线程块被调度到一个 SM 上执行，块内的所有线程共享这块快速内存。整个过程是：从 HBM 读取数据到共享内存，做计算（可能涉及线程间通信），再写回 HBM。

Triton 的核心思想就是让你以线程块为单位思考，而不是每个线程。

从第一张结构图到一次 Block 的执行

可以把 GPU 想象成一座工厂：CPU 负责下达任务，GPU 负责大量并行计算。GPU 内部有许多流式多处理器（Streaming Multiprocessors，SMs），每个 SM 都有自己的计算核心、寄存器和 shared memory；所有 SM 共享 GPU 的 global memory（显存）。SM 可以理解成一个能够独立工作的“车间”，但它不是 CPU 核心的简单放大版：一个 SM 会同时管理许多线程和多个线程块。

程序员通常不会只启动一个线程，而是启动一个由大量线程组成的 Grid。Grid 又被划分为许多 Block，每个 Block 由若干 Thread 组成：

text

Grid：整个计算任务
└── Block：任务的一小块
    └── Thread：处理一个或几个元素

例如，对于向量加法 c[i] = a[i] + b[i]，如果有 8192 个元素、每个 Block 有 32 个线程，就需要 256 个 Block。一个 Block 内的线程可以合作并访问同一块 shared memory；不同 Block 之间通常不能直接通信，因此它们可以被 GPU 以任意顺序调度。

GPU 启动任务后，会把 Block 动态分配给 SM：

text

Block 0 ──→ SM 0
Block 1 ──→ SM 1
Block 2 ──→ SM 0

一个 Block 必须整体放在一个 SM 上，不能拆到多个 SM，因为 Block 内的线程需要共享 shared memory，并且可能需要同步。如果 SM 资源（寄存器、shared memory 等）足够，一个 SM 可以同时容纳多个 Block；资源不足时，只能容纳较少的 Block。

Block 被放入 SM 后，硬件会把线程分成多个 Warp。一个 Warp 通常包含 32 个线程，因此 128 个线程的 Block 会被分成 4 个 Warp：

text

Block 0
├── Warp 0：Thread 0–31
├── Warp 1：Thread 32–63
├── Warp 2：Thread 64–95
└── Warp 3：Thread 96–127

GPU 的指令调度基本以 Warp 为单位：同一 Warp 中的线程执行相同的指令，但处理不同的数据。以向量加法为例，每个线程根据自己的 Block ID 和 Thread ID 定位元素：

python

i = block_id * threads_per_block + thread_id
c[i] = a[i] + b[i]

因此，一个 Block 的执行过程可以概括为：

text

CPU 发起 kernel
→ GPU 创建 Grid
→ 调度器把 Block 分配给 SM
→ SM 把 Block 分成 Warp
→ Warp 中的线程读取数据、执行计算、写回结果
→ Block 完成后，SM 继续处理下一个 Block

这一过程解释了后面几个性能概念的来源：Block 是否能同时驻留在 SM 上，取决于资源使用量；Warp 中的线程是否访问连续地址，会影响内存合并；Warp 是否走不同分支，会产生控制发散。

三、编程模型与硬件的交互：那些决定性能的细节

编程模型本身很优雅，但性能对硬件极其敏感。以下是五个你必须理解的关键概念。

1. Warp 与控制发散

在硬件层面，线程块中的线程被分成 warp，每个 warp 有 32 个线程。一个 warp 内的所有线程必须以锁步（lockstep）方式执行同一条指令。如果遇到分支（if A else B），warp 只能先执行 A 路径，再执行 B 路径，这是串行化的，称为控制发散（Control Divergence），应该尽量避免。

warp 的存在还有一个重要目的：隐藏延迟。当某个 warp 因为 HBM 读取而阻塞时（这可能需要 100 个周期），SM 可以零成本切换到另一个可执行的 warp。这正是 GPU 通过大量并发来掩盖内存延迟的设计哲学。

2. 占用率

每个线程最多使用 255 个寄存器，而每个 SM 的寄存器总数是固定的（如 B200 为 65,000）。因此，每个线程使用的寄存器越多，SM 能同时容纳的线程就越少，占用率（Occupancy）就越低。

但这不一定是坏事。课程中有一个具体例子：128 个线程的块，每个线程用 160 个寄存器，SM 最多只能同时运行 3 个块（12 个 warp），占用率仅为 18%。如果这是瓶颈，你可以通过**线程粗化（Thread Coarsening）**来调整——让每个线程处理多个元素，而不是一个。这样线程数减少，但每个线程做更多工作，可能反而更高效。

3. 存储体冲突

共享内存被划分为 32 个存储体（Banks），每个存储体 4 字节宽。每个周期，每个存储体最多只能被一个线程访问。如果多个线程试图访问同一个存储体的不同位置，访问就会被串行化，这就是存储体冲突（Bank Conflicts）。

最坏情况：32 个线程同时访问矩阵的第一列，就会造成 32 路存储体冲突。这在矩阵乘法中几乎是不可避免的（你必须访问行和列）。解决方案是 swizzling——重新排列共享内存的存储布局，使得冲突概率降低。

4. 内存合并

当 warp 中的 32 个线程访问 HBM 时，访问会被合并成 128 字节的事务（相当于一个缓存行）。如果所有线程访问连续内存（如第 1 行第 0 列、第 0 行第 1 列……），这就是完全合并，一次就能抓取全部数据。反之，访问列则会浪费大量带宽。

注意，这与存储体冲突很相似，但它们是两个不同的限制：一个涉及共享内存，一个涉及 HBM。

5. 线程块占用率与波次量化

线程块以波次（waves）调度到 SM 上。B200 有 148 个 SM，如果你启动 160 个线程块，第一波是 148 个，第二波只有 12 个，剩下大量 SM 空闲。这就是波次量化问题（Wave Quantization）。理想情况下，线程块的数量应该能整除 SM 的数量。

四、基准测试与性能分析：成功的方法论

课程反复强调一个成功公式：对代码进行基准测试和性能分析，做出修改，然后再测试。关键是，要在编写内核之前就测量代码，找出瓶颈在哪。

基准测试的基本功

课程用矩阵乘法演示了从零开始做基准测试的方法。有三个注意事项：

预热（Warmup）：一些操作是惰性编译的，第一次运行的时间不能代表稳定状态。
多次计时：考虑方差，通常取平均值。
使用 CUDA 事件而非 time.time()：CUDA 事件能提供更精确的 GPU 时间，并且需要同步（torch.cuda.synchronize()）确保所有异步操作完成。

一个有趣的观察是：当矩阵维度小于约 2000 时，时间几乎是常数；之后呈三次方增长。这是因为 GPU 是为大型矩阵乘法设计的，小矩阵无法充分利用硬件。

性能分析：看穿 PyTorch 的伪装

PyTorch 内置了性能分析器。当你分析 a + b 时，会看到一个名为 kernel at CUDA functor add 的内核——这告诉你，即使是最简单的加法，底层也是一个 CUDA 内核。

更有趣的是矩阵乘法：a @ b（2048×2048）会调用 cutlass3x_sm100_simt_sgemm_f32_f32_64x64x16...，而 128×128 的矩阵乘法调用的内核名却不同（32x32x16）。这些名字其实透露了实现细节：

cutlass 是 NVIDIA 的线性代数 CUDA 库。
sm100 对应 Blackwell 架构（B200）。
f32 是数据类型。
64x64x16 是分块（tile）的形状。

这意味着：同样的 PyTorch 代码，不同维度会调用不同的底层内核，而你根本无法直接控制这些选择。

五、GeLU 案例：内核融合的巨大威力

让我们将基准测试和性能分析应用到 GeLU 激活函数上。课程对比了三种实现：

Naïve 实现：直接用 PyTorch 公式写出 0.5 * x * (1 + torch.tanh(0.79788456 * (x + 0.044715 * x^3)))。
内置实现：torch.nn.functional.gelu(x, approximate="tanh")。
编译实现：对 naïve 版本调用 torch.compile。

基准测试中，naïve 实现明显慢于内置和编译版本。课程借此强调：即使数学表达式相同，kernel 数量与中间张量读写也会显著改变性能。

性能分析揭示了原因：

naïve 实现调用了多个内核（加法、乘法、tanh 等），每个内核都从 HBM 读数据、算完写回，下一个内核再从 HBM 读。这就是没有内核融合（No Kernel Fusion），导致大量 HBM 往返。
内置版本是一个单独的 GeLU 内核，因为它太常用，所以有人专门写了一个融合内核放进标准库。
编译版本（torch.compile）生成的也是一个单一的 Triton 内核——编译器自动识别出整个计算图，合并成一个融合操作。

结论是明确的：内核融合（Kernel Fusion）将多次 HBM 读写减少为一次读、一次写，这是性能提升的关键。

六、Triton 入门：以线程块为单位的思考

CUDA 的思维模型是“每个线程做什么”：你写一段代码，通过线程 ID 定位自己，然后执行操作。这给了你细粒度控制，但你得自己管理共享内存、同步等细节。

Triton 的思维模型是“每个线程块做什么”。你指定每个块要执行的计算，Triton 编译器负责把块映射到具体的 warp 和线程。对于本课程来说，Triton 已经足够强大。

Triton 代码看起来非常像 Python，但有三个关键差异：

没有返回值，只有显式的读写：你需要分配输出张量，内核往里写。
指针运算：输入输出是整数指针，你需要计算偏移量。
块（Block）的概念：每个块处理数据的一个片段。

第一个 Triton 内核：GeLU

python

@triton.jit
def triton_gelu_kernel(x_ptr, y_ptr, num_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)      # 当前块 ID
    start = pid * BLOCK_SIZE         # 该块的起始位置
    offsets = start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < num_elements    # 处理边界

    x = tl.load(x_ptr + offsets, mask=mask)
    # 计算 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)

流程很清晰：醒来（获取 PID）→ 定位数据（计算偏移量）→ 读取（tl.load）→ 计算 → 写回（tl.store）。

查看生成的 PTX（并行线程执行，GPU 的中间汇编）代码，你会发现：

ld.global.* 和 st.global.* 是内存读写指令。
每个线程实际处理 8 个元素（线程粗化），编译器自动做了这个优化。
%ctaid.x 是块索引，%tid.x 是线程索引——代码只编译一次，每个线程通过索引区分自己。

七、Softmax：归约操作与行列设计

Softmax 不再是逐元素操作，它需要对矩阵的每一行进行归约（计算最大值、求和）。这里的关键设计决策是：让每一行成为一个线程块。因为行与行之间不需要共享内存，块之间可以完全独立。

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)
    col_offsets = tl.arange(0, BLOCK_SIZE)

    x_ptrs = x_ptr + row_idx * x_row_stride + col_offsets
    x_row = tl.load(x_ptrs, mask=col_offsets < num_cols, other=float("-inf"))

    x_row = x_row - tl.max(x_row, axis=0)
    numerator = tl.exp(x_row)
    denominator = tl.sum(numerator, axis=0)
    y_row = numerator / denominator

    y_ptrs = y_ptr + row_idx * y_row_stride + col_offsets
    tl.store(y_ptrs, y_row, mask=col_offsets < num_cols)

这里有一个新元素：tl.max 和 tl.sum 在块内做归约。Triton 自动处理了线程间的通信，你不需要手动同步。

对比朴素 PyTorch 实现，Triton 版本大幅减少了 HBM 读写次数（从 5MN 次读取和 3MN 次写入降到 MN 读取和 MN 写入），这正是性能提升的核心。

八、行求和：当数据装不下一个块

当矩阵的行超过一个块能容纳的列数时（比如 4096 列，但块大小只有 1024），就需要一种新的策略。课程以行求和为例演示了这个问题。

核心思路是循环遍历多个块片（Tiles）：

python

@triton.jit
def row_sum_kernel(x_ptr, out_ptr, N, BLOCK_SIZE: tl.constexpr):
    row = tl.program_id(0)
    acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32)

    for start in range(0, N, BLOCK_SIZE):    # 遍历所有块片
        cols = start + tl.arange(0, BLOCK_SIZE)
        mask = cols < N
        x = tl.load(x_ptr + row * N + cols, mask=mask, other=0.0)
        acc += x

    result = tl.sum(acc, axis=0)    # 最终归约
    tl.store(out_ptr + row, result)

每个线程保持一个累加器，逐块片处理一行数据。最后，通过 tl.sum(acc, axis=0) 将所有线程的累加器合并成一个标量。

这是从逐元素到分块的过渡——数据不再能一次性装进块里，你需要引入循环和更复杂的索引管理。

九、矩阵乘法：分块（Tiling）的经典实践

矩阵乘法是深度学习的命根子，也是本讲的压轴戏。我们来看三种方法的演进：

朴素方法

固定某个 (m, n) 位置，对每个 k 从 HBM 读入 A[m, k] 和 B[k, n]，乘累加。这是正确但极其低效的：读取次数是 M×K×N 量级，算术强度（操作数 / 传输字节数）是常数 O(1)，完全不可接受。

理想化方法

把整个 A 和 B 都加载到共享内存，然后计算 C。这样读取次数降到 M×K + K×N，算术强度达到理想的 O(N)。但问题是：A 和 B 通常太大，无法放入共享内存。

分块方法（Tiling）

这是真正的解法。把输出矩阵 C 分成多个输出块（tiles），每个线程块负责一个输出块。对于每个 (A 的行块, B 的列块)：

加载对应的 A 块和 B 块到共享内存。
执行块级矩阵乘法（tl.dot）。
累加到部分和中。

python

@triton.jit
def matmul_relu_kernel(a_ptr, b_ptr, c_ptr, M, N, K, ...):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    indices_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    indices_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    indices_k = tl.arange(0, BLOCK_K)

    a_ptrs = a_ptr + indices_m[:, None] * stride_am + indices_k[None, :] * stride_ak
    b_ptrs = b_ptr + indices_k[:, None] * stride_bk + indices_n[None, :] * stride_bn

    acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)

    for k in range(0, K, BLOCK_K):
        a = tl.load(a_ptrs, mask=...)
        b = tl.load(b_ptrs, mask=...)
        acc += tl.dot(a, b)          # 核心：块级矩阵乘法
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    acc = tl.maximum(acc, 0.0)       # 内核融合：在写回前应用 ReLU
    tl.store(c_ptrs, acc, mask=...)

这个实现的额外好处是内核融合：在写回 HBM 之前，你可以对累加结果应用任何逐元素函数（如 ReLU）。这避免了额外的内核启动和 HBM 往返。

这里的算术强度提升到分块大小的量级，虽然达不到理想的 O(N)，但如果分块足够大，效果已经很好。

十、课程总结与学习路径

这一讲在整个课程中的位置非常关键：它填补了从“宏观理解 GPU”到“实际编写内核”之间的鸿沟。上一讲（第五讲）建立算法与硬件的桥梁，这一讲则让你手握 Triton 这一工具，为接下来的课程（多 GPU 编程、Flash Attention 实现）打下基础。

核心要点可以归纳为几个层次：

编程模型保证正确性：PyTorch、Triton、PTX 三个层次提供了从高层抽象到底层控制的不同选择。
硬件细节决定性能：SM 数量、warp 锁步执行、存储体冲突、内存合并、占用率——这些都会影响最终速度。
基准测试和性能分析帮助你找到瓶颈：永远先测量再优化。
Triton 让线程块思考变得自然：读入共享内存 → 操作（可以融合）→ 写回 HBM。

从 GeLU（逐元素）到 Softmax（行内归约）、行求和（块片循环）再到矩阵乘法（真正的分块），难度逐步递增。理解了这些例子，你就掌握了实现 Flash Attention 所需的全部要素——它本质上就是 Softmax 和矩阵乘法在分块策略下的融合。

下次我们将从单个 GPU 转向多个 GPU，探讨数据并行和模型并行等更深层次的分布式训练技术。但在此之前，请确保你已经掌握了本讲的核心：用 Triton 以线程块为单位思考，让硬件的每个部分都物尽其用。

参考资料

官方视频：CS336 2026 Lecture 6
课程讲义：CS336 · Language Modeling from Scratch（Lecture 6: Kernels, Triton）
GPU 架构入门讲义：Stanford CS149 · GPU Architecture & CUDA Programming

CS336 2026 Lecture 6：完整中文逐字稿

视频：官方视频
来源：英文人工字幕提取后翻译为中文；保留段落起始时间戳，只修正明显的断句、重复和语序。
说明：这是逐字稿，不是讲义摘要；相邻且语义连续的片段已合并为阅读段落；术语保留必要英文。

正文

[00:05.240] 欢迎大家回来。周一的时候，Tatsu 非常出色地概述了 GPU、如何思考性能问题以及 GPU 带来的各种特性。这节课我们将继续深入，更仔细地研究代码，编写一些 Triton 内核，并进行一些基准测试和性能分析。那么首先，为了回顾一下 GPU 的运作方式，这里是一个典型 GPU 的简化示意图。有内存，然后是实际的 GPU 芯片。请大家记住 GPU 的特点。当然，我们讨论的是 NVIDIA GPU，但 Tatsu 也提到了 CPU，还有 AMD 等等。但让我们聚焦于 NVIDIA GPU。我们经历了多代产品，从 A100 到 H100，再到 B200。每一代都有一个 GPU，它拥有一定数量的流式多处理器（Streaming Multiprocessors，SMs）。SM 的数量大约在 100 到 200 个之间。这个数字变化不大。

[01:30.250] 在每个单独的 SM 内部，有一组寄存器。比如 B200 每个 SM 有 65,000 个寄存器，总计 256K。这个数字变化也不大。此外，还有 L1 缓存和共享内存（shared memory）。请记住，这两者是同一块物理内存，共享内存是可以由你控制的。L1 缓存你不能控制。这是每个 SM 私有的。两者的大小也在同一个数量级。然后是 L2 缓存，它不是每个 SM 私有的，而是整个芯片共享的。它要大一些。最后是高带宽内存（High Bandwidth Memory，HBM），容量非常大。你会看到，真正增长很快的是这个数字。好的。除了容量，你还可以考虑带宽的情况。基本上，带宽与容量是负相关的。寄存器非常快，L1 稍慢一些，L2 更慢，而高带宽内存是最慢的。

[02:50.850] 尽管每秒 8 TB 的速度在宏观尺度上也不算慢。这就是你脑海中应该有的主要层次结构。大容量内存速度慢、距离远但空间大。而像寄存器和 L1 这样的快速内存位于 SM 上，它们是本地的、快速的，但容量小。好的，这就是我们所处的环境。那么，如何对 GPU 进行编程呢？编程模型如下。这里有线程（threads）的概念，每个线程在数据的一小部分上执行一段代码。你可以把这里的每一个箭头想象成一个线程。这些线程被组织成线程块（thread blocks），也称为并发线程数组（Concurrent Thread Arrays，CTAs）。这是一个线程组。最后，你有一个由线程块组成的集合，叫做网格（grid）。当你启动一个内核时，你实际上是启动了一个由线程和线程块组成的网格，让它们全部并行地同时执行某些计算。

[03:59.690] 好的，这是简化后的模型。还有一些其他的情况。例如，H100 和 B200 还有线程块簇（thread block clusters），它们是线程块的集合，支持一定量的分布式共享内存。此外，B200 还有张量内存。这是为了张量核心（tensor cores）准备的，其速度介于寄存器和共享内存之间。其中有些东西对程序员是不可见的，但它们存在于硬件中。不过在这节入门课上，我们不会深入讨论这些。好的，你可能会问一个问题，为什么我们需要线程块？为什么我们不能只有一个线程网格，每个线程直接取一块数据然后执行操作呢？如果你只想做逐元素（element-wise）的操作，这通常没问题，我们稍后会看到。例如，GeLU 就是一个逐元素应用的激活函数。

[04:54.460] 线程在这里很自然。每个线程处理一个元素。就像对数据集中的每个 i 执行 Fi 一样。但对于涉及线程间通信的操作，比如 softmax 或矩阵乘法，这种视图就远远不够了。原因是，当然，如果你愿意承担读写 HBM 的代价，那它可能是够用的。因为那样你仍然可以——比如，在矩阵乘法中我们稍后会看到——你可以让每个单元格里的每个元素都去计算矩阵的那个元素。它可以直接读写 HBM。但正如我们注意到的，HBM 非常慢，所以这不是一个好策略。因此，我们应该做的是使用共享内存，它位于 SM 本地。线程块让你能够考虑的是一组线程，所有这些线程都将访问这个共享内存。

[06:05.080] 因此，你可以把一个线程块想象成被调度到一个流式多处理器上执行。然后它做它的事情。它要做的是从 HBM 读取大量数据，然后进行处理，处理过程可能涉及线程间通过共享内存进行的通信，最后把结果写回去。好的。这是一个关键点。稍后我们会讨论分块（tiling），这是整个游戏的核心。好的，在 Triton 中——这实际上是我们现在编写内核的主要方式——我们将自然地以线程块为单位进行思考。而且我认为一旦你习惯了用线程块的思维来思考，这会变得非常有意义，并且让你的生活轻松很多，我们稍后会看到。好的，编程模型本身我认为是相当简单的。有线程、线程块和网格。我认为事情变得更复杂的地方在于编程模型和硬件之间的交互。

[07:23.210] 编程模型实际上非常好。它提供了对硬件的抽象。当你编写一个内核时，你只需要知道有一堆线程块。你定义它们，然后定义块内每个线程或线程块需要执行的计算。这些就像是写 Python 代码一样。所以这部分不难。事实上，如果你不关心性能——如果你只关心正确性，你只需要知道这些。但在实践中，性能对硬件非常敏感。所以你需要真正深入地理解硬件才能获得高性能。事实上，我们讨论 GPU 和内核的全部原因就是你想榨取性能。所以这里有两个层面：你试图理解计算，这仅仅是编程模型的一部分。但运行速度将在很大程度上取决于硬件。

[08:26.140] 好的，我将给出一些例子来说明为什么这很重要，其中一些可能是对 Tatsu 所讲内容的回顾。但希望能强化这些概念。我想大概会有五个不同的例子。这只是为了让你对这些考虑因素有个初步了解。比如有个叫 warp 的概念，我之前没提到。在这种非常简化的视图中，warp 并不是编程模型这个清晰图景的一部分。你有线程和线程块或网格。技术上，你也可以超越 warp 来考虑，但仅仅编程时你不必这样做。那么什么是 warp？每个线程块，记住，是一组线程。但这些线程实际上被分组成 warp。每个 warp 基本上有 32 个线程。例如，如果一个线程块有 64 个线程，那么就有两个 warp，前 32 个和后 32 个。

[09:29.840] 正如 Tatsu 上次提到的，一个 warp 内的所有线程，所有这些 T，必须在 SM 上以锁步（lockstep）方式执行同一条指令。每个周期，它们必须执行完全相同的指令。而控制发散（control divergence）是指 warp 内不同的线程需要执行不同的指令。例如，如果你有分支语句，如果满足某个条件，则执行A，否则执行B。那么接下来发生的是，在同一个线程束中，你只能先让那些执行A的线程工作，然后再执行B。所以这会变成串行化的。因此这既糟糕又低效，这就是为什么通常要尽量避免使用分支。我认为线程束中一个很酷的地方在于，一个流式多处理器实际上运行着多个线程束。有一个线程束调度器，还有一些即将运行的驻留线程束。每个线程束都有自己的——线程拥有寄存器。SM可以在它们之间以零成本进行切换。

[10:40.610] 没错，这在一般情况下是不成立的，比如在CPU上就不是这样。但这样设计的目的就是为了隐藏延迟。这一点很重要，因为其中一个线程束，例如，正在从HBM读取数据，记得那非常昂贵，可能需要100个周期左右。你不能就闲坐在那里等着那个线程束，让它无所事事。你应该立即切换到另一个线程束，它实际上可以执行一些张量核心操作。好的，这就是关于线程束需要注意的一点。线程束出现的另一个原因是占用率这个概念。更具体地说，是线程束占用率。每个线程，仅仅从硬件限制来说，每个线程最多只能使用255个寄存器。好的，那么接下来发生的是，SM拥有的寄存器数量是固定的。所以每个线程使用的寄存器越多，你能拥有的线程就越少。这只是数学问题。这可能会降低你的占用率。但这不一定就是坏事，因为如果你有更少的线程。

[11:54.790] 但每个线程做更多的工作，这实际上可能是好的。所以占用率是你可以测量的东西，但占用率不一定越大越好。因为这里还有其他一些权衡。事实上，一个你可能想要更少线程的例子是线程粗化这种思路，比如说你有一个逐元素操作。你可以让一个线程只对一个元素执行操作。好的，这样你就有很多线程。但你也可以让每个线程处理多个元素。比如一个常数，也许是八个。这样会给你更少的线程，这使得调度等工作更简单，但每个线程会做更多的工作。所以如果你的线程非常轻量，也许你想让它们变“胖”一点。好的。这里只是举一个可能发生情况的例子。假设你有一个线程块，一个包含128个线程的线程块。每个线程使用160个寄存器。好的。这里存在一个硬件限制。

[13:06.270] B200每个SM拥有的寄存器上限是65,000个。它还有一个限制，规定了一次最多不能运行超过64个线程束。那么你就可以做一些简单的数学计算来算出这里的占用率。即每个线程的寄存器数量。这必须小于255。这没问题。每个块使用的寄存器数量是每个块的线程数乘以每个线程的寄存器数。大约是20,000个。这意味着你的SM最多只能同时运行三个块。这对应于12个线程束。而由于最大线程束数是64，这意味着你的占用率是18%。所以我们只运行了总线程束数量的18%。这是因为每个线程使用了大量的寄存器。好的，这是一个例子，说明在这种情况下，内存——我认为这里指的是寄存器资源——在计算能力方面限制了你的发挥。

[14:21.420] 好的，这里还有一个例子。叫做存储体冲突。这适用于共享内存。记住，共享内存就像L1缓存。它在SM上。硬件的工作方式是，共享内存被划分为32个存储体，每个存储体四个字节宽。好的，这些是我的32个存储体，这是我的内存。每个存储体在这里有很多元素。这里有个限制，每个时钟周期，每个存储体只能被一个线程访问，最多一个线程。假设访问的位置不同。例如，你不能同时访问这个位置和那个位置。这意味着如果你有多个线程试图访问同一个存储体，这些访问就必须被串行化。这就叫做存储体冲突。在最坏情况的例子中，想象这是一个按这种方式排列的矩阵。你有一些操作，涉及到32个线程，你会想，哇，32这个数字太棒了。

[15:36.340] 我可以大规模地并行处理它，然后你试图让这些线程都去访问第一列。结果它们只能排队等候。这就是一个32路存储体冲突，是你可能遇到的最糟糕的情况。现在你可能会说，好吧，那你为什么要那样做呢？当然啦，换成访问行不就行了。但这是不可避免的，因为例如，如果你在做矩阵乘法，你必须访问一个矩阵的行和另一个矩阵的列。而且有时你会进行转置，所以你并不能总能通过选择访问顺序来解决问题。对于逐元素操作，这没问题。你可以按任何顺序访问。但如果你在做矩阵乘法，你不能控制每个矩阵是行主序还是列主序。这里有一些解决方案，我就不深入讲了。有一种叫做swizzling的技术，它会重新排列你的共享内存，这样当你在访问时，就可以避免存储体冲突。

[16:37.120] 好的，所以这是另一个需要考虑的因素。当你进行性能分析时，你可以查看存储体冲突，查看占用率。你可以看到发生了什么。最后一点——实际上，还有两件事。一个是内存合并，你们已经讨论过了。所以我只是快速提醒大家这是关于什么的。当一个线程束中的32个线程试图访问HBM时，内存访问实际上会被合并成一个128字节的事务，也就是所谓的缓存行。它会一次性去获取所有这些数据。所以想象一下你的内存是这样排列的。这是32个。在最好的情况下，也就是所谓的完全合并，所有线程都在访问同一个缓存行。那么你有线程1访问M00，线程2访问M01，依此类推。这意味着你可以一次性抓取整个缓存行。

[17:49.230] 是这样的，而如果你去访问列，那么你会抓取大量的内存到这里，但这些内存你根本用不到。第二行也一样，以此类推。所以这感觉上跟存储体冲突很相似，但这是两个非常不同的限制。一个跟共享内存有关，另一个跟HBM有关。好的。好的，最后一件事，线程块占用率。记住，线程块是被调度到SM上执行的。但我们生活在一个有限的——逻辑上，你可以定义任意多的线程块。但物理上在芯片上，你只有一定数量的SM。例如，148个。所以如果你启动了160个线程块，那么你只能调度其中的148个。然后你必须等待它们完成，然后再调度剩下的12个。但当你调度这12个时会发生什么，就是有一堆SM只是闲着没事做。这就叫做低占用率。

[18:55.640] 当最后一批线程块的数量少于线程块总数最大值时就会发生。好的，所以总的来说，也许让线程块的数量能够整除SM的数量是个好主意。也许在这里总结一下，这个编程模型非常优雅，你有一个线程块网格，线程块中包含了个别的线程。在内存方面，HBM对所有人都是全局的。共享内存是线程块局部的。寄存器是线程局部的。但同样，硬件的所有细节——线程束、存储体冲突、内存合并、占用率——确实决定了最终的性能表现。而且这些细节中的很多都很难搞清楚，因为，我的意思是，性能分析工具能告诉你一堆信息。但你必须确切地知道有多少个 SM，以及所有东西的确切尺寸。有时候调度器会做一些你实际上无法控制的事情。

[20:14.060] 所以它比编程模型要稍微混乱一些。好的，我这里就停下来，等大家提问。嗯？好，一个快速的问题。那么 Cam，有没有一种场景，你可以让它们共享，或者——比如，在我们看到的那个问题里，有 140 个 SM [听不清] 块，对吧。而且我们必须用两种不同的方式启动。有没有办法让一个块共享一个 SM？嗯，所以问题是，一个块能不能共享一个 SM？我觉得问题是，如果你做的事情是对的，那么——我想这取决于块本身。如果你有一个块，比如说，它使用了 SM 上大部分的 Tensor Core，那么把另一个块放上去并不会真正加快速度。我认为从根本上说，这里存在一个参差不齐或不均匀的问题。因为块必须保持在一起。所以你不能把这里的负载分散到那边去。你能做的其实是改变你的块大小。也就是说，你改变块的数量，这样就不会出现尾部的空闲。

[21:36.040] 还有其他问题吗？好，那我们就继续。希望大家现在对 GPU 感觉更熟悉了一些。那么现在我要讲一下基准测试和性能分析。我可能实际上在内容方面不会讲太多。但我确实想强调一下这里的哲学，那就是，这是一个成功的秘诀。你对代码进行基准测试和性能分析。你做出修改，然后再对你的代码进行基准测试和性能分析。我之所以在教你们尝试和更早之前进行基准测试和性能分析，是因为你应该总是在编写内核之前，先去测量代码中正在发生的事情，找出瓶颈在哪里。好的，基准测试基本上就是事情花了多长时间。它只给你一个端到端的时间。它不会告诉你在哪里花费了时间。但尽管如此，它还是非常有用的，因为归根结底，那才是你真正关心的事情，即事情运行了多久。

[22:58.650] 而且它还给你，因为它把一切都浓缩成一个数字，你可以看到事物是如何扩展的，比如说，随维度变化。所以有一个很好的基准测试工具。但是因为这门课是从零开始的语言模型，所以我打算从零开始做这件事。但我这么做只是为了强调基准测试中的一些注意事项。假设我们有这个操作，矩阵乘法。那么 run_operation 就是这个包装器，它基本上会实例化两个大小为 dimension 乘 dimension 的随机方阵。然后它返回一个执行该操作的函数。好的，matmul，如果你调用这个，基本上就是执行那两个随机矩阵的矩阵乘法。好的，随机矩阵生成了，它就像把两者相乘一样。好的，那你如何进行基准测试呢？最天真的做法就是开始计时，运行它，然后停止计时。

[24:01.910] 但有一些需要注意的事情。一个是，永远记得要做预热。我想我在第二次课上提到过这个，但我认为值得强调。这是因为，如果有些东西是惰性编译的，你想确保那个东西——那个时间不会被算进去。因为大多数时候，你关心的是某个东西有多快，因为你打算一遍又一遍地运行它，所以初始条件并不真正重要。然后你通常需要多次计时，因为存在一些方差。正确的计时方法是使用这些 CUDA 事件——一个开始事件和一个结束事件，你可以在它们上面调用 record。实际进行计算，然后执行结束事件的 record。然后记得要同步，等待 CUDA 线程完成，因为如果你有异步操作，GPU 上的所有事情都在发生。这是一个同步屏障。然后你记录时间。

[25:12.360] 好的，然后你可以再来一次，然后再来一次。然后在这里，我们只是取平均值。如果你非常讲究，你可能想看看整个分布，比如 P95 或者什么。但我们在这里只做平均值。然后，使用基准测试你可以做的一件事是，比如说，扩大你的矩阵，看看时间如何变化。你可以看到，在这种情况下，矩阵乘法，正如预期的那样，它应该是以三次方的速度增长。但请注意，这里有一个缺陷，就是直到你达到接近 2000 维的矩阵之前，事情基本上是不变的。对吧，这是因为，正如我们讨论过的，这些 GPU 的形状是为相当大的矩阵乘法而构建的。如果你有一个 2 乘 2 的矩阵，那将会非常低效。好的。那么快速说一下性能分析。性能分析告诉你时间实际花在了哪里。希望你们所有人都熟悉性能分析并且正在做这件事。也许不太明显的是，性能分析，即使你不关心时间，也能帮助你弄清楚引擎盖下实际发生了什么。

[26:44.690] 因为特别是使用这些高级语言，你写一些代码，然后它运行了，你得到一个结果。有时候，了解一下实际正在发生什么是有好处的。所以 PyTorch 有一个内置的性能分析器。在你们的作业中，你们会使用 Nsight，它能提供更多细节。但为了节省时间，我们这里就跳过它了。那么让我们从，如果你在 PyTorch 里把两个数加一下，或者——抱歉，两个张量，开始。同样，run_operation 创建两个随机矩阵并应用该操作。所以在性能分析中，我正在预热。然后我把这个上下文——性能分析的上下文——放进去，然后执行 run。

[27:36.320] 然后让我们看看在 PyTorch 中仅仅做 a 加 b 的分析结果看起来是什么样的。通常在 PyTorch 中编程时，你可能不会去想它。就像，好吧，这两个张量直接就加起来了。那么在引擎盖下实际发生了什么呢？对吧，如果你看看它调用的东西，有这个很长的名字，kernel at CUDA functor add。所以这基本上就是一个把两个张量相加的内核。时间本身不会那么有趣，因为我只是在做加法。所以那会占用 100% 的时间。但这告诉你，嗯，在引擎盖下有一个叫做 add 的东西。那矩阵乘法呢？

[28:28.010] 在 PyTorch 中，我正在做 a 乘以 b。类似地，有一个很长的名字描述了这种特定的矩阵乘法内核。Cutlass f32, f32, 64x64x16，等等。请注意，如果你改变维度——那么现在我做一个 128 乘 128 的矩阵乘法——你会得到实际上不同的一个内核。如果你仔细看，这个是 64x64x16，这个是 32x32x16。好的，所以在引擎盖下，PyTorch 看起来你似乎只是在做乘法。但在引擎盖下，可能发生各种各样的事情。所以，观察。在这里，你可以看到哪些 CUDA 内核实际上被调用了。

[29:18.520] 这些通常是那些名字很长的。而且不同的 CUDA 内核会根据张量维度被调用。这个名字实际上也告诉了你一些关于实现的信息。所以这个名字。cutlass 是 NVIDIA 的线性代数 CUDA 库。SM100 对应 Blackwell 架构。所以这是一个专门为 Blackwell 设计的内核。这是 f32。然后 64x64x16 是块（tile）的形状，我们稍后会在讨论矩阵乘法上下文中的分块（tiling）时讲到。好的，这就是基准测试和分析（benchmarking profiling）。记住一定要做这件事。我想我们在作业里也让你做了，所以你别无选择。好的，那让我们把这种方法应用到另一个例子上，也就是 GeLU。回忆一下，GeLU 激活函数就是这个函数，它是一种常用的非线性函数。

[30:32.800] 通常，它会用这个 tanh 近似来逼近，这样对计算机更友好。好的，那么如果在 PyTorch 中天真地实现 GeLU，可以像这样写。你拿一个张量，直接把这个公式写进 PyTorch 就行了。这没问题，你会得到一些数值。PyTorch 也有内置版本。如果你调用 functional GeLU，也能得到那个版本。你可以检查它们是否一致。你可以运行这两个 GeLU 版本，确保在随机输入下结果相同。还有一件大家可能已经发现的事情。如果你还没遇到过，这是件值得了解的重要事情。那就是你可以对任意 PyTorch 函数调用 torch.compile，它会生成另一个函数，做完全相同的事情。

[31:47.810] 好的，所以这场比赛里我们有三位选手。我们有 naïve 实现（naive implementation），内置实现（built-in implementation），以及编译实现（compiled implementation）。我们可以对它们进行基准测试。Naïve 版本耗时三微秒，我猜这里是 3.75。内置版本快得多，编译版本也快得多，但可能没有内置版本那么快。好的，那这里发生了什么？为什么这些实现性能差异这么大？它们都计算出相同的结果，但性能特征却截然不同。这时我们可以调出分析器（profiler），看看底层到底发生了什么。如果你看这个 naïve GeLU，然后做个 profile，这个视图有点问题，所以我切换到另一个视图。

[32:51.560] 你可以看到 profile 显示了时间是如何被消耗的。这里有一堆不同的内核（kernel），比如 binary functor unary add，tanh 这里也是个 kernel。这对应于这样一个事实：在 PyTorch 中，当我们写出 PyTorch 表达式时，看计算图，计算图中的每个基本操作实际上都由一个 kernel 来实现。对。而它慢的原因在于，当你启动一个 kernel 时，它必须从 HBM 读取数据，一路拉到你的 SM 上，进行计算，再写回去。然后下一个 kernel 再从 HBM 里取出来，又写回去，如此反复。

[33:43.180] 所以你在进行大量来回的读写操作。因为在两次 kernel 调用之间，数据必须回到 HBM。再看看内置版本在做什么，其实没什么特别的。它有一个 GeLU CUDA kernel 实现，就是一个单独实现 GeLU 的 kernel。那这又是为什么呢？因为人们都在用 GeLU。所以有人为它写了个 kernel，放进了标准库里。我是说，这没什么神奇的。而编译版本这里真的很有意思。我不打算过多展开它到底怎么工作的，但我确实认为这是个很迷人的话题——你可以拿一个看似普通的实现，还记得吗，它在 PyTorch 里，有一个计算图。然后跑一个编译器。基本上，如果你看它的底层，它其实只是一个单独的 kernel。这是因为编译器推断出要去看这个计算图，然后用 Triton 把这个 kernel 写出来。

[34:59.142] 所以你可以看到这实际上是一个 Triton kernel。好，naïve 实现需要多个 kernel，意味着多次读写 HBM，没有任何 kernel 融合，所以慢。内置编译版本只有一个 kernel，基本上 GeLU 里的所有操作都被融合到一个 kernel 里了。所以你每个元素只从 HBM 读一次，写一次。而且你可以看到编译出来的 kernel 是 Triton kernel。好的，这也许是个不错的过渡，来聊聊这个 Triton 到底是怎么回事。嗯？[听不清] 让我想想。我觉得既然它上面写着 CUDA kernel implementation，那应该是有人用 CUDA 写的。

[36:03.236] 好的，还有其他问题吗？嗯？为什么 Triton kernel 比这个快？为什么 Triton kernel 更快？其实在这个例子里，Triton kernel 并不更快。哦，抱歉。编译版本是一个 Triton kernel，而它比内置版本要慢。对。我记得去年我讲这个例子的时候，差距其实更接近一些。但这些东西总是在变，而且非常依赖硬件。而且我觉得这些代码都没有做极致的优化。这只是为了给你一个总体的概念。好，那我们就来写一些 Triton kernel 吧。先回顾一下我们的编程模型。你有一堆线程，组织成线程组、线程块（thread blocks），上面还有一个线程块网格。如果你用 CUDA 来写——CUDA 最初是 NVIDIA 开发的，多年来它一直是编写 kernel 的标准方式——那么思维模型是：每个线程做什么？

[37:22.190] 你写一段代码，本质上是有一个 ID 来标识你指的是哪个线程，然后它就执行这段代码。这样做的好处是它和底层实际发生的事情非常接近。而且它是细粒度的，给你非常精细的控制。但这也有缺点，那就是记住，所有这些线程都在一个线程块里。有些操作需要它们之间进行通信。那会发生什么呢？它们必须同步（synchronize）。它们同时从 HBM 读取数据。然后必须同步，进行计算。你得自己处理这些簿记工作。所以如果你做的是全逐元素操作，CUDA 完全没问题，根本无所谓。但随着操作越来越复杂，你就开始需要一些抽象带来的价值了。

[38:32.140] 所以 Triton 是 OpenAI 开发的。我觉得到现在它已经相当标准了。你基本上指定每个线程块做什么。对于这个课程来说，尤其是你刚开始的时候，它通常已经足够强大了。如果你真的想去挖掘最新硬件的每一个新特性，它可能不会给你完全充分的灵活性。但我们现在不用操心这个。在 Triton 里的概念框架是，你思考的是：一个块（block）做什么？一个块会把数据加载到共享内存，对它进行操作，再写回全局内存。所以在某种意义上，这些块是介于思考单个元素在做什么以及思考整体操作之间的一个中间点。而在 PyTorch 中，你基本上定义这些巨大的矩阵，然后说把它们乘起来。这就是原子操作。你大量思考的是：我怎样才能把各种计算组织成大型的矩阵乘法？

[39:34.070] 对的，在某种程度上，Triton 是那种大型矩阵乘法和逐元素操作之间的混合体，我们很快就会看到。好，那我们从 GeLU 这个例子开始吧。我们定义一个 8,000 维的向量，开始写一些 Triton 代码。好，Triton 基本上就是你在写 Python。你们有多少人以前写过 Triton？举手给我看看。好的好的。不太多，这其实挺好，因为这样你们就不会觉得无聊了。好，这是普通的 PyTorch 代码。对吧，这里面没有 Triton。我只是在准备。那我接下来要做的是，拿这个张量，然后分配一个输出张量。

[40:35.890] 好的，因为在 Triton 里，我们不再以函数式思维来思考了。我们想的只是搬移数据。你必须显式地读写。所以没有返回值。我要分配一个输出张量，然后这个 kernel 会往里写。这个张量理论上可以任意大，对吧。通常我没法把整个张量都放进一个SM里，因为它太大了。所以我要把它拆分成块。接下来我要做的是，你可以把这个x想象成这样一个数组，我要把它切分成多个块。元素总数是8000。我现在先把块大小设为1024。然后我就有八个块。好，接下来我会用这种有点特殊的语法来调用内核。这是Triton GeLU内核。方括号里的内容基本上告诉了我网格的形状。

[41:54.630] 基本上，这表示网格有num blocks个块。然后对于每一个块，我会调用这个函数，也就是Triton GeLU内核，稍后我会详细讲。我会传入x、y、元素总数和块大小。好，很遗憾我没法一步步完整追踪这个过程，所以我就直接给你们看这里的代码长什么样。让我把这个删掉。好，这大概是你能想到的最简单的内核了。当你看到这个triton gelu内核时，注意前面我们用的是x和y，现在这些是指针。你可以把它们理解为整数，就是地址。你需要习惯这一点。然后我们有元素总数和块大小，实际上就是从这边传进来的。好，理解方式是这样的：对于每一个块，我们都会调用这个函数。

[43:20.240] 实际发生的过程是，我做的第一件事是“醒来”。这个块醒来后问自己：我是谁？程序ID——基本上PID就是用来标识块的。所以对于块1、2、3，PID分别是0、1、2。好，接下来我得弄清楚我要处理哪些数据。那就是PID乘以块大小。这是对x指针的偏移量。如果PID是0，那么起始位置是0。如果PID是1，起始位置是块大小。如果PID是2，那就是两倍的块大小，以此类推。好，现在我要确定我操作的区间范围。offsets就是起始位置加上……这是Triton库里的函数。一个从0到块大小的范围，概念上给你整数0到块大小减1。所以offsets基本上就是——比如说块1，它会是从块大小、块大小加1，一直到两倍块大小减1。

[44:39.700] 好，在这个例子中，元素总数恰好能被块大小整除。但一般情况下并不总是这样。所以你会经常在训练代码中看到这种掩码操作。它的意思是，有时候如果张量只到这里为止，那么我就会形成一个掩码，在这个位置之前都是真，之后都是假。对于不是最后一个块的块来说，掩码就全是一。好，现在我已经完成了准备工作。接下来我要做的是读取数据。这基本上就是指针运算。记住，x指针是一个整数，它指定了x所在的内存位置。我要把offsets加到这个指针上，这样就得到了前块大小个元素。根据掩码，如果我被遮住了，那我就不读取。然后现在你可以把这个看作一个向量。接着你做正常的计算。

[45:49.040] 然后得到y。y和x大小相同。然后你执行tl store，把y指针加上offsets，存储y和掩码。好，这个load从HBM加载数据，做了一些计算，然后写回HBM。好，我先讲到这里，看看大家对这第一个Triton内核有什么问题。在后面的同学？[听不清] 对。区别在于，和CUDA的区别是什么？对于这种逐元素操作，看起来几乎一样，对吧。事实上，我觉得CUDA甚至更简单，因为它确实是逐元素的。你醒来，识别自己的线程。然后你只操作那个元素。这里唯一的区别是这是向量化版本，你有一个块，你操作的是整个块。稍后我们会看到，当你要做的不只是逐元素操作时，操作块会比CUDA麻烦得多。

[47:14.530] 有几个问题。后面的同学。如果你想用张量核心，这要怎么关联？如果你想使用张量核心，这跟它有什么关系？我稍后会给你们看这段代码实际编译成什么。也许我会回到这个问题。但简短的回答是，你不控制那个。硬件自己决定数据放哪里。还有什么问题？你能不能实际走一遍从HBM到共享内存再到寄存器的tf.load过程？数据实际流动的每一步是怎样的？好，问题是，当这条语句执行时，实际发生了什么？简单来说，这在某种意义上是个“谎言”。对吧，并不是说GPU真的在GPU上调用Triton库并执行这段代码。这基本上是为了让我们理解而用来描述计算的。编译器接收这段代码，我们稍后会看到，把它写成叫做ptx的东西，然后它才真正干活。

[48:32.990] 这就是机制上发生的事情，我稍后会讲。如果你问的是概念上的问题，理解方式是这样的：x指针是HBM中的一个内存位置。这基本上指定了一段内存地址范围。load获取这些内存位置并返回与之关联的数据。这里我有一个局部变量叫x。在实践中，这通常是一个寄存器或共享内存。Triton自己决定放在哪里。[听不清]——在线程级别。当这个被调用时，其实不太清楚。我是不是在之前就已经安排好哪些数据会放到共享内存、哪些放到寄存器，还是说现在才发生，那就太晚了？那时候所有线程就都空闲在那儿等着内存数据灌进来。我敢肯定它是准备好执行了，所以我们不会干坐在那儿，无所事事地等着数据从HBM进来。对，对。问题是，这到底什么时候真正执行？这不会阻塞吗？

[50:01.370] 让我试着在给你们看ptx的时候回到这个问题，希望到时候能提供更多上下文来解释实际发生了什么。好，还有别的问题吗？好，所有内核看起来都差不多是这样。我确实想确保大家都理解这个通用的形式，那就是通常你有输入、输出。你醒来，弄清楚你要处理哪个索引。你读取。做一些计算，然后写回HBM。好，让我们简单聊聊ptx。让我看看，这个能不能行？这个链接能打开吗？好，当你写Triton时，编译器会生成ptx代码，这是GPU的中间汇编语言。我显然不会逐一讲所有这些，但为了让你们对这个有个感觉。让我先从一个观察说起。有几件事。这现在是单个线程实际在做什么。不是线程块，因为那已经被编译掉了。

[51:35.710] 如果你看——这里有几个注释。这个LD global基本上就是从HBM加载到一些寄存器。寄存器用r表示整数寄存器，fr表示浮点寄存器。然后你看到像mov 0到r5、mov 0到r6这样的语句。然后有乘法运算。你要把这个寄存器乘以这个常量，然后把结果放到这个寄存器里。所以这段代码被执行了。在最后，你应该能看到global store。这就是写回HBM。所以这段代码被执行了。这实际上就是在线程上真正执行的代码。而 Triton 基本上就是在这个之上再抽象一层。另外一件值得注意的事情是，你会看到所有这些块（block），它们是怎么回事呢。实际上，这里我提到了线程粗化（thread coarsening），也就是说这是一个线程（thread）。但它并不是处理单个元素，而是实际处理八个元素。

[52:59.630] 编译器决定说，嗯，其实这个线程做的事情很轻量，工作量不大，那我们就把它加粗一点吧。所以看一下 PTX 代码能让你对底层发生的事情有一些体会。嗯？所以它是为每个单独的——[听不清] 所以问题是，它是为每个线程都生成这个吗？这个代码只编译一次。每个线程运行的都是同一份代码。线程区分自己的方式在于，这份代码会传入基本上就是线程 ID 的东西。所以这里的 cta.x 是块索引（block index），而 tid.x 是线程索引（thread index）。所以基本上，如果这段代码在运行，它告诉你我在哪个块里。而 tid.x 告诉你我是那个块里的哪个线程。好的，还有其他问题吗？我在想怎么——好的。好，那这基本上就是你们第一个 Triton 内核的大致样子，从 HBM 加载数据，计算，再写回 HBM。

[54:42.920] 然后我们看了 PTX，那是在底层真正发生的、比较粗糙的细节。但 PTX 里仍然有很多东西是没有指定的。比如说，操作在哪个 SM 上执行，warp 怎么安排等等。很多这些东西都是由硬件控制的，所以你根本看不到。嗯？所以 PTX 代码是由编译器生成的，不是通常你自己会去写的？它看起来像是汇编代码。对，所以问题是，PTX 是不是只是编译器生成的？确实有些人会直接写 PTX，如果你们真的觉得自己比编译器更厉害的话。而且我觉得 NVIDIA 的编译器总体来说已经相当成熟了，但其他一些加速器可能发展得没那么完善。我觉得有时候你确实需要深入底层，进行更多的手动干预。但一般来说，是的，你不需要这样做。嗯？[听不清] 我的工作会被调度到某个 SM 上。

[55:59.050] 我会执行到那个 tl.load，这几乎就像一个 CPU 的陷阱调用（trap call），我只是在等待某些事情发生。然后我就闲下来了，其他工作就会被调度过来顶替我。然后它会执行操作。现在那个加载完成了，然后那个工作会被重新调度回来，我继续执行。是不是——嗯，没错。所以让我重复一下这个问题或者说这个评论。如果你看一下这个 Triton 代码，这些语句，这个加载（load），会阻塞若干个周期。然后这是在某个 SM 上的某个 warp 的某个线程上运行的。记住，SM 同时运行多个 warp。所以当执行到那个点的时候，它就可以找另一个 warp 来运行。然后当这个加载完成后，warp 调度器会回来接管。你为什么要[听不清]？抱歉。哦，抱歉。我记得从[听不清]那里。对，所以问题是，为什么是四个 warp 调度器？

[57:09.800] 我不太确定这背后的确切原因。好，现在我们来看一些其他的例子。也许就快速预览一下。我们还会做三个例子。GeLU 是最简单的形式。尽管这个计算本身挺复杂的，但它只是逐元素操作。所以从概念上讲，在本次讲座的语境下，它其实非常简单。现在我们来看 Softmax，这里我们就要做归约（reduction）了。但在这种情况下，我们先考虑每一行能放进一个块（block）里的情况。然后我们再考虑一行放不进一个块的情况，接着我们会一路做到矩阵乘法（matmul）。然后希望到了那个点，你们就已经具备了完成作业和实现 Flash Attention 所需的全部要素了。好的。到目前为止，我们看了逐元素操作。现在让我们想想那些会对多个值进行聚合的操作。

[58:16.200] 回忆一下 Softmax 是做什么的。我们就把它想象成一个矩阵。你对矩阵的每一行进行指数化和归一化。好的，这用于注意力机制中，也用于生成输出概率。总的来说是件好事。那么我们就从一个朴素的实现开始，并跟踪一下发生了什么。这里，我定义了一个张量。这里是朴素实现。实际上，我觉得这是作业一里的内容。好的，总之。这里我有一个 M 乘 N 的矩阵。所以对于每一行，我需要计算——我要计算每一行的最大值。这是为了数值稳定性。我会减去最大值。然后逐元素做指数运算。然后我求和，计算每一行的归一化常数。然后做除法。就这样。好的。如果你们数一下读写次数，记住，这只是普通的 PyTorch。

[59:47.180] 这是一个不同的内核。这也是一个不同的内核。除非你调用 torch.compile，不然这些会是不同的操作。每个操作都会从 HBM 读写、读写。好的，所以你有五次 MN 的读取，三次 MN 的写入。而原则上，你应该只需要少得多的次数。好的，这就是一段代码。能理解。大家都应该熟悉 Softmax 在做什么。好的。那让我们写下 Triton 内核。从某种意义上说，形式会非常相似。就像在 GeLU 里一样。对吧，一旦你搭好了脚手架，核心计算看起来就和朴素版本非常像了。所以我们打算做的是，我们说每一行就是一个块。好的。为什么我要把每一行做成一个块呢？因为，记住，每一行我得做归一化和求和。

[01:00:55.010] 所以它不是逐元素的。Softmax 不是逐元素的。但它是按行进行的。所以块之间不交互。而且块之间没有共享内存，所以这没问题。块之间没有共享内存。那么在每一行内部，我就做一些操作。好的，让我们看看会发生什么。同样，在 Python 里，我会分配我的输出张量。我有一个 M 乘 N 的矩阵。我会把块大小（blocksize）定义为基本上就是列数。为了保险起见，取下一个 2 的幂。然后块的数量就是行数。好的，然后我会调用这个内核。那有多少个块呢？M 个。对吧，每一行一个。对于每个块，我会传入输入指针、输出指针。然后我还会传入这些 stride，它们告诉我要往下移动多少。好的，让我实际给你们展示一下那个 Softmax 内核。

[01:02:23.860] 好的，那这是什么样子的呢？我启动了。我在某个特定的行上。然后这会给我从 0 到块大小（blocksize）的所有列，我假设这就是所有列。我要读取。所以基本上，我需要弄清楚从哪里读取内存，那就是我数据的起始位置。加上我在哪一行。而行 stride 基本上告诉我每一行——基本上，这就是列数。它会告诉我需要往下走多远。然后这个 x 指针基本上就是我需要加载的所有数据的地址。我把它们加载进来。然后在这里，我做了这个操作，如果它被掩码掉了，我就把它设为负无穷大，因为对于 Softmax 操作来说，这就相当于 0。然后这部分本质上和朴素的 Softmax 是一样的。我就是计算最大值，减去它，做指数运算，求和，然后除法。

[01:03:44.070] 然后我把它写回去。好的。嗯。是的。[听不清] 是的，所以这是一个每个块都可以覆盖整行的版本。你可以看到 Triton 让这件事变得非常容易，对吧。因为这就好像你在写普通的 PyTorch 代码一样。所以如果所有东西都能装进一个块里，你基本上——问题是，如果你能把它装进一个块里，你几乎就可以直接写普通的 PyTorch。嗯？如果我的列数、行数比块大小大怎么办？是的，我们马上会讲到这个。一般来说，列数、行数会远大于块大小。所以我们回头再讲这个，是的。所以如果你想对每一列做 Softmax，你是不是必须得……[听不清]？

[01:04:52.530] 问题是，如果你想按列做 Softmax。我觉得那应该没问题，因为在这里，我们在跟踪这些指针，对吧。指针可以是任何东西。基本上，我觉得你只需要改变这里的步长，让它去访问列就行了。实际上，应该是在这里。对于列偏移，你只需要把它乘以行步长，我觉得。好的，我们继续。好了，那么现在，热身到矩阵乘法。假设你的行无法装进一个块里。好的，那在这种情况下你怎么办？比如说，如果你有一行有 4,000 列，但块大小只有 1,024。所以你必须在这里做点什么。策略是这样的。我们把这行拆分成多个块片。在这种情况下，会有四个块片。每个线程会遍历这些块片并累加一个总和。最后，我们会通过把每个线程产生的结果全部加起来来做归约。

[01:06:17.530] 我给你们看一个这样的例子。现在我从 Softmax 切换到行求和，因为这样更容易思考。好的，那么在这里，好吧，这个没什么意思。内置的行求和功能会做你所期望的事情。所以行求和操作基本上就是接收一个矩阵，然后计算每一行的总和，好的。所以这里没什么让人意外的。好的，那么我们在这里要做的事情，概念上如下。好的，每个块仍然负责一行。这部分没有改变。所以假设块一，行一。我醒过来，我做什么？我要记住，现在，有块片。假设块片 0 是第 0 到 3 列，块片 1 是第 4 到 7 列。然后块片 2 是第 8 到 11 列。好的，那么我在这里要做的是，我现在要开始迭代了。所以首先，我基本上处理这个，所有线程。基本上，每个线程保持一个累加器并处理第一个块片。

[01:07:32.040] 然后它继续到下一个块片，并把当前元素加到累加器里。所以这里，我放的是 3，1，4，1。然后在第二次循环迭代中，我要把 5 加到这上面。我得到 8。我把 9 加到这个上面，得到 10。2 加到 4 上，6 加到 1 上，对吧。所以这四个线程中的每一个都在累加各自的东西。然后在最后，我再做一次循环，针对块片 2。我把 5 和 3 加到各自的累加器上。所以在一天结束时，我有了这个，一个累加器向量。然后那实际上就是——那个，我直接把它加起来就行。好的，让我们看看代码是什么样的。我就直接调用行求和内核。让我们看看。好的，这就是它的样子。醒过来。我在一个特定的行上。这就是一行看起来的样子。有一个块片，第二个块片，第三个块片，依此类推。

[01:08:46.399] 所以 n 是那一行的元素数量。而 blocksize 是这个块片的大小。好的。记住，blocksize 是线程数。我要处理的数据比块大小要大。我现在只需要迭代地处理它。所以我要遍历所有的块片。所以 start 会从 0 到 blocksize 到 2 倍的 blocksize，以此类推。所以每次我在跨越时，我都在进行调用。基本上是获取特定块片的偏移量。然后我要从 HBM 加载数据，并把它加到累加值上。这个累加值要么在寄存器里，要么在共享内存里。好的。然后最后，在我遍历完所有块片后，我处理完整行，我基本上得到——对于每个线程，我都有一个累加器，记录那个线程获取了什么。我就可以直接做一个求和得到一个标量，然后把它写出去。

[01:10:05.865] 好的，这比之前稍微复杂一点，因为现在我们在一个线程内部有一个 for 循环。当你的数据无法装进一个块时，这是必要的。嗯？[听不清] 以防我们没有足够的寄存器。我们必须把它放在——是的，所以问题是，我觉得如果你能控制累加器驻留在哪里。至少在这个 Triton 程序中，你不会显式地说明。这取决于 Triton 编译器来决定把它放在哪里。但一般来说，如果块大小足够大，它就必须放到共享内存里。好的。好了，这讲得通吗？我想，为了确保大家都能跟上。记得在 GeLU 中，我们也把一行拆分成很多块。但那些是块。对吧，所以那些块中的每一个都是独立处理的块。这些不是块。这些是块片。块对应的是整行。而块基本上必须处理所有的块片。

[01:11:28.560] 所以这就是它开始看起来不像 PyTorch 的地方，因为你不能以一种漂亮的——不是所有东西都能装进共享内存。嗯？[听不清] ——也存储在共享内存里？是的。让我们假设累加器存储在共享内存中。是的。好的。好了，那我们现在也许可以进入我们的压轴戏了，也就是矩阵乘法。好的，矩阵乘法，当然是深度学习的命根子。它已经被优化到了极致。而且在某种意义上，它是一个非常基础的操作。你拿两个矩阵，把它们相乘。我在这里要加一点小变化。我要做一个矩阵乘法后面接一个 ReLU，就只是为了好玩。好的。嗯，这种情况确实会发生，对吧。因为如果你有一个线性层，那就是一个矩阵乘法。然后你应用一个 ReLU 激活。所以这不是完全没有来由的。

[01:12:37.750] 但我稍后会展示我为什么这么做。好的，那么你怎么构建一个矩阵乘法内核呢？好了，这里有一个朴素的方法。好的，那么这是我的，比如说，一个矩阵。我用 A 乘以 B。我试图写出矩阵 C。A 是 M 乘 K，B 是 K 乘 N。所以 C 是 M 乘 N。好的，那么我要做的是，固定这些元素中的一个。假设 M 等于 1，N 等于 2。好的，所以我在处理——实际上，我们做 M 等于 1，N 等于 1。所以我在计算 C5。然后，基本上对于每一个 k，所以我要，是的，我要遍历 k，遍历 A 的行和 B 的列。我要从 HBM 读取。

[01:13:37.040] 我要把它们相乘，累加起来。最后，我要写到这个和这个元素。所以那是一个有效的矩阵乘法内核。好的。那它有什么问题？它是正确的。但如果你看看它做了多少次读写，这不太好。对吧，基本上对于每一个 M、N 和 K，我都得从 HBM 读取。所以读取次数是 M 乘 K 乘 N 的数量级。写入次数。我的意思是，这无所谓，但这是瓶颈，对吧。如果你记得第二讲，如果你看你做的操作数除以传输的字节数，那就是算术强度，你希望它高。所以操作数是 M 乘 K 乘 N 数量级。然后读取次数也是同样的数量级。所以算术强度是个常数，这不太好。好的，所以如果你仔细看，你会发现有很多冗余的读取。

[01:14:49.310] 那么想象一下计算 C4。你需要读取 A4、A5 和 A6。而如果你计算 C5，你还需要再次读取这些内容。所以如果你能只读取一次，就能真正节省读取次数，这非常好。那么让我们尝试用共享内存来实现这一点。这里是理想化的方法。我要把 A 和 B 全部加载到共享内存中。然后我就直接计算 C。如果我能做到这一点，那么——现在我不再有这个三次方数量的读取，我只有二次方数量的读取，这意味着我得到的算术强度是 n 量级的，在第二讲中，我说过这是一种你能期望的理想情况。好的，所以如果你能做到这一点，那就太好了，因为基本上——没有冗余读取。你把所有内容读取一次到共享内存中。你做计算，然后写回去。好的，但这种理想化方法的问题是什么？问题是 A 和 B 通常太大，无法放入共享内存。

[01:16:09.166] 对吧。那么，你现在需要做什么呢？好的，这里的思路就是非常经典的分块（tiling）思想。这个想法基本上就是，尽可能多地把你能够放进去的内容放入共享内存。所以从某种意义上说，它看起来像朴素方法——从全局来看它看起来像朴素方法，但从局部来看它看起来像理想化方法，这是理解它的方式。这里是示意图。我想 Tatsu 也展示过这个。那么我们要做的是，我们取这个矩阵 C，然后不再——记住，朴素方法只是说对于每个元素，我要计算它。但现在我要说对于每个分块（tile），我要做这件事。所以我把它分解成多个分块。每个这样的分块将是一个线程块（thread block）。我将有一组线程负责计算这个。另一个不同的分块将由另一个线程块完全独立地计算。

[01:17:14.090] 好的，所以我要——想象我在 Triton 中。我醒来，我看着这个线程块。那么我需要做什么？嗯，从某种意义上说，这和朴素方法是一样的，对于 A 的每一行分块，我将扫描各行。对于 B 的每一列分块，我将加载对应的 A 分块和对应的——比如说我在这里——以及从 B 中对应的分块到共享内存中。我要将这两个相乘。这就像理想化方法一样。我要把它累加到部分和中。而这一切都放在共享内存中。然后在我完成这里所有行和列的扫描之后，我终于可以将这个输出分块写入 HBM。好的，这就是概念上发生的事情。那么这里的算术强度现在上升到分块大小的量级。所以你通常无法达到 n 量级，因为这要求你把所有内容都放入共享内存。但如果你的分块很大，那也还不错。

[01:18:36.860] 好的，作为一个额外的好处。当你反正要为这个编写内核时，有时候如果你想应用一个逐元素激活函数，很简单，只需在最后加上它。这就是内核融合（kernel fusion）。好的，很快地看一下这里的实现。那么——哦。好的，先提醒一下关于步幅（strides）的内容，因为这会涉及到。张量是一个多维数组。但在内存中，它是线性化的。张量的步幅告诉你基本上如何将多维索引（比如行、列）映射到实际索引。基本上你要做的就是将行乘以行步幅，再加上列乘以列步幅。好的，那么在这种情况下，每次你前进到下一行，你在内存中移动 4 个位置。每次你前进一列，你移动 1 个。

[01:19:42.600] 如果是转置的，那就是翻转的。好的，那么这个内核看起来是什么样的？好的，启动部分没什么意思。所以你醒来后你在分块 M 和 N 上。就是说，好的，我负责计算 c 矩阵，但针对 M、N 分块，有一堆索引操作，你可以——我就简单带过了，但它很直接，只是有一点复杂。你只需要跟踪索引。所以这基本上告诉你我在看矩阵 A 的哪些行。哪些行——哦抱歉，我在看 B 的哪些列。然后这就是从 1 到 K 的数字。然后我要获取 A 和 B 在我分块位置的指针。然后我要设置这个累加器矩阵。它将在共享内存中。它是 M 乘 N 的。然后这将看起来像行归约（row reduction），对吧。你有这个对所有分块的求和。但不是我仅仅横跨分块，我现在横跨 A 的行分块，同时也向下遍历 B 的列分块。我加载 A，这个小矩阵，我加载 B，这个小的 B

[01:21:13.820] 分块。然后我执行这个点积。所以记住，当东西在共享内存中时，事情看起来像 PyTorch，我可以直接说 matmul，它就会执行操作。然后我前进到 A 的下一个行分块和 B 的下一个列分块。额外的好处是，如果我想应用逐元素的非线性，我不如在这里做。在我写到 HBM 之前，我可以对它做任何操作。然后最后，我直接写出去。好的，所以有一些索引你需要留意。但希望算法的形式是清晰的。好的，关于这个有什么问题吗？好吧。也许我做个总结，你们可以之后再问我。那么今天，我们讨论了有一个编程模型，就是说你讨论的是 PyTorch 或者 Triton 或者 ptx。这是程序员可以控制的部分。对吧，甚至 ptx，你可以编写 ptx，你可以控制每一个——想要多特殊化就多特殊化。

[01:22:38.560] 但这不是全貌，因为现实是你的代码必须在硬件上运行，而且只有有限数量的 SM，有限数量的存储体（banks），内存和寄存器的大小也是有限的。所有都是有限的。所以你带着你巨大的矩阵和 transformer 进来，想要——它必须适应硬件的约束。所以这就是为什么基准测试（benchmarking）和性能分析（profiling）非常重要，用来理解硬件的复杂性如何转化为性能。我们讨论了 Triton，我觉得这是一种相当不错的、简洁的语言来思考线程块。希望到现在你们能体会到，思考线程块比思考单个线程更容易，因为你不需要显式地考虑线程同步或共享内存。思考方式是你要弄清楚——你有你的计算。你把它分解成线程块，你只需要从共享内存中读取。做一些操作，然后写回 HBM。

[01:23:50.380] 然后我们看到了一些难度递增的示例。逐元素是最容易的。然后是跨行归约。还有不适合单行的归约。我们引入了初步的分块。然后矩阵乘法是真正做分块的经典示例。好的，关于如何编程单个 GPU，这就是我要说的全部内容。下次，我们将转向更多 GPU，讨论多 GPU 编程。好的，有问题。所以如果我想编写内核，有哪些选择？我是否需要尝试它们——对于我的每个选择，我能多接近最优水平？所以问题是，Triton 的替代方案是什么？这里有一个权衡——每种语言都有归纳偏置，对吧。这使得某些事情更容易，某些事情更难。所以我们做的大部分工作——Triton 是由训练 transformer 的人构建的。

[01:24:56.350] 那么，凡是涉及 transformer 的部分，我觉得相对都会比较容易。当然，极端情况下，你总是可以下探到 ptx 层面去手写。但我不会建议把这作为第一步。这里还有不少其他库，比如 thunder kittens。还有 cute、delirious 这类 DSL，它们能给你提供不同层级的抽象——不一定是在整个技术栈上简单地上移或下移。它们只是带来不同的特性。嗯？你提到的是高维张量放大，还是高维张量处理？最好的方法是什么？我在想的是，我能不能把我整个传感器数据都加载出来，把这两个传感器的数据同时放到一大堆线程或线程块上让 GPU 并行计算？还是说像我们刚才这里做的，每次取一个元素会更好，对吧。

[01:26:06.210] 比如我传感器里的某一个具体分量，处理完再最终写回去。嗯，简单来说，这个问题——我们可能也得收尾了——就是，是一次性全部读入好，还是说一次处理单个元素更好？我觉得脱离具体场景很难给这个一个确切的答案。这完全取决于计算的性质。也许我们之后可以私下再聊。好的，那今天就到这里，下次见。

CS336 2026 · Lecture 6 · Kernels, Triton

视频：官方视频
字幕：en-US · manual captions
处理：连续字幕合并为自然段；每段保留首条字幕的起始时间戳
状态：原语言字幕稿，待校对

正文

[00:05.240] So welcome back, everyone. So on Monday, Tatsu did an excellent job of describing a high level overview of GPUs and how to think about performance and all the quirks that come with GPUs. This lecture is going to be a continuation of that, where we're going to dive more deeply into the code, write some Triton kernels, and as well as do some benchmarking and profiling. So just to start, just to refresh what's going on with GPUs, here is a simplified diagram of what a typical GPU looks like. There's memory, and then the actual GPU chip. And just to keep in your head what are the characteristics of a GPU. So of course, we're talking about NVIDIA GPUs, but Tatsu also talked about CPUs, and there's AMD and so on. But focusing on NVIDIA GPUs. We've had multiple generations, from A100s, to H100s, to B200s. And in each generation, you have a GPU. It has a number of Streaming Multiprocessors, or SMs. And the number of SMs is about 100, between 100 and 200. So that hasn't really changed that much.

[01:30.250] And then within the individual SM, there's a set of registers. So B200s have 65,000 registers for a total of 256K per SM. And this also hasn't changed that much. In addition, there's a L1 cache as well as shared memory. So remember, these are the same memory, shared memory you can control. L1, you can't. This is per SM. And this is in the same order ballpark of size. And then there's L2 cache. And this is not per SM, this is for the whole chip. And this is a bit larger. And then finally, you have a High Bandwidth Memory, or HBM, which is large. And you see that this is the number that's actually going up quite a bit. OK. And then in addition to size, you can think about what's happening with the bandwidth. And essentially, it's inversely correlated. So registers are very fast, L1 is slightly less fast, L2 is less fast, and high bandwidth memory is the slowest.

[02:50.850] Although eight terabytes a second is still not that slow in the grand scheme of things. So this is the main hierarchy you should have in your head. Large memory is slow and far but big. And fast memory like registers and L1 resides on the SM. It's local and it's fast, but small. OK, that's the environment we're dealing with. OK, so then, how do you program a GPU? So the program model is as follows. So there's notions of threads, which each thread executes a piece of code on a small part of the data. So each one of these arrows, you can think about it as a thread. These threads are organized into thread blocks, also known as Concurrent Thread Arrays, or CTAs. And this is a group of threads. And then finally, you have a collection of thread blocks which forms the grid. So when you launch a kernel, you're basically launching a grid of threads and thread blocks to all parallel simultaneously do some computation.

[03:59.690] OK, so this is simplified. There's some other things that are happening. For example, there's H100s and B200s also have thread block clusters, which are clusters of thread blocks that enable some amount of distribution memory. Also, B200s have tensor memory. Now for tensor cores, which is somewhere between registers and shared memory. Some of these things are invisible to the programmer, but they are in the hardware. But we won't worry about them for this intro lecture. OK, so one thing you might ask about is, why do we have thread blocks? Why can't we just have a grid of threads, and each thread just takes a piece of data and does something? And this would be normally fine if all you want to do is element-wise operations, which we'll see. For example, GeLU is an activation function that applies element-wise.

[04:54.460] And threads are pretty natural. Each thread processes one element. So it's like Fi for i equals ranging over your data set. But for operations that involve communicating with threads, such as softmax or matrix multiplication, this view isn't really enough. And the reason is that, well, it could be enough if you were willing to pay the cost of writing and reading HBM. Because then you can still have-- for example, in a matrix we'll see later, you can just have every element in a cell. It goes and computes that element of the matrix. And it can just read and write HBM. But as we've noticed, HBM is very slow, so this would not be a good strategy. So instead, what we should do is use shared memory, which is local to SM. And what the thread block allows you to think about is a collection of threads that are all going to access this shared memory.

[06:05.080] So consequently, you can think about a thread block, one of these thread blocks is being scheduled on one of these semantic-- not semantic, streaming multiprocessors. And it does its thing. And what it's going to do is it's going to read a bunch of data from HBM and then process it where the processing might involve communication between the threads via the shared memory, and then writes it back out. OK. So this is a critical piece that comes up. Later, we'll talk about tiling, and that's the whole game here. OK, in Triton, in fact, which we'll see as a primary way that we're going to write kernels these days, we're going to think natively in thread blocks. And I think once you get the hang of thinking of thread blocks, this makes a lot of sense and makes your life a lot easier, as we'll see. OK, so the programming model is fairly simple, I think. There's threads, thread blocks, and grids. I think things get a bit more complicated, is the interaction between the programming model and the hardware.

[07:23.210] So the programming model is actually very nice. It provides abstraction of the hardware. When you can write a kernel, all you have to know is that there's a bunch of thread blocks. You define them and then define what computations each thread within the block or the thread block has to do. And those are just writing-- it's like writing Python. So that part isn't hard. And in fact, if you don't care about-- if you just care about correctness, that's all you need to know. But in practice, the performance is very sensitive to the hardware. And so you need to really deeply understand the hardware to obtain high performance. And in fact, the whole reason we're talking about GPUs and kernels is that you're trying to squeeze out performance. So there's these two levels here where you're trying to understand the computation, and that is only a part of the programming model. But how fast it runs is going to be strongly dependent on the hardware.

[08:26.140] OK, so I'm going to give you some examples of why this matters, and some of this will be review from what Tatsu talked about. But hopefully it will reinforce some of these concepts. So I think there will be around five different examples. And this is just to give you a flavor of the considerations. So there's something called warps, which I didn't talk about. So in the very simplified view, warps are not really part of this clean picture of the programming model. You have threads and thread blocks or grids. Technically, you can also exceed the warps, but you don't have to when you're just programming. So what is a warp? So each thread blocks, remember, is a collection of threads. But these threads are actually grouped into warps. And there's basically 32 threads per warp. So for example, if you have 64 threads in a thread block, there's two warps, the first 32 and the second 32.

[09:29.840] And as Tatsu mentioned last time, all threads within a warp, all of these T's, must execute the same instruction in lockstep on SM. So every cycle, they have to execute exactly the same instruction. And control divergence is when different threads within a warp need to actually execute different instructions. For example, if you have branching, if something, then A else B. And then what happens is that you can only do the A's in your threads in your warp, and then you do the B. So this gets sequential. So this is bad and inefficient, and that's why branching is something you want to generally avoid. One thing that's, I think, really cool about warps is that a SM actually is running multiple warps. There's a warp scheduler, and there's a bunch of resident warps that are about to run. Each warp has some-- threads have registers. And the SM can switch between them with zero cost.

[10:40.610] Right, this is not true in general, like on a CPU. But this way that it's designed is to hide latency. So this is important, because one of the warps is, for example, reading from HBM, which remember is very expensive, can take 100 cycles or something. You don't want to just sit around waiting for that warp to do nothing. You switch immediately to another warp where it can actually do some tensor core operations. OK, so that's one thing to note about the warp. And another reason why warp comes up is this idea of occupancy. More specifically, warp occupancy. So each thread, just the hardware constraint says that each thread can use at most 255 registers. OK, and so what happens is that the SM has a fixed number of registers. So the more registers each thread is using, the fewer threads you can have. That's just math. And that can reduce your occupancy. But that's not necessarily bad, because if you have fewer threads.

[11:54.790] But each thread is doing more work, that can actually be good. So occupancy is something that you can measure, but it's not necessarily the larger, the better. Because there are some other tradeoffs here. And in fact, an example where you might want to have fewer threads is this idea of thread coarsening, where let's say you have an element-wise operation and you can have one thread just perform work on one element. OK, so then you have a lot of threads. But you can also say each thread processes multiple elements. Like a constant, maybe eight. So that gives you fewer threads, which makes scheduling and these things easier, but each thread is doing more work. So if your threads are very light, maybe you want to fatten them up a little bit. OK. So just as an example of what might happen here. So suppose you have a block, a thread block with 128 threads in it. And each thread is using 160 registers. OK. So there's a hardware constraint.

[13:06.270] So the B200s have maximum of registers of 65,000 per SM. And it also has a constraint that says you can't run more than 64 warps at a time. So then you can go and do some simple math to compute what is the occupancy here. So number of registers per thread. This has to be less than 255. That's fine. The number of registers that you're using per block is the number of threads per block times the number of registers per thread. That's about 20,000. That means you can run at most three blocks on your SM concurrently. And then that corresponds to 12 warps. And because the number of maximum warps is 64, that means you have an occupancy of 18%. So we're only running 18% the total number of warps you have. And this is because you have a lot of registers use per thread. OK, so this is an example of where I guess this is memory, is constraining you in terms of how much compute you can do.

[14:21.420] OK, here's another example. Something called bank conflicts. This applies to shared memory. So remember, shared memory is like L1. It's on an SM. And just the way the hardware works is that shared memory is divided into 32 banks, each one four bytes wide. OK, so these are my 32 banks, and this is my memory here. Each bank has many elements here. So there's a constraint that says every clock cycle, each bank can only be accessed by one thread, at most one thread. Assuming that's not a location. So for example, if you can't access this location and this location at the same time. Which means that if you have multiple threads that are trying to access the same bank, the accesses have to be serialized. And this is called a bank conflict. In the worst case example, imagine this were a matrix that was laid out like this. And you had some operation where 32 threads, you think, wow, this 32 is so great.

[15:36.340] I can massively parallelize this, and you try to all access this first column. They're just going to just wait in line. And this is a 32-way bank conflict, which is the worst possible setting you can be in. Now you could say, well, OK, why did you do that? Just like access rows, of course. But this is unavoidable, because, for example, if you're doing a matmul, you have to access rows of one matrix and columns of another matrix. And sometimes you do transposes, so you can't always just get away with choosing the order in which you go down. For element-wise operations, it's fine. You can go in any order. But if you're doing matmuls, you can't control for every matrix which row major or column major you're doing. There are some solutions here, which I won't get into. Something called swizzling, which arranges your shared memory so that when you're going through, you can avoid bank conflicts.

[16:37.120] OK, so that's another consideration. And when you're profiling, you can look at the bank conflicts, you can look at occupancy. And you can see what's happening. A final note which-- actually, two more things. So memory coalescing, which you talked about. So I'll just really quickly remind people what this is about. So when you have 32 threads in a warp and they try to access HBM, the memory accesses actually get combined into a transaction of 128 bytes, which are called cache lines. And it goes and fetches it all at once. So imagine you have memory laid out like this. This is 32. And in the best case, which is called full coalescing, all the threads are accessing the same cache line. So you have thread 1 accessing M00, thread 2 accessing M01, and so on. Which means that all at once, you're going to grab this entire cache line.

[17:49.230] Right, and whereas if you go down the columns, then you're going to fetch a lot of memory here, which you're not going to use. Same for the second row and so on. And so this feels similar to bank conflicts, but it's a very different constraint. And this is about shared memory, and this is about HBM. OK. OK, so final thing, block occupancy. So thread blocks are, remember, scheduled onto SMS. But we live in a finite-- logically, you can define as many thread blocks as you want. But physically on chip, you only have a certain number of SMs. For example, 148. So if you launch 160 thread blocks, then you can only schedule 148 of them. And then you have to wait until they're done, and then you schedule the 12. But what happens if you schedule 12 is that there's a bunch of SMs are just not doing anything. So that's called low occupancy.

[18:55.640] When the last wave of thread blocks gets fewer than the total maximum number of thread blocks. OK, so in general, maybe it's a good idea to make the number of thread blocks divide the SMs. So maybe just to summarize here, there is a very elegant programming model where you have a grid of thread blocks and the thread blocks have individual threads in them. And in terms of the memory, HBM is global to everyone. Shared memory is local to a thread block. And registers are local to a thread. But again, all the details of the hardware-- warps, bank conflicts, memory coalescing, occupancy-- really determine performance. And many of these details are hard to know, because, I mean, the profiling tells you a bunch of information. But you have to know exactly how many SMs there are and exactly the sizing of everything. Sometimes the scheduler does something you don't really have control over what it's doing.

[20:14.060] So it's a little messier than the programming model. OK, I will stop there for questions. Yeah? Yes, a quick question. So Cam, is there a scenario in which you one can share with them or-- like, for example, in the problem that we see 140 SM [INAUDIBLE] blocks, right. And we have to launch into two different ways. Is there no way in which a block can share an SM? Yeah, so the question is, can a block share an SM? I think the issue is if you're doing things right, then the-- I guess it depends on the block. If you have a block that's for example, using most of the tensor cores on the SM, then putting another block there isn't really going to speed things up. I think fundamentally, there is this jagged problem here of unevenness. Because the blocks have to stay together. So you can't take this and spread it out over here. The thing you would do is to change your blocksize. So you change the number of blocks so you don't get this tail here.

[21:36.040] Any other questions? OK, so let's move on. Hopefully you guys are getting more comfortable with GPUs now. So now I'm going to talk about benchmarking and profiling. I'm not going to actually maybe say that much in terms of content. But I do want to emphasize the philosophy here, which is, here's a recipe for success. You benchmark and profile your code. You make changes, and you benchmark your profile, your code again. And the reason I'm doing benchmarking and profiling as opposed to teaching you try and earlier is because you should always just measure what's going on in your code and figure out what the bottlenecks are before you start writing kernels. OK, so benchmarking is basically how long things take. And it gives you just this end to end time. It doesn't tell you where things are spent. But nonetheless, it's pretty useful, because ultimately, that's the thing that you care about, how long things are running for.

[22:58.650] And it also gives you, because it distills things into one number, you can see how things are scaling, let's say, with dimension. So there's a nice tool for benchmarking. But because this class is language models from scratch, I'm going to do it from scratch. But I think I'm doing it just to highlight a few gotchas with benchmarking. So let's say we have this operation, matrix multiplication. So run operation 2 is this wrapper that basically instantiates two random square matrices of size dimension by dimension. And then it returns a function to perform the operation. OK, so matmul, if you call this, basically does the matmul of those two random matrices. OK, the random matrices are generated, it's just like multiplying the two. OK, so how do you benchmark? So the naive thing to do is just start time, run it, and then stop time.

[24:01.910] But there's a few things. One is that always remember to do your warm-up. I think I mentioned this on the second lecture, but I think it's worth emphasizing. And this is because if you have some things are lazily compiled and you want to just make sure that that thing is-- that time doesn't factor into. Because most of the time, you care about how fast something is, because you're going to run it over and over again, so the initial conditions don't really matter. And then you often want to time it multiple times because there is some variance. The proper way to time things is to use these CUDA events that-- a start event and an end event, which you can call record on. Actually do the computation and then hit the end event record. And then remember to synchronize to wait for the CUDA threads to finish, because everything on the GPU is happening if you have an async. And this is a synchronization barrier. And then you record the time.

[25:12.360] OK, and then you can do this again, and then again. And then here, we're just taking the average. If you are being very particular, you probably want to maybe look at the whole distribution, the P95 or whatever. But we'll just do the average here. And then one thing that you can do with benchmarking is you can, let's say, scale up your matrices and see how the time changes. And you can see that in this case, matrix multiplication, as expected, it should grow cubically. But notice that there is this flaw where up until you get up to almost 2,000 dimensional matrices, things are basically a constant. Right, and this is because, as we've discussed, the shapes of these GPUs are built for fairly large matrix multiplications. And if you have a 2 by 2 matrix, it's going to be very inefficient. OK. So really quickly on profiling. Profiling tells you where time is actually spent. Hopefully all of you are familiar with profiling and are doing it. Maybe less obviously, profiling, even if you don't care about the time, helps you figure out what's actually happening under the hood.

[26:44.690] Because especially with these high level languages, you write some code, and then it runs and you get some result. And sometimes it can be just good to understand what's actually going on. So PyTorch has a built-in profiler. In your assignment, you'll use Nsight, which gives you more details. But we're going to skip that in the interest of time. So let's start with just if you add two numbers and-- or sorry, two tensors in PyTorch, again, a run operation creates two random matrices and applies the operation. So in profiling, I'm warming up. And then I'm just putting this context of-- the profiling context and doing the run.

[27:36.320] And then let's look at what the profile looks like for just a plus b in PyTorch. Normally doing PyTorch, you probably don't think about it. It's like, OK, well, these two tensors just get added. So what's actually going on underneath the hood? Right, so if you look at the things that it's calling, there is this long name, kernel at CUDA functor add. So this is basically a kernel that adds to tensors. And the times aren't going to be that interesting, because I'm only adding. So that's going to take 100% of the time. But this tells you that, well, underneath the hood there is this thing called add. What about matmul?

[28:28.010] In PyTorch, I'm doing a at b. Similarly, there is this long name that describes this particular matmul kernel. Cutlass f32, f32, 64x64x16, and so on. Notice that if you change the dimensions-- so now I'm doing a 128 by 128 matmul-- you get actually a different one. If you look closely, that this is 64x64x16, this is 32x32x16. OK, so underneath the hood, PyTorch, it looks like you're doing just at. But underneath the hood, there could be all sorts of things that are happening. So observations. Here, you can see which CUDA kernels are actually being called.

[29:18.520] These are generally the ones that have the long names. And different CUDA kernels are invoked depending on the tensor dimensions. The name actually tells you something about the implementation as well. So this name. So cutlass is the NVIDIA CUDA library for linear algebra. SM100 corresponds to the Blackwell architecture. So this is a kernel that's specifically designed for Blackwell. This is f32. And then 64x64x16 is the shape of the tile, which we're going to talk about later when we talk about tiling in the context of matmuls. OK, so that's benchmarking profiling. Just remember to do it. I think we make you do it on the assignment, so you have no choice. OK. So let's apply this to another example here on the GeLU. So remember, the GeLU activation function is this function, which is a typical non-linearity that's used.

[30:32.800] Often, it's approximated with this tanh approximation, which is more computer friendly. OK, so naively, if you implement the GeLU, you can do it in PyTorch just like this. So you take a tensor, you just take this equation and you just put it into PyTorch, OK. That's fine. And you get some number. PyTorch also has a built in. If you call functional GeLU, you can get that version as well. And you can check that they're the same. You can run the two GeLU versions and make sure that the answer is the same for a random input. There's also something that people probably have discovered. If you haven't encountered this, this is an important thing to know. Is that you can take any PyTorch function, you call torch.compile on it, and it generates another function. And it does the same thing.

[31:47.810] OK, so we have three horses in this race. We have the naive implementation. We have the built-in implementation. We have the compiled implementation. So we can benchmark them. Naive takes three. I guess this is 3.75. Built-in is much faster and compiled is much faster as well, but not maybe as quite as fast as built-in. OK. So what's happening here? Why are these things different? They all compute the same answer, but they have wildly different performance characteristics. And so here we can pull up the profiler and see what's actually going on underneath the hood. So if you look at the naive GeLU here and just do a profile, this-- something's wrong with this view, so I'm going to go to this.

[32:51.560] So you see that the profile shows how time is being spent. There's a bunch of different kernels on this binary functor unary add. tanh is a kernel here. And this corresponds to the fact that in PyTorch, when we write the PyTorch expression, you look at the computation graph, and each primitive in the computation graph is actually realized in the kernel. Right. And the reason this is slow is that when you launch a kernel, the kernel has to read from HBM. Pull it all the way over to your SM, do the computation, and write it back. And then the kernel picks that up from HBM and then writes it back, and so on and so forth.

[33:43.180] So you're doing a lot of reads and writes back and forth. Because between kernel invocations, things have to go back to HBM. So if you look at what the built-in is doing, this is actually not that interesting. There's this GeLU CUDA kernel implementation, which is a single kernel that just implements the GeLU. And so why does this exist? Well, because people use GeLU. So someone wrote a kernel for it and put it in the standard library. I mean, there's nothing magical. So compilation is really interesting here. I'm not going to say too much about how it works, but it's really, I think, a fascinating topic where you can take a naive implementation, which, remember, it's in PyTorch, it has a computation graph. And run a compiler. Which basically, if you look at what's underneath the hood, it's just a single kernel. And this is because it's figured out to look at the computation graph and essentially write that kernel in Triton.

[34:59.142] So you can see that this is actually a Triton kernel. OK, so naive implementation. Multiple kernels requires multiple reads and writes to and from HBM. There's no kernel fusion here. This is slow. The built-in compiler version, there's one kernel. Basically all the operation in the GeLU have been fused together into one kernel. So you read from HBM once, you write to HBM once per element. And you see that the compiled kernel is a Triton kernel. OK, so that maybe is a good segue to talk about what this Triton thing is all about. Yeah? [INAUDIBLE] Let's see. I mean, I guess it says CUDA kernel implementation, so I imagine someone wrote it in CUDA.

[36:03.236] OK, any other questions? Yeah? Why is Triton kernel faster than this one? So why is the Triton kernel faster? So a Triton kernel is actually not faster in this case. Oh, sorry. The compiled kernel is one Triton kernel, and this is slower than the built-in. Yeah. I think last year when I did this, it was actually closer. But these things change, and it's very hardware dependent. And I think none of this is terribly optimized. This is just giving you the general idea here. OK, so let's write some Triton kernels. So remember our programming model. You have a bunch of threads organized by Thread groups, a thread blocks, and there's a grid of thread blocks. And if you were to write in CUDA, which was originally developed by NVIDIA, and that's been for years the thing that you do when you write kernels, the mental model is, what does each thread do?

[37:22.190] So you write a piece of code which essentially has some ID that identifies which thread you're talking about, and it just executes the code. So the nice thing about this is that it's very closely related to what is actually happening underneath the hood. And it's fine-grained, gives you fine-grained control. But there's cons here, which is that, remember, all these threads are in a thread block. And some operations require them to communicate. So what has to happen is that they have to synchronize. And they read from HBM all at once. And they have to synchronize and do the computation. And you have to basically do that bookkeeping. So if you were doing all element-wise operations, this CUDA is just fine. It doesn't really matter. But as you get more complex operations, then try and provide some value of the abstraction.

[38:32.140] So Triton was developed by OpenAI. I think by now, it's been pretty standard. You basically specify what each thread block does. Generally it's powerful enough, especially for this class that you're getting started. If you really go and want to exploit every single new feature of the latest hardware, it might not give you the full flexibility. But let's not worry about that. And the conceptual framework to think about in Triton is that you're thinking about, what does a block do? A block is going to load data into shared memory, operate it on, and write it back to global memory. So in some sense, these blocks are intermediate point between thinking about what individual elements are doing as well as thinking about the general operation. And so in PyTorch, you basically define these huge matrices and you say multiply them together. And that's the atomic operation. And a lot of what you're thinking about is, how do I get things into big matmuls?

[39:34.070] Right, and at some level, Triton is a hybrid between that and the individual elements, as we'll see. OK, so let's start with the GeLU example here. So let's define an 8,000 dimensional vector and start writing some Triton. OK, so Triton is basically going to be you write Python. How many of you have written Triton before? Just a show of hands. OK, OK. Not that many, which is good, because then you won't be bored. OK, so this is normal PyTorch. Right, there's no Triton here. But I'm just preparing. And so what I'm going to do is I take this tensor, and I'm going to allocate an output tensor.

[40:35.890] OK, because in Triton, we're not thinking functionally anymore. We're just thinking about moving. You have to explicitly read and write. So there's no returning value. So I'm going to allocate the output tensor, which I'm going to-- the kernel's going to write to. And then so this tensor can be arbitrarily big, right. And I can't generally fit this all into one SM, because it's just too big. So I need to break it up into blocks. So what I'm going to do is you can think about this x as being this array. I'm going to chop it up into blocks. So the total number of elements is 8,000. I'm just going to set the blocksize to 1,024 for now. And then I have eight blocks. OK. So then I'm going to use this weird syntax to call the kernel. So Triton GeLU kernel. This basically in brackets tells me essentially the shape of the grid.

[41:54.630] So basically, this says the grid has num blocks. Blocks. And I'm going to, for every one of these blocks, invoke this function, Triton GeLU kernel, which I'll talk about in a bit. I'm going to pass in an x, y num elements and the blocksize. OK. All right, so unfortunately, I'm not going to be able to trace through this, so I'll just show you what the code looks like here. So let me get rid of that. OK, so this is probably the simplest kernel you can imagine. So when you are looking at the triton gelu kernel, now this is before we had x and y. So now these are pointers. Right, you can think about these as just integers, they're addresses. So you have to get comfortable with that. And then we have the number of elements and the blocksize, which are basically actually passed in from here. OK, so the way to think about it is that for every block, we're going to have this function being called.

[43:20.240] And what happens is the first thing I'm going to do is I'm going to wake up. This block wakes up and say, who am I? Well, the program ID is-- basically the PID is basically identifying the block. So PID would be 0 for this block 1, 2, and 3 here. OK, so now I have to figure out what data I'm going to operate on. So that's PID times blocksize. So this is the offset into x pointer. So if PID is 0, then start is 0. If PID is 1, then start is blocksize. If PID is 2, then it's 2 times blocksize, and so on. OK. So now, I'm going to figure out the span I'm operating on. So offsets is start plus. This is the Triton library. A range 0 blocksize, which conceptually gives you the integers 0 through blocksize minus 1. So offsets is going to be essentially for-- let's say block 1, it's going to be blocksize blocksize plus 1 all the way to blocks 2 times blocksize minus 1.

[44:39.700] OK, so in this case, num elements divides by blocksize. But in general, that's not the case. So you'll often see in train code, there's this masking. Which says, well, sometimes if the-- let's say the tensor only goes up to here. Then I'm going to form a mask, which is going to be true up until that point and false after that point. And for blocks that are not the final block, it's going to be just all ones. OK, so now I've done the setup. What I do is I read. And this is basically pointer arithmetic. So x pointer, remember, is the integer that specifies the memory location of where x is. I'm going to add offsets to that, which gives me the first blocksize number of elements there. According to the mask, if I'm masked out, then I don't read it. And then now I think you can just think about this a vector. And then you do your normal computation.

[45:49.040] And then you get y. And y is the same size as x. And then you do tl store y pointer plus offsets y and the mask. OK, so this load loads from HBM, does some stuff, and then it writes back to HBM. OK, I'm going to stop there and take any questions about this first Triton kernel. Yeah? [INAUDIBLE] Yeah. So the difference is, what is the difference between CUDA? And for this element wise, it looks pretty much the same, right. And in fact, CUDA is even, I think, simpler, because it really is element wise. You wake up, and you identify the thread. And then you just operate on that element. Now the only thing here is that it's like the vectorized version, where you have a block and you operate on that block. Later, we'll see that operating on blocks is doing more than-- if you do something more than element-wise, then CUDA is going to be a lot more annoying to work with.

[47:14.530] So a bunch of questions. Yeah, back there. How does this relate to if you wanted to use the tensor units? So how is this related to if you want to use the tensor units? So I'll later show you what this code actually compiles down into. And maybe I'll get back to that. But the short answer is that you don't control that. The hardware figures out where to put things. Yeah? Can you actually walk through what's happening at tf.load from HBM to share to register? What's the step by step where data is actually flowing? Yeah, so the question is, what is actually happening when this statement is executed? So the short version is that this is in some sense a lie. Right, it's not like the GPU actually calls the Triton library on the GPU and it's actually executing this code. This is basically for our consumption to specify the computation. The compiler takes this, as we'll see later, and writes it into something called ptx, and then it will actually do the work.

[48:32.990] So that's mechanically what's happening, which I'll get to later. If you're asking about the conceptual question, the way to think about this is that this-- x pointer is a memory location in HBM. And this basically specifies a range of memory locations. And load takes those memory locations and returns the data associated with them. Here, I have a local variable called x. In practice, this is going to be generally a register or shared memory. Triton figures out what to do there. [INAUDIBLE] --at the threat level. When this is invoked, it's not really clear. Did I line up what is going to be in shared and going to be at the register level beforehand, or is it happening now, which is way late? And so then so that's just where all the threads are sitting idle to get a memory to flush in. I'm sure it's ready to execute so that we're not sitting there twiddling our thumbs, just waiting for data to come in from HBM. Yeah, yeah. So the question is, when does this actually get executed, and doesn't this block?

[50:01.370] Let me try to come back to this question when I show you the ptx, and maybe hopefully it will provide a bit more context on what's happening. OK, any other questions? OK, so all the kernels are going to look something like this. So I do want to make sure people understand just the general form, which is that you have generally your inputs, your outputs. You wake up, you figure out which index you're going to look at. You read. You do some stuff, and then you write to HBM. OK. So let's talk about ptx briefly. So let's see, does this work? Does this link work? OK, so when you write Triton, the compiler generates ptx code, which is this intermediate assembly language for GPUs. I am obviously not going to go through all of this, but just to give you the flavor of what this looks like. Let me actually start with the observation. So a few things. So this is now what a thread is actually doing. Not a thread block, because that's been compiled away.

[51:35.710] And so if you look at-- a few notes here. So this LD global is basically saying load from HBM into some registers. And the registers are denoted as the r are integer registers, fr, floating point registers. And then you have statements like mov 0 into r5, mov 0 into r6, and so on. And then you have multiplication. You're going to multiply this register by this constant and then put it into this. So this code gets executed. And at the bottom, you should see global store. So this is writing back into the HBM. So this gets executed. This is actually the code that gets executed on a thread. And Triton is basically a layer up above that. One other thing to notice is that you see all these blocks and what's going on there. And what's going on there is I alluded to thread coarsening, which is that this is one thread. But rather than processing a single element, it's actually processing eight elements.

[52:59.630] So the compiler decided that, well, actually, this thread is pretty lightweight. It doesn't do that much, so let's just try to thicken it up a little bit. So looking at ptx can give you some of appreciation for what's going on underneath the hood. Yeah? So it makes this for each individual-- [INAUDIBLE] So the question is, does it make this for each thread? So this is compiled once. And it's the same piece of code that each thread runs. And the way that the thread distinguishes itself is that this piece of code gets passed basically the thread ID. So here, the cta.x is the block index, and tid.x is the thread index. So this basically, if this piece of code is running, this tells you which block I'm in. And tid.x tells you which thread inside that block I'm in. OK. Any other questions? And I'm trying to figure out how to-- OK. OK, so that's generally a flavor of your first Triton kernel load from HBM compute right back to HBM.

[54:42.920] And then we saw the ptx, which is the grungy what actually happens underneath the hood. There's still a lot of things that are not specified in ptx. Like, for example, which SM things are operating on and the warps and everything. A lot of those are hardware controlled, so you don't even see. Yeah? So ptx code is generated by the compiler, not something that you would normally code? It looks like it's assembly code. Yeah, so the question is, is ptx just generated by the compiler? So there are people who do write ptx if you really think you're better than the compiler. And I think the NVIDIA compilers are generally pretty mature, but some other accelerators that are less developed. I think sometimes you just have to reach in and actually handhold a bit more. But generally, yeah, you shouldn't need to do that. Yeah? [INAUDIBLE] My work will get scheduled onto an SM.

[55:59.050] I'll get to that tf.load and it's almost like a CPU trap call where I'm just waiting for something to happen. And now I am turning my thumbs and some other work will get scheduled over me. And then it will do operations. Now that you have to upload is done, then I'll reschedule that work and then I'll continue on. Is that-- Yeah, that's right. So just to repeat the question or the comment. So if you look at this triangle code, these statements, this load, is going to block for some number of cycles. And then so this is running on some thread on some warp on some SM. And remember, the SM is running multiple warps at the same time. So when you get to that point, it can just find another warp to run. And then when this is done, the warp scheduler comes back and takes over. Why do you [INAUDIBLE]? Sorry. Oh, sorry. I recall from [INAUDIBLE]. Yeah, so the question is, why four warp schedules?

[57:09.800] I don't exactly the reason behind that. OK, so now let's go through some other examples. Maybe just as a quick preview. We're going to do three more examples. GeLU is the simplest form. Even though the computation has a lot of messiness. It's just element wise. So conceptually, in the context of this lecture, is actually very simple. Now we're going to look at Softmax, where now you're going to have to do reduction. But in this case, we're going to think about the case where row fits on a block. And then we're going to consider the case where row doesn't fit on a block, and then we're going to do go up to matmul. And then hopefully by that point, you'll have all the ingredients that you need to do the assignment and implement flash attention. OK. So so far, we looked at element wise operations. Now let's think about other operations that aggregate over multiple values.

[58:16.200] So remember what the Softmax does. Let's just think of as a matrix. You exponentiate and normalize each row of a matrix. OK, so this is used in intention and used in generating output probabilities. Generally a good thing to do. So let's just start with a naive implementation and just keep track of what's happening. So here, I'm defining this tensor. And here's the naive implementation. Actually, I think this is on assignment one. OK, anyway. So here, I have an M by N matrix. And so I need to, for every row, compute the-- I'm going to compute the max of each row. And this is for numerical stability. I'm going to subtract off the max. And then I'm going to exponentiate element-wise. And I'm going to sum and compute the normalization constant for every row. And I'm going to divide. And then that's it. OK. So if you count the number of reads and writes, remember, this is just plain PyTorch.

[59:47.180] So this is a different kernel. This is a different kernel. And unless you call to torch and compile, these are going to be different operations. And each operation is going to read and write, read and write from HBM. OK, so you have five MN reads, three MN writes. And in principle, you should really only have much fewer. OK, so this is a piece of code. Makes sense. Everyone should be familiar what the Softmax is doing. OK. So let's write down the Triton kernel. And in some sense, the form is going to be very similar. Just like in the GeLU. Right, once you do the scaffolding, the core computation looks very much like the naive version. So what we're going to do is we're going to say each row is a block. OK. And why do I make each row a block? Well, because now, remember, each row, I have to normalize and sum.

[01:00:55.010] So it's not element-wise. So Softmax is not an element wise. But it is row wise. So the blocks don't interact. And the blocks don't have shared memory, so that's fine. There's no shared memory across blocks. So then within each row, I'm just going to do some stuff. OK, so let's see what happens. So again, I'm going to, in Python, I'm going to allocate my output tensor. I have this M by N matrix. I'm going to define the blocksize as basically the number of columns. Go to the next power of 2 for good luck. And then the number of blocks is just the number of rows. OK, and I'm going to call this kernel. So how many blocks are there? M. Right, one for each row. And each block, I'm going to pass the input pointer, the output pointer. And then I'm going to also pass these strides which tell me how far to move down. OK, let me actually go and show you that Softmax kernel, OK.

[01:02:23.860] OK, so what does this look like here? So I wake up. I'm on a particular row, OK. And this is going to give me all the columns from 0 to the number of blocksize, which I'm assuming that's all the columns. I'm going to read. So basically, I need to figure out where to read from memory, and that's going to be the start of my data. Plus which row. And the row stride basically gives me every row is-- basically, this is the number of columns, essentially. It's going to tell me where how far to go down. And then this is x pointers is basically the addresses of all the data I need to load. I load them up. And here, I do this thing where if it's masked out, then I put minus infinity, because that's going to be the equivalent of a 0 for the Softmax operation. And then this part is essentially the same as the naive Softmax. I'm just going to compute the max, subtract it off, exponentiate, sum, and divide.

[01:03:44.070] And then I write it back. OK. Yeah. [INAUDIBLE] Yeah, so this is the version where each block can just span the entire row. And you can see that Triton makes this very easy, right. Because this is as if you were just writing normal PyTorch code. So if everything fits in a block, you basically-- the thing is, if you can fit it into a block, you can just write normal PyTorch, almost. Yeah? What if my number of columns, the number of rows, are bigger than blocksize? Yeah, so we'll get to that. Now the number of columns, rows, in general, it's going to be much larger than the blocksize. So we'll come back to that, yeah. So if you wanted to do a Softmax per column, do you have to do [INAUDIBLE]?

[01:04:52.530] So the question is if you want to do Softmax by column. I think that should be fine, because here, we're tracking these pointers, right. So the pointers can be anything. Basically, I think all you would have to do is change the stride here to basically access the columns. Actually, it would be here. The column offsets, you would just multiply this by the row stride, I think. OK, let's move on. All right, so now, warming up to the matmul. Suppose that your row doesn't fit into a block. OK, so what do you do in this case? So for example, if you have, let's say, a row that's 4,000 columns, but the blocksize is only 1,024. So you have to do something here. So here's the strategy. We're going to break up the row into tiles. So in this case, there's going to be four tiles. And each thread is going to iterate over the tiles and accumulate a sum. And then finally at the end, we're going to do the reduction by summing everything that each thread produced.

[01:06:17.530] So I'll show you an example of this. And now I'm switching from Softmax to row sum because it's just easier to think about. OK, so here, OK, well, this is not very interesting. The built-in row sum does what you would expect. So the row sum operation just basically takes a matrix and then computes the sum of each row, OK. So just no surprises there. OK, so conceptually what we're going to do here is as follows. OK, so each block is still in charge of one row. So that part hasn't changed. So suppose block one, row one. I wake up, and what do I do? I'm going to remember, now, there's tiles. And suppose tile 0 is columns 0 through 3, tile 1 is columns 4 through 7. And then tile 2 is columns 8 through 11. OK, so what I'm going to do here is that I'm going to iterate now. So first, I basically process this, all the threads. Basically, each thread keeps a kind of accumulator and processes the first tile.

[01:07:32.040] And it's going to move on to the next tile, and it's going to add the current element to the accumulator. So here, I'm putting 3, 1, 4, 1. And then the second loop iteration, I'm going to add 5 to this. I get 8. I add 9 to this, I get 10. 2 to 4, and 6 to 1, right. So each of these four threads is going to be accumulating its own thing. And then at the end, I do another for tile 2. I add 5 and 3 to the respective accumulators. So at the end of the day, I have this, a vector of accumulators. And then that is actually-- that, I can just sum up. OK, so let's see what the code looks like here. So I'm just going to call the row sum kernel. And let's see. OK, so here's what it looks like. So wake up. I am on a particular row. And this is what one row looks like. There's one tile, a second tile, a third tile, and so on.

[01:08:46.399] So n is the number of elements of that row. And blocksize is the size of this tile. OK. So remember, blocksize is the number of threads. I'm processing more data than the larger than the blocksize. I just have to do it iteratively now. So I'm going to loop over all of the tiles. So start is going to go from 0 to blocksize to 2 times blocksize, and so on and so forth. So each time I'm jumping across, I am doing calls. Basically get the offsets at the particular tile. And then I'm going to load the data from HBM, and I'm going to add it to the accumulation. So this accumulation is going to be either in registers or shared memory. OK. And then finally after I loop over all the tiles, I process the whole row, I get basically-- for every thread I have accumulator of what that thread has picked up. And I can just do a sum to get a scalar, and I write it out.

[01:10:05.865] OK, so this is a little bit more complicated than before, because now we have a for loop within a thread. And this is necessary when your data doesn't fit within a block. Yeah? [INAUDIBLE] Just in case we don't have enough register. We have to put it in-- Yeah, so the question is I think if you can control where the accumulator resides. At least in this Triton program, you don't explicitly say. And this is up to the Triton compiler to figure out where to put it. But in general, if the blocksize is large enough, it has to go in shared memory. OK. All right, does this make sense? Just to, I think, make sure people are on the same page. Remember when GeLU, we also split a row into a bunch of pieces. But those were blocks. Right, and so each of those pieces was a block that was processed independently. These are not blocks. These are tiles. The block corresponds to this whole row. And the block has to basically process all the tiles.

[01:11:28.560] And so this is where it starts to not look like PyTorch, because you're not able to process all your data in one nice kind of-- not everything fits into shared memory. Yeah? [INAUDIBLE] --also stored in shared memory? Yeah. Let's say the accumulator is stored in shared memory. Yeah. OK. All right, so let's now maybe go on to our finale, which is matmul. OK, so matmul, of course, is the bread and butter of deep learning. It's been optimized to death. And it's, in some sense, very fundamental operation. So you take two matrices, you multiply them. I'm going to add a little bit of a twist here. I'm going to do a matmul followed by a ReLU, just because, just for kicks. OK. Well, this happens, right. Because if you have one linear layer, it's a matmul. And then you apply a ReLU activation. So this is not completely out of nowhere.

[01:12:37.750] But I'll show you later why I did this. OK, so how do you build a matmul kernel? All right, so here's a naive approach. OK, so here's my, let's say, a matrix. I'm multiplying A times B. And I'm trying to write out the matrix C. And A is M by K, B is K by N. So C is M by N. OK, so what I'm going to do is I'm going to fix one of these elements. Let's say M equals 1, N equals 2. OK, so I'm processing-- actually, let's do M equals 1, N equals 1. So I'm doing C5. And then, basically for every k, so I'm going to, yeah, I'm going to iterate over the k, the rows of A and the columns of B. I'm going to read from HBM.

[01:13:37.040] I'm going to multiply them, accumulate that. And at the end, I'm going to write out to this and this element. So that is a valid matmul kernel. OK. So what's wrong with it? It's correct. But if you look at how many reads and writes it's doing, this is not good. Right, so basically for every M and N and K, I have to read from HBM. So it's on the order of M times K times N reads. Number of writes. I mean, it doesn't matter, but this is the bottleneck, right. And if you remember from the second lecture, if you look at the number of operations you're doing divided by the number of bytes that were transferred, that's the arithmetic intensity which you want to be high. So the number of operations is M times K times N order. And then the number of reads is also the same. So arithmetic intensity is a constant, which is not good. OK, so if you look carefully, you notice that there is a lot of redundant reads.

[01:14:49.310] So imagine computing C4. You needed to read A4, A5, and A6. And if you compute C5, you're going to have to read those over again as well. So if you can just read those once, then you really save on reads, and that's great. So let's try to use shared memory to do that. So here's the idealized approach. I'm going to load all of A and B into shared memory. And then I'm going to just compute C. So if I can do that, then I get-- now I don't have this cubic number of reads, I only get quadratic number of reads, which means I get the error intensity of order n, which in the second lecture, I said was a kind of ideal thing you could hope for. OK, so if you can do that, that's great, because you basically-- there's no redundant reads. You read everything once into shared memory. You do the computation and you write it back. OK, but what's the problem with this idealized approach. The problem is that A and B are usually too large to fit into shared memory.

[01:16:09.166] Right. So then, what do you have to do now? OK, so the idea here is this very classic idea of tiling. And the idea basically is to, well, fit as much as you can to shared memory as you can, essentially. So in some sense, it's going to look like the naive-- it's going to globally look like the naive approach, but locally it look like the idealized approach is the way to think about it. So here's the picture. I think Tatsu showed this as well. So what we're going to do is we're going to take this matrix C, and instead-- remember, the naive approach just said for every element, I'm going to compute it. But now I'm going to say for every tile, I'm going to do it. So I break it up into tiles. And each of these tiles is going to be a thread block. And I'm going to have a bunch of threads that is responsible for computing this. Now a different tile is going to be computed completely separately by another thread block.

[01:17:14.090] OK, so I'm going to-- imagine I'm in Triton. I wake up. I'm looking at this thread block. So what do I have to do? Well, in some sense, it's the same as the naive approach where, for every row tile of A, I'm going to go scan across the rows. And for every column tile of B, I'm going to load the corresponding tile A and the corresponding-- so say I'm here-- and the corresponding tile from B into shared memory. I'm going to multiply these two together. And that's going to be like the idealized approach. And I'm going to accumulate that in the partial sum. And this is all sitting in shared memory. And then after I finish all the sweep of the row and the column here, then I can finally write this output tile to HBM. OK, so that's the conceptually what's happening. And then the arithmetic intensity here now goes up to order tile size. So you can generally not reach order n, because that would require you to fit everything into shared memory. But if your tiles are big, then that's still not too bad.

[01:18:36.860] OK, so just as a bonus. While you're doing all writing a kernel for this anyway, sometimes if you want to apply an element-wise activation function, it's very easy to just put it on at the end. And this is kernel fusion. OK, so very quickly, the implementation here. So-- oops. OK, so just a reminder about strides since this is going to show up. So a tensor is a multidimensional array. But in memory, it's linearized. And strides of a tensor tells you basically how to map from a multidimensional index such as a row, a column, into an actual index. And basically what you do is multiply the row by the stride, the plus the column by the stride of the column. OK, so in this case, every time you go advance to the next row, you go four positions in your memory. And every time you advance a column, you go one.

[01:19:42.600] And if it were to transpose, it would be flipped. OK, so what does this kernel look like here? OK, so the launch is not interesting. So you wake up and you are on tile M and N. So it's like OK, I'm responsible for computing the c matrix, but for the M comma N tile there's a bunch of index manipulation which you can-- I'll just gloss through, but it's straightforward but a little bit. You just have to track the indices. So this basically tells you which rows of A, the matrix A, I'm looking at. Which rows of-- oh sorry, which columns of B I'm looking at. And then this is just the numbers 1 through K. Then I'm going to get the pointers into A and B at my tile location. And then I'm going to set up this accumulator matrix. This is going to be in shared memory. It's M by N. And then this is going to look like the row reduction, right. You have this sum over all the tiles. But instead of just going across the tiles, I'm now going across the row tiles and also simultaneously down the column tiles of B. I load A, the small matrix, I load B, the small B

[01:21:13.820] tile. And then I perform this dot. So remember, whenever things are in shared memory, things look like PyTorch and I can just say, matmul it, and it will do the thing. And then I advance to the next row tile of A and the next column tile of B. The bonus is that if I wanted to apply element-wise non-linearity, I might as well do that here. Before I write it out to HBM, I can do some any operation on it. And then finally, I just write it out. OK, so there's some indices that you have to pay attention to. But hopefully, the form of the algorithm is clear. OK. Any questions about that? All right. Maybe I'll just summarize, and you can ask me later. So today, we talked about there's a programming model, which is you're talking about either PyTorch or Triton or ptx. This is what the kind of programmer can control. Right, and even ptx. You can write ptx and you can control every-- specialize however you want.

[01:22:38.560] But this is not the full picture, because the reality is that your code has to run on hardware, and there's only a finite number of SMs, a finite number of banks, and sizing of memory and registers. All are finite. So you come in with your big matrix and transformer, and want to-- it has to fit in with the constraints of the hardware. So that's why benchmarking and profiling are really important to understand how the messiness of the hardware translates to performance. We talked about Triton, which is, I think, a pretty nice, neat language to think about thread blocks. Hopefully by now you can appreciate that things are easier to think about thread blocks than individual threads, because you don't have to think about explicitly synchronizing threads or doing shared memory. And the way to think about it is that you figure out-- you have your computation. You break it down into the thread blocks where you just need to read from shared memory. Do some stuff, and write it back into HBM.

[01:23:50.380] And then we saw some examples of increasing difficulty. Element wise is the easiest. And then reduction over a row. Reduction where it doesn't fit into a row. And we introduced the baby tiling. And then matmul is the canonical example where you actually do tiling. OK, so that's all I'm going to say about how to program a single GPU. Next time, we're going to go to more GPU and talk about multi-GPU programming. Yeah, question. So if I want to write down kernels, what are the alternatives? Do I have to try them and-- for each of my choices that I have, how close can I get to the optimal level? So the question is, what are alternatives to Triton? So there's a trade off between-- every language has an inductive bias, right. That makes certain things easier and certain things harder. So most of what we do-- Triton was built by people who train transformers.

[01:24:56.350] So anything involving transformers, I think it's going to be relatively easy there. Of course, in the extreme, you can always go to ptx and write that. But I wouldn't advise that as a first step. There are a bunch of other libraries. There's thunder kittens. There's cute, delirious DSLs that allow you to give you-- they're not necessarily comparable either up or down the stack. They just give you different characteristics. Yeah? Do you have a high dimensional tensor magnification or high dimensional tensor processing? What is the best approach? What I'm thinking is can I load unload the whole sensor, both of the sensors on bunch of threads or thread blocks on my GPU at the same time and compute them? Or would it be better to take each element, just what we did here, right.

[01:26:06.210] A particular component of my sensor, then process it and then write it back in the end. Yeah, so briefly, the question-- and we should probably wrap up-- is, is it better to read all at once or maybe process individual elements at a time? I think it's hard to answer this in the abstract. It depends on the nature of the computation. Maybe we can talk offline about that. OK, all right. See you next time.

Kernels, Triton

课程简介 ​

CS336 Lecture 6：Kernels、Triton 与 GPU 性能优化 ​

学习目标 ​

课程主线 ​

一、GPU 硬件回顾 ​

1.1 内存层次结构 ​

1.2 编程模型 ​

三个层次 ​

为什么需要线程块？ ​

1.3 编程模型与硬件的交互 ​

Warp：锁步执行的线程组 ​

占用率（Occupancy） ​

存储体冲突（Bank Conflicts） ​

内存合并（Memory Coalescing） ​

块占用率（Block Occupancy） ​

二、基准测试与性能分析 ​

2.1 方法论：成功配方 ​

2.2 基准测试（Benchmarking）：测量端到端时间 ​

正确做法 ​

基准测试的用途：观察扩展性 ​

2.3 性能分析（Profiling）：定位瓶颈 ​

PyTorch 内置 profiler ​

观察结果 ​

三、案例研究：GeLU 的三种实现 ​

3.1 GeLU 函数 ​

3.2 三种实现 ​

3.3 为什么性能差异这么大？ ​

四、Triton 入门 ​

4.1 CUDA vs Triton ​

4.2 Triton 的概念框架 ​

五、Triton 内核示例 ​

5.1 GeLU：逐元素操作 ​

5.2 底层发生了什么：PTX ​

5.3 Softmax：按行归约 ​

5.4 Row Sum：一行装不下的归约 ​

5.5 矩阵乘法：Tiling 的艺术 ​

朴素方法的问题 ​

理想化方法 ​

Tiling 方法（分块） ​

六、与课程作业的关联 ​

七、常见误区 ​

八、复习清单 ​

从 PyTorch 到 Triton：理解 GPU 内核编写的核心思维 ​

导语 ​

一、GPU 硬件回顾：容量与带宽的权衡 ​

二、编程模型：线程、线程块与网格 ​

为什么需要线程块？ ​

从第一张结构图到一次 Block 的执行 ​

三、编程模型与硬件的交互：那些决定性能的细节 ​

1. Warp 与控制发散 ​

2. 占用率 ​

3. 存储体冲突 ​

4. 内存合并 ​

5. 线程块占用率与波次量化 ​

四、基准测试与性能分析：成功的方法论 ​

基准测试的基本功 ​

性能分析：看穿 PyTorch 的伪装 ​

五、GeLU 案例：内核融合的巨大威力 ​

六、Triton 入门：以线程块为单位的思考 ​

第一个 Triton 内核：GeLU ​

七、Softmax：归约操作与行列设计 ​

八、行求和：当数据装不下一个块 ​

九、矩阵乘法：分块（Tiling）的经典实践 ​

朴素方法 ​

理想化方法 ​

分块方法（Tiling） ​

十、课程总结与学习路径 ​

参考资料 ​

CS336 2026 Lecture 6：完整中文逐字稿 ​

正文 ​

CS336 2026 · Lecture 6 · Kernels, Triton ​

正文 ​

课程简介

CS336 Lecture 6：Kernels、Triton 与 GPU 性能优化

学习目标

课程主线

一、GPU 硬件回顾

1.1 内存层次结构

1.2 编程模型

三个层次

为什么需要线程块？

1.3 编程模型与硬件的交互

Warp：锁步执行的线程组

占用率（Occupancy）

存储体冲突（Bank Conflicts）

内存合并（Memory Coalescing）

块占用率（Block Occupancy）

二、基准测试与性能分析

2.1 方法论：成功配方

2.2 基准测试（Benchmarking）：测量端到端时间

正确做法

基准测试的用途：观察扩展性

2.3 性能分析（Profiling）：定位瓶颈

PyTorch 内置 profiler

观察结果

三、案例研究：GeLU 的三种实现

3.1 GeLU 函数

3.2 三种实现

3.3 为什么性能差异这么大？

四、Triton 入门

4.1 CUDA vs Triton

4.2 Triton 的概念框架

五、Triton 内核示例

5.1 GeLU：逐元素操作

5.2 底层发生了什么：PTX

5.3 Softmax：按行归约

5.4 Row Sum：一行装不下的归约

5.5 矩阵乘法：Tiling 的艺术

朴素方法的问题

理想化方法

Tiling 方法（分块）

六、与课程作业的关联

七、常见误区

八、复习清单

从 PyTorch 到 Triton：理解 GPU 内核编写的核心思维

导语

一、GPU 硬件回顾：容量与带宽的权衡

二、编程模型：线程、线程块与网格

为什么需要线程块？

从第一张结构图到一次 Block 的执行

三、编程模型与硬件的交互：那些决定性能的细节

1. Warp 与控制发散

2. 占用率

3. 存储体冲突

4. 内存合并

5. 线程块占用率与波次量化

四、基准测试与性能分析：成功的方法论

基准测试的基本功

性能分析：看穿 PyTorch 的伪装

五、GeLU 案例：内核融合的巨大威力

六、Triton 入门：以线程块为单位的思考

第一个 Triton 内核：GeLU

七、Softmax：归约操作与行列设计

八、行求和：当数据装不下一个块

九、矩阵乘法：分块（Tiling）的经典实践

朴素方法

理想化方法

分块方法（Tiling）

十、课程总结与学习路径

参考资料

CS336 2026 Lecture 6：完整中文逐字稿

正文

CS336 2026 · Lecture 6 · Kernels, Triton

正文