""" ============================================================================= 第 6 讲 · 代码 3：带归约的核函数 —— 融合 softmax 与分块行求和 ============================================================================= 对应讲义：lectures/lecture06_kernels_triton.md 第 6、7 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 逐元素运算里，每个元素互不相干。更常见也更难的是**归约**（reduction）： 一行的结果依赖这一行的所有元素，比如求和、求最大值、softmax、LayerNorm/RMSNorm。 1. 融合 softmax：每个程序负责矩阵的一行，把整行读进来， 在片上完成"求最大值 → 减去最大值 → exp → 求和 → 相除"，只读一次、写一次。 朴素的 PyTorch 实现要读写很多遍（CS336 课上的计数：读 5MN+M、写 3MN+2M）。 2. 分块行求和：如果一行太长、一个块装不下怎么办？在程序内部**循环**， 一块一块地读、累加到寄存器里，最后再归约 —— 这是"分块"最简单的形式， 也是第 5 讲 FlashAttention 中"沿 key 方向循环"的原型。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 03_triton_reductions.py （有 NVIDIA GPU + triton 时用真正的 Triton；否则用 mini_triton 模拟器。） """ import torch REAL_TRITON = False try: import triton import triton.language as tl REAL_TRITON = torch.cuda.is_available() except ImportError: pass if not REAL_TRITON: import mini_triton as triton tl = triton.language DEVICE = "cuda" if REAL_TRITON else "cpu" def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # ============================================================================= # 第 1 部分：融合 softmax # ============================================================================= @triton.jit def softmax_kernel(x_ptr, y_ptr, n_cols, x_row_stride, y_row_stride, BLOCK_SIZE: tl.constexpr): """ 对矩阵的每一行做 softmax。网格有 n_rows 个程序，第 i 个程序负责第 i 行。 【参数】 n_cols : 每行元素个数 x_row_stride : 相邻两行在内存中相隔多少个元素（行优先连续矩阵就是 n_cols） BLOCK_SIZE : ≥ n_cols 的 2 的幂（整行一次读进来） 【边界】 - 被掩掉的位置用 -inf 填充：exp(-inf) = 0，不影响最大值和求和。 - 一行必须能放进一个块（受寄存器/共享内存限制，大约几万个元素以内）。 更长的行需要在程序内循环 + 在线 softmax（第 5 讲的技巧）。 """ row = tl.program_id(0) cols = tl.arange(0, BLOCK_SIZE) mask = cols < n_cols x = tl.load(x_ptr + row * x_row_stride + cols, mask=mask, other=float("-inf")) x = x - tl.max(x, axis=0) # 数值稳定：减去最大值 num = tl.exp(x) den = tl.sum(num, axis=0) tl.store(y_ptr + row * y_row_stride + cols, num / den, mask=mask) def triton_softmax(x): """ 对二维张量的每一行做 softmax。 【参数】x : (n_rows, n_cols)，最后一维连续 【返回】同形状的 softmax 结果 """ assert x.dim() == 2 and x.stride(1) == 1 n_rows, n_cols = x.shape y = torch.empty_like(x) block = triton.next_power_of_2(n_cols) softmax_kernel[(n_rows,)](x, y, n_cols, x.stride(0), y.stride(0), BLOCK_SIZE=block) return y def naive_softmax(x): """ 朴素的 softmax（CS336 课上的例子），注释里是 M×N 矩阵每一步的显存读写： x_max = x.max(dim=1) 读 MN，写 M z = x - x_max[:, None] 读 MN + M，写 MN num = exp(z) 读 MN，写 MN den = num.sum(dim=1) 读 MN，写 M y = num / den[:, None] 读 MN + M，写 MN 合计：读 5MN + 2M，写 3MN + 2M （课件写的是读 5MN + M；最后一步读分母的 M 个数算不算进去，只是计数口径的差别。） """ x_max = x.max(dim=1).values z = x - x_max[:, None] num = torch.exp(z) den = num.sum(dim=1) return num / den[:, None] def demo_softmax(): """演示 1：融合 softmax。""" print_header("演示 1：融合 softmax —— 每个程序负责一行") print(f" 运行模式：{'真正的 Triton（GPU）' if REAL_TRITON else 'mini_triton 模拟器（CPU）'}") torch.manual_seed(0) M, N = 300, 781 # N 故意不是 2 的幂 x = torch.randn(M, N, device=DEVICE) * 5 if not REAL_TRITON: triton.reset_stats() y = triton_softmax(x) print(f" ({M}, {N}) 矩阵：与 torch.softmax 的最大差异 {(y - torch.softmax(x, dim=1)).abs().max().item():.1e}；" f"每行和为 1：{torch.allclose(y.sum(1), torch.ones(M, device=DEVICE))}") MN = M * N if not REAL_TRITON: s = triton.get_stats() print(f" 融合核函数：{s['programs']} 个程序（每行一个），读 {s['loaded']:,} = MN，写 {s['stored']:,} = MN") print(f" 朴素实现（估算）：读 5MN + 2M = {5 * MN + 2 * M:,}，写 3MN + 2M = {3 * MN + 2 * M:,}") print(f" → 融合后显存读写总量约为朴素实现的 {2 * MN / (8 * MN + 4 * M):.0%}，理论上约快 4 倍。") # 边界：-inf 填充的重要性 print(""" 【边界】为什么被掩掉的位置要填 -inf 而不是 0？ 如果填 0，这些位置会参与 max（当整行都是负数时，最大值会错误地变成 0）， 也会参与求和（exp(0) = 1），分母就错了。填 -inf 则 exp(-inf) = 0，对结果毫无影响。""") # ============================================================================= # 第 2 部分：分块行求和 # ============================================================================= @triton.jit def row_sum_kernel(x_ptr, out_ptr, n_cols, row_stride, BLOCK_SIZE: tl.constexpr): """ 每行求和，行可以比 BLOCK_SIZE 长：在程序内部一块一块地循环。 【思路】 acc 是一个长度为 BLOCK_SIZE 的累加器（在寄存器里）， 每次读入这一行的下一段 BLOCK_SIZE 个元素，逐位置加到 acc 上； 循环结束后，再把 acc 的 BLOCK_SIZE 个部分和加起来，得到整行的和。 【边界】最后一段可能不满，用 mask 并填 0（加 0 不影响求和）。 """ row = tl.program_id(0) acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32) for start in range(0, n_cols, BLOCK_SIZE): cols = start + tl.arange(0, BLOCK_SIZE) mask = cols < n_cols acc += tl.load(x_ptr + row * row_stride + cols, mask=mask, other=0.0) tl.store(out_ptr + row, tl.sum(acc, axis=0)) def triton_row_sum(x, block_size=1024): """二维张量按行求和。""" assert x.dim() == 2 and x.stride(1) == 1 out = torch.empty(x.shape[0], device=x.device, dtype=torch.float32) row_sum_kernel[(x.shape[0],)](x, out, x.shape[1], x.stride(0), BLOCK_SIZE=block_size) return out def demo_row_sum(): """演示 2：分块行求和。""" print_header("演示 2：行太长怎么办？—— 在程序内循环（分块）") torch.manual_seed(0) x = torch.randn(64, 5000, device=DEVICE) for block in [256, 1024, 8192]: out = triton_row_sum(x, block) loops = -(-5000 // block) print(f" BLOCK_SIZE = {block:>4}：每个程序循环 {loops:>2} 次，与 x.sum(1) 的最大差异 " f"{(out - x.sum(1)).abs().max().item():.1e}") print(""" BLOCK_SIZE 太小：循环次数多，每次读的数据少；太大：占用太多寄存器，占用率下降（第 5 讲）。 这种"在程序内沿一个维度循环、把部分结果累加在寄存器里"的模式，就是分块矩阵乘法（04 号文件） 和 FlashAttention（第 5 讲）的基本骨架。""") def main(): demo_softmax() demo_row_sum() print("\n全部演示运行完毕。") if __name__ == "__main__": main()