Skip to content

fix(models): accumulate streamed logprobs in place to avoid O(n^2) copies#3762

Merged
seratch merged 1 commit into
openai:mainfrom
anxkhn:patch-6
Jul 9, 2026
Merged

fix(models): accumulate streamed logprobs in place to avoid O(n^2) copies#3762
seratch merged 1 commit into
openai:mainfrom
anxkhn:patch-6

Conversation

@anxkhn

@anxkhn anxkhn commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

When Chat Completions streaming is used with logprobs enabled, the stream handler
rebuilt the entire accumulated output-text logprobs list on every content delta:

existing_logprobs = state.text_content_index_and_output[1].logprobs or []
state.text_content_index_and_output[1].logprobs = existing_logprobs + output_logprobs

list + list allocates and copies a brand-new list of the full accumulated length
for each delta, so over an n-token stream this is O(n^2) allocation and copying
(the adjacent ... .text += delta.content on the line above already grows in
place). The accumulator was introduced in #2134 using + where an in-place extend
was intended.

This extends the existing list in place instead, keeping the None-provider guard
so non-OpenAI providers (for example LiteLLM), which may leave logprobs unset,
still behave exactly as before. The accumulated contents and their order are
unchanged; only the per-delta work drops from O(n) to O(1), making the whole stream
O(n).

The converted logprobs are a fresh, unaliased list[Logprob] from
ChatCmplHelpers.convert_logprobs_for_output_text, and the response part's
logprobs is process-local (initialized to [] at part creation), so extending it
in place is safe and observably identical to the previous concatenation.

Test plan

  • Added test_stream_response_accumulates_logprobs_across_many_deltas in
    tests/models/test_openai_chatcompletions_stream.py: streams several
    logprob-bearing content deltas and asserts the final output-text part accumulates
    every token's logprob in order (guarding that the in-place refactor preserves
    ordered accumulation). This complements the existing
    test_stream_response_includes_logprobs.
  • make format (no changes), make lint, and make typecheck (pyright clean) all
    pass on the change.
  • uv run pytest tests/models/ passes (all model tests, including both streaming
    logprob tests).

Issue number

Checks

  • I've added new tests, if relevant
  • I've run .agents/skills/code-change-verification/scripts/run.sh
  • I've confirmed all verification steps pass
  • If using Codex, I've run /review before submitting this PR

…pies

The Chat Completions stream handler rebuilt the entire accumulated
output-text logprobs list on every content delta via
`existing_logprobs + output_logprobs`, allocating and copying O(n^2)
elements over an n-token stream when logprobs are enabled.

Extend the existing list in place instead, keeping the `None`-provider
guard. Behavior is unchanged: the converted logprobs are a fresh,
unaliased list and the part's `logprobs` is process-local, so the
accumulated result is identical while the work is now linear.

Add a regression test that streams several logprob-bearing deltas and
asserts the final output-text part accumulates every token in order.
@seratch seratch added this to the 0.18.x milestone Jul 9, 2026
@seratch seratch merged commit 1b0c7c2 into openai:main Jul 9, 2026
9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants