Partially Conditioned Patch Parallelism for Accelerated Diffusion Model Inference

TL;DR

PCPP tackles the latency bottleneck of high-resolution diffusion model inference by partially conditioning patch-level denoising on local and neighboring context, enabled by asynchronous point-to-point communication. By replacing AllGather with neighbor-to-neighbor exchanges and restricting attention to partial neighboring keys/values, it achieves substantial communication reductions (≈$70\%$) and faster inference (≈$2.36\sim 8.02\times$) on $4$–$8$ GPUs, with a trade-off in image quality. The approach demonstrates the feasibility of rapid, multi-GPU diffusion inference for interactive, high-resolution generation, offering meaningful energy and latency benefits while highlighting the need for smarter context selection to preserve quality. The work suggests promising directions for scaling diffusion acceleration to practical real-time applications and invites further refinement via adaptive context, integration with control mechanisms, and broader hardware exploration.

Abstract

Diffusion models have exhibited exciting capabilities in generating images and are also very promising for video creation. However, the inference speed of diffusion models is limited by the slow sampling process, restricting its use cases. The sequential denoising steps required for generating a single sample could take tens or hundreds of iterations and thus have become a significant bottleneck. This limitation is more salient for applications that are interactive in nature or require small latency. To address this challenge, we propose Partially Conditioned Patch Parallelism (PCPP) to accelerate the inference of high-resolution diffusion models. Using the fact that the difference between the images in adjacent diffusion steps is nearly zero, Patch Parallelism (PP) leverages multiple GPUs communicating asynchronously to compute patches of an image in multiple computing devices based on the entire image (all patches) in the previous diffusion step. PCPP develops PP to reduce computation in inference by conditioning only on parts of the neighboring patches in each diffusion step, which also decreases communication among computing devices. As a result, PCPP decreases the communication cost by around $70\%$ compared to DistriFusion (the state of the art implementation of PP) and achieves $2.36\sim 8.02\times$ inference speed-up using $4\sim 8$ GPUs compared to $2.32\sim 6.71\times$ achieved by DistriFusion depending on the computing device configuration and resolution of generation at the cost of a possible decrease in image quality. PCPP demonstrates the potential to strike a favorable trade-off, enabling high-quality image generation with substantially reduced latency.

Figures (9)

• Figure 1: Example of images generated using Partially Conditioned Patch Parallelism (PCPP) with varying partial value. The images are generated with 4 GPUs using PCPP. Due to the use of classifier-free guidance, 2 GPUs are used in the denoising process to generate pixels directly. The image is separated into upper and lower patches computed by dedicated GPUs parallelly. To provide a clear comparison, warm-up synchronization steps are not applied, i.e., the computation of patches is separated at the very first step. The blue/green region refers to the activation used to compute the upper/lower region. The activation can consist of fresh activation of the local patch and partial stale activation from the last step of the neighbor patch, depending on the partial value used.
• Figure 2: Partially Conditioned Patch Parallelism (PCPP) asynchronous communication design. This diagram captures the computation and recv_buffer state at $t = 9$ and layer $3$ for device rank $2$ prior to executing the self-attention computations. The device is confirming the completion of async recv of the previous layer's activation from its adjacent rank; once this is completed, $A_{2,10}$ in the recv_buffer from rank $1$ will be overwritten by $A_{2,9}$ for subsequent use at time step $t = 8$. In the meantime, it initiates an (async send) of its own activation $A_{3,9}$ to its neighboring rank, concurrent with the self-attention computation. As described in Section \ref{['sec:PCA']}, the input to the self-attention computation is $A_{2,9}^{(2)}$ and the two neighboring stale activations from last step $A_{3,10}^{(1)}$ and $A_{3,10}^{(3)}$.
• Figure 3: Overview of partially conditioned attention. The local patch receives partial $Q_{NP}$s and $K_{NP}$s from the neighboring patches. The size of partial keys and values is determined by the parameter $partial$, whose effect is shown in Figure \ref{['fig:2-device']}. $K$ and $V$ are then assembled using $Q_{LP}$, $Q_{NP}$s and $K_{LP}$, $K_{NP}$s respectively, instead of using all the patches from the entire image. LP means local patch, NP means neighboring patch. Black arrows refer to data transfer on the local device, and blue arrows refer to communication between patches on different devices.
• Figure 4: Inference latency for generating one image with 50-step DDIM sampler Song2020DenoisingDI by resolution and device configuration. Our method sets the partial value for partially conditioned attention to 0.3 for 4 devices and 0.8 for 8 devices. The red numbers show the speed-up achieved compared to inference on 1 device.
• Figure 5: Visual comparison of sample images generated by three prompts that work relatively well with PCPP. More examples are available in Appendix \ref{['appx:generation']}.