实测Stable Diffusion + TensorRT推理速度提升3.8倍
在AI绘画应用逐渐从“能画”迈向“快画”的今天，响应速度成了用户体验的分水岭。你有没有遇到过这样的场景：输入一段提示词，点击生成，然后盯着进度条等上四五秒——对用户来说，这已经接近注意力的临界点。而如果是在一个高并发的AIGC服务平台上，这种延迟直接决定了每台GPU每分钟能服务多少用户。
正是在这种背景下，我们将目光投向了TensorRT——NVIDIA为深度学习推理量身打造的“性能加速器”。通过一次完整的工程实践，我们成功将Stable Diffusion在RTX 3090上的端到端图像生成时间从4.2秒压缩至1.1秒，提速达3.8倍。这不是理论值，而是真实负载下的实测结果。
这个数字背后，是模型结构、硬件特性与底层优化技术的深度协同。接下来，我们就拆开来看，为什么是UNet拖慢了整个流程？TensorRT是如何“动手术式”优化它的？以及最终这套方案能在哪些场景真正落地？
Stable Diffusion看似是一个模型，实际上是由多个子模块串联而成的复杂流水线：文本编码器（CLIP）、去噪网络（UNet）、图像解码器（VAE）。每一环都不可或缺，但它们的计算开销却极不均衡。
以标准512×512图像生成为例，在默认20步去噪迭代下：
- CLIP Text Encoder：约60ms
- VAE Decoder：约80ms
- UNet × 20次调用：超过3秒
这意味着，超过70%的时间花在了UNet的反复推理上。即便其他部分再怎么优化，只要UNet还是“原生态”，整体性能就很难有质的飞跃。
而UNet本身的结构也注定了它是性能瓶颈——它包含大量卷积层、注意力机制和跳跃连接，每一层之间都有显存读写和CUDA kernel启动的开销。在PyTorch这类动态框架中，这些操作几乎是逐层独立执行的，导致频繁的上下文切换和内存搬运。
这时候，TensorRT的价值就凸显出来了。它不像ONNX Runtime那样只是做个格式转换，而是像一位精通GPU架构的“编译器老手”，把整个计算图重新组织、融合、量化，最终生成一个高度定制化的推理引擎。
举个直观的例子：原始UNet中有成百上千个Conv + Bias + SiLU这样的小组合。在PyTorch里，这三个算子会分别调度三次kernel；而在TensorRT中，它们会被自动合并为一个复合算子，只需一次kernel launch即可完成，大大减少了GPU调度开销和中间缓存访问。
更进一步，TensorRT还会根据目标GPU的具体型号（比如你的卡是A100还是RTX 4090），选择最适合的CUDA内核实现方式。同样是FP16卷积，不同SM架构下的最优算法可能完全不同。TensorRT内置了庞大的内核库，并通过自动调优机制选出最佳组合，这种“因卡制宜”的策略，是通用框架难以企及的优势。
我们实际构建TensorRT引擎的核心代码如下：
import tensorrt as trt TRT_LOGGER = trt.Logger(trt.Logger.WARNING) builder = trt.Builder(TRT_LOGGER) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser = trt.OnnxParser(network, TRT_LOGGER) with open("stable_diffusion_unet.onnx", "rb") as model: if not parser.parse(model.read()): print("ERROR: Failed to parse ONNX file.") exit() config = builder.create_builder_config() config.max_workspace_size = 1 << 30 # 1GB临时空间 config.set_flag(trt.BuilderFlag.FP16) # 启用半精度 # 支持动态batch size profile = builder.create_optimization_profile() profile.set_shape("input", min=(1, 4, 64, 64), opt=(2, 4, 64, 64), max=(4, 4, 64, 64)) config.add_optimization_profile(profile) engine_bytes = builder.build_serialized_network(network, config) with open("sd_unet.engine", "wb") as f: f.write(engine_bytes)这段脚本虽然不长，但完成了最关键的“脱胎换骨”过程：
- 把ONNX模型加载进来；
- 开启FP16模式，让计算单元吞吐翻倍；
- 定义动态shape profile，允许运行时灵活调整batch size；
- 最终输出一个.engine文件——这是一个可以直接被部署的二进制推理镜像，不再依赖Python环境。
值得注意的是，这个过程是离线进行的，通常只需要执行一次。一旦引擎构建完成，后续每次启动都不需要重新编译，避免了线上服务冷启动的延迟问题。
构建完引擎后，真正的推理环节也需要精心设计。我们使用PyCUDA封装了一个高效的异步推理类：
import pycuda.driver as cuda import numpy as np class TrtInferenceRunner: def __init__(self, engine_path): with open(engine_path, "rb") as f: runtime = trt.Runtime(TRT_LOGGER) engine = runtime.deserialize_cuda_engine(f.read()) self.context = engine.create_execution_context() self.allocate_buffers(engine) def allocate_buffers(self, engine): self.inputs = [] self.outputs = [] for binding in engine: size = trt.volume(engine.get_binding_shape(binding)) * engine.num_optimization_profiles dtype = trt.nptype(engine.get_binding_dtype(binding)) host_mem = cuda.pagelocked_empty(size, dtype) device_mem = cuda.mem_alloc(host_mem.nbytes) if engine.binding_is_input(binding): self.inputs.append({'host': host_mem, 'device': device_mem}) else: self.outputs.append({'host': host_mem, 'device': device_mem}) def infer(self, input_data): np.copyto(self.inputs[0]['host'], input_data.ravel()) cuda.memcpy_htod_async(self.inputs[0]['device'], self.inputs[0]['host']) self.context.execute_async_v2( bindings=[i['device'] for i in self.inputs] + [o['device'] for o in self.outputs], stream_handle=pycuda.autoinit.context.get_current_stream().handle ) cuda.memcpy_dtoh_async(self.outputs[0]['host'], self.outputs[0]['device']) return self.outputs[0]['host'].reshape(self.context.get_binding_shape(1))这里的关键在于execute_async_v2和CUDA流的配合使用。它使得数据传输、计算和返回可以重叠进行，有效隐藏了部分I/O延迟。尤其是在多步去噪循环中，前一步的结果还没完全拷回CPU，下一步的输入就已经开始上传，形成了类似流水线的并行效果。
回到最初的问题：到底是什么带来了3.8倍的提速？
我们可以从三个维度来拆解：
- 算子层面：层融合减少了约60%的kernel调用次数；
- 精度层面：FP16使显存带宽需求减半，计算吞吐翻倍；
- 系统层面：异步执行+CUDA流提升了GPU利用率，避免空转。
三者叠加，才实现了从4.2秒到1.1秒的跨越。
当然，这一切也不是没有代价。例如，INT8量化虽然能进一步提速，但在UNet这种对细节敏感的模型上容易引入伪影，需谨慎使用。我们也做过测试：在Text Encoder上启用INT8几乎无损，但在UNet上会导致部分复杂提示词生成失真。因此，我们的建议是——关键路径保精度，非核心模块冲性能。
另一个常被忽视的点是显存管理。原生PyTorch在每次推理后并不会立即释放显存，导致batch size稍大就会OOM。而TensorRT引擎内部采用了内存池机制，复用已分配的显存块，结合set_max_workspace_size控制峰值占用，最终使显存消耗下降52%，支持更大的批量处理。
这直接带来了商业价值上的变化：原本一台服务器只能服务2路并发请求，现在可以轻松承载4路以上。对于云服务商而言，这意味着单位算力成本下降近一半。
至于跨平台兼容性，必须明确一点：TensorRT引擎不具备可移植性。你在A100上编译的.engine文件无法在T4或消费级显卡上运行，因为底层优化是针对特定SM架构生成的。所以理想的工作流应该是——在目标设备上统一构建引擎，或者通过CI/CD流水线为不同硬件预生成多个版本。
最后值得一提的是，这套优化思路并不局限于Stable Diffusion。只要是基于Transformer+CNN结构的扩散模型，都可以借鉴类似的拆解→导出→量化→重构流水线的方式进行加速。甚至一些视频生成模型（如Stable Video Diffusion），其UNet变体同样适用。
当AI生成进入“毫秒级响应”时代，拼的早已不是谁的模型更大，而是谁能把算力榨得更干。TensorRT在这场效率竞赛中扮演的角色，更像是一个“极限压榨者”——它不改变模型能力，却能让同样的模型跑出不一样的速度。
而3.8倍的提速，也不只是一个数字，它意味着：
- 用户等待时间从“喝口咖啡”变成“眨个眼”；
- 单张GPU每天可处理的请求数翻了几番；
- 原本只能云端运行的服务，现在也能跑在高端笔记本上。
未来随着TensorRT对稀疏注意力、动态分辨率等新特性的支持不断完善，这类优化的空间还会继续扩大。也许很快，我们会看到本地化AIGC工具成为标配，就像今天的Photoshop插件一样自然融入创作流程。
这条路的技术底座，正在一点点夯实。
