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$÷$ours) and up to $\boldsymbol{1.19\times}$ higher throughput (ours$÷$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.

翻译：基于张量并行的大语言模型分布式推理，即便在通过高速GPU互连NVLink连接的设备上，通信开销仍可达20%。现有研究提出通过将计算分解为细粒度子任务并与通信重叠来降低此类开销。然而，在vLLM、SGLang及TensorRT-LLM等系统的张量并行服务中，默认并未启用任何这类技术。其原因在于，为支持低延迟服务，每次迭代处理的token数量通常保持较小规模，而分解此类轻量级工作负载以实现通信重叠反而会导致性能下降。此外，通信过程本身占用大量原本可用于计算的流式多处理器，进一步加剧了开销。我们提出TokenWeave——首个能对低至1024 token长度的张量并行模型推理实现高效计算-通信重叠的系统。TokenWeave识别出此前被忽略的关键算子RMSNorm，通过实现新型融合AllReduce-RMSNorm核函数对其进行优化，并同步优化通信过程。该核函数利用现代GPU（如Hopper、Blackwell）支持的NVSHARP/Multimem特性，在8×H100 DGX系统上仅需2-8个流式多处理器即可高效协同完成通信与RMSNorm计算。评估表明，在多种模型与工作负载下，TokenWeave可实现最高1.28倍的延迟加速比（基线÷本系统）及最高1.19倍的吞吐量提升（本系统÷基线）。在多种场景中，TokenWeave甚至优于完全消除通信的等效模型。开源代码见https://github.com/microsoft/tokenweave。