分布式推理框架 xDit

1. xDiT 简介

xDiT 是一个为大规模多 GPU 集群上的 Diffusion **Transformers（DiTs）**设计的可扩展推理引擎。它提供了一套高效的并行方法和 GPU 内核加速技术，以满足实时推理需求。

1.1 DiT 和 LLM

**DiT（Diffusion Transformers）**是文生图与文生视频的核心网络结构。

DiT 和 LLM 推理任务特点不一样，这导致二者软件设计哲学也有差异：

1. LLM 有 Prefill 和 Decode 两阶段，分别是计算密集和访存密集的；而 DiT 的计算和 Prefill 相似是计算密集的。所以 Batching、Paged Attention 等为了增加 Decode 阶段计算密度的设计在 DiT 中不需要，其实也没有 KVCache 不断 Append 增长的现象。

2. LLM 模型很大，而序列长度有限，比如 ShareGPT 的任务生成顶多几百个 token。而 DiT 正好反过来，模型不大，推理时序列长度很长，甚至 Million 级别。这导致张量并行这种 vLLM 中第一优先级的并行方法，在 DiT 中性能很低。

3. LLM 模型大多收敛到微调 Llama 架构，而 DiT 架构则呈现出较大的差异性。例如，SD3、Latte 等网络与 Meta DiT 论文中的 DiT Block 在算子设计上存在显著差异，且 DiT Block 之间的连接方式也各不相同，常采用类似 U-Net 的 U 型结构进行链接。因此，DiT 模型的代码相对 LLM 不够简洁，往往涉及很多文件和函数，对其进行模型层面的改造也更具挑战性。

1.2 背景与挑战

DiT 模型在生成长视觉序列时的推理延迟显著增加，例如，Sora 模型的推理延迟可达数分钟，单个 GPU 无法满足实际应用的延迟要求。因此，DiT 模型的部署必然是多卡并行处理。Attention 机制导致的计算时间与序列长度呈二次增长，为了达到同样延迟指标，GPU 数应该随着序列长度平方项增长的，所有部署 Sore 很大概率需要多机多卡大集群执行单个视频生成，DiT 的并行推理扩展性至关重要。尽管大语言模型（LLMs）中常用的并行化技术，如张量并行和序列并行，通信量和激活参数正相关，在扩散模型中由于大激活小参数特点，这些技术效率不高。通信成本往往超过了并行计算的好处。因此，DiT 模型的部署仍然需要配备高带宽互连（如 NVLink+RDMA）的 GPU 集群。

随着 DiTs 中输入上下文长度的增加，Attention 机制的计算需求呈二次方增长，因此多 GPU 和多机器部署对于满足在线服务的实时需求至关重要。

1.3Diffusion Model 推理原理和特性

尽管 DiT 和 LLM 的结构相似都是 Transformers 架构，看起来工程挑战类似。

**扩散模型的训练过程：**给一个图片，经过很多步骤，每一步加噪声最后变成一个全是噪音的图片。训练过程就是预测一个 Noise Predictor 的监督学习任务，输入是中间状态图片、输入指令和 timestep 来 ground truth 就是该 timestep 的噪声。

**扩散模型的推理过程：**给一个噪声，通过 Noise Predictor 来通过多次去噪，最后变成一个有意义的图片。每一次去噪就是一个 Diffusion Step。

通过描述可知，扩散模型推理和 LLM 的有一个很大的差异点，就是扩散模型是重复计算很多相同的 Diffusion Step，而且连续的 Diffusion Step 之间输入数据和激活状态之间存在的高度相似性，称之为**输入时间冗余（Input Temporal Redundancy）。**这种冗余性表明，在进行相邻时间 Diffusion Step 去噪时，可以利用前一步骤的过期 Activation 来顶替当前步骤生成的 Activation 做一些计算，可能结果也不会差太多。

在 DiT 推理中，鉴于结构的相似性，可以用 LLM 中非常成熟的**张量并行（TP）和序列并行（SP），**也可以使用一些利用扩散模型 Input Temporal Redundancy 特性的并行方法，比如 DistriFusion 和 PipeFusion。

1.4xDiT Overview

正如 vLLM 让 transformers 具备 Batching 的能力，xDiT 让 diffusers 有并行的能力。考虑到 DiT 的特点，首先提出了一系列创新的并行方式，并可以相互混合。更重要的是，xDiT 提供了一套优雅的开发接口，针对性解决了 DiT 模型更改难度高的问题，这套开发接口尽可能可能复用 diffusers 的现有逻辑，开发者通过一些 wrapper 实现复杂的混合并行，实现高效地扩展方法。将之前的 PipeFusion 项目升级成了通用的 DiT 推理引擎，命名为 **xDiT。（**向 vLLM 致敬，X表示Scalable，希望 xDiT 成为 DiT 推理领域的 vLLM）

vLLM=huggingface transformers+Batching，xDiT=huggingface diffusers+Parallel

xDiT 是一个为大规模多 GPU 集群上的 Diffusion **Transformers（DiTs）**设计的可扩展推理引擎。它提供了一套高效的并行方法和 GPU 内核加速技术，以满足实时推理需求。

xDiT 有如下核心能力：

第一，DiT 主干网络混合并行：xDiT 支持任意混合四种基础并行策略方法。基础并行方法有：

• Pipefusion Parallel：TeraPipe 方式流水线并行方法，在弱互联，比如 pcie/以太网的网络硬件上有优势，比如以太网互联的多机多卡情境下效率远高于其他方法；

• Sequence Parallel：混合序列并行，使用 ring & ulysses sequence parallel 混合方式 USP；

• Data Parallel：在输入多个 prompt 或单个 prompt 生成多张图片时，在 image 间并行处理；

• CFG Parallel, a.k.a Split Batch：在模型使用 classifier free guidance（CFG）时可以开启，并行度恒定为 2。

xDiT 中这四种并行方式可以以任何形式排列组合，混合策略在通信量和通信拓扑两方面可以取得最佳效果，使 xDiT 可达到近似线性的扩展性。

第二，**Parallel VAE：**针对扩散模型后处理的解码 VAE 模块在高分辨率图像生成时 OOM 问题，xDiT 实现了 Patch Parallel 版本的 VAE。

第三，**简单灵活的开发接口：**xDiT 轻松帮助用户支持新的 DiT 模型。它尽可能复用开源社区 diffuser，完成对模型推理的并行化改造。这套接口也是让 xDiT 从之前的不同并行方法集合变成一个灵活开发工具。通过 xDiT 接口，实测 15 分钟可以让 Pixart-Sigma 模型并行起来！

1.5 xDit 并行方法

• 多GPU****并行方法：
  • PipeFusion：利用扩散模型的特性，通过位移Patch实现Patch级流水线并行
  • Unified Sequence Parallel (USP)：将 DeepSpeed-Ulysses 和 Ring-Attention 结合起来的并行方法
  • Hybrid Parallel：混合并行，可配置多种并行方法以优化通信模式，适应底层网络硬件
  • CFG Parallel：也称为 Split Batch，在使用无分类器指导（CFG）时激活，保持恒定的并行度为2
  • Parallel VAE：并行 VAE
• 单GPU加速方法：
  • 编译加速：xDiT 提供了 GPU 内核加速，以提高推理速度，torch.compile 和 onediff 这些编译加速与并行化方法结合使用
    • torch.compile：PyTorch 的一个特性，利用自动图优化、操作融合以及内核选择等技术，加速PyTorch模型的执行效率。在 xDiT 中，torch.compile 用于提升 GPU 上的单 GPU 性能
    • onediff：用于 PyTorch 的自动微分编译器，通过优化编译过程来提高模型的执行速度和效率
  • DiTFastAttn：xDiT还为单GPU加速提供了DiTFastAttn，它可以通过利用扩散模型不同步骤之间的冗余来降低注意力层的计算成本

xDiT 支持多种 Diffusion Transformers 模型，包括 CogVideoX、Flux.1、HunyuanDiT、SD3、Pixart、Latte 等。

针对以上不同的模型，xDiT目前实现了不同的方式对它进行并行推理加速。例如，Latte是文生视频模型，xDiT目前实现了USP方式对它进行并行推理加速，PipeFusion还在开发中，使用混合序列并行(ulysses_degree=2, ring_degree=4)可以获得最佳性能。CogVideo 是一个文本到视频的模型，xDiT 目前整合了 USP 技术（包括 Ulysses 注意力和 Ring 注意力）和 CFG 并行来提高推理速度，同时 PipeFusion 的工作正在开发中。

在使用不同GPU数目时，最佳的并行方案都是不同的，这说明了多种并行和混合并行的重要性。 例如，HunyuanDiT最佳的并行策略在不同GPU规模时分别是：在2个GPU上，使用ulysses_degree=2；在4个GPU上，使用cfg_parallel=2, ulysses_degree=2；在8个GPU上，使用cfg_parallel=2, pipefusion_parallel=4。torch.compile带来的加速效果也很可观。stable-diffusion-3最佳的并行策略在不同GPU规模时分别是：在2个GPU上，使用cfg_parallel=2；在4个GPU上，使用cfg_parallel=2, pipefusion_parallel=2；在8个GPU上，使用cfg_paralle