EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Fix stuck in get_new_prefill_batch (#948)

Browse Source
This commit is contained in:
Liangsheng Yin
2024-08-06 01:05:58 -07:00
committed by GitHub
parent ffb15744b5
commit ad56e68495
1 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/sglang/srt/managers/tp_worker.py
View File

@@ -364,12 +364,13 @@ class ModelTpServer:
# Compute matched prefix length # Compute matched prefix length
for req in self.waiting_queue: for req in self.waiting_queue:
req.input_ids = req.origin_input_ids + req.output_ids req.input_ids = req.origin_input_ids + req.output_ids
prefix_indices, last_node = self.tree_cache.match_prefix( try_match_ids = req.input_ids
rid=req.rid,
key=req.input_ids,
)
if req.return_logprob: if req.return_logprob:
prefix_indices = prefix_indices[: req.logprob_start_len] try_match_ids = req.input_ids[: req.logprob_start_len]
# NOTE: the prefix_indices must always be aligned with last_node
prefix_indices, last_node = self.tree_cache.match_prefix(
rid=req.rid, key=try_match_ids
)
req.extend_input_len = len(req.input_ids) - len(prefix_indices) req.extend_input_len = len(req.input_ids) - len(prefix_indices)
req.prefix_indices = prefix_indices req.prefix_indices = prefix_indices
req.last_node = last_node req.last_node = last_node