AoiZora优化扩散模型推理,让视频生成更快更稳。
AoiZora: Topology-Aware Auto-Parallel Optimization for Inference of Diffusion Transformers

- 通过编译器感知拓扑,智能分配计算任务到TPU芯片
- 在TPU v5e上将Wan 2.1单步去噪延迟降低1.42倍
- 无需改动模型代码,适合追求低延迟视频生成的开发者
视频扩散模型已成为重要的生成服务负载,但每段视频生成需在大时空潜在空间上进行多次去噪迭代,单设备难以实现低延迟推理。因此,去噪步骤通常分布于多个加速器上,TPU子切片成为理想的部署方案。然而,现有自动并行系统仅在逻辑设备网格中搜索,忽略物理互联拓扑对性能的影响,导致大量性能损失。为此,我们提出AoiZora,一种面向TPU子切片的低延迟视频扩散推理拓扑感知自动并行优化框架。其核心思想是通过编译流程中的不同阶段,将逻辑分片与物理布局重新关联:先在低成本预编译IR中剔除弱候选方案,再对剩余方案使用编译后的HLO与拓扑感知通信模型进行物理布局排序,最终在标准编译路径中实现最优计划。该方法不改变模型代码、编译下放、集体内核和网络路由。在TPU v5e子切片上,AoiZora相较现有方案将Wan 2.1单步去噪延迟最高降低1.42倍。
原文摘要 · Abstract (English)
Video diffusion has quickly grown into a key generative serving workload, yet producing each clip demands many denoising iterations over large spatio-temporal latents, which puts low-latency inference out of reach on a single device. A denoising step is therefore typically distributed across multiple accelerators, and TPU sub-slices have become an attractive and practical fabric for doing so. Current auto-parallel systems, however, search almost exclusively over logical device meshes and disregard how a chosen sharding is actually laid out on the physical TPU interconnect -- an oversight that leaves large, topology-dependent performance on the table. We address this gap with AoiZora, a compiler-mediated topology planner built for low-latency video diffusion inference on TPU sub-slices. Its guiding principle is to reconnect logical sharding with physical placement by drawing on different points in the compilation flow: AoiZora first eliminates weak sharding candidates from inexpensive pre-compilation IRs, then compiles only the ones that survive and orders their physical placements using compiled HLO together with a topology-aware communication model. The winning plan is realized along the ordinary compiler path, leaving model code, compiler lowering, collective kernels, and network routing entirely intact. On TPU v5e sub-slices, AoiZora reduces Wan 2.1 one-step denoising latency by as much as 1.42x relative to existing solutions.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。