From dba6fb3d308d19e8068380fe34ce3292cd7a67ce Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Sat, 21 Mar 2026 00:18:48 -0700 Subject: [PATCH] Fix streaming logprobs corruption caused by shared mutable list reference (#21030) --- .../srt/entrypoints/openai/serving_chat.py | 26 ++++++++--------- .../entrypoints/openai/serving_completions.py | 28 ++++++------------ .../sglang/srt/managers/tokenizer_manager.py | 1 + .../openai_server/basic/test_openai_server.py | 29 +++++++++---------- 4 files changed, 36 insertions(+), 48 deletions(-) diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 1c8f4d94f..67cfef0a9 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -650,22 +650,19 @@ class OpenAIServingChat(OpenAIServingBase): routed_experts[index] = content["meta_info"].get("routed_experts", None) # Handle logprobs - finish_reason = content["meta_info"].get("finish_reason", None) choice_logprobs = None if request.logprobs: n_prev_token = n_prev_tokens.get(index, 0) - total_output_logprobs = len( - content["meta_info"]["output_token_logprobs"] - ) - # When finish_reason is set and all logprobs have been sent, - # any remaining text is just buffered text being flushed by the - # detokenizer (it holds back text at word boundaries). Return None - # for logprobs since no new tokens were generated for this text. - if n_prev_token < total_output_logprobs or finish_reason is None: + total_output_logprobs = content["meta_info"][ + "output_token_logprobs_length" + ] + if n_prev_token < total_output_logprobs: choice_logprobs = self._process_streaming_logprobs( - content, n_prev_token + content, n_prev_token, total_output_logprobs ) n_prev_tokens[index] = total_output_logprobs + + finish_reason = content["meta_info"].get("finish_reason", None) finish_reason_type = finish_reason["type"] if finish_reason else None # Track finish_reason for each index @@ -1174,15 +1171,18 @@ class OpenAIServingChat(OpenAIServingBase): return ToolCallProcessingResult(None, text, finish_reason) def _process_streaming_logprobs( - self, content: Dict[str, Any], n_prev_token: int + self, + content: Dict[str, Any], + n_prev_token: int, + total_output_logprobs: int, ) -> ChoiceLogprobs: """Process logprobs for streaming response""" logprobs = to_openai_style_logprobs( output_token_logprobs=content["meta_info"]["output_token_logprobs"][ - n_prev_token: + n_prev_token:total_output_logprobs ], output_top_logprobs=content["meta_info"].get("output_top_logprobs", [])[ - n_prev_token: + n_prev_token:total_output_logprobs ], ) diff --git a/python/sglang/srt/entrypoints/openai/serving_completions.py b/python/sglang/srt/entrypoints/openai/serving_completions.py index 2554fa738..0bda21907 100644 --- a/python/sglang/srt/entrypoints/openai/serving_completions.py +++ b/python/sglang/srt/entrypoints/openai/serving_completions.py @@ -244,32 +244,22 @@ class OpenAIServingCompletion(OpenAIServingBase): input_top_logprobs = None n_prev_token = n_prev_tokens.get(index, 0) - total_output_logprobs = len( - content["meta_info"]["output_token_logprobs"] - ) - output_logprobs_slice = content["meta_info"][ - "output_token_logprobs" - ][n_prev_token:] - finish_reason_for_logprobs = content["meta_info"]["finish_reason"] - - # When finish_reason is set and all logprobs have been sent, - # any remaining text is just buffered text being flushed by the - # detokenizer (it holds back text at word boundaries). Return None - # for logprobs since no new tokens were generated for this text. + total_output_logprobs = content["meta_info"][ + "output_token_logprobs_length" + ] if ( - len(output_logprobs_slice) == 0 - and finish_reason_for_logprobs is not None - and input_token_logprobs is None + n_prev_token < total_output_logprobs + or input_token_logprobs is not None ): - logprobs = None - else: logprobs = to_openai_style_logprobs( input_token_logprobs=input_token_logprobs, input_top_logprobs=input_top_logprobs, - output_token_logprobs=output_logprobs_slice, + output_token_logprobs=content["meta_info"][ + "output_token_logprobs" + ][n_prev_token:total_output_logprobs], output_top_logprobs=content["meta_info"].get( "output_top_logprobs", [] - )[n_prev_token:], + )[n_prev_token:total_output_logprobs], ) n_prev_tokens[index] = total_output_logprobs diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 1f55d2df9..d7a324afd 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1719,6 +1719,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi meta_info["input_token_logprobs"] = state.input_token_logprobs meta_info["output_token_logprobs"] = state.output_token_logprobs + meta_info["output_token_logprobs_length"] = len(state.output_token_logprobs) # 2. Handle top logprobs if top_logprobs_num > 0: diff --git a/test/registered/openai_server/basic/test_openai_server.py b/test/registered/openai_server/basic/test_openai_server.py index 17fc78764..cdc39f29a 100644 --- a/test/registered/openai_server/basic/test_openai_server.py +++ b/test/registered/openai_server/basic/test_openai_server.py @@ -160,23 +160,20 @@ class TestOpenAIServer(CustomTestCase): is_first = is_firsts.get(index, True) if logprobs: - # When finish_reason is set, logprobs may be None if this chunk - # only contains buffered text being flushed (no new tokens generated). - # The detokenizer holds back text at word boundaries during streaming. - if response.choices[0].logprobs is not None: + assert response.choices[0].logprobs, f"no logprobs in response" + assert isinstance( + response.choices[0].logprobs.tokens[0], str + ), f"{response.choices[0].logprobs.tokens[0]} is not a string" + if not (is_first and echo): assert isinstance( - response.choices[0].logprobs.tokens[0], str - ), f"{response.choices[0].logprobs.tokens[0]} is not a string" - if not (is_first and echo): - assert isinstance( - response.choices[0].logprobs.top_logprobs[0], dict - ), f"top_logprobs was not a dictionary" - ret_num_top_logprobs = len( - response.choices[0].logprobs.top_logprobs[0] - ) - # FIXME: Sometimes, some top_logprobs are missing in the return value. The reason is that some output id maps to the same output token and duplicate in the map - # assert ret_num_top_logprobs == logprobs, f"{ret_num_top_logprobs} vs {logprobs}" - assert ret_num_top_logprobs > 0, f"ret_num_top_logprobs was 0" + response.choices[0].logprobs.top_logprobs[0], dict + ), f"top_logprobs was not a dictionary" + ret_num_top_logprobs = len( + response.choices[0].logprobs.top_logprobs[0] + ) + # FIXME: Sometimes, some top_logprobs are missing in the return value. The reason is that some output id maps to the same output token and duplicate in the map + # assert ret_num_top_logprobs == logprobs, f"{ret_num_top_logprobs} vs {logprobs}" + assert ret_num_top_logprobs > 0, f"ret_num_top_logprobs was 0" if is_first: if echo: