[Rust Frontend] Support OpenAI n choices

rust/src/server/src/routes/openai/chat_completions.rs

@@ -15,7 +15,9 @@ use axum::extract::State;
 use axum::http::HeaderMap;
 use axum::response::sse::Event;
 use axum::response::{IntoResponse, Response};
-use futures::{Stream, StreamExt as _, pin_mut};
+use futures::future::try_join_all;
+use futures::stream::{BoxStream, select_all};
+use futures::{Stream, StreamExt as _, TryStreamExt as _, pin_mut};
 use serde_json::Value;
 use thiserror_ext::AsReport as _;
 use tracing::{debug, error, info, trace};
@@ -38,6 +40,9 @@ use crate::routes::openai::chat_completions::types::{
 use crate::routes::openai::utils::logprobs::{
     decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps,
 };
+use crate::routes::openai::utils::parallel_sampling::{
+    choice_request_id, choice_seed, merge_choice_usage,
+};
 use crate::routes::openai::utils::types::{
     ChatLogProbs, FunctionCallDelta, FunctionCallResponse, StreamResponseEnvelope, ToolCall,
     ToolCallDelta, Usage,
@@ -80,44 +85,134 @@ pub async fn chat_completions(
     let created = unix_timestamp();
     let api_server_options = state.api_server_options;
 
-    let chat_stream =
-        match state.chat.chat(prepared.chat_request).instrument(request_span.clone()).await {
-            Ok(stream) => stream,
+    let mut chat_streams = Vec::with_capacity(prepared.n as usize);
+    for index in 0..prepared.n {
+        let chat_request = match choice_chat_request(&prepared.chat_request, index, prepared.n) {
+            Ok(chat_request) => chat_request,
+            Err(error) => return error.into_response(),
+        };
+        match state.chat.chat(chat_request).instrument(request_span.clone()).await {
+            Ok(stream) => chat_streams.push(stream),
             Err(error) => {
                 return chat_submit_error("failed to submit chat request", error).into_response();
             }
-        };
+        }
+    }
 
     if stream {
-        let chunk_stream = chat_completion_chunk_stream(
-            chat_stream,
-            prepared.request_id,
-            prepared.response_model,
-            created,
-            api_server_options,
-            prepared.options,
-        );
+        let chunk_streams = chat_streams
+            .into_iter()
+            .zip(0..)
+            .map(|(chat_stream, index)| {
+                chat_completion_chunk_stream(
+                    chat_stream,
+                    prepared.request_id.clone(),
+                    prepared.response_model.clone(),
+                    created,
+                    api_server_options,
+                    prepared.options.clone(),
+                )
+                .map_ok(move |chunk| with_choice_index(chunk, index))
+                .boxed()
+            })
+            .collect();
+        let chunk_stream = merge_choice_chunk_streams(chunk_streams);
         let sse_stream = chat_completion_sse_stream(chunk_stream).instrument(request_span);
 
         sse_response(sse_stream, api_server_options.sse_keep_alive_interval)
     } else {
-        let response = match collect_chat_completion(
-            chat_stream,
-            prepared.request_id,
-            prepared.response_model,
-            created,
-            api_server_options,
-            prepared.options,
-        )
+        let responses = match try_join_all(chat_streams.into_iter().map(|chat_stream| {
+            collect_chat_completion(
+                chat_stream,
+                prepared.request_id.clone(),
+                prepared.response_model.clone(),
+                created,
+                api_server_options,
+                prepared.options.clone(),
+            )
+        }))
         .instrument(request_span.clone())
         .await
         {
-            Ok(response) => response,
+            Ok(responses) => responses,
             Err(error) => return error.into_response(),
         };
 
-        Json(response).into_response()
+        Json(merge_choice_responses(responses)).into_response()
+    }
+}
+
+/// Build the child chat request for one of `n` choices.
+fn choice_chat_request(base: &ChatRequest, index: u32, n: u32) -> Result<ChatRequest, ApiError> {
+    let mut chat_request = base.clone();
+    chat_request.request_id = choice_request_id(&base.request_id, index, n);
+    chat_request.sampling_params.seed = choice_seed(base.sampling_params.seed, index)?;
+    Ok(chat_request)
+}
+
+/// Set the choice index on every choice of one streamed chunk.
+fn with_choice_index(
+    mut chunk: ChatCompletionStreamResponse,
+    index: u32,
+) -> ChatCompletionStreamResponse {
+    for choice in &mut chunk.choices {
+        choice.index = index;
+    }
+    chunk
+}
+
+/// Interleave per-choice chunk streams into one response stream.
+///
+/// Each choice's usage-only chunk (usage set, no choices) is held back and
+/// folded into one final usage chunk, so clients see a single usage chunk
+/// after all choices finish.
+#[try_stream]
+async fn merge_choice_chunk_streams(
+    streams: Vec<BoxStream<'static, Result<ChatCompletionStreamResponse, ApiError>>>,
+    mut y: TryYielder<ChatCompletionStreamResponse, ApiError>,
+) -> Result<(), ApiError> {
+    let mut chunks = select_all(streams);
+    let mut merged_usage_chunk: Option<ChatCompletionStreamResponse> = None;
+
+    while let Some(chunk) = chunks.next().await {
+        let mut chunk = chunk?;
+        if !chunk.choices.is_empty() || chunk.usage.is_none() {
+            y.yield_ok(chunk).await;
+            continue;
+        }
+        if let Some(merged) = merged_usage_chunk.take() {
+            chunk.usage = merged
+                .usage
+                .zip(chunk.usage)
+                .map(|(usage, choice_usage)| merge_choice_usage(usage, choice_usage));
+        }
+        merged_usage_chunk = Some(chunk);
+    }
+
+    if let Some(usage_chunk) = merged_usage_chunk {
+        y.yield_ok(usage_chunk).await;
+    }
+    Ok(())
+}
+
+/// Fold per-choice responses into one response with indexed choices.
+///
+/// Response-level fields such as prompt logprobs and `kv_transfer_params` come
+/// from the first choice.
+fn merge_choice_responses(responses: Vec<ChatCompletionResponse>) -> ChatCompletionResponse {
+    let mut responses = responses.into_iter();
+    let mut merged = responses.next().expect("n must be at least 1");
+    for response in responses {
+        merged.usage = merged
+            .usage
+            .zip(response.usage)
+            .map(|(usage, choice_usage)| merge_choice_usage(usage, choice_usage));
+        merged.choices.extend(response.choices);
+    }
+    for (choice, index) in merged.choices.iter_mut().zip(0..) {
+        choice.index = index;
     }
+    merged
 }
 
 async fn collect_chat_completion(

rust/src/server/src/routes/openai/chat_completions/convert.rs

@@ -30,6 +30,8 @@ pub(super) struct PreparedRequest {
     pub request_id: String,
     /// Public model ID echoed back to the client.
     pub response_model: String,
+    /// Number of choices to sample for this request.
+    pub n: u32,
     /// Public response rendering options for route-layer helpers.
     pub options: ResponseOptions,
     /// Lowered chat request for `vllm-chat`.
@@ -205,6 +207,7 @@ pub(super) fn prepare_chat_request(
     Ok(PreparedRequest {
         request_id,
         response_model,
+        n: request.n.unwrap_or(1),
         options: ResponseOptions {
             include_usage,
             include_continuous_usage,

rust/src/server/src/routes/openai/chat_completions/validate.rs

@@ -24,10 +24,6 @@ pub(super) fn validate_request_compat(
         );
     }
 
-    if request.n.unwrap_or(1) > 1 {
-        bail_invalid_request!(param = "n", "Only n=1 is supported.");
-    }
-
     if request.top_logprobs.is_some() && !request.logprobs {
         bail_invalid_request!(
             param = "top_logprobs",

rust/src/server/src/routes/openai/completions.rs

@@ -15,7 +15,9 @@ use axum::extract::State;
 use axum::http::HeaderMap;
 use axum::response::sse::Event;
 use axum::response::{IntoResponse, Response};
-use futures::{Stream, StreamExt as _, pin_mut};
+use futures::future::try_join_all;
+use futures::stream::{BoxStream, select_all};
+use futures::{Stream, StreamExt as _, TryStreamExt as _, pin_mut};
 use serde_json::Value;
 use thiserror_ext::AsReport as _;
 use tracing::{debug, error, info, trace};
@@ -41,6 +43,9 @@ use crate::routes::openai::completions::types::{
     CompletionChoice, CompletionResponse, CompletionSseChunk, CompletionStreamChoice,
     CompletionStreamResponse,
 };
+use crate::routes::openai::utils::parallel_sampling::{
+    choice_request_id, choice_seed, merge_choice_usage,
+};
 use crate::routes::openai::utils::types::LogProbs;
 use crate::routes::openai::utils::usage::ContinuousUsage;
 use crate::routes::openai::utils::validated_json::ValidatedJson;
@@ -87,49 +92,133 @@ pub async fn completions(
 
     let created = unix_timestamp();
     let api_server_options = state.api_server_options;
-    let text_stream = match state
-        .chat
-        .text()
-        .generate(prepared.text_request)
-        .instrument(request_span.clone())
-        .await
-    {
-        Ok(stream) => stream,
-        Err(error) => {
-            return text_submit_error("failed to submit completion request", error).into_response();
+    let mut text_streams = Vec::with_capacity(prepared.n as usize);
+    for index in 0..prepared.n {
+        let text_request = match choice_text_request(&prepared.text_request, index, prepared.n) {
+            Ok(text_request) => text_request,
+            Err(error) => return error.into_response(),
+        };
+        match state.chat.text().generate(text_request).instrument(request_span.clone()).await {
+            Ok(stream) => text_streams.push(stream),
+            Err(error) => {
+                return text_submit_error("failed to submit completion request", error)
+                    .into_response();
+            }
         }
-    };
+    }
 
     if stream {
-        let chunk_stream = completion_chunk_stream(
-            text_stream,
-            prepared.request_id,
-            prepared.response_model,
-            created,
-            api_server_options,
-            prepared.options,
-        );
+        let chunk_streams = text_streams
+            .into_iter()
+            .zip(0..)
+            .map(|(text_stream, index)| {
+                completion_chunk_stream(
+                    text_stream,
+                    prepared.request_id.clone(),
+                    prepared.response_model.clone(),
+                    created,
+                    api_server_options,
+                    prepared.options.clone(),
+                )
+                .map_ok(move |chunk| with_choice_index(chunk, index))
+                .boxed()
+            })
+            .collect();
+        let chunk_stream = merge_choice_chunk_streams(chunk_streams);
         let sse_stream = completion_sse_stream(chunk_stream).instrument(request_span);
 
         sse_response(sse_stream, api_server_options.sse_keep_alive_interval)
     } else {
-        let response = match collect_completion(
-            text_stream,
-            prepared.request_id,
-            prepared.response_model,
-            created,
-            api_server_options,
-            prepared.options,
-        )
+        let responses = match try_join_all(text_streams.into_iter().map(|text_stream| {
+            collect_completion(
+                text_stream,
+                prepared.request_id.clone(),
+                prepared.response_model.clone(),
+                created,
+                api_server_options,
+                prepared.options.clone(),
+            )
+        }))
         .instrument(request_span.clone())
         .await
         {
-            Ok(response) => response,
+            Ok(responses) => responses,
             Err(error) => return error.into_response(),
         };
 
-        Json(response).into_response()
+        Json(merge_choice_responses(responses)).into_response()
+    }
+}
+
+/// Build the child text request for one of `n` choices.
+fn choice_text_request(base: &TextRequest, index: u32, n: u32) -> Result<TextRequest, ApiError> {
+    let mut text_request = base.clone();
+    text_request.request_id = choice_request_id(&base.request_id, index, n);
+    text_request.sampling_params.seed = choice_seed(base.sampling_params.seed, index)?;
+    Ok(text_request)
+}
+
+/// Set the choice index on every choice of one streamed chunk.
+fn with_choice_index(mut chunk: CompletionSseChunk, index: u32) -> CompletionSseChunk {
+    if let CompletionSseChunk::Chunk(response) = &mut chunk {
+        for choice in &mut response.choices {
+            choice.index = index;
+        }
+    }
+    chunk
+}
+
+/// Interleave per-choice chunk streams into one response stream.
+///
+/// Each choice's usage chunk is held back and folded into one final usage
+/// chunk, so clients see a single usage chunk after all choices finish.
+#[try_stream]
+async fn merge_choice_chunk_streams(
+    streams: Vec<BoxStream<'static, Result<CompletionSseChunk, ApiError>>>,
+    mut y: TryYielder<CompletionSseChunk, ApiError>,
+) -> Result<(), ApiError> {
+    let mut chunks = select_all(streams);
+    let mut merged_usage_chunk: Option<CompletionStreamResponse> = None;
+
+    while let Some(chunk) = chunks.next().await {
+        match chunk? {
+            CompletionSseChunk::Usage(mut usage_chunk) => {
+                if let Some(merged) = merged_usage_chunk.take() {
+                    usage_chunk.usage = merged
+                        .usage
+                        .zip(usage_chunk.usage)
+                        .map(|(usage, choice_usage)| merge_choice_usage(usage, choice_usage));
+                }
+                merged_usage_chunk = Some(usage_chunk);
+            }
+            chunk => y.yield_ok(chunk).await,
+        }
+    }
+
+    if let Some(usage_chunk) = merged_usage_chunk {
+        y.yield_ok(CompletionSseChunk::Usage(usage_chunk)).await;
+    }
+    Ok(())
+}
+
+/// Fold per-choice responses into one response with indexed choices.
+///
+/// Response-level fields such as `kv_transfer_params` come from the first
+/// choice.
+fn merge_choice_responses(responses: Vec<CompletionResponse>) -> CompletionResponse {
+    let mut responses = responses.into_iter();
+    let mut merged = responses.next().expect("n must be at least 1");
+    for response in responses {
+        merged.usage = merged
+            .usage
+            .zip(response.usage)
+            .map(|(usage, choice_usage)| merge_choice_usage(usage, choice_usage));
+        merged.choices.extend(response.choices);
+    }
+    for (choice, index) in merged.choices.iter_mut().zip(0..) {
+        choice.index = index;
     }
+    merged
 }
 
 async fn collect_completion(