[serving] replace O(n²) stream_buffer string concat with integer offset (#22606)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Alex Nails
2026-04-14 14:48:44 -07:00
committed by GitHub
co-authored by Claude Opus 4.6 gemini-code-assist[bot]
parent 36891ab514
commit 8092431316
2 changed files with 11 additions and 10 deletions
@@ -649,7 +649,7 @@ class OpenAIServingChat(OpenAIServingBase):
# State tracking for streaming # State tracking for streaming
is_firsts = {} is_firsts = {}
stream_buffers = {} stream_offsets = {}
n_prev_tokens = {} n_prev_tokens = {}
has_tool_calls = {} has_tool_calls = {}
finish_reasons = {} finish_reasons = {}
@@ -737,13 +737,13 @@ class OpenAIServingChat(OpenAIServingBase):
yield f"data: {chunk.model_dump_json()}\n\n" yield f"data: {chunk.model_dump_json()}\n\n"
stream_started = True stream_started = True
stream_buffer = stream_buffers.get(index, "") offset = stream_offsets.get(index, 0)
if self.tokenizer_manager.server_args.incremental_streaming_output: if self.tokenizer_manager.server_args.incremental_streaming_output:
# content["text"] is already the incremental delta # content["text"] is already the incremental delta
delta = content["text"] delta = content["text"]
else: else:
delta = content["text"][len(stream_buffer) :] delta = content["text"][offset:]
stream_buffers[index] = stream_buffer + delta stream_offsets[index] = len(content["text"])
# Handle reasoning content # Handle reasoning content
if self.reasoning_parser and request.separate_reasoning: if self.reasoning_parser and request.separate_reasoning:
@@ -214,7 +214,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
created = int(time.time()) created = int(time.time())
# State tracking for streaming # State tracking for streaming
stream_buffers = {} stream_offsets = {}
n_prev_tokens = {} n_prev_tokens = {}
# Usage tracking # Usage tracking
@@ -249,9 +249,10 @@ class OpenAIServingCompletion(OpenAIServingBase):
hidden_states[index] = content["meta_info"].get("hidden_states", None) hidden_states[index] = content["meta_info"].get("hidden_states", None)
routed_experts[index] = content["meta_info"].get("routed_experts", None) routed_experts[index] = content["meta_info"].get("routed_experts", None)
stream_buffer = stream_buffers.get(index, "") is_first_chunk = index not in stream_offsets
offset = stream_offsets.get(index, 0)
# Handle echo for first chunk # Handle echo for first chunk
if not stream_buffer: # The first chunk if is_first_chunk: # The first chunk
if request.echo: if request.echo:
echo_text = self._get_echo_text(request, index) echo_text = self._get_echo_text(request, index)
text = echo_text + text text = echo_text + text
@@ -260,7 +261,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
logprobs = None logprobs = None
if request.logprobs is not None: if request.logprobs is not None:
# The first chunk and echo is enabled. # The first chunk and echo is enabled.
if not stream_buffer and request.echo: if is_first_chunk and request.echo:
input_token_logprobs = content["meta_info"][ input_token_logprobs = content["meta_info"][
"input_token_logprobs" "input_token_logprobs"
] ]
@@ -301,8 +302,8 @@ class OpenAIServingCompletion(OpenAIServingBase):
n_prev_tokens[index] = total_output_logprobs n_prev_tokens[index] = total_output_logprobs
# Generate delta # Generate delta
delta = text[len(stream_buffer) :] delta = text[offset:]
stream_buffers[index] = stream_buffer + delta stream_offsets[index] = len(content["text"])
finish_reason = content["meta_info"].get("finish_reason", None) 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