Motivation

Official OpenAI API and vLLM supports stream: true with parallel sampling (n > 1), and multi-prompt completions (len(prompt) > 1). However current version of SGLang does not support this.

The goal here is to support requests with stream: true AND (n > 1 OR len(prompt) > 1).

Modification

This PR:

• Updates tokenizer_manager._handle_batch_request function to support streaming response when enabled. This has been implemented by re-using tokenizer_manager._wait_for_response, which is currently being used for tokenizer_manager._handle_single_request.
• Index of the response is determined by i * batch_index + parallel_sample_index, which is the same behavior of official OpenAI API and vLLM.
• Updates tokenizer_manager._wait_for_response to be able to handle batch request as well.
• Updates openai.v1_(chat_|)completions functions to be able to correctly handle batched streaming responses.
• Moves stream = False assertion from tokenizer_manager, which was asserting all the batch requests, to openai.process_batch, which will assert only the /v1/batches requests.
• Adds relevant tests.

Notes to reviewers

I have tested this on my local environment, and will be rolling out to few of our production GPU nodes soon for further testing.

Checklist

• Before submitting a PR for review, make sure it has passed verification in your local development environment at least.
• Ensure pre-commit pre-commit run --all-files or other linting tools are used to fix potential lint issues.
• Confirm that modifications are covered by complete unit tests. If not, please add more unit tests for correctness.
• Modify documentation as needed, such as docstrings or example tutorials.