Search papers, labs, and topics across Lattice.
This paper introduces CoCoDiff, a distributed inference engine that optimizes collective communications in Diffusion Transformers (DiTs) under Ulysses sequence parallelism. CoCoDiff leverages tile-aware parallel all-to-all communication, V-first scheduling, and V-major selective communication to overlap communication with computation and exploit temporal redundancy between denoising steps. Experiments on the Aurora supercomputer demonstrate an average speedup of 3.6x and a peak speedup of 8.4x across various DiT models and node configurations.
Overlapping communication with computation in distributed Diffusion Transformer inference can yield up to 8.4x speedups, challenging the assumption that collectives are a fixed bottleneck.
Diffusion Transformers (DiTs) are increasingly adopted in scientific computing, yet growing model sizes and resolutions make distributed multi-GPU inference essential. Ulysses sequence parallelism scales DiT inference but introduces frequent all-to-all collectives that dominate latency. Overlapping these with computation is difficult due to tight data dependencies, large message volumes, and asymmetric interconnect bandwidths. We introduce CoCoDiff, a distributed DiT inference engine exploiting two observations: (1) V requires only linear projection while Q/K need additional normalization and RoPE, creating opportunities to overlap V's communication with Q/K computation; (2) adjacent denoising steps produce similar tensors, yielding temporal redundancy. CoCoDiff introduces three mechanisms: Tile-Aware Parallel All-to-all (TAPA) decomposes collectives into topology-aligned phases; V-First scheduling hides V's communication behind Q/K computation; and V-Major selective communication transmits only active projections on slow interconnects. On the Aurora supercomputer with four DiT models across 1-8 nodes (up to 96 Intel GPU tiles), CoCoDiff achieves an average speedup of 3.6x, peaking at 8.4x.
