arXiv:2607.23264cs.DCcs.AI2026-07

发现并利用推理中的通信延迟空隙,提升扩散模型的计算效率。

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

论文配图:X-Stage: An Overlooked Pipeline Stage for Communication-Computation Overlap in DiT Inference
图 1 · 摘自论文原文
  • 识别出通信后阶段(X-Stage)的隐藏延迟,可影响性能
  • 通过轻量模型预测通信压力,实现1.62倍的加速
  • 适合优化大规模分布式推理系统的工程师

细粒度、设备发起的通信使分布式扩散变换器(DiT)推理中的持久化GPU内核能发出远程存储,并将数据移动与张量核心计算重叠。现有系统仅调度通信发出和接收数据可用的时间,忽略了发出后的进度,导致发送方背压难以预测。我们识别出一个软件可见的通信后阶段——X-Stage。在配备最新NVIDIA架构的八卡节点上测量显示,短时远程存储突发会随发出者恢复工作而迅速耗尽,而持续注入则会耗尽有限的未完成容量,延迟后续发出。一个基于无背压发出时间、有效耗尽率和未完成容量参数化的轻量级Burst-Gap模型,可预测发出开销、突发间恢复时间及背压出现时机。基于该模型,我们重构了两个通信-计算融合内核。对DeepGEMM MegaMoE,跨专家波交错执行Linear-1和Linear-2操作,在集中远程存储突发之间插入计算,84种配置下几何平均提速1.18倍,最大提速1.62倍。对Ulysses序列并行注意力,将注意力后All-to-All与FlashAttention以块粒度融合,让输出块拥有者发出远程存储并立即恢复计算,无需专用通信线程或流多处理器。FlashAttention-3和FlashAttention-4相比串行执行分别达到最高1.43x和1.42x的发送方可见加速,长序列下稳态时间趋近于FlashAttention原版。结果确立了通信后进度作为可测量调度杠杆,用于调控突发、避免背压、隐藏发送方开销。

原文摘要 · Abstract (English)

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.

分布式推理通信优化DiT性能调度

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。