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

• URL: http://arxiv.org/abs/2607.23264v1
• Date: Sat, 25 Jul 2026 16:04:36 GMT
• Title: X-Stage: An Overlooked Pipeline Stage for Communication-Computation Overlap in DiT Inference
• Abstract summary: Device-initiated communication lets persistent GPU kernels in distributed diffusion transformer (DiT) inference issue remote stores.<n>Existing systems schedule when communication is issued and when received data becomes consumable, but omit post-issue progress before remote-visible completion.<n>We identify X-Stage, a software-visible post-issue pipeline stage.
• Score: 18.883833948875488
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• VC-Attention: Value Smoothing and Softmax Casting for Low-bit Attention [47.30761996669105]
  We propose VC-Attention, a training-free low-bit attention framework.<n>It addresses both by smoothing value tokens by a probability Cast. V-Smooth reorders value tokens by lightweight clustering online, so the tokens in a fused hardware block quantize well together.<n>It maps log-domain scores directly to E4M3 probability codes with one fused multiply-add, eliminating the FP32 exponential and the format conversion.
  arXiv  Detail & Related papers  (2026-09-14T16:17:35Z)
• Faster Than Flash: Exploiting Attention Sparsity for Efficient Long-Context Decoding [51.48425605758328]
  We present Faster Flash Decoding (FFD), a novel hardware-algorithm co-design framework designed to break the memory wall in long-context decoding.<n>FFD integrates the selector and computer into a fully fused kernel, replacing external metadata indices with content-aware scanning.<n>It achieves up to 11.6x kernel-level speedup and scaling to 256K context length, with 2.37x end-to-end throughput improvement.
  arXiv  Detail & Related papers  (2026-08-31T15:13:20Z)
• An Efficient vLLM-Based Inference Pipeline for Unified Audio Understanding and Generation [63.50393485685279]
  We present a vLLM-based inference pipeline for unified speech understanding and generation.<n>We overcome the shared intuition-Free Guidance (CFG) halves throughput.<n>By co-scheduling paired conditional and unconditional requests within a continuous batch, our CFG implementation sustains 80% of non-CFG throughput.
  arXiv  Detail & Related papers  (2026-07-02T12:55:11Z)
• Scaling Parallel Sequence Models to Foundation-Scale Vision Encoders [123.58723804218151]
  Vision foundation models are bottlenecked by the quadratic cost of self-attention.<n>We present C-GSPN based on 2D spatial propagation propagation.
  arXiv  Detail & Related papers  (2026-05-30T14:29:43Z)
• Expert Streaming: Accelerating Low-Batch MoE Inference via Multi-chiplet Architecture and Dynamic Expert Trajectory Scheduling [29.998645673111536]
  Mixture-of-Experts is a promising approach for edge AI with low-batch inference.<n>MoE sparsity and dynamic gating shift distributed strategies toward much finer runtime granularity.<n>Fully Sharded Expert Data Parallelism is a parallelization paradigm specifically architected for low-batch MoE computation.
  arXiv  Detail & Related papers  (2026-03-29T10:44:55Z)
• Spava: Accelerating Long-Video Understanding via Sequence-Parallelism-aware Approximate Attention [63.69228529380251]
  Spava is a sequence-parallel framework with optimized attention for long-video inference.<n>Spava delivers speedups of 12.72x, 1.70x, and 1.18x over FlashAttn, ZigZagRing, and APB, without notable performance loss.
  arXiv  Detail & Related papers  (2026-01-29T09:23:13Z)
• Polychronous Wave Computing: Timing-Native Address Selection in Spiking Networks [0.0]
  Polychronous Wave Computing is a timing-native address-selection primitive.<n>Spike times are phase-encoded in a rotating frame and processed by a programmable multiport interferometer.<n>We derive the operating envelope imposed by phase wrapping and mutual coherence.
  arXiv  Detail & Related papers  (2026-01-19T14:12:01Z)
• Tiny, On-Device Decision Makers with the MiniConv Library [0.0]
  Reinforcement learning (RL) has achieved strong results, but deploying visual policies on resource-constrained edge devices remains challenging.<n>We introduce a split-policy architecture in which a small on-device encoder transforms each observation into a compact feature tensor that is transmitted to a remote policy head.<n>We evaluate across an NVIDIA Jetson Nano, a Raspberry Pi 4B, and a Raspberry Pi Zero 2 W, reporting learning results, on-device execution behaviour under sustained load, and end-to-end decision latency and scalability measurements under bandwidth shaping.
  arXiv  Detail & Related papers  (2025-12-17T00:53:30Z)
• Efficient and Adaptable Overlapping for Computation and Communication via Signaling and Reordering [13.185314408519107]
  Generative models have achieved remarkable success across various applications, driving the demand for multi-GPU computing.<n>We propose FlashOverlap, which utilizes a novel signaling mechanism.<n>Experiments show that FlashOverlap achieves up to 1.65x speedup through overlap, outperforming existing works in most cases.
  arXiv  Detail & Related papers  (2025-04-28T06:37:57Z)
• Task-Oriented Feature Compression for Multimodal Understanding via Device-Edge Co-Inference [54.53508601749513]
  We propose a task-oriented feature compression (TOFC) method for multimodal understanding in a device-edge co-inference framework.<n>To enhance compression efficiency, multiple entropy models are adaptively selected based on the characteristics of the visual features.<n>Results show that TOFC achieves up to 52% reduction in data transmission overhead and 63% reduction in system latency.
  arXiv  Detail & Related papers  (2025-03-17T08:37:22Z)
• Turbocharge Speech Understanding with Pilot Inference [0.9699101045941684]
  This paper sets to accelerate modern speech understanding on resource-constrained edge devices. It takes a hybrid approach: to speed up on-device execution; to offload inputs that are beyond the device's capacity. Our prototype, called PASU, is tested on Arm platforms with 6 - 8 cores: it delivers SOTA accuracy; it reduces the end-to-end latency by 2x and reduces the offloading needs by 2x.
  arXiv  Detail & Related papers  (2023-11-22T17:14:18Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.