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

Raja Gond, Nipun Kwatra, Ramachandran Ramjee (Microsoft, per the paper's open-source repository)

• Why: Tensor-parallel AllReduce costs ~20% even over NVLink, and — usefully documented here — none of the major serving engines enable overlap by default because it backfires at small token counts. This paper explains why and fixes it.
• What: The first system enabling efficient compute–communication overlap for tensor-parallel inference at batch sizes as small as 1024 tokens: up to 1.28× lower latency and 1.19× higher throughput across models and workloads on an 8×H100 DGX, in several settings beating an equivalent model with all communication removed. Open source (from abstract only).
• How: Identifies RMSNorm as the overlooked fusion partner and implements a fused AllReduce–RMSNorm kernel that exploits the NVSHARP/Multimem hardware feature, using only 2–8 SMs for communication on Hopper/Blackwell-class GPUs.
• Where applicable: Requires modern NVIDIA GPU interconnects (NVLink + NVSHARP/Multimem on Hopper/Blackwell) — the mechanism is silicon-specific. The diagnosis of why overlap is off by default in vLLM/SGLang/TensorRT-LLM is broadly instructive.