[Rust Frontend] Support beam search (use_beam_search, length_penalty)

.buildkite/test_areas/rust_frontend.yaml

@@ -46,7 +46,7 @@ steps:
   # The Qwen3.5-35B FP8 MTP case needs two GPUs.
   - pytest -v -s entrypoints/openai/chat_completion/test_thinking_token_budget.py -k "not qwen35"
   # The Rust frontend drops prompt_tokens_details when nothing was cached, rejects batched
-  # prompts, and supports only n=1 (so no parallel sampling or beam search).
+  # prompts, and supports n > 1 only with beam search (so no parallel sampling).
   - pytest -v -s entrypoints/openai/completion/test_completion.py -k "not test_single_completion and not test_prompt_logprobs_completion and not test_parallel and not test_batch_completions"
   - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py --timeout=900 --timeout-method=thread
   - pytest -v -s entrypoints/openai/test_chunked_prompt.py

rust/src/chat/src/error.rs

@@ -114,6 +114,12 @@ pub enum Error {
         #[source]
         error: BoxedError,
     },
+    #[error(
+        "Tool calling that constrains output is not supported with beam search: \
+         `tool_choice` `required` or a named function, strict tools or \
+         `--tool-strict-level`, or a tool parser that always constrains output."
+    )]
+    OutputGrammarWithBeamSearch { param: Option<&'static str> },
     #[error(transparent)]
     Text(#[from] vllm_text::Error),
     #[error(transparent)]
@@ -132,6 +138,7 @@ impl Error {
             | Self::DuplicateToolName { .. }
             | Self::ToolChoiceRequiresTools
             | Self::ToolChoiceFunctionNotFound { .. }
+            | Self::OutputGrammarWithBeamSearch { .. }
             | Self::UnsupportedChatRole { .. }
             | Self::UntrustedRequestChatTemplate
             | Self::UnsupportedRequestChatTemplate
@@ -153,6 +160,7 @@ impl Error {
     pub fn param(&self) -> Option<&'static str> {
         match self {
             Self::Text(error) => error.param(),
+            Self::OutputGrammarWithBeamSearch { param } => *param,
             Self::UnsupportedRequestChatTemplate => Some("chat_template"),
             _ => None,
         }

rust/src/chat/src/lib.rs

@@ -70,7 +70,7 @@ use vllm_engine_core_client::protocol::dtype::ModelDtype;
 use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
 use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputBackend;
 use vllm_llm::Llm;
-use vllm_text::{Prompt, TextLlm, TextRequest, TextRequestProcessor};
+use vllm_text::{BeamSearchParams, Prompt, TextLlm, TextRequest, TextRequestProcessor};
 
 /// Chat request preparation shared by inference and render-only frontends.
 pub struct ChatRequestProcessor {
@@ -288,27 +288,49 @@ impl ChatRequestProcessor {
     /// for lowering.
     pub async fn prepare(
         &self,
-        mut request: ChatRequest,
+        request: ChatRequest,
         text_processor: &TextRequestProcessor,
     ) -> Result<PreparedChatRequest> {
+        let (text_request, mut output_processors, prompt_text) =
+            self.prepare_for_outputs(request, text_processor, 1).await?;
+        Ok(PreparedChatRequest {
+            text_request,
+            output_processor: output_processors.pop().expect("one output processor requested"),
+            prompt_text,
+        })
+    }
+
+    /// Like [`Self::prepare`], with one output processor per output sequence.
+    async fn prepare_for_outputs(
+        &self,
+        mut request: ChatRequest,
+        text_processor: &TextRequestProcessor,
+        num_outputs: usize,
+    ) -> Result<(TextRequest, Vec<DynChatOutputProcessor>, Option<String>)> {
         request.validate()?;
         self.map_reasoning_effort(&mut request);
 
-        let mut output_processor = self.backend.new_chat_output_processor(
-            &mut request,
-            NewChatOutputProcessorOptions {
-                tool_call_parser: &self.tool_call_parser,
-                reasoning_parser: &self.reasoning_parser,
-                tool_strict_level: self.tool_strict_level,
-            },
-        )?;
+        let mut output_processors = (0..num_outputs)
+            .map(|_| {
+                self.backend.new_chat_output_processor(
+                    &mut request,
+                    NewChatOutputProcessorOptions {
+                        tool_call_parser: &self.tool_call_parser,
+                        reasoning_parser: &self.reasoning_parser,
+                        tool_strict_level: self.tool_strict_level,
+                    },
+                )
+            })
+            .collect::<Result<Vec<_>>>()?;
         let (mut text_request, prompt_text) = self.prepare_text_request(request).await?;
 
-        // Initialize the output processor from the final prompt token IDs and apply
-        // the output grammar to the request, if any.
+        // Initialize the output processors from the final prompt token IDs and
+        // apply the output grammar to the request, if any.
         let prompt_token_ids = text_processor.tokenize_in_place(&mut text_request)?;
-        output_processor.initialize(prompt_token_ids)?;
-        let grammar = output_processor.build_output_grammar()?;
+        for output_processor in &mut output_processors {
+            output_processor.initialize(prompt_token_ids)?;
+        }
+        let grammar = output_processors[0].build_output_grammar()?;
         // Output grammars are xgrammar structural tags. With another backend the
         // request keeps its own constraint and parsing runs unconstrained.
         match text_processor.structured_outputs_backend() {
@@ -326,11 +348,7 @@ impl ChatRequestProcessor {
             }
         }
 
-        Ok(PreparedChatRequest {
-            text_request,
-            output_processor,
-            prompt_text,
-        })
+        Ok((text_request, output_processors, prompt_text))
     }
 }
 
@@ -495,6 +513,52 @@ impl ChatLlm {
         )
     }
 
+    /// Run beam search and return one finished chat stream per beam, best first.
+    pub async fn beam_search(
+        &self,
+        request: ChatRequest,
+        params: BeamSearchParams,
+    ) -> Result<Vec<ChatEventStream>> {
+        let requested_constraint = request.sampling_params.structured_outputs.is_some();
+        let grammar_param = match request.tool_context.tool_choice {
+            ChatToolChoice::Required | ChatToolChoice::Function { .. } => Some("tool_choice"),
+            _ if !request.tool_context.effective_tools.is_empty() => Some("tools"),
+            _ => None,
+        };
+        let (text_request, output_processors, prompt_text) = self
+            .processor
+            .prepare_for_outputs(
+                request,
+                self.text.request_processor(),
+                params.beam_width.get() as usize,
+            )
+            .await?;
+        // The engine cannot resume a tool output grammar from a beam's
+        // tokens. TODO: enforce it per beam with a frontend grammar matcher, as
+        // Python's offline `LLM.beam_search` does.
+        if !requested_constraint && text_request.sampling_params.structured_outputs.is_some() {
+            return Err(Error::OutputGrammarWithBeamSearch {
+                param: grammar_param,
+            });
+        }
+        let request_id = text_request.request_id.clone();
+        let (sampling, beams) = self.text.beam_search(text_request, params).await?;
+
+        // A failed search returns fewer streams than processors.
+        beams
+            .into_iter()
+            .zip(output_processors)
+            .map(|(beam, output_processor)| {
+                let structured_stream =
+                    output_processor.process(beam.map_err(Error::from).boxed())?;
+                Ok(
+                    ChatEventStream::new(request_id.clone(), sampling, structured_stream)
+                        .with_prompt_text(prompt_text.clone()),
+                )
+            })
+            .collect()
+    }
+
     /// Abort in-flight requests by their external (user-supplied) request ids.
     pub async fn abort(&self, external_ids: &[String]) -> Result<()> {
         self.text.abort(external_ids).await?;

rust/src/chat/src/stream.rs

@@ -9,9 +9,7 @@ use futures::Stream;
 use trait_set::trait_set;
 use vllm_engine_core_client::protocol::output::RequestSpecDecodeMetrics;
 use vllm_llm::{RequestTimestamps, WeightVersionSpan};
-use vllm_text::{
-    DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs, EffectiveSampling,
-};
+use vllm_text::{DecodedLogprobs, DecodedPromptLogprobs, EffectiveSampling};
 
 use crate::FinishReason;
 use crate::error::{Error, Result};
@@ -103,7 +101,8 @@ impl ChatEventStream {
         let mut prompt_logprobs = None;
         let mut prompt_token_ids: Arc<[u32]> = Arc::from([]);
         let prompt_text = self.prompt_text.take();
-        let mut logprob_positions: Vec<DecodedPositionLogprobs> = Vec::new();
+        // Kept when empty: a beam can stop at EOS before returning any token.
+        let mut output_logprobs: Option<DecodedLogprobs> = None;
         let mut token_ids: Vec<u32> = Vec::new();
         while let Some(event) = self.next().await.transpose()? {
             match event {
@@ -120,7 +119,10 @@ impl ChatEventStream {
                     token_ids: delta_ids,
                 } => {
                     if let Some(logprobs) = logprobs {
-                        logprob_positions.extend(logprobs.positions);
+                        output_logprobs
+                            .get_or_insert_with(|| DecodedLogprobs { positions: vec![] })
+                            .positions
+                            .extend(logprobs.positions);
                     }
                     token_ids.extend(delta_ids);
                 }
@@ -139,9 +141,7 @@ impl ChatEventStream {
                         prompt_token_ids,
                         prompt_text,
                         prompt_logprobs,
-                        logprobs: (!logprob_positions.is_empty()).then_some(DecodedLogprobs {
-                            positions: logprob_positions,
-                        }),
+                        logprobs: output_logprobs,
                         token_ids,
                         usage,
                         timestamps,

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

@@ -26,8 +26,9 @@ use vllm_chat::{
 };
 use vllm_engine_core_client::protocol::output::StopReason;
 use vllm_llm::WeightVersionSpan;
+use vllm_text::BeamSearchParams;
 
-use self::convert::{ResponseOptions, prepare_chat_request};
+use self::convert::{PreparedRequest, ResponseOptions, prepare_chat_request};
 pub(crate) use self::types::ChatCompletionRequest;
 use crate::config::{ApiServerOptions, OutputLogConfig};
 use crate::error::{ApiError, bail_server_error, chat_submit_error, engine_error, server_error};
@@ -94,6 +95,21 @@ pub async fn chat_completions(
     // Like Python, report reasoning-token details only when a reasoning parser
     // is configured to count them.
     let count_reasoning_tokens = state.chat.reasoning_parser_name().is_some();
+    if let Some(params) = prepared.beam_search {
+        return match beam_search_chat_completion(
+            &state,
+            prepared,
+            params,
+            created,
+            count_reasoning_tokens,
+        )
+        .instrument(request_span)
+        .await
+        {
+            Ok(response) => Json(response).into_response(),
+            Err(error) => error.into_response(),
+        };
+    }
 
     let mut chat_stream =
         match state.chat.chat(prepared.chat_request).instrument(request_span.clone()).await {
@@ -139,6 +155,49 @@ pub async fn chat_completions(
     }
 }
 
+async fn beam_search_chat_completion(
+    state: &AppState,
+    prepared: PreparedRequest,
+    params: BeamSearchParams,
+    created: u64,
+    count_reasoning_tokens: bool,
+) -> Result<ChatCompletionResponse, ApiError> {
+    let beams = (state.chat.beam_search(prepared.chat_request, params))
+        .await
+        .map_err(|error| chat_submit_error("failed to run beam search", error))?;
+    let mut merged: Option<ChatCompletionResponse> = None;
+    for beam in beams {
+        let response = collect_chat_completion(
+            beam,
+            prepared.request_id.clone(),
+            prepared.response_model.clone(),
+            created,
+            state.api_server_options,
+            count_reasoning_tokens,
+            prepared.options.clone(),
+        )
+        .await?;
+        match &mut merged {
+            None => merged = Some(response),
+            Some(merged) => {
+                if let Some(usage) = &mut merged.usage
+                    && let Some(choice_usage) = &response.usage
+                {
+                    usage.add_choice(choice_usage);
+                }
+                merged.choices.extend(response.choices);
+            }
+        }
+    }
+    let mut merged = merged.ok_or_else(|| server_error!("beam search returned no beams"))?;
+    for (index, choice) in (0..).zip(&mut merged.choices) {
+        choice.index = index;
+    }
+    // A search spans many engine requests; Python reports no metrics either.
+    merged.metrics = None;
+    Ok(merged)
+}
+
 async fn collect_chat_completion(
     stream: ChatEventStream,
     request_id: String,

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

@@ -7,16 +7,17 @@ use vllm_chat::{
     ChatMessage as VllmChatMessage, ChatOptions, ChatRequest, ChatTool, ChatToolChoice,
     GenerationPromptMode, ResolvedToolContext, SamplingParams,
 };
+use vllm_text::BeamSearchParams;
 
 use super::types::ChatCompletionRequest;
 use super::validate;
 use crate::error::{ApiError, bail_invalid_request, chat_submit_error, text_submit_error};
 use crate::lora::LoraModelResolution;
-use crate::routes::openai::utils::resolve_generation_prompt_truncation;
 use crate::routes::openai::utils::structured_outputs::convert_from_response_format;
 use crate::routes::openai::utils::types::{
     ChatMessage, ContentPart, MessageContent, Tool, ToolChoice, ToolChoiceValue, ToolReference,
 };
+use crate::routes::openai::utils::{resolve_beam_search, resolve_generation_prompt_truncation};
 use crate::utils::{
     ResolvedRequestContext, convert_logit_bias, merge_ec_transfer_params, merge_kv_transfer_params,
     resolve_session_id,
@@ -32,6 +33,8 @@ pub(super) struct PreparedRequest {
     pub response_model: String,
     /// Public response rendering options for route-layer helpers.
     pub options: ResponseOptions,
+    /// Set when the request uses beam search.
+    pub beam_search: Option<BeamSearchParams>,
     /// Lowered chat request for `vllm-chat`.
     pub chat_request: ChatRequest,
 }
@@ -74,6 +77,8 @@ pub(super) fn prepare_chat_request(
     ctx: ResolvedRequestContext,
 ) -> Result<PreparedRequest, ApiError> {
     validate::validate_request_compat(&request, &lora_resolution.model_names)?;
+    let beam_search =
+        resolve_beam_search(request.use_beam_search, request.n, request.length_penalty)?;
 
     let prompt_truncation = resolve_generation_prompt_truncation(
         request.truncate_prompt_tokens,
@@ -218,6 +223,7 @@ pub(super) fn prepare_chat_request(
     Ok(PreparedRequest {
         request_id,
         response_model,
+        beam_search,
         options: ResponseOptions {
             include_usage,
             include_continuous_usage,

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

@@ -143,7 +143,7 @@ pub struct ChatCompletionRequest {
     pub repetition_penalty: Option<f32>,
 
     /// Length penalty for beam search
-    pub length_penalty: Option<f32>,
+    pub length_penalty: Option<f64>,
 
     /// Specific token IDs to use as stop conditions
     pub stop_token_ids: Option<Vec<u32>>,