TokenWeave: Efficient Compute-Communication Overlap for Distributed LLM Inference

TL;DR

TokenWeave tackles persistent compute-communication overheads in distributed LLM inference by splitting the input tokens into two wave-aware halves and overlapping the communication of one half with the computation of the other. It also reorders RMSNorm and introduces a fused AllReduce--RMSNorm kernel that leverages Multimem on Hopper/Blackwell GPUs to minimize SM usage and to overlap memory-bound RMSNorm with other work. Empirical results show up to $1.29\times$ latency reductions and $1.26\times$ throughput gains across multiple models and workloads, with occasional surpassing of an equivalent all-communication-removed model due to the fusion. The contributions are practical and deployable, integrated into vLLM-V1 and available publicly, advancing efficient distributed LLM inference through hardware-aware fusion and coarse-grained overlap.

Abstract

Distributed inference of large language models (LLMs) can introduce overheads of up to 20% even over GPUs connected via high-speed interconnects such as NVLink. Multiple techniques have been proposed to mitigate these overheads by decomposing computations into finer-grained tasks and overlapping communication with sub-tasks as they complete. However, fine-grained decomposition of a large computation into many smaller computations on GPUs results in overheads. Furthermore, the communication itself uses many streaming multiprocessors (SMs), adding to the overhead. We present TokenWeave to address these challenges. TokenWeave proposes a Token-Splitting technique that divides the tokens in the inference batch into two approximately equal subsets in a wave-aware manner. The communication of one subset is then overlapped with the computation of the other. In addition, TokenWeave optimizes the order of the layer normalization computation with respect to communication operations and implements a novel fused AllReduce--RMSNorm kernel that carefully leverages Multimem instruction support available on Hopper and Blackwell NVIDIA GPUs. These optimizations allow TokenWeave to perform communication and RMSNorm using only 2-8 SMs. Moreover, our kernel enables the memory-bound RMSNorm to be overlapped with the other batch's computation, providing additional gains. Our evaluations demonstrate up to 1.29x speedup in latency and 1.26x higher throughput across multiple models and workloads. In several settings, TokenWeave results in better performance compared to an equivalent model with all communication removed.

Figures (16)

• Figure 1: Communication overhead of AllReduce for three models vs. sequence length on an $8\times$H100 DGX (error bar shows std. dev. across 5 runs). Even with NVLink/NVSHARP, communication overheads can be over $20\%$.
• Figure 2: Inference latency of Llama-3.3-70B on 8xH100 DGX for various sequence lengths. vllm-Multimem corresponds to vLLM with an optimized AllReduce implementation using Multimem and NVSHARP support. vllm-nocomm is a counterfactual baseline corresponding to only the computation time without any communication. The dotted lines show performance normalized to the vllm-Multimem baseline. TokenWeave achieves up to $1.29\times$ speedup. Even at shorter sequence lengths, TokenWeave provides significant gains, for example, $1.18\times$ at a sequence length of $1K$ tokens, while prior schemes incur overheads. At sequence lengths $\ge 4K$, TokenWeave outperforms vllm-nocomm by not only recovering the complete communication overhead but also providing additional gains due to our fused AllReduce--RMSNorm kernel.
• Figure 3: RMSNorm overhead for different models vs sequence length on an $8\times$H100 DGX (error bar shows std. dev. across 5 runs). RMSNorm performed after AllReduce has non-trivial overheads and can range from $5-9\%$.
• Figure 4: Latency of executing a single AllReduce and RMSNorm operation for $8192$ dimension hidden size, $bf16$ on $8\times$H100 DGX across varying sequence lengths for three approaches. Simple reordering of RMSNorm between ReduceScatter (RS) and AllGather (AG) is worse than performing RMSNorm after AllReduce except at high sequence lengths as the overheads of splitting AllReduce (AR) into RS and AG eat into the gains. Our fused AllReduce--RMSNorm kernel results in up to $1.40\times$ improvement over the entire sequence length range.
• Figure 5: Splitting AllReduce (AR) into ReduceScatter (RS) and AllGather (AG) can result in non-trivial overheads. Shown are the individual times and the relative performance (line plots) of these operations on $8\times$H100 DGX. All runs are with a hidden size of $8192$ with $bf16$.