DiTango：利用注意力空间局部性的低成本并行扩散生成

以下译文用于快速探索和学习，技术术语按需要保留英文；正式引用和精确表述请以原论文为准。

第 1 页

DiTango：具有选择性注意力状态复用的低成本并行扩散生成 Yuyang Chen Runxin Zhong Zan Zong chen-yy20@sjtu.edu.cn zhongrx24@mails.tsinghua.edu.cn zongzan@ustb.edu.cn 上海交通大学, 清华大学 北京科技大学 上海人工智能实验室 中国北京 北京 上海, 中国 北京, 中国

Hengjie Li Yuyang Jin Jidong Zhai lihengjie@pjlab.org.cn jinyuyang@tsinghua.edu.cn zhaijidong@tsinghua.edu.cn 上海人工智能实验室 清华大学 清华大学 上海, 中国 北京, 中国 北京, 中国

摘要 60 Ring_CP Comp. Head_CP Comp. TP Comp. 1 GPU: 64min Ring_CP Comm. Head_CP Comm. TP Comm. 人工智能生成内容的最新进展推动了扩散Transformer (DiTs) 在高分辨率、长 50 内容生成中的广泛应用。虽然并行化技术加速了扩散推理，但它们面临显著的可扩 (min) 40

30 持续时间 展性挑战，因为在多节点环境中存在过多的通信开销。 20 35.7 35.9 36.0 28.1 29.0 30.7 13.7 17.4 14.0 19.6 21.0 我们观察到，上下文并行 (CP) 中的序列分区表现出明显的异质性：空间上邻近的分区对注意力计算结果的贡献更大。通过将这种异质模式映射到分层通信拓扑，我们可以以较低的通信成本访问高贡献分区。这一见解激发了我们新颖的选择性注意力状态机制，该机制在去噪步骤之间策略性地平衡部分注意力计算和历史结果复用。 我们提出了 DiTango1，一个用于 DiT 生成的高效并行框架。DiTango 具有一个锚点引导的状态选择规划器，用于优化每个分区的计算复用决策，并辅以编排高效状态中心操作的运行时。这种设计在保持生成质量的同时实现了卓越的系统效率。 在流行扩散模型上的实验评估表明，DiTango 在多节点设置中实现了高达 1.9 倍的端到端加速和 3.2 倍的注意力加速，并具有近乎线性的扩展性，同时保持了与最先进方法相当的生成质量。 10 18.9 18.1 18.5 9.4 9.2 9.72026 0 2 4 8 16 24 32 GPU 数量 图 1：每节点 8 个 GPU 的多 GPU 部署的并行推理性能分解。

1 引言 扩散模型已成为人工智能生成内容 (AIGC) 的基础范式，支撑着最先进的图像和视频合成系统。随着长序列生成（例如，具有更高空间和时间分辨率的更长视频）的趋势，扩散Transformer (DiTs) [23] 已成为主导骨干网络，推动了闭源模型（例如，Sora [1]，Kling [13]）和开源模型（例如，Wan 2.2 [28]，HunyuanVideo [12]，CogVideoX [32]）的快速进步。 尽管生成质量取得了令人印象深刻的进展，但基于 DiT 的生成器在推理时仍然慢得令人望而却步。例如，Wan2.2 [28] 在单个 NVIDIA A100 GPU 上生成 5 秒 720p 视频需要超过一小时。这种延迟很大程度上由 DiT 架构中的完整注意力计算主导：与通过 KV 缓存摊销注意力成本的自回归 LLM 不同，DiT 必须在每个去噪步骤中对整个时空序列计算完整注意力。 从系统角度来看，并行推理为加速基于 DiT 的扩散生成提供了一种有前景的方法。然而，由于扩散生成的固有单样本性质（批量大小 = 1），数据并行 (DP) 和流水线并行 (PP) 提供的收益有限。因此，实际部署主要利用两种策略：张量并行 (TP) [25] 和上下文并行 (CP) [8, 11, 17]。TP 在 GPU 间分片模型权重，并使用集合操作组装中间结果，而 CP [8, 11, 17] 沿序列维度对令牌进行分区，并在注意力期间进行远程键值分区访问的通信。当注意力计算主导

• 计算机系统组织 → 并行架构；• 计算方法 → 机器学习。

关键词 扩散，并行，长上下文 arXiv:2607.15650v1

1 这是被 HPDC '26 接受发表的论文的预印本。最终认证版本将在 10.1145/3806645.3807581 处提供。代码可在 https://github.com/thu-pacman/ChituDiffusion 获取。

第 2 页

Yuyang Chen, Runxin Zhong, Zan Zong, Hengjie Li, Yuyang Jin, and Jidong Zhai

Timestep 0-9 Timestep 11-39 Timestep 40-49 In this paper, we propose DiTango, an efficient and scalable gen- Absence 10 1 eration system in distributed DiT inference. DiTango exploits the 10 2 heterogeneous contribution in CP-partitioned attention through Partition 10 3 KV a selective computation strategy: it prioritizes cost-effective KV with 10 4 pairs with high contribution and low communication costs, while MSE 0 1 2 3 4 5 6 7 reusing attention states (partial attention computation results) from Partition Distance = |Query_idx – KV_idx| previous denoising steps for less critical computations. As the two- partition example illustrated in Fig.3, the conventional full attention Figure 2: Attention computation contribution is highly un- (Fig.3a) performs costly communication to access remote KV parti- equal with distance-based decay across different CP parti- tion in step 1, despite their trivial contribution to the final result. In tions. contrast, DiTango (Fig. 3b) selectively computes crucial partition locally while bypassing expensive communication by reusing atten- tion states from previous denoising steps for trivial contributions. This design elegantly preserves the mathematical properties of at- the runtime, these data transfers can be effectively overlapped with tention computation while significantly reducing communication independent computation. overhead. However, these parallel inference solutions still fall short of However, implementing efficient attention state reuse presents practical requirements. Our profiling of Wan-14B[28] generating a two fundamental challenges: 81-frame 720p video (75,600 tokens) on NVIDIA H20s reveals that (1) Complex Selection Space – The reuse decision for each even with 8 GPUs in a single node, generation takes around 10 state must balance multiple factors: computational contri- minutes. More critically, scaling to multiple nodes often degrades bution, communication overhead, and reuse error. Making performance due to communication bottlenecks, as shown in Fig. 1. optimal selections across numerous states to achieve both This performance degradation stems from the communication- good performance and quality becomes jointly complex. intensive nature of both TP and CP. TP requires frequent collective (2) Irregular Access Patterns – Uneven selection disrupts operations for sharded model weights, while CP variants either the structured pattern in context parallelism, resulting in need All-to-All exchanges (Head-CP) or multiple rounds of P2P interleaved local memory accesses and heterogeneous com- transfers (Ring-CP) to access global Key-Value sequence. In multi- munications that are inherently difficult to parallelize. node settings, limited inter-node bandwidth severely constrains both collective and ring communication patterns. Furthermore, the To address these challenges, DiTango introduces two key com- compute-communication overlap that proves effective within a ponents: single node becomes significantly less efficient across nodes, exac- (1) Anchor-guided Selection Planner – By theoretically mod- erbating the end-to-end performance impact. eling attention state errors and designing anchor step-guided To address the performance bottleneck in parallel inference, we error prediction, we precisely identify states with minimal identify a key pattern in CP-partitioned sequences for Diffusion impact for reuse to optimize the performance-quality trade- inference: attention contributions from different partitions exhibit off. strong spatial locality. As shown in Fig. 2, we conduct a systematic (2) State-centric Parallel Runtime – We schedule attention analysis on different models and prompts under CP size=8 by mea- states with similar characteristics in groups and design a suring the Mean Squared Error (MSE) of attention outputs when series of state-centric manipulations for efficient computa- excluding specific KV partitions. Our findings reveal that partitions tion and retrieval. A specialized pipeline orchestrates these contribute highly unequally to the final attention results. For in- operations to maximize parallel computation and commu- stance, excluding the nearest partition (distance= 0) yields MSE nication efficiency. of 10−2, while distant partitions (distance≥2) contribute merely 10−4—a difference of 2-3 orders of magnitude. Attention contribu- Our contributions are summarized as follows: tions systematically decay with spatial distance that KV Partitions • We identify and analyze the spatial locality in attention closer to the query (smaller |Q_idx −KV_idx|) consistently dom- contributions during DiT inference, revealing an oppor- inate the computation, demonstrating clear spatial locality. This tunity to selectively reuse attention states based on their spatial locality persists across diverse models, different DiT layers contribution significance and communication costs. and various input prompts—indicating it is an intrinsic property of • We design and implement DiTango, a high-performance diffusion attention rather than an artifact of specific configurations. generation system that leverages theoretical error modeling This observation unveils a compelling optimization opportunity: to make precise reuse decisions while maintaining efficient by aligning locality priorities with the distributed system topol- execution through group-wise state management. ogy—mapping high-contribution partitions to nodes with efficient • Extensive evaluations on Wan2.1 [28] and HunyaunVideo [12] intra-node communication while placing low-contribution ones demonstrate that DiTango achieves near-linear scalabil- across nodes—we can strategically compute only the most cost- ity across 32 GPU cards, delivering up to 1.9× end-to-end effective partitions and reuse historical results for less critical ones. speedup and 3.2× core attention speedup while maintaining This contribution-aware approach enables substantial efficiency superior generation quality compared to state-of-the-art gains without sacrificing generation quality. frameworks [7, 15, 27, 38].

2

第 3 页

DiTango：通过选择性注意力状态复用实现低成本并行扩散生成

GPU0 时间线 GPU0 时间线 计算 KV 访问 计算 计算 复用 先前 Q 0 Q 0 Q 0 旁路 注意力状态 发送 KV0 KV 0 接收 KV1 KV 1 KV 0 KV 访问 关键注意力 平凡注意力 关键注意力 获取平凡状态

组合 组合 注意力状态 0 1 0 注意力状态 0 1 0 输出分区 输出分区

(a) 全序列分块注意力（基线）：KV 访问带来大量通信但贡献甚微。 (b) DiTango 结合选择性序列计算：复用 访问带来大量通信，但贡献很小。 先前去噪步骤的注意力状态以绕过 昂贵的操作

图 3：全序列注意力与通过状态复用进行选择性计算的 DiTango 之间的比较。

2 背景与动机 MLP 注意力 Y = f(X) Q, K, V = g(Y ) 逐头注意力 2.1 背景 s X W1 s Y W2 s Q K V s Z 2.1.1 扩散 Transformer。扩散 Transformer (DiTs)[23] 通过迭代去噪过程生成内容，通常需要 20-50 步。在每一步中，相同的 transformer 模型处理潜在表示，逐渐将随机噪声细化为高质量内容。与传统 U-Net 架构不同，DiTs 利用基于 transformer 的设计，通过其注意力机制擅长处理长序列。最近的 DiT 架构[12, 20, 26, 28, 32]广泛采用了 3D 全注意力，将空间和时间维度统一为单个序列进行处理。虽然这种统一的注意力方法显著提高了生成质量和时间一致性，但注意力操作主导了计算成本，在所有去噪步骤中消耗了超过 70% 的总生成时间。

(a) 张量并行

2.1.2 扩散生成的并行方法。并非所有并行方法都适合 DiT 推理。数据并行 (DP) 和流水线并行 (PP) 对扩散生成无效，因为单样本推理模式阻止了批级并行化。这使得 DiT 加速主要有两种策略： 如图 4(a) 所示，张量并行 (TP) 将模型权重跨设备分片。线性层需要 All-Reduce 操作来聚合部分结果，而注意力模块需要 All-Gather 通信来收集逐头结果。尽管 TP 内存效率高，但它 suffers from 频繁的同步开销。 上下文并行 (CP) 通过将序列跨设备分区来更好地解决长序列生成问题，消除了线性层中的通信，仅需要键值 (KV) 远程访问用于注意力。出现了两种主要的 CP 变体： Head-CP [11]（图 4(b)）采用 all-to-all 通信重新分发张量布局，实现局部化的头计算。 Ring-CP[17]（图 4(c)）保持序列级分区，但通过环形通信交换 KV 块，允许每个 GPU 逐步访问完整序列。像统一序列并行 [8] 这样的最新工作提供了框架，可以根据工作负载特征在这些策略之间动态切换。

(b) 具有逐头注意力的上下文并行

2.1.3 扩散模块特征复用。特征复用已成为扩散模型中一种流行的优化技术，利用连续去噪步骤通常产生相似中间特征这一观察结果。这种相似性使得通过战略性地缓存和重用先前步骤的特征来实现无需训练但有损加速成为可能，从而有效地绕过某些计算依赖。不同的方法针对模型的不同方面：Delta-DiT [2] 重用逐层计算特征，金字塔注意力广播 (PAB)[38] 专注于注意力输出

3

第 4 页

Yuyang Chen, Runxin Zhong, Zan Zong, Hengjie Li, Yuyang Jin, and Jidong Zhai

(a) Partition Computation Contribution (b) Partition Communication Topo. as a simpler distance-dependent pattern when observed at the par-

0 0 tition level. 1 1 2 2 3 3 4 Intra Node Cross Node Index Index 4 5 5 2 This computational pat- 2.2.3 Insight: System-Pattern Alignment. 6 6 7 7 8 tern naturally aligns with the hierarchical communication topology 8 9 9 3 Chunk Chunk 10 log(Contribution 10 in distributed systems: 11 11 Cross Node Intra Node 12 12 13 4 13 • Local: KV stored in local memory with negligible access over- 14 14 Query Score) Query 15 15 0 1 2 3 4 5 6 7 8 9 101112131415 0 1 2 3 4 5 6 7 8 9 101112131415 head KV Chunk Index KV Chunk Index • Intra-node: KV from different GPUs within the same node via high-bandwidth NVLink Figure 5: Spatial locality of CP attention in typical cluster: • Cross-node: KV from remote nodes via low-bandwidth Infini- (a) KV partition contribution scores and (b) communication Band, incurring substantial cost costs across partitions, demonstrating natural alignment be- By mapping spatially adjacent partitions to hardware units with tween computational importance and system topology. efficient communication, we can achieve perfect alignment between computation importance and communication efficiency. This align- reuse, while TeaCache[15] and TaylorSeer[18] dynamically reuse ment presents a compelling opportunity: prioritize high-importance, DiT’s overall outputs based on input patterns. However, feature low-cost computations while replacing expensive operations of mini- reuse inevitably introduces approximation errors, making the de- mal contribution with historical results. This insight enables us to sign of reuse strategies crucial for maintaining generation quality. explore previously unreachable optimization space in the scalability- quality trade-off: by prioritizing computations with high contribution- 2.2 Motivation to-cost ratios while reusing historical results for expensive yet less critical operations, we can maximize communication efficiency with2.2.1 Limitation: The Scalability-Quality Dilemma. Recent gen- minimal quality impact.eration frameworks have pursued higher performance by com- A straightforward approach would be caching KV partitions frombining parallel methods with feature reuse strategies. For exam- previous steps. However, this faces a critical memory bottleneck:ple, VideoSys [27] integrates Dynamic-SP [37] parallelization with storing all KV partitions requires 2 × 𝐿× 𝐻× 𝐷× 𝑁× 𝐶memory,PAB [38]-based attention output reuse, while SGLang-Diffusion [39] where 𝐿=sequence length, 𝐻=heads, 𝐷=hidden size, 𝑁=layers, andcombines Unified-SP [8] parallelization with Cache-DiT’s TaylorSeer [18]- 𝐶=CFG batch size. For Wan-14B (2 × 75600 × 40 × 128 × 40 × 2 inbased reuse mechanism. These approaches demonstrate impressive BF16), this reaches 58GB per GPU, causing OOM when combinedacceleration within single-node environments. However, their per- with model weights and activations.formance deteriorates significantly in multi-node scenarios due to limited cross-node communication bandwidth. 2.2.4 Solution: Attention State Reuse. To leverage this system- This cross-node communication bottleneck creates a fundamen- pattern alignment while addressing the memory challenge, we tal dilemma: to maintain reasonable performance scaling, these introduce Attention State, commonly used in partitioned parallel frameworks must adopt more aggressive feature reuse strategies to attention computation [5, 16, 17, 33], as our optimization medium. reduce communication overhead. However, such aggressive reuse For sequence partition 𝑖, an Attention State AS𝑡(𝑖) at timestep 𝑡 inevitably leads to substantial quality degradation. Conversely, at- comprises output OUT𝑡(𝑖) and log-sum-exp LSE𝑡(𝑖): tempting to preserve generation quality by limiting feature reuse results in poor scaling efficiency due to increased cross-node com- munication. This inherent trade-off between distributed scalability LSE𝑡(𝑖) = log ∑︁ exp(q𝑡· k𝑗) and generation quality highlights the limitations of treating par- OUT𝑡(𝑖) 𝑗∈I𝑖 AS𝑡(𝑖) = , (1) exp(q𝑡· k𝑗)allelization and feature reuse as independent optimizations, sug- LSE𝑡(𝑖) OUT𝑡(𝑖) = ∑︁ v𝑗 gesting the need for a more integrated approach to distributed exp(LSE𝑡(𝑖)) 𝑗∈I𝑖 generation acceleration. where q, k, v are query, key, and value vectors. 2.2.2 Analysis: Spatial Locality of Partition Contribution. To under- A key property of Attention States is their composability—they stand this challenge, we analyze the computation contribution of 16 can be composed associatively and commutatively. For partitions 𝑖 partitions by measuring output differences when omitting specific and 𝑗: KV computations, averaged across DiT layers and timesteps using 50 prompts on Wan2.1. Fig. 5(a) reveals a striking pattern of Spatial Locality: computational importance strongly correlates with spatial proximity. Query-Key pairs from nearby partitions contribute sig- nificantly more to attention outputs, with importance diminishing as positional distance increases. This Spatial Locality pattern, consis- tent across different models and inputs, offers a natural opportunity for optimization when considering system topology. Theoretically, this pattern demonstrates how fine-grained token- level attention sparsity studied in previous works [30], manifests

第 5 页

DiTango：通过选择性注意力状态复用实现低成本并行扩散生成

去噪步骤 t-1 拓扑感知序列分区 键值对 跨 GPU 发送/接收 GPU0 0 1 2 3 4 5 6 7 DiTango 0 1 2 3

序列级注意力 输出 组级 锚点 误差建模 查询 组织 步骤 0 0 1 2 3 0 ASE(i) = w(i, t) ⋅δ(i) 0 1 2-3 注意力状态 组合与存储 4-5 6-7 去噪步骤 t 决策：计算 | 复用 键值对 GPU0 跨 GPU 发送/接收 注意力状态缓存 锚点引导的选择规划 (§4) 0 1 2 ⊕3

序列级注意力 复用 查询 输出 以状态为中心的 运行时编排 操作 0 0 1 2 ⊕3 0 注意力状态 …

图 6：选择性注意力状态复用机制（CP=4，rank=0）。 图 7：DiTango 概览

从远程分区（𝐴𝑆𝑡−1(2) 和 𝐴𝑆𝑡−1(3)）组合状态并 缓存结果。 流，利用无依赖的复用操作来掩盖 在步骤 𝑡 中，我们仅为本地和邻近分区计算新状态 跨节点通信延迟并最大化资源利用率 （𝐴𝑆𝑡(0) 和 𝐴𝑆𝑡(1)），同时复用缓存的组合状态 (𝐴𝑆𝑡−1(2 ⊕3)) 用于远程分区，从而有效规避了键值对 传输和注意力计算的开销。 这种基于注意力状态的方法相比原始的键值对复用提供了两个关键优 势：(1) 通过消除远程分区的通信和计算成本，显著加速；(2) 通过紧凑的状态表示（仅为键值对大小的一半）和灵活的状态组合，提高了内存效率。

3 DiTango 概览 我们提出了 DiTango，这是一个并行推理系统，通过选择性注意力状态复用优化扩散生成。如图 7 所示，DiTango 有两个关键组件：锚点引导的选择规划 (§4)。基于对注意力状态计算误差传播的分析，DiTango 对复用引入的累积误差进行建模，并采用周期性锚点步骤进行在线决策。在锚点步骤，系统在所有分区上执行新鲜计算以重置误差累积，同时利用观察到的注意力权重分布来预测后续步骤的复用误差。通过组级组织——将序列划分为与设备拓扑对齐的连续组——以及误差预算约束，规划器为每个组生成平衡精度和效率的计算/复用决策。以状态为中心的并行运行时 (§5)。为了处理由计算/复用决策引起的异构计算和通信模式，DiTango 将注意力操作解耦为专门的状态中心原语：跨组传输通过设备间的对称点对点交换整合所需的键值对分区；组内组合通过本地化设备组内的环形通信计算新鲜的注意力状态；动态组组合在检测到内存压力时自适应地合并缓存状态。运行时编排这些操作，跨越异步计算和通信流，利用无依赖的复用操作来掩盖跨节点通信延迟并最大化整个管道的资源利用率。 总之，这些组件使 DiTango 能够利用扩散注意力中的空间局部性，同时在分布式加速器上保持准确性保证和高效的并行执行。

4 锚点引导的选择规划 4.1 误差建模 4.1.1 注意力状态误差传播。我们首先分析注意力状态中的误差如何通过组合传播。考虑两个具有误差 𝛿OUT𝑖 和 𝛿LSE𝑖 的注意力状态，其中 𝑖∈{1, 2}。 定义归一化权重： (3) 𝑤𝑖= 锚点引导的选择规划 (§4)。基于误差 𝑒LSE𝑖 𝑒LSE1 + 𝑒LSE2 传播分析，DiTango 对复用引入的累积误差进行建模，并采用周期性锚点步骤进行在线决策。在锚点步骤，系统在所有分区上执行新鲜计算以重置误差累积，同时利用观察到的注意力权重分布来预测后续步骤的复用误差。通过组级组织——将序列划分为与设备拓扑对齐的连续组——以及误差预算约束，规划器为每个组生成平衡精度和效率的计算/复用决策。 组合操作传播误差如下： (4) 𝛿OUT = 𝑤1𝛿OUT1 + 𝑤2𝛿OUT2 + 𝑤1𝑤2(OUT1 −OUT2)(𝛿LSE1 −𝛿LSE2) 𝛿LSE = 𝑤1𝛿LSE1 + 𝑤2𝛿LSE2 对于 𝑛 个分区，累积输出误差满足： 𝑛 ∥𝛿OUT∥2 ≤ ∑︁ 𝑤𝑖∥𝛿OUT𝑖∥2 𝑖=1 (5) 以状态为中心的并行运行时 (§5)。为了处理由计算/复用决策引起的异构计算和通信模式，DiTango 将注意力操作解耦为专门的状态中心原语：跨组传输通过设备间的对称点对点交换整合所需的键值对分区；组内组合通过本地化设备组内的环形通信计算新鲜的注意力状态；动态组组合在检测到内存压力时自适应地合并缓存状态。运行时编排这些操作，跨越异步计算和通信 + ∑︁ 𝑤𝑖𝑤𝑗∥OUT𝑖−OUT𝑗∥2|𝛿LSE𝑖−𝛿LSE𝑗| 𝑖<𝑗 计算和通信模式，DiTango 将注意力操作解耦为专门的状态中心原语：跨组传输通过设备间的对称点对点交换整合所需的键值对分区；组内组合通过本地化设备组内的环形通信计算新鲜的注意力状态；动态组组合在检测到内存压力时自适应地合并缓存状态。运行时编排这些操作，跨越异步计算和通信 第二个项是二阶项，可以安全地忽略：(1) 当一个分区占主导时，乘积 𝑤𝑖𝑤𝑗 很小，且 (2) |𝛿LSE𝑖−𝛿LSE𝑗|（误差之差）通常远小于单个误差。由此可得： 𝑛 ∥𝛿OUT∥2 ≈ ∑︁ 𝑤𝑖∥𝛿OUT𝑖∥2 (6) 𝑖=1 5

第 6 页

Yuyang Chen, Runxin Zhong, Zan Zong, Hengjie Li, Yuyang Jin, and Jidong Zhai

Part 0 (local) Part 1 Part 2 Part 3 and communication costs, leading to consistent selection decisions. Layer 10 Layer 20 Consequently, the 𝑝 states are divided into 𝑝/𝑔 groups, each con- 1.0 1.0 taining 𝑔 states. In particular, because the local state is handled1 separately, the intra-node group containing it has a reduced size of 0.5 0.5 Prompt w(i,t) 𝑔−1. 0.0 0.0 0 10 20 30 40 0 10 20 30 40 Organizing partitions as group introduce following benefits: 1.0 1.0 Grouping strategy naturally reduces cross-term errors in Equa- 2 tion 5 as spatially adjacent states have similar outputs. It also en- 0.5 0.5 ables efficient unified scheduling and memory management through Prompt w(i,t) 0.0 0.0 group-level attention state composition. Building upon this group 0 10 20 30 40 0 10 20 30 40 Timestep Timestep organization and our error model, we now present the selection strategy that determines which attention state groups to compute Figure 8: Partition importance 𝑤(𝑖,𝑡) with different prompts or reuse at each timestep. and layers. 4.2.2 Strategy Formulation. For each timestep 𝑡, we maintain a set of attention state groups G = {𝐺0,𝐺1, …,𝐺𝑚−1} where 𝑚= ⌈𝑝/𝑔⌉. Our goal is to determine for each group 𝐺𝑖:4.1.2 Online Attention State Error Model. Based on Equation 6, we model the Attention State Error (ASE) of partition 𝑖’s cached state AS𝑡𝑐(𝑖) at current timestep 𝑡as: ( COMPUTE compute fresh state Decision(𝐺𝑖,𝑡) = REUSE(𝑡𝑐,𝑖) reuse cached state from 𝑡𝑐,𝑖 ASE(𝑖,𝑡,𝑡𝑐) = 𝑤(𝑖,𝑡) · 𝛿(𝑖,𝑡,𝑡𝑐) (7) (11) where 𝑡𝑐 is the cache timestep, 𝑤(𝑖,𝑡) captures partition impor- where 𝑡𝑐,𝑖 denotes the last computation timestep for group 𝐺𝑖. tance, and 𝛿(𝑖,𝑡,𝑡𝑐) models temporal error growth over cache age The strategy balances quality (bounded error) and efficiency (maxi- 𝜏= 𝑡−𝑡𝑐. mized reuse). Partition Importance Weight. The weight 𝑤(𝑖,𝑡) corresponds to 4.2.3 Error-based Selection. For each group 𝐺𝑖 at timestep 𝑡, we 𝑤𝑖 in Equation 6. Since computing exact 𝑤(𝑖,𝑡) = 𝑒LSE𝑡(𝑖)/Í𝑘𝑒LSE𝑡(𝑘) compute its aggregated Attention State Error by summing errors requires full attention, we leverage the empirical observation that from all partitions within the group: LSE distributions remain stable over 𝜏≤5-10 steps (Figure 8). We update 𝑤(𝑖,𝑡) only at anchor steps 𝑡𝑎—timesteps with full attention computation: ASE(𝐺𝑖,𝑡) = ∑︁ 𝑤(𝑗,𝑡) · 𝛿(𝑗,𝑡,𝑡𝑐,𝑖) (12) 𝑒LSE𝑡𝑎(𝑖) Given an error threshold 𝜖, the selection decision follows: 𝑤(𝑖,𝑡) = 𝑤(𝑖,𝑡𝑎) = , ∀𝑡∈(𝑡𝑎,𝑡𝑎+ 𝜏max] (8) Í𝑛𝑘=1 𝑒LSE𝑡𝑎(𝑘) ( COMPUTE if ASE(𝐺𝑖,𝑡) > 𝜖 Decision(𝐺𝑖,𝑡) = (13) Age Penalty Error. Computing 𝛿(𝑖,𝑡,𝑡𝑐) = ∥AS𝑡(𝑖) −AS𝑡𝑐(𝑖)∥2 REUSE(𝑡𝑐,𝑖) otherwise directly contradicts our reuse goal. Instead, we exploit that the local partition 𝑖loc (the diagonal partition, always computed) serves as a 4.2.4 Anchor Step Enforcement. To prevent unbounded error ac- zero-cost error indicator. cumulation and ensure periodic weight updates (Eq. 8 and 9), we At anchor step 𝑡𝑎, we compute the scale ratio: enforce anchor steps where all groups perform full computation. ∥AS𝑡𝑎(𝑖)∥2 An anchor step is triggered when: 𝛼(𝑖) = (9) ∥AS𝑡𝑎(𝑖loc)∥2 True if ∀𝐺𝑖: ASE(𝐺𝑖,𝑡) > 𝜖 For timesteps 𝑡> 𝑡𝑎, we extrapolate partition 𝑖’s error from the IsAnchor(𝑡) =  True if 𝛿local(𝑡) > 𝜖local (14) local state change: False otherwise 𝛿(𝑖,𝑡,𝑡𝑎) = 𝛼(𝑖) · ∥AS𝑡(𝑖loc) −AS𝑡𝑎(𝑖loc)∥2 (10)  where 𝛿local(𝑡) is the relative error of the local state. Unlike Our evaluation guarantees that remote states evolve proportion- prior static thresholds, 𝜖 is adapted online as a percentile of the ally to the local state, scaled by their relative magnitudes. Since global anchor correction errors to match the target acceleration AS𝑡(𝑖loc) is computed every step, Equation 10 incurs nearly zero ad- ratio, while the static 𝜏max is replaced by this dynamic local drift ditional cost while achieving strong empirical accuracy (𝑅2 > 0.95, condition. At anchor steps, we: (1) force all groups to compute, (2) Section 6.5). update partition importance weights 𝑤(𝑗,𝑡), (3) update scale ratios 𝛼(𝑗), and (4) reset cache timestamps 𝑡𝑐,𝑖←𝑡. 4.2 Group-wise Compute & Reuse Selection Algorithm 1 summarizes our selection strategy. At each timestep, 4.2.1 Group-wise organization. With a Context Parallelism (CP) size of 𝑝, KV tensors are we first compute local attention and predict errors for all groups distributed across GPUs, each corresponding to an attention stare. (lines 2-4). Then we check anchor conditions (lines 5-7). If anchoring To achieve an efficient planning granularity, we group spatially is triggered, all groups compute and metadata is updated (lines 8- adjacent states. This is motivated by the observed Spatial Locality 13); otherwise, groups selectively compute or reuse based on error that adjacent partitions exhibit similar computational contributions thresholds (lines 14-21).

6

第 7 页

DiTango：通过选择性注意力状态复用实现低成本并行扩散生成

算法 1 注意力状态选择策略 1: 输入：时间步 $t$，组 $G$，缓存比率 $R$，局部阈值 $\epsilon_{local}$ 2: 计算 $\delta_{local}(t)$ 3: for each group $G_i \in G$ do 4: 计算 $ASE(G_i, t) \leftarrow \sum_{j \in G_i} w(j, t) \cdot \delta(j, t, t_{c,i})$ 5: end for 6: all_exceed $\leftarrow \forall G_i: ASE(G_i, t) > \epsilon$ 7: local_drift $\leftarrow \delta_{local}(t) > \epsilon_{local}$ 8: is_anchor $\leftarrow$ all_exceed $\lor$ local_drift 9: if is_anchor then 10: for each group $G_i$ do 11: Decision($G_i, t$) $\leftarrow$ COMPUTE 12: end for 13: 根据 $R$ 和锚点误差更新 $\epsilon$ 14: 为所有分区 $j$ 更新 $w(j, t)$ 和 $\alpha(j)$ 15: $t_{c,i} \leftarrow t$ for all $G_i$ 16: else 17: for each group $G_i$ do 18: if $ASE(G_i, t) > \epsilon$ then 19: Decision($G_i, t$) $\leftarrow$ COMPUTE, $t_{c,i} \leftarrow t$ 20: else 21: Decision($G_i, t$) $\leftarrow$ REUSE($t_{c,i}$) 22: end if 23: end for 24: end if 25: return {Decision($G_i, t$)}

5 以状态为中心的注意力运行时

基于选择阶段（第 4 节）得出的计算/复用决策，运行时必须高效执行异构分区访问模式，同时最大化计算与通信的重叠并管理内存约束。我们设计了一套以状态为中心的操纵（manipulations），并在流水线运行时系统中协调它们。

5.1 以状态为中心的操纵

5.1.1 跨组 KV 传输。当选择组 $G_i$ 进行计算（COMPUTE）时，我们必须从远程设备收集其 KV 分区，以启用本地注意力计算。在时间步 $t$，令 $C_t = \{G_i | \text{Decision}(G_i, t) = \text{COMPUTE}\}$ 为需要新鲜计算（fresh computation）的组集合。由于我们的计算/复用策略在所有设备上全局一致，每个 GPU 可以根据组内分区距离独立确定其通信对等体。

我们充分利用分层带宽：节点内传输利用 NVLink，而节点间传输使用 InfiniBand，确保对可用互连的高效利用。

5.1.2 组状态计算。一旦 KV 分区传输完成，每个组独立通过基于环的组合注意力（ring-based composed attention）计算其新的注意力状态。如图 10(b) 所示，组 $G_i$ 内的设备形成一个通信环，KV 分区在环中循环流动，同时注意力计算继续进行。

计算由 $g$ 个连续步骤组成（其中 $g = |G_i|$ 为组大小）。在每一步中，每个设备：(1) 使用其当前的本地 KV 分区和查询计算部分注意力；(2) 使用公式 4 中的组合操作，将此结果与其累积的注意力状态进行增量组合；(3) 将其 KV 分区发送给环中的下一个设备，同时从前一个设备接收一个新的分区。

这种环注意力模式将 $(g-1)$ 次 KV 传输与注意力计算重叠。此操作的目标是在每个参与设备上高效地为组 $G_i$ 生成完整的组合注意力状态 $AS_t(G_i)$。经过 $g$ 步后，每个设备已处理组中的所有分区并获得完整的组合状态，随后将其存储在注意力状态缓存中以供未来复用。基于组的计算边界再次利用分层带宽，因为环循环发生在由跨组 KV 传输建立的局部设备集内。

5.1.3 动态组状态组合。DiTango 持续监控系统内存使用情况。当内存消耗接近预设阈值 $M_{max}$ 时，系统触发动态组状态组合（Dynamic Group State Compose）以防止内存溢出故障。此操作将现有的注意力状态合并为更粗粒度的组，从而指数级减少内存开销。

给定当前组大小 $g$ 和缩放因子 $k$（通常 $k=2$），该操作为：

$$ \text{GroupCompose}(g \to kg): AS_t(G_i) \oplus AS_t(G_{i+1}) \oplus \dots \oplus AS_t(G_{i+k-1}) \to AS_t(G'_{\lfloor i/k \rfloor}) \quad (15) $$

此合并递归进行，直到内存使用量降至 $M_{max}$ 以下或发生锚点步骤（此时重新计算所有状态并重置为原始组大小 $g$）。由于注意力状态的可组合属性（公式 4），组合状态仍然有效，可用于复用。

5.2 运行时编排

如图 10(a) 所示，跨组 KV 传输通过对称的点对点（P2P）通信进行操作。对于每个组 $G_i \in C_t$，持有 $G_i$ 分区的设备交换其 KV 对：每个 GPU 执行 SendRecv 操作，同时向一个对等体发送其本地 KV 分区，并从另一个对等体接收不同的分区。例如，在一个跨越节点 0 和节点 1 的组中，GPU 0 将其分区发送给 GPU 4（跨节点），而 GPU 4 发送给 GPU 0，从而形成对称交换模式。

此操作的目标是将每个所需的 KV 分区重新定位到最接近其消费位置的设备上，从而将计算选定的组的所有分区整合在一个局部设备集中。通过在组粒度上组织通信，

DiTango 运行时通过双异步流——计算和通信——编排以状态为中心的操纵，以最大化重叠同时限制内存使用。图 9 说明了流水线执行流程。

内存感知初始化。在每个时间步，运行时首先调用动态组状态组合来监控内存消耗。如果缓存的注意力状态接近阈值 $M_{max}$，系统会将现有状态合并为更粗的组，指数级减少内存开销。如果组大小 $g$ 发生变化，这可能会触发计算/复用决策的重新规划，确保选择策略在更新的粒度下仍然有效。

7

第 8 页

Yuyang Chen, Runxin Zhong, Zan Zong, Hengjie Li, Yuyang Jin, and Jidong Zhai

动态组 选择 操作 状态组合 局部注意力 规划 组状态计算 组状态计算

计算 流 注意力 (0) 复用 (1) 注意力 (2) ⊕ 注意力 (3) ⊕ 复用 (6,7) 注意力 (4) ⊕ 注意力 (5) ⊕

通信 跨组变换 组内点对点 跨组变换 组内点对点 流

分区状态 04 15 26 37 04 15 26 37 04 15 26 37 04 15 26 37 04 15 26 37 04 15 26 37

未处理 已计算 已复用

图 9：运行时编排时间线。计算流和通信流异步执行，将局部注意力与跨组传输重叠，并对多个组状态计算进行流水线处理。

灵活的复用调度。一个关键的设计见解是利用状态复用操作来掩盖通信延迟。与依赖于先前 KV 访问完成的“新鲜注意力状态计算”不同，复用操作——直接从本地内存加载缓存的 $AS_{t-1}(G_j)$ ——没有数据依赖关系，可以立即执行。运行时策略性地调度这些轻量级的复用操作，以填充跨组传输中由跨节点通信引起的计算气泡，从而在整个流水线中保持高 GPU 利用率。 节点 0 节点 1 节点 0 节点 1 (a) 跨组 KV 传输 (b) 组状态计算 组 组 组 组 注意力状态计算 注意力通信 注意力通信 从本地内存加载——没有数据依赖关系，可以立即执行。运行时策略性地调度这些轻量级的复用操作，以填充跨组传输中由跨节点通信引起的计算气泡，从而在整个流水线中保持高 GPU 利用率。 1 点对点 5 1 5 注意力通信 注意力通信 组 组 组 组 注意力通信 注意力通信 复用操作 2 点对点 6 2 6 通信 3 点对点 7 3 7 利用率。 图 9 中的“分区状态”可视化展示了 DiTango 注意力的 i : 位于 GPU i 中的 KV 分区 进度：每个分区从“未处理”（灰色）转变为通过组状态 计算“已计算”（绿色）或通过缓存状态加载“已复用”（黄色）。在 图 10：以状态为中心的通信操作。(a) 跨组 KV 传输：对称的点对点交换重新组织设备间的 KV 分区，以将所需数据共置。(b) 组状态计算：组内 KV 循环使得每个组内能够进行组合注意力状态计算。 具有 8 个分区的示例中，分区 0、2、3、4、5 被计算，而 1、6、7 被复用。我们的选择算法（算法 1）保证在每个时间步结束时，所有 $n$ 个分区都通过新鲜计算或缓存复用得到覆盖，确保最终注意力输出的完整性。

局部注意力和选择规划。计算流随后在查询的头分区上执行局部注意力，这与通信流的第一个跨组 KV 传输重叠。同时，运行时为下一个时间步执行选择规划，根据累积的错误预算（算法 1）确定哪些组需要计算或复用。这种规划与计算的重叠隐藏了决策延迟。 组状态计算流水线。对于每个组 $G_i \in C_t$（被选中进行计算的组集合），运行时依次执行： (1) 驱逐陈旧状态：从缓存中移除 $AS(G_i)$，因为它将被新鲜状态 $AS_t(G_i)$ 替换。 (2) 组状态计算：在由跨组传输建立的本地化设备集内执行基于环的组合注意力（第 5.1.2 节）。通信流执行 $(g-1)$ 次组内点对点传输，与计算流中的注意力计算重叠。 (3) 存储新鲜状态：缓存 $AS_t(G_i)$ 以供将来复用。 该流水线对所有 $|C_t|$ 个被选中计算的组重复执行。运行时动态调度状态复用操作——直接加载缓存的 $AS_{t-1}(G_j)$ 用于复用的组——以填充跨组传输中由跨节点通信延迟引起的计算气泡。这种灵活的交错确保了整个流水线的高 GPU 利用率。

6 评估 DiTango 使用 7K 行 Python 代码实现 (https://github.com/thu-pacman/ChituDiffusion)，确保与现代视频生成模型无缝集成。我们评估了其性能和质量改进，随后进行了详细的消融研究以剖析驱动这些增益的关键因素。 6.1 评估设置 平台。我们在 4 个节点上评估 DiTango，每个节点配备 8 个 NVIDIA H20 GPU（148 TFLOPS FP16，96GB），节点内通信使用 NVLink（900GB/s），节点间通信使用 InfiniBand（400Gbps）。我们的软件栈包括 PyTorch 2.5.0、Python 3.10.16、CUDA 12.4 和 GCC 11.4。 工作负载。我们使用 VBench Leaderboard [10] 上最好的开源模型评估 DiTango：Wan2.1-14B、Wan2.1-1.3B [28] 和 HunyuanVideo [12]。表 1 总结了它们的关键特征。张量形状定义为 [batch size, sequence length, head number, hidden size]。Wan2.1 模型使用无分类器引导 [9]，潜在批量大小设置为 2。为了内存高效的推理，我们在 Wan2.1-14B 中将 T5 编码器卸载到 CPU。 基线。我们将 DiTango 与五种最先进的基线进行评估：张量并行 [25]，表现最好的上下文并行 (CP)

第 9 页

DiTango：通过选择性注意力状态复用实现低成本并行扩散生成

TP CP xDiT VideoSys SGLD DiTango

Wan-1.3B Wan-14B HunyuanVideo 1 GPU: 3272 1 GPU: 68715 1 GPU: 98544 (ms) 2000 20000 2856 20683 18742 22358 26147 20000 326742501325013 29923 32456 33684 34125 10000 1.6× 1.9× 1.1× 1.5× 1.9× 1.1× 1.2× 2.1× Latency 0 0 0 4 8 16 24 32 4 8 16 24 32 4 8 16 24 32 Number of GPUs Number of GPUs Number of GPUs

(a) DiT End-to-end latency.

400 (ms) 50 63.5 335.8337.2336.0336.6 361.6361.6 200 200 25 2.0× 1.9× 1.6× 2.1× 2.6× 1.1× 1.7× 1.7× 3.2× Latency 0 0 0 4 8 16 24 32 4 8 16 24 32 4 8 16 24 32 Number of GPUs Number of GPUs Number of GPUs

(b) Core Attention latency.

图11：DiTango的性能评估。运行时间过长的基线方法的条形图被截断，其执行时间标记在条形图上。DiTango条形图上方的数字显示其在多节点推理中相对于最佳基线的加速比。

表1：评估模型规格 DiTango在分布式计算环境中的特定有效性， 其中通信开销通常成为显著瓶颈。 Model Size Type Frames Reso. Tensor Shape Wan2.1-14B 14B FP16 81 720×1280 [2, 76K, 40, 128] HunyuanVideo 13B BF16 129 720×1280 [1, 119K, 48, 64] 尽管DiTango在Wan-14B和Hun- Wan2.1-1.3B 1.3B FP16 81 832×480 [2, 33K, 12, 128] yuanVideo上表现出强大的可扩展性，但在Wan- GPUs上，其性能在超过24个GPU时恶化，因为有限的计算量不足以掩盖节点内 通信延迟，导致出现通信受限行为。 i 介于序列级CP [17]和头级CP [11]之间的并行变体， 通过xDiT [7]实现的混合CP [8]，以及两种基于特征复用的 6.3 核心注意力性能与效率基于有损加速系统 – VideoSys [27, 38] 使用PAB (Pyramid Attention Broadcast, 预热/冷却=4, 广播 (金字塔注意力广播, warm-up/cool-down=4, broadcast 范围=3) 和 SGLang-Diffusion [39]（记为SGLD），后者 range=3) 和 SGLang-Diffusion [39]（记为SGLD），后者 集成了Cache-DiT [6]。 整合了Cache-DiT [6]。

6.2 DiT端到端性能 尽管DiTango在绕过约三分之一注意力计算的激进优化策略下， 其性能优势甚至超过了VideoSys [27]。SGLang-Diffusion [39] 由于 我们对端到端DiT推理延迟和核心注意力机制性能进行了全面评估。 共享的并行注意力实现，其核心注意力性能与xDiT [8]相当，但其优化 在端到端性能方面，DiTango显示出显著的效率提升，在32个H20 GPU上 仅限于绕过粗粒度的DiT步骤，而非细粒度的注意力操作。 执行Wan-14B [28]推理时，相比最先进的基线方法实现了高达1.9×的加速 比（图11a）。虽然现有的并行方法如TP、CP和xDiT通过完整注意力计算 保持模型准确性，但DiTango通过选择性操作绕过实现了更优越的性能， 且未损害生成质量。 在单节点场景（≤8个GPU）中，DiTango的处理时间略长于传统的 有损加速框架VideoSys和SGLang-Diffusion，这主要是由于其更高的 计算比率旨在保留生成质量。然而，在多节点部署中，DiTango通过有效 绕过跨节点通信同时保持高质量生成结果，优于所有基线方法。这证明了 DiTango在分布式计算环境中的特定有效性，其中通信开销通常成为显著 瓶颈。

9

第 10 页

Yuyang Chen, Runxin Zhong, Zan Zong, Hengjie Li, Yuyang Jin, and Jidong Zhai

TP CP xDiT DiTango 设计，使得在给定时间约束内能够进行更有效的计算。 Wan-14B HunyuanVideo 这些架构优势在质量-性能权衡空间中带来了显著更优的帕累托前沿， 80 80 如图 13(b) 所示，DiTango 在平衡生成质量与计算效率方面始终优于 TFLOP/s 60 60 现有方法。 40 40 注意力 20 20 4 8 16 24 32 4 8 16 24 32 减少步数 VideoSys SGLD DiTango GPU 数量 GPU 数量 (a) 计算-质量效率 (b) 性能-质量权衡 22 1.0 图 12：不同 GPU 配置下核心注意力计算效率（以 FLOP/s 衡量）。 20 0.9 注意，VideoSys 和 SGLD 被省略，因为它们使用了与各自基础实现 18 0.8VideoSys and SGLD is omitted as they utilize the same under- (1-LPIPS) PSNR 16 0.7 相同的底层注意力机制（分别为 CP 和 xDiT）。 14 质量 0.6(CP and xDiT respectively). 12 0.5 40 60 80 10 15 20 25 30 35 40 计算比率 (%) 加速比 (vs 1 GPU) 延迟 视觉质量 方法 加速比 ↑ LPIPS ↓ SSIM ↑ PSNR ↑ VBench(%) ↑ 图 13：在 Wan- 14B 上使用 32-GPU 推理的质量分析：(a) 不同计算 Wan-1.3B 1.20× – – – 82.34 比率下的质量比较；(b) 不同复用配置下的质量-加速比权衡。 VideoSys 3.13× 0.288 0.676 14.08 81.37 SGLD 2.87× 0.220 0.728 15.77 81.92 DiTango 6.08× 0.214 0.734 16.34 82.21 Wan-14B 7.02× – – – 81.53 6.5 性能分解 VideoSys 11.25× 0.215 0.743 16.36 80.14 我们分析了 DiTango 中关键组件的有效性。 SGLD 16.51× 0.261 0.707 15.02 79.86 DiTango 31.15× 0.137 0.806 18.55 80.79 误差建模准确性。为了验证我们误差建模方法的有效性，我们在跨越不同时间步、层和提示的多样化样本上评估了预测准确性。 HunyuanVideo 6.26× – – – 78.22 如图 14 所示，我们的误差模型表现出显著的准确性，R² 值为 0.9792，皮尔逊相关系数为 0.9899，表明预测误差与实际误差之间存在强线性关系。这种高相关性验证了 DiTango 的两个关键设计原则：(1) 不同秩的注意力状态误差确实与局部状态误差高度相关，(2) 我们的锚点引导的选择机制有效地捕捉了时间窗口内的误差传播模式。较低的 RMSE (32.46) 和 MAE (21.88) 值进一步证实，我们的模型能够可靠地预测误差幅度，从而在平衡计算效率和生成质量的规划器选择中做出明智决策。

表 2：不同模型的性能和质量比较。加速比是将 32-GPU 并行推理与单 GPU 基线进行比较得出的。

6.4 生成质量与权衡 运行时效率。我们分析了 Di- Tango 的状态中心注意力机制在 32-GPU 分布式设置中不同状态组大小下的运行时效率。如图 15 所示，Di- 我们使用多项指标全面评估 DiTango 的生成质量：PSNR、SSIM 和 LPIPS 用于与原始模型输出进行保真度比较，VBench Score [10] 用于整体生成质量评估。如表 2 所示，DiTango 在大多数质量指标上始终取得优越的性能，并在所有评估模型上提供显著更高的加速比。值得注意的是，包括 TP、CP 和 xDiT 在内的基线方法通过无损并行化保持了原始模型的性能，因此作为具有最优 VBench 分数的质量参考，但没有可比较的 PSNR、SSIM 或 LPIPS 测量值。

我们深入分析了导致 DiTango 优越生成性能的因素。如图 13(a) 所示，DiTango 通过两个关键机制在等效计算预算下实现更高的生成质量：(1) 其贡献感知方法精确识别并保留最关键的计算，以及 (2) 其通信高效的设计使得在给定时间约束内能够进行更有效的计算。这些架构优势在质量-性能权衡空间中带来了显著更优的帕累托前沿，如图 13(b) 所示，DiTango 在平衡生成质量与计算效率方面始终优于现有方法。

这种卓越的计算效率归因于我们有效的计算-通信重叠策略和分层通信设计。在组大小为 4 和 8 时，DiTango 通过最小化跨节点通信展现出最佳性能。然而，在组大小为 16 时，尽管仍保持比基线更好的效率，但由于必要的组间状态通信，性能受到轻微影响。值得注意的是，DiTango 额外组件（规划器执行和内存操作）的开销保持最小，如分解图中较小的“其他”部分所示。

10

第 11 页

DiTango：通过选择性注意力状态复用实现低成本并行扩散生成

加速比 18.5 31.15 18.5 19 R² = 0.9819 30 PSNR 1200 Pearson = 0.9912 (×) 22.73 18 (dB) RMSE = 30.5322 20 17 16.4 MAE = 21.0671 14.02 13.67 1000 N = 2520 16 15.6 10 加速比 15PSNR 误差 800 0 14 完整结果复用 带随机 带选择规划器 带选择规划器 600 状态选择 + 状态中心运行时 预测 400 图 16：消融研究展示了不同 DiTango 组件对加速比和生成质量的影响 200 y = x (PSNR) 在 32 块 GPU 上进行 Wan-14B 推理。 0 拟合: y = 1.01x + 0.29 0 500 1000 实际误差 在我们的实验设置中，我们评估了四种具有 递增系统复杂性的配置。全注意力 结果复用的基线方法展示了粗粒度复用的局限性 图 14：DiTango 误差建模准确性。 策略，显示出适度的加速比和生成质量。 引入带有随机选择的细粒度状态复用显示出 有限的改进，因为其有效性受到非结构化 计算 通信 其他 通信模式和不可控的误差积累的阻碍。 DiTango_4 65.4% 31.1% 28.4 ms 引入我们的选择规划器标志着生成质量的显著提 DiTango_8 65.7% 30.2% 48.5 ms 高，同时通过更明智的状态选择也提升了加速比。 DiTango_16 48.6% 47.5% 128.6 ms 最后，同时具备选择规划器和状态中心运行时 编排的完整 DiTango 实现实现了最佳性能，同时保持高 SGLD 37.0% 63.0% 325.2 ms 生成质量，展示了这些组件如何协同工作以克服 CP 23.1% 76.9% 579.5 ms 可扩展性与质量之间的权衡。 0 25 50 75 100 注意力运行时分解 (%) 7 相关工作 图 15：运行时分解分析展示了计算、通信和其他操作 在不同方法和组大小下的比例，以及在 32-GPU 分布式 推理中。 扩散生成的并行化。传统的 3D 并行化策略 像 Megatron-LM [25] 在单样本扩散生成中效果有限。 由于批次大小本质上为 1，数据并行（Data Parallelism）和流水线并行（Pipeline Parallelism）都无法被高效利用。 虽然 AsyncDiff [4] 引入了基于时间步的流水线，但它根本 额外开销。尽管调度逻辑复杂，DiTango 依赖于多样本批处理，因此不适用于单样本场景。 引入了极小的系统开销。在计算方面，锚点 引导的选择规划器产生的开销可以忽略不计。在锚点 步骤期间，额外的统计分析是在归一化 压缩表示上执行的，而不是全精度张量，与标准全精度 计算相比，仅增加不到 1% 的计算量。在内存方面，维护注意力状态 缓存确实需要额外的容量。例如，以全精度缓存 必要的注意力状态对于像 HunyuanVideo 和 Wan2.1-14B 这样的大模型来说，大约消耗 30GB 的额外内存。然而，与大型语言模型不同，DiT 拥有 相对较小的参数规模。因此，这种额外的 内存开销完全在现代 GPU 架构的容量范围内。第 5 节中的 内存感知管理足够稳健，可以确保实际部署时不会受到 内存不足（out-of-memory）约束。 有损扩散加速。扩散生成本质上 计算密集但对噪声具有鲁棒性，因此有损加速 是实现显著加速且视觉退化可忽略不计的主流技术。模型级方法如蒸馏 [19, 22, 24] 可以将推理步骤减少一个数量级， 但它们需要大量的额外数据和训练。在算子

11

第 12 页

余阳·陈，润心·钟，赞·宗，恒杰·李，余阳·金，及继东·翟

量化 [35, 36] 和稀疏性 [30, 31] 内核提供了高效的硬件级加速；这些技术与 DiTango 正交，并可与其无缝集成。

基于缓存的方法 [2, 3, 15, 21, 34, 38] 利用了扩散步骤之间的时间冗余，通过重用相似的中间特征来跳过相应的计算。与蒸馏相比，它们无需训练且易于部署。然而，其粗粒度的重用策略和对启发式配置的依赖严重限制了其加速潜力。这种限制在长视频生成中尤为明显，因为特征动态变化迅速，需要更细粒度的适应性。

8 结论

本质上，DiTango 是一个并行缓存框架，在算法设计与系统编排之间取得了优雅的平衡，在不牺牲生成质量的情况下带来了显著的性能提升。它解决了分布式扩散推理中因带宽有限而导致的可扩展性瓶颈。通过选择性重用注意力状态，DiTango 在保留关键计算的同时，减轻了昂贵的分区访问开销。

其核心在于在线选择规划器与高效运行时引擎的结合，实现了对序列分区的细粒度计算重用控制。我们的研究结果表明，这种细粒度对于实现系统可扩展性和生成保真度至关重要，这一点在典型 GPU 集群上，与最先进方法相比，在质量损失最小的情况下实现了显著的速度提升得到了验证。

展望未来，DiTango 在大型生成模型服务中展现出巨大的潜力。它在异构集群中尤为有效，能够将生成延迟降低至秒级，为实时和交互式生成应用铺平了道路。

尽管取得了这些进展，我们承认 DiTango 目前存在两个局限性。首先，在高带宽环境（例如单节点设置）中，跨节点通信不再是主要瓶颈，此时专注于纯计算减少的方法可能会产生更高的吞吐量；因此，在实践中，结合这两种范式的混合部署策略是最优的。其次，尽管 DiTango 包含了内存感知管理，但注意力状态缓存仍消耗大量容量，这可能会限制生产环境中的 GPU 利用率。未来的工作将致力于通过缓存压缩技术或智能 CPU 卸载机制来解决这一内存占用问题，以进一步释放并行扩散系统的潜力。

致谢

我们要感谢匿名审稿人的深刻见解。本研究部分得到了教育部基础与交叉学科突破计划（JYB2025XDXM910）、国家杰出青年科学基金（Grant 62225206）、国家自然科学基金（Grant 62532006, U23A6007）以及北京市自然科学基金（Grant L242017, L243001）的支持。翟继东为本文通讯作者。

参考文献

[1] Tim Brooks, Bill Peebles, Connor Holmes, Amos Storkey, Alexei A. Efros, Andreas Terzis, Abhinav Gupta, Devi Parikh, Douwe Kiela, Gabriel Synnaeve, Hannaneh Hajishirzi, Ilya Sutskever, James Zung, Joelle Pineau, Luke Metz, Mira Murati, Pranav Shyam, Rohun Kulkarni, Ruth Fong, Vedant Misra, Yufei Guo, Adrià Recasens, Alexander Papiez, Alexandre Lebrun, Arthur Mensch, Avel Guénin, Bowen Baker, Brandon Houghton, Brian Tanner, Briane Paul V. Samson, Chang Chen, Christopher Clark, Cory McLean, David Martinez-Rubio, David Schnurr, Desu Narendra, Eli Moser, Ethan Perez, Igor Mordatch, J.R. Harris, Joey Faulkner, John Baumgartner, Karl Cobbe, Liam Fedus, Madeleine Thompson, Mark Chen, Mayur Mudigonda, Mimee Xu, Noemi Dreymann, Teddy Lee, Theresa Yoon, Timothy Lillicrap, Torrey Fellbaum, Trevor Darrell, Yuntao Bai, and Yutian Chen. 2024. Video generation models as world simulators. Technical Report. OpenAI. https://openai.com/research/video-generation-models-as-world-simulators

[2] Pengtao Chen, Mingzhu Shen, Peng Ye, Jianjian Cao, Chongjun Tu, Christos-Savvas Bouganis, Yiren Zhao, and Tao Chen. 2024. Delta-DiT: A Training-Free Acceleration Method Tailored for Diffusion Transformers. arXiv:2406.01125 [cs.CV] https://arxiv.org/abs/2406.01125

[3] Yuyang Chen, Linqian Zeng, Yijin ZHou, Hengjie Li, and Jidong Zhai. 2026. Jano: Adaptive Diffusion Generation with Early-stage Convergence Awareness. arXiv:2603.00519 [cs.CV] https://arxiv.org/abs/2603.00519

[4] Zigeng Chen, Xinyin Ma, Gongfan Fang, Zhenxiong Tan, and Xinchao Wang. 2024. AsyncDiff: Parallelizing Diffusion Models by Asynchronous Denoising. In Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, Vancouver, BC, Canada, December 10 – 15, 2024, Amir Globersons, Lester Mackey, Danielle Belgrave, Angela Fan, Ulrich Paquet, Jakub M. Tomczak, and Cheng Zhang (Eds.). http://papers.nips.cc/paper_files/paper/2024/hash/ad15848baa3932c0d2deabf0e11d1dcd-Abstract-Conference.html

[5] Tri Dao. 2024. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net. https://openreview.net/forum?id=mZn2Xyh9Ec

[6] vipshop.com DefTruth. 2025. cache-dit: A PyTorch-native and Flexible Inference Engine with Hybrid Cache Acceleration and Parallelism for DiTs. https://github.com/vipshop/cache-dit.git Open-source software available at https://github.com/vipshop/cache-dit.git.

[7] Jiarui Fang, Jinzhe Pan, Xibo Sun, Aoyu Li, and Jiannan Wang. 2024. xDiT: an Inference Engine for Diffusion Transformers (DiTs) with Massive Parallelism. CoRR abs/2411.01738 (2024). doi:10.48550/ARXIV.2411.01738 arXiv:2411.01738

[8] Jiarui Fang and Shangchun Zhao. 2024. A Unified Sequence Parallelism Approach for Long Context Generative AI. arXiv preprint arXiv:2405.07719 (2024).

[9] Jonathan Ho and Tim Salimans. 2022. Classifier-Free Diffusion Guidance. CoRR abs/2207.12598 (2022). doi:10.48550/ARXIV.2207.12598 arXiv:2207.12598

[10] Ziqi Huang, Yinan He, Jiashuo Yu, Fan Zhang, Chenyang Si, Yuming Jiang, Yuanhan Zhang, Tianxing Wu, Qingyang Jin, Nattapol Chanpaisit, Yaohui Wang, Xinyuan Chen, Limin Wang, Dahua Lin, Yu Qiao, and Ziwei Liu. 2024. VBench: Comprehensive Benchmark Suite for Video Generative Models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition.

[11] Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He. 2023. DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models. CoRR abs/2309.14509 (2023). doi:10.48550/ARXIV.2309.14509 arXiv:2309.14509

[12] Weijie Kong, Qi Tian, Zijian Zhang, Rox Min, Zuozhuo Dai, Jin Zhou, Jiangfeng Xiong, Xin Li, Bo Wu, Jianwei Zhang, Kathrina Wu, Qin Lin, Junkun Yuan, Yanxin Long, Aladdin Wang, Andong Wang, Changlin Li, Duojun Huang, Fang Yang, Hao Tan, Hongmei Wang, Jacob Song, Jiawang Bai, Jianbing Wu, Jinbao Xue, Joey Wang, Kai Wang, Mengyang Liu, Pengyu Li, Shuai Li, Weiyan Wang, Wenqing Yu, Xinchi Deng, Yang Li, Yi Chen, Yutao Cui, Yuanbo Peng, Zhentao Yu, Zhiyu He, Zhiyong Xu, Zixiang Zhou, Zunnan Xu, Yangyu Tao, Qinglin Lu, Songtao Liu, Daquan Zhou, Hongfa Wang, Yong Yang, Di Wang, Yuhong Liu, Jie Jiang, and Caesar Zhong. 2024. HunyuanVideo: A Systematic Framework For Large Video Generative Models. CoRR abs/2412.03603 (2024). doi:10.48550/ARXIV.2412.03603 arXiv:2412.03603

[13] Kuaishou Technology. 2025. Kling – Community for short video & livestream. https://kling.kuaishou.com/en Accessed: 2025-03-07.

[14] Muyang Li, Tianle Cai, Jiaxin Cao, Qinsheng Zhang, Han Cai, Junjie Bai, Yangqing Jia, Kai Li, and Song Han. 2024. DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2024, Seattle, WA, USA, June 16-22, 2024. IEEE, 7183–7193. doi:10.1109/CVPR52733.2024.00686

[15] Feng Liu, Shiwei Zhang, Xiaofeng Wang, Yujie Wei, Haonan Qiu, Yuzhong Zhao, Yingya Zhang, Qixiang Ye, and Fang Wan. 2024. Timestep Embedding Tells: It’s Time to Cache for Video Diffusion Model. CoRR abs/2411.19108 (2024). doi:10.48550/ARXIV.2411.19108 arXiv:2411.19108

12

第 13 页

DiTango：通过选择性注意力状态复用实现低成本并行扩散生成

[16] Hao Liu 和 Pieter Abbeel。2023。用于大上下文模型的块级并行 Transformer。在 Advances in Neural Information Processing Systems 36: 第 36 届神经信息处理系统年度会议论文集 (NeurIPS 2023) 中，2023 年 12 月 10-16 日，美国新奥尔良，Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, 和 Sergey Levine (编)。http://papers.nips.cc/paper_files/paper/2023/hash/1bfd87d2d92f0556819467dc08034f76-Abstract-Conference.html [17] Hao Liu, Matei Zaharia, 和 Pieter Abbeel。2023。具有块级 Transformer 的环形注意力以实现近无限上下文。CoRR abs/2310.01889 (2023)。doi:10.48550/ARXIV.2310.01889 arXiv:2310.01889 [18] Jiacheng Liu, Chang Zou, Yuanhuiyi Lyu, Junjie Chen, 和 Linfeng Zhang。2025。从复用到预测：使用 TaylorSeers 加速扩散模型。arXiv 预印本 arXiv:2503.06923 (2025)。 [19] Qiang Liu。2022。Rectified Flow：一种保持边缘分布的最优传输方法。arXiv:2209.14577 [stat.ML] https://arxiv.org/abs/2209.14577 [20] Guoqing Ma, Haoyang Huang, Kun Yan, Liangyu Chen, Nan Duan, Shengming Yin, Changyi Wan, Ranchen Ming, Xiaoniu Song, Xing Chen, Yu Zhou, Deshan Sun, Deyu Zhou, Jian Zhou, Kaijun Tan, Kang An, Mei Chen, Wei Ji, Qiling Wu, Wen Sun, Xin Han, Yanan Wei, Zheng Ge, Aojie Li, Bin Wang, Bizhu Huang, Bo Wang, Brian Li, Changxing Miao, Chen Xu, Chenfei Wu, Chenguang Yu, Dapeng Shi, Dingyuan Hu, Enle Liu, Gang Yu, Ge Yang, Guanzhe Huang, Gulin Yan, Haiyang Feng, Hao Nie, Haonan Jia, Hanpeng Hu, Hanqi Chen, Haolong Yan, Heng Wang, Hongcheng Guo, Huilin Xiong, Huixin Xiong, Jiahao Gong, Jianchang Wu, Jiaoren Wu, Jie Wu, Jie Yang, Jiashuai Liu, Jiashuo Li, Jingyang Zhang, Junjing Guo, Junzhe Lin, Kaixiang Li, Lei Liu, Lei Xia, Liang Zhao, Liguo Tan, Liwen Huang, Liying Shi, Ming Li, Mingliang Li, Muhua Cheng, Na Wang, Qiaohui Chen, Qinglin He, Qiuyan Liang, Quan Sun, Ran Sun, Rui Wang, Shaoliang Pang, Shiliang Yang, Sitong Liu, Siqi Liu, Shuli Gao, Tiancheng Cao, Tianyu Wang, Weipeng Ming, Wenqing He, Xu Zhao, Xuelin Zhang, Xianfang Zeng, Xiaojia Liu, Xuan Yang, Yaqi Dai, Yanbo Yu, Yang Li, Yineng Deng, Yingming Wang, Yilei Wang, Yuanwei Lu, Yu Chen, Yu Luo, Yuchu Luo, Yuhe Yin, Yuheng Feng, Yuxiang Yang, Zecheng Tang, Zekai Zhang, Zidong Yang, Binxing Jiao, Jiansheng Chen, Jing Li, Shuchang Zhou, Xiangyu Zhang, Xinhao Zhang, Yibo Zhu, Heung-Yeung Shum, 和 Daxin Jiang。2025。Step-Video-T2V 技术报告：视频基础模型的实践、挑战与未来。arXiv:2502.10248 [cs.CV] https://arxiv.org/abs/2502.10248 [21] Xinyin Ma, Gongfan Fang, Michael Bi Mi, 和 Xinchao Wang。2024。学习缓存：通过层缓存加速扩散 Transformer。arXiv:2406.01733 [cs.LG] [22] Chenlin Meng, Robin Rombach, Ruiqi Gao, Diederik P. Kingma, Stefano Ermon, Jonathan Ho, 和 Tim Salimans。2022。关于引导扩散模型的蒸馏。arXiv e-prints, 文章 arXiv:2210.03142 (2022 年 10 月), arXiv:2210.03142 页。doi:10.48550/arXiv.2210.03142 arXiv:2210.03142 [cs.CV] [23] William Peebles 和 Saining Xie。2023。可扩展的 Transformer 扩散模型。在 IEEE/CVF 国际计算机视觉会议 (ICCV 2023) 中，2023 年 10 月 1-6 日，法国巴黎。IEEE, 4172–4182。doi:10.1109/ICCV51070.2023.00387 [24] Axel Sauer, Dominik Lorenz, Andreas Blattmann, 和 Robin Rombach。2023。对抗性扩散蒸馏。arXiv:2311.17042 [cs.CV] https://arxiv.org/abs/2311.17042 [25] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, 和 Bryan Catanzaro。2019。Megatron-LM：使用模型并行性训练数十亿参数语言模型。CoRR abs/1909.08053 (2019)。arXiv:1909.08053 http://arxiv.org/abs/1909.08053 [26] Genmo Team。2024。Mochi 1。https://github.com/genmoai/models。 [27] VideoSys Team。2024。VideoSys：一个简单高效的视频生成系统。https://github.com/NUS-HPC-AI-Lab/VideoSys [28] Team Wan, Ang Wang, Baole Ai, Bin Wen, Chaojie Mao, Chen-Wei Xie, Di Chen, Feiwu Yu, Haiming Zhao, Jianxiao Yang, Jianyuan Zeng, Jiayu Wang, Jingfeng Zhang, Jingren Zhou, Jinkai Wang, Jixuan Chen, Kai Zhu, Kang Zhao, Keyu Yan, Lianghua Huang, Mengyang Feng, Ningyi Zhang, Pandeng Li, Pingyu Wu, Ruihang Chu, Ruili Feng, Shiwei Zhang, Siyang Sun, Tao Fang, Tianxing Wang, Tianyi Gui, Tingyu Weng, Tong Shen, Wei Lin, Wei Wang, Wei Wang, Wenmeng Zhou, Wente Wang, Wenting Shen, Wenyuan Yu, Xianzhong Shi, Xiaoming Huang, Xin Xu, Yan Kou, Yangyu Lv, Yifei Li, Yijing Liu, Yiming Wang, Yingya Zhang, Yitong Huang, Yong Li, You Wu, Yu Liu, Yulin Pan, Yun Zheng, Yuntao Hong, Yupeng Shi, Yutong Feng, Zeyinzi Jiang, Zhen Han, Zhi-Fan Wu, 和 Ziyu Liu。2025。Wan：开放且先进的大规模视频生成模型。arXiv 预印本 arXiv:2503.20314 (2025)。 [29] Jiannan Wang, Jiarui Fang, Aoyu Li, 和 PengCheng Yang。2024。PipeFusion：用于扩散 Transformer 模型推理的位移块流水线并行。CoRR abs/2405.14430 (2024)。doi:10.48550/ARXIV.2405.14430 arXiv:2405.14430 [30] Haocheng Xi, Shuo Yang, Yilong Zhao, Chenfeng Xu, Muyang Li, Xiuyu Li, Yujun Lin, Han Cai, Jintao Zhang, Dacheng Li, Jianfei Chen, Ion Stoica, Kurt Keutzer, 和 Song Han。2025。Sparse VideoGen：利用时空稀疏性加速视频扩散 Transformer。CoRR abs/2502.01776 (2025)。doi:10.48550/ARXIV.2502.01776 arXiv:2502.01776 [31] Shuo Yang, Haocheng Xi, Yilong Zhao, Muyang Li, Jintao Zhang, Han Cai, Yujun Lin, Xiuyu Li, Chenfeng Xu, Kelly Peng, 等。2025。Sparse VideoGen2：加速

13