X-Stage: An Overlooked Pipeline Stage for Communication-Computation Overlap in DiT Inference

2026-07-25

AI Summary

This work addresses the unpredictable backpressure in existing distributed DiT inference systems, which stems from neglecting the intermediate phase between communication initiation and remote visibility completion, thereby limiting effective overlap of communication and computation. To resolve this, we explicitly model this phase as a schedulable X-Stage and introduce a lightweight Burst-Gap model to characterize bursty remote memory behavior. Leveraging this insight, we redesign fused kernels to mitigate backpressure and enhance overlap efficiency. Key techniques include fine-grained device-side communication initiation, persistent GPU kernels, interleaved expert wave execution, and tile-level fusion of FlashAttention with All-to-All. Experiments across 84 configurations show that our DeepGEMM MegaMoE fused kernel achieves a 1.18× geometric mean and up to 1.62× speedup, while Ulysses attention with FlashAttention-3/4 attains peak speedups of 1.43× and 1.42×, respectively, with long-sequence steady-state performance approaching the compute-only baseline.

Abstract

Fine-grained, device-initiated communication lets persistent GPU kernels in distributed diffusion transformer (DiT) inference issue remote stores and overlap data movement with Tensor Core computation. Existing systems schedule when communication is issued and when received data becomes consumable, but omit post-issue progress before remote-visible completion, making sender backpressure hard to predict. We identify X-Stage, a software-visible post-issue pipeline stage. Measurements on an eight-GPU node with a recent NVIDIA architecture show that short remote-store bursts drain as the issuer resumes work, whereas sustained injection exhausts finite outstanding capacity and delays later issues. A lightweight Burst-Gap model parameterized by backpressure-free issue time, effective drain rate, and outstanding capacity predicts issue overhead, recovery between bursts, and the onset of backpressure. Guided by the model, we redesign two communication-computation fused kernels. For DeepGEMM MegaMoE, interleaving Linear-1 and Linear-2 work across expert waves places computation between concentrated remote-store bursts, yielding a 1.18x geometric-mean and 1.62x maximum kernel speedup over the Expert-Wave baseline across 84 configurations. For Ulysses sequence-parallel attention, tile-granular fusion of the post-attention All-to-All with FlashAttention lets an output-tile owner issue remote stores and resume computation without a dedicated communication warp or streaming multiprocessor. FlashAttention-3 and FlashAttention-4 reach maximum sender-visible speedups of 1.43x and 1.42x over serial execution, and at long sequences their steady-state times approach those of FlashAttention alone. These results establish post-issue progress as a measurable scheduling lever for shaping bursts, avoiding backpressure, and hiding sender-side overhead.

Problem

• communication-computation overlap
• backpressure
• pipeline stage
• distributed inference
• DiT

Innovation

• X-Stage
• communication-computation overlap
• persistent GPU kernels
• backpressure modeling
• DiT inference