TokenWeave: Efficient Compute-Communication Overlap for Distributed LLM Inference

1. Problem Statement

Distributed inference of Large Language Models (LLMs) using Tensor Parallelism (TP) faces significant performance bottlenecks due to communication overheads, even on high-speed interconnects like NVLink.

• Communication Overhead: In TP deployments, communication (specifically AllReduce) accounts for 9–23% of end-to-end inference latency.
• The "Small Batch" Challenge: Existing techniques to overlap communication with computation (e.g., decomposing work into smaller tiles or tokens) are not enabled by default in major serving systems (vLLM, SGLang, TensorRT-LLM). This is because decomposing small workloads (typical in low-latency inference, e.g., <2K tokens) often introduces compute inefficiencies (wave quantization effects) that outweigh the gains from overlapping communication.
• Resource Contention: Communication operations often consume a large number of Streaming Multiprocessors (SMs) (16–20+), starving the compute kernels of resources.
• RMSNorm Redundancy: The RMSNorm operation, often overlooked, adds 4–9% overhead. In standard TP, it is performed redundantly on every GPU after AllReduce, despite all GPUs having identical data at that point.

2. Methodology: TokenWeave

TokenWeave is a system designed to enable efficient compute-communication overlap for tensor-parallel inference, even for small batch sizes (down to 1024 tokens). It employs three core innovations:

A. Wave-Aware Smart Token Splitting

Instead of naively splitting a batch into two equal halves, TokenWeave uses a wave-aware splitting strategy.

• Mechanism: It partitions the input batch into two subsets (prefix and suffix) such that the total number of execution "waves" (groups of CTAs) required for both subsets does not exceed the waves required for the original unsplit batch.
• Goal: This minimizes wave quantization overheads, ensuring that splitting the work does not degrade compute efficiency.
• Execution: The system uses two CUDA streams (Compute and Communication). While the AllReduce (fused with RMSNorm) runs on the first split, the attention computation runs on the second split, and vice versa.

B. RMSNorm Reordering and Fusion

TokenWeave identifies RMSNorm as a critical optimization target and reorders the execution flow.

• Reordering: Instead of performing AllReduce followed by RMSNorm, TokenWeave decomposes AllReduce into ReduceScatter and AllGather.
  1. ReduceScatter: Each GPU receives its shard of the tensor.
  2. Local RMSNorm: Each GPU performs RMSNorm on its local shard before the final aggregation. Since the data is already partitioned, this eliminates redundant computation across GPUs.
  3. AllGather: The normalized shards are gathered.
• Fused Kernel: A novel Fused AllReduce–RMSNorm kernel is implemented. It fuses ReduceScatter, RMSNorm, and AllGather into a single operation.
  • Memory Optimization: By computing variance directly on the ReduceScatter result in SM registers, it eliminates one HBM read. By writing normalized values directly to Multimem addresses for AllGather, it eliminates one HBM write.
  • SM Efficiency: This kernel leverages NVSHARP/Multimem (hardware features on Hopper/Blackwell GPUs) to perform communication and normalization using only 2–8 SMs, leaving the majority of SMs free for computation.

C. Adaptive Overlap Strategy

TokenWeave dynamically decides whether to use full splitting/overlap or just the fused kernel based on the batch size:

• Large Batches: Enables full two-way splitting and overlap.
• Small Batches: Disables splitting (to avoid overhead) but still applies the Fused AllReduce–RMSNorm kernel to reduce memory traffic and eliminate redundant computation.

3. Key Contributions

1. RMSNorm Optimization: First to identify and optimize RMSNorm in TP inference via reordering and a custom fused kernel, reducing HBM traffic and redundant compute.
2. Wave-Aware Splitting: A novel token-splitting algorithm that prevents compute efficiency loss due to wave quantization, enabling overlap even for small token batches (1K–2K).
3. Hardware-Aware Efficiency: Leverages NVSHARP/Multimem to reduce communication SM usage from ~20 to 2–8, preventing resource starvation.
4. System Integration: The first system to successfully enable compute-communication overlap for tensor-parallel serving in low-latency settings, integrated into the vLLM framework.

4. Experimental Results

Evaluations were conducted on 8×H100 DGX systems using models like Llama-3.3-70B, Qwen2.5-72B, and Mixtral-8x22B.

• Latency Improvement:
  • Achieves up to 1.28× speedup in latency compared to an optimized baseline (vLLM with Multimem).
  • Significant gains even at short sequence lengths (e.g., 1.2× at 1K tokens).
  • Outperforms state-of-the-art overlap solutions like TileLink, which only show gains at larger batch sizes (≥4K) and incur overheads at smaller sizes.
• Throughput Improvement:
  • Delivers up to 1.19× higher throughput on real-world traces (ShareGPT, arXiv).
  • In some settings, TokenWeave outperforms a theoretical baseline with zero communication (vLLM-nocomm) because its fused kernel optimizes the memory-bound RMSNorm operation better than the baseline handles pure computation.
• Kernel Efficiency:
  • The fused AllReduce–RMSNorm kernel provides 1.34–1.39× speedup over sequential execution across all token sizes.
  • It achieves near-optimal performance using only 4–8 SMs.

5. Significance

• Enables Low-Latency Distributed Inference: By solving the "small batch" problem, TokenWeave makes distributed TP inference viable for interactive applications where latency is critical, without requiring massive batch sizes.
• Hardware Utilization: It maximizes the utility of modern GPU features (NVSHARP, Multimem), shifting the bottleneck from communication overhead to pure computation.
• Practical Impact: The source code is open-sourced and integrated into vLLM, offering immediate performance benefits for production LLM serving systems without requiring changes to model architecture.
• Paradigm Shift: It challenges the assumption that communication overlap is only beneficial for large training batches or expert-parallelism (MoE) setups, proving it is essential for standard tensor-parallel inference as well.