304 lines
13 KiB
Diff
304 lines
13 KiB
Diff
From 5d8dd1f7f99f4c74bc302dd9ef6d32c00c41f7a8 Mon Sep 17 00:00:00 2001
|
|
From: ziang663 <1902123738@qq.com>
|
|
Date: Fri, 24 Jul 2026 03:15:10 +0800
|
|
Subject: [PATCH] fix(disagg): honor PP consensus for bootstrap and prealloc
|
|
|
|
---
|
|
python/sglang/srt/disaggregation/decode.py | 50 +++++++++++++++----
|
|
python/sglang/srt/disaggregation/prefill.py | 46 +++++++++++------
|
|
python/sglang/srt/disaggregation/utils.py | 29 ++++++++++-
|
|
.../sglang/srt/managers/scheduler_pp_mixin.py | 8 +--
|
|
.../test_decode_queue_cleanup.py | 2 +
|
|
...test_priority_scheduling_disaggregation.py | 1 +
|
|
.../mem_cache/test_decode_radix_lock_ref.py | 1 +
|
|
7 files changed, 106 insertions(+), 31 deletions(-)
|
|
|
|
diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py
|
|
index 5224920c4481..a664fbe85119 100644
|
|
--- a/python/sglang/srt/disaggregation/decode.py
|
|
+++ b/python/sglang/srt/disaggregation/decode.py
|
|
@@ -55,6 +55,7 @@
|
|
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,
|
|
@@ -325,6 +326,7 @@ def __init__(
|
|
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
|
|
@@ -719,23 +721,40 @@ def resume_retracted_reqs(
|
|
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
|
|
|
|
@@ -856,11 +875,22 @@ def _resolve_pending_reqs(self) -> None:
|
|
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 = []
|
|
diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py
|
|
index 1f4c947fa3fa..51d6e2dba7f2 100644
|
|
--- a/python/sglang/srt/disaggregation/prefill.py
|
|
+++ b/python/sglang/srt/disaggregation/prefill.py
|
|
@@ -45,6 +45,7 @@
|
|
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,
|
|
)
|
|
@@ -333,13 +334,15 @@ def _process_req(self, req: Req) -> None:
|
|
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 = []
|
|
@@ -352,21 +355,32 @@ def pop_bootstrapped(
|
|
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:
|
|
diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py
|
|
index 8d741804d875..7893d588248d 100644
|
|
--- a/python/sglang/srt/disaggregation/utils.py
|
|
+++ b/python/sglang/srt/disaggregation/utils.py
|
|
@@ -5,7 +5,16 @@
|
|
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
|
|
@@ -34,6 +43,24 @@
|
|
_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):
|
|
diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py
|
|
index 0a55601d4128..f3e1f76f4ad1 100644
|
|
--- a/python/sglang/srt/managers/scheduler_pp_mixin.py
|
|
+++ b/python/sglang/srt/managers/scheduler_pp_mixin.py
|
|
@@ -811,8 +811,8 @@ def process_bootstrapped_queue(
|
|
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 @@ def process_prealloc_queue(self: Scheduler, prealloc_rids: Optional[List[str]]):
|
|
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 [
|
|
diff --git a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py
|
|
index e136befc4d54..0a904792b5af 100644
|
|
--- a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py
|
|
+++ b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py
|
|
@@ -40,6 +40,7 @@ def test_prealloc_abort_clears_receiver_before_removing_request(self):
|
|
decode_req = SimpleNamespace(req=req, kv_receiver=receiver)
|
|
|
|
queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
|
|
+ queue.pp_size = 1
|
|
queue.queue = [decode_req]
|
|
queue.pending_reqs = []
|
|
queue.retracted_queue = []
|
|
@@ -86,6 +87,7 @@ def __eq__(self, other):
|
|
decode_req = SimpleNamespace(req=req, kv_receiver=receiver)
|
|
|
|
queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
|
|
+ queue.pp_size = 1
|
|
queue.queue = [decode_req]
|
|
queue.pending_reqs = [decode_req] # same object, dual ownership
|
|
queue.retracted_queue = []
|
|
diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py
|
|
index a5f052c7311e..762cb5489ccb 100644
|
|
--- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py
|
|
+++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py
|
|
@@ -119,6 +119,7 @@ def _new_decode_req(self, rid: str, priority: int, *, failed: bool = False):
|
|
|
|
def _new_queue(self, decode_reqs, *, low_priority_values_first: bool = False):
|
|
queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
|
|
+ queue.pp_size = 1
|
|
queue.queue = list(decode_reqs)
|
|
queue.pending_reqs = []
|
|
queue.retracted_queue = []
|
|
diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
|
|
index 7cba0af202b7..4abdfdc1b933 100644
|
|
--- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
|
|
+++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
|
|
@@ -294,6 +294,7 @@ def test_full_transfer_failure(self):
|
|
|
|
def test_pop_preallocated_rechecks_budget_after_lock(self):
|
|
queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
|
|
+ queue.pp_size = 1
|
|
|
|
req = MagicMock()
|
|
req.rid = "req-1"
|