chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Remove overlap thread (#11210)

Browse Source 
Co-authored-by: Lianmin Zheng <15100009+merrymercy@users.noreply.github.com>
Co-authored-by: Hanming Lu <69857889+hanming-lu@users.noreply.github.com>
This commit is contained in:
Liangsheng Yin
2025-10-07 20:12:12 +08:00
committed by GitHub
parent 24bc3fb0f9
commit 1519a89cfd
14 changed files with 280 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
python/sglang/srt/disaggregation/prefill.py
View File

@@ -321,6 +321,8 @@ class SchedulerDisaggregationPrefillMixin:
self.result_queue = deque()
while True:
self.launch_last_batch_sample_if_needed()
recv_reqs = self.recv_requests()
self.process_input_requests(recv_reqs)
self.waiting_queue.extend(
@@ -368,7 +370,6 @@ class SchedulerDisaggregationPrefillMixin:
self: Scheduler,
batch: ScheduleBatch,
result: GenerationBatchResult,
launch_done: Optional[threading.Event] = None,
) -> None:
"""
Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue
@@ -379,31 +380,30 @@ class SchedulerDisaggregationPrefillMixin:
next_token_ids,
extend_input_len_per_req,
extend_logprob_start_len_per_req,
copy_done,
) = (
result.logits_output,
result.next_token_ids,
result.extend_input_len_per_req,
result.extend_logprob_start_len_per_req,
result.copy_done,
)
if copy_done is not None:
copy_done.synchronize()
logprob_pt = 0
# Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue
if self.enable_overlap:
# wait
logits_output, next_token_ids, _ = self.tp_worker.resolve_last_batch_result(
launch_done
)
else:
next_token_ids = result.next_token_ids.tolist()
if batch.return_logprob:
if logits_output.next_token_logprobs is not None:
logits_output.next_token_logprobs = (
logits_output.next_token_logprobs.tolist()
)
if logits_output.input_token_logprobs is not None:
logits_output.input_token_logprobs = tuple(
logits_output.input_token_logprobs.tolist()
)
next_token_ids = result.next_token_ids.tolist()
if batch.return_logprob:
if logits_output.next_token_logprobs is not None:
logits_output.next_token_logprobs = (
logits_output.next_token_logprobs.tolist()
)
if logits_output.input_token_logprobs is not None:
logits_output.input_token_logprobs = tuple(
logits_output.input_token_logprobs.tolist()
)
hidden_state_offset = 0
for i, (req, next_token_id) in enumerate(