# TokenWeave: Efficient Compute-Communication Overlap for Distributed LLM Inference

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

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 44 references that are in the map.
- First posted on arXiv 16 May 2025.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [Mixtral of Experts](https://mlsanity.com/paper/2401.04088/) (2024)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [SGLang: Efficient Execution of Structured Language Model Programs](https://mlsanity.com/paper/2312.07104/) (2023)
- [DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving](https://mlsanity.com/paper/2401.09670/) (2024)
- [Splitwise: Efficient generative LLM inference using phase splitting](https://mlsanity.com/paper/2311.18677/) (2023)
- [Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve](https://mlsanity.com/paper/2403.02310/) (2024)
- [A Discourse-Aware Attention Model for Abstractive Summarization of Long Documents](https://mlsanity.com/paper/1804.05685/) (2018)
- [OpenChat: Advancing Open-source Language Models with Mixed-Quality Data](https://mlsanity.com/paper/2309.11235/) (2023)
- [FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion](https://mlsanity.com/paper/2406.06858/) (2024)
- [Breaking the Computation and Communication Abstraction Barrier in Distributed Machine Learning Workloads](https://mlsanity.com/paper/2105.05720/) (2021)
- [Triton-distributed: Programming Overlapping Kernels on Distributed AI Systems with the Triton Compiler](https://mlsanity.com/paper/2504.19442/) (2025)
- [TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives](https://mlsanity.com/paper/2503.20313/) (2025)
- [Lancet: Accelerating Mixture-of-Experts Training via Whole Graph Computation-Communication Overlapping](https://mlsanity.com/paper/2404.19429/) (2024)

## Built on by (most cited first, 14 of 14)

- [LLM-42: Enabling Determinism in LLM Inference with Verified Speculation](https://mlsanity.com/paper/2601.17768/) (2026)
- [MPK: A Compiler and Runtime for Mega-Kernelizing Tensor Programs](https://mlsanity.com/paper/2512.22219/) (2025)
- [xLLM Technical Report](https://mlsanity.com/paper/2510.14686/) (2025)
- [Kareus: Joint Reduction of Dynamic and Static Energy in Large Model Training](https://mlsanity.com/paper/2601.17654/) (2026)
- [FlashAccel: Leveraging High-Bandwidth Flash (HBF) for High-Throughput LLM Inference](https://mlsanity.com/paper/2607.10186/) (2026)
- [Rearchitecting Datacenter Lifecycle for AI: A TCO-Driven Framework](https://mlsanity.com/paper/2509.26534/) (2025)
- [CUCo: An Agentic Framework for Compute and Communication Co-design](https://mlsanity.com/paper/2603.02376/) (2026)
- [Balancing Sustainability And Performance: The Role Of Small-Scale LLMs In Agentic Artificial Intelligence Systems](https://mlsanity.com/paper/2601.19311/) (2026)
- [A Switch-Centric In-Network Architecture for Accelerating LLM Inference in Shared-Memory Network](https://mlsanity.com/paper/2603.28239/) (2026)
- [Rethinking Network Topologies for Cost-Effective Mixture-of-Experts LLM Serving](https://mlsanity.com/paper/2605.00254/) (2026)
- [DynaFlow: Transparent and Flexible Intra-Device Parallelism via Programmable Operator Scheduling](https://mlsanity.com/paper/2605.21603/) (2026)
- [ASAP: A Disaggregated and Asynchronous Inference System for MoE Prefill](https://mlsanity.com/paper/2606.22541/) (2026)
- [Weave: Fine-Grained Dynamic SM Scheduling in an MoE Megakernel for Compute-Communication Overlap](https://mlsanity.com/paper/2609.21483/) (2026)
- [Echo: Error-Compensated 4-Bit Tensor-Parallel Communication for LLM Training and Inference](https://mlsanity.com/paper/2610.14277/) (2026)

## Abstract (arXiv)

Distributed inference of large language models (LLMs) using tensor parallelism can introduce communication overheads of 20% even over GPUs connected via NVLink, a high-speed GPU interconnect. Several techniques have been proposed to mitigate these overheads by decomposing computations into smaller tasks and overlapping communication with these subtasks. However, none of these techniques are turned on by default during tensor-parallel serving in systems like vLLM, SGLang and TensorRT-LLM. This is because the number of tokens processed per iteration is typically kept small to support low-latency serving, and decomposing such smaller workloads to enable communication overlap results in worse performance. Further, the communication itself uses many streaming multiprocessors (SMs) that would otherwise be available for computation, increasing overhead. We present TokenWeave, the first system to enable efficient compute-communication overlap for tensor-parallel model inference for token lengths as small as 1024. TokenWeave identifies RMSNorm, a previously overlooked operation, as crucial and optimizes it along with communication by implementing a novel fused AllReduce--RMSNorm kernel. Further, this kernel leverages the NVSHARP/Multimem feature available on modern GPUs (e.g., Hopper, Blackwell) to jointly perform communication and RMSNorm efficiently using only 2-8 streaming multiprocessors (SMs) on an 8×H100 DGX system. Our evaluations demonstrate up to boldsymbol1.28× speedup in latency (baselinedivours) and up to boldsymbol1.19× higher throughput (oursdivbaseline) across multiple models and workloads. In several settings, TokenWeave delivers better performance than an equivalent model with all communication removed. The source code is available at https://github.com/microsoft/tokenweave.