[RFC] Intra-node shared memory (SHM) optimizations for communication operators on CPUs

Overview and Motivation

Tensor parallel (TP) inference is a key optimization and well adopted in LLM inference on CPUs. Most use cases for LLM benchmark/deployment focus on one node host (intra-node), where there are many numbers of CPU sockets/NUMAs/sub-NUMAs available.

Referring to the practices of well-known libraries including DeepSpeed , Sglang, vLLM and TGI, for intra-node TP cases, shared memory (SHM) optimizations bring much benefit for communication operators including allreduce, allgather and allgather_into_tensor.

In this RFC, we propose a low-latency SHM-based optimization for Gloo, intending to cover all key communication operators which are mostly used in LLM inference. And as the CPU default communication backend in PyTorch, this SHM optimization can bring significant improvement when working with PyTorch Tensor parallel solution.

Design:

1. Intra-node check and fallback

# Check if an intra-node case
if  local_size >= 0 and local_size == word_size  
 ->   SHM_lmpl

# Fallback
else ->  Gloo Ring_lmpl

2. OP register and dispatch

3. dtype scope:

In PyTorch integration, Gloo uses datatype defined in PyTorch, including c10::Half, c10::Float and c10::Bfloat16. Thus, our SHM impl is also targeted at c10::Half, c10::Float, and c10::Bfloat16.

4. Shared memory design

In general, there are 3 steps for SHM optimization (take SHM allreduce as example):
a. copy input to shared memory buffer.
b. do sum operation on reduce buffer (or gather for allgather).
c. copy result back to each rank

Compared to ring allreduce, it uses one united shared memory buffer for calculation in intra-node case. 
SHM allreduce won't use other collective primitives such as send, receive like in ring allreduce.
In that case it's faster and more efficient.

Implementation (Work in progress):

• allreduce. https://github.com/pytorch/gloo/pull/458 （under review)
• allgather.
• allgather_into_tensor.

Follow-ups:

• SHM path to work lazy init mode

Start by reviewing the RFC design and the linked Gloo pull request 458, which covers the allreduce work and is under review. Then examine how the remaining allgather, allgather_into_tensor, and lazy-init follow-ups fit the proposed SHM dispatch and fallback paths. Done means the planned operators and lazy-init path are implemented and reviewed.