[MODEL] Fix handling of multiple channels for gpt-oss with speculative decoding #26291
Conversation
|
This pull request has merge conflicts that must be resolved before it can be |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request effectively addresses a bug in handling multiple channels for gpt-oss with speculative decoding, particularly in streaming mode. The changes introduce a more robust mechanism by tracking the state for each token, grouping them by channel and recipient, and then constructing the delta messages. This ensures that channel switches within a single decoding step are handled correctly. The related logging improvements are also a good addition, providing more comprehensive output. I have one suggestion to improve code readability and reduce the risk of future bugs by removing a magic number.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
| if group_text: | ||
| tool_messages.append(DeltaToolCall( | ||
| index=base_index - 1, # Use the index of the current tool call |
There was a problem hiding this comment.
Within the Harmony streaming branch, base_index is derived from harmony_parser.messages to count previously completed tool calls, the same value the prior implementation used directly for both opening a call and streaming additional arguments. The new logic subtracts one when emitting argument deltas (index=base_index - 1). When a tool call spans multiple decoding iterations, base_index remains 0 (no completed calls yet), so subsequent argument chunks are sent with index=-1, which clients cannot associate with the active tool call. This breaks streaming for any tool call whose arguments arrive over multiple chunks.
Useful? React with 👍 / 👎.
astralord
force-pushed
the
fix-gpt-oss-with-speculative-decoding-handle-multiple-channels
branch
2 times, most recently
from
a953d9c to
ede4584
Compare
astralord
force-pushed
the
fix-gpt-oss-with-speculative-decoding-handle-multiple-channels
branch
from
704867c to
3c1bf55
Compare
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
| groups = [] | ||
| for channel, recipient, text in token_states: | ||
| if not text: | ||
| continue | ||
| if groups and groups[-1]['channel'] == channel and groups[-1]['recipient'] == recipient: | ||
| groups[-1]['text'] += text | ||
| else: |
There was a problem hiding this comment.
The new grouping logic skips any token whose last_content_delta is empty before building groups. When a speculative chunk switches channels (e.g., starts a new commentary function call) the parser often changes current_channel/current_recipient using structural tokens that have no textual delta. Previously, the code still emitted a DeltaToolCall whenever prev_recipient != cur_recipient even if delta_text was empty, so clients immediately saw the tool name. With the if not text: continue filter, those structural tokens are ignored and no DeltaToolCall is streamed until arguments appear; calls with empty arguments are never surfaced at all. This regression breaks streaming compatibility for Harmony tool calls that begin with metadata-only tokens.
Useful? React with 👍 / 👎.
astralord
force-pushed
the
fix-gpt-oss-with-speculative-decoding-handle-multiple-channels
branch
from
3c1bf55 to
e1f14dd
Compare
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
| # next_tool_index tracks the index for the next NEW tool call | ||
| next_tool_index = base_index | ||
|
|
There was a problem hiding this comment.
next_tool_index is initialised from base_index and only increased when prev_recipient != group_recipient. If this chunk first streams more arguments for an already-active tool call and then starts a second tool call, the continuation leaves next_tool_index unchanged, so the new call is emitted with the same index as the ongoing call and its arguments are also written under that index. The OpenAI streaming API expects tool call indexes to be unique and monotonically increasing, so this mislabels calls and causes clients to merge two separate calls. Consider advancing next_tool_index when a continuation is emitted or computing base_index to include the in-flight call.
Useful? React with 👍 / 👎.
Signed-off-by: Aleksandr Samarin <[email protected]>
Signed-off-by: Aleksandr Samarin <[email protected]>
Signed-off-by: Aleksandr Samarin <[email protected]>
Signed-off-by: Aleksandr Samarin <[email protected]>
Signed-off-by: Aleksandr Samarin <[email protected]>
Signed-off-by: Aleksandr Samarin <[email protected]>
Signed-off-by: Aleksandr Samarin <[email protected]>
astralord
force-pushed
the
fix-gpt-oss-with-speculative-decoding-handle-multiple-channels
branch
from
e4f6360 to
3ad1d7b
Compare
…-multiple-channels
|
@codex review |
|
Codex Review: Didn't find any major issues. You're on a roll. ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request provides a much-needed fix for handling multiple channels in gpt-oss streaming with speculative decoding. The previous implementation had a flaw where it only considered the state after the last token in a chunk, which could lead to data loss or incorrect message construction if the channel or recipient changed within the chunk.
The new approach is robust and correctly handles this complex scenario. Key improvements include:
- Tracking the state (channel, recipient, delta) for each individual token.
- Grouping consecutive tokens with the same state for efficient processing.
- Refactoring the logic to build a single, comprehensive
DeltaMessagethat can contain content, reasoning, and tool calls from a single chunk. - Improving the indexing logic for tool calls, correctly handling calls that span across multiple streamed chunks.
- Enhancing logging to be more comprehensive.
The changes significantly increase the correctness and reliability of streaming for gpt-oss models. The implementation is well-structured, and the added complexity is justified by the problem it solves. I don't see any issues with the proposed changes.
Hi team!
Purpose
We've noticed that that the recent PR doesn't fully fix gpt-oss + streaming + speculative-decoding issue, for example generated messages end abruptly. This happens because multiple tokens can relate to different channels (e.g.
<final>,<analysis>,None) in one decoding stage. This PR handles it.Test Plan
Test Result