# X-Stage: Modeling Post-Issue Backpressure in GPU Communication--Computation Fusion

arXiv: https://arxiv.org/abs/2607.23264 · HTML: https://mlsanity.com/paper/2607.23264/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 16 of its 32 references that are in the map.
- First posted on arXiv 25 Jul 2026.

## Builds on (most cited first, 16 of 16)

- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models](https://mlsanity.com/paper/2309.14509/) (2023)
- [MegaBlocks: Efficient Sparse Training with Mixture-of-Experts](https://mlsanity.com/paper/2211.15841/) (2022)
- [FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling](https://mlsanity.com/paper/2603.05451/) (2026)
- [USP: A Unified Sequence Parallelism Approach for Long Context Generative AI](https://mlsanity.com/paper/2405.07719/) (2024)
- [Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts](https://mlsanity.com/paper/2502.19811/) (2025)
- [Breaking the Computation and Communication Abstraction Barrier in Distributed Machine Learning Workloads](https://mlsanity.com/paper/2105.05720/) (2021)
- [TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives](https://mlsanity.com/paper/2503.20313/) (2025)
- [FSMoE: A Flexible and Scalable Training System for Sparse Mixture-of-Experts Models](https://mlsanity.com/paper/2501.10714/) (2025)
- [ParallelKittens: Systematic and Practical Simplification of Multi-GPU AI Kernels](https://mlsanity.com/paper/2511.13940/) (2025)
- [Efficient and Adaptable Overlapping for Computation and Communication via Signaling and Reordering](https://mlsanity.com/paper/2504.19519/) (2025)
- [Demystifying NVSHMEM: A System-Level Analysis on Symmetric Memory and Device-Initiated Operations in GPU Communication](https://mlsanity.com/paper/2606.05951/) (2026)
- [ReaLB: Real-Time Load Balancing for Multimodal MoE Inference](https://mlsanity.com/paper/2604.19503/) (2026)

## 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.
