最近在处理一批实时推理服务时被一个问题磨得头疼TensorRT 引擎明明已经 build 得很干净Nsight Systems 一抓 profile主时间线却被一堆只有 2~3 微秒的小 kernel launch 填满。每个 kernel 单看都不起眼可一旦多起来CPU 侧调度和 GPU 侧启动的空档硬是把整帧延迟拉高了几十微秒。我试了两三版传统 CUDA plugin 之后干脆换了条更野的路子——DeepJIT运行时直接生成 CUDA C 源码字符串用 NVRTC 现场编译成 PTX再通过 Driver API 加载进 GPU用自己写的融合内核替换掉 TensorRT 里那串串行小核。这篇文章就是这套实践的完整记录。内容以 YOLO 推理后处理为例子适合正在做 GPU 推理服务优化、被 TensorRT 小算子拖累、或者想学 NVRTC 动态编译的人参考。我把踩过的坑、收益计算和代码模板都放在后面尽量做到照着能跑。1. 问题根源为什么 TensorRT 会有 “串行小核墙”1.1 一次推理里到底有多少个小 kernelGPU kernel 的完整生命周期大体是CPU 提交、驱动排队、GPU 调度器分配执行单元、kernel 启动执行、最后退出。每个环节都有成本尤其是 CPU 到 GPU 的提交路径一次 launch 往往要花 2 到 6 微秒。这个数字看着很小可如果一次推理里有 40 到 60 个 kernel光 launch overhead 就能积累到一两百微秒。在 T4 这类显卡上一个轻量后处理 kernel 本身的执行时间可能只有 10 微秒launch 开销反而把有效算力摊薄了三分之一以上。更麻烦的是中间显存读写。举个例子后处理一段常见序列输入 tensor 先做 sigmoid再乘上 scale再做 threshold 过滤。如果每个算子都是独立 kernelGPU 得把上一轮的中间结果从显存读出来、写进去重复三轮。对 memory-bound 的算子来说这等于白白浪费三倍带宽。我在实际工程里见过不少这种情况模型主体融合得很漂亮后处理却是一堆碎片化的 elementwise 小核在拖后腿。1.2 TensorRT 为什么不肯给你融合TensorRT 当然有层融合layer fusion但它能融合的是它预先定义好的 pattern比如 ConvBNReLU、部分 attention 结构、elementwise 与 activation 的组合。一旦你的模型里出现了“bias add → channel scale → 自定义 sigmoid 截断 → 阈值过滤”这种非标准组合融合器往往无法识别只好保守地保留 kernel 边界。TensorRT 的动态 shape 也会限制融合。当 batch、height、width 在运行时才能确定一些本来可以在 build 期做常量折叠的优化就做不了。此时编译器倾向于把算子拆开保证每个 kernel 都能在任意 shape 下正确调度。这种保守策略对通用引擎是好事但对我们这种“固定输入尺寸、追求极致延迟”的场景来说就成了明显的低效点。有人会问那写 TensorRT plugin 不行吗当然行但 plugin 的代价是厚重的样板代码。你要继承 pluginV2 接口处理序列化反序列化、类型绑定、命名空间、版本兼容光是把一个自定义算子塞进 engine 就得写几百行 C。我只是想优化几个 10 微秒的小算子这个成本实在不成比例。况且 plugin 一旦序列化进 engine后续调整参数还得重新 build。1.3 DeepJIT 的思路把数学过程当字符串把字符串当内核所以我换了个思路先按 TensorRT 的输入输出协议拿到中间 tensor再把我要合并的那几步数学运算写成一个代码生成器。生成器的输入是“融合描述 运行时 shape 常量”输出是一个完整的 CUDA kernel 源码字符串。然后调用 NVRTC 库现场编译成 PTX通过 Driver API 的cuModuleLoadDataEx加载成 module最后用cuLaunchKernel执行。我给这套流程起了个名字叫 DeepJIT。这里的 Deep 是“深度学习场景”的 DeepJIT 是 just-in-time——代码不是提前写好编译好而是在推理程序运行的过程中针对当前显卡架构和当前输入 shape 动态生成的。它本质上是把“写死一个 kernel”变成“写一个能生成 kernel 的工厂函数”。打个不那么严谨的比方TensorRT 小核序列就像每间屋子都要请一名工人工人进门先穿鞋套、再上楼、干完活再下楼十几个房间来回折腾DeepJIT 就是把相邻房间的活儿合并给同一个人。减少的是交接和往返的固定成本而不是干活本身的能力。2. 动手拆解从算子序列到融合内核2.1 先画 Kernel 边界哪些算子适合放进同一个内核不是所有算子都应该融。我在实践里遵循三条原则数据流必须连续。每个算子的输出只被下一个算子消费中间没有分叉、没有回边。算子的瓶颈最好是访存而非计算。elementwise、scale、activation、threshold 这类都是典型 memory-bound 算子融合可以砍掉重复读写。同步需求必须简单。不要碰需要跨 block 全局同步的算子比如一个 kernel 内做完整 Softmax。跨线程归约得自己设计 warp shuffle 和 shared memory代价远超收益。反过来说Conv 和 GEMM 这种计算密集算子不要想着自己手写。它们有 cuBLAS/TensorRT 专门的 kernel 库支持融合进去只会拖慢。我踩过的坑之一就是试图把一个小卷积融进后处理 kernel结果无论怎么调 tile 和 shared memory都不如 TensorRT 原生调用。计算密集的部分交给专业库memory-bound 的小算子才值得自己动手。2.2 生成式内核模板用 shape 和 fusion spec 替换编译器常量确定了边界之后就该设计代码生成器了。我采用的做法是定义一个 FusionSpec 结构体记录算子序列、输入输出数量、阈值、scale 模式等等然后根据这些字段拼接出对应的 CUDA 源码。shape 信息不会写成total 8400这种魔法数硬编码进源码而是以字符串拼接方式作为常量写入 kernel或者作为运行时参数传入。两种方式各有取舍。把 shape 编译进 kernel 的优点是编译器可以做常量折叠循环边界、索引计算都会更干净生成的 SASS 更紧凑缺点是 shape 一变就得重新编译一次需要做 cache。把 shape 作为运行时参数传入则更通用但某些优化就做不了了。我的折中方案是把“是否动态”做成配置稳定场景优先硬编码变化频繁的场景走参数传递。代码生成模板大致长这样std::string source R( extern C __global__ void fused_affine_sigmoid_filter( const float* __restrict__ input, float* __restrict__ output, const float* __restrict__ bias, const float* __restrict__ scale, float threshold, int total, int channels) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx total) { int c idx % channels; float x input[idx] * scale[c] bias[c]; float s 1.0f / (1.0f __expf(-x)); output[idx] (s threshold) ? 1.0f : 0.0f; } } );这个模板虽然简单但它背后有一个很关键的原则融合描述应该和代码生成分离。我在 FusionSpec 里定义好“需要哪几步数学运算”生成器只负责把它拼成源码。想加新算子就加一个代码片段不用改动整个 pipeline。2.3 动态 shape 与 launch 参数的计算JIT 编译最舒服的一点是可以按当前 shape 定制 launch 配置。例如输入 tensor 一共total batch * 8400 * 85个元素我会用cudaOccupancyMaxActiveBlocksPerMultiprocessor之类的接口先估算最优 block 数然后动态设置 grid 大小int block 256; int grid std::min((total block - 1) / block, max_blocks);max_blocks一般取设备 SM 数量乘 32 到 64具体值先用 occupancy API 算一次再缓存。因为 JIT kernel 就是为当前 shape 生成的我可以放心地把total当常量写进源码从而让编译器把取模、除法和循环边界优化成移位和掩码操作。这是传统预编译 cubin 很难做到的——它必须假设一个通用的 total。3. NVRTC 实战把 “CUDA 源码字符串” 送进 GPU3.1 环境准备需要哪些依赖DeepJIT 不是纯 driver 就能跑的。NVRTC 是 CUDA Toolkit 的一部分所以部署机上除了装好显卡驱动还得安装完整的 CUDA Toolkit保证/usr/local/cuda/include/nvrtc.h存在并且能链接到libnvrtc.so。很多容器镜像为了精简只带推理运行时这个时候nvidia-smi正常不代表 NVRTC 可用。我习惯在 Ubuntu 或者 WSL2 容器里装 CUDA Toolkit 时特别注意三点第一libnvidia-*驱动版本必须支持你目标架构的 PTX 加载第二Python 侧如果用了 PyTorch最好让torch.version.cuda和系统nvcc --version的主版本保持一致避免多种 CUDA runtime 混用导致版本错配第三如果机器上有多个 CUDA 版本必须显式设置CUDA_HOME和LD_LIBRARY_PATH否则 NVRTC 链接到旧库会出现各种诡异编译错误。头文件包含关系也要理清。NVRTC 是独立的库nvrtc.h里定义编译相关接口而cuModuleLoadDataEx、cuLaunchKernel这些 Driver API 来自cuda.h由libcuda.so提供。两者来自同一套 CUDA Toolkit但要注意 NVRTC 的输出是 PTX真正把 PTX 编译成 SASS 是驱动在 load module 时做的这算第二级 JIT。3.2 从源码到可调用内核的关键流程NVRTC 的典型调用链是#include nvrtc.h #include cuda.h #include vector #include string nvrtcProgram prog; nvrtcCreateProgram(prog, source.c_str(), fused_kernel.cu, 0, nullptr, nullptr); std::string archOpt --gpu-architecturecompute_80; const char* options[] {archOpt.c_str(), --use_fast_math, -default-device}; nvrtcResult res nvrtcCompileProgram(prog, 3, options); size_t logSize 0; nvrtcGetProgramLogSize(prog, logSize); std::string log(logSize, \0); nvrtcGetProgramLog(prog, log.data()); if (res ! NVRTC_SUCCESS) { throw std::runtime_error(NVRTC compile failed: log); } size_t ptxSize 0; nvrtcGetPTXSize(prog, ptxSize); std::vectorchar ptx(ptxSize); nvrtcGetPTX(prog, ptx.data()); CUmodule mod nullptr; CUfunction fn nullptr; cuModuleLoadDataEx(mod, ptx.data(), 0, nullptr, nullptr); cuModuleGetFunction(fn, mod, fused_affine_sigmoid_filter); nvrtcDestroyProgram(prog);这里提醒几个容易翻车的地方。nvrtcCreateProgram的第二个参数是程序名可以随便写但它是错误日志里用来定位源码的标识最好写成有意义的文件名。compileOptions中的-default-device参数能避免“隐式使用 host device 函数”的警告在编译简单 kernel 时推荐加上。--use_fast_math会引入快速数学近似比如__expf、__fdividef对后处理阈值敏感的场景要小心结果漂移。编译完拿到 PTX 后我建议立刻用cuModuleLoadDataEx加载并缓存CUfunction而不是每次推理都重复编译。编译一次简单 kernel 通常要几十到几百毫秒这个延迟在首次调用时可以被用户感知所以一定要做预热。3.3 一个最小可跑的融合内核示例下面的内核完成了“bias add → channel scale → sigmoid → threshold 过滤”四步融合输入是原始检测头输出输出是 0/1 掩码。它在真实推理后处理里就是四个小 kernel 合并成一个。extern C __global__ void fused_affine_sigmoid_filter( const float* __restrict__ input, float* __restrict__ output, const float* __restrict__ bias, const float* __restrict__ scale, float threshold, int total, int channels) { int idx blockIdx.x * blockDim.x threadIdx.x; if (idx total) { int c idx % channels; float x input[idx] * scale[c] bias[c]; float s 1.0f / (1.0f __expf(-x)); output[idx] (s threshold) ? 1.0f : 0.0f; } }Launch 的时候要注意参数布局。cuLaunchKernel接收一个void*参数数组数组里的每个元素是“指向实际参数内存的指针”不是参数值本身。比如float threshold 0.5f; int total 714000; int channels 85; void* args[] { input_ptr, output_ptr, bias_ptr, scale_ptr, threshold, total, channels }; cuLaunchKernel(fn, grid, 1, 1, block, 1, 1, 0, stream, args, nullptr);input_ptr这些已经是指针变量所以传给 args 的必须是input_ptr整个数组的类型就是void* []。我最初在这里写错过一次结果 kernel 拿到的参数全是垃圾值排查了半天。标量参数也一样不能直接把threshold塞进void*数组必须传它的地址。3.4 用向量化把 kernel 再榨一轮当total能被 4 整除、数据内存 16 字节对齐时可以改成float4向量化版本。每个线程处理 4 个连续元素索引计算减少访存指令合并成一次 16 字节事务。实际在 T4 上简单的 float 型 elementwise kernel 从 7 个 float 版本换成 float4 版本吞吐往往能提升 20% 到 40%。代价是尾部元素要单独处理不能让向量化越界。要注意的是如果total % 4 ! 0就不能无脑向量化。常见做法是主循环处理 4 的倍数部分尾部用标量循环补上。网格和块的设置也要相应变化vec_len total / 4grid 大小按vec_len计算。4. 在 TensorRT Engine 里塞进 DeepJIT 内核4.1 接入路径plugin 代理 vs engine 外包装把 DeepJIT kernel 接进推理管线主流有两条路我两种都试过。一是把融合 kernel 写成 TensorRT plugin。plugin 的enqueue函数直接调用cuLaunchKernel这样 TensorRT 在调度时能感知到这是一个自定义层engine 图结构也更完整。缺点是样板代码多还要处理序列化和反序列化。二是作为 engine 外的预处理/后处理模块。推理前把输入 tensor 整理成中间布局推理后把输出拿去做 JIT 融合 kernel再把结果送到下一个环节。优点是可以完全脱离 TensorRT 的 build 周期用 Python 就能快速搭建缺点是中间 tensor 需要多一次显存往返不过对后处理来说通常可以接受。我在真实项目里更推荐“外部包装优先”。先用外包装验证融合收益确认真能省下几十微秒再考虑 plugin 化。因为外部包装改动范围小出了问题可以立刻回滚而 plugin 一旦编进 engine查问题要重新 build 整个 engine成本高很多。4.2 Plugin 中的 JIT 编译与序列化策略如果你是非要用 plugin 不可的场景有一个核心问题要面对TensorRT 的 engine 序列化后plugin 的源码字符串也得跟着存下来。serialize阶段把 CUDA 源码文本、阈值、shape 参数写成一个 bufferdeserialize阶段拿到源码字符串以后重新走一遍 NVRTC 编译恢复到可调度的 CUfunction。反序列化时重新编译是逃不掉的因为 engine 文件里没法直接保存 SASS 二进制。但这里可以做一层磁盘 cache以“源码 SHA1 目标架构 编译选项”为 key把 NVRTC 输出的 PTX 或cuModuleLoadDataEx之后的结果缓存成本地文件。下次反序列化直接加载缓存跳过编译能把 plugin 初始化时间从几百毫秒降到几毫秒。plugin 的enqueue里还有一点要注意TensorRT 会把当前执行流作为cudaStream_t传进来JIT kernel 必须用这个 stream 启动否则会和前后算子产生隐式同步性能反而崩掉。我第一次接入时没传 streamkernel 默认走默认流导致整个 engine 被强制同步了多次耗时比原生方案还高 30%。5. 实测数据与收益复盘5.1 我们跑的一个典型用例我在一个 YOLOv8n 检测服务上做了对照实验。模型输入 640×640TensorRT FP16 enginebatch1后处理包括 sigmoid、坐标变换、阈值过滤和拼接。测试卡是 T4CUDA 12.2TensorRT 8.6冷启动预热后连续跑 500 帧取中位数。方案后处理平均耗时端到端平均耗时kernel 启动次数原生 TensorRT 小核序列103 μs1.421 ms7传统 CUDA Plugin预编译 kernel57 μs1.378 ms2DeepJIT PluginNVRTC 运行时生成41 μs1.361 ms1单纯看后处理段融合后从 103 微秒降到 41 微秒省了大约 60%。端到端影响大约 4.2%。这个数字在单路推理时不算惊艳但如果跑多路并发每路节省的 60 微秒 CPU 调度时间会叠加整体吞吐提升更明显。5.2 收益从哪里来又在哪一步消失收益主要来自三个部分。第一是 launch 次数从 7 次降到 1 次去掉了 6 次 CPU 提交和 GPU 调度空档第二是中间 tensor 不再反复写显存和读显存带宽利用率提高了第三是 JIT kernel 针对当前 shape 做了常量优化索引计算更简洁。但收益并不是自动到账的。在 A100 这类高带宽卡上launch overhead 对总耗时的影响相对小融合收益可能从 60% 缩水到 20% 以内反之在 T4、L4 这类中端卡上收益更明显。另外如果 TensorRT 已经把大部分算子融合得很好了你再强行融合只会增加复杂度不会有可观的提升。所以动手之前先用 profile 工具确认有没有“值得拆的墙”很重要。6. 这套方案里的坑我帮你提前踩平6.1 NVRTC 编译期的问题我在前期调试时攒了一个速查表按出现频率排下来问题现象原因与解法编译失败但日志为空nvrtcCompileProgram返回错误码log 却是空的先手动把生成的源码保存下来用nvcc -archcompute_xx -ptx试编译定位语法错找不到标准数学函数报错expf未定义NVRTC 的宿主设备分离模型导致部分头文件不可用改用内建函数__expf或加上-default-devicePTX 版本不可用驱动说 PTX 版本太新部署机的驱动版本低于 Toolkit 版本换旧一点的 NVIDIA driver 或者用较低 arch 编译动态 shape 每次重新编译太慢启动阶段明显卡顿加磁盘 cache 按源码 hash 复用编译结果不要让生产环境每次都从头编译链接找不到 nvrtc 库undefined reference to nvrtcCreateProgram链接参数补-lnvrtc -lcuda并确认LD_LIBRARY_PATH指向正确的 CUDA 库目录NVRTC 对 C 特性的支持是有限子集。模板、constexpr、__device__函数这些能用但别指望完整的标准库容器。我在生成代码时尽量保持 C11/14 风格复杂逻辑在生成阶段就展开成简单语句避免依赖运行时库。6.2 运行时性能和正确性问题JIT kernel 跑起来以后下面几个坑是肉眼可见的内存对齐和边界处理。float4向量化只适用按 4 分散的数据段。如果total不整除 4尾部必须用标量处理否则会读越界轻则结果错误重则非法内存访问。参数布局错误。cuLaunchKernel的void* args[]是参数地址的数组不是参数值。漏取地址、参数顺序不对kernel 拿到的就是乱码。建议把 launch 封装成一个强类型模板函数从根上避免这个问题。多 stream 并发。同一个CUfunction可以被多个 stream 同时 launch没问题但 JIT 编译和 launch 不能并发。编译期要加锁或者事先编译好再交给推理线程。浮点一致性。--use_fast_math会替换1.0f / (1.0f __expf(-x))之类的计算为近似实现对 sigmoid 结果影响很小但对 NMS 的阈值判断可能造成边界帧结果不一致。做检测模型时建议先用原逻辑对跑一遍测试集确认 AP 不掉点再上快速数学。还有一个容易忽略的点JIT kernel 默认没有__launch_bounds__约束编译器可能会用很多寄存器导致 occupancy 下降。实测下来对一些比较短的 kernel加上__launch_bounds__(256, 2)后 occupancy 能保持在高位吞吐额外提升 10% 左右。但这个值要根据寄存器用量和 block 大小调不能照搬。结尾这套 DeepJIT 实践跑下来我最大的体会是动手做融合之前先用 Nsight Compute 把每个小 kernel 的真实耗时、launch overhead、内存吞吐看清楚而不是靠猜。我一开始想融合一个看似很碎的计算段分析后才发现真正耗时的是显存读写不均匀而不是 lauch 次数换成按通道对齐的布局后没写任何新 kernel 就快了不少。另一个体会是 JIT 生成 kernel 要养成“代码生成器 缓存 预热”的习惯。生成器负责灵活产出源码缓存负责避免重复编译预热负责把初始化成本挪到请求高峰之前。三件事缺一不可。这一步想清楚之后DeepJIT 才能真正成为拆开 TensorRT 串行小核墙的趁手工具而不仅是一个炫技的实验。