AoiZora: Topology-Aware Auto-Parallel Optimization for Inference of Diffusion Transformers

2026-06-16

AI Summary

Existing automated parallelization systems for video diffusion model inference often disregard the physical topology of TPUs, resulting in suboptimal communication efficiency and high latency. This work proposes a compiler-mediated, topology-aware parallelization optimization method that, for the first time, integrates the physical interconnect topology of TPUs into the automated parallelism search process. By leveraging pre-compilation IR filtering and HLO-level communication modeling, the approach jointly optimizes logical sharding strategies and physical device placement. Evaluated on TPU v5e sub-slices, the method enables highly efficient deployment, reducing the per-step denoising latency of Wan 2.1 by up to 1.42× compared to state-of-the-art baselines.

Abstract

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.

Problem

• auto-parallel
• topology-aware
• diffusion transformers
• TPU sub-slices
• inference optimization

Innovation

• topology-aware
• auto-parallel
• diffusion transformers
• TPU sub-slices
• compiler-mediated optimization