想象一家深夜食堂的老板:他最头疼的不是炒菜,而是“重新开火”——每来一桌客人,后厨就要从头准备一遍:穿围裙、调灶火、摆调料,等一切就绪,锅才真正热起来。GPU 编程的世界里也有同样的烦恼:每次调用一个 kernel(内核函数),GPU 都要走一遍启动流程,传递参数、调度线程,然后才轮到真正的计算。数据量小、调用频繁时,启动开销甚至比计算本身还贵。 “持久化 kernel”(persistent kernel)的思路,就是请几位大厨常驻厨房:一次开火,客人一桌一桌来,大厨们循环颠勺,把所有“重新准备”的功夫全省掉。Triton 是一门面向 GPU 的编程语言,它用很少的语法就能实现这种“驻场大厨”模式。这篇文章会用后厨的视角,把 persistent kernel 的设计原理、实现方法和适用场景一次讲透。

岗位关联

对 AI 算子优化工程师而言,持久化 kernel 是削减启动开销、提升 GPU 吞吐量的关键手段,掌握它能让高频小算子的性能获得数量级提升。

目录
  1. 可视化演示
  2. 内核启动开销
  3. 持久化设计思想
  4. 硬件资源约束
  5. Triton语法实现
  6. 适用场景判断
  7. 工程实践要点
  8. AI 应用场景
  9. Python 演示
  10. 常见误区
  11. 面试题

可视化演示

动态演示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总=N内核×(T启动+T计算)\displaystyle T_{总} = N_{内核} \times (T_{启动} + T_{计算})
总耗时等于内核数量乘以“启动加计算”。内核启动得越频繁,被“烧炉子”浪费掉的时间就越多。
  1. T_总 总耗时整个任务从发出第一个 kernel 到全部完成的墙钟时间。
    T总T_{总}
  2. N_内核 内核个数任务被拆成多少个独立 kernel,每个 kernel 都要重新走一遍启动流程。
    N内核N_{内核}
  3. T_启动 启动开销与数据量几乎无关的固定成本,传递参数、调度线程都会消耗它。
    T启动T_{启动}
  4. T_计算 计算时间真正在 SM 上执行加减乘除等运算所花的时间。
    T计算T_{计算}
  5. 乘号 线性叠加启动开销随内核数量线性累积,小任务启动次数越多,浪费越惊人。
    ×\times

持久化设计思想

食堂老板琢磨出一个省钱的妙招:与其每桌客人重新叫一批厨师,不如让几位大厨签“驻场协议”——一直待在厨房里,客人一桌接一桌,大厨们循环颠勺,只在最后收工时离场。持久化 kernel 就是 GPU 世界的“驻场大厨”:程序启动一个 kernel,把网格尺寸(grid)设得足够小,让生成的线程块恰好能同时驻留在 GPU 上;这些线程块不退出,而是用一个循环把多段数据依次处理完。

从软件视角看,持久化 kernel 把“N 次 kernel 启动”压缩成了“1 次启动 + M 段循环计算”。每次循环里,线程块从全局内存搬运一小块数据,算完再搬下一块,就像大厨从传菜口一道道接单。由于线程块一直活着,寄存器、共享内存里的中间结果可以跨循环复用,省下的不止是启动开销。

把 N 次启动压缩成一次,节省下来的都是实打实的时间。但驻场大厨不是想请几位就请几位——厨房的灶台、橱柜和站脚空间都有上限。到底能驻留多少个线程块?下一节我们聊聊硬件资源的“天花板”。

T持久=T启动+M×T计算块\displaystyle T_{持久} = T_{启动} + M \times T_{计算块}
对比上一节的公式:启动开销从出现 N 次降为只出现 1 次,代价是循环多跑 M 轮。
  1. T_持久 持久化总耗时采用常驻线程块方案后,整个任务的总墙钟时间。
    T持久T_{持久}
  2. T_启动 单次启动开销整个任务只需要支付一次内核启动成本,这是持久化的核心收益。
    T启动T_{启动}
  3. M 循环轮数每个线程块循环处理的数据段数量,由总数据量除以块数再除以块大小决定。
    MM
  4. T_计算块 单段计算时间循环体内处理一段数据块(如一个 BLOCK_SIZE)所花的时间。
    T计算块T_{计算块}

硬件资源约束

大厨们常驻厨房,首先得问一个问题:厨房装得下吗?一间后厨的灶台数量、橱柜大小、站立空间都是固定的,驻场大厨多了,后面的人只能排队等位。GPU 的每个 SM(流式多处理器,相当于一个独立灶间)同样资源有限:寄存器文件、共享内存、线程槽位,每一项都有硬上限。

如果一个线程块需要 32 个寄存器和 16KB 共享内存,而 SM 上只有 64K 寄存器——那它最多同时驻留 2 个块。持久化 kernel 的网格大小必须小于或等于所有 SM 能容纳的块数总和,否则“常驻”就成了空话——多余的块只能等前面的块跑完才能上 SM,调度反而更低效。

资源约束就像后厨的消防验收单,提醒我们持久化不是无脑把网格设小,而是要在“驻留得下”与“跑得满”之间取平衡。摸清 SM 的天花板之后,下一步就是动手写 Triton 代码——把“驻场协议”落到纸面上。

N驻留=min⁡(R总r块,S总s块,N上限)\displaystyle N_{驻留} = \min\left( \frac{R_{总}}{r_{块}}, \frac{S_{总}}{s_{块}}, N_{上限} \right)
可驻留的线程块数 = 寄存器资源、共享内存资源、网格上限三者中的最小值。短板决定你能请几位“驻场大厨”。
  1. N_驻留 驻留块数GPU 上能同时存活并执行的线程块数量。
    N驻留N_{驻留}
  2. R_总 总寄存器数单个 SM 的寄存器文件总容量,例如 Ampere 架构为 65536 个 32 位寄存器。
    R总R_{总}
  3. r_块 每块寄存器需求一个线程块所有线程需要的寄存器总量,由编译器根据代码压栈确定。
    r块r_{块}
  4. S_总 总共享内存单个 SM 上可用的共享内存总量,通常为 100KB 量级。
    S总S_{总}
  5. s_块 每块共享内存每个线程块申请的共享内存大小。
    s块s_{块}
  6. min 取最小约束三个约束取最小值,就像厨房里既受灶台数限制,又受站立空间限制。
    min⁡(⋯ )\min(\cdots)

Triton语法实现

签了驻场协议还不够,还得有张清晰的“排班表”:几位大厨、每人负责哪几桌、按什么顺序上菜。Triton 写持久化 kernel 也只需要三步:把网格(grid)设成一个小到能驻留的整数;在 kernel 内部写一个 for 循环遍历多段数据;用 tl.constexpr 把循环段数作为编译期常量传进去。

网格大小 P 的选择很关键——它通常等于 GPU 上所有 SM 能同时驻留的线程块数,比如 8 个 SM 各驻留 1 块,P 就取 8。每个线程块循环处理的段数 M 由总数据量决定:总元素数除以“块数 × 块大小”。块太瘦,循环次数多、循环判断开销占比高;块太胖,占用率不足,SM 闲等。

排班表一旦写好,大厨们就按部就班地循环出菜。不过,任何管理制度都有适用边界——不是所有餐馆都适合请驻场大厨。什么时候该用、什么时候不该用,下一节我们算一笔经济账。

M=⌈NB×P⌉\displaystyle M = \left\lceil \frac{N}{B \times P} \right\rceil
M 是每个常驻线程块要循环处理的段数:总元素数 N 除以单块单次处理量(B×P)后向上取整。数据再多,也只是让每个大厨多颠几勺。
  1. N 总元素数需要被计算的全部数据元素个数。
    NN
  2. B 块内处理量一个线程块单次循环处理的元素数,即 Triton 中的 BLOCK_SIZE。
    BB
  3. P 驻留块总数持久化 kernel 的网格大小,应小于 SM 可同时容纳的块数。
    PP
  4. B×P 单轮总处理量全体线程块并行跑完一轮循环时共同处理的元素数。
    B×PB \times P
  5. 取整向上取整,保证最后剩下的零星数据也有人处理,不丢元素。
    ⌈⋯ ⌉\lceil \cdots \rceil

适用场景判断

请驻场大厨不是免费的——虽然省了开火功夫,但大厨们的工位长期占着,别的活就插不进来。所以聪明的老板会先算一笔账:到底“重新开火”的花费占多大比例?如果一道菜要炖 3 小时,那开火那几分钟根本不值得操心;如果每位客人都只点一盘花生米,那开火功夫就比炒菜还贵,必须请驻场大厨。

GPU 上也一样:kernel 启动开销与计算时间的比值,是判断要不要上持久化 kernel 的标尺。当核函数体很小(比如一次向量加、一次逐元素激活),而调用频率很高(比如在训练循环里每步都跑)时,启动开销占比可能超过 50%,持久化改造收益巨大。反之,如果单次 kernel 已经在跑毫秒级的大矩阵乘法,启动开销占比不足 1%,就没必要折腾。

算清启动占比这笔账,你就能在“普通模式”与“驻场模式”之间从容选择。而一旦决定请驻场大厨,真正考验运营能力的细节才刚刚开始——负载均衡、工具分析,这些工程实践要点我们放到最后一节。

R=T启动T启动+T计算\displaystyle R = \frac{T_{启动}}{T_{启动} + T_{计算}}
启动开销占比 R 越接近 1,说明时间越浪费在“开火”上,持久化 kernel 的价值越大。
  1. R 占用比例启动开销占单次内核总时间的比重,衡量“开火”的浪费程度。
    RR
  2. 分子 T_启动固定启动开销,与数据量无关。
    T启动T_{启动}
  3. 分母 总时间启动开销与计算时间之和。
    T启动+T计算T_{启动} + T_{计算}
  4. 比值含义R 大于 50% 时,说明一半以上时间没在算,应该考虑持久化或算子融合。
    ⋯⋯\frac{\cdots}{\cdots}

工程实践要点

真把驻场大厨请回家,运营细节开始冒头:按“前 100 桌归一号大厨”分单,一号忙到冒烟,二号却闲到擦灶台——这就是负载不均。GPU 上的惯用解法是“交错分配”:让每个线程块轮流取一块数据,第 t 轮取第 \(pid + t \times P\) 块,这样每轮所有人都忙,不会出现有的块早早收工、有的块还在苦熬的场面。

另一个要点是分清“持久化”与“融合”。算子融合(kernel fusion)把多个算子合并成一个 kernel,同样能减少启动次数——它解决的是“内核个数”问题;持久化解决的是“单个内核循环复用硬件”的问题。二者不但不冲突,还常常组合使用:先融合减少内核个数,再持久化让每个内核循环吃满数据。

最后,一切优化都要用数据说话。用 Nsight Compute 这类性能剖析工具测量启动耗时与 SM 占用率,改造前后分别记录墙钟时间做对比,只有当收益量化可见时,“请驻场大厨”这笔投入才算真正划算。

运营一家“驻场大厨”的后厨,讲究的是排班均衡、菜单设计与成本核算。把这套思路带回 GPU 世界,你手里的 Triton 就不再只是“每次重新点火”的普通灶台,而是一间随时待命、循环出菜的高效厨房。

offset=(pid+t×P)×B,t=0,1,…,M−1\displaystyle \text{offset} = (pid + t \times P) \times B, \quad t = 0, 1, \dots, M-1
每个常驻线程块按“块编号 + 轮次 × 块总数”交错取数据,保证所有块同步忙碌,负载均衡。
  1. pid 当前块编号线程块在网格中的唯一编号,相当于大厨的工号。
    pidpid
  2. t 循环轮次当前处理的第几段数据,取 0 到 M-1。
    tt
  3. P 块总数驻留线程块的总数,作为轮次之间的步长。
    PP
  4. B 块大小每段数据的元素个数,决定偏移量的粒度。
    BB
  5. 交错分配同一轮中相邻的块取相邻的数据段,避免某一块独扛大量数据。
    (pid+t×P)×B(pid + t \times P) \times 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 演示

python37 行
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()
目的:这个 Python 示例实现了一个 Triton 持久化 kernel——网格只含少量线程块,每个块用循环处理多段数据,一次启动完成整个向量加。
  1. grid = (n // (block * stages),)网格固定为少量线程块,让它小到足以全部驻留 GPU。
  2. for i in range(stages):kernel 内循环,一个常驻块连续处理多段数据,不再反复启动。
  3. num_blocks = tl.num_programs(0)获取网格块的数目,作为交错取块的步长。
  4. offset = (pid + i * num_blocks) * block按块编号交错取数据,保证负载均衡且互不重叠。
  5. mask = idx < n边界掩码,防止最后一段数据越界访问。
  6. BLOCK: tl.constexpr, STAGES: tl.constexpr编译期常量,让编译器做循环展开等优化。
  7. torch.cuda.synchronize()等待 GPU 全部执行完,再校验结果,避免异步误判。
  8. 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 的核心设计思想是什么?

单选题 · easy

  • A.让内核在 GPU 上反复重新启动
  • B.让线程块驻留并循环处理数据← 正确答案
  • C.通过扩大 L2 缓存区提升性能
  • D.减少线程数量降低资源占用率

答案:B。持久化 kernel 的关键是线程块常驻 SM,通过 for 循环依次处理多段数据,把 N 次启动压成 1 次。A 是传统多次启动;C 属于缓存优化;D 与持久化无直接关系。选 B。

Q2. 关于 kernel 启动开销,下列说法正确的是?

单选题 · easy

  • A.开销主要发生在数值计算当中
  • B.每次启动都有固定调度开销← 正确答案
  • C.启动开销与数据量大小成正比
  • D.启动开销小到几乎可以忽略

答案:B。启动开销是与数据量无关的固定成本,包含参数传递、线程调度、缓存准备等,通常为微秒级。A、C、D 均与事实不符,故选 B。

Q3. 下面哪些属于决定 SM 驻留线程块数量的硬件资源?

多选题 · easy

  • A.寄存器文件的总量上限← 正确答案
  • B.共享内存的容量大小限制← 正确答案
  • C.硬盘转速影响内核驻留
  • D.每个 SM 的线程驻留上限← 正确答案

答案:ABD。驻留块数由寄存器、共享内存和线程槽位共同决定。硬盘转速与 GPU 内核驻留无关,C 错误。故答案为 A、B、D。

Q4. Triton 中把每个线程块要循环处理的段数传给内核,通常用什么方式?

单选题 · mid

  • A.用 tl.constexpr 参数声明段数← 正确答案
  • B.在 tl.load 加载时指定段数
  • C.通过 program_id 返回值指定
  • D.在启动配置中指定网格段数

答案:A。循环段数是编译期常量,适合用 tl.constexpr 传入,便于编译器展开优化。B 的 load 只负责加载数据,C 返回块编号,D 指定网格大小,都不是传段数的正确方式。

Q5. 持久化 kernel 中,如何避免两个线程块处理同一段数据?

单选题 · mid

  • A.使用原子加锁保护每块数据
  • B.用块编号约束循环访问偏移← 正确答案
  • C.用共享内存标记数据属主
  • D.增大网格尺寸避免数据冲突

答案:B。最常用的是用 pid 构造不相交的偏移,如 offset=(pid+t×P)×B,使每段数据只属于一个块。A 是动态取任务的另类技巧,C 中块间共享内存不可见,D 不能消除冲突。

Q6. 持久化 kernel 能设置的网格大小上限取决于什么?

单选题 · mid

  • A.任务需要的数据块总量大小
  • B.寄存器与共享内存的资源量← 正确答案
  • C.输入数据的浮点精度高低
  • D.编译器优化选项版本差异

答案:B。网格上限由硬件的资源天花板决定——寄存器总量与共享内存容量限制了每个 SM 能驻留的块数,所有 SM 的驻留能力之和就是网格上限。A 影响循环段数而非网格上限,C、D 与驻留无关。

Q7. 用 Triton 编写持久化 kernel 通常需要做哪些事?

多选题 · mid

  • A.网格设为少量常驻线程块← 正确答案
  • B.内核内部添加循环处理段← 正确答案
  • C.用 constexpr 指定循环段数← 正确答案
  • D.必须依赖内联 CUDA 汇编

答案:ABC。三步走:网格设小、kernel 内 for 循环、constexpr 传段数。Triton 是高层 DSL,无需手写 CUDA 汇编,D 错误。

Q8. 持久化 kernel 中跨线程块协作可能死锁,根本原因是什么?

单选题 · extreme

  • A.数据量太大导致显存分配失败
  • B.块间缺少全局同步会导致互等← 正确答案
  • C.共享内存过小导致数据溢出错误
  • D.网格太大触发看门狗超时重置

答案:B。线程块之间没有内置全局同步原语,若用自旋锁等机制协作,两个块互相等待对方释放资源就可能永久互等,形成死锁。A、C、D 属于资源与稳定性问题,不是死锁的根因。

Q9. 某 kernel 启动开销 4μs、计算 1μs,需跑 100 次。改为持久化后总耗时约为?

单选题 · hard

  • 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 工程落地时需要注意的实践要点?

多选题 · hard

  • A.交错分配数据实现负载均衡← 正确答案
  • B.与算子融合搭配组合使用← 正确答案
  • C.用性能剖析工具实测收益← 正确答案
  • D.网格越大驻留效果才越好

答案:ABC。交错取块保均衡、与融合互补、用 Nsight 实测,都是正确实践。D 与驻留约束相悖,块数超过上限只会排队,错误。