生成式 AI 工作负载正迅速超出单块 GPU 的显存与算力预算。对于构建媒体生成管线的推理开发者而言，难点在于跨多设备扩展时，仍要保住 TensorRT 为生产部署提供的关键优化，例如内核融合、显存规划和量化。
多设备推理支持是 TensorRT 11.0 引入的新特性，它把原生高性能多 GPU 推理带入 TensorRT 运行时，使面向边缘设备的多设备生产部署成为可能。把它与 Torch-TensorRT 结合使用，开发者可以转换并部署庞大的 PyTorch 模型，在框架之外突破单设备的显存与算力限制。
传输层由 NVIDIA NCCL 承担。NVIDIA 集合通信库（NCCL）提供高性能的多 GPU、多节点集合操作，支撑着跨数千块 GPU 的大规模模型训练。NCCL 会针对给定拓扑自动选择最优传输方式，把 NVIDIA NVLink、NVSwitch、PCIe 和 InfiniBand 抽象在统一接口之后。通过直接集成 NCCL，TensorRT 在多设备推理场景中继承了这套传输优化。新的多设备特性覆盖了 NVIDIA NCCL 分布式集合操作的完整集合，包括 AllReduce、Broadcast、Reduce、AllGather、ReduceScatter、AlltoAll、Gather 与 Scatter。
分布式推理可以用多种并行策略表达，各自在显存节省、算力扩展和通信开销之间有不同的取舍，最常见的是张量并行和上下文并行。在张量并行中，单个层的权重被切分到多块 GPU 上，每块 GPU 计算该层矩阵乘法的一个分片，再通过集合通信汇合部分结果得到完整输出。这降低了每设备的权重显存占用，当单个层的权重超过单块 GPU 显存时，它往往成为自然乃至唯一的选择，与输入序列长度或批量大小无关。在 Transformer 块中，列并行的投影（例如 QKV 和 MLP 上投影）与行并行的投影（注意力输出与 MLP 下投影）配对，使每个块只需一次 AllReduce，从而把通信开销控制在有界范围内。
在上下文并行中，输入序列沿序列维度被切分到多块 GPU 上，每块 GPU 只处理序列的一个切片，集合通信在需要时（例如注意力计算过程中）让全局序列可用。上下文并行对长序列工作负载尤为有效——注意力计算随序列长度呈二次增长，是算力与显存的主要消耗方。它同样非常适合扩散模型和 DiT 模型，因为其双向注意力绕开了因果掩码带来的负载不均衡问题。
TensorRT 11.0 引入了各种并行策略所需的 IDistCollectiveLayer 原语支持。在生成式媒体管线中，基于扩散的图像与视频生成会把很大一部分算力和显存预算花在作用于长 token 序列的注意力块内：一张高分辨率图像的隐变量或一段多帧视频，每个块可产生数万个 token 的序列，而注意力开销随序列长度二次增长，因此上下文并行直接针对了这一主导性成本。
上下文并行把序列划分到多块 GPU 上，每个 rank 处理与其序列分区对应的查询切片。一种直接的实现方式是 AllGather KV：各 rank 在计算本地注意力之前，通过 AllGather 集合操作交换各自的键（K）与值（V）分片，使每个 rank 都能对完整序列做注意力计算，代价是每个注意力块增加一次集合通信，而本地 Q × Kᵀ 矩阵乘法的规模随 rank 数成比例缩小。对视频和高分辨率图像扩散而言，这种取舍会随去噪步骤不断累积出有利效果，每一步的通信开销被序列维度的 AllGather 限制在有界范围内。
WeChat
Profile