最近压测一个视频分析项目T4 上跑 TensorRT 加速的 YOLO 640 检测单路延迟看起来能接受可一旦往多路扩展帧率就上不去。抓了几轮 profile发现问题根本不在主干推理而在 TensorRT 引擎里那二十多个串行小核上——每次 launch 几微秒积少成多把整个流水线拖成了“排队过闸”。这个局我见过不少次这次干脆试了试 DeepJIT 路线手写 CUDA 内核把 TensorRT 生成的串行小核墙从中间拆掉。文章写给正在做模型部署和推理优化的人尤其是已经在用 TensorRT、又对性能不满、却不想折腾整套自定义插件的朋友。核心思路很简单TensorRT 当主力JIT 编译的自定义 CUDA 内核当“拆墙工”把瓶颈算子从引擎里掏出来自己写然后动态编译进推理管线。这篇文章把我的完整过程、踩坑记录和实测数据都摊开讲能帮你少走不少弯路。1. 先搞清楚“串行小核墙”到底挡在哪1.1 TensorRT 明明很会融合怎么还会串行TensorRT 的看家本事是图融合——把 Conv BN ReLU 这种常见组合并成一个 kernel减少从 CPU 侧发起的 kernel launch 次数。这个方向没问题GPU 最怕的就是“一个 kernel 干一丁点活然后灰溜溜退出”一次 launch 的空洞开销就有 3~10 微秒。所以融合得越狠延迟越低。但融合不是万能的。它要满足很多前提算子之间的数据布局能对齐、计算类型能匹配、图的形状是静态的或至少是可推导的还得有对应的融合 kernel 实现。碰到 decode、gather、sort、NMS 前后处理这种形状跳动大、逻辑又偏“非主流”的算子时TensorRT 往往不融合而是老老实实生成一堆很小的 kernel逐个执行。这个行为本质上也不是 bug它是保底策略——生成器宁可保守也不能产生错误结果。问题是对于追求极致吞吐的部署场景保守就等于串行墙。我这次项目里的模型主干推理只用了 4ms 左右但后面跟着十多个 100us 到 1ms 不等的残存小核总量 1.8ms占比超过 30%。在 T4 这种卡上这不是小数。1.2 串行小核墙的两种形态我把实际遇到的“墙”分成两类方便对症下药。第一类是图内残存的小 kernel 串行 launch。TensorRT 引擎里还留着 slice、gather、transpose、cast 这类算子每个都是单独的 kernel顺着执行。你说它们能不能合并在特定情况下能但 TensorRT 插件机制非常重改一个算子要重新 build engine、重新做序列化开发节奏太慢。第二类是单个 kernel 内部的伪并行。TensorRT 生成的某些 kernel 看起来在线程里跑但算法本质是串行的——例如遍历一整张 feature map 的所有 anchor每个线程只负责一小块连续数据却要同步等一个全局循环结束。结果就是占了一堆 SM却没有真正把并行度跑满。后者比前者更坑因为它在 profile 里只显示一个 kernel不拆开看内部逻辑根本发现不了。打个比方串行小核墙就像快递分拣流水线明明有十条通道但每个包裹必须在同一个闸口过秤一条通道堵住了后面九条全在等。TensorRT 负责把包裹仓库盖得很漂亮但仓库入口那个老旧闸口它不管。2. 为什么选 DeepJIT 这类方案而不是老实写 TensorRT plugin2.1 TensorRT plugin 的老问题我最早也想走插件路线毕竟 TensorRT 官方文档里写着“自定义算子请用 plugin”。但实际动手才发现这套机制对性能调优来说是个负担。插件要继承 IPluginV2DynamicExt 这类接口实现一大堆方法getOutputDimensions、enqueue、serialize、deserialize、getSerializationSize、supportsFormatCombination……光是把这些钩子填对就能耗掉一晚上。而且 build engine 和 inference 是两套生命周期你在 enqueue 里写的 kernel 有什么问题必须等到 engine 跑起来才能看到中途想改 kernel 逻辑就得重新编译整个插件、重新 build engine循环非常长。更烦的是 TensorRT 版本升级后插件 API 频繁变动8.x 时代一套写法9.x 又换一套。每次升级都在补接口。对一个“尝鲜”性质、想快速验证手写 kernel 是否有效的探索阶段来说这种重流程会直接劝退。2.2 DeepJIT 方式带来的自由DeepJIT 的核心就一句话用运行时编译把自定义 CUDA 内核动态注入推理流程让 kernel 的修改不需要跨过“重新 build engine”这道鸿沟。做法上不需要引入什么神秘框架CUDA 生态里现成的材料足够用 NVRTC 库把 CUDA C 源码字符串在生产环境动态编译成 PTX用 cuModuleLoad / cuModuleGetFunction把 PTX 加载进当前进程直接拿到 kernel 函数句柄如果是 Python 侧做原型验证pynvrtc 或者 torch.utils.cpp_extension 也能做类似的事生产环境可以退到 cuLaunchKernel 或 CUDA Graph把自定义内核和 TensorRT 的执行流无缝拼在一起。我实际是把主推理图留在 TensorRT 手里把 decode 和残存小核的活全部掏出来自己写。这样 TensorRT 负责它擅长的密集卷积计算我负责那些需要精细控制性能的边界算子。两边通过同一个 CUDA stream 串起来数据不用拷回 CPU全程留在显存里。顺带一提这套方案对环境的要求不算苛刻。我用的 CUDA 12.8 cuDNN 9.xTensorRT 10.xUbuntu 22.04都是当前比较常见的组合。如果你是 CUDA 11.x 的旧项目NVRTC 接口也基本没变迁移成本很低。3. 定位问题用工具把“墙”找出来3.1 Nsight Systems 抓时间轴在动手拆墙之前得先把“墙”量化出来。我强烈建议每一步都基于 profile 数据而不是感觉。这里用 Nsight Systems 最简单一条命令跑完整个推理流程nsys profile --gpu-metrics-deviceall -o yolov8_t4 python run_inference.py跑完后用 nsys stats 导出 kernel 时间分布或者直接打开 Nsight Systems GUI 看 GPU 时间轴。在时间轴上会看到一串很短的 kernel——它们并排连在一起每一个只有几百微秒甚至几十微秒但 gap 和 launch 延迟像心跳一样规律出现。这部分就是串行小核墙的物理形态。我当时还额外做了个对比实验把注意力集中在引擎中间那段看两件大事。第一从第一个 kernel 到最后一个 kernel 的墙钟时间是多少第二主干卷积大 kernel 之间的间隙总共占了多少。结论是两个数据加起来 1.8ms而我手写 decode kernel 的目标就是把这个数字压到 0.4ms 以下。如果你手头没有 Nsight也可以用 CUDA Events 在代码里手动打点cudaEventRecord(start); engine-enqueueV2(buffers, stream, nullptr); cudaEventRecord(stop); cudaEventSynchronize(stop); float ms 0; cudaEventElapsedTime(ms, start, stop);这种方式精度有限但胜在简单能快速确认“瓶颈到底在引擎内还是引擎外”。3.2 换算业务指标多少路视频被拖累定位问题只是第一步得把技术损耗翻译成业务损失才能说服自己(和对面的需求方)值得花这个精力。这次项目场景和网上大家讨论得很多的“T4 单卡能跑多少路 1080p25 的 YOLO 640 检测”是同一类问题。单路算力预算按 40ms 算(25 帧)如果引擎延迟 4.2ms 后处理 1.8ms总共 6ms单卡理想上限是 6~7 路。这里面后处理占的 1.8ms本质就是串行小核墙造成的浪费。优化后变 4.2ms 0.4ms 4.6ms40ms 预算里能塞下 8 路路数上限提升了 25% 以上。对一台跑几十路视频的服务器来说就等于少买两张卡。这还没算延迟降低对首帧响应、追焦实时性的改善。所以别小看那几毫秒在推理服务里延迟每一毫秒都是钱。4. 手写 CUDA 内核从 decode 算子开刀4.1 我挑了这个算子的原因YOLO 系列模型的 decode 阶段(把网格预测值变成实际 box 坐标)是个典型瓶颈。TensorRT 对这类逻辑的处理一直偏保守——不是不能做而是因为后面通常还跟 NMS、confidence 过滤这些形状动态的算子引擎很难对整段做激进的融合优化。结果就是生成一堆小 kernel每个干一点活数据在显存里倒来倒去。很多人图省事直接把 decode 放回 CPU 做。640×640 输入、三个尺度、总共约 8400 个 anchor看起来量不大但在 CPU 上每个 box 要算 sigmoid、坐标缩放、类别打分几百路一起跑时 CPU 占用立刻爆炸。这是我这次坚决在 GPU 上解决的原因。还有一个动机是网上常见的 CUDA decode 代码大多是“每个线程处理一个 anchor”的朴素写法理论并行度有 8400但实际因为内存访问跨度大、分支多跑下来并不快。我知道一个更好的写法能显著压缩耗时正好借这个机会验证一下。4.2 一个能打的 CUDA kernel 长什么样手写内核前先立几条规矩全局遍历用 grid-stride loop保证任意 grid 配置都能正确跑完每个线程一次处理 4 个 anchor而不是 1 个这样能摊薄索引计算并且利用指令级并行(ILP)尽量用 vectorized load读 4 个 float 当一次 float4指针全部加__restrict__让编译器知道你不会有别名冲突。核心 decode 部分大致长这样__global__ void decode_kernel( const float* __restrict__ pred, // (num_anchors, 4 num_classes 1) 连续排布 float* __restrict__ boxes, // 输出: x1, y1, x2, y2, score, class_id float conf_threshold, int num_anchors, int num_classes) { int idx (blockIdx.x * blockDim.x threadIdx.x) * 4; if (idx num_anchors) return; float4 vals[4]; #pragma unroll for (int i 0; i 4 idx i num_anchors; i) { int a idx i; int base a * (4 num_classes 1); float cx pred[base 0]; float cy pred[base 1]; float w pred[base 2]; float h pred[base 3]; float obj pred[base 4]; int best_cls 0; float best_score 0.0f; #pragma unroll 8 for (int c 0; c num_classes; c) { float s pred[base 5 c] * obj; if (s best_score) { best_score s; best_cls c; } } if (best_score conf_threshold) { // 转换为 x1,y1,x2,y2与输入图像尺寸相关 boxes[out_base] cx - w * 0.5f; boxes[out_base 1] cy - h * 0.5f; boxes[out_base 2] cx w * 0.5f; boxes[out_base 3] cy h * 0.5f; boxes[out_base 4] best_score; boxes[out_base 5] (float)best_cls; } } }这段代码看起来简单但几个细节决定了它和普通版本的区别。一是把“类别打分×obj 置信度”融合成一次循环不在 kernel 内部分两个阶段减少重复读内存二是循环内没有复杂的分支只有连续比较非常利于 GPU 的分支预测和编译器向量化三是每个线程负责 4 个 anchor流水线内部可以同时进行多组独立的乘加操作不会因为等待乘法器空转。实际编译时我还加了一行#pragma unroll 8让编译器展开类别循环——类别数通常 80展开 8 次是个较稳妥的平衡既能减少循环开销又不至于让代码体积膨胀到指令缓存装不下。4.3 阈值与 box 后处理的取舍在 decode 的 kernel 里顺便做 confidence 过滤是个很诱人的优化但这里有个精密的取舍。好处显而易见无效 box 不写入显存后续 NMS 的数据量大幅减少甚至 NMS 本身可以省掉一大部分计算。坏处是如果你把conf_threshold设得太高可能过早丢掉低置信度但经过 NMS 后最终被保留的框。比如两个重叠框一个 0.4 分一个 0.9 分NMS 会把 0.4 的抑制掉但如果 decode 阶段就过滤掉 0.4结果是一样的。可如果两个框分属不同类别NMS 在不同类别间通常不做抑制这时提前过滤就会误伤。实操上我的建议是decode 阶段只做“安全过滤”即过滤掉置信度极低的框(比如低于 0.1)真正的精确阈值留给 NMS 阶段。这样既减少了无效数据量又不会影响最终结果。这个 0.1 是我实测下来不改变 mAP 的保守值不同模型可能要微调。另外一个注意点如果要做“只输出前 k 个框”这种带跨 block 压缩的操作最简单用atomicAdd维护一个全局 count或者分两步先统计有效框数量再压缩写入。直接在 decode kernel 里混合做原子操作和过滤容易把性能吃掉我一般避免在 8400 个 anchor 这么小的规模上做复杂跨 block 协作。5. 把自定义内核接进 DeepJIT 流程5.1 JIT 编译与模块加载手写 kernel 只是第一步真正让开发体验“DeepJIT”起来的关键是运行时编译。我拿 NVRTC 把上面的 CUDA 源码直接编译成 PTX然后通过 Driver API 加载整个流程在 C 里就是这个样子#include nvrtc.h #include cuda.h std::string source load_kernel_source(decode.cu); nvrtcProgram prog; nvrtcCreateProgram(prog, source.c_str(), decode.cu, 0, nullptr, nullptr); const char* opts[] {--stdc17, -use_fast_math, --gpu-architecturesm_75}; nvrtcCompileProgram(prog, 3, opts); size_t ptx_size 0; nvrtcGetPTX(prog, ptx); nvrtcDestroyProgram(prog); CUmodule module; cuModuleLoadData(module, ptx); CUfunction kernel; cuModuleGetFunction(kernel, module, decode_kernel);之后每次修改 kernel 源码只需要重新跑一遍编译-加载不用退出进程、不用重建 TensorRT engine。这就是我标题里说“尝鲜”的来源——DeepJIT 不是某个固定产品而是一套让你能快速迭代自定义内核工作流的总称。生产环境里首次编译 PTX 通常要几百毫秒甚至更久这在每次启动时不可接受。我的处理是加一层编译缓存第一次编译成功后把 PTX 二进制写到磁盘下次启动直接cuModuleLoadData加载跳过 NVRTC。这样既保留了修改源码后“热更新”的灵活性又不会拖慢启动速度。5.2 和 TensorRT 引擎如何搭讲完了 JIT 侧再说和 TensorRT 的协作方式。我试过两种搭法。第一种是“把 decode 挪到墙外”TensorRT engine 只输出原始预测张量decodeNMS 完全由我的 CUDA kernel 接管。这是最干净的切分引擎结构最简单自定义内核拥有完全的数据布局自由。缺点是模型本身不再是一个端到端引擎外部调用方需要知道如何处理原始输出封装层要多做一层。第二种是“显式嵌入同一条 stream”TensorRT enqueue 后紧接着在同一个 CUDA stream 上 launch decode kernel。这样保证引擎和自定义内核的 GPU 工作天然排队不会并发乱序同时避免了多次 CPU→GPU 同步。实测这种方式性能最优延迟和第一种几乎一样但工程上更容易集成进现有推理管线。我最终选择了第二种。最后还可以用 CUDA Graph 把它们固化cudaGraphBegin(); // 先记录 TensorRT enqueue再记录自定义 kernel launch cudaGraphInstantiate(graphExec, graph, nullptr, nullptr, 0); // 以后每次推理只需要 graphLaunch 一次把多个 launch 录进一张图后CPU 侧的 launch 开销会大幅下降。实测中原本 20 多次小 kernel launch 的 CPU 开销几乎归零GPU 上的串行等待也明显减少。5.3 多 stream 的坑在继续之前提醒一个容易踩的坑不要为了表面上的“并行”而开一堆 CUDA stream 跑这些自定义内核。T4 这种卡的计算单元不少但显存带宽有限。解码、NMS、缩放这些算子大多吃带宽而不是吃算力你开两个 stream 同时跑它们会抢带宽互相拖慢最终总吞吐反而下降。正确做法是在单张卡上用一个主 stream 把所有工作串好用 CUDA Graph 固化然后通过多进程或多卡去扩展路数。我最早测试时图省事给每路视频开一个 stream结果显卡占用率不高、延迟却没降下来。后来把所有路的工作塞进同一个 stream Graph显存带宽利用率上去了整体吞吐立刻改善。这就是典型的“看着并行、实际串行抢资源”陷阱。6. 实测效果与踩坑记录6.1 替换前后数据折腾完一轮直接上对比数据。同一台机器、同一模型、同样的 T4batch size 1YOLOv8s 640×640阶段优化前(ms)优化后(ms)TensorRT 主干推理4.24.2图内残存小核总计1.20.3decode 后处理0.60.1整体延迟(单路)6.04.6延迟降了 23%但更关键的是把可支持路数从 6 路拉到了 8 路(按 1080p25 计算)。对一个 32 路的边缘盒子来说相当于少买四分之一的卡。单独说 decode kernel 的性能从原来的 0.6ms 降到了 0.1ms而且这还是在同时做了 confidence 过滤的情况下。这个提升主要来自三点vectorized load、每线程多 anchor、以及去掉了 kernel 内的分支发散。6.2 踩坑一bank conflict 与 shared memory 布局第一次写 kernel 时我天真地把所有中间结果先存进 shared memory然后再读出来计算。结果性能不仅没提升反而比直接读 global 还慢。用 Nsight Compute 一看shared memory 的 bank conflict 严重得吓人——访问步幅恰好是 2 个 bank导致吞吐直接砍半。原因是我按 anchor 顺序存中间值而后续按类别顺序读取时相邻线程访问的地址跨越了固定步长落到同一个 bank 上。解决办法有两种一是给 shared memory 数组加 padding把步幅错开二是干脆不用 shared memory改成让每个线程重复读少量 global 数据利用 L1/L2 缓存兜底。我在这个 kernel 里选了第二种因为数据量不大、访问模式规律L1 命中率很好比折腾 shared memory 更省事。这里也想吐槽一句bank conflict 是优化 CUDA kernel 最容易踩的坑而且不跑 profiler 基本看不出来。优化这类代码时Nsight Compute 是标配不要靠猜。6.3 踩坑二NVRTC 优化选项并非白给NVRTC 编译选项中有一堆看似美好的开关最典型的是-use_fast_math。它会把expf、sinf这类函数替换成精度略低的硬件近似指令速度确实快不少但对部分算子来说精度损失会传导到最终结果。我的 decode 里用了 sigmoid内部是 expNMS 又对 box 坐标敏感。开启-use_fast_math后测试集 mAP 掉了约 0.2 个点肉眼虽然看不出差别但对线上模型来说这是不可接受的。最后我选择不使用全局 fast math只在 sigmoid 那条路径上手动写了一个精度和速度平衡的近似版本。另一个坑是--gpu-architecturesm_75这种硬编码。如果你把 PTX 在别的卡上加载要么无法运行要么需要 JIT 再编译。生产环境最好针对目标机器写死但开发环境建议用compute_75这类兼容符号避免换个机器就崩。6.4 踩坑三别用 vectorized load 踩了未对齐地址vectorized load(float4)要求地址 16 字节对齐。TensorRT 输出的张量通常是按最大对齐分配的问题不大。但我自己在自定义缓冲区里分配输出时曾经图省事直接malloc导致后续float4访存全部错位出现随机崩溃和极其诡异的结果。排查了很久才发现是cudaMalloc返回的内存天然满足大地址对齐而malloc只保证最基础的对齐。之后统一改回cudaMalloc或者用posix_memalign问题立刻消失。凡是涉及 vectorized 访存内存分配一定要显式对齐别赌编译器帮你处理。还有个细节是__builtin_assume_aligned可以用它告诉编译器指针已经是 16 字节对齐让编译器放心生成ld.global.v4.f32指令。这算是个微优化但配合__restrict__经常能带来 5%~10% 的额外加速。7. 最后再分享一个小技巧这种优化做完后别急着结束。把 JIT 编译和 CUDA Graph 封装成一个“补丁层”是很有价值的工作——后续再遇到 TensorRT 生成的其它低效小核你只需要写一个 CUDA kernel、插进同一套加载流程就能解决。我后来把同样的方式用在了模型的 NMS 阶段以及一个很冷门的 Resize 算子身上效果都不错。我个人体会是TensorRT 是出色的主力引擎但你永远要保留对它“生成策略”的审视能力。遇到串行小核墙不要急着推翻整套技术栈先用 profiler 量化再挑一两个真正吃时间的算子用 JIT 手写内核替换。这种“发动机不动只换零件”的策略在工程上最稳、见效最快、风险最低。