diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index d23f45282..af4738005 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/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() diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index dc73c6e65..2daac51cd 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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() diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index af6a62fdc..ea7247e45 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/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 = [ diff --git a/test/registered/unit/managers/test_scheduler_on_idle_load.py b/test/registered/unit/managers/test_scheduler_on_idle_load.py index 985b4526b..9cba96e57 100644 --- a/test/registered/unit/managers/test_scheduler_on_idle_load.py +++ b/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)