X-Stage: Modeling Post-Issue Backpressure in GPU Communication--Computation Fusion
Not on the map. Why?
Abstract (arXiv)
Fine-grained, device-initiated communication allows fused GPU kernels to issue remote stores directly from their compute pipelines, a pattern increasingly used in expert parallelism (EP), tensor parallelism (TP), and Ulysses-style sequence parallelism (UP). Existing designs reason about where communication is issued and when remote data becomes ready, but lack a quantitative model of the sender-side interval after a remote store is accepted and before it becomes visible at the destination. This interval determines whether communication remains decoupled from computation or backpressures it. We identify X-Stage, a software-visible post-issue stage with finite decoupling. Downstream pressure can dissipate while the issuer resumes useful work, whereas sustained injection consumes X-Stage headroom and eventually stalls the compute pipeline. We characterize this behavior and build a calibrated model that predicts whether remote-store arrivals accumulate backpressure or recover during intervening computation. Guided by the model, we reshape bursty arrivals when they would exhaust X-Stage headroom and exploit natural compute windows when headroom can recover concurrently. Evaluation across representative EP, TP, and UP workloads shows up to 1.62x fused-kernel, 1.75x end-to-end, and 1.43x sender-visible speedup, respectively. Microbenchmarks further validate the model's predictions of backlog accumulation, recovery, and sender-side backpressure.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Builds on 16 of its 32 references that are in the map. First posted on arXiv 25 Jul 2026.
Builds on
16 of 32 references.
References · 16
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
- 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,551
- 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,829
- 2024FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precisioncited by 507
- 2023DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Modelscited by 239
- 2022MegaBlocks: Efficient Sparse Training with Mixture-of-Expertscited by 204
- 2024USP: A Unified Sequence Parallelism Approach for Long Context Generative AIcited by 58
- 2026FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scalingcited by 58
- 2025Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Expertscited by 46
- 2021Breaking the Computation and Communication Abstraction Barrier in Distributed Machine Learning Workloadscited by 40
- 2025TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitivescited by 23
- 2025FSMoE: A Flexible and Scalable Training System for Sparse Mixture-of-Experts Modelscited by 17
4 more
- 2025ParallelKittens: Systematic and Practical Simplification of Multi-GPU AI Kernelscited by 11
- 2025Efficient and Adaptable Overlapping for Computation and Communication via Signaling and Reorderingcited by 9
- 2026Demystifying NVSHMEM: A System-Level Analysis on Symmetric Memory and Device-Initiated Operations in GPU Communicationcited by 4
- 2026ReaLB: Real-Time Load Balancing for Multimodal MoE Inferencecited by 1
Built on by
Not cited yet.