Video diffusion transformers (DiTs) process tens of thousands of spatio-temporal tokens per denoising step, so low-latency inference and training commonly spread them across several accelerators. Auto-parallelization systems price a sharding on its logical graph, but its cost depends on the program the compiler generates and, on direct networks such as TPU sub-slices, on which chip runs each logical rank; both are visible only after compilation, while the chip assignment must be fixed before it. We present AoiZora, a compiler-agnostic auto-parallelization framework built on two IRs, one for the program a partitioner shards and one for the program a compiler emits; it chooses the logical mesh, the layout of every tensor and operation, and the device assignment of an unsharded model by pricing its compiled programs, without executing any candidate. At the core of AoiZora is evidence, a compiler-neutral record of a compiled program that leaves its placement open, so that one compilation can be priced under many placements by a calibrated model of the accelerator network and the placement can determine which sharding is selected. On TPU v5e, AoiZora reduces the denoising latency of Wan 2.1 and HunyuanVideo by 1.23x and 1.32x on average over expert configurations and Alpa, finding each plan within 10 minutes; it also speeds up training on TPU v5p by 1.11x on 4 chips and 1.21x on 8 chips on average and parallelizes PyTorch programs on GPUs with the same core.

翻译：暂无翻译