TensorRT-LLM 中对 wan 加速流程与方法

前言

本文详细拆解NVIDIA TensorRT-LLM 官方对 Wan 系列视频生成模型的完整加速实现，而非 Wan-AI 原生代码。TensorRT-LLM 凭借其深度优化的 CUDA kernel、灵活的并行策略和高效的缓存机制，提升了 Wan 2.1/2.2 系列模型的推理速度。

补充：目前TensorRT-LLM 官方未有直接使用的包，需要自己进行源码编译

很多开发者在使用visual_gen_wan_i2v.py脚本时，往往只知道通过命令行参数开启各种加速功能，却不清楚这些技术在 TensorRT-LLM 代码中究竟是如何工作的，以及它们在整个推理流程中处于什么位置。本文将从代码执行的完整流程出发，为你详细拆解所有 14 项核心加速技术的精确流程位置 、关键代码入口 和底层实现原理。

TensorRT-LLM/examples/visual_gen/README.md at main · NVIDIA/TensorRT-LLM

一、整体架构概览

TensorRT-LLM 的 Wan 推理系统采用了主从分布式架构，同一个 Python 脚本既支持单机单卡运行，也支持多机多卡集群部署。整个系统可以分为 7 个核心层级，从用户入口一直到最终的视频输出：

复制代码

用户脚本(visual_gen_wan_i2v.py)
    ↓
VisualGen(配置入口 + 协调器)
    ↓
DiffusionRemoteClient(分布式客户端)
    ↓
Worker进程 × N(总GPU数)
    ↓
PipelineLoader(模型加载 + TRT编译 + 预热)
    ↓
WanImageToVideoPipeline(核心推理逻辑)
    ↓
DiffusionResponse(视频输出)

核心设计原则 ：所有加速技术的开关在VisualGenArgs阶段就已经完全确定，后续流程不再动态变更。这种设计保证了系统的稳定性和可预测性，也为 TensorRT-LLM 的静态图编译和优化提供了基础。

二、核心加载流程：PipelineLoader

PipelineLoader 是整个加速系统的 "总控中心"，所有 TensorRT-LLM 特有的加速技术的初始化和配置都在这里完成。它的执行流程分为 6 个关键步骤：

STEP 1-4：基础准备

解析 Checkpoint：支持本地路径和 HuggingFace Hub ID
加载模型配置 ：读取model_index.json和config.json，合并用户参数
构建设备网格 ：创建 4D 设备网格[cfg, tp, cp, ulysses]，初始化所有 NCCL 通信组
Meta 初始化：先创建 meta tensor (零显存)，再分配 CUDA 显存，避免瞬间显存峰值

STEP 5：加载权重与 TRT 量化

复制代码

# 关键代码位置：PipelineLoader.load() :205-220
load_transformer_weights()
load_weights()
if dynamic_weight_quant:
    # TensorRT-LLM DynamicLinearWeightLoader
    # CPU端完成BF16 → FP8/NVFP4量化
    # 直接生成TRT兼容的权重格式
    DynamicLinearWeightLoader.quantize_and_load()
load_standard_components()  # VAE、TextEncoder不量化

技术映射：

FP8/NVFP4 量化：在此处执行，TensorRT-LLM 特有的动态量化实现，CPU 端一次性量化，节省 GPU 显存带宽

STEP 6：后处理、TRT 编译与预热

复制代码

# 关键代码位置：PipelineLoader.load() :228-248
setup_parallel_vae()          # 替换为TRT优化的带halo exchange的并行VAE
post_load_weights()           # 注册TeaCache和Cache-DiT钩子
torch_compile()               # JIT编译非TRT部分的代码
warmup()                      # dummy推理触发TRT编译和kernel预热

技术映射：

并行 VAE：在此处替换为 TensorRT-LLM 优化的并行 VAE 层
TeaCache/Cache-DiT：在此处注册 TRT 兼容的钩子函数
torch.compile：在此处编译非 TRT 部分的 Python 代码
Warmup 预热：在此处执行，触发 TensorRT-LLM 的最终编译和所有缓存

三、推理主循环与加速技术分布

模型加载完成后，进入serve_forever()主循环，等待并处理用户请求。核心推理逻辑在WanImageToVideoPipeline.forward()中，这也是大多数 TensorRT-LLM 加速技术实际生效的地方。

3.1 完整推理流程

复制代码

def forward(self, prompt, image, ...):
    # (A) 文本编码
    text_embeds = _encode_prompt(prompt)
    
    # (B) 图像编码(Wan2.1用CLIP，Wan2.2不用)
    image_embeds = _encode_image(image)
    
    # (C) 准备Latent
    latents, condition = _prepare_latents(image, noise)
    
    # (D) 去噪循环(核心)
    for t in timesteps:
        # 1. CFG并行：正负prompt合并
        # 2. TeaCache：判断是否跳过当前步
        # 3. CUDA Graph：replay已录制的TRT kernel序列
        # 4. forward_fn：TRT transformer推理
        #    → 两阶段去噪：根据时间步选择TRT引擎
        #    → TRT注意力后端调度
        #    → Ulysses/Attention2D并行
        # 5. Cache-DiT：逐block缓存判断
        # 6. 调度器更新latents
        # 7. Wan2.2-5B：后步图像固定
        latents = scheduler.step(noise_pred, t, latents)
    
    # (E) TRT并行VAE解码
    video = vae.decode(latents)
    
    return video

3.2 加速技术完整映射表

这是本文最核心的部分，我将所有 TensorRT-LLM 特有的加速技术精确映射到对应的流程位置和关键代码：

加速技术	流程位置	关键代码入口
FP8/NVFP4 量化	PipelineLoader STEP 5 → load_weights	`DynamicLinearWeightLoader`
TRT 注意力后端	WanBlock → Attention → backend dispatch	`TrtllmAttention`(默认)
Ulysses 序列并行	transformer.forward → attention wrapper	`UlyssesAttention(all-to-all)`
Attention2D 并行	transformer.forward → attention wrapper	`Attention2DAttention(2D mesh)`
CFG 并行	denoise () 循环内部	`_denoise_step_cfg_parallel(all-gather)`
TRT 并行 VAE	PipelineLoader STEP 6 + forward VAE 解码	`ParallelVAE_Wan(halo exchange)`
TeaCache	denoise () 循环头部 + post_load_weights 注册	`TeaCacheHook(距离阈值判断)`
Cache-DiT	denoise () 循环内 + post_load_weights 注册	`DBCache(逐block残差判断)`
两阶段去噪	forward_fn 模型选择逻辑	`current_t >= boundary → transformer_2`
CUDA Graph	denoise () 循环内	`CUDAGraphRunner.replay()`
torch.compile	PipelineLoader STEP 6	`torch.compile(block)`
Warmup 预热	PipelineLoader STEP 6	dummy 推理触发 TRT 编译 + 预热
逐 patch 2D 时间步	pipeline_wan.py forward	Wan2.2 TI2V-5B 专属
后步图像固定	pipeline_wan.py _pin_i2v_first_frame	Wan2.2 TI2V-5B 专属
融合 QK Norm+RoPE	代码中存在，默认关闭	`fuse_qk_norm_rope=False`

重要说明：

逐 patch 2D 时间步和后步图像固定是Wan 2.2 TI2V-5B 专属技术，不在 I2V A14B 的流程中
融合 QK Norm+RoPE 虽然代码中存在，但 Wan 默认关闭，因为 flashinfer 在 full-dim norm 下更快
TensorRT-LLM 中TrtllmAttention是默认后端，Vanilla 和 FlashAttn4 仅作为兜底或特殊场景使用

四、加速技术兼容矩阵

不同版本的 Wan 模型在 TensorRT-LLM 中支持的加速技术有所不同，下表是完整的兼容情况：

加速技术	Wan 2.1 T2V	Wan 2.1 I2V	Wan 2.2 T2V (A14B)	Wan 2.2 I2V (A14B)	Wan 2.2 TI2V-5B
TeaCache	✅	✅	❌	❌	❌
Cache-DiT	✅ 单 transformer	✅ 单 transformer	✅ 双 transformer	✅ 双 transformer	✅ 单 transformer
逐 patch 2D 时间步	❌	❌	❌	❌	✅
后步图像固定	❌	❌	❌	❌	✅ (I2V)
CFG 并行	✅	✅	✅	✅	✅
Ulysses/Attention2D	✅	✅	✅	✅	✅
TRT 并行 VAE	✅	✅	✅	✅	✅
CUDA Graph	✅	✅	✅ (SharedGraphPool)	✅ (SharedGraphPool)	✅
torch.compile	✅	✅	✅	✅	✅
FP8/NVFP4 量化	✅	✅	✅	✅	✅
融合 QK Norm+RoPE	可选 (默认关)	可选 (默认关)	可选 (默认关)	可选 (默认关)	可选 (默认关)

五、TensorRT-LLM 关键加速技术深度解析

5.1 量化压缩：FP8 vs NVFP4

量化是解决大模型显存瓶颈最直接有效的方法。TensorRT-LLM 对 Wan 模型的量化支持最为完善，全部采用动态量化方式，不需要预量化模型：

模式	位宽	分块方式	显存节省	适用场景
FP8 Per-Tensor	8-bit E4M3	整个矩阵一个 scale	~50%	对精度要求高
FP8 Blockwise	8-bit E4M3	128×128 分块	~50%	精度与速度平衡
NVFP4	4-bit E2M1	16 元素分块	~75%	显存严重不足

TensorRT-LLM 特有的实现细节：

FP8 使用torch.ops.tensorrt_llm.quantize_e4m3_per_tensor
NVFP4 使用torch.ops.trtllm.fp4_quantize，并启用AutoTuner自动搜索最优 kernel
所有量化在CPU 端完成，直接生成 TRT 引擎可直接加载的权重格式
支持混合精度量化，可以对不同层使用不同的量化精度

5.2 TRT 注意力后端：核心优化引擎

在 TensorRT-LLM 中，TrtllmAttention是默认的注意力计算引擎，也是所有加速技术的核心：

TrtllmAttention：TensorRT-LLM 深度优化的注意力 kernel，支持：
- SageAttention 逐块 INT8 Q/K 量化
- 融合 QKV 计算
- 自动处理自注意力和交叉注意力
- 与 Ulysses/Attention2D 并行无缝集成
FlashAttn4Attention：仅作为 Attention2D 的前置依赖，因为需要输出 LSE (log-sum-exp)
VanillaAttention：仅作为兜底方案，当 TRT kernel 不支持某些特殊配置时使用

关键行为：

TRTLLM 不支持 cross-attention 时会自动降级到 VANILLA
自注意力使用融合 QKV，cross-attention 使用分离 QKV
SageAttention 通过sage_attn_num_elts_per_blk_q/k和sage_attn_qk_int8参数控制

5.3 并行策略：从单卡到多机

TensorRT-LLM 的并行设计非常灵活，支持多种并行策略的组合，并且所有并行通信都通过 NCCL 优化：

CFG 并行

解决问题：CFG 需要同时推理正负两条 prompt，计算量翻倍
原理：将正负分支分配到不同 GPU，最后通过 all-gather 合并结果
优势：单卡计算量减半，可与其他并行度正交组合

Ulysses 序列并行

核心思想：序列切分 + all-to-all 通信，在 "序列完整 / 头分散" 和 "序列分散 / 头完整" 间切换
约束：ulysses_size 必须整除注意力头数 (Wan 头数 12，支持 1/2/3/4/6/12 卡)

Attention2D 上下文并行

核心思想：序列切分到二维 mesh，Q 在行组 all-gather，K/V 在列组 all-gather
优势：无头数约束，通信效率更高 (O (N/√P) vs Ulysses 的 O (N))
约束：必须配合 FA4 后端，暂不能与 Ulysses 组合

TRT 并行 VAE

解决问题：VAE 编解码在高分辨率视频上是严重瓶颈
原理：沿空间维度切分，边界使用 halo exchange 解决卷积依赖
TensorRT-LLM 特有的优化：所有 halo exchange 通信都通过 TRT kernel 实现，没有 Python 开销

5.4 缓存加速：TeaCache vs Cache-DiT

缓存加速是通过消除扩散过程中的冗余计算来提升速度，是目前最有效的单卡加速技术。TensorRT-LLM 对这两种技术都做了深度优化：

TeaCache

核心洞察：相邻时间步的输入变化很小，很多步的 transformer forward 是冗余的
原理：计算 time embedding 的 L1 距离，小于阈值则跳过当前步
限制：Wan 2.2 不支持，因为架构改动导致 time embedding 变化模式不同

Cache-DiT

核心优势：逐 block 粒度判断，而非整个 transformer 跳过
三种子策略：
- DBCache：残差差异判断
- TaylorSeer：高阶泰勒展开预测
- SCM：预定义步骤级计算策略
Wan 2.2 双 transformer 适配：低噪声 expert 使用更保守的缓存策略

5.5 工程优化：CUDA Graph vs torch.compile

这两项技术都是为了消除 Python 和 CUDA kernel launch 的开销：

特性	CUDA Graph	torch.compile
原理	录制 replay TRT kernel 序列	JIT 编译融合 kernel
灵活性	shape 必须完全一致	可处理动态 shape
加速来源	消除 launch 开销	融合 kernel + 消除 launch
启动开销	低 (一次 capture)	高 (需要编译)
Wan 默认	关	开

TensorRT-LLM 特有的优化：

支持SharedGraphPool，多个 TRT 引擎可以共享同一个 CUDA Graph 的显存
CUDA Graph 与 TRT 引擎无缝集成，不需要额外的代码修改
与 TeaCache/Cache-DiT 完全兼容，缓存的步骤会自动跳过 graph replay

六、实际部署最佳实践

根据不同的使用场景，我为你整理了 TensorRT-LLM 部署 Wan 模型的最佳加速技术组合：

6.1 单卡部署

显存不够：开启 NVFP4 量化
速度优先：Cache-DiT + torch.compile
极致速度：Cache-DiT + CUDA Graph (关闭 torch.compile)
质量优先：FP8 量化 + torch.compile

6.2 多卡部署

2 卡：CFG 并行 (cfg_size=2)
4 卡：CFG 并行 (2) + Ulysses (2) 或 CFG 并行 (2) + Attention2D (2×1)
8 卡：CFG 并行 (2) + Ulysses (4) 或 CFG 并行 (2) + Attention2D (2×2)
大规模部署：CFG 并行 (2) + Attention2D (8×4) = 64 卡

6.3 生产环境部署

所有加速技术全开
提前完成 Warmup 预热，触发所有 TRT 编译
使用 SharedGraphPool 共享 CUDA Graph 显存
针对常用分辨率和帧数预编译 TRT 引擎
启用 AutoTuner，针对你的硬件生成最优 kernel

七、总结

TensorRT-LLM 对 Wan 系列模型的加速系统是一个分层互补的完整体系，而不是一堆孤立技术的简单堆砌：

内存瓶颈层：FP8/NVFP4 量化，让大模型能在有限显存上运行
计算分布式层：CFG 并行、Ulysses/Attention2D、并行 VAE，将计算负载分散到多卡
计算冗余消除层：TeaCache/Cache-DiT，跳过不必要的计算
工程优化层：CUDA Graph、torch.compile、融合算子，消除系统开销
架构特化层：逐 patch 时间步、图像固定，针对 I2V 场景的专门优化

理解这些技术在 TensorRT-LLM 代码流程中的精确位置和相互关系，不仅能帮助你更好地部署 Wan 模型，也能为你自己的 AIGC 项目提供宝贵的参考。

如果你觉得这篇文章对你有帮助，欢迎点赞收藏，也欢迎在评论区交流讨论。后续我会继续深入拆解 TensorRT-LLM 对其他视频生成模型的加速实现，敬请关注。