[Rust Frontend] Support n > 1 choices

rust/src/chat/src/lib.rs

@@ -63,6 +63,7 @@ mod request;
 mod stream;
 
 use std::collections::HashMap;
+use std::num::NonZeroU32;
 
 use tracing::warn;
 use vllm_engine_core_client::EngineCoreClient;
@@ -282,33 +283,42 @@ impl ChatRequestProcessor {
 
     /// Prepare one chat request without submitting it to an engine.
     ///
-    /// The returned text request is already tokenized: the output processor
-    /// is initialized from the final prompt token IDs and its grammar, if any,
-    /// is applied to the request before it is handed to `text_processor` again
-    /// for lowering.
+    /// The returned text request is already tokenized: the output processors
+    /// are initialized from the final prompt token IDs and their grammar, if
+    /// any, is applied to the request before it is handed to `text_processor`
+    /// again for lowering. The prompt is rendered once; each of the `n`
+    /// samples gets its own output processor, since parsing is stateful.
     pub async fn prepare(
         &self,
         mut request: ChatRequest,
+        n: NonZeroU32,
         text_processor: &TextRequestProcessor,
     ) -> Result<PreparedChatRequest> {
         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..n.get())
+            .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. The processors of
+        // one request build the same grammar, so the first one's is applied.
         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() {
@@ -328,7 +338,7 @@ impl ChatRequestProcessor {
 
         Ok(PreparedChatRequest {
             text_request,
-            output_processor,
+            output_processors,
             prompt_text,
         })
     }
@@ -338,8 +348,8 @@ impl ChatRequestProcessor {
 pub struct PreparedChatRequest {
     /// The tokenized text request, with any output grammar applied.
     pub text_request: TextRequest,
-    /// Request-scoped processor for the decoded output.
-    pub output_processor: DynChatOutputProcessor,
+    /// Request-scoped processors for the decoded output, one per sample.
+    pub output_processors: Vec<DynChatOutputProcessor>,
     /// The rendered prompt text, when the request asked for it and the
     /// renderer produced text.
     pub prompt_text: Option<String>,
@@ -476,23 +486,39 @@ impl ChatLlm {
 
     /// Render, tokenize, and submit one chat request.
     pub async fn chat(&self, request: ChatRequest) -> Result<ChatEventStream> {
+        let mut streams = self.chat_n(request, NonZeroU32::MIN).await?;
+        Ok(streams.pop().expect("one event stream per sample"))
+    }
+
+    /// Render and tokenize one chat request once, submit it as `n` parallel
+    /// samples, and return one event stream per sample, in index order.
+    pub async fn chat_n(
+        &self,
+        request: ChatRequest,
+        n: NonZeroU32,
+    ) -> Result<Vec<ChatEventStream>> {
         let PreparedChatRequest {
             text_request,
-            output_processor,
+            output_processors,
             prompt_text,
-        } = self.processor.prepare(request, self.text.request_processor()).await?;
+        } = self.processor.prepare(request, n, self.text.request_processor()).await?;
         let request_id = text_request.request_id.clone();
         let prepared = self.text.request_processor().prepare(text_request)?;
         let sampling = prepared.effective_sampling();
-        let decoded_stream =
-            self.text.generate_prepared(prepared).await?.map_err(Error::from).boxed();
-
-        let structured_stream = output_processor.process(decoded_stream)?;
-
-        Ok(
-            ChatEventStream::new(request_id, sampling, structured_stream)
-                .with_prompt_text(prompt_text),
-        )
+        let decoded_streams = self.text.generate_prepared_n(prepared, n).await?;
+
+        output_processors
+            .into_iter()
+            .zip(decoded_streams)
+            .map(|(output_processor, decoded_stream)| {
+                let structured_stream =
+                    output_processor.process(decoded_stream.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.

rust/src/cmd/src/cli.rs

@@ -37,8 +37,9 @@ use vllm_managed_engine::data_parallel::{DataParallelLbArgs, DataParallelTopolog
 use vllm_server::{
     AccessLogConfig, ApiServerOptions, ChatTemplateContentFormatOption, Config, CoordinatorMode,
     CorsConfig, CustomHistogramBuckets, DEFAULT_KEEP_ALIVE_TIMEOUT, DEFAULT_LOG_STATS_INTERVAL,
-    DEFAULT_REQUEST_BODY_LIMIT_BYTES, HttpListenerMode, LoraModulePath, OutputLogConfig,
-    OutputLogging, ParserSelection, ReasoningEffort, RenderConfig, RendererSelection,
+    DEFAULT_MAX_N_SEQUENCES, DEFAULT_REQUEST_BODY_LIMIT_BYTES, HttpListenerMode, LoraModulePath,
+    OutputLogConfig, OutputLogging, ParserSelection, ReasoningEffort, RenderConfig,
+    RendererSelection,
 };
 use vllm_text::backend::hf::{HfOverrides, ModelFileSources};
 
@@ -143,6 +144,10 @@ pub struct RenderArgs {
         default_value = "0"
     )]
     skip_model_name_validation: bool,
+    /// Upper bound on a request's `n` (default 16384). Like Python, it is set
+    /// only through `VLLM_MAX_N_SEQUENCES`.
+    #[arg(long, hide = true, env = "VLLM_MAX_N_SEQUENCES", value_parser = parse_int)]
+    max_n_sequences: Option<i64>,
     /// Select the tool-call parser. Use `auto` to infer from the model or
     /// `none` to disable parsing.
     #[arg(long, default_value_t)]
@@ -216,6 +221,7 @@ impl RenderArgs {
             hf_overrides: self.hf_overrides,
             served_model_name: self.served_model_name,
             skip_model_name_validation: self.skip_model_name_validation,
+            max_n_sequences: self.max_n_sequences.unwrap_or(DEFAULT_MAX_N_SEQUENCES),
             host: self.host,
             port: self.port,
             tool_call_parser: self.tool_call_parser,
@@ -577,6 +583,12 @@ pub struct SharedRuntimeArgs {
     #[serde(skip)]
     pub skip_model_name_validation: bool,
 
+    /// Upper bound on a request's `n` (default 16384). Like Python, it is set
+    /// only through `VLLM_MAX_N_SEQUENCES`.
+    #[arg(long, hide = true, env = "VLLM_MAX_N_SEQUENCES", value_parser = parse_int)]
+    #[serde(skip)]
+    pub max_n_sequences: Option<i64>,
+
     /// The model name(s) used in the API. If multiple names are provided, the
     /// server will respond to any of the provided names. The model name in the
     /// model field of a response will be the first name in this list. If not
@@ -761,6 +773,11 @@ impl SharedRuntimeArgs {
         if let Some(value) = lookup("VLLM_SKIP_MODEL_NAME_VALIDATION") {
             self.skip_model_name_validation = parse_truthy_flag(&value)?;
         }
+        if let Some(value) = lookup("VLLM_MAX_N_SEQUENCES") {
+            let max_n_sequences = parse_int(&value)
+                .map_err(|error| format!("invalid VLLM_MAX_N_SEQUENCES: {error}"))?;
+            self.max_n_sequences = Some(max_n_sequences);
+        }
         Ok(())
     }
 
@@ -963,6 +980,7 @@ impl SharedRuntimeArgs {
                 Some(0) => None,
                 Some(limit) => Some(limit),
             },
+            max_n_sequences: self.max_n_sequences.unwrap_or(DEFAULT_MAX_N_SEQUENCES),
         }
     }
 
@@ -1072,11 +1090,16 @@ impl FromStr for LogStatsInterval {
     }
 }
 
+/// Parse an integer environment variable as Python's `int(value)`, for
+/// decimal integers in the `i64` range.
+fn parse_int(value: &str) -> Result<i64, String> {
+    value.trim().parse().map_err(|_| format!("`{value}` is not an integer"))
+}
+
 /// Parse an integer flag environment variable as Python's `bool(int(value))`,
 /// for decimal integers in the `i64` range.
 fn parse_int_flag(value: &str) -> Result<bool, String> {
-    let value: i64 = value.trim().parse().map_err(|_| format!("`{value}` is not an integer"))?;
-    Ok(value != 0)
+    Ok(parse_int(value)? != 0)
 }
 
 /// Parse a boolean flag environment variable as Python's

rust/src/cmd/src/cli/tests.rs

@@ -64,6 +64,9 @@ fn render_args_build_config_without_tls() {
         // hidden flag.
         "--skip-model-name-validation",
         "True",
+        // Set through `VLLM_MAX_N_SEQUENCES`, which binds this hidden flag.
+        "--max-n-sequences",
+        " 4 ",
     ])
     .unwrap();
 
@@ -87,6 +90,7 @@ fn render_args_build_config_without_tls() {
     assert_eq!(config.max_model_len, Some(32768));
     assert_eq!(config.served_model_name, ["qwen"]);
     assert!(config.skip_model_name_validation);
+    assert_eq!(config.max_n_sequences, 4);
     assert_eq!(config.tool_call_parser, ParserSelection::Auto);
     assert_eq!(config.reasoning_parser, ParserSelection::Auto);
     assert_eq!(config.renderer, RendererSelection::DeepSeekV32);
@@ -105,6 +109,7 @@ fn render_args_allow_omitted_max_model_len() {
 
     assert_eq!(config.max_model_len, None);
     assert!(!config.skip_model_name_validation);
+    assert_eq!(config.max_n_sequences, 16384);
 }
 
 #[test]
@@ -268,6 +273,7 @@ fn serve_args_forward_python_flags_with_separator() {
                         disable_request_id_randomization: false,
                         http_max_request_body_size: None,
                         skip_model_name_validation: false,
+                        max_n_sequences: None,
                         served_model_name: [],
                         allowed_origins: JsonStringList(
                             [],
@@ -768,15 +774,16 @@ fn serve_args_forward_env_only_settings_to_frontend_config() {
     assert_eq!(config.log_stats_interval, Some(Duration::from_secs(10)));
     assert!(!config.disable_request_id_randomization);
     assert!(!config.api_server_options.skip_model_name_validation);
+    assert_eq!(config.api_server_options.max_n_sequences, 16384);
     assert_eq!(
         config.api_server_options.request_body_limit_bytes,
         Some(32 * 1024 * 1024)
     );
 
     // Set through `VLLM_LOG_STATS_INTERVAL`,
     // `VLLM_DISABLE_REQUEST_ID_RANDOMIZATION`,
-    // `VLLM_HTTP_MAX_REQUEST_BODY_SIZE`, and
-    // `VLLM_SKIP_MODEL_NAME_VALIDATION`, which bind these hidden flags.
+    // `VLLM_HTTP_MAX_REQUEST_BODY_SIZE`, `VLLM_SKIP_MODEL_NAME_VALIDATION`,
+    // and `VLLM_MAX_N_SEQUENCES`, which bind these hidden flags.
     let config = parse(&[
         "--log-stats-interval",
         "2.5",
@@ -786,10 +793,13 @@ fn serve_args_forward_env_only_settings_to_frontend_config() {
         "1048576",
         "--skip-model-name-validation",
         "1",
+        "--max-n-sequences",
+        " 8 ",
     ]);
     assert_eq!(config.log_stats_interval, Some(Duration::from_millis(2500)));
     assert!(config.disable_request_id_randomization);
     assert!(config.api_server_options.skip_model_name_validation);
+    assert_eq!(config.api_server_options.max_n_sequences, 8);
     assert_eq!(
         config.api_server_options.request_body_limit_bytes,
         Some(1024 * 1024)
@@ -809,6 +819,15 @@ fn serve_args_forward_env_only_settings_to_frontend_config() {
     ])
     .unwrap_err();
     assert_eq!(error.kind(), clap::error::ErrorKind::ValueValidation);
+    let error = Cli::try_parse_from([
+        "vllm-rs",
+        "serve",
+        "Qwen/Qwen3-0.6B",
+        "--max-n-sequences",
+        "1e3",
+    ])
+    .unwrap_err();
+    assert_eq!(error.kind(), clap::error::ErrorKind::ValueValidation);
 }
 
 #[test]
@@ -825,6 +844,7 @@ fn args_json_reads_env_only_settings_from_env() {
     assert_eq!(args.log_stats_interval(), Some(Duration::from_secs(10)));
     assert!(!args.disable_request_id_randomization);
     assert!(!args.api_server_options().skip_model_name_validation);
+    assert_eq!(args.api_server_options().max_n_sequences, 16384);
     assert_eq!(
         args.api_server_options().request_body_limit_bytes,
         Some(32 * 1024 * 1024)
@@ -835,11 +855,13 @@ fn args_json_reads_env_only_settings_from_env() {
         ("VLLM_DISABLE_REQUEST_ID_RANDOMIZATION", "1"),
         ("VLLM_HTTP_MAX_REQUEST_BODY_SIZE", "1048576"),
         ("VLLM_SKIP_MODEL_NAME_VALIDATION", " TRUE "),
+        ("VLLM_MAX_N_SEQUENCES", " 8 "),
     ])
     .unwrap();
     assert_eq!(args.log_stats_interval(), Some(Duration::from_millis(500)));
     assert!(args.disable_request_id_randomization);
     assert!(args.api_server_options().skip_model_name_validation);
+    assert_eq!(args.api_server_options().max_n_sequences, 8);
     assert_eq!(
         args.api_server_options().request_body_limit_bytes,
         Some(1024 * 1024)
@@ -855,13 +877,15 @@ fn args_json_reads_env_only_settings_from_env() {
         apply(&[("VLLM_DISABLE_REQUEST_ID_RANDOMIZATION", "true")]).unwrap_err(),
         apply(&[("VLLM_HTTP_MAX_REQUEST_BODY_SIZE", "32MiB")]).unwrap_err(),
         apply(&[("VLLM_HTTP_MAX_REQUEST_BODY_SIZE", "-1")]).unwrap_err(),
+        apply(&[("VLLM_MAX_N_SEQUENCES", "1e3")]).unwrap_err(),
     ];
     expect![[r#"
         [
             "invalid VLLM_LOG_STATS_INTERVAL: `ten` is not a number of seconds",
             "invalid VLLM_DISABLE_REQUEST_ID_RANDOMIZATION: `true` is not an integer",
             "invalid VLLM_HTTP_MAX_REQUEST_BODY_SIZE: `32MiB` is not a number of bytes",
             "invalid VLLM_HTTP_MAX_REQUEST_BODY_SIZE: `-1` is not a number of bytes",
+            "invalid VLLM_MAX_N_SEQUENCES: `1e3` is not an integer",
         ]
     "#]]
     .assert_debug_eq(&errors);
@@ -1826,6 +1850,7 @@ fn frontend_args_accept_json() {
                         disable_request_id_randomization: false,
                         http_max_request_body_size: None,
                         skip_model_name_validation: false,
+                        max_n_sequences: None,
                         served_model_name: [],
                         allowed_origins: JsonStringList(
                             [],
@@ -2754,6 +2779,7 @@ fn serve_args_accept_handshake_aliases() {
                         disable_request_id_randomization: false,
                         http_max_request_body_size: None,
                         skip_model_name_validation: false,
+                        max_n_sequences: None,
                         served_model_name: [],
                         allowed_origins: JsonStringList(
                             [],
@@ -3050,6 +3076,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() {
                 request_body_limit_bytes: Some(
                     33554432,
                 ),
+                max_n_sequences: 16384,
             },
             cors: CorsConfig {
                 allow_origins: [],
@@ -3178,6 +3205,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() {
                 request_body_limit_bytes: Some(
                     33554432,
                 ),
+                max_n_sequences: 16384,
             },
             cors: CorsConfig {
                 allow_origins: [],
@@ -3419,6 +3447,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present
                 request_body_limit_bytes: Some(
                     33554432,
                 ),
+                max_n_sequences: 16384,
             },
             cors: CorsConfig {
                 allow_origins: [],