🤖 machine learning

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

AoiZora is a compiler-mediated topology planner that optimizes low-latency video diffusion inference on TPU sub-slices by integrating physical interconnect awareness into the auto-parallel search process, achieving up to a 1.42x latency reduction for Wan 2.1 compared to existing solutions.

Original authors: Kaijian Wang, Yuanyuan Xu, Fanjiang Ye, Ye Cao, Jingwei Zuo, T. S. Eugene Ng, Yarong Mu, Yuke Wang

Published 2026-06-17

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Kaijian Wang, Yuanyuan Xu, Fanjiang Ye, Ye Cao, Jingwei Zuo, T. S. Eugene Ng, Yarong Mu, Yuke Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to organize a massive, high-stakes relay race to create a video. The "runners" are computer chips (specifically Google's TPU chips), and the "baton" is the data needed to generate each frame of the video.

The video generation process is like a long, repetitive loop where the runners have to clean up a noisy image over and over again (called "denoising") until it looks clear. Because the video is huge and the process is slow, you can't do it on just one chip; you need a whole team of chips working together.

The Problem: The "Seating Chart" Mistake
Currently, when engineers set up this team, they focus on the logical plan: "Runner A will handle the background, Runner B will handle the sky, and Runner C will handle the people." This is like assigning tasks on a whiteboard.

However, they often ignore the physical reality: Where exactly are these runners sitting in the stadium?

In a real stadium (the physical TPU network), some runners are sitting right next to each other and can pass a baton instantly. Others are on opposite sides of the field, and passing a baton takes time and causes traffic jams on the shared paths between them.

The paper found that even if you give the runners the exact same tasks (the same logical plan), simply changing where they sit in the stadium can make the race 16.6% slower or faster. The main culprit isn't the distance they run, but the fact that some runners are forced to use the same narrow hallway (a shared network link) at the same time, causing a traffic jam.

The Solution: AoiZora (The Smart Coach)
The authors built a new tool called AoiZora (named after the Japanese words for "Blue Sky"). Think of AoiZora as a super-smart coach who doesn't just write the play; it also designs the seating chart based on the actual stadium layout.

AoiZora solves a tricky problem: Checking every possible seating arrangement is too slow and expensive. If you tried to test every single way to seat the runners, you'd spend more time planning than actually running the race.

How AoiZora Works (The Two-Stage Strategy)
AoiZora uses a clever "filter and refine" approach, like a hiring process:

Stage 1: The Quick Resume Screen (Pre-Compilation)
First, AoiZora looks at the "resumes" (the logical task plans) using a cheap, fast method. It doesn't build the whole stadium yet; it just checks if the tasks make sense on paper. It quickly throws away the bad plans that would clearly fail (like trying to split a task that can't be split). This is fast and cheap, like a 30-second interview.
Stage 2: The Full Simulation (Post-Compilation)
For the few top candidates that survived the first round, AoiZora does the hard work. It builds a detailed simulation of the actual stadium (the physical TPU chips) and runs a full test. It sees exactly where the traffic jams will happen on the specific network links. It then picks the single best seating arrangement that avoids those jams.

The Result
By using this two-step method, AoiZora finds the perfect combination of "what to do" and "where to do it."

The Analogy: Imagine you are organizing a dinner party.
- Old Way: You decide who brings what dish (logical plan) and tell them to sit anywhere. Sometimes, the person bringing the heavy soup has to walk past the person carrying the fragile cake, and they bump into each other, spilling everything.
- AoiZora Way: You first quickly rule out guests who can't bring a dish. Then, for the remaining guests, you look at the actual table layout and seat the soup-bearer right next to the kitchen door and the cake-bearer on the other side, ensuring no one bumps into anyone else.

The Bottom Line
On Google's TPU chips, AoiZora made video generation 1.42 times faster than previous methods. It didn't change the video game itself or the rules of the race; it just figured out the smartest way to seat the runners so they could pass the baton without tripping over each other. This allows for faster, lower-latency video generation without needing to rewrite the underlying software.

Technical Summary: AoiZora

Problem Statement
Video diffusion inference, particularly for models like Wan 2.1, has become a critical generative workload. However, the inference process is computationally intensive, requiring hundreds of thousands of spatio-temporal tokens and repeated denoising steps over large activation tensors. This necessitates distributed inference across multiple accelerators. While current auto-parallel systems (e.g., those based on JAX/XLA) effectively search for optimal logical device meshes (determining how tensors are sharded), they largely ignore the physical placement of these logical ranks on the hardware.

On TPU sub-slices (contiguous rectangular subsets of a TPU pod), the inter-chip interconnect (ICI) is a nearest-neighbor mesh without wrap-around links in all dimensions. Consequently, the same logical sharding strategy can result in significantly different end-to-end latencies depending on how the logical ranks are mapped to physical chips. The paper observes that this "placement effect" alone can cause up to a 16.6% variation in one-step denoising latency on a v5e-8 slice, primarily due to shared-link contention rather than hop count. Existing systems bind logical ranks to physical devices early in the workflow (often via default ordering) before the compiler exposes the concrete collective communication graph, leaving topology-dependent performance unoptimized.

Methodology: A Two-Stage Compiler-Mediated Planner
AoiZora addresses this gap by placing topology-aware planning at the boundary between compiler analysis and execution. It leverages the "cost/fidelity asymmetry" of the JAX/XLA compilation pipeline: early Intermediate Representations (IRs) are cheap but lack concrete communication details, while fully compiled High-Level Optimizations (HLO) are expensive but expose the exact collective graph and device groups.

The planner operates in two distinct stages:

Stage 1: Pre-Compilation Logical Sharding Pruning
- Goal: Efficiently eliminate inefficient logical sharding candidates.
- Mechanism: AoiZora enumerates legal workload-axis assignments (e.g., tensor, sequence, data parallelism) and lowers them to StableHLO/Shardy IRs.
- Analysis: It constructs a prediction DAG that includes explicit communications and synthetic implicit communication nodes (e.g., for resharding, sharded reductions, and opaque-kernel boundaries) to estimate logical communication volume.
- Scoring: It uses a placement-oblivious throughput bound ( $Q_1$ ) to score candidates. This stage filters out shardings with poor memory footprints or excessive logical communication, retaining only a small "Top-K" set of survivors.
Stage 2: Post-Compilation Topology Ranking
- Goal: Select the optimal physical placement for the surviving logical shardings.
- Mechanism: The Top-K logical candidates are fully compiled by XLA to generate concrete HLO.
- Analysis: The system enumerates distinct physical axis orders (mapping logical axes to TPU sub-slice coordinates). It parses the compiled HLO to extract the exact collective graph, including replica groups, payload sizes, and dependencies.
- Scoring: A topology-aware cost model ( $Q_2$ ) ranks each sharding-placement pair. This model reconstructs physical routes for collectives, estimates link loads, and penalizes placements that cause contention on shared ICI links.
- Selection: The pair with the lowest estimated cost is selected.

Key Contributions

Characterization of Placement Sensitivity: The authors demonstrate that on TPU sub-slices, physical placement is a first-class optimization variable. They show that identical logical shardings can yield up to 16.6% latency differences due to shared-link contention.
Compiler-Fidelity-Aware Search: The paper identifies and exploits the asymmetry between cheap pre-compilation IRs (suitable for high-recall pruning) and expensive compiled HLO (necessary for accurate topology ranking).
AoiZora System Design: A compiler-mediated planner that separates logical pruning from physical ranking. Crucially, it integrates with the standard JAX/XLA stack without modifying model code, lowering rules, collective kernel implementations, or network routing. It outputs a sharding strategy and physical axis order that the standard compiler then executes.
Evaluation: Extensive experiments on Wan 2.1 across TPU v5e sub-slices (4, 8, and 16 chips) and various video resolutions (480p/720p) and frame counts.

Experimental Results
Evaluated on Wan 2.1, AoiZora consistently outperforms two baselines:

Existing Deployment: A hand-tuned, expert-configured parallelism strategy (using FSDP, sequence, and tensor parallelism) with fixed physical placement.
Logical-Search Baseline: An Alpa-inspired system that optimizes logical sharding using an abstract cost model but does not perform post-compilation physical placement ranking.

Performance Gains: AoiZora achieves speedups of up to 1.42× over the existing solution and 1.37× over the logical-search baseline. Geometric mean speedups are 1.09× and 1.11×, respectively.
Fidelity: The pre-compilation stage successfully retains the measured best logical sharding in all tested settings (100% top-1 recall), while the post-compilation stage achieves high rank correlation ( $\rho_s$ up to 0.985) between predicted and measured latency.
Overhead: The offline planning cost is practical, ranging from ~47 seconds (v5e-4) to ~376 seconds (v5e-16), which is amortized over repeated inference requests of the same shape.

Significance and Claims
The paper claims that AoiZora bridges the gap between logical parallelism and physical TPU topology without requiring changes to the underlying serving stack. By deferring physical placement decisions until the compiler has materialized the concrete collective graph, AoiZora avoids the "oversight" of current auto-parallel systems that treat logical meshes as complete abstractions.

The authors position AoiZora as a practical solution for the specific constraints of TPU sub-slices, where the loss of wrap-around links makes topology-aware placement critical. They assert that their approach allows model developers to remain agnostic to topology details while achieving near-optimal inference latency for video diffusion workloads. The work suggests that for dense video DiT architectures with static tensor shapes, a one-time offline planning step can significantly improve serving performance by resolving the contention issues inherent in fixed physical placements.

Technical Summary: AoiZora

More like this