想象一家深夜食堂的老板:他最头疼的不是炒菜,而是“重新开火”——每来一桌客人,后厨就要从头准备一遍:穿围裙、调灶火、摆调料,等一切就绪,锅才真正热起来。GPU 编程的世界里也有同样的烦恼:每次调用一个 kernel(内核函数),GPU 都要走一遍启动流程,传递参数、调度线程,然后才轮到真正的计算。数据量小、调用频繁时,启动开销甚至比计算本身还贵。 “持久化 kernel”(persistent kernel)的思路,就是请几位大厨常驻厨房:一次开火,客人一桌一桌来,大厨们循环颠勺,把所有“重新准备”的功夫全省掉。Triton 是一门面向 GPU 的编程语言,它用很少的语法就能实现这种“驻场大厨”模式。这篇文章会用后厨的视角,把 persistent kernel 的设计原理、实现方法和适用场景一次讲透。
对 AI 算子优化工程师而言,持久化 kernel 是削减启动开销、提升 GPU 吞吐量的关键手段,掌握它能让高频小算子的性能获得数量级提升。
可视化演示
动态演示Triton持久化Kerne的核心机制,通过源头/管道/节点/终端四种模式展示关键概念与交互流程。
本节配有交互动效(进入视口后自动加载,约 3KB)。
Triton持久化Kerne:持久化k 持久化kernel(persi 持久化kernel(persi 核心机制可视化 Triton持久化Kerne整体架构与关键组件
内核启动开销
深夜食堂的老板最怕“重新开火”:新客人落座,后厨就得从头忙活一遍——穿围裙、验调料、等锅热,这一套流程还没炒出一粒米。GPU 上的普通 kernel 启动正是如此:每次调用内核,驱动程序都要完成参数传递、上下文准备、线程调度等固定动作,然后计算才开始。
在 Triton 中,一次 kernel 调用就对应一次完整的启动流程。启动开销是一笔与数据量几乎无关的固定成本:它包含参数传输、网格调度、缓存准备等环节,通常以微秒计。对于单次只算几百个元素的小算子,这笔开销常常比实际计算还大——就像为了炒一碟花生米,却烧了整整一炉大火。
当任务被拆成 N 个内核时,启动开销会随 N 线性累积。把每次启动开销记为 \(T_{启动}\)、每次计算时间记为 \(T_{计算}\),总耗时就是下面的式子。规律很直观:N 越大、\(T_{计算}\) 越小,启动开销在总耗时中的比重就越骇人。
所以,聪明的优化方向不是把炉火烧得更旺,而是想办法少开几次火——这正是持久化 kernel 的出发点。带着“开火成本”的账本,我们走进下一节,看看大厨们如何在厨房里“驻场”,把 N 次开火压缩成一次。
- T_总 总耗时整个任务从发出第一个 kernel 到全部完成的墙钟时间。
- N_内核 内核个数任务被拆成多少个独立 kernel,每个 kernel 都要重新走一遍启动流程。
- T_启动 启动开销与数据量几乎无关的固定成本,传递参数、调度线程都会消耗它。
- T_计算 计算时间真正在 SM 上执行加减乘除等运算所花的时间。
- 乘号 线性叠加启动开销随内核数量线性累积,小任务启动次数越多,浪费越惊人。
持久化设计思想
食堂老板琢磨出一个省钱的妙招:与其每桌客人重新叫一批厨师,不如让几位大厨签“驻场协议”——一直待在厨房里,客人一桌接一桌,大厨们循环颠勺,只在最后收工时离场。持久化 kernel 就是 GPU 世界的“驻场大厨”:程序启动一个 kernel,把网格尺寸(grid)设得足够小,让生成的线程块恰好能同时驻留在 GPU 上;这些线程块不退出,而是用一个循环把多段数据依次处理完。
从软件视角看,持久化 kernel 把“N 次 kernel 启动”压缩成了“1 次启动 + M 段循环计算”。每次循环里,线程块从全局内存搬运一小块数据,算完再搬下一块,就像大厨从传菜口一道道接单。由于线程块一直活着,寄存器、共享内存里的中间结果可以跨循环复用,省下的不止是启动开销。
把 N 次启动压缩成一次,节省下来的都是实打实的时间。但驻场大厨不是想请几位就请几位——厨房的灶台、橱柜和站脚空间都有上限。到底能驻留多少个线程块?下一节我们聊聊硬件资源的“天花板”。
- T_持久 持久化总耗时采用常驻线程块方案后,整个任务的总墙钟时间。
- T_启动 单次启动开销整个任务只需要支付一次内核启动成本,这是持久化的核心收益。
- M 循环轮数每个线程块循环处理的数据段数量,由总数据量除以块数再除以块大小决定。
- T_计算块 单段计算时间循环体内处理一段数据块(如一个 BLOCK_SIZE)所花的时间。
硬件资源约束
大厨们常驻厨房,首先得问一个问题:厨房装得下吗?一间后厨的灶台数量、橱柜大小、站立空间都是固定的,驻场大厨多了,后面的人只能排队等位。GPU 的每个 SM(流式多处理器,相当于一个独立灶间)同样资源有限:寄存器文件、共享内存、线程槽位,每一项都有硬上限。
如果一个线程块需要 32 个寄存器和 16KB 共享内存,而 SM 上只有 64K 寄存器——那它最多同时驻留 2 个块。持久化 kernel 的网格大小必须小于或等于所有 SM 能容纳的块数总和,否则“常驻”就成了空话——多余的块只能等前面的块跑完才能上 SM,调度反而更低效。
资源约束就像后厨的消防验收单,提醒我们持久化不是无脑把网格设小,而是要在“驻留得下”与“跑得满”之间取平衡。摸清 SM 的天花板之后,下一步就是动手写 Triton 代码——把“驻场协议”落到纸面上。
- N_驻留 驻留块数GPU 上能同时存活并执行的线程块数量。
- R_总 总寄存器数单个 SM 的寄存器文件总容量,例如 Ampere 架构为 65536 个 32 位寄存器。
- r_块 每块寄存器需求一个线程块所有线程需要的寄存器总量,由编译器根据代码压栈确定。
- S_总 总共享内存单个 SM 上可用的共享内存总量,通常为 100KB 量级。
- s_块 每块共享内存每个线程块申请的共享内存大小。
- min 取最小约束三个约束取最小值,就像厨房里既受灶台数限制,又受站立空间限制。
Triton语法实现
签了驻场协议还不够,还得有张清晰的“排班表”:几位大厨、每人负责哪几桌、按什么顺序上菜。Triton 写持久化 kernel 也只需要三步:把网格(grid)设成一个小到能驻留的整数;在 kernel 内部写一个 for 循环遍历多段数据;用 tl.constexpr 把循环段数作为编译期常量传进去。
网格大小 P 的选择很关键——它通常等于 GPU 上所有 SM 能同时驻留的线程块数,比如 8 个 SM 各驻留 1 块,P 就取 8。每个线程块循环处理的段数 M 由总数据量决定:总元素数除以“块数 × 块大小”。块太瘦,循环次数多、循环判断开销占比高;块太胖,占用率不足,SM 闲等。
排班表一旦写好,大厨们就按部就班地循环出菜。不过,任何管理制度都有适用边界——不是所有餐馆都适合请驻场大厨。什么时候该用、什么时候不该用,下一节我们算一笔经济账。
- N 总元素数需要被计算的全部数据元素个数。
- B 块内处理量一个线程块单次循环处理的元素数,即 Triton 中的 BLOCK_SIZE。
- P 驻留块总数持久化 kernel 的网格大小,应小于 SM 可同时容纳的块数。
- B×P 单轮总处理量全体线程块并行跑完一轮循环时共同处理的元素数。
- 取整向上取整,保证最后剩下的零星数据也有人处理,不丢元素。
适用场景判断
请驻场大厨不是免费的——虽然省了开火功夫,但大厨们的工位长期占着,别的活就插不进来。所以聪明的老板会先算一笔账:到底“重新开火”的花费占多大比例?如果一道菜要炖 3 小时,那开火那几分钟根本不值得操心;如果每位客人都只点一盘花生米,那开火功夫就比炒菜还贵,必须请驻场大厨。
GPU 上也一样:kernel 启动开销与计算时间的比值,是判断要不要上持久化 kernel 的标尺。当核函数体很小(比如一次向量加、一次逐元素激活),而调用频率很高(比如在训练循环里每步都跑)时,启动开销占比可能超过 50%,持久化改造收益巨大。反之,如果单次 kernel 已经在跑毫秒级的大矩阵乘法,启动开销占比不足 1%,就没必要折腾。
算清启动占比这笔账,你就能在“普通模式”与“驻场模式”之间从容选择。而一旦决定请驻场大厨,真正考验运营能力的细节才刚刚开始——负载均衡、工具分析,这些工程实践要点我们放到最后一节。
- R 占用比例启动开销占单次内核总时间的比重,衡量“开火”的浪费程度。
- 分子 T_启动固定启动开销,与数据量无关。
- 分母 总时间启动开销与计算时间之和。
- 比值含义R 大于 50% 时,说明一半以上时间没在算,应该考虑持久化或算子融合。
工程实践要点
真把驻场大厨请回家,运营细节开始冒头:按“前 100 桌归一号大厨”分单,一号忙到冒烟,二号却闲到擦灶台——这就是负载不均。GPU 上的惯用解法是“交错分配”:让每个线程块轮流取一块数据,第 t 轮取第 \(pid + t \times P\) 块,这样每轮所有人都忙,不会出现有的块早早收工、有的块还在苦熬的场面。
另一个要点是分清“持久化”与“融合”。算子融合(kernel fusion)把多个算子合并成一个 kernel,同样能减少启动次数——它解决的是“内核个数”问题;持久化解决的是“单个内核循环复用硬件”的问题。二者不但不冲突,还常常组合使用:先融合减少内核个数,再持久化让每个内核循环吃满数据。
最后,一切优化都要用数据说话。用 Nsight Compute 这类性能剖析工具测量启动耗时与 SM 占用率,改造前后分别记录墙钟时间做对比,只有当收益量化可见时,“请驻场大厨”这笔投入才算真正划算。
运营一家“驻场大厨”的后厨,讲究的是排班均衡、菜单设计与成本核算。把这套思路带回 GPU 世界,你手里的 Triton 就不再只是“每次重新点火”的普通灶台,而是一间随时待命、循环出菜的高效厨房。
- pid 当前块编号线程块在网格中的唯一编号,相当于大厨的工号。
- t 循环轮次当前处理的第几段数据,取 0 到 M-1。
- P 块总数驻留线程块的总数,作为轮次之间的步长。
- B 块大小每段数据的元素个数,决定偏移量的粒度。
- 交错分配同一轮中相邻的块取相邻的数据段,避免某一块独扛大量数据。
AI 应用场景
大模型训练中的高频小算子
阶段:算子开发与性能优化
解决的问题:在大模型训练循环里,LayerNorm、激活函数等小算子每步都执行,传统方式下每次 kernel 启动都伴随数微秒固定开销,累积起来浪费大量 GPU 时间。用 persistent kernel 让线程块常驻,在循环中反复处理每步数据,启动开销只付一次,训练吞吐显著提升。
具体案例:某团队训练 7B 参数 LLM,前向里逐元素算子单次仅 3μs 计算,启动却要 5μs;改造为 persistent kernel 后,这类算子累计耗时下降约 40%。
科学计算中的迭代求解器
阶段:高性能计算优化
解决的问题:流体模拟、电磁场求解等迭代算法,每一轮迭代都在相同形状的网格上执行相同计算。传统方式每轮都要重新启动 kernel;persistent kernel 一次启动、循环迭代,中间结果还可留在寄存器与共享内存中复用。
具体案例:二维波动方程模拟迭代 1000 步,网格 512×512,持久化改造后总时长下降 35%,主要收益来自省去 999 次启动及数据重载。
推荐系统在线推理中的小算子
阶段:推理服务优化
解决的问题:在线 CTR 预估每请求都要做 embedding 查表与聚合,运算量小但调用频繁。persistent kernel 让常驻线程块循环处理批量请求,配合低延迟推理框架,可明显降低单请求延迟并提高 GPU 利用率。
具体案例:广告推荐服务单日请求量过亿,将 embedding 聚合改为 persistent kernel 后,P99 延迟降低 28%,单卡承载 QPS 提升近一倍。
Python 演示
import tritonimport triton.language as tlimport torch @triton.jitdef persistent_add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr, STAGES: tl.constexpr): pid = tl.program_id(0) num_blocks = tl.num_programs(0) # 常驻线程块循环处理 STAGES 段数据, 交错取块保证负载均衡 for i in range(STAGES): offset = (pid + i * num_blocks) * BLOCK idx = offset + tl.arange(0, BLOCK) mask = idx < n x = tl.load(x_ptr + idx, mask=mask) y = tl.load(y_ptr + idx, mask=mask) tl.store(out_ptr + idx, x + y, mask=mask) def main(): n = 1 << 20 block = 1024 stages = 8 grid = (n // (block * stages),) # 网格固定为少量线程块, 常驻 GPU, 由内核内循环消化全部数据 x = torch.randn(n, device='cuda') y = torch.randn(n, device='cuda') out = torch.empty_like(x) persistent_add_kernel[grid](x, y, out, n, BLOCK=block, STAGES=stages) torch.cuda.synchronize() assert torch.allclose(out, x + y, atol=1e-6) print("persistent kernel 校验通过: 一次启动完成全部计算") if __name__ == "__main__": main()
grid = (n // (block * stages),)网格固定为少量线程块,让它小到足以全部驻留 GPU。for i in range(stages):kernel 内循环,一个常驻块连续处理多段数据,不再反复启动。num_blocks = tl.num_programs(0)获取网格块的数目,作为交错取块的步长。offset = (pid + i * num_blocks) * block按块编号交错取数据,保证负载均衡且互不重叠。mask = idx < n边界掩码,防止最后一段数据越界访问。BLOCK: tl.constexpr, STAGES: tl.constexpr编译期常量,让编译器做循环展开等优化。torch.cuda.synchronize()等待 GPU 全部执行完,再校验结果,避免异步误判。assert torch.allclose(out, x + y, atol=1e-6)用全量比对验证持久化 kernel 的正确性。
常见误区
误区我学会了 persistent kernel 后,恨不得把所有算子都改成循环模式,总觉得只要把循环写进 kernel,性能就一定会提升。
正解持久化 kernel 省的是启动开销,如果单次计算本身就很大,启动占比不足 1%,改造收益微乎其微。而且常驻线程块长期占用 SM 资源,可能挤掉其他并发任务。是否值得,必须先算启动占比这笔账。
为何常见:教程案例往往挑启动开销占比高的“典型样本”,读者容易以偏概全地把技巧当万能药。
误区普通 kernel 里 grid 越大并行度越高,那 persistent kernel 也把网格设得越大越好,块多了总归能多干活。
正解超过 SM 驻留上限的线程块根本无法同时执行,只能在队列里排队,反而增加调度压力。持久化 kernel 的网格应约等于 SM 可驻留块数,让每个块都“落地生根”循环干活。
为何常见:普通 kernel 的“块越多越快”观念根深蒂固,习惯性地被带入了持久化设计。
误区线程块都常驻在 GPU 上了,那我是不是可以让它们像共享内存一样直接互访数据?
正解线程块之间的数据隔离是硬件事实,块间没有全局同步原语。跨块协作必须借助原子操作或全局内存,稍不留神就会死锁。“常驻”不等于“共享”,这是两码事。
为何常见:把软件的“常驻”语义和硬件的“共享”能力混为一谈,忽视了线程块彼此隔离的架构设计。
误区既然 persistent kernel 能减少启动开销,那 kernel fusion 就不需要了,一个就够。
正解两者削减启动的路径不同:融合减少内核个数,持久化让单个内核循环复用硬件。实践中常常先融合减少内核,再持久化吃满数据,是互补关系而非替代关系。
为何常见:两者都能降低启动开销,表面效果相似,容易让人误以为功能重复。
面试题
Q1. persistent kernel 的核心设计思想是什么?
- A.让内核在 GPU 上反复重新启动
- B.让线程块驻留并循环处理数据← 正确答案
- C.通过扩大 L2 缓存区提升性能
- D.减少线程数量降低资源占用率
答案:B。持久化 kernel 的关键是线程块常驻 SM,通过 for 循环依次处理多段数据,把 N 次启动压成 1 次。A 是传统多次启动;C 属于缓存优化;D 与持久化无直接关系。选 B。
Q2. 关于 kernel 启动开销,下列说法正确的是?
- A.开销主要发生在数值计算当中
- B.每次启动都有固定调度开销← 正确答案
- C.启动开销与数据量大小成正比
- D.启动开销小到几乎可以忽略
答案:B。启动开销是与数据量无关的固定成本,包含参数传递、线程调度、缓存准备等,通常为微秒级。A、C、D 均与事实不符,故选 B。
Q3. 下面哪些属于决定 SM 驻留线程块数量的硬件资源?
- A.寄存器文件的总量上限← 正确答案
- B.共享内存的容量大小限制← 正确答案
- C.硬盘转速影响内核驻留
- D.每个 SM 的线程驻留上限← 正确答案
答案:ABD。驻留块数由寄存器、共享内存和线程槽位共同决定。硬盘转速与 GPU 内核驻留无关,C 错误。故答案为 A、B、D。
Q4. Triton 中把每个线程块要循环处理的段数传给内核,通常用什么方式?
- A.用 tl.constexpr 参数声明段数← 正确答案
- B.在 tl.load 加载时指定段数
- C.通过 program_id 返回值指定
- D.在启动配置中指定网格段数
答案:A。循环段数是编译期常量,适合用 tl.constexpr 传入,便于编译器展开优化。B 的 load 只负责加载数据,C 返回块编号,D 指定网格大小,都不是传段数的正确方式。
Q5. 持久化 kernel 中,如何避免两个线程块处理同一段数据?
- A.使用原子加锁保护每块数据
- B.用块编号约束循环访问偏移← 正确答案
- C.用共享内存标记数据属主
- D.增大网格尺寸避免数据冲突
答案:B。最常用的是用 pid 构造不相交的偏移,如 offset=(pid+t×P)×B,使每段数据只属于一个块。A 是动态取任务的另类技巧,C 中块间共享内存不可见,D 不能消除冲突。
Q6. 持久化 kernel 能设置的网格大小上限取决于什么?
- A.任务需要的数据块总量大小
- B.寄存器与共享内存的资源量← 正确答案
- C.输入数据的浮点精度高低
- D.编译器优化选项版本差异
答案:B。网格上限由硬件的资源天花板决定——寄存器总量与共享内存容量限制了每个 SM 能驻留的块数,所有 SM 的驻留能力之和就是网格上限。A 影响循环段数而非网格上限,C、D 与驻留无关。
Q7. 用 Triton 编写持久化 kernel 通常需要做哪些事?
- A.网格设为少量常驻线程块← 正确答案
- B.内核内部添加循环处理段← 正确答案
- C.用 constexpr 指定循环段数← 正确答案
- D.必须依赖内联 CUDA 汇编
答案:ABC。三步走:网格设小、kernel 内 for 循环、constexpr 传段数。Triton 是高层 DSL,无需手写 CUDA 汇编,D 错误。
Q8. 持久化 kernel 中跨线程块协作可能死锁,根本原因是什么?
- A.数据量太大导致显存分配失败
- B.块间缺少全局同步会导致互等← 正确答案
- C.共享内存过小导致数据溢出错误
- D.网格太大触发看门狗超时重置
答案:B。线程块之间没有内置全局同步原语,若用自旋锁等机制协作,两个块互相等待对方释放资源就可能永久互等,形成死锁。A、C、D 属于资源与稳定性问题,不是死锁的根因。
Q9. 某 kernel 启动开销 4μs、计算 1μs,需跑 100 次。改为持久化后总耗时约为?
- A.从五百微秒降到一百零四微秒← 正确答案
- B.从五百微秒降到四百微秒
- C.从四百微秒降到两百微秒
- D.从一百零四微秒降到五微秒
答案:A。传统总耗时=100×(4+1)=500μs;持久化后只付一次启动,总耗时=4+100×1=104μs。选 A。B 只省了 100μs,说明把启动误算成 1μs;C、D 的初始值都不对。
Q10. 下列哪些是持久化 kernel 工程落地时需要注意的实践要点?
- A.交错分配数据实现负载均衡← 正确答案
- B.与算子融合搭配组合使用← 正确答案
- C.用性能剖析工具实测收益← 正确答案
- D.网格越大驻留效果才越好
答案:ABC。交错取块保均衡、与融合互补、用 Nsight 实测,都是正确实践。D 与驻留约束相悖,块数超过上限只会排队,错误。