minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Fix: add grammar sync in PP for structured output (#30747)

Browse Source 
Signed-off-by: Jing Wang <jingwang96@qq.com>
Co-authored-by: ziang663 <119752791+ziang663@users.noreply.github.com>
Co-authored-by: Chao Shi <chao.shi@alibaba-inc.com>
This commit is contained in:
Jing Wang
2026-07-12 02:54:50 +08:00
committed by GitHub
co-authored by ziang663 Chao Shi
parent 9b4bb415dd
commit ed554aac17
6 changed files with 239 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/constrained/grammar_manager.py
+110 -42
View File
@@ -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] = []
python/sglang/srt/distributed/communication_tags.py
+15
View File
@@ -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")
python/sglang/srt/distributed/parallel_state.py
+6 -2
View File
@@ -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()
python/sglang/srt/mem_cache/hiradix_cache.py
+9 -2
View File
@@ -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)
python/sglang/srt/mem_cache/unified_radix_cache.py
+9 -2
View File
@@ -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)