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

Fix missed hisparse release and stale field cleanup in pause retract (#30674)

Browse Source
This commit is contained in:
fzyzcjy
2026-07-15 14:31:11 +08:00
committed by GitHub
parent b6cc897fea
commit 1967b9ec99
2 changed files with 98 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+42 -26
View File
@@ -4106,33 +4106,49 @@ class Scheduler(
if not self.running_batch.is_empty():
self.running_batch.filter_batch()
if len(self.running_batch.reqs) != 0:
# Decode-side retract always rebootstraps (recomputes the KV from
# the prefill), so skip the device->host KV offload that release_req
# would otherwise do; the offloaded copy would be immediately
# discarded. Non-decode modes ignore offload_kv (they never offload).
retracted_reqs = self.running_batch.reqs
retract_all(
reqs=retracted_reqs,
server_args=self.server_args,
req_to_token_pool=self.running_batch.req_to_token_pool,
token_to_kv_pool_allocator=self.running_batch.token_to_kv_pool_allocator,
tree_cache=self.running_batch.tree_cache,
hisparse_coordinator=self.running_batch.hisparse_coordinator,
offload_kv=False,
)
self.running_batch.reqs = []
for req in retracted_reqs:
if self.disaggregation_mode == DisaggregationMode.DECODE:
if req.output_ids:
req.pd_rebootstrap_forced_output_id = req.output_ids.pop()
req.pd_rebootstrap_in_progress = True
req.time_stats.set_retract_time()
self.disagg_decode_prealloc_queue.hold_rebootstrap(req)
else:
self._add_request_to_queue(req)
self.running_batch.batch_is_full = False
retracted_reqs = list(self.running_batch.reqs)
if (
self.chunked_req is not None
and not self.chunked_req.finished()
and self.chunked_req not in retracted_reqs
and self.disaggregation_mode != DisaggregationMode.PREFILL
):
retracted_reqs.append(self.chunked_req)
if retracted_reqs:
# Decode-side retract always rebootstraps (recomputes the KV from
# the prefill), so skip the device->host KV offload that release_req
# would otherwise do; the offloaded copy would be immediately
# discarded. Non-decode modes ignore offload_kv (they never offload).
retract_all(
reqs=retracted_reqs,
server_args=self.server_args,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
tree_cache=self.tree_cache,
hisparse_coordinator=self.hisparse_coordinator,
offload_kv=False,
)
self.running_batch.reqs = []
for req in retracted_reqs:
if self.disaggregation_mode == DisaggregationMode.DECODE:
if req.output_ids:
req.pd_rebootstrap_forced_output_id = req.output_ids.pop()
req.pd_rebootstrap_in_progress = True
req.time_stats.set_retract_time()
self.disagg_decode_prealloc_queue.hold_rebootstrap(req)
else:
self._add_request_to_queue(req)
self.running_batch.batch_is_full = False
# In disagg-PREFILL, keep a live mid-chunk chunked_req rather than retract it:
# freeing its KV under a live disagg KV-sender crashes pop_bootstrapped or
# sends freed/reused KV to decode. Kept, it resumes prefill after the pause.
# TODO(disagg-prefill-retract): tear the sender down (abort + release metadata
# buffer + reset pending_bootstrap) before freeing KV, then retract for real.
# Until then a weight-update pause leaves stale-weight prefix KV (off-policy).
if self.disaggregation_mode != DisaggregationMode.PREFILL:
self.chunked_req = None
# Surface the paused state to dashboards immediately. The scheduler