Address review: add push_ar to benchmark and backend logging

- Add PushAllReduce to benchmark_device_communicators.py for
  comparing against other allreduce implementations
- Add PUSH_AR to _log_all_reduce_backend_selection in
  cuda_communicator.py for visibility in dispatch logging

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Alexander Matveev <alexm-redhat@dgx-b200-02.mgmt.accl-001.lab.rdu2.dc.redhat.com>
This commit is contained in:
Alexander Matveev
2026-06-15 18:09:24 -04:00
co-authored by Claude Opus 4.6
parent bc50e5fc2e
commit fd44100bb0
2 changed files with 35 additions and 0 deletions
@@ -33,6 +33,7 @@ from vllm.distributed.device_communicators.custom_all_reduce import CustomAllred
from vllm.distributed.device_communicators.flashinfer_all_reduce import (
FlashInferAllReduce,
)
from vllm.distributed.device_communicators.push_all_reduce import PushAllReduce
from vllm.distributed.device_communicators.pynccl import (
PyNcclCommunicator,
register_nccl_symmetric_ops,
@@ -80,6 +81,7 @@ class CommunicatorBenchmark:
# Initialize communicators
self.custom_allreduce = None
self.push_ar_comm = None
self.pynccl_comm = None
self.symm_mem_comm = None
self.symm_mem_comm_multimem = None
@@ -106,6 +108,23 @@ class CommunicatorBenchmark:
)
self.custom_allreduce = None
try:
self.push_ar_comm = PushAllReduce(
group=self.cpu_group,
device=self.device,
max_size=self.max_size_override,
)
if not self.push_ar_comm.disabled:
logger.info("Rank %s: PushAllReduce initialized", self.rank)
else:
logger.info("Rank %s: PushAllReduce disabled", self.rank)
self.push_ar_comm = None
except Exception as e:
logger.warning(
"Rank %s: Failed to initialize PushAllReduce: %s", self.rank, e
)
self.push_ar_comm = None
try:
self.pynccl_comm = PyNcclCommunicator(
group=self.cpu_group, device=self.device
@@ -216,6 +235,19 @@ class CommunicatorBenchmark:
)
)
if self.push_ar_comm is not None:
comm = self.push_ar_comm
communicators.append(
(
"push_ar",
lambda t, c=comm: c.all_reduce(t),
lambda t, c=comm: c.should_use(t),
comm.capture(),
{},
None,
)
)
if self.pynccl_comm is not None:
comm = self.pynccl_comm
communicators.append(
@@ -224,6 +224,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
"NCCL_SYMM_MEM",
"QUICK_REDUCE",
"FLASHINFER",
"PUSH_AR",
"CUSTOM",
"SYMM_MEM",
"PYNCCL",
@@ -257,6 +258,8 @@ class CudaCommunicator(DeviceCommunicatorBase):
enabled_ar_backends.append("QUICK_REDUCE")
if self.fi_ar_comm is not None and not self.fi_ar_comm.disabled:
enabled_ar_backends.append("FLASHINFER")
if self.push_ar_comm is not None:
enabled_ar_backends.append("PUSH_AR")
if self.ca_comm is not None and not self.ca_comm.disabled:
enabled_ar_backends.append("CUSTOM")
if self.symm_mem_comm is not None and not self.symm_mem_comm.disabled: