From fd44100bb0c2a49483e5d2aa682d0d571fa67d69 Mon Sep 17 00:00:00 2001
From: Alexander Matveev
 <alexm-redhat@dgx-b200-02.mgmt.accl-001.lab.rdu2.dc.redhat.com>
Date: Mon, 15 Jun 2026 18:07:47 -0400
Subject: [PATCH] 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>
---
 .../kernels/benchmark_device_communicators.py | 32 +++++++++++++++++++
 .../device_communicators/cuda_communicator.py |  3 ++
 2 files changed, 35 insertions(+)

diff --git a/benchmarks/kernels/benchmark_device_communicators.py b/benchmarks/kernels/benchmark_device_communicators.py
index 24e22023b91..0753ebf1913 100644
--- a/benchmarks/kernels/benchmark_device_communicators.py
+++ b/benchmarks/kernels/benchmark_device_communicators.py
@@ -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(
diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py
index 325a9ee37f1..a2407e37aa6 100644
--- a/vllm/distributed/device_communicators/cuda_communicator.py
+++ b/vllm/distributed/device_communicators/cuda_communicator.py
@@ -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: