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

[PD+PP] Honor PP consensus for bootstrap and prealloc (#31869)

Browse Source 
Co-authored-by: Chao Shi <chao.shi@alibaba-inc.com>
This commit is contained in:
ziang663
2026-07-30 01:19:17 +08:00
committed by GitHub
co-authored by Chao Shi
parent 62d0f81f16
commit eefb434d17
7 changed files with 106 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+40 -10
View File
@@ -55,6 +55,7 @@ from sglang.srt.disaggregation.utils import (
is_dsv4_c128_online_enabled,
is_mla_backend,
poll_and_all_reduce,
poll_and_all_reduce_pp,
poll_and_all_reduce_with_staging,
prepare_abort,
setup_state_kv_args,
@@ -327,6 +328,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
self.bootstrap_port = bootstrap_port
self.max_total_num_tokens = max_total_num_tokens
self.pp_rank = pp_rank
self.pp_size = scheduler.ps.pp_size
self.num_reserved_decode_tokens = num_reserved_decode_tokens
self.transfer_backend = transfer_backend
# Queue for requests pending pre-allocation
@@ -727,23 +729,40 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
return resumed_reqs
def _update_handshake_waiters(
self, rids_to_check: Optional[List[str]] = None
self,
rids_to_check: Optional[List[str]] = None,
pp_good_rids: Optional[List[str]] = None,
pp_bad_rids: Optional[List[str]] = None,
) -> None:
if not self.queue:
return
# Still poll if any receiver was aborted, otherwise it stays stuck.
if all(decode_req.waiting_for_input for decode_req in self.queue) and not any(
decode_req.kv_receiver.conclude_state == KVPoll.Failed
for decode_req in self.queue
if (
self.pp_size <= 1
and all(decode_req.waiting_for_input for decode_req in self.queue)
and not any(
decode_req.kv_receiver.conclude_state == KVPoll.Failed
for decode_req in self.queue
)
):
return
polls = poll_and_all_reduce(
[decode_req.kv_receiver for decode_req in self.queue], self.gloo_group
)
if self.pp_size > 1:
polls = poll_and_all_reduce_pp(
(decode_req.req.rid for decode_req in self.queue),
KVPoll.WaitingForInput,
pp_good_rids,
pp_bad_rids,
)
else:
polls = poll_and_all_reduce(
[decode_req.kv_receiver for decode_req in self.queue], self.gloo_group
)
for i, (decode_req, poll) in enumerate(zip(self.queue, polls)):
for decode_req, poll in zip(self.queue, polls):
if poll is None:
continue
if rids_to_check is not None and decode_req.req.rid not in rids_to_check:
continue
@@ -864,11 +883,22 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
decode_req.kv_receiver.init(prefill_dp_rank)
def pop_preallocated(
self, rids_to_check: Optional[List[str]] = None
self,
rids_to_check: Optional[List[str]] = None,
pp_good_rids: Optional[List[str]] = None,
pp_bad_rids: Optional[List[str]] = None,
) -> Tuple[List[DecodeRequest], List[DecodeRequest]]:
"""Pop the preallocated requests from the pending queue (FIFO)."""
is_pp_mode = self.pp_size > 1
if is_pp_mode and (pp_good_rids is None or pp_bad_rids is None):
raise ValueError("PP consensus is required when pp_size > 1")
if is_pp_mode and rids_to_check is not None:
raise ValueError("rids_to_check cannot be used in PP mode")
self._resolve_pending_reqs()
self._update_handshake_waiters(rids_to_check)
self._update_handshake_waiters(rids_to_check, pp_good_rids, pp_bad_rids)
if is_pp_mode:
rids_to_check = pp_good_rids + pp_bad_rids
failed_reqs = []
preallocated_reqs = []
python/sglang/srt/disaggregation/prefill.py
+30 -16
View File
@@ -45,6 +45,7 @@ from sglang.srt.disaggregation.utils import (
is_dsv4_c128_online_enabled,
is_mla_backend,
poll_and_all_reduce_attn_cp_tp_group,
poll_and_all_reduce_pp,
prepare_abort,
setup_state_kv_args,
)
@@ -365,13 +366,15 @@ class PrefillBootstrapQueue:
def pop_bootstrapped(
self,
return_failed_reqs: bool = False,
rids_to_check: Optional[List[str]] = None,
) -> List[Req]:
pp_good_rids: Optional[List[str]] = None,
pp_bad_rids: Optional[List[str]] = None,
) -> List[Req] | tuple[List[Req], List[Req]]:
"""
pop the reqs which has finished bootstrapping
return_failed_reqs: For PP, on rank 0, also return the failed reqs to notify the next rank
rids_to_check: For PP, on rank > 0, check the rids from the previous rank has consensus with the current rank.
pp_good_rids: RIDs that PP consensus determined as WaitingForInput.
pp_bad_rids: RIDs that PP consensus determined as Failed.
"""
bootstrapped_reqs = []
@@ -384,21 +387,32 @@ class PrefillBootstrapQueue:
else:
return [], []
polls = poll_and_all_reduce_attn_cp_tp_group(
[req.disagg_kv_sender for req in self.queue],
self.scheduler.attn_cp_cpu_group,
self.scheduler.attn_tp_cpu_group,
)
if self.pp_size > 1:
polls = poll_and_all_reduce_pp(
(req.rid for req in self.queue),
KVPoll.WaitingForInput,
pp_good_rids,
pp_bad_rids,
)
uncovered = [i for i, poll in enumerate(polls) if poll is None]
if uncovered:
local_polls = poll_and_all_reduce_attn_cp_tp_group(
[self.queue[i].disagg_kv_sender for i in uncovered],
self.scheduler.attn_cp_cpu_group,
self.scheduler.attn_tp_cpu_group,
)
for i, local_poll in zip(uncovered, local_polls):
if local_poll == KVPoll.Failed:
polls[i] = KVPoll.Failed
else:
polls = poll_and_all_reduce_attn_cp_tp_group(
[req.disagg_kv_sender for req in self.queue],
self.scheduler.attn_cp_cpu_group,
self.scheduler.attn_tp_cpu_group,
)
for i, (req, poll) in enumerate(zip(self.queue, polls)):
if (
rids_to_check is not None
and req.rid not in rids_to_check
and poll != KVPoll.Failed
):
# In PP mode, successful bootstrap still requires cross-rank
# consensus. Local failures are terminal and must be drained
# even if an earlier PP rank has already removed the request.
if poll is None:
continue
if poll == KVPoll.Failed:
python/sglang/srt/disaggregation/utils.py
+28 -1
View File
@@ -5,7 +5,16 @@ import random
from collections import deque
from contextlib import nullcontext
from enum import Enum
from typing import TYPE_CHECKING, List, Literal, Optional, Tuple, Type, overload
from typing import (
TYPE_CHECKING,
Iterable,
List,
Literal,
Optional,
Tuple,
Type,
overload,
)
import numpy as np
import torch
@@ -39,6 +48,24 @@ FAKE_BOOTSTRAP_HOST = "2.2.2.2"
_IS_HIP = is_hip()
def poll_and_all_reduce_pp(
rids: Iterable[str],
ready_poll: int,
pp_good_rids: Optional[List[str]] = None,
pp_bad_rids: Optional[List[str]] = None,
) -> List[Optional[int]]:
"""Map authoritative PP consensus to poll states without polling again."""
if pp_good_rids is None or pp_bad_rids is None:
raise ValueError("PP consensus is required")
good_rids = set(pp_good_rids)
bad_rids = set(pp_bad_rids)
return [
KVPoll.Failed if rid in bad_rids else ready_poll if rid in good_rids else None
for rid in rids
]
def get_dsa_seed_metadata_dim(hf_config) -> int:
"""Return the model-defined PD seed width, independent of local spec mode."""
if not getattr(hf_config, "index_share_for_mtp_iteration", False):
python/sglang/srt/managers/scheduler_pp_mixin.py
+4 -4
View File
@@ -811,8 +811,8 @@ class SchedulerPPMixin:
good_reqs, failed_reqs = (
self.disagg_prefill_bootstrap_queue.pop_bootstrapped(
return_failed_reqs=True,
rids_to_check=good_consensus_bootstrapped_rids
+ bad_consensus_bootstrapped_rids,
pp_good_rids=good_consensus_bootstrapped_rids,
pp_bad_rids=bad_consensus_bootstrapped_rids,
)
)
self.waiting_queue.extend(good_reqs)
@@ -1417,8 +1417,8 @@ class SchedulerPPMixin:
bad_consensus_prealloc_rids,
) = prealloc_rids
good_reqs, failed_reqs = self.disagg_decode_prealloc_queue.pop_preallocated(
rids_to_check=good_consensus_prealloc_rids
+ bad_consensus_prealloc_rids,
pp_good_rids=good_consensus_prealloc_rids,
pp_bad_rids=bad_consensus_prealloc_rids,
)
self.disagg_decode_transfer_queue.extend(good_reqs)
return [