最近压测一个视频分析项目，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 plugin
2.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.x，TensorRT 10.x，Ubuntu 22.04，都是当前比较常见的组合。如果你是 CUDA 11.x 的旧项目，NVRTC 接口也基本没变，迁移成本很低。
3. 定位问题：用工具把“墙”找出来
3.1 Nsight Systems 抓时间轴
在动手拆墙之前，得先把“墙”量化出来。我强烈建议每一步都基于 profile 数据，而不是感觉。这里用 Nsight Systems 最简单，一条命令跑完整个推理流程：
nsys profile --gpu-metrics-device=all -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.6ms，40ms 预算里能塞下 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[] = {"--std=c++17", "-use_fast_math", "--gpu-architecture=sm_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 只输出原始预测张量，decode+NMS 完全由我的 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 替换前后数据
折腾完一轮，直接上对比数据。同一台机器、同一模型、同样的 T4，batch size 1，YOLOv8s 640×640：
| 阶段 | 优化前(ms) | 优化后(ms) |
|---|---|---|
| TensorRT 主干推理 | 4.2 | 4.2 |
| 图内残存小核总计 | 1.2 | 0.3 |
| decode 后处理 | 0.6 | 0.1 |
| 整体延迟(单路) | 6.0 | 4.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（内部是 exp），NMS 又对 box 坐标敏感。开启-use_fast_math后，测试集 mAP 掉了约 0.2 个点，肉眼虽然看不出差别，但对线上模型来说这是不可接受的。最后我选择不使用全局 fast math，只在 sigmoid 那条路径上手动写了一个精度和速度平衡的近似版本。
另一个坑是--gpu-architecture=sm_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 手写内核替换。这种“发动机不动，只换零件”的策略，在工程上最稳、见效最快、风险最低。
