[PD] Skip singleton transfer-status all-reduces (#40003)

Co-authored-by: Pranjal Shankhdhar <pranjalssh@users.noreply.github.com>
Co-authored-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>

python/sglang/srt/disaggregation/utils.py

@@ -185,6 +185,9 @@ def _apply_metadata_gate(polls, decode_reqs, metadata_buffers) -> None:
 
 def _all_reduce_polls(polls: List[int], group: dist.ProcessGroup) -> List[int]:
     """MIN-reduce poll states so no rank commits ahead of its peers."""
+    if dist.get_world_size(group) == 1:
+        return polls
+
     tensor_to_reduce = torch.tensor(polls, dtype=torch.uint8, device="cpu")
     dist.all_reduce(tensor_to_reduce, op=dist.ReduceOp.MIN, group=group)
     return tensor_to_reduce.tolist()

python/sglang/srt/managers/scheduler.py

@@ -4888,9 +4888,12 @@ class Scheduler(
                 self.load_publisher.publish_load_stat(
                     self.load_inquirer.get_loads, force=True, snapshot=snapshot
                 )
-            if self.enable_hicache_storage:
-                # Storage workers need the GIL between I/O calls. Yield while
-                # there is no GPU batch so polling cannot starve their acks.
+            if (
+                self.enable_hicache_storage
+                or self.disaggregation_mode != DisaggregationMode.NULL
+            ):
+                # Storage and transfer workers need the GIL between I/O calls.
+                # Singleton PD polls no longer yield through a collective.
                 time.sleep(0)
             return
         self.metrics_reporter.record_scheduler_idle()

test/registered/unit/disaggregation/test_disaggregation_wire.py

@@ -2,12 +2,13 @@ import struct
 import threading
 import unittest
 from types import SimpleNamespace
-from unittest.mock import Mock, patch
+from unittest.mock import ANY, Mock, call, patch
 
 import numpy as np
 import torch
+import torch.distributed as dist
 
-from sglang.srt.disaggregation.base.conn import KVArgs, StateType
+from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll, StateType
 from sglang.srt.disaggregation.common.conn import CommonKVManager
 from sglang.srt.disaggregation.common.staging_buffer import (
     StagingAllocator,
@@ -36,6 +37,9 @@ from sglang.srt.disaggregation.utils import (
     build_transfer_entry_pairs,
     compute_mamba_state_slice_byte_blocks,
     get_qsa_pending_state_indices,
+    poll_and_all_reduce,
+    poll_and_all_reduce_attn_cp_tp_group,
+    poll_and_all_reduce_with_staging,
     setup_state_kv_args,
     should_send_replicated_state,
 )
@@ -158,6 +162,117 @@ class TestDisaggregationWire(unittest.TestCase):
         self.assertEqual(unpack_list_of_buffers(pack_list_of_buffers(bufs)), bufs)
 
 
+class TestPollCollectives(CustomTestCase):
+    def setUp(self):
+        super().setUp()
+        failure_prob = patch.object(
+            envs.SGLANG_TEST_DISAGG_FAILURE_PROB, "get", return_value=0
+        )
+        failure_prob.start()
+        self.addCleanup(failure_prob.stop)
+
+    def test_tp_cp_consensus_skips_only_singleton_groups(self):
+        """Poll states retain TP/CP consensus without singleton collectives."""
+        local = [KVPoll.Success, KVPoll.Success, KVPoll.Success]
+        tp_peers = [KVPoll.Success, KVPoll.Transferring, KVPoll.Success]
+        cp_peers = [KVPoll.Bootstrapping, KVPoll.Success, KVPoll.Failed]
+        for tp_size, cp_size in [(4, 1), (2, 2), (1, 4), (1, 1)]:
+            with self.subTest(tp_size=tp_size, cp_size=cp_size):
+                tp, cp = object(), object()
+                sizes = {tp: tp_size, cp: cp_size}
+                peers = {tp: tp_peers, cp: cp_peers}
+                pollers = [Mock(poll=Mock(return_value=state)) for state in local]
+
+                def reduce(tensor, op, group):
+                    self.assertEqual(op, dist.ReduceOp.MIN)
+                    if sizes[group] > 1:
+                        tensor.copy_(
+                            torch.minimum(
+                                tensor, torch.tensor(peers[group], dtype=tensor.dtype)
+                            )
+                        )
+
+                expected = local.copy()
+                expected_calls = []
+                for group in [tp, cp]:
+                    if sizes[group] > 1:
+                        expected = [min(a, b) for a, b in zip(expected, peers[group])]
+                        expected_calls.append(
+                            call(ANY, op=dist.ReduceOp.MIN, group=group)
+                        )
+                with (
+                    patch.object(dist, "get_world_size", side_effect=sizes.__getitem__),
+                    patch.object(dist, "all_reduce", side_effect=reduce) as all_reduce,
+                ):
+                    result = poll_and_all_reduce_attn_cp_tp_group(pollers, cp, tp)
+
+                self.assertEqual(result, expected)
+                self.assertEqual(all_reduce.call_args_list, expected_calls)
+                for poller in pollers:
+                    poller.poll.assert_called_once_with()
+
+    def test_singleton_polling_needs_no_tensor_or_collective(self):
+        """A single rank can poll transfers without allocating a reduction tensor."""
+        states = [KVPoll.Failed, KVPoll.Transferring, KVPoll.Success]
+        pollers = [Mock(poll=Mock(return_value=state)) for state in states]
+        with (
+            patch.object(dist, "get_world_size", return_value=1),
+            patch.object(dist, "all_reduce") as all_reduce,
+            patch.object(
+                torch, "tensor", side_effect=AssertionError("No tensor needed")
+            ),
+        ):
+            self.assertEqual(poll_and_all_reduce(pollers, object()), states)
+        all_reduce.assert_not_called()
+
+    def test_singleton_still_waits_for_decode_metadata(self):
+        pollers = [Mock(poll=Mock(return_value=KVPoll.Success)) for _ in range(2)]
+        requests = [
+            SimpleNamespace(
+                req=SimpleNamespace(bootstrap_host="127.0.0.1"),
+                metadata_buffer_index=index,
+            )
+            for index in range(2)
+        ]
+        metadata = SimpleNamespace(bootstrap_room=torch.tensor([[0], [42]]))
+        with (
+            patch.object(dist, "get_world_size", return_value=1),
+            patch.object(dist, "all_reduce") as all_reduce,
+        ):
+            result = poll_and_all_reduce(pollers, object(), requests, metadata)
+        self.assertEqual(result, [KVPoll.Transferring, KVPoll.Success])
+        all_reduce.assert_not_called()
+
+    def test_singleton_still_advances_and_waits_for_staging(self):
+        request = SimpleNamespace(
+            kv_receiver=Mock(
+                require_staging=True, poll=Mock(return_value=KVPoll.Success)
+            )
+        )
+        staging = Mock(
+            is_done=Mock(return_value=False), is_failed=Mock(return_value=False)
+        )
+        with (
+            patch.object(dist, "get_world_size", return_value=1),
+            patch.object(dist, "all_reduce") as all_reduce,
+        ):
+            result = poll_and_all_reduce_with_staging([request], staging, object())
+        self.assertEqual(result, [KVPoll.Transferring])
+        staging.advance_scatter.assert_called_once_with(request)
+        all_reduce.assert_not_called()
+
+    def test_singleton_preserves_failure_injection(self):
+        poller = Mock(poll=Mock(return_value=KVPoll.Success))
+        with (
+            patch.object(envs.SGLANG_TEST_DISAGG_FAILURE_PROB, "get", return_value=1),
+            patch.object(dist, "get_world_size", return_value=1),
+            patch.object(dist, "all_reduce") as all_reduce,
+        ):
+            result = poll_and_all_reduce([poller], object())
+        self.assertEqual(result, [KVPoll.Failed])
+        all_reduce.assert_not_called()
+
+
 class TestCPReplicatedStateTransfer(unittest.TestCase):
     def test_only_nonzero_cp_ranks_without_layer_split_skip_state(self):
         cases = [

test/registered/unit/managers/test_scheduler_on_idle_load.py

@@ -14,6 +14,7 @@ from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
 
 maybe_stub_sgl_kernel()
 
+from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.managers.scheduler import Scheduler
 
 register_cpu_ci(est_time=12, suite="base-a-test-cpu")
@@ -24,6 +25,7 @@ class TestOnIdleStallPublish(CustomTestCase):
         s = Scheduler.__new__(Scheduler)
         s.scheduler_stage_metrics = None
         s.enable_hicache_storage = False
+        s.disaggregation_mode = DisaggregationMode.NULL
         s.maybe_send_health_check_signal = MagicMock()
         s.is_fully_idle = MagicMock(return_value=False)  # stalled, not idle
         s.publish_load_snapshot = MagicMock(return_value=None)