Fix streaming logprobs corruption caused by shared mutable list reference (#21030)

This commit is contained in:
Lianmin Zheng
2026-03-21 00:18:48 -07:00
committed by GitHub
parent d089db0563
commit dba6fb3d30
4 changed files with 36 additions and 48 deletions
@@ -650,22 +650,19 @@ class OpenAIServingChat(OpenAIServingBase):
routed_experts[index] = content["meta_info"].get("routed_experts", None) routed_experts[index] = content["meta_info"].get("routed_experts", None)
# Handle logprobs # Handle logprobs
finish_reason = content["meta_info"].get("finish_reason", None)
choice_logprobs = None choice_logprobs = None
if request.logprobs: if request.logprobs:
n_prev_token = n_prev_tokens.get(index, 0) n_prev_token = n_prev_tokens.get(index, 0)
total_output_logprobs = len( total_output_logprobs = content["meta_info"][
content["meta_info"]["output_token_logprobs"] "output_token_logprobs_length"
) ]
# When finish_reason is set and all logprobs have been sent, if n_prev_token < total_output_logprobs:
# 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:
choice_logprobs = self._process_streaming_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 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 finish_reason_type = finish_reason["type"] if finish_reason else None
# Track finish_reason for each index # Track finish_reason for each index
@@ -1174,15 +1171,18 @@ class OpenAIServingChat(OpenAIServingBase):
return ToolCallProcessingResult(None, text, finish_reason) return ToolCallProcessingResult(None, text, finish_reason)
def _process_streaming_logprobs( 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: ) -> ChoiceLogprobs:
"""Process logprobs for streaming response""" """Process logprobs for streaming response"""
logprobs = to_openai_style_logprobs( logprobs = to_openai_style_logprobs(
output_token_logprobs=content["meta_info"]["output_token_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", [])[ output_top_logprobs=content["meta_info"].get("output_top_logprobs", [])[
n_prev_token: n_prev_token:total_output_logprobs
], ],
) )
@@ -244,32 +244,22 @@ class OpenAIServingCompletion(OpenAIServingBase):
input_top_logprobs = None input_top_logprobs = None
n_prev_token = n_prev_tokens.get(index, 0) n_prev_token = n_prev_tokens.get(index, 0)
total_output_logprobs = len( total_output_logprobs = content["meta_info"][
content["meta_info"]["output_token_logprobs"] "output_token_logprobs_length"
) ]
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.
if ( if (
len(output_logprobs_slice) == 0 n_prev_token < total_output_logprobs
and finish_reason_for_logprobs is not None or input_token_logprobs is not None
and input_token_logprobs is None
): ):
logprobs = None
else:
logprobs = to_openai_style_logprobs( logprobs = to_openai_style_logprobs(
input_token_logprobs=input_token_logprobs, input_token_logprobs=input_token_logprobs,
input_top_logprobs=input_top_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=content["meta_info"].get(
"output_top_logprobs", [] "output_top_logprobs", []
)[n_prev_token:], )[n_prev_token:total_output_logprobs],
) )
n_prev_tokens[index] = total_output_logprobs n_prev_tokens[index] = total_output_logprobs
@@ -1719,6 +1719,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
meta_info["input_token_logprobs"] = state.input_token_logprobs meta_info["input_token_logprobs"] = state.input_token_logprobs
meta_info["output_token_logprobs"] = state.output_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 # 2. Handle top logprobs
if top_logprobs_num > 0: if top_logprobs_num > 0:
@@ -160,23 +160,20 @@ class TestOpenAIServer(CustomTestCase):
is_first = is_firsts.get(index, True) is_first = is_firsts.get(index, True)
if logprobs: if logprobs:
# When finish_reason is set, logprobs may be None if this chunk assert response.choices[0].logprobs, f"no logprobs in response"
# only contains buffered text being flushed (no new tokens generated). assert isinstance(
# The detokenizer holds back text at word boundaries during streaming. response.choices[0].logprobs.tokens[0], str
if response.choices[0].logprobs is not None: ), f"{response.choices[0].logprobs.tokens[0]} is not a string"
if not (is_first and echo):
assert isinstance( assert isinstance(
response.choices[0].logprobs.tokens[0], str response.choices[0].logprobs.top_logprobs[0], dict
), f"{response.choices[0].logprobs.tokens[0]} is not a string" ), f"top_logprobs was not a dictionary"
if not (is_first and echo): ret_num_top_logprobs = len(
assert isinstance( response.choices[0].logprobs.top_logprobs[0]
response.choices[0].logprobs.top_logprobs[0], dict )
), f"top_logprobs was not a dictionary" # 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
ret_num_top_logprobs = len( # assert ret_num_top_logprobs == logprobs, f"{ret_num_top_logprobs} vs {logprobs}"
response.choices[0].logprobs.top_logprobs[0] assert ret_num_top_logprobs > 0, f"ret_num_top_logprobs was 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 is_first:
if echo: if echo: