Technical Summary: Understanding and Improving Communication Performance in Multi-node LLM Inference

Problem Statement

As Large Language Models (LLMs) increase in size, their memory footprints often exceed the capacity of a single GPU, necessitating distributed inference across multiple nodes. While model parallelism strategies like Tensor Parallelism (TP) and Pipeline Parallelism (PP) are well-studied for single-node inference, their performance in multi-node settings remains under-explored. Multi-node inference introduces significant challenges due to higher inter-node latencies compared to intra-node connections (e.g., NVLink).

The authors identify that while TP is often superior for memory-bound, decode-heavy workloads, it suffers from substantial communication overheads due to frequent all-reduce operations. Conversely, Hybrid Parallelism (HP), which combines TP and PP, reduces communication overhead but fails to effectively reduce computation time in decode-heavy regimes due to the nature of small matrix multiplications. A critical bottleneck identified is the sub-optimal performance of standard NCCL all-reduce implementations for the small message sizes (128 KB to 2 MB) typical of LLM decode phases, particularly when scaling across nodes.

Methodology

The authors conducted a systematic performance study using a combination of existing state-of-the-art inference engines and a custom research prototype:

1. Experimental Setup:

   • Hardware: Experiments were conducted on the Perlmutter system (NVIDIA A100 GPUs, HPE Slingshot-11 interconnect) and the Vista system (NVIDIA GH200 GPUs, InfiniBand interconnect).
   • Models: Llama 3.1 70B and 405B (Instruct) in bf16 precision.
   • Workloads: Two primary regimes were tested: "Prefill-heavy" (large prompt lengths) and "Decode-heavy" (long generation sequences).
   • Engines: The study utilized vLLM and SGLang alongside YALIS, a custom open-source inference engine developed by the authors to facilitate controlled experimentation in multi-node HPC environments. YALIS features a unified model definition layer, Torch Compile-based execution, and AxoNN-based tensor parallelism.

2. Analysis Approach:

   • The authors performed strong-scaling experiments, measuring end-to-end batch latency as GPU counts increased from single-node to 32-node configurations.
   • They decomposed execution time into computation (Matmul, Other Comp.), communication, and idle time to identify bottlenecks.
   • Synthetic GEMM benchmarks were used to isolate the effects of micro-batching (PP) versus tensor splitting (TP) on matrix multiplication kernels.

3. Proposed Solution (NVRAR):

   • To address the communication bottleneck, the authors developed NVRAR, a hierarchical all-reduce algorithm based on recursive doubling using NVSHMEM.
   • Design: NVRAR operates in three phases:
     1. Intra-node Reduce-Scatter: Uses NCCL to reduce data within a node.
     2. Inter-node Recursive-Doubling: Uses NVSHMEM non-blocking RMA primitives (put_nbi) to perform all-reduce across nodes using a recursive doubling topology.
     3. Intra-node All-Gather: Uses NCCL to gather the reduced data back to all GPUs within a node.
   • Optimizations:
     • Chunked Non-blocking Communication: Overlaps computation and communication by processing data in chunks.
     • Fused Payloads: Combines data and synchronization flags into a single 8-byte payload to avoid expensive explicit signaling primitives on Slingshot/InfiniBand.
     • Sequence-Number Synchronization: Uses unique sequence numbers for global synchronization to avoid the latency overhead of global barriers (fence/quiet) while ensuring correctness.

Key Contributions

1. Systematic Performance Study: The paper provides a detailed breakdown of TP vs. HP scaling behaviors in multi-node settings, identifying that TP is generally superior for decode-heavy workloads but is limited by communication overheads, while HP struggles to reduce computation time in these regimes.
2. YALIS Engine: The development of YALIS, a research-oriented inference engine designed for instrumentability and controlled experimentation in multi-node HPC environments.
3. NVRAR Algorithm: The proposal and implementation of NVRAR, a custom all-reduce algorithm optimized for the small-message regime characteristic of LLM decode phases. It leverages NVSHMEM to achieve lower latency than standard NCCL implementations on Slingshot and InfiniBand networks.
4. Integration and Evaluation: The integration of NVRAR into YALIS and vLLM, demonstrating end-to-end performance improvements for large models (Llama 3.1 405B) and Mixture-of-Experts (MoE) models (Qwen3-235B).

Results

• Scaling Behavior: Both TP and HP exhibit poor strong scaling in multi-node settings. HP outperforms TP in compute-bound (prefill-heavy) regimes, while TP outperforms HP in memory-bound (decode-heavy) regimes.
• Communication Bottlenecks: NCCL all-reduce scales poorly across nodes for small messages (128 KB–2 MB), often performing worse than MPI-based solutions in these specific regimes.
• NVRAR Performance:
  • Microbenchmarks: NVRAR achieves 1.9×–3.6× lower latency than NCCL for message sizes between 128 KB and 2 MB on HPE Slingshot and InfiniBand interconnects.
  • End-to-End Inference: Integrated into YALIS and vLLM, NVRAR reduces end-to-end batch latency by up to 1.72× for the Llama 3.1 405B model in decode-heavy workloads using tensor parallelism.
  • MoE Models: NVRAR also improves throughput for MoE models (e.g., Qwen3-235B) by accelerating the TP component of the parallelism strategy, achieving up to 1.14× higher throughput compared to NCCL-based configurations.

Significance

The paper claims that as LLMs grow larger and inference moves to multi-node clusters, communication efficiency becomes the primary determinant of performance. The authors demonstrate that standard communication libraries (NCCL) are not optimized for the specific message sizes and latency constraints of LLM decode phases. By introducing NVRAR, they provide a targeted solution that significantly reduces communication overhead, enabling more efficient scaling of tensor parallelism for large models. This work highlights the necessity of custom communication primitives for next-generation distributed inference systems and validates the utility of hierarchical, NVSHMEM-based approaches over standard collective communication libraries in specific HPC interconnect environments.