[MODEL] Fix handling of multiple channels for gpt-oss with speculative decoding

[astralord]

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

Server command:

vllm serve openai/gpt_oss_20b--speculative-config '{"method": "eagle3", "model": <name-of-your-draft-model>}'

Test script streaming_client.py:

#!/usr/bin/env python3

import asyncio
import sys
from typing import List, Dict

import httpx
from openai import AsyncOpenAI


class StreamingClient:
    def __init__(self, api_url: str = "http://127.0.0.1:8000/v1", api_key: str = "EMPTY"):
        self.client = AsyncOpenAI(
            api_key=api_key,
            base_url=api_url,
            timeout=httpx.Timeout(timeout=300.0, connect=5.0)
        )

    async def send_request(
        self,
        messages: List[Dict[str, str]],
        model: str = "openai/gpt-oss-20b",
        streaming: bool = True,
        **kwargs
    ) -> str:
        """Send a request and return only the generated content."""
        response = await self.client.chat.completions.create(
            messages=messages,
            model=model,
            stream=streaming,
            **kwargs
        )

        if streaming:
            generated_text = ""
            async for chunk in response:
                if chunk.choices and chunk.choices[0].delta.content:
                    generated_text += chunk.choices[0].delta.content
            return generated_text
        else:
            if response.choices and response.choices[0].message.content:
                return response.choices[0].message.content
            return ""

    async def send_multiple_requests(
        self,
        prompts: List[List[Dict[str, str]]],
        model: str = "openai/gpt-oss-20b",
        streaming: bool = True,
        **kwargs
    ) -> List[str]:
        """Send multiple requests concurrently and return generated contents."""
        tasks = [
            asyncio.create_task(self.send_request(messages, model, streaming, **kwargs))
            for messages in prompts
        ]
        return await asyncio.gather(*tasks)


async def main():
    test_prompts = [
        [{"role": "user", "content": "Explain quantum computing in simple terms."}],
        [{"role": "user", "content": "Write a Python function to calculate fibonacci numbers."}],
        [{"role": "user", "content": "What are the benefits of renewable energy?"}],
        [{"role": "user", "content": "Describe the process of photosynthesis."}]
    ]

    api_url = "http://127.0.0.1:8000/v1"
    model = "openai/gpt-oss-20b"
    streaming = len(sys.argv) > 1 and sys.argv[1].lower() == "streaming"

    generation_params = {
        "max_tokens": 256,
        "temperature": 0.7,
        "top_p": 0.9
    }

    mode = "streaming" if streaming else "non-streaming"
    print(f"Starting {mode} requests to {api_url}")
    print(f"Using model: {model}")
    print(f"Number of requests: {len(test_prompts)}\n")

    client = StreamingClient(api_url=api_url)
    results = await client.send_multiple_requests(test_prompts, model, streaming, **generation_params)

    for i, content in enumerate(results, 1):
        print(f"--- Response {i} ---")
        print(content)
        print()


if __name__ == "__main__":
    asyncio.run(main())

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @astralord.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[gemini-code-assist]

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]

💡 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 👍.

[astralord]

@codex review

[chatgpt-codex-connector]

💡 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 👍.

[astralord]

@codex review

[chatgpt-codex-connector]

💡 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 👍.

[astralord]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. You're on a roll.

ℹ️ 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 👍.

[astralord]

/gemini review

[gemini-code-assist]

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 DeltaMessage that 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.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @astralord.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[qandrew]

Hi @astralord , just curious were you planning on merging this in still? If so I can do a review :)

[astralord]

@qandrew Hi! It's great to hear, please, take a look :)

[qandrew]

Hi @astralord , overall looks good. Can you add the results of streaming_client.py to the PR description, and also add a unit test?

[qandrew]

do we need Any here?