diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index d3e369bbe3..14b949c87d 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -12,6 +12,7 @@ mod unsupported; use std::collections::HashMap; use std::ffi::{OsStr, OsString}; +use std::num::NonZeroU32; use std::path::PathBuf; use std::time::Duration; @@ -268,6 +269,13 @@ pub struct SharedRuntimeArgs { #[arg(long, value_parser = clap::value_parser!(i32).range(-1..), allow_negative_numbers = true)] #[serde(default)] pub max_logprobs: Option, + /// The interval (or buffer size) for streaming in terms of token length. + /// A smaller value (1) makes streaming smoother by sending each token + /// immediately, while a larger value (e.g., 10) reduces host overhead and + /// may increase throughput by batching multiple tokens before sending. + #[arg(long, default_value_t = NonZeroU32::MIN)] + #[serde(default = "default_stream_interval")] + pub stream_interval: NonZeroU32, /// Maximum time to wait for active requests to drain during shutdown. #[arg(long, default_value_t = 0)] #[serde(default)] @@ -547,6 +555,7 @@ impl SharedRuntimeArgs { lora_modules: self.lora_modules, chat_template_content_format: self.chat_template_content_format, max_logprobs: self.max_logprobs, + stream_interval: self.stream_interval, api_server_options, cors, tls, @@ -609,6 +618,7 @@ impl SharedRuntimeArgs { lora_modules: self.lora_modules, chat_template_content_format: self.chat_template_content_format, max_logprobs: self.max_logprobs, + stream_interval: self.stream_interval, api_server_options, cors, tls, @@ -646,6 +656,10 @@ fn default_engine_ready_timeout_secs() -> u64 { 600 } +fn default_stream_interval() -> NonZeroU32 { + NonZeroU32::MIN +} + fn default_cors_wildcard() -> JsonStringList { JsonStringList(vec!["*".to_string()]) } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index aa164bc618..57780027de 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -189,6 +189,7 @@ fn serve_args_forward_python_flags_with_separator() { renderer: Auto, language_model_only: false, max_logprobs: None, + stream_interval: 1, shutdown_timeout: 0, http_timeout_keep_alive: None, chat_template: None, @@ -373,6 +374,34 @@ fn serve_args_forward_profiler_config_to_managed_engine() { ); } +#[test] +fn serve_args_forward_stream_interval_to_frontend_config() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--stream-interval", + "4", + ]) + .unwrap(); + + let Command::Serve(args) = cli.command else { + panic!("expected serve args"); + }; + let frontend_config = args.to_frontend_config("tcp://127.0.0.1:62100".to_string()); + assert_eq!(frontend_config.stream_interval.get(), 4); + + let error = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--stream-interval", + "0", + ]) + .unwrap_err(); + assert_eq!(error.kind(), clap::error::ErrorKind::ValueValidation); +} + #[test] fn serve_args_forward_max_logprobs_to_frontend_and_managed_engine() { let cli = Cli::try_parse_from([ @@ -1030,6 +1059,7 @@ fn frontend_args_accept_json() { renderer: Auto, language_model_only: false, max_logprobs: None, + stream_interval: 1, shutdown_timeout: 0, http_timeout_keep_alive: None, chat_template: None, @@ -1702,6 +1732,7 @@ fn serve_args_accept_handshake_aliases() { renderer: Auto, language_model_only: false, max_logprobs: None, + stream_interval: 1, shutdown_timeout: 0, http_timeout_keep_alive: None, chat_template: None, @@ -1866,6 +1897,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { lora_modules: [], chat_template_content_format: Auto, max_logprobs: None, + stream_interval: 1, api_server_options: ApiServerOptions { enable_log_requests: false, enable_prompt_tokens_details: false, @@ -1961,6 +1993,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { lora_modules: [], chat_template_content_format: Auto, max_logprobs: None, + stream_interval: 1, api_server_options: ApiServerOptions { enable_log_requests: false, enable_prompt_tokens_details: false, @@ -2100,6 +2133,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present lora_modules: [], chat_template_content_format: Auto, max_logprobs: None, + stream_interval: 1, api_server_options: ApiServerOptions { enable_log_requests: false, enable_prompt_tokens_details: false, diff --git a/rust/src/cmd/src/cli/unsupported.rs b/rust/src/cmd/src/cli/unsupported.rs index f65f766e84..0a14d1b841 100644 --- a/rust/src/cmd/src/cli/unsupported.rs +++ b/rust/src/cmd/src/cli/unsupported.rs @@ -359,10 +359,6 @@ pub struct EngineUnsupportedArgs { #[arg(long)] pub collect_detailed_traces: Option, - /// The interval (or buffer size) for streaming in terms of token length. - #[arg(long)] - pub stream_interval: Option, - /// Maximum number of requests that can be in-flight (waiting or running) /// at the same time. When the limit is reached, new requests are rejected /// with HTTP 503 so the client can retry on another instance. diff --git a/rust/src/engine-core-client/examples/external_engine_logprobs.rs b/rust/src/engine-core-client/examples/external_engine_logprobs.rs index d4ec116f17..e0da15623c 100644 --- a/rust/src/engine-core-client/examples/external_engine_logprobs.rs +++ b/rust/src/engine-core-client/examples/external_engine_logprobs.rs @@ -87,8 +87,9 @@ fn build_request( } async fn wait_for_final_output( - mut stream: vllm_engine_core_client::EngineCoreOutputStream, + stream: vllm_engine_core_client::EngineCoreOutputStream, ) -> Result { + let mut stream = stream.into_outputs(); while let Some(output) = stream.next().await { let output = output.context("failed to receive engine-core output")?; if output.finished() { diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index b3516936fe..d8b436c646 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::BTreeMap; +use std::num::NonZeroU32; use std::sync::Arc; use std::time::Duration; @@ -25,10 +26,11 @@ use crate::runtime::{BackgroundShutdownRuntime, build_zmq_runtime}; use crate::transport::{self, ConnectedEngine}; pub(crate) mod imp; +mod output_channel; mod state; mod stream; -pub use stream::{EngineCoreOutputStream, EngineCoreStreamOutput}; +pub use stream::{EngineCoreOutputStream, EngineCoreStreamDelivery, EngineCoreStreamOutput}; /// How the frontend acquires its request/response transport with Python /// `EngineCoreProc`s. @@ -614,6 +616,12 @@ fn validate_lora_capabilities(engines: &[ConnectedEngine]) -> Result<()> { impl EngineCoreClient { /// Add a new request to the engine and return a per-request raw output /// stream. + /// + /// With `sampling_params.stream_interval` above one, outputs after the + /// first are delivered to the stream in batches of at least that many new + /// tokens, and the terminal output is delivered immediately together with + /// any held-back outputs. The stream still yields every raw output in + /// order; batching only reduces how often the consuming task is woken. pub async fn call(&self, mut req: EngineCoreRequest) -> Result { req.client_index = self.config.client_index; req.validate()?; @@ -628,8 +636,15 @@ impl EngineCoreClient { let request_id = req.request_id.clone(); let lora_name = req.lora_request.as_ref().map(|lora| lora.lora_name.clone()); let data_parallel_rank = req.data_parallel_rank; - let (engine_id, rx) = - self.inner.register_request(request_id.clone(), lora_name, data_parallel_rank)?; + let stream_interval = (req.sampling_params.as_ref()) + .and_then(|params| params.stream_interval) + .unwrap_or(NonZeroU32::MIN); + let (engine_id, rx) = self.inner.register_request( + request_id.clone(), + lora_name, + data_parallel_rank, + stream_interval, + )?; // Construct the output stream first before actually sending the request to the engine. // This ensures that cancelling the future will properly clean up the request registration diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index c8abb95bbd..0a4af1a5ee 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::{BTreeMap, BTreeSet}; +use std::num::NonZeroU32; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; @@ -15,7 +16,9 @@ use tracing::{debug, info, trace, warn}; use vllm_metrics::METRICS; use zeromq::RouterSendHalf; -use crate::client::state::{OutputReceiver, RequestRegistry, UtilityReceiver, UtilityRegistry}; +use crate::client::state::{ + OutputReceiver, OutputSender, RequestRegistry, UtilityReceiver, UtilityRegistry, +}; use crate::client::stream::EngineCoreStreamOutput; use crate::client::{AbortCause, AbortRequest}; use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output}; @@ -111,12 +114,13 @@ impl ClientInner { request_id: String, lora_name: Option, data_parallel_rank: Option, + stream_interval: NonZeroU32, ) -> Result<(EngineId, OutputReceiver)> { let mut registry = self.request_reg.lock(); if registry.is_closed() { return Err(self.closed_error()); } - registry.register(request_id, lora_name, data_parallel_rank) + registry.register(request_id, lora_name, data_parallel_rank, stream_interval) } /// Allocate the next utility `call_id` and register its waiting receiver. @@ -159,13 +163,15 @@ impl ClientInner { Ok(registry.abortable_request_ids(request_ids)) } - /// Obtain stream senders for a whole engine output batch with one registry - /// lock acquisition. - pub fn take_senders_for_outputs<'a>( + /// Send a whole engine output batch to the request streams with one + /// registry lock acquisition. Returns the outputs whose requests are no + /// longer tracked. + pub fn send_outputs( &self, - outputs: impl IntoIterator, - ) -> Vec>>> { - self.request_reg.lock().senders_for_outputs(outputs) + outputs: impl IntoIterator, + ) -> Vec { + let mut registry = self.request_reg.lock(); + outputs.into_iter().filter_map(|output| registry.send_output(output)).collect() } /// Remove a batch of requests that have finished or aborted, returning @@ -173,7 +179,7 @@ impl ClientInner { pub fn finish_requests<'a>( &self, request_ids: impl IntoIterator, - ) -> Vec>> { + ) -> Vec { self.request_reg.lock().finish_many(request_ids) } @@ -214,7 +220,7 @@ impl ClientInner { // Notify all ongoing requests that the client is closed. for sender in request_senders { - let _ = sender.send(Err(Error::Shared(persistent_error.clone()))); + sender.send_error(Error::Shared(persistent_error.clone())); } for sender in utility_senders { let _ = sender.send(Err(Error::Shared(persistent_error.clone()))); @@ -413,6 +419,16 @@ pub(crate) async fn run_abort_loop( } } +/// Number of tokens one output accounts for in `vllm:iteration_tokens_total`. +fn output_iteration_tokens(output: &EngineCoreOutput) -> u64 { + let mut tokens = output.new_token_ids.len() as u64; + if let Some(prefill) = &output.prefill_stats { + // The engine emits prefill_stats once, on the first output. + tokens += u64::from(prefill.num_computed_tokens); + } + tokens +} + /// Background loop that listens for engine-core outputs and dispatches them to /// the corresponding request streams based on their `request_id`. pub(crate) async fn run_output_dispatcher_loop( @@ -435,29 +451,20 @@ pub(crate) async fn run_output_dispatcher_loop( match outputs { EngineCoreOutputs::RequestBatch(batch) => { let has_outputs = !batch.outputs.is_empty(); - let mut iteration_tokens = 0_u64; - let senders = inner.take_senders_for_outputs(&batch.outputs); - for (output, sender) in batch.outputs.into_iter().zip(senders) { - let request_id = output.request_id.clone(); - let Some(sender) = sender else { - debug!(request_id, "dropping output for inactive request"); - continue; - }; - - iteration_tokens += output.new_token_ids.len() as u64; - if let Some(prefill) = &output.prefill_stats { - // The engine emits prefill_stats once, on the first output. - iteration_tokens += u64::from(prefill.num_computed_tokens); - } - - let wrapped_output = EngineCoreStreamOutput { - engine_index: batch.engine_index, - timestamp: batch.timestamp, - output, - }; - if sender.send(Ok(wrapped_output)).is_err() { - debug!(request_id, "request output stream receiver dropped"); - } + let mut iteration_tokens: u64 = + batch.outputs.iter().map(output_iteration_tokens).sum(); + let (engine_index, timestamp) = (batch.engine_index, batch.timestamp); + let outputs = batch.outputs.into_iter().map(|output| EngineCoreStreamOutput { + engine_index, + timestamp, + output, + }); + for output in inner.send_outputs(outputs) { + debug!( + request_id = output.request_id, + "dropping output for inactive request" + ); + iteration_tokens -= output_iteration_tokens(&output); } if has_outputs diff --git a/rust/src/engine-core-client/src/client/output_channel.rs b/rust/src/engine-core-client/src/client/output_channel.rs new file mode 100644 index 0000000000..7b58d3dc2f --- /dev/null +++ b/rust/src/engine-core-client/src/client/output_channel.rs @@ -0,0 +1,321 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Per-request output channel between the output dispatcher and +//! [`EngineCoreOutputStream`](crate::EngineCoreOutputStream). +//! +//! With a stream interval above one, the sender holds back the outputs after +//! the first one until they carry at least that many new tokens, then delivers +//! them in one channel message, which the stream yields as one +//! [`EngineCoreStreamDelivery`](crate::EngineCoreStreamDelivery). The request +//! task is then woken once per delivery instead of once per engine step. +//! +//! Outputs are never merged, reordered, or dropped here; the consumer merges a +//! delivery into one output (see `vllm-llm`). Merging here would hide the raw +//! outputs that request metrics observe one by one (per-step timestamps for +//! inter-token latency, prefill stats, lifecycle events), and would have to +//! combine wire-level fields such as undecoded logprobs tensors and opaque +//! values. + +use std::num::NonZeroU32; + +use tokio::sync::mpsc; +use tracing::{debug, warn}; + +use crate::client::stream::{EngineCoreStreamDelivery as Delivery, EngineCoreStreamOutput}; +use crate::error::{Error, Result}; + +/// Create the output channel for one request. +pub(crate) fn output_channel(stream_interval: NonZeroU32) -> (OutputSender, OutputReceiver) { + let (tx, rx) = mpsc::unbounded_channel(); + let sender = OutputSender { + tx, + stream_interval: stream_interval.get() as usize, + delivered_first: false, + buffered: Vec::new(), + buffered_tokens: 0, + }; + (sender, rx) +} + +/// Receiving half of a request's output channel. +pub(crate) type OutputReceiver = mpsc::UnboundedReceiver>; + +/// Sending half of a request's output channel, owned by its registry entry. +#[derive(Debug)] +pub(crate) struct OutputSender { + tx: mpsc::UnboundedSender>, + /// Minimum number of new tokens per delivery after the first one. + stream_interval: usize, + /// Whether the first output has been delivered. + delivered_first: bool, + /// Outputs held back since the last delivery. + buffered: Vec, + /// Number of new tokens in `buffered`. + buffered_tokens: usize, +} + +impl OutputSender { + /// Send one output, or hold it back until the held-back outputs reach the + /// stream interval. The first output and the terminal output are always + /// delivered immediately, the latter together with any held-back outputs. + pub fn send(&mut self, output: EngineCoreStreamOutput) { + let finished = output.finished(); + if self.buffered.is_empty() + && (finished || !self.delivered_first || self.stream_interval == 1) + { + self.delivered_first = true; + self.deliver(Delivery::One(output)); + return; + } + + self.buffered_tokens += output.new_token_ids.len(); + self.buffered.push(output); + if finished || self.buffered_tokens >= self.stream_interval { + self.flush(); + } + } + + /// Deliver the held-back outputs, if any. + pub fn flush(&mut self) { + if self.buffered.is_empty() { + return; + } + self.buffered_tokens = 0; + let outputs = std::mem::take(&mut self.buffered); + self.deliver(Delivery::Many(outputs)); + } + + /// Deliver the held-back outputs, then terminate the stream with `error`. + pub fn send_error(mut self, error: Error) { + self.flush(); + let _ = self.tx.send(Err(error)); + } + + fn deliver(&self, delivery: Delivery) { + if let Err(mpsc::error::SendError(Ok(delivery))) = self.tx.send(Ok(delivery)) { + debug!( + request_id = delivery.first().request_id, + "request output stream receiver dropped" + ); + } + } +} + +impl Drop for OutputSender { + fn drop(&mut self) { + // Every path that ends a live request delivers held-back outputs first + // (terminal output, `flush`, or `send_error`). Outputs left here are lost. + if !self.buffered.is_empty() && !self.tx.is_closed() { + warn!( + request_id = self.buffered[0].request_id, + count = self.buffered.len(), + "dropping request output sender with held-back outputs" + ); + } + } +} + +#[cfg(test)] +mod tests { + use std::num::NonZeroU32; + + use expect_test::expect; + use tokio::sync::mpsc::error::TryRecvError; + + use super::*; + use crate::protocol::output::{EngineCoreFinishReason, EngineCoreOutput}; + + fn output(token_ids: &[u32]) -> EngineCoreStreamOutput { + EngineCoreStreamOutput { + engine_index: 0, + timestamp: 0.0, + output: EngineCoreOutput { + request_id: "req".to_string(), + new_token_ids: token_ids.to_vec(), + ..EngineCoreOutput::default() + }, + } + } + + fn finished(token_ids: &[u32]) -> EngineCoreStreamOutput { + let mut output = output(token_ids); + output.output.finish_reason = Some(EngineCoreFinishReason::Length); + output + } + + fn channel(stream_interval: u32) -> (OutputSender, OutputReceiver) { + output_channel(NonZeroU32::new(stream_interval).unwrap()) + } + + /// Render a delivery as the token ids of its outputs. + fn token_ids(delivery: Delivery) -> Vec> { + delivery.into_iter().map(|output| output.output.new_token_ids).collect() + } + + /// Drain all pending channel messages, rendering each delivery as the + /// token ids of its outputs. + fn deliveries(receiver: &mut OutputReceiver) -> Vec>> { + let mut deliveries = Vec::new(); + loop { + match receiver.try_recv() { + Ok(Ok(delivery)) => deliveries.push(token_ids(delivery)), + Ok(Err(error)) => panic!("unexpected error delivery: {error:?}"), + Err(TryRecvError::Empty | TryRecvError::Disconnected) => return deliveries, + } + } + } + + #[test] + fn stream_interval_one_delivers_every_output() { + let (mut sender, mut receiver) = channel(1); + sender.send(output(&[1])); + sender.send(output(&[2, 3])); + sender.send(finished(&[4])); + + expect![[r#" + [ + [ + [ + 1, + ], + ], + [ + [ + 2, + 3, + ], + ], + [ + [ + 4, + ], + ], + ] + "#]] + .assert_debug_eq(&deliveries(&mut receiver)); + } + + #[test] + fn stream_interval_batches_outputs_after_the_first() { + let (mut sender, mut receiver) = channel(3); + // The first output is delivered immediately; later ones are held back + // until they carry at least 3 tokens, counted across multi-token + // outputs. The terminal output flushes whatever is held back. + for token_ids in [&[1][..], &[2], &[3], &[4, 5], &[6], &[7, 8, 9, 10], &[11]] { + sender.send(output(token_ids)); + } + sender.send(finished(&[12])); + + expect![[r#" + [ + [ + [ + 1, + ], + ], + [ + [ + 2, + ], + [ + 3, + ], + [ + 4, + 5, + ], + ], + [ + [ + 6, + ], + [ + 7, + 8, + 9, + 10, + ], + ], + [ + [ + 11, + ], + [ + 12, + ], + ], + ] + "#]] + .assert_debug_eq(&deliveries(&mut receiver)); + } + + #[test] + fn terminal_output_without_held_back_outputs_is_delivered_alone() { + let (mut sender, mut receiver) = channel(2); + sender.send(output(&[1])); + sender.send(output(&[2, 3])); + sender.send(finished(&[])); + + expect![[r#" + [ + [ + [ + 1, + ], + ], + [ + [ + 2, + 3, + ], + ], + [ + [], + ], + ] + "#]] + .assert_debug_eq(&deliveries(&mut receiver)); + } + + #[test] + fn flush_delivers_held_back_outputs() { + let (mut sender, mut receiver) = channel(4); + sender.send(output(&[1])); + sender.send(output(&[2])); + sender.flush(); + sender.flush(); + + expect![[r#" + [ + [ + [ + 1, + ], + ], + [ + [ + 2, + ], + ], + ] + "#]] + .assert_debug_eq(&deliveries(&mut receiver)); + } + + #[test] + fn error_follows_held_back_outputs() { + let (mut sender, mut receiver) = channel(4); + sender.send(output(&[1])); + sender.send(output(&[2])); + sender.send_error(Error::EngineCoreDead); + + let mut items = Vec::new(); + while let Ok(item) = receiver.try_recv() { + items.push(item.map(|delivery| token_ids(delivery).len()).ok()); + } + // One delivery for the first output, one for the held-back output, + // then the error. + assert_eq!(items, [Some(1), Some(1), None]); + } +} diff --git a/rust/src/engine-core-client/src/client/state.rs b/rust/src/engine-core-client/src/client/state.rs index a8f0fa3c24..9c183b87e4 100644 --- a/rust/src/engine-core-client/src/client/state.rs +++ b/rust/src/engine-core-client/src/client/state.rs @@ -2,12 +2,15 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::num::NonZeroU32; use std::sync::atomic::{AtomicU64, Ordering}; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::oneshot; use tracing::trace; use crate::EngineId; +use crate::client::output_channel::output_channel; +pub(crate) use crate::client::output_channel::{OutputReceiver, OutputSender}; use crate::client::stream::EngineCoreStreamOutput; use crate::error::{Error, Result}; use crate::protocol::output::{EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput}; @@ -15,8 +18,6 @@ use crate::protocol::stats::SchedulerStats; use crate::protocol::utility::UtilityOutput; use crate::transport::ConnectedEngine; -pub type OutputSender = mpsc::UnboundedSender>; -pub type OutputReceiver = mpsc::UnboundedReceiver>; pub type UtilitySender = oneshot::Sender>; pub type UtilityReceiver = oneshot::Receiver>; @@ -137,18 +138,22 @@ impl RequestRegistry { /// When `data_parallel_rank` is provided, the request is routed directly to /// the engine at that rank index, bypassing load balancing. Otherwise /// the engine with the lowest routing score is chosen. + /// + /// `stream_interval` sets how many new tokens the channel batches into one + /// delivery after the first output; see [`OutputSender::send`]. pub fn register( &mut self, request_id: String, lora_name: Option, data_parallel_rank: Option, + stream_interval: NonZeroU32, ) -> Result<(EngineId, OutputReceiver)> { if self.requests.contains_key(&request_id) { return Err(Error::DuplicateRequestId { request_id }); } let engine_id = self.choose_engine_for_request(data_parallel_rank)?; - let (tx, rx) = mpsc::unbounded_channel(); + let (tx, rx) = output_channel(stream_interval); let lora = lora_name.map(|adapter_name| LoraRequestState { adapter_name, phase: LoraPhase::Waiting, @@ -213,17 +218,30 @@ impl RequestRegistry { by_engine } - /// Obtain the stream sender for one output. If it indicates the request is - /// finished, it will be removed from the registry. - pub fn sender_for_output(&mut self, output: &EngineCoreOutput) -> Option { - self.apply_lora_events(output); + /// Send one output to its request stream. If it indicates the request is + /// finished, the request will be removed from the registry. Returns the + /// output back if its request is no longer tracked. + /// + /// The send happens under the registry lock: the sender is stateful (see + /// [`OutputSender::send`]) and owned by the registry entry, and sending on + /// its unbounded channel never blocks. + pub fn send_output( + &mut self, + output: EngineCoreStreamOutput, + ) -> Option { + self.apply_lora_events(&output.output); if output.finished() { - self.remove(output.request_id.as_str()).map(|tracked| tracked.0) + match self.remove(output.request_id.as_str()) { + Some((mut sender, _)) => sender.send(output), + None => return Some(output), + } } else { - self.requests - .get(output.request_id.as_str()) - .map(|tracked| tracked.sender.clone()) + match self.requests.get_mut(output.request_id.as_str()) { + Some(tracked) => tracked.sender.send(output), + None => return Some(output), + } } + None } /// Advance the request's LoRA scheduling phase from the engine-core events @@ -267,17 +285,8 @@ impl RequestRegistry { (running, waiting) } - /// Obtain stream senders for a whole engine output batch under one - /// registry lock. Finished outputs are removed before returning. - pub fn senders_for_outputs<'a>( - &mut self, - outputs: impl IntoIterator, - ) -> Vec> { - outputs.into_iter().map(|output| self.sender_for_output(output)).collect() - } - /// Remove a batch of requests that have finished or aborted, returning - /// their stream senders. + /// their stream senders after delivering any held-back outputs. pub fn finish_many<'a>( &mut self, request_ids: impl IntoIterator, @@ -285,6 +294,10 @@ impl RequestRegistry { request_ids .into_iter() .filter_map(|request_id| self.remove(request_id.as_str()).map(|tracked| tracked.0)) + .map(|mut sender| { + sender.flush(); + sender + }) .collect() } @@ -325,7 +338,7 @@ impl RequestRegistry { ) -> Vec { let mut aborted = Vec::new(); for request_id in request_ids { - let Some((sender, engine_id)) = self.remove(request_id) else { + let Some((mut sender, engine_id)) = self.remove(request_id) else { continue; }; let output = EngineCoreStreamOutput { @@ -337,7 +350,7 @@ impl RequestRegistry { ..EngineCoreOutput::default() }, }; - let _ = sender.send(Ok(output)); + sender.send(output); aborted.push(request_id.clone()); } aborted @@ -469,11 +482,13 @@ impl UtilityRegistry { #[cfg(test)] mod tests { use std::collections::BTreeSet; + use std::num::NonZeroU32; use crate::EngineId; use crate::client::state::{ - EngineLoadSnapshot, EngineRoutingState, RequestRegistry, UtilityRegistry, + EngineLoadSnapshot, EngineRoutingState, OutputReceiver, RequestRegistry, UtilityRegistry, }; + use crate::client::stream::EngineCoreStreamOutput; use crate::mock_engine::default_ready_response; use crate::protocol::output::{ EngineCoreEvent, EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput, @@ -508,6 +523,14 @@ mod tests { } } + fn stream_output(output: EngineCoreOutput) -> EngineCoreStreamOutput { + EngineCoreStreamOutput { + engine_index: 0, + timestamp: 0.0, + output, + } + } + fn adapter_names(values: &[&str]) -> BTreeSet { values.iter().map(|name| (*name).to_string()).collect() } @@ -515,8 +538,9 @@ mod tests { #[test] fn registry_rejects_duplicate_request_ids() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); - registry.register("req-1".to_string(), None, None).unwrap(); - let error = registry.register("req-1".to_string(), None, None).unwrap_err(); + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap(); + let error = + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap_err(); assert!(matches!( error, crate::error::Error::DuplicateRequestId { request_id } if request_id == "req-1" @@ -526,25 +550,99 @@ mod tests { #[test] fn registry_removes_finished_request_on_output() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); - registry.register("req-1".to_string(), None, None).unwrap(); + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap(); - let sender = registry.sender_for_output(&EngineCoreOutput { + let unrouted = registry.send_output(stream_output(EngineCoreOutput { request_id: "req-1".to_string(), finish_reason: Some(EngineCoreFinishReason::Length), ..Default::default() - }); + })); - assert!(sender.is_some()); + assert!(unrouted.is_none()); assert!(!registry.contains("req-1")); } + /// Drain the outputs that are ready on `receiver` without waiting, + /// rendering each as its token ids and finish reason. + fn ready_outputs( + receiver: &mut OutputReceiver, + ) -> Vec<(Vec, Option)> { + let mut outputs = Vec::new(); + while let Ok(delivery) = receiver.try_recv() { + for output in delivery.unwrap() { + outputs.push((output.output.new_token_ids, output.output.finish_reason)); + } + } + outputs + } + + #[test] + fn registry_abort_delivers_held_back_outputs_before_terminal_output() { + let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); + let (_, mut receiver) = registry + .register("req-1".to_string(), None, None, NonZeroU32::new(4).unwrap()) + .unwrap(); + + for token_id in 1..=3 { + let unrouted = registry.send_output(stream_output(EngineCoreOutput { + request_id: "req-1".to_string(), + new_token_ids: vec![token_id], + ..Default::default() + })); + assert!(unrouted.is_none()); + } + // Only the first output is delivered; the next two are held back + // below the stream interval. + expect_test::expect![[r#" + [ + ( + [ + 1, + ], + None, + ), + ] + "#]] + .assert_debug_eq(&ready_outputs(&mut receiver)); + + registry.abort_many(&["req-1".to_string()], 0.0); + expect_test::expect![[r#" + [ + ( + [ + 2, + ], + None, + ), + ( + [ + 3, + ], + None, + ), + ( + [], + Some( + Abort, + ), + ), + ] + "#]] + .assert_debug_eq(&ready_outputs(&mut receiver)); + } + #[test] fn registry_tracks_lora_phases_from_engine_events() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); registry - .register("req-lora".to_string(), Some("adapter-a".to_string()), None) + .register( + "req-lora".to_string(), + Some("adapter-a".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); - registry.register("req-plain".to_string(), None, None).unwrap(); + registry.register("req-plain".to_string(), None, None, NonZeroU32::MIN).unwrap(); // Registered but not yet scheduled: counted as waiting. The non-LoRA // request never shows up. @@ -554,7 +652,7 @@ mod tests { ); // Queued then scheduled in one output: running. - drop(registry.sender_for_output(&output_with_events( + registry.send_output(stream_output(output_with_events( "req-lora", &[EngineCoreEventType::Queued, EngineCoreEventType::Scheduled], None, @@ -565,7 +663,7 @@ mod tests { ); // Preempted: back to waiting. - drop(registry.sender_for_output(&output_with_events( + registry.send_output(stream_output(output_with_events( "req-lora", &[EngineCoreEventType::Preempted], None, @@ -576,7 +674,7 @@ mod tests { ); // Finished: dropped from tracking entirely. - drop(registry.sender_for_output(&output_with_events( + registry.send_output(stream_output(output_with_events( "req-lora", &[EngineCoreEventType::Scheduled], Some(EngineCoreFinishReason::Stop), @@ -591,18 +689,33 @@ mod tests { fn registry_unions_lora_adapters_across_requests() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); registry - .register("req-a1".to_string(), Some("adapter-a".to_string()), None) + .register( + "req-a1".to_string(), + Some("adapter-a".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); registry - .register("req-a2".to_string(), Some("adapter-a".to_string()), None) + .register( + "req-a2".to_string(), + Some("adapter-a".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); registry - .register("req-b".to_string(), Some("adapter-b".to_string()), None) + .register( + "req-b".to_string(), + Some("adapter-b".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); // One of adapter-a's requests starts running while the other waits: // the adapter appears in both sets. - drop(registry.sender_for_output(&output_with_events( + registry.send_output(stream_output(output_with_events( "req-a1", &[EngineCoreEventType::Scheduled], None, @@ -620,7 +733,7 @@ mod tests { fn registry_counts_only_active_lora_requests() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); - registry.register("req-plain".to_string(), None, None).unwrap(); + registry.register("req-plain".to_string(), None, None, NonZeroU32::MIN).unwrap(); assert_eq!(registry.active_lora_requests(), 0); assert_eq!( registry.lora_adapter_states(), @@ -632,6 +745,7 @@ mod tests { "req-lora-a".to_string(), Some("adapter-a".to_string()), None, + NonZeroU32::MIN, ) .unwrap(); registry @@ -639,6 +753,7 @@ mod tests { "req-lora-b".to_string(), Some("adapter-b".to_string()), None, + NonZeroU32::MIN, ) .unwrap(); assert_eq!(registry.active_lora_requests(), 2); @@ -661,7 +776,12 @@ mod tests { fn registry_clears_lora_count_on_close() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); registry - .register("req-lora".to_string(), Some("adapter-a".to_string()), None) + .register( + "req-lora".to_string(), + Some("adapter-a".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); assert_eq!(registry.active_lora_requests(), 1); @@ -677,7 +797,12 @@ mod tests { fn registry_drops_lora_tracking_on_abort() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); registry - .register("req-lora".to_string(), Some("adapter-a".to_string()), None) + .register( + "req-lora".to_string(), + Some("adapter-a".to_string()), + None, + NonZeroU32::MIN, + ) .unwrap(); drop(registry.finish_many(&["req-lora".to_string()])); @@ -691,8 +816,8 @@ mod tests { #[test] fn registry_closes_all_requests_on_failure() { let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]); - registry.register("req-1".to_string(), None, None).unwrap(); - registry.register("req-2".to_string(), None, None).unwrap(); + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap(); + registry.register("req-2".to_string(), None, None, NonZeroU32::MIN).unwrap(); let senders = registry.close(); @@ -708,9 +833,12 @@ mod tests { connected_engine(engine_0.clone()), connected_engine(engine_1.clone()), ]); - let (chosen_0, _) = registry.register("req-1".to_string(), None, None).unwrap(); - let (chosen_1, _) = registry.register("req-2".to_string(), None, None).unwrap(); - let (chosen_0_again, _) = registry.register("req-3".to_string(), None, None).unwrap(); + let (chosen_0, _) = + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap(); + let (chosen_1, _) = + registry.register("req-2".to_string(), None, None, NonZeroU32::MIN).unwrap(); + let (chosen_0_again, _) = + registry.register("req-3".to_string(), None, None, NonZeroU32::MIN).unwrap(); assert_eq!(chosen_0, engine_0); assert_eq!(chosen_1, engine_1); @@ -737,9 +865,12 @@ mod tests { connected_engine(engine_1.clone()), ]); - let (chosen_0, _) = registry.register("req-1".to_string(), None, None).unwrap(); - let (chosen_1, _) = registry.register("req-2".to_string(), None, None).unwrap(); - let (chosen_0_again, _) = registry.register("req-3".to_string(), None, None).unwrap(); + let (chosen_0, _) = + registry.register("req-1".to_string(), None, None, NonZeroU32::MIN).unwrap(); + let (chosen_1, _) = + registry.register("req-2".to_string(), None, None, NonZeroU32::MIN).unwrap(); + let (chosen_0_again, _) = + registry.register("req-3".to_string(), None, None, NonZeroU32::MIN).unwrap(); assert_eq!(chosen_0, engine_0); assert_eq!(chosen_1, engine_1); @@ -809,7 +940,8 @@ mod tests { } )); - let (chosen, _) = registry.register("req-stats".to_string(), None, None).unwrap(); + let (chosen, _) = + registry.register("req-stats".to_string(), None, None, NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_1); } @@ -824,7 +956,9 @@ mod tests { for (rank, running) in counts.into_iter().enumerate() { for index in 0..running { let request_id = format!("old-{rank}-{index}"); - registry.register(request_id.clone(), None, Some(rank as u32)).unwrap(); + registry + .register(request_id.clone(), None, Some(rank as u32), NonZeroU32::MIN) + .unwrap(); old_requests.push(request_id); } registry.apply_scheduler_counts( @@ -839,7 +973,8 @@ mod tests { let mut distribution = [0; 2]; for index in 0..14 { - let (engine, _) = registry.register(format!("new-{index}"), None, None).unwrap(); + let (engine, _) = + registry.register(format!("new-{index}"), None, None, NonZeroU32::MIN).unwrap(); distribution[engine.engine_index().unwrap() as usize] += 1; } distributions.push(distribution); @@ -865,7 +1000,8 @@ mod tests { ); } let request_id = format!("req-{index}"); - let (engine, _) = registry.register(request_id.clone(), None, None).unwrap(); + let (engine, _) = + registry.register(request_id.clone(), None, None, NonZeroU32::MIN).unwrap(); chosen.push(engine.engine_index().unwrap()); drop(registry.finish_many(&[request_id])); } @@ -890,12 +1026,15 @@ mod tests { }, ); for index in 0..4 { - registry.register(format!("pinned-{index}"), None, Some(0)).unwrap(); + registry + .register(format!("pinned-{index}"), None, Some(0), NonZeroU32::MIN) + .unwrap(); } // Even unchanged counts replace the optimistic estimate for that engine. registry.apply_scheduler_counts(0, snapshot); - let (after_refresh, _) = registry.register("after-refresh".into(), None, None).unwrap(); + let (after_refresh, _) = + registry.register("after-refresh".into(), None, None, NonZeroU32::MIN).unwrap(); // An older, empty snapshot must not erase the five admitted requests. for rank in [0, 1] { registry.apply_scheduler_counts( @@ -906,7 +1045,8 @@ mod tests { }, ); } - let (after_empty, _) = registry.register("after-empty".into(), None, None).unwrap(); + let (after_empty, _) = + registry.register("after-empty".into(), None, None, NonZeroU32::MIN).unwrap(); expect_test::expect!["(Some(0), Some(1))"].assert_eq(&format!( "{:?}", (after_refresh.engine_index(), after_empty.engine_index()) @@ -927,10 +1067,11 @@ mod tests { }, ); } - registry.register("pinned".into(), None, Some(9)).unwrap(); + registry.register("pinned".into(), None, Some(9), NonZeroU32::MIN).unwrap(); let mut chosen = Vec::new(); for index in 0..3 { - let (engine, _) = registry.register(format!("auto-{index}"), None, None).unwrap(); + let (engine, _) = + registry.register(format!("auto-{index}"), None, None, NonZeroU32::MIN).unwrap(); chosen.push(engine.engine_index().unwrap()); } expect_test::expect!["[2, 5, 2]"].assert_eq(&format!("{chosen:?}")); @@ -948,15 +1089,18 @@ mod tests { ]); // Explicitly target rank 2 (third engine). - let (chosen, _) = registry.register("req-1".to_string(), None, Some(2)).unwrap(); + let (chosen, _) = + registry.register("req-1".to_string(), None, Some(2), NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_2); // Explicitly target rank 0 (first engine). - let (chosen, _) = registry.register("req-2".to_string(), None, Some(0)).unwrap(); + let (chosen, _) = + registry.register("req-2".to_string(), None, Some(0), NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_0); // Explicitly target rank 1. - let (chosen, _) = registry.register("req-3".to_string(), None, Some(1)).unwrap(); + let (chosen, _) = + registry.register("req-3".to_string(), None, Some(1), NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_1); } @@ -970,11 +1114,12 @@ mod tests { ]); // Load-balance: first two go to engine_0 and engine_1. - registry.register("req-lb-0".to_string(), None, None).unwrap(); + registry.register("req-lb-0".to_string(), None, None, NonZeroU32::MIN).unwrap(); // Now engine_0 has 1 in-flight. Without dp_rank, next would go to engine_1. // But with dp_rank=0, it should still go to engine_0. - let (chosen, _) = registry.register("req-dp".to_string(), None, Some(0)).unwrap(); + let (chosen, _) = + registry.register("req-dp".to_string(), None, Some(0), NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_0); } @@ -985,7 +1130,9 @@ mod tests { connected_engine(EngineId::from_engine_index(1)), ]); - let error = registry.register("req-1".to_string(), None, Some(2)).unwrap_err(); + let error = registry + .register("req-1".to_string(), None, Some(2), NonZeroU32::MIN) + .unwrap_err(); assert!(matches!( error, crate::error::Error::InvalidDataParallelRank { @@ -1000,10 +1147,13 @@ mod tests { let engine_3 = EngineId::from_engine_index(3); let mut registry = RequestRegistry::new(&[connected_engine(engine_3.clone())]); - let (chosen, _) = registry.register("req-ok".to_string(), None, Some(3)).unwrap(); + let (chosen, _) = + registry.register("req-ok".to_string(), None, Some(3), NonZeroU32::MIN).unwrap(); assert_eq!(chosen, engine_3); - let error = registry.register("req-bad".to_string(), None, Some(0)).unwrap_err(); + let error = registry + .register("req-bad".to_string(), None, Some(0), NonZeroU32::MIN) + .unwrap_err(); assert!(matches!( error, crate::error::Error::InvalidDataParallelRank { diff --git a/rust/src/engine-core-client/src/client/stream.rs b/rust/src/engine-core-client/src/client/stream.rs index ba408f2c87..6f363f6101 100644 --- a/rust/src/engine-core-client/src/client/stream.rs +++ b/rust/src/engine-core-client/src/client/stream.rs @@ -5,8 +5,9 @@ use std::ops::Deref; use std::pin::Pin; use std::task::{Context, Poll}; -use futures::Stream; use futures::stream::FusedStream; +use futures::{Stream, StreamExt as _}; +use itertools::Either; use thiserror_ext::AsReport as _; use tokio::sync::mpsc; use tracing::{debug, error, warn}; @@ -40,6 +41,68 @@ impl Deref for EngineCoreStreamOutput { } } +/// Consecutive outputs of one request, delivered to its stream together and in +/// order. Outputs are grouped by the request's `stream_interval` (see +/// [`EngineCoreClient::call`]) but not merged; consumers merge them after +/// observing each raw output. +/// +/// [`EngineCoreClient::call`]: crate::EngineCoreClient::call +#[derive(Debug)] +#[allow( + clippy::large_enum_variant, + reason = "single outputs stay inline, avoiding an allocation per engine step" +)] +pub enum EngineCoreStreamDelivery { + /// A single output, sent without being held back. + One(EngineCoreStreamOutput), + /// Outputs held back by the stream interval, in order. Never empty. + Many(Vec), +} + +impl EngineCoreStreamDelivery { + /// Return the first output. + pub fn first(&self) -> &EngineCoreStreamOutput { + match self { + Self::One(output) => output, + Self::Many(outputs) => outputs.first().expect("deliveries are never empty"), + } + } + + /// Return the last output, the only one that can be terminal. + pub fn last(&self) -> &EngineCoreStreamOutput { + match self { + Self::One(output) => output, + Self::Many(outputs) => outputs.last().expect("deliveries are never empty"), + } + } + + /// Iterate over the outputs in order. + pub fn iter( + &self, + ) -> Either< + std::iter::Once<&EngineCoreStreamOutput>, + std::slice::Iter<'_, EngineCoreStreamOutput>, + > { + match self { + Self::One(output) => Either::Left(std::iter::once(output)), + Self::Many(outputs) => Either::Right(outputs.iter()), + } + } +} + +impl IntoIterator for EngineCoreStreamDelivery { + type Item = EngineCoreStreamOutput; + type IntoIter = + Either, std::vec::IntoIter>; + + fn into_iter(self) -> Self::IntoIter { + match self { + Self::One(output) => Either::Left(std::iter::once(output)), + Self::Many(outputs) => Either::Right(outputs.into_iter()), + } + } +} + /// Stream of raw engine-core outputs for one request. /// /// The stream yields only [`EngineCoreStreamOutput`] values whose embedded @@ -79,10 +142,21 @@ impl EngineCoreOutputStream { pub fn engine_index(&self) -> u32 { self.engine_index } + + /// Flatten the stream into one item per raw output, for consumers that do + /// not care how outputs are grouped into deliveries. + pub fn into_outputs(self) -> impl Stream> + Send + Unpin { + self.flat_map(|item| { + futures::stream::iter(match item { + Ok(delivery) => Either::Left(delivery.into_iter().map(Ok)), + Err(error) => Either::Right(std::iter::once(Err(error))), + }) + }) + } } impl Stream for EngineCoreOutputStream { - type Item = Result; + type Item = Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { if self.is_terminated() { @@ -93,10 +167,11 @@ impl Stream for EngineCoreOutputStream { Poll::Pending => Poll::Pending, Poll::Ready(Some(item)) => { match &item { - Ok(output) => { + Ok(delivery) => { // If the output indicates the request is finished, mark the stream as // terminated with cleanly-finished state and expect no more outputs to - // come. + // come. Only the last output of a delivery can be terminal. + let output = delivery.last(); if output.finished() { if output.finish_reason == Some(EngineCoreFinishReason::Error) { error!( diff --git a/rust/src/engine-core-client/src/lib.rs b/rust/src/engine-core-client/src/lib.rs index 7b60206194..8f79e227de 100644 --- a/rust/src/engine-core-client/src/lib.rs +++ b/rust/src/engine-core-client/src/lib.rs @@ -14,7 +14,7 @@ mod transport; pub use client::{ AbortCause, CoordinatorMode, EngineCoreClient, EngineCoreClientConfig, EngineCoreOutputStream, - EngineCoreStreamOutput, TransportMode, + EngineCoreStreamDelivery, EngineCoreStreamOutput, TransportMode, }; pub use error::{Error, Result}; pub use transport::{ENGINE_CORE_DEAD_SENTINEL, EngineId}; diff --git a/rust/src/engine-core-client/src/protocol/sampling.rs b/rust/src/engine-core-client/src/protocol/sampling.rs index 577dcb651f..51e80d0f66 100644 --- a/rust/src/engine-core-client/src/protocol/sampling.rs +++ b/rust/src/engine-core-client/src/protocol/sampling.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::{BTreeSet, HashMap}; +use std::num::NonZeroU32; use serde::{Deserialize, Serialize}; use serde_default::DefaultFromSerde; @@ -161,6 +162,11 @@ pub struct EngineCoreSamplingParams { /// Number of prompt tokens to skip from returned routed-expert data. /// A value of zero returns routing data for the entire prompt. pub routed_experts_prompt_start: u32, + /// Number of newly generated tokens to batch into each streamed output. + /// The first and final outputs are always emitted immediately. Engine-core + /// ignores it; [`EngineCoreClient::call`](crate::EngineCoreClient::call) + /// batches output deliveries by it. + pub stream_interval: Option, } impl EngineCoreSamplingParams { @@ -195,6 +201,7 @@ impl EngineCoreSamplingParams { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } } } diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index dbd3fa69ce..0ad58b0271 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -755,8 +755,8 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { ) .await; - let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap(); - let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap(); + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap().into_outputs(); + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap().into_outputs(); let final_1 = timeout(Duration::from_secs(1), stream_1.next()) .await @@ -778,7 +778,7 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { tokio::time::sleep(Duration::from_millis(100)).await; - let mut stream_3 = client.call(sample_request_with_id("req-3")).await.unwrap(); + let mut stream_3 = client.call(sample_request_with_id("req-3")).await.unwrap().into_outputs(); let final_3 = timeout(Duration::from_secs(1), stream_3.next()) .await .unwrap() @@ -933,7 +933,7 @@ async fn coordinator_accepts_stats_only_outputs() { ) .await; - let mut stream = client.call(sample_request_with_id("req-stats")).await.unwrap(); + let mut stream = client.call(sample_request_with_id("req-stats")).await.unwrap().into_outputs(); let final_output = timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap(); assert_eq!(final_output.request_id, "req-stats"); @@ -1026,8 +1026,8 @@ async fn client_fail_closes_when_main_output_path_receives_dp_control() { assert!(client.ready_responses()[0].max_model_len > 0); assert_eq!(client.vllm_version(), "test-vllm-version"); - let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap(); - let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap(); + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap().into_outputs(); + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap().into_outputs(); let error_2 = timeout(Duration::from_secs(1), stream_2.next()) .await @@ -1103,7 +1103,7 @@ async fn duplicate_request_ids_are_rejected_without_sending_a_second_add() { ) .await; - let mut stream = client.call(sample_request()).await.unwrap(); + let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); let error = match client.call(sample_request()).await { Ok(_) => panic!("expected duplicate request error"), Err(error) => error, @@ -1170,7 +1170,7 @@ async fn finished_requests_without_final_output_is_treated_as_unexpected_close() ) .await; - let mut stream = client.call(sample_request()).await.unwrap(); + let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); let error = timeout(Duration::from_secs(1), stream.next()) .await .unwrap() @@ -1233,7 +1233,7 @@ async fn dropping_a_live_stream_triggers_abort() { ) .await; - let mut stream = client.call(sample_request()).await.unwrap(); + let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); let first = timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap(); assert_eq!(first.new_token_ids, vec![99]); drop(stream); @@ -1316,7 +1316,7 @@ async fn dropping_multiple_live_streams_aborts_all_in_a_burst() { // emits outputs, then drain the first token from each stream. let mut streams = Vec::new(); for id in request_ids { - streams.push(client.call(sample_request_with_id(id)).await.unwrap()); + streams.push(client.call(sample_request_with_id(id)).await.unwrap().into_outputs()); } for stream in streams.iter_mut() { let first = timeout(Duration::from_secs(1), stream.next()).await.unwrap().unwrap().unwrap(); @@ -1363,8 +1363,8 @@ async fn dispatcher_failure_propagates_to_streams_and_future_calls() { ) .await; - let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap(); - let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap(); + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap().into_outputs(); + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap().into_outputs(); let error_1 = timeout(Duration::from_secs(1), stream_1.next()) .await @@ -1862,7 +1862,7 @@ async fn client_decodes_multipart_logprob_outputs() { ) .await; - let stream = client.call(sample_request()).await.unwrap(); + let stream = client.call(sample_request()).await.unwrap().into_outputs(); let outputs = stream.collect::>().await; assert_eq!(outputs.len(), 1); @@ -1950,7 +1950,7 @@ async fn client_sends_large_multimodal_tensor_as_aux_frame() { ) .await; - let outputs = client.call(request).await.unwrap().collect::>().await; + let outputs = client.call(request).await.unwrap().into_outputs().collect::>().await; assert_eq!(outputs.len(), 1); assert!(outputs[0].is_ok()); @@ -2079,8 +2079,8 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() { assert_eq!(client.ready_responses().len(), 2); assert_eq!(client.engine_identities()[0], b"engine-0"); - let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap(); - let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap(); + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap().into_outputs(); + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap().into_outputs(); assert_eq!( timeout(Duration::from_secs(1), engine_0_seen_rx.recv()).await.unwrap().unwrap(), "req-1" @@ -2100,7 +2100,7 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() { assert_eq!(final_1.new_token_ids, vec![10]); assert_eq!(final_1.finish_reason, Some(EngineCoreFinishReason::Length)); - let mut stream_3 = client.call(sample_request_with_id("req-3")).await.unwrap(); + let mut stream_3 = client.call(sample_request_with_id("req-3")).await.unwrap().into_outputs(); assert_eq!( timeout(Duration::from_secs(1), engine_0_seen_rx.recv()).await.unwrap().unwrap(), "req-3" @@ -2271,8 +2271,8 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() { assert!(client.is_sleeping().await.unwrap()); - let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap(); - let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap(); + let mut stream_1 = client.call(sample_request_with_id("req-1")).await.unwrap().into_outputs(); + let mut stream_2 = client.call(sample_request_with_id("req-2")).await.unwrap().into_outputs(); client .abort(&[ @@ -2803,6 +2803,7 @@ fn python_msgpack_fixtures_match_rust_encoding() { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, }, ); @@ -3266,7 +3267,7 @@ async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends .await; tokio::time::sleep(Duration::from_millis(50)).await; - let mut stream = client.call(sample_request()).await.unwrap(); + let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); let wakeup = timeout( Duration::from_secs(1), @@ -3352,7 +3353,7 @@ async fn bootstrapped_external_coordinator_running_state_suppresses_wakeup() { send_external_coordinator_publish(&mut stats_socket, &(Value::Nil, 5_u32, true)).await; tokio::time::sleep(Duration::from_millis(50)).await; - let mut stream = client.call(sample_request()).await.unwrap(); + let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); assert!( timeout( diff --git a/rust/src/llm/src/lib.rs b/rust/src/llm/src/lib.rs index 015624b409..7e9b5e0ba5 100644 --- a/rust/src/llm/src/lib.rs +++ b/rust/src/llm/src/lib.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::num::NonZeroU32; + use tracing::Span; use vllm_engine_core_client::EngineCoreClient; @@ -34,6 +36,7 @@ use crate::request_metrics::RequestMetricsTracker; pub struct Llm { client: EngineCoreClient, randomize_request_id: bool, + stream_interval: NonZeroU32, stats_logger: Option, inflight: InflightRequests, } @@ -45,6 +48,7 @@ impl Llm { Self { client, randomize_request_id: true, + stream_interval: NonZeroU32::MIN, stats_logger: None, inflight: InflightRequests::new(), } @@ -71,6 +75,14 @@ impl Llm { self } + /// Set the frontend-level stream interval: the minimum number of newly + /// generated tokens batched into each streamed output after the first one. + /// A request's own `stream_interval` can only raise it. + pub fn with_stream_interval(mut self, stream_interval: NonZeroU32) -> Self { + self.stream_interval = stream_interval; + self + } + /// Expose the underlying engine-core client for low-level utility/admin /// calls. pub fn engine_core_client(&self) -> &EngineCoreClient { @@ -79,7 +91,16 @@ impl Llm { /// Submit one tokenized generate request and return a per-request output /// stream. - pub async fn generate(&self, req: GenerateRequest) -> Result { + pub async fn generate(&self, mut req: GenerateRequest) -> Result { + // Clamp the request's stream interval to the frontend-level one, like + // Python. Engine-core ignores it; the client batches deliveries by it. + let stream_interval = (req.sampling_params.stream_interval) + .map_or(self.stream_interval, |interval| { + interval.max(self.stream_interval) + }); + req.sampling_params.stream_interval = + (stream_interval > NonZeroU32::MIN).then_some(stream_interval); + let prepared = req.prepare(self.randomize_request_id)?; let prompt_token_ids = prepared.prompt_token_ids().into(); let external_request_id = prepared diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index 4dde96f2e2..b6ccb77175 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -16,7 +16,9 @@ use vllm_engine_core_client::protocol::output::{ }; use vllm_engine_core_client::protocol::sampling_mask::SamplingMask; use vllm_engine_core_client::protocol::tensor::WireNdArray; -use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream}; +use vllm_engine_core_client::{ + AbortCause, EngineCoreOutputStream, EngineCoreStreamDelivery, EngineCoreStreamOutput, +}; use crate::error::Result; use crate::inflight::RequestGuard; @@ -195,6 +197,35 @@ impl GenerateOutput { pub fn finished(&self) -> bool { self.finish_reason.is_some() } + + /// Append a later output of the same request, as if both were produced in + /// one step. Per-token data is concatenated; per-output fields take the + /// later output's value, like Python's `RequestState.make_request_output()` + /// does for the output that flushes a stream interval. + fn merge(&mut self, next: GenerateOutput) { + debug_assert!( + next.prompt_info.is_none(), + "only the first output carries prompt info" + ); + self.token_ids.extend(next.token_ids); + if let Some(next_logprobs) = next.logprobs { + match &mut self.logprobs { + Some(logprobs) => logprobs.positions.extend(next_logprobs.positions), + None => self.logprobs = Some(next_logprobs), + } + } + if let Some(next_mask) = next.sampling_mask { + match &mut self.sampling_mask { + Some(mask) => mask.rows.extend(next_mask.rows), + None => self.sampling_mask = Some(next_mask), + } + } + self.cached_token_count = self.cached_token_count.max(next.cached_token_count); + self.finish_reason = next.finish_reason; + self.kv_transfer_params = next.kv_transfer_params; + self.ec_transfer_params = next.ec_transfer_params; + self.spec_decode_metrics = next.spec_decode_metrics; + } } #[cfg(any(test, feature = "test-util"))] @@ -266,18 +297,27 @@ impl GenerateOutputStream { pub fn request_id(&self) -> &str { self.raw_stream.request_id() } -} -impl Stream for GenerateOutputStream { - type Item = Result; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let raw = match ready!(Pin::new(&mut self.raw_stream).poll_next(cx)) { - Some(Ok(raw)) => raw, - Some(Err(error)) => return Poll::Ready(Some(Err(error.into()))), - None => return Poll::Ready(None), + /// Merge the outputs of one delivery into one output, like Python + /// `RequestState.make_request_output()` does for outputs held back by + /// `stream_interval`. Request metrics still observe every raw output. + fn merge_delivery(&mut self, delivery: EngineCoreStreamDelivery) -> Result { + // Match rather than iterate, so the common single output is converted + // without first being moved into an iterator. + let mut outputs = match delivery { + EngineCoreStreamDelivery::One(output) => return self.convert_output(output), + EngineCoreStreamDelivery::Many(outputs) => outputs.into_iter(), }; + let first = outputs.next().expect("deliveries are never empty"); + let mut merged = self.convert_output(first)?; + for raw in outputs { + merged.merge(self.convert_output(raw)?); + } + Ok(merged) + } + /// Convert one raw engine-core output. + fn convert_output(&mut self, raw: EngineCoreStreamOutput) -> Result { let received_at = current_unix_timestamp_secs(); self.request_metrics.observe_output(raw.timestamp, received_at, &raw.output); @@ -301,11 +341,11 @@ impl Stream for GenerateOutputStream { if let Some(mask) = sampling_mask.as_ref() && mask.rows.len() != raw.new_token_ids.len() { - return Poll::Ready(Some(Err(crate::Error::SamplingMaskTokenCountMismatch { + return Err(crate::Error::SamplingMaskTokenCountMismatch { request_id: raw.request_id, token_count: raw.new_token_ids.len(), row_count: mask.rows.len(), - }))); + }); } let cached_token_count = raw .prefill_stats @@ -331,7 +371,20 @@ impl Stream for GenerateOutputStream { spec_decode_metrics: raw.spec_decode_metrics, }; - Poll::Ready(Some(Ok(output))) + Ok(output) + } +} + +impl Stream for GenerateOutputStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let delivery = match ready!(Pin::new(&mut self.raw_stream).poll_next(cx)) { + Some(Ok(delivery)) => delivery, + Some(Err(error)) => return Poll::Ready(Some(Err(error.into()))), + None => return Poll::Ready(None), + }; + Poll::Ready(Some(self.merge_delivery(delivery))) } } diff --git a/rust/src/llm/tests/generate.rs b/rust/src/llm/tests/generate.rs index 813171555b..e1205ede99 100644 --- a/rust/src/llm/tests/generate.rs +++ b/rust/src/llm/tests/generate.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::BTreeSet; +use std::num::NonZeroU32; use std::sync::{Arc, Once}; use std::time::Duration; @@ -308,6 +309,111 @@ async fn generate_streams_outputs() { llm.shutdown().await.unwrap(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn generate_merges_outputs_by_stream_interval() { + init_tracing(); + let ipc = IpcNamespace::new().unwrap(); + let handshake_address = ipc.handshake_endpoint(); + let engine_id = b"engine-stream-interval".to_vec(); + + // One logprobs position per token, so merged outputs can be checked + // against the logprobs of their concatenated tokens. + fn token_logprobs(token_ids: &[u32]) -> Logprobs { + Logprobs { + positions: token_ids + .iter() + .map(|&token_id| PositionLogprobs { + entries: vec![TokenLogprob { + token_id, + logprob: -0.5, + rank: 1, + }], + }) + .collect(), + } + } + + let (shutdown_tx, engine_task) = spawn_mock_engine_task( + handshake_address.clone(), + engine_id.clone(), + |dealer, push| { + Box::pin(async move { + let add = recv_engine_message(dealer).await; + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); + + // One engine step per batch. + let steps: [(&[u32], Option); 6] = [ + (&[1], None), + (&[2], None), + (&[3], None), + (&[4, 5], None), + (&[6], None), + (&[7], Some(EngineCoreFinishReason::Length)), + ]; + for (token_ids, finish_reason) in steps { + send_outputs( + push, + RequestBatchOutputs { + outputs: vec![request_output_with_logprobs( + &request.request_id, + token_ids.to_vec(), + finish_reason, + Some(token_logprobs(token_ids)), + None, + )], + ..Default::default() + } + .into(), + ) + .await; + } + }) + }, + ); + + // The request's interval raises the frontend-level one. + let llm = connect_async_llm_with_ipc(handshake_address, 7, "test-model", &ipc) + .await + .with_stream_interval(NonZeroU32::new(2).unwrap()); + let mut request = sample_generate_request("req-stream-interval", 7); + request.sampling_params.stream_interval = NonZeroU32::new(3); + let outputs: Vec<_> = llm + .generate(request) + .await + .unwrap() + .map(|output| output.unwrap()) + .collect() + .await; + + let _ = shutdown_tx.send(()); + engine_task.await.unwrap(); + llm.shutdown().await.unwrap(); + + // The first output is yielded immediately, later ones once they carry at + // least 3 tokens, and the terminal one flushes the rest. + let shapes: Vec<_> = outputs + .iter() + .map(|output| { + ( + output.prompt_info.is_some(), + output.token_ids.clone(), + output.finish_reason.clone(), + ) + }) + .collect(); + assert_eq!( + shapes, + [ + (true, vec![1], None), + (false, vec![2, 3, 4, 5], None), + (false, vec![6, 7], Some(FinishReason::Length)), + ] + ); + for output in &outputs { + assert_eq!(output.logprobs, Some(token_logprobs(&output.token_ids))); + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn collect_output_aggregates_raw_tokens_logprobs_and_terminal_metadata() { init_tracing(); diff --git a/rust/src/mock-engine/src/tests.rs b/rust/src/mock-engine/src/tests.rs index d82ceac083..0b053cbbb1 100644 --- a/rust/src/mock-engine/src/tests.rs +++ b/rust/src/mock-engine/src/tests.rs @@ -116,7 +116,7 @@ async fn chunk_size_one_outputs_one_token_per_update() { let ipc = IpcNamespace::new().expect("ipc namespace"); let handshake_address = ipc.handshake_endpoint(); let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await; - let mut stream = client.call(sample_request("req-1", 3)).await.expect("call"); + let mut stream = client.call(sample_request("req-1", 3)).await.expect("call").into_outputs(); let first = stream.next().await.expect("first").expect("first ok"); assert_eq!(first.new_token_ids.len(), 1); @@ -147,7 +147,7 @@ async fn chunk_size_clips_final_output_to_max_tokens() { let ipc = IpcNamespace::new().expect("ipc namespace"); let handshake_address = ipc.handshake_endpoint(); let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 4).await; - let mut stream = client.call(sample_request("req-clip", 6)).await.expect("call"); + let mut stream = client.call(sample_request("req-clip", 6)).await.expect("call").into_outputs(); let first = stream.next().await.expect("first").expect("first ok"); assert_eq!(first.new_token_ids.len(), 4); @@ -165,7 +165,11 @@ async fn abort_cancels_active_request_and_emits_terminal_output() { let ipc = IpcNamespace::new().expect("ipc namespace"); let handshake_address = ipc.handshake_endpoint(); let (client, shutdown, task) = connect_with_mock(handshake_address, 1, 1).await; - let mut stream = client.call(sample_request("req-abort", 1_000_000)).await.expect("call"); + let mut stream = client + .call(sample_request("req-abort", 1_000_000)) + .await + .expect("call") + .into_outputs(); let first = stream.next().await.expect("first").expect("first ok"); assert_eq!(first.finish_reason, None); diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index 43192ab78f..a391f56cae 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::HashMap; +use std::num::NonZeroU32; use std::time::Duration; use anyhow::{Context, Result, bail}; @@ -81,6 +82,7 @@ async fn main() -> Result<()> { lora_modules: Vec::new(), chat_template_content_format: ChatTemplateContentFormatOption::Auto, max_logprobs: None, + stream_interval: NonZeroU32::MIN, api_server_options: ApiServerOptions::default(), cors: CorsConfig::default(), tls: None, diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index 8e51f1440b..3f95b5891c 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -3,6 +3,7 @@ use std::collections::HashMap; use std::fmt; +use std::num::NonZeroU32; use std::str::FromStr; use std::time::Duration; @@ -258,6 +259,10 @@ pub struct Config { /// Optional maximum number of top log probabilities accepted by the /// frontend. `None` delegates to the text layer default. pub max_logprobs: Option, + /// Minimum number of newly generated tokens batched into each streamed + /// output after the first one. Requests can raise it with their own + /// `stream_interval`. + pub stream_interval: NonZeroU32, /// HTTP/API-server behavior switches. pub api_server_options: ApiServerOptions, /// CORS settings applied to every HTTP response. diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index f54f8ce3e4..d5999cd3e7 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -123,7 +123,9 @@ async fn build_state(config: &Config) -> Result> { .await .context("failed to connect to engine core")?; - let llm = Llm::new(client).with_log_stats(!config.disable_log_stats); + let llm = Llm::new(client) + .with_log_stats(!config.disable_log_stats) + .with_stream_interval(config.stream_interval); let text = TextLlm::new(llm, text_backend).with_max_logprobs(config.max_logprobs); let chat = ChatLlm::new(text, chat_backend) diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 4f80dda6c7..46aaafbc69 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -179,6 +179,7 @@ pub(super) fn prepare_chat_request( merge_ec_transfer_params(request.vllm_xargs, request.ec_transfer_params.as_ref()), request.kv_transfer_params.as_ref(), ), + stream_interval: request.stream_interval, }, chat_options: ChatOptions { generation_prompt_mode, @@ -443,6 +444,7 @@ fn convert_tool_choice(tool_choice: &ToolChoice) -> Result(request).is_err()); + } + #[test] fn chat_http_request_defaults_missing_or_null_model() { for model in [None, Some(serde_json::Value::Null)] { @@ -1004,6 +1015,7 @@ mod tests { frequency_penalty: Some(0.3), presence_penalty: Some(0.4), repetition_penalty: Some(1.1), + stream_interval: NonZeroU32::new(4), ..base_request() }; @@ -1019,6 +1031,7 @@ mod tests { frequency_penalty: Some(0.3), presence_penalty: Some(0.4), repetition_penalty: Some(1.1), + stream_interval: NonZeroU32::new(4), ..VllmSamplingParams::default() }; assert_eq!(prepared.chat_request.sampling_params, expected); diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs index 13c7c7233f..204f3bab5a 100644 --- a/rust/src/server/src/routes/openai/chat_completions/types.rs +++ b/rust/src/server/src/routes/openai/chat_completions/types.rs @@ -3,6 +3,7 @@ use std::collections::HashMap; use std::fmt; +use std::num::NonZeroU32; use std::sync::Arc; use serde::{Deserialize, Serialize}; @@ -256,6 +257,12 @@ pub struct ChatCompletionRequest { /// Parameters for detecting repetitive N-gram patterns in output tokens pub repetition_detection: Option, + + /// Number of tokens to batch into each streamed chunk. Raises the server's + /// `--stream-interval` for this request. Values below the server setting + /// are clamped up to it. The first and last chunks are always sent + /// immediately. Ignored for non-streaming requests. + pub stream_interval: Option, } impl Default for ChatCompletionRequest { @@ -323,6 +330,7 @@ impl Default for ChatCompletionRequest { ec_transfer_params: None, vllm_xargs: None, repetition_detection: None, + stream_interval: None, } } } diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index 0b37f84057..33fc6bbe0a 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -144,6 +144,7 @@ pub(super) fn prepare_completion_request( merge_ec_transfer_params(request.vllm_xargs, request.ec_transfer_params.as_ref()), request.kv_transfer_params.as_ref(), ), + stream_interval: request.stream_interval, }, decode_options: TextDecodeOptions { skip_special_tokens: request.skip_special_tokens, @@ -211,6 +212,8 @@ fn completion_echo_text( #[cfg(test)] mod tests { + use std::num::NonZeroU32; + use axum::http::HeaderMap; use serde_json::json; use validator::Validate; @@ -400,7 +403,8 @@ mod tests { "presence_penalty": 0.3, "repetition_penalty": 1.1, "ignore_eos": true, - "skip_special_tokens": false + "skip_special_tokens": false, + "stream_interval": 4 })) .expect("parse request"); @@ -436,6 +440,10 @@ mod tests { ); assert!(prepared.text_request.sampling_params.ignore_eos); assert!(!prepared.text_request.decode_options.skip_special_tokens); + assert_eq!( + prepared.text_request.sampling_params.stream_interval, + NonZeroU32::new(4) + ); } #[test] diff --git a/rust/src/server/src/routes/openai/completions/types.rs b/rust/src/server/src/routes/openai/completions/types.rs index 2a69013706..b10f8a23d7 100644 --- a/rust/src/server/src/routes/openai/completions/types.rs +++ b/rust/src/server/src/routes/openai/completions/types.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::HashMap; +use std::num::NonZeroU32; use std::sync::Arc; use serde::{Deserialize, Serialize}; @@ -114,6 +115,12 @@ pub struct CompletionRequest { /// Parameters for detecting repetitive N-gram patterns in output tokens pub repetition_detection: Option, + /// Number of tokens to batch into each streamed chunk. Raises the server's + /// `--stream-interval` for this request. Values below the server setting + /// are clamped up to it. The first and last chunks are always sent + /// immediately. Ignored for non-streaming requests. + pub stream_interval: Option, + /// Length penalty for beam search pub length_penalty: Option, diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 2b09070049..9fa8c78f4f 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -124,6 +124,7 @@ pub fn lower_sampling_params( structured_outputs, skip_reading_prefix_cache, vllm_xargs, + stream_interval, } = sampling_params; validate_logprobs( @@ -214,6 +215,7 @@ pub fn lower_sampling_params( skip_reading_prefix_cache, extra_args: vllm_xargs, routed_experts_prompt_start: 0, + stream_interval, }; validate_resolved_sampling_params(¶ms)?; validate_vocab_range(¶ms, &sampling_limits)?; @@ -341,6 +343,7 @@ fn merge_unique_token_ids( #[cfg(test)] mod tests { use std::collections::{BTreeSet, HashMap}; + use std::num::NonZeroU32; use serial_test::file_serial; use vllm_engine_core_client::protocol::kv_hints::{KvHintAction, KvHintsEnvelope}; @@ -693,6 +696,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -746,6 +750,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -921,6 +926,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -992,6 +998,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -1056,6 +1063,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -1401,6 +1409,7 @@ mod tests { skip_reading_prefix_cache: None, extra_args: None, routed_experts_prompt_start: 0, + stream_interval: None, } "#]] .assert_debug_eq(¶ms); @@ -1430,6 +1439,26 @@ mod tests { assert_eq!(prepared.generate_request.request_id, "text-1"); } + #[test] + fn lower_text_request_carries_stream_interval_in_sampling_params() { + let mut request = sample_request(); + request.sampling_params.stream_interval = NonZeroU32::new(4); + + let prepared = lower_text_request( + request, + vec![1, 2, 3], + sample_sampling_hints(), + sample_sampling_limits(), + &stub_tokenizer(), + ) + .unwrap(); + + assert_eq!( + prepared.generate_request.sampling_params.stream_interval, + NonZeroU32::new(4) + ); + } + #[test] fn lower_text_request_passes_kv_hints_through() { let hints = KvHintsEnvelope { diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs index 08758748ac..b00d8f3b2f 100644 --- a/rust/src/text/src/request.rs +++ b/rust/src/text/src/request.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use std::collections::HashMap; +use std::num::NonZeroU32; use enum_as_inner::EnumAsInner; use serde::{Deserialize, Serialize}; @@ -147,6 +148,11 @@ pub struct SamplingParams { pub skip_reading_prefix_cache: Option, /// Additional request parameters for custom extensions. pub vllm_xargs: Option>, + /// Number of newly generated tokens to batch into each streamed output. + /// Raises the interval above the frontend-level `--stream-interval`; + /// values below it are clamped up to it. The first and final outputs are + /// always emitted immediately. + pub stream_interval: Option, } #[allow(clippy::derivable_impls)] // more explicit @@ -180,6 +186,7 @@ impl Default for SamplingParams { structured_outputs: None, skip_reading_prefix_cache: None, vllm_xargs: None, + stream_interval: None, } } }