diff --git a/python/sglang/srt/constrained/grammar_manager.py b/python/sglang/srt/constrained/grammar_manager.py index 5442cb5d3..b039020fd 100644 --- a/python/sglang/srt/constrained/grammar_manager.py +++ b/python/sglang/srt/constrained/grammar_manager.py @@ -12,6 +12,7 @@ from sglang.srt.constrained.base_grammar_backend import ( create_grammar_backend, ) from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject +from sglang.srt.distributed.communication_tags import P2PTag from sglang.srt.environ import envs if TYPE_CHECKING: @@ -48,6 +49,10 @@ class GrammarManager: self.grammar_sync_size = scheduler.dp_tp_group.world_size self.grammar_sync_entry = scheduler.dp_tp_group.first_rank self.is_grammar_sync_entry = scheduler.dp_tp_group.is_first_rank + self.pp_rank = scheduler.ps.pp_rank + self.pp_size = scheduler.ps.pp_size + self.pp_group = scheduler.pp_group + self.grammar_pp_sync_work_list = [] self.SGLANG_GRAMMAR_POLL_INTERVAL = envs.SGLANG_GRAMMAR_POLL_INTERVAL.get() self.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS = ( @@ -64,6 +69,43 @@ class GrammarManager: def has_waiting_grammars(self) -> bool: return len(self.grammar_queue) > 0 + def _drain_pp_sync_work(self): + for p2p_work in self.grammar_pp_sync_work_list: + p2p_work.work.wait() + self.grammar_pp_sync_work_list.clear() + + def _pp_sync_ready_failed( + self, + ready_req_idxs: set[int], + failed_req_idxs: set[int], + ) -> tuple[set[int], set[int]]: + """ + Synchronize ready/failed grammar request indexes across the PP pipeline. + + PP0 provides the data. Each later PP rank receives it from the previous + rank and asynchronously forwards it to the next rank. + """ + if self.pp_size <= 1 or self.pp_group is None: + return ready_req_idxs, failed_req_idxs + + self._drain_pp_sync_work() + data = (ready_req_idxs, failed_req_idxs) + if self.pp_rank > 0: + data = self.pp_group.recv_object( + src=self.pp_rank - 1, + tag=P2PTag.GRAMMAR_PP_SYNC, + ) + if self.pp_rank + 1 < self.pp_size: + self.grammar_pp_sync_work_list.extend( + self.pp_group.send_object( + data, + dst=self.pp_rank + 1, + async_send=True, + tag=P2PTag.GRAMMAR_PP_SYNC, + ) + ) + return data + def abort_requests(self, recv_req: AbortReq): for req in self.grammar_queue: if recv_req.abort_all or req.rid.startswith(recv_req.rid): @@ -143,60 +185,86 @@ class GrammarManager: """ Move requests whose grammar objects are ready from grammar_queue to waiting_queue. - Rank i returns two sets ready_reqs_i, failed_reqs_i - ready_reqs_all = all_gather(ready_reqs_i) - failed_reqs_all = all_gather(failed_reqs_i) + For PP0, DP/TP group rank i returns two sets ready_reqs_i, + failed_reqs_i. ready_reqs_all = all_gather(ready_reqs_i) within + PP0's DP/TP group. failed_reqs_all = all_gather(failed_reqs_i) + within PP0's DP/TP group. ready_reqs = intersect(ready_reqs_all) failed_reqs = union(failed_reqs_all) + + PP0 then propagates the synced result to later PP ranks. Later PP + ranks receive and apply the propagated ready/failed decision. """ assert self.grammar_backend ready_req_idxs: set[int] = set() failed_req_idxs: set[int] = set() - # Poll for ready requests - start_time = time.perf_counter() - while time.perf_counter() - start_time < self.SGLANG_GRAMMAR_POLL_INTERVAL: + if self.pp_rank == 0: + # Poll for ready requests + start_time = time.perf_counter() + while time.perf_counter() - start_time < self.SGLANG_GRAMMAR_POLL_INTERVAL: + for i, req in enumerate(self.grammar_queue): + if i in ready_req_idxs: + continue + + if ( + req.finished() or req.grammar is None + ): # It is aborted by AbortReq + ready_req_idxs.add(i) + continue + + assert isinstance(req.grammar, futures.Future), f"{req=}" + if req.grammar.done(): + ready_req_idxs.add(i) + + if len(ready_req_idxs) == len(self.grammar_queue): + break + + # Sleep a bit to avoid busy waiting + time.sleep(self.SGLANG_GRAMMAR_POLL_INTERVAL / 10) + + # Check failed requests for i, req in enumerate(self.grammar_queue): - if i in ready_req_idxs: - continue + if i not in ready_req_idxs: + # grammar_wait_ct is only updated on PP0; later PP ranks + # receive PP0's ready/failed decision through PP sync. + self.grammar_queue[i].grammar_wait_ct += 1 + if ( + self.grammar_queue[i].grammar_wait_ct + >= self.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS + ): + # Timeout after max poll iterations + # The actual waiting time is SGLANG_GRAMMAR_MAX_POLL_ITERATIONS * max(SGLANG_GRAMMAR_POLL_INTERVAL, GPU_forward_batch_latency) + failed_req_idxs.add(i) - if req.finished() or req.grammar is None: # It is aborted by AbortReq - ready_req_idxs.add(i) - continue - - assert isinstance(req.grammar, futures.Future), f"{req=}" - if req.grammar.done(): - ready_req_idxs.add(i) - - # Sleep a bit to avoid busy waiting - time.sleep(self.SGLANG_GRAMMAR_POLL_INTERVAL / 10) - - # Check failed requests - for i, req in enumerate(self.grammar_queue): - if i not in ready_req_idxs: - self.grammar_queue[i].grammar_wait_ct += 1 - if ( - self.grammar_queue[i].grammar_wait_ct - >= self.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS - ): - # Timeout after max poll iterations - # The actual waiting time is SGLANG_GRAMMAR_MAX_POLL_ITERATIONS * max(SGLANG_GRAMMAR_POLL_INTERVAL, GPU_forward_batch_latency) - failed_req_idxs.add(i) - - # Sync ready and failed requests across all ranks - if self.grammar_sync_size == 1: + # Sync ready and failed requests across all TP ranks in PP0. + if self.grammar_sync_size == 1: + synced_ready_req_idxs = ready_req_idxs + synced_failed_req_idxs = failed_req_idxs + else: + all_gather_output = [None] * self.grammar_sync_size + torch.distributed.all_gather_object( + all_gather_output, + (ready_req_idxs, failed_req_idxs), + group=self.grammar_sync_group, + ) + synced_ready_req_idxs = set.intersection( + *[x[0] for x in all_gather_output] + ) + synced_failed_req_idxs = set.union(*[x[1] for x in all_gather_output]) + else: synced_ready_req_idxs = ready_req_idxs synced_failed_req_idxs = failed_req_idxs - else: - all_gather_output = [None] * self.grammar_sync_size - torch.distributed.all_gather_object( - all_gather_output, - (ready_req_idxs, failed_req_idxs), - group=self.grammar_sync_group, - ) - synced_ready_req_idxs = set.intersection(*[x[0] for x in all_gather_output]) - synced_failed_req_idxs = set.union(*[x[1] for x in all_gather_output]) + + # Propagate PP0's grammar queue decision to later PP ranks. + ( + synced_ready_req_idxs, + synced_failed_req_idxs, + ) = self._pp_sync_ready_failed( + synced_ready_req_idxs, + synced_failed_req_idxs, + ) # Return ready requests return_reqs: List[Req] = [] diff --git a/python/sglang/srt/distributed/communication_tags.py b/python/sglang/srt/distributed/communication_tags.py new file mode 100644 index 000000000..9b8cf90e0 --- /dev/null +++ b/python/sglang/srt/distributed/communication_tags.py @@ -0,0 +1,15 @@ +from enum import IntEnum, unique + + +@unique +class P2PTag(IntEnum): + """ + Tags reserved for point-to-point communication protocols. + + Communications introduced outside existing scheduler loops need explicit + tags to avoid being consumed by unrelated send/recv paths. + """ + + DEFAULT = 0 + HIRADIX_PP_SYNC = int.from_bytes(b"PpHi", byteorder="big") + GRAMMAR_PP_SYNC = int.from_bytes(b"PpGr", byteorder="big") diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 72a87476d..d5ef2a767 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -1284,6 +1284,7 @@ class GroupCoordinator: obj: Any, dst: int, async_send: bool = False, + tag: int = 0, ) -> List[P2PWork]: """ Send the input object list to the destination rank. @@ -1314,6 +1315,7 @@ class GroupCoordinator: size_tensor, self.ranks[dst], group=self.cpu_group, + tag=tag, ) if async_send: p2p_work.append(P2PWork(size_work, size_tensor)) @@ -1322,6 +1324,7 @@ class GroupCoordinator: object_tensor, self.ranks[dst], group=self.cpu_group, + tag=tag, ) if async_send: p2p_work.append(P2PWork(object_work, object_tensor)) @@ -1331,6 +1334,7 @@ class GroupCoordinator: def recv_object( self, src: int, + tag: int = 0, ) -> Any: """Receive the input object list from the source rank.""" """NOTE: `src` is the local rank of the source rank.""" @@ -1345,7 +1349,7 @@ class GroupCoordinator: # Receive object size # We have to use irecv here to make it work for both isend and send. work = torch.distributed.irecv( - size_tensor, src=self.ranks[src], group=self.cpu_group + size_tensor, src=self.ranks[src], group=self.cpu_group, tag=tag ) work.wait() @@ -1357,7 +1361,7 @@ class GroupCoordinator: ) work = torch.distributed.irecv( - object_tensor, src=self.ranks[src], group=self.cpu_group + object_tensor, src=self.ranks[src], group=self.cpu_group, tag=tag ) work.wait() diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index bfbe91ffe..733479550 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Tuple import torch from sglang.srt.disaggregation.kv_events import StorageMedium +from sglang.srt.distributed.communication_tags import P2PTag from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, @@ -281,14 +282,20 @@ class HiRadixCache(RadixCache): return if self.pp_rank > 0: torch.distributed.recv( - data, group_src=self.pp_rank - 1, group=self.pp_group, tag=2 + data, + group_src=self.pp_rank - 1, + group=self.pp_group, + tag=P2PTag.HIRADIX_PP_SYNC, ) if self.pp_rank + 1 < self.pp_size: # Make a copy of data, so that the caller is safe to modify `data` after this call. # This is cheap, as _pp_sync is not to be used for transmitting large data. copy_of_data = data.clone() send_work = torch.distributed.isend( - copy_of_data, group_dst=self.pp_rank + 1, group=self.pp_group, tag=2 + copy_of_data, + group_dst=self.pp_rank + 1, + group=self.pp_group, + tag=P2PTag.HIRADIX_PP_SYNC, ) self.work_list.append(send_work) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 0384627e5..9bb07beb4 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Iterator, NamedTuple, Optional, TypeVar import torch from sglang.srt.disaggregation.kv_events import StorageMedium +from sglang.srt.distributed.communication_tags import P2PTag from sglang.srt.environ import envs from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, @@ -428,12 +429,18 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): return if self.pp_rank > 0: torch.distributed.recv( - data, group_src=self.pp_rank - 1, group=self.pp_group, tag=2 + data, + group_src=self.pp_rank - 1, + group=self.pp_group, + tag=P2PTag.HIRADIX_PP_SYNC, ) if self.pp_rank + 1 < self.pp_size: copy_of_data = data.clone() send_work = torch.distributed.isend( - copy_of_data, group_dst=self.pp_rank + 1, group=self.pp_group, tag=2 + copy_of_data, + group_dst=self.pp_rank + 1, + group=self.pp_group, + tag=P2PTag.HIRADIX_PP_SYNC, ) self.work_list.append(send_work) diff --git a/test/registered/unit/constrained/test_grammar_manager.py b/test/registered/unit/constrained/test_grammar_manager.py index fbc59457b..cb1d793a6 100644 --- a/test/registered/unit/constrained/test_grammar_manager.py +++ b/test/registered/unit/constrained/test_grammar_manager.py @@ -25,6 +25,7 @@ from sglang.srt.constrained.base_grammar_backend import ( ) from sglang.srt.constrained.grammar_manager import GrammarManager from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject +from sglang.srt.distributed.communication_tags import P2PTag from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(2.0, "base-a-test-cpu") @@ -45,6 +46,9 @@ def _make_scheduler(grammar_backend_name="none", skip_tokenizer=False): scheduler.dp_tp_group.world_size = 1 scheduler.dp_tp_group.first_rank = 0 scheduler.dp_tp_group.is_first_rank = True + scheduler.ps.pp_rank = 0 + scheduler.ps.pp_size = 1 + scheduler.pp_group = None return scheduler @@ -627,6 +631,92 @@ class TestGetReadyGrammarRequests(unittest.TestCase): self.assertEqual(len(mgr.grammar_queue), 0) +class _FakePPSendWork: + def __init__(self): + self.waited = False + self.work = self + + def wait(self): + self.waited = True + + +class _FakePPGroup: + def __init__(self, recv_data=None): + self.recv_data = recv_data + self.recv_calls = [] + self.send_calls = [] + + def recv_object(self, *, src, tag): + self.recv_calls.append((src, tag)) + return self.recv_data + + def send_object(self, data, *, dst, async_send, tag): + self.send_calls.append((data, dst, async_send, tag)) + return [_FakePPSendWork()] + + +class TestGrammarManagerPPSync(unittest.TestCase): + """Test PP synchronization of grammar ready/failed indexes.""" + + def _make_mgr_for_pp(self, pp_rank, pp_size, pp_group): + scheduler = _make_scheduler() + scheduler.server_args.skip_tokenizer_init = True + scheduler.ps.pp_rank = pp_rank + scheduler.ps.pp_size = pp_size + scheduler.pp_group = pp_group + mgr = GrammarManager(scheduler) + mgr.grammar_backend = MagicMock(spec=BaseGrammarBackend) + return mgr + + def test_pp0_sends_ready_failed_without_recv(self): + pp_group = _FakePPGroup() + mgr = self._make_mgr_for_pp(pp_rank=0, pp_size=3, pp_group=pp_group) + + data = mgr._pp_sync_ready_failed({1}, {3}) + + self.assertEqual(data, ({1}, {3})) + self.assertEqual(pp_group.recv_calls, []) + self.assertEqual( + pp_group.send_calls, + [(({1}, {3}), 1, True, P2PTag.GRAMMAR_PP_SYNC)], + ) + + def test_middle_pp_rank_receives_and_forwards_pp0_result(self): + pp0_data = ({1, 2}, {4}) + pp_group = _FakePPGroup(recv_data=pp0_data) + mgr = self._make_mgr_for_pp(pp_rank=1, pp_size=3, pp_group=pp_group) + + data = mgr._pp_sync_ready_failed(set(), set()) + + self.assertEqual(data, pp0_data) + self.assertEqual(pp_group.recv_calls, [(0, P2PTag.GRAMMAR_PP_SYNC)]) + self.assertEqual( + pp_group.send_calls, + [(pp0_data, 2, True, P2PTag.GRAMMAR_PP_SYNC)], + ) + + def test_last_pp_rank_receives_without_forwarding(self): + pp0_data = ({0}, {2}) + pp_group = _FakePPGroup(recv_data=pp0_data) + mgr = self._make_mgr_for_pp(pp_rank=2, pp_size=3, pp_group=pp_group) + + data = mgr._pp_sync_ready_failed(set(), set()) + + self.assertEqual(data, pp0_data) + self.assertEqual(pp_group.recv_calls, [(1, P2PTag.GRAMMAR_PP_SYNC)]) + self.assertEqual(pp_group.send_calls, []) + + def test_pp_sync_drains_previous_async_send_work(self): + pp_group = _FakePPGroup() + mgr = self._make_mgr_for_pp(pp_rank=0, pp_size=2, pp_group=pp_group) + work = _FakePPSendWork() + mgr.grammar_pp_sync_work_list = [work] + + mgr._pp_sync_ready_failed({1}, set()) + + self.assertTrue(work.waited) + + class TestStrictReasoningPaths(unittest.TestCase): """Test _enable_strict_thinking code paths in GrammarManager."""