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

feat: return partial generation results when aborting requests in waiting queue (#11673)

Browse Source
This commit is contained in:
Yuhong Guo
2025-10-29 22:03:00 +08:00
committed by GitHub
parent 750940ae36
commit caa5d2967c
6 changed files with 168 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/sglang/srt/environ.py
View File

@@ -135,6 +135,7 @@ class Envs:
# Scheduler: memory leak test
SGLANG_TEST_RETRACT = EnvBool(False)
SGLANG_TEST_RETRACT_INTERVAL = EnvInt(3)
SGLANG_TEST_RETRACT_NO_PREFILL_BS = EnvInt(2 ** 31)
SGLANG_ENABLE_RUNTIME_MEM_LEAK_CHECK = EnvBool(False)
# Scheduler: new token ratio hyperparameters

7
python/sglang/srt/managers/scheduler.py
View File

@@ -198,6 +198,7 @@ logger = logging.getLogger(__name__)
# Test retract decode for debugging purposes
TEST_RETRACT = envs.SGLANG_TEST_RETRACT.get()
TEST_RETRACT_INTERVAL = envs.SGLANG_TEST_RETRACT_INTERVAL.get()
TEST_RETRACT_NO_PREFILL_BS = envs.SGLANG_TEST_RETRACT_NO_PREFILL_BS.get()
GRAMMAR_TIMEOUT = float(os.environ.get("SGLANG_GRAMMAR_TIMEOUT", 300))
@@ -1657,6 +1658,12 @@ class Scheduler(
# Get priority queue
self.policy.calc_priority(self.waiting_queue)
if TEST_RETRACT and running_bs > TEST_RETRACT_NO_PREFILL_BS:
# If we are testing retraction and the running batch size exceeds
# TEST_RETRACT_NO_PREFILL_BS, we skip the prefill to keep the requests
# in the waiting queue.
return None
# Prefill policy
adder = PrefillAdder(
self.page_size,

130
python/sglang/srt/managers/tokenizer_manager.py
View File

@@ -1487,6 +1487,51 @@ class TokenizerManager(TokenizerCommunicatorMixin):
if self.crash_dump_folder and state.finished and state.obj.log_metrics:
self.record_request_for_crash_dump(state, out_dict)
def add_logprob_to_meta_info(
self,
meta_info: dict,
state: ReqState,
top_logprobs_num: int,
token_ids_logprob: List[int],
return_text_in_logprobs: bool,
):
meta_info["input_token_logprobs"] = self.detokenize_logprob_tokens(
state.input_token_logprobs_val,
state.input_token_logprobs_idx,
return_text_in_logprobs,
)
meta_info["output_token_logprobs"] = self.detokenize_logprob_tokens(
state.output_token_logprobs_val,
state.output_token_logprobs_idx,
return_text_in_logprobs,
)
if top_logprobs_num > 0:
meta_info["input_top_logprobs"] = self.detokenize_top_logprobs_tokens(
state.input_top_logprobs_val,
state.input_top_logprobs_idx,
return_text_in_logprobs,
)
meta_info["output_top_logprobs"] = self.detokenize_top_logprobs_tokens(
state.output_top_logprobs_val,
state.output_top_logprobs_idx,
return_text_in_logprobs,
)
if token_ids_logprob is not None:
meta_info["input_token_ids_logprobs"] = self.detokenize_top_logprobs_tokens(
state.input_token_ids_logprobs_val,
state.input_token_ids_logprobs_idx,
return_text_in_logprobs,
)
meta_info["output_token_ids_logprobs"] = (
self.detokenize_top_logprobs_tokens(
state.output_token_ids_logprobs_val,
state.output_token_ids_logprobs_idx,
return_text_in_logprobs,
)
)
def convert_logprob_style(
self,
meta_info: dict,
@@ -1513,16 +1558,6 @@ class TokenizerManager(TokenizerCommunicatorMixin):
state.output_token_logprobs_idx.extend(
recv_obj.output_token_logprobs_idx[recv_obj_index]
)
meta_info["input_token_logprobs"] = self.detokenize_logprob_tokens(
state.input_token_logprobs_val,
state.input_token_logprobs_idx,
return_text_in_logprobs,
)
meta_info["output_token_logprobs"] = self.detokenize_logprob_tokens(
state.output_token_logprobs_val,
state.output_token_logprobs_idx,
return_text_in_logprobs,
)
if top_logprobs_num > 0:
if len(recv_obj.input_top_logprobs_val) > 0:
@@ -1538,16 +1573,6 @@ class TokenizerManager(TokenizerCommunicatorMixin):
state.output_top_logprobs_idx.extend(
recv_obj.output_top_logprobs_idx[recv_obj_index]
)
meta_info["input_top_logprobs"] = self.detokenize_top_logprobs_tokens(
state.input_top_logprobs_val,
state.input_top_logprobs_idx,
return_text_in_logprobs,
)
meta_info["output_top_logprobs"] = self.detokenize_top_logprobs_tokens(
state.output_top_logprobs_val,
state.output_top_logprobs_idx,
return_text_in_logprobs,
)
if token_ids_logprob is not None:
if len(recv_obj.input_token_ids_logprobs_val) > 0:
@@ -1563,18 +1588,14 @@ class TokenizerManager(TokenizerCommunicatorMixin):
state.output_token_ids_logprobs_idx.extend(
recv_obj.output_token_ids_logprobs_idx[recv_obj_index]
)
meta_info["input_token_ids_logprobs"] = self.detokenize_top_logprobs_tokens(
state.input_token_ids_logprobs_val,
state.input_token_ids_logprobs_idx,
return_text_in_logprobs,
)
meta_info["output_token_ids_logprobs"] = (
self.detokenize_top_logprobs_tokens(
state.output_token_ids_logprobs_val,
state.output_token_ids_logprobs_idx,
return_text_in_logprobs,
)
)
self.add_logprob_to_meta_info(
meta_info,
state,
state.obj.top_logprobs_num,
state.obj.token_ids_logprob,
return_text_in_logprobs,
)
def detokenize_logprob_tokens(
self,
@@ -1759,26 +1780,33 @@ class TokenizerManager(TokenizerCommunicatorMixin):
return
state = self.rid_to_state[recv_obj.rid]
state.finished = True
abort_message = recv_obj.abort_reason or "Abort in waiting queue"
finish_reason = {
"type": "abort",
"message": abort_message,
}
if recv_obj.finished_reason:
out = {
"meta_info": {
"id": recv_obj.rid,
"finish_reason": recv_obj.finished_reason,
},
}
else:
out = {
"text": "",
"meta_info": {
"id": recv_obj.rid,
"finish_reason": {
"type": "abort",
"message": "Abort before prefill",
},
"prompt_tokens": 0,
"completion_tokens": 0,
},
}
finish_reason = recv_obj.finished_reason
meta_info = {"id": recv_obj.rid, "finish_reason": finish_reason}
is_stream = getattr(state.obj, "stream", False)
if getattr(state.obj, "return_logprob", False):
self.add_logprob_to_meta_info(
meta_info,
state,
state.obj.top_logprobs_num,
state.obj.token_ids_logprob,
state.obj.return_text_in_logprobs
and not self.server_args.skip_tokenizer_init,
)
output_ids = state.output_ids
meta_info["completion_tokens"] = len(output_ids)
out = {
"text": state.text,
"output_ids": [output_ids[-1]] if is_stream else output_ids,
"meta_info": meta_info,
}
state.out_list.append(out)
state.event.set()