forked from Karylab-cklius/vllm
Alexander MatveevandAlexander Matveev
e3b4fdaf5d
Port SGLang's push-based 2-buffer allreduce protocol into vLLM as a new communicator backend for small-message reductions. The push protocol eliminates the two explicit cross-GPU NVLink barrier round-trips used by the existing barrier-based CustomAllreduce, replacing them with a sentinel-based data arrival detection mechanism and double-buffered epoch alternation. Key advantages over the barrier-based approach: - Zero barriers: data arrival IS the synchronization (positive-zero sentinel) - Single NVLink round-trip instead of two barrier exchanges + remote reads - All SMs active (SM_count CTAs vs 2 CTAs) for higher NVLink bandwidth - No cudaMemcpy to IPC staging buffer in eager mode - PDL (griddepcontrol) support for kernel overlap on sm_90+ The new PushAllReduce is inserted in the CudaCommunicator dispatch chain above the existing CustomAllreduce for messages below a size threshold (~720 KB at TP=8). Larger messages continue to use the barrier-based path. The existing CustomAllreduce code is not modified. Measured results on DeepSeek-V4-Pro (61 layers, TP=8, 8x NVIDIA B200, BS=1, decode with ISL=4, OSL=33024): - Throughput: +2.14% (84.06 vs 82.30 tokens/s) - TPOT: -2.09% (11.90 vs 12.15 ms/token) Correctness verified via lm_eval gsm8k 5-shot with no regression (exact_match delta within statistical noise). The feature can be disabled at runtime via VLLM_DISABLE_PUSH_ALLREDUCE=1 to fall back to the barrier-based path. Signed-off-by: Alexander Matveev <amatveev@redhat.com>
64 lines
1.6 KiB
Python
64 lines
1.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""
|
|
Worker helper for push allreduce unit tests.
|
|
Run via torch.multiprocessing.spawn from test_push_all_reduce.py.
|
|
|
|
Provides init/teardown helpers that create separate gloo (CPU) and
|
|
nccl (device) process groups for PushAllReduce (which needs gloo for
|
|
IPC handle exchange) and NCCL reference reduction (which needs nccl).
|
|
"""
|
|
|
|
import os
|
|
import socket
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
|
|
def find_free_port() -> int:
|
|
"""Find a free TCP port for distributed init."""
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
s.bind(("", 0))
|
|
return s.getsockname()[1]
|
|
|
|
|
|
# Global references to process groups created by init_groups
|
|
_cpu_group = None
|
|
_nccl_group = None
|
|
|
|
|
|
def init_groups(rank: int, world_size: int, port: int):
|
|
"""Initialize gloo (CPU) and nccl process groups.
|
|
|
|
PushAllReduce uses the gloo group for IPC handle exchange.
|
|
NCCL group is used for reference allreduce.
|
|
"""
|
|
global _cpu_group, _nccl_group
|
|
|
|
os.environ["MASTER_ADDR"] = "localhost"
|
|
os.environ["MASTER_PORT"] = str(port)
|
|
|
|
torch.cuda.set_device(rank)
|
|
|
|
dist.init_process_group(
|
|
backend="gloo", rank=rank, world_size=world_size
|
|
)
|
|
_cpu_group = dist.group.WORLD
|
|
|
|
# Create a separate NCCL group for reference allreduce
|
|
_nccl_group = dist.new_group(backend="nccl")
|
|
|
|
|
|
def get_cpu_group():
|
|
return _cpu_group
|
|
|
|
|
|
def get_nccl_group():
|
|
return _nccl_group
|
|
|
|
|
|
def teardown():
|
|
"""Clean up distributed groups."""
|
|
dist.destroy_process_group()
|