想象你是这家 GPU 大饭店的厨房总管。食材仓库(HBM)在楼下,备菜案板(SRAM)就在手边,而炉台上的每位厨师都算得飞快。可老厨师有个坏习惯:每做好一步,都要把半成品端回楼下仓库,下一步再跑下去端上来。炒一盘 Linear 配 GELU,要跑三四趟;调 Attention 的三款酱料 Q、K、V,更是来回折腾。菜没出几道,腿先跑细了——整个后厨的瓶颈,从来不在火候,而在跑腿。 Triton 算子融合,就是给后厨立的一条新规矩:能留在案板上做完的,绝不下楼。把 Linear 和 GELU 放进同一个 kernel 一口锅炒,中间结果留在 SRAM;把 Q、K、V 三次投影合成一次大矩阵乘法,调好酱料再分装。省掉的,正是 GPU 上最贵的东西——访存带宽。 这篇文章会从“为什么要融合”讲起,带你认识 Triton 的块级编程思维,亲手实现 Linear+GELU 融合,再攻下 Attention 的 QKV 投影融合与在线 Softmax,最后给出工程验收清单。备好围裙,我们开始。
算子融合是 AI 编译与性能优化的核心杠杆:理解它,你才能把同样的 GPU 资源榨出成倍的吞吐,是算子开发工程师从“能跑”迈向“跑得快”的必修课。
可视化演示
本节配有交互动效(进入视口后自动加载,约 3KB)。
算子融合前后对比:访存开销变化 \u2717 未融合 (3个独立kernel)\u2193 6次访存\u2713 融合后 (1个kernel)\u2717 未融合 (3个独立kernel)\u2717 未融合 (3个独立kernel)\u2193 6次访存\u2713 融合后 (1个kernel)\u2717 未融合 (3个独立kernel) \u2717 未融合 (3个独立kernel)\u2193 6次访存\u2713 融合后 (1个kernel)\u2717 未融合 (3个独立kern...
融合背景动机
老厨师做青椒炒肉,流程是这样的:下楼取肉,切好,端回仓库;下楼取青椒,切好,端回仓库;再把两份半成品端上来,下锅。每一步“端盘子”都要穿过又长又挤的走廊。GPU 里的走廊就是 HBM 总线,带宽再大,也架不住来回倒腾。
计算单元在芯片上飞速运转,但数据从 HBM 搬进寄存器,走的是片外总线,延迟和能耗是片上 SRAM 的几十倍。如果两个算子之间存在大块中间结果,每写一次 HBM 再读一次,就等于白跑两趟仓库。
把时间开销写成公式,就一目了然:
所以融合的本质不是少算,而是少搬。Linear+GELU 融合后,中间结果从 HBM 的写读两次,变成 SRAM 上的原地变换;Attention 融合后,QKV 三次投影合并为一次。跑腿少了,上菜自然就快。
- T — 搬运总耗时这是我们要优化的目标。融合前后,火候(计算量)几乎不变,变的只有 T。
- D — 搬运数据量所有中间结果进出 HBM 的字节总数。融合把这个分子压到最小。
- B — 访存带宽走廊每秒能通过的人数,由硬件决定。我们能改变的是 D,而不是 B。
Triton核心概念
新主厨 Triton 拿到一本《块级菜谱》:他不用操心整个饭店有几张桌子,只需写清楚“一桌菜”怎么做——比如一桌四人的分量怎么切配、怎么下锅。写好后,饭店自动把这套流程复制到每一桌。这就是块级编程(Block-Level Programming)。
在 Triton 里,我们自己写的这个“一桌菜”叫 kernel;每一桌是一个 program,由 program_id 唯一标识。核心里用 tl.arange 生成块内编号,用 tl.load / tl.store 读写数据,用 tl.dot 做矩阵乘。编译器负责把块内的循环变成 GPU 上的向量化指令,并自动调度到所有计算单元上。
一桌处理多大的块,用网格公式描述:
有了“一桌菜”的思维,融合就变得自然:你只需要在这一桌的案板(SRAM)上连续完成多道工序,中间不换桌、不下楼。
- M / N — 输出规模输出矩阵的总行数和总列数,对应整个后厨要服务的桌数。
- B_M / B_N — 块大小每个 program 一次处理的行数和列数,对应一桌坐几位客人。
- 向上取整⌈⌉ 表示有余数就多开一桌,确保所有元素都有程序负责。
- 总 program 数 G两者的乘积就是整个 grid 的规模,GPU 会把它铺到所有 SM 上执行。
Linear算子剖析
Linear 层在神经网络里,相当于“调底料”:输入 X 是成筐的食材,权重 W 是配方表,偏置 b 是最后的点睛之盐。每一份输出,都是把某一行食材和一列配方逐项相乘再累加的结果。
Triton 内核里,我们把输出划分成 BLOCK_M x BLOCK_N 的块,每个 program 负责一块;外层沿输出网格铺开,内层按 BLOCK_K 分小段,用 tl.dot 计算并累加。累加器 acc 在寄存器里一直接力,直到 K 维全部算完。
底料调好了,传统做法是盛出来送回仓库,等下一个算子 GELU 再端上来。而我们要做的,是让它留在锅里直接调味——这就是下一节的 GELU 融合。
- X_{m,k} — 输入元素第 m 行输入的第 k 个特征,来自上一层激活后的结果。
- W_{n,k} — 权重元素第 n 个输出神经元对应的第 k 个权重。
- K 次累加把所有 k 从 0 到 K-1 的乘积加起来,这就是一次内积。
- b_n — 偏置给第 n 个输出加上一个可学习的常数项。
GELU激活融合
GELU 是一个“看人下菜”的调味师:输入为正,基本原样保留;输入为负,狠狠压小;输入接近 0,适当削弱。它不像 ReLU 那样一刀切(负的全变 0),而是平滑地筛出“这条食材值得保留”的概率。
由于 GELU 是逐元素的纯函数(输入一个数,输出一个数),它不需要跨行跨列的“社交”,是最适合融合的调味料。把它直接接在 Linear 的累加器后面,对 acc 的每个元素套公式即可。
在代码里这只是一行:acc = acc * 0.5 * (1.0 + tl.erf(acc * 0.70710678))。累加完成立刻调味,锅都不用挪。如果拆成两个 kernel,这一行前后就是两次 HBM 往返——你算算省了多少趟。
- x³ 微小修正括号里的 0.044715x³ 是对高斯累积函数的泰勒近似补充。
- tanh 压缩把修正后的值压到 (-1, 1) 区间,模拟正态分布的尾部形状。
- 1 + tanh(……)把区间平移到 (0, 2),相当于一个“保留概率”的未缩放版本。
- 0.5x 整体缩放最后乘以 x 本身,保证输入越大、保留越多。
QKV投影融合
Attention 需要三种酱料:查询 Q、键 K、值 V。对同一个输入 X,憨厚的旧做法是分三次投影:先调 Q,送回仓库;再调 K,送回仓库;最后调 V。三次都从同一个 X 出发,等于把一份底料端下楼三次。
聪明的主厨会把三张配方表横向拼成一张大表 W_qkv,一次矩阵乘同时得到 Q、K、V 三个酱料桶,再按列切分。由于 GPU 的矩阵乘单元天然擅长大批量计算,合并后三个小矩阵乘变成一个大矩阵乘,利用率往往更高。
切分好的 Q、K、V 同样不需要下楼——它们就在案板(SRAM)上,直接进入下一道注意力工序。
- W_qkv — 拼接权重把 W_q、W_k、W_v 按列拼接得到的权重大矩阵。
- 一次矩阵乘X 与 W_qkv 转置相乘,得到 (M, 3d) 的拼接结果,三份投影共用一次数据搬运。
- b_qkv — 拼接偏置三个偏置向量拼接后,同样一次广播加完。
- 按列切片把结果切成 Q、K、V 三份,分别对应前 d、中 d、后 d 列。
注意力内核融合
注意力就像给每位客人上菜前做一次“搭配确认”:Q 问“这道菜想要什么”,K 说“这道菜有什么”,两者一比得到相似度,再用 softmax 归一化成权重,最后按权重把 V(食材)组合起来。关键难点是:完整 softmax 需要看到整行所有分数,而块式计算一次只看一部分。
FlashAttention 的思路是“边尝边调整”:每算完一个分块,用当前最大值更新运行最大值 m,并用指数缩放修正旧的累加和,最后统一除。这样既不需要先算出完整的 QKᵀ 矩阵,也无需把整行分数写回 HBM,正是融合思想在注意力上的应用。
配合在线 softmax 的运行最大值技巧,即使完整的 QKᵀ 矩阵永远不在 HBM 里出现,数值结果也和标准实现完全一致。这就像厨师边尝边加盐——不需要等所有菜上齐才定咸淡。
- QKᵀ — 相似度矩阵查询与键两两内积,得到 (M, N) 的得分表。
- 除以 √d_k防止维度大时点积数值过大,把 softmax 推向梯度消失的饱和区。
- softmax 归一化把每行得分变成和为 1 的权重。
- 乘以 V按权重把值向量加权累加,得到最终输出。
工程实践验收
新菜谱好不好,得上菜计时。工程上,融合 kernel 上线前要做三件事:测速、调参、验精度。缺一不可,否则可能“看起来快了,结果错了”。
测速用端到端延迟对比;调参主要调节 BLOCK_M / BLOCK_N / BLOCK_K 三个旋钮,让块大小与寄存器压力平衡;精度验证则把融合 kernel 与 PyTorch 参考实现逐元素比对,统计最大绝对误差。
回到后厨:同样的菜,出锅时间从三趟仓库变成一趟案板。指标不会说谎——把加速比、误差报告贴在墙上,整个团队都对“一锅出”心服口服。融合不是一个魔法,而是一套可测量、可验证的工程方法。
- T_before — 融合前耗时未融合的基线版本端到端耗时。
- T_after — 融合后耗时融合成单个 kernel 后的端到端耗时。
- 加速比大于 1 说明融合生效;一般 Linear+GELU 融合可拿到 1.2-2x 的收益。
AI 应用场景
LLM 推理服务 Prefill 加速
阶段:模型推理部署
解决的问题:大模型生成每个 token 都要执行 QKV 投影和注意力计算。传统做法中,投影结果矩阵频繁写回 HBM,显存带宽被吃满,kernel 启动次数多达七八次。将 QKV 投影融合为一次大矩阵乘,再把 Attention 用在线 Softmax 分块执行后,中间矩阵不再落地,首 token 延迟显著下降。
具体案例:部署 7B 参数模型时,输入 2048 个 token 的 prefill 阶段,QKV 投影融合配合分块注意力,使首 token 延迟从约 1.2 秒降至 0.8 秒,吞吐提升约 30%。这就是每个 LLM 推理服务都值得做的“案板革命”。
高性能算子库内核开发
阶段:算子开发与优化
解决的问题:框架自带算子覆盖通用场景,但对融合 GELU 的多层感知机块、动态形状的注意力等特殊结构,存在明显冗余访存。手写 Triton 融合 kernel 能把“仓库往返”变成“案板流转”。
具体案例:为训练框架开发 fused Linear+GELU kernel,在 2000 次迭代的 MLP 基准上,单卡训练吞吐提升 18%,同时代码量只有 CUDA 方案的十分之一。该 kernel 被封装成自定义算子,经自动测试后合入内部算子库。
Python 演示
import torchimport tritonimport triton.language as tl @triton.jitdef linear_gelu_fused_kernel( X_ptr, W_ptr, B_ptr, Y_ptr, M, N, K, stride_xm, stride_xk, stride_wn, stride_wk, stride_ym, stride_yn, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,): pid_m = tl.program_id(0) pid_n = tl.program_id(1) offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) offs_k = tl.arange(0, BLOCK_K) x_ptrs = X_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk w_ptrs = W_ptr + offs_n[None, :] * stride_wn + offs_k[:, None] * stride_wk acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) for k in range(0, tl.cdiv(K, BLOCK_K)): x = tl.load(x_ptrs) w = tl.load(w_ptrs) acc = tl.dot(x, w, acc) x_ptrs += BLOCK_K * stride_xk w_ptrs += BLOCK_K * stride_wk b = tl.load(B_ptr + offs_n) acc = acc + b[None, :] acc = acc * 0.5 * (1.0 + tl.erf(acc * 0.70710678)) y_ptrs = Y_ptr + offs_m[:, None] * stride_ym + offs_n[None, :] * stride_yn tl.store(y_ptrs, acc) def linear_gelu_fused(x, w, b): M, K = x.shape N = w.shape[0] y = torch.empty((M, N), device=x.device, dtype=x.dtype) grid = (triton.cdiv(M, 64), triton.cdiv(N, 64)) linear_gelu_fused_kernel[grid]( x, w, b, y, M, N, K, x.stride(0), x.stride(1), w.stride(0), w.stride(1), y.stride(0), y.stride(1), BLOCK_M=64, BLOCK_N=64, BLOCK_K=32, ) return y if __name__ == "__main__": torch.manual_seed(0) M, K, N = 256, 128, 256 x = torch.randn((M, K), device="cuda", dtype=torch.float32) w = torch.randn((N, K), device="cuda", dtype=torch.float32) b = torch.randn((N,), device="cuda", dtype=torch.float32) y = linear_gelu_fused(x, w, b) ref = torch.nn.functional.gelu(x @ w.T + b) print("max abs err:", (y - ref).abs().max().item())
@triton.jit把 Python 函数编译成 GPU kernel 的装饰器,相当于“一桌菜”的菜谱。tl.program_id(0)拿到当前 program 在网格中的坐标,即桌号。tl.arange(0, BLOCK_M)生成块内的行编号,对应这一桌的座位号。acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)初始化累加器,底料锅先烧热。acc = tl.dot(x, w, acc)分块矩阵乘并累加,内层循环沿 K 维接力。b = tl.load(B_ptr + offs_n)加载偏置向量,相当于最后的点睛之盐。acc = acc * 0.5 * (1.0 + tl.erf(acc * 0.70710678))GELU 累加后原地调味,中间结果不落地。tl.store(y_ptrs, acc)只写一次最终结果,不再为中间张量跑仓库。grid = (triton.cdiv(M, 64), triton.cdiv(N, 64))按 64 的块算出网格大小,多开一桌补余数。
常见误区
误区我以为算子融合就是在 Python 里先算 Linear 再算 GELU,把两段计算逻辑写在一起就算融合了,代码看着紧凑就行。
正解真正的融合是让中间结果留在 SRAM 或寄存器里,不写回 HBM。就算代码写在同一个 kernel 中,如果中间结果仍然落回全局内存,性能不会有本质提升。融合的关键是数据流,不是代码排版。
为何常见:名字里的“融合”容易让人联想到代码层面的合并,而忽略了它真正的目标是减少访存。
误区我总觉得 BLOCK_M 和 BLOCK_N 越大,一个块处理的活越多,线程利用率越高,性能就一定越好。
正解块太大会撑爆寄存器和共享内存,导致 occupancy 下降,甚至无法编译。要参考目标 GPU 的硬件规格,让块体积与资源容量匹配。案板就那么大,硬塞一吨食材只会放不下。
为何常见:“一次多干点”的直觉很自然,但忽略了 GPU 资源上限这个硬约束。
误区PyTorch 2.0 不是能自动融合算子吗?那我还花时间学 Triton 干嘛,交给编译器就好了。
正解torch.compile 对常见图模式有效,但对涉及切片重组的 QKV 投影、跨自定义算子的复杂融合,自动编译器的收益有限。手写 Triton 能精确控制块大小、数据布局和融合边界,是高级性能优化的必经之路。
为何常见:宣传材料常夸大自动编译器的适用范围,实际项目里“编译器没融好”的案例比比皆是。
误区GELU 有 erf 和 tanh 两种近似形式,我觉得算出来差不多,随便挑一个用,反正都是激活函数。
正解两种形式的数值差异大约在 1e-3 量级,训练时可接受;但推理时如果做精度对齐或量化校准,必须与参考实现保持一致。融合 kernel 里选哪种公式,直接影响端到端误差报告。
为何常见:误差在浮点上看起来很小,容易忽略“基准不一致”会引发的调试灾难。
面试题
Q1. 算子融合带来的最主要收益是什么?
- A.减少中间结果的访存开销← 正确答案
- B.减少模型的训练参数数目
- C.让代码结构显得更加简洁
- D.提升数据类型的表示精度
答案:A。融合最核心的收益是减少中间结果在 HBM 与计算单元之间的往返,对应后厨里“少跑仓库”。参数总量没变,代码风格只是副产物,精度也未改变,因此 B、C、D 都不对。
Q2. Linear 与 GELU 融合后,中间结果存储在哪里?
- A.被写回 HBM 全局内存等待
- B.留在 SRAM 供下一步使用← 正确答案
- C.暂存在 CPU 主存队列里
- D.直接丢弃再重新计算一次
答案:B。融合后 Linear 的累加结果留在 SRAM(案板)上,GELU 直接读取并原地变换。A 是未融合的旧行为,C 跨设备搬运更慢,D 丢弃数据会导致结果错误,故选 B。
Q3. 以下哪些场景适合做算子融合?
- A.Linear 输出被 GELU 消费← 正确答案
- B.QKV 三个投影共用输入← 正确答案
- C.两个算子间数据量巨大← 正确答案
- D.两个算子完全互不依赖
答案:ABC。融合适合“生产者—消费者”紧邻且中间数据量大的情况。A 是典型配对,B 共用输入可合并成一次投影,C 数据量越大省得越多;D 无依赖的算子合并没有收益,反而可能增加寄存器压力,不选。
Q4. Attention 中除以根号 d_k 的主要目的是什么?
- A.防止 QK 点积过大进入饱和区← 正确答案
- B.使点积结果保持精确为零
- C.增加模型参数量提升表达力
- D.减少键矩阵的显存占用大小
答案:A。点积随维度线性增大,除以根号 d_k 把它压回单位量级,避免 softmax 进入指数饱和区导致梯度消失。B、C、D 都与缩放的真实目的无关,故选 A。
Q5. Triton 中块大小参数的选择需要考虑哪些因素?
- A.GPU 寄存器与共享内存容量← 正确答案
- B.输入矩阵的具体形状尺寸← 正确答案
- C.目标 GPU 的硬件架构版本← 正确答案
- D.训练优化器的学习率大小
答案:ABC。块大小受硬件资源上限(寄存器/SRAM)、输入形状与架构特性共同约束,案板能放多少得看案板本身。学习率是训练超参数,与 kernel 的网格划分无关,D 不选。
Q6. FlashAttention 中在线 Softmax 解决的核心问题是什么?
- A.无需保存完整分数矩阵即可归一化← 正确答案
- B.让注意力分数结果更加精确
- C.把 Softmax 替换成线性激活函数
- D.消除 QKV 三个矩阵的参数冗余
答案:A。在线 Softmax 用运行最大值和运行累加和修正,使分块计算也能得到与全局 softmax 完全一致的结果,因此不需要把完整 (M, N) 分数矩阵留在 HBM。B、C、D 都不是它要解决的核心问题。
Q7. 融合后 GELU 计算发生在哪个时刻?
- A.线性累加完成之后立刻执行← 正确答案
- B.输入数据拷贝完成之前执行
- C.权重矩阵初始化过程中执行
- D.结果写回 HBM 之后才执行
答案:A。融合内核中 GELU 在 acc 完成所有 K 维累加并加上偏置后、写回 Y 之前执行,即“累加完立刻调味”。B、C 在计算前无从激活,D 写回后执行已失去融合意义,故选 A。
Q8. 验证融合算子正确性可以采用哪些方法?
- A.与 PyTorch 参考实现逐元素比对← 正确答案
- B.检查最大绝对误差是否可接受← 正确答案
- C.用多组随机输入统计误差分布← 正确答案
- D.直接查看代码缩进是否工整
答案:ABC。工程上融合 kernel 必须与参考实现逐元素比对、检查最大误差,并用多组随机输入验证误差分布稳定。D 只关乎代码风格,无法验证数值正确性,不选。