TokenWeave: Efficient Compute-Communication Overlap for Distributed LLM Inference

TL;DR

TokenWeave is a system that enables efficient compute-communication overlap for distributed LLM inference at small token lengths, significantly reducing latency and increasing throughput.

Contribution

It introduces a novel fused AllReduce--RMSNorm kernel leveraging GPU features, achieving improved performance for tensor-parallel inference.

Findings

Up to 1.28x latency speedup over baseline.

Up to 1.19x higher throughput compared to baseline.

Effective for token lengths as small as 1024.

Abstract

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\times$H100 DGX system. Our evaluations demonstrate up to $\boldsymbol{1.28\times}$ speedup in latency (baseline$\div$ours) and up to $\boldsymbol{1.19\times}$ higher throughput (ours$\div$baseline) 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.